diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..0e652a81 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,3 @@ +**/__pycache__ +**/*.pyc +**/.git diff --git a/verify_moe_e2e.py b/verify_moe_e2e.py new file mode 100644 index 00000000..74a8bc5a --- /dev/null +++ b/verify_moe_e2e.py @@ -0,0 +1,211 @@ +#!/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())