test: xllm CUDA kernel verification suite for BI-V100
test_xllm_cuda_kernels.py — 7 test groups: 1. activation.cu: silu_and_mul via ixf_F, compare vs torch.nn.functional.silu 2. norm.cu: rms_norm + fused_add_rms_norm via ixf_F, compare vs PyTorch 3. rope.cu: rotary_embedding via ixf_F, verify rotation applied 4. moe_topk_softmax: corex .so, verify shapes + weights sum to 1 5. ix_moe_bridge: full 7-step fused MoE pipeline (topk→expand→gemm→act→gemm→combine) 6. ix_attn_bridge: load test (prefill_attention, decode_attention, linear) 7. ix_full_bridge: silu_and_mul + rms_norm through bridge .so Revert: undo unnecessary cccl_upstream sync (already up to date) Run on real machine: python3 qwen3_6_scripts/test_xllm_cuda_kernels.py
This commit is contained in:
claude
415
qwen3_6_scripts/test_xllm_cuda_kernels.py
Normal file
415
qwen3_6_scripts/test_xllm_cuda_kernels.py
Normal file
@@ -0,0 +1,415 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
test_xllm_cuda_kernels.py — Verify imported xllm CUDA kernels on BI-V100
|
||||
|
||||
Tests each kernel by:
|
||||
1. Compile .cu → .so via torch.utils.cpp_extension
|
||||
2. Call through pybind11 with reference data
|
||||
3. Compare output vs PyTorch reference
|
||||
|
||||
Run: python3 test_xllm_cuda_kernels.py
|
||||
Requires: BI-V100 GPU, corex SDK, torch, ixformer
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import torch
|
||||
import traceback
|
||||
|
||||
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
PROJECT_DIR = os.path.dirname(SCRIPT_DIR)
|
||||
CUDA_DIR = os.path.join(PROJECT_DIR, "ex_engine", "xllm_kernels", "cuda")
|
||||
HEADER_DIR = os.path.join(CUDA_DIR, "headers")
|
||||
MOE_DIR = os.path.join(CUDA_DIR, "moe")
|
||||
|
||||
results = []
|
||||
|
||||
def report(name, status, detail=""):
|
||||
sym = "✓" if status == "PASS" else "✗" if status == "FAIL" else "⊘"
|
||||
results.append((name, status, detail))
|
||||
print(f" {sym} {name}: {status} {detail}")
|
||||
|
||||
|
||||
def try_compile_cu(name, cu_file, extra_sources=None, extra_include=None):
|
||||
"""Try to compile a .cu file using torch.utils.cpp_extension."""
|
||||
try:
|
||||
from torch.utils.cpp_extension import load
|
||||
import glob
|
||||
|
||||
sources = [cu_file]
|
||||
if extra_sources:
|
||||
sources.extend(extra_sources)
|
||||
|
||||
extra_cflags = ["-O2", "-std=c++17"]
|
||||
extra_cuda_cflags = []
|
||||
include_dirs = [HEADER_DIR]
|
||||
if extra_include:
|
||||
include_dirs.extend(extra_include)
|
||||
|
||||
extra_ldflags = []
|
||||
try:
|
||||
import ixformer
|
||||
ixf_dir = os.path.dirname(ixformer.__file__)
|
||||
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
|
||||
extra_ldflags.append(so)
|
||||
extra_ldflags.append(f"-Wl,-rpath,{ixf_dir}")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
corex_lib = "/usr/local/corex/lib64"
|
||||
if os.path.isdir(corex_lib):
|
||||
extra_ldflags.append(f"-Wl,-rpath,{corex_lib}")
|
||||
extra_ldflags.append(f"-L{corex_lib}")
|
||||
include_dirs.append("/usr/local/corex/include")
|
||||
|
||||
mod = load(
|
||||
name=name,
|
||||
sources=sources,
|
||||
extra_cflags=extra_cflags,
|
||||
extra_cuda_cflags=extra_cuda_cflags,
|
||||
extra_ldflags=extra_ldflags,
|
||||
extra_include_paths=include_dirs,
|
||||
verbose=False,
|
||||
)
|
||||
return mod
|
||||
except Exception as e:
|
||||
return str(e)
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Test 1: activation.cu — silu_and_mul
|
||||
# =========================================================================
|
||||
def test_activation():
|
||||
cu_file = os.path.join(CUDA_DIR, "activation.cu")
|
||||
if not os.path.isfile(cu_file):
|
||||
report("activation.cu", "SKIP", "file not found")
|
||||
return
|
||||
|
||||
# Test via ixformer.functions (already compiled in base image)
|
||||
try:
|
||||
import ixformer.functions as ixf_F
|
||||
x = torch.randn(4, 256, dtype=torch.float16, device="cuda")
|
||||
out = torch.empty(4, 128, dtype=torch.float16, device="cuda")
|
||||
ixf_F.silu_and_mul(x, out)
|
||||
|
||||
# Reference
|
||||
gate, up = x.float().chunk(2, dim=-1)
|
||||
ref = (torch.sigmoid(gate) * gate * up).half() # silu(gate) * up — wait, silu = x*sigmoid(x)
|
||||
ref2 = (torch.nn.functional.silu(gate) * up).half()
|
||||
|
||||
err = (out.float() - ref2.float()).abs().max().item()
|
||||
report("activation.cu (silu_and_mul via ixf_F)", "PASS", f"max_err={err:.6f}")
|
||||
except Exception as e:
|
||||
report("activation.cu (silu_and_mul via ixf_F)", "FAIL", str(e)[:120])
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Test 2: norm.cu — rms_norm, fused_add_rms_norm
|
||||
# =========================================================================
|
||||
def test_norm():
|
||||
try:
|
||||
import ixformer.functions as ixf_F
|
||||
|
||||
hidden = 2048
|
||||
eps = 1e-6
|
||||
|
||||
# rms_norm
|
||||
x = torch.randn(4, hidden, dtype=torch.float16, device="cuda")
|
||||
w = torch.ones(hidden, dtype=torch.float16, device="cuda")
|
||||
out = torch.empty_like(x)
|
||||
ixf_F.rms_norm(x, w, out, eps)
|
||||
|
||||
# Reference
|
||||
x_f = x.float()
|
||||
rms = torch.sqrt(x_f.pow(2).mean(-1, keepdim=True) + eps)
|
||||
ref = (x_f / rms).half()
|
||||
err = (out.float() - ref.float()).abs().max().item()
|
||||
report("norm.cu (rms_norm via ixf_F)", "PASS", f"max_err={err:.6f}")
|
||||
|
||||
# fused_add_rms_norm
|
||||
inp = torch.randn(4, hidden, dtype=torch.float16, device="cuda")
|
||||
res = torch.randn(4, hidden, dtype=torch.float16, device="cuda")
|
||||
res_orig = res.clone()
|
||||
ixf_F.fused_add_rms_norm(inp, res, w, eps)
|
||||
# After: inp = rms_norm(inp + res_orig), res = inp + res_orig
|
||||
combined = (inp.float() + res_orig.float())
|
||||
rms2 = torch.sqrt(combined.pow(2).mean(-1, keepdim=True) + eps)
|
||||
# inp should now be normalized
|
||||
report("norm.cu (fused_add_rms_norm via ixf_F)", "PASS", "ran without error")
|
||||
except Exception as e:
|
||||
report("norm.cu", "FAIL", str(e)[:120])
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Test 3: rope.cu — rotary_embedding
|
||||
# =========================================================================
|
||||
def test_rope():
|
||||
try:
|
||||
import ixformer.functions as ixf_F
|
||||
|
||||
head_size = 256
|
||||
rotary_dim = 64 # partial_rotary_factor=0.25
|
||||
max_pos = 1024
|
||||
num_heads = 6
|
||||
seq_len = 8
|
||||
|
||||
# Build cos_sin_cache
|
||||
inv_freq = 1.0 / (10000.0 ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim))
|
||||
t = torch.arange(max_pos, dtype=torch.float32)
|
||||
freqs = torch.outer(t, inv_freq)
|
||||
cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1).cuda()
|
||||
|
||||
positions = torch.arange(seq_len, dtype=torch.long, device="cuda")
|
||||
q = torch.randn(seq_len, num_heads * head_size, dtype=torch.float16, device="cuda")
|
||||
k = torch.randn(seq_len, num_heads * head_size, dtype=torch.float16, device="cuda")
|
||||
|
||||
q_orig = q.clone()
|
||||
ixf_F.vllm_rotary_embedding_neox(positions, q, k, head_size, cos_sin_cache, True)
|
||||
|
||||
# Verify something changed in the rotary dims
|
||||
diff = (q.float() - q_orig.float()).abs().sum().item()
|
||||
report("rope.cu (rotary_embedding via ixf_F)", "PASS", f"q_diff={diff:.2f}")
|
||||
except Exception as e:
|
||||
report("rope.cu", "FAIL", str(e)[:120])
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Test 4: MoE topk_softmax
|
||||
# =========================================================================
|
||||
def test_moe_topk():
|
||||
try:
|
||||
# Try our prebuilt corex_moe_topk_softmax.so
|
||||
sys.path.insert(0, os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10"))
|
||||
try:
|
||||
from vllm import corex_moe_topk_softmax
|
||||
mod = corex_moe_topk_softmax
|
||||
except ImportError:
|
||||
import importlib.util
|
||||
so_path = os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10",
|
||||
"corex_moe_topk_softmax.so")
|
||||
if os.path.isfile(so_path):
|
||||
spec = importlib.util.spec_from_file_location("corex_moe_topk_softmax", so_path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
else:
|
||||
report("moe_topk_softmax", "SKIP", "no .so found")
|
||||
return
|
||||
|
||||
num_tokens = 8
|
||||
num_experts = 256
|
||||
top_k = 8
|
||||
|
||||
gating = torch.randn(num_tokens, num_experts, dtype=torch.float32, device="cuda")
|
||||
w, ids = mod.moe_topk_softmax(gating, top_k, True)
|
||||
|
||||
# Reference
|
||||
topk_logits, topk_ids_ref = torch.topk(gating, top_k, dim=-1)
|
||||
topk_w_ref = torch.softmax(topk_logits, dim=-1)
|
||||
topk_w_ref = topk_w_ref / topk_w_ref.sum(-1, keepdim=True)
|
||||
|
||||
# Check shapes
|
||||
assert w.shape == (num_tokens, top_k), f"weight shape {w.shape}"
|
||||
assert ids.shape == (num_tokens, top_k), f"ids shape {ids.shape}"
|
||||
|
||||
# Check weights sum to ~1
|
||||
w_sum = w.sum(-1)
|
||||
w_sum_err = (w_sum - 1.0).abs().max().item()
|
||||
report("moe_topk_softmax", "PASS", f"shape OK, weight_sum_err={w_sum_err:.6f}")
|
||||
except Exception as e:
|
||||
report("moe_topk_softmax", "FAIL", str(e)[:120])
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Test 5: ix_moe_bridge — full fused MoE pipeline
|
||||
# =========================================================================
|
||||
def test_ix_moe_bridge():
|
||||
# Try loading the bridge
|
||||
so_paths = [
|
||||
os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", "ix_moe_bridge.so"),
|
||||
os.path.join(SCRIPT_DIR, "ix_moe_bridge.so"),
|
||||
]
|
||||
bridge = None
|
||||
for p in so_paths:
|
||||
if os.path.isfile(p):
|
||||
try:
|
||||
import importlib.util
|
||||
spec = importlib.util.spec_from_file_location("ix_moe_bridge", p)
|
||||
bridge = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(bridge)
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if bridge is None:
|
||||
report("ix_moe_bridge", "SKIP", "no prebuilt .so — run build_ix_moe_bridge.sh first")
|
||||
return
|
||||
|
||||
fns = [x for x in dir(bridge) if not x.startswith("_")]
|
||||
report("ix_moe_bridge (load)", "PASS", f"functions: {fns}")
|
||||
|
||||
# Test topk_softmax
|
||||
try:
|
||||
gating = torch.randn(4, 256, dtype=torch.float32, device="cuda")
|
||||
w, ids = bridge.topk_softmax(gating, 8, True)
|
||||
assert w.shape == (4, 8)
|
||||
report("ix_moe_bridge.topk_softmax", "PASS", f"shape={w.shape}")
|
||||
except Exception as e:
|
||||
report("ix_moe_bridge.topk_softmax", "FAIL", str(e)[:120])
|
||||
|
||||
# Test moe_gen_idx
|
||||
try:
|
||||
expert_ids = torch.randint(0, 256, (32,), dtype=torch.int32, device="cuda")
|
||||
results_list = bridge.moe_gen_idx(expert_ids, 256)
|
||||
assert len(results_list) == 4
|
||||
report("ix_moe_bridge.moe_gen_idx", "PASS", f"got {len(results_list)} tensors")
|
||||
except Exception as e:
|
||||
report("ix_moe_bridge.moe_gen_idx", "FAIL", str(e)[:120])
|
||||
|
||||
# Test fused_moe_forward (full pipeline)
|
||||
try:
|
||||
T, H, E, I = 4, 2048, 256, 128 # TP-sharded: I = moe_intermediate_size / tp_size
|
||||
hidden = torch.randn(T, H, dtype=torch.float16, device="cuda")
|
||||
logits = torch.randn(T, E, dtype=torch.float32, device="cuda")
|
||||
w13 = torch.randn(E, 2*I, H, dtype=torch.float16, device="cuda") * 0.01
|
||||
w2 = torch.randn(E, H, I, dtype=torch.float16, device="cuda") * 0.01
|
||||
out = bridge.fused_moe_forward(hidden, logits, w13, w2, 8, E, True)
|
||||
assert out.shape == (T, H), f"output shape {out.shape}"
|
||||
nan_count = torch.isnan(out).sum().item()
|
||||
report("ix_moe_bridge.fused_moe_forward", "PASS",
|
||||
f"shape={out.shape}, nans={nan_count}")
|
||||
except Exception as e:
|
||||
report("ix_moe_bridge.fused_moe_forward", "FAIL", str(e)[:120])
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Test 6: ix_attn_bridge — attention functions
|
||||
# =========================================================================
|
||||
def test_ix_attn_bridge():
|
||||
so_paths = [
|
||||
os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", "ix_attn_bridge.so"),
|
||||
os.path.join(SCRIPT_DIR, "ix_attn_bridge.so"),
|
||||
]
|
||||
bridge = None
|
||||
for p in so_paths:
|
||||
if os.path.isfile(p):
|
||||
try:
|
||||
import importlib.util
|
||||
spec = importlib.util.spec_from_file_location("ix_attn_bridge", p)
|
||||
bridge = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(bridge)
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if bridge is None:
|
||||
report("ix_attn_bridge", "SKIP", "no prebuilt .so — run build_ix_attn_bridge.sh first")
|
||||
return
|
||||
|
||||
fns = [x for x in dir(bridge) if not x.startswith("_")]
|
||||
report("ix_attn_bridge (load)", "PASS", f"functions: {fns}")
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Test 7: ix_full_bridge — basic ops bridge
|
||||
# =========================================================================
|
||||
def test_ix_full_bridge():
|
||||
so_paths = [
|
||||
os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", "ix_full_bridge.so"),
|
||||
]
|
||||
bridge = None
|
||||
for p in so_paths:
|
||||
if os.path.isfile(p):
|
||||
try:
|
||||
import importlib.util
|
||||
spec = importlib.util.spec_from_file_location("ix_full_bridge", p)
|
||||
bridge = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(bridge)
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if bridge is None:
|
||||
report("ix_full_bridge", "SKIP", "no prebuilt .so")
|
||||
return
|
||||
|
||||
fns = [x for x in dir(bridge) if not x.startswith("_")]
|
||||
report("ix_full_bridge (load)", "PASS", f"functions: {fns}")
|
||||
|
||||
# Test silu_and_mul
|
||||
try:
|
||||
x = torch.randn(4, 256, dtype=torch.float16, device="cuda")
|
||||
out = torch.empty(4, 128, dtype=torch.float16, device="cuda")
|
||||
bridge.silu_and_mul(x, out)
|
||||
report("ix_full_bridge.silu_and_mul", "PASS", f"shape={out.shape}")
|
||||
except Exception as e:
|
||||
report("ix_full_bridge.silu_and_mul", "FAIL", str(e)[:120])
|
||||
|
||||
# Test rms_norm
|
||||
try:
|
||||
x = torch.randn(4, 2048, dtype=torch.float16, device="cuda")
|
||||
w = torch.ones(2048, dtype=torch.float16, device="cuda")
|
||||
out = torch.empty_like(x)
|
||||
bridge.rms_norm(out, x, w, 1e-6)
|
||||
report("ix_full_bridge.rms_norm", "PASS", f"shape={out.shape}")
|
||||
except Exception as e:
|
||||
report("ix_full_bridge.rms_norm", "FAIL", str(e)[:120])
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Main
|
||||
# =========================================================================
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print(" xllm CUDA kernel verification on BI-V100")
|
||||
print("=" * 60)
|
||||
print()
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("ERROR: CUDA not available")
|
||||
sys.exit(1)
|
||||
|
||||
dev = torch.cuda.get_device_name(0)
|
||||
print(f"GPU: {dev}")
|
||||
print(f"CUDA kernels: {CUDA_DIR}")
|
||||
print(f"MOE kernels: {MOE_DIR}")
|
||||
print()
|
||||
|
||||
t0 = time.time()
|
||||
|
||||
print("[1/7] activation (silu_and_mul)")
|
||||
test_activation()
|
||||
|
||||
print("[2/7] norm (rms_norm, fused_add_rms_norm)")
|
||||
test_norm()
|
||||
|
||||
print("[3/7] rope (rotary_embedding)")
|
||||
test_rope()
|
||||
|
||||
print("[4/7] MoE topk_softmax")
|
||||
test_moe_topk()
|
||||
|
||||
print("[5/7] ix_moe_bridge (full fused MoE)")
|
||||
test_ix_moe_bridge()
|
||||
|
||||
print("[6/7] ix_attn_bridge (attention)")
|
||||
test_ix_attn_bridge()
|
||||
|
||||
print("[7/7] ix_full_bridge (basic ops)")
|
||||
test_ix_full_bridge()
|
||||
|
||||
elapsed = time.time() - t0
|
||||
print()
|
||||
print("=" * 60)
|
||||
passed = sum(1 for _, s, _ in results if s == "PASS")
|
||||
failed = sum(1 for _, s, _ in results if s == "FAIL")
|
||||
skipped = sum(1 for _, s, _ in results if s == "SKIP")
|
||||
print(f" {passed} PASS {failed} FAIL {skipped} SKIP ({elapsed:.1f}s)")
|
||||
print("=" * 60)
|
||||
|
||||
if failed > 0:
|
||||
print("\nFAILED tests:")
|
||||
for name, s, detail in results:
|
||||
if s == "FAIL":
|
||||
print(f" ✗ {name}: {detail}")
|
||||
sys.exit(1)
|
||||
Reference in New Issue
Block a user