dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity
Files
900ae0b1efc19e46cd313eb3ba7eb694fa596f6b
project_6/qwen3_6_scripts/test_xllm_cuda_kernels.py
claude 3a2cfc87c9 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
2026-08-14 08:01:23 +00:00

416 lines
15 KiB
Python
Raw Blame History

#!/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 View Git Blame Copy Permalink