fix: add .dockerignore (exclude __pycache__) + e2e MoE verification script

.dockerignore

@@ -0,0 +1,3 @@
+**/__pycache__
+**/*.pyc
+**/.git

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())