212 lines
7.2 KiB
Python
212 lines
7.2 KiB
Python
#!/usr/bin/env python3
|
|
"""End-to-end MoE forward path verification on single BI-V100.
|
|
|
|
Simulates Qwen3.5 MoE dimensions:
|
|
hidden_size=2048, num_experts=256, top_k=8, intermediate=128
|
|
w13: (256, 256, 2048), w2: (256, 2048, 128)
|
|
|
|
Tests the full chain:
|
|
1. topk_softmax kernel (router_logits → topk_weights, topk_ids)
|
|
2. moe_compute_index kernel (topk_ids → sorted order)
|
|
3. Per-expert GEMM (F.linear through sorted experts)
|
|
4. Weighted combine (output)
|
|
|
|
Compares kernel-accelerated path vs pure PyTorch path.
|
|
|
|
Run: python3 verify_moe_e2e.py
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import time
|
|
import importlib.util
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def load_so(name, so_path):
|
|
if not os.path.exists(so_path):
|
|
return None
|
|
try:
|
|
spec = importlib.util.spec_from_file_location(name, so_path)
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
return mod
|
|
except Exception as e:
|
|
print(f"[WARN] Failed to load {so_path}: {e}")
|
|
return None
|
|
|
|
|
|
def pure_pytorch_moe(hidden_states, router_logits, w13, w2, top_k):
|
|
"""Exact copy of qwen3_5.py _pure_pytorch_experts prefill path."""
|
|
T = hidden_states.shape[0]
|
|
topk_logits, topk_ids = torch.topk(router_logits.float(), top_k, dim=-1)
|
|
topk_weights = torch.softmax(topk_logits, dim=-1).to(hidden_states.dtype)
|
|
|
|
out = torch.zeros_like(hidden_states)
|
|
flat_eids = topk_ids.reshape(-1)
|
|
order = torch.argsort(flat_eids, stable=True)
|
|
sorted_tok_ids = torch.arange(
|
|
T, device=topk_ids.device).repeat_interleave(top_k)[order]
|
|
sorted_weights = topk_weights.reshape(-1)[order]
|
|
expert_counts = torch.bincount(flat_eids, minlength=w13.shape[0]).tolist()
|
|
|
|
start = 0
|
|
for eid, count in enumerate(expert_counts):
|
|
end = start + count
|
|
if count == 0:
|
|
start = end
|
|
continue
|
|
tok_ids = sorted_tok_ids[start:end]
|
|
tokens = hidden_states[tok_ids]
|
|
gate_up = F.linear(tokens, w13[eid])
|
|
gate, up = gate_up.chunk(2, dim=-1)
|
|
act = F.silu(gate) * up
|
|
expert_out = F.linear(act, w2[eid])
|
|
weights = sorted_weights[start:end].unsqueeze(-1)
|
|
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype))
|
|
start = end
|
|
return out
|
|
|
|
|
|
def kernel_moe(hidden_states, router_logits, w13, w2, top_k,
|
|
topk_mod, index_mod):
|
|
"""Kernel-accelerated MoE path."""
|
|
T = hidden_states.shape[0]
|
|
|
|
# Step 1: topk_softmax kernel
|
|
topk_weights, topk_ids = topk_mod.moe_topk_softmax(
|
|
router_logits.float(), top_k, True)
|
|
topk_ids = topk_ids.to(torch.int64)
|
|
topk_weights = topk_weights.to(hidden_states.dtype)
|
|
|
|
# Step 2: moe_compute_index kernel
|
|
flat_eids = topk_ids.reshape(-1)
|
|
src_dst, dst_src, expert_sizes = index_mod.moe_compute_index(
|
|
flat_eids, w13.shape[0])
|
|
sorted_tok_ids = torch.arange(
|
|
T, device=topk_ids.device).repeat_interleave(top_k)[dst_src.long()]
|
|
sorted_weights = topk_weights.reshape(-1)[dst_src.long()]
|
|
expert_counts = expert_sizes.tolist()
|
|
|
|
# Step 3: Per-expert GEMM (same as PyTorch — this is the bottleneck)
|
|
out = torch.zeros_like(hidden_states)
|
|
start = 0
|
|
for eid, count in enumerate(expert_counts):
|
|
end = start + count
|
|
if count == 0:
|
|
start = end
|
|
continue
|
|
tok_ids = sorted_tok_ids[start:end]
|
|
tokens = hidden_states[tok_ids]
|
|
gate_up = F.linear(tokens, w13[eid])
|
|
gate, up = gate_up.chunk(2, dim=-1)
|
|
act = F.silu(gate) * up
|
|
expert_out = F.linear(act, w2[eid])
|
|
weights = sorted_weights[start:end].unsqueeze(-1)
|
|
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype))
|
|
start = end
|
|
return out
|
|
|
|
|
|
def main():
|
|
print("=" * 60)
|
|
print("BI-V100 MoE end-to-end verification")
|
|
print("=" * 60)
|
|
|
|
if not torch.cuda.is_available():
|
|
print("FATAL: No CUDA device")
|
|
return 1
|
|
|
|
# Load kernels
|
|
prebuilt = os.path.join(os.path.dirname(os.path.abspath(__file__)),
|
|
"qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10")
|
|
topk_mod = load_so("corex_moe_topk_softmax",
|
|
os.path.join(prebuilt, "corex_moe_topk_softmax.so"))
|
|
index_mod = load_so("corex_moe_index_combine",
|
|
"/tmp/moe_test/corex_moe_index_combine.so")
|
|
|
|
if topk_mod is None:
|
|
print("[FAIL] Cannot load topk_softmax .so")
|
|
return 1
|
|
if index_mod is None:
|
|
print("[FAIL] Cannot load index_combine .so — run verify_moe_index_combine.py first")
|
|
return 1
|
|
|
|
print(f"[OK] Both kernel modules loaded")
|
|
|
|
# Qwen3.5 MoE dimensions (TP=4 sharded)
|
|
hidden_size = 2048
|
|
num_experts = 256
|
|
top_k = 8
|
|
inter_per_partition = 128 # moe_intermediate_size / tp_size
|
|
|
|
torch.manual_seed(42)
|
|
|
|
# --- Test 1: Single token (decode) ---
|
|
print("\n--- Test 1: 1 token (decode path) ---")
|
|
h = torch.randn(1, hidden_size, device="cuda", dtype=torch.float16)
|
|
router = torch.randn(1, num_experts, device="cuda", dtype=torch.float16)
|
|
w13 = torch.randn(num_experts, 2 * inter_per_partition, hidden_size,
|
|
device="cuda", dtype=torch.float16) * 0.01
|
|
w2 = torch.randn(num_experts, hidden_size, inter_per_partition,
|
|
device="cuda", dtype=torch.float16) * 0.01
|
|
|
|
ref_out = pure_pytorch_moe(h, router, w13, w2, top_k)
|
|
kern_out = kernel_moe(h, router, w13, w2, top_k, topk_mod, index_mod)
|
|
|
|
diff = (ref_out.float() - kern_out.float()).abs().max().item()
|
|
print(f" Max diff: {diff:.8f}")
|
|
print(f" Match: {diff < 0.01}")
|
|
|
|
# --- Test 2: 32 tokens (prefill) ---
|
|
print("\n--- Test 2: 32 tokens (prefill path) ---")
|
|
h = torch.randn(32, hidden_size, device="cuda", dtype=torch.float16)
|
|
router = torch.randn(32, num_experts, device="cuda", dtype=torch.float16)
|
|
|
|
ref_out = pure_pytorch_moe(h, router, w13, w2, top_k)
|
|
kern_out = kernel_moe(h, router, w13, w2, top_k, topk_mod, index_mod)
|
|
|
|
diff = (ref_out.float() - kern_out.float()).abs().max().item()
|
|
rel_diff = diff / (ref_out.float().abs().max().item() + 1e-8)
|
|
print(f" Max abs diff: {diff:.8f}")
|
|
print(f" Relative diff: {rel_diff:.8f}")
|
|
print(f" Match: {rel_diff < 0.01}")
|
|
|
|
# --- Test 3: Performance comparison ---
|
|
print("\n--- Performance: 32 tokens prefill ---")
|
|
h = torch.randn(32, hidden_size, device="cuda", dtype=torch.float16)
|
|
router = torch.randn(32, num_experts, device="cuda", dtype=torch.float16)
|
|
|
|
# Warmup
|
|
for _ in range(5):
|
|
pure_pytorch_moe(h, router, w13, w2, top_k)
|
|
kernel_moe(h, router, w13, w2, top_k, topk_mod, index_mod)
|
|
torch.cuda.synchronize()
|
|
|
|
N = 20
|
|
torch.cuda.synchronize()
|
|
t0 = time.perf_counter()
|
|
for _ in range(N):
|
|
pure_pytorch_moe(h, router, w13, w2, top_k)
|
|
torch.cuda.synchronize()
|
|
pt_ms = (time.perf_counter() - t0) / N * 1000
|
|
|
|
torch.cuda.synchronize()
|
|
t0 = time.perf_counter()
|
|
for _ in range(N):
|
|
kernel_moe(h, router, w13, w2, top_k, topk_mod, index_mod)
|
|
torch.cuda.synchronize()
|
|
kern_ms = (time.perf_counter() - t0) / N * 1000
|
|
|
|
print(f" PyTorch: {pt_ms:.1f} ms")
|
|
print(f" Kernel: {kern_ms:.1f} ms")
|
|
print(f" Speedup: {pt_ms/kern_ms:.2f}x")
|
|
|
|
print("\n" + "=" * 60)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|