under test, not sure no errors

verify_moe_index_combine.py

@@ -0,0 +1,186 @@
+#!/usr/bin/env python3
+"""Verify moe_compute_index + moe_combine_result on real BI-V100.
+
+Step 1: Compile corex_moe_index_combine.cu → .so
+Step 2: Test moe_compute_index vs PyTorch argsort+bincount
+Step 3: Test moe_combine_result vs PyTorch weighted sum
+Step 4: End-to-end MoE prefill path benchmark
+
+Run: python3 verify_moe_index_combine.py
+"""
+
+import sys
+import os
+import time
+import torch
+import torch.nn.functional as F
+
+def compile_kernel():
+    """Compile the .so using corex clang++."""
+    script_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                              "qwen3_6_scripts")
+    build_sh = os.path.join(script_dir, "build_corex_moe_index_combine.sh")
+    # Use a temp vllm root for testing
+    tmp_root = "/tmp/moe_test"
+    os.makedirs(tmp_root, exist_ok=True)
+    ret = os.system(f"bash {build_sh} {tmp_root} 2>&1")
+    so_path = os.path.join(tmp_root, "corex_moe_index_combine.so")
+    if ret != 0 or not os.path.exists(so_path):
+        print(f"[FAIL] Compilation failed (exit={ret})")
+        return None
+    print(f"[OK] Compiled: {so_path}")
+
+    import importlib.util
+    spec = importlib.util.spec_from_file_location(
+        "corex_moe_index_combine", so_path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def pytorch_compute_index(expert_ids_flat, num_experts):
+    """Reference: what qwen3_5.py does in prefill path."""
+    order = torch.argsort(expert_ids_flat, stable=True)
+    expert_counts = torch.bincount(
+        expert_ids_flat, minlength=num_experts)
+    # dst_src[i] = which flat_idx goes to position i (sorted order)
+    dst_src = torch.arange(len(expert_ids_flat),
+                           device=expert_ids_flat.device)[order]
+    # src_dst[flat_idx] = position in sorted order
+    src_dst = torch.empty_like(order)
+    src_dst[order] = torch.arange(len(order), device=order.device)
+    return src_dst, dst_src, expert_counts
+
+
+def pytorch_combine(expert_outputs, weights, topk, num_tokens, H):
+    """Reference: weighted sum of expert outputs."""
+    # expert_outputs: (N*topk, H), weights: (N, topk)
+    out = expert_outputs.view(num_tokens, topk, H)
+    w = weights.unsqueeze(-1)  # (N, topk, 1)
+    return (out * w).sum(dim=1)  # (N, H)
+
+
+def main():
+    print("=" * 60)
+    print("BI-V100 moe_compute_index + moe_combine verification")
+    print("=" * 60)
+
+    if not torch.cuda.is_available():
+        print("FATAL: No CUDA device")
+        return 1
+
+    mod = compile_kernel()
+    if mod is None:
+        return 1
+
+    # ---- Test 1: moe_compute_index ----
+    print("\n--- Test 1: moe_compute_index (256 experts, 32 tokens, top_k=8) ---")
+    num_tokens = 32
+    num_experts = 256
+    topk = 8
+    torch.manual_seed(42)
+    # Simulate topk routing: each token picks 8 experts
+    topk_ids = torch.randint(0, num_experts, (num_tokens, topk),
+                             device="cuda", dtype=torch.int64)
+    flat_ids = topk_ids.reshape(-1)  # (256,)
+
+    # Kernel
+    kern_src_dst, kern_dst_src, kern_sizes = mod.moe_compute_index(
+        flat_ids, num_experts)
+
+    # PyTorch reference
+    ref_src_dst, ref_dst_src, ref_sizes = pytorch_compute_index(
+        flat_ids, num_experts)
+
+    # Compare sizes (must match exactly)
+    sizes_match = torch.equal(kern_sizes.cpu(), ref_sizes.cpu().to(torch.int32))
+    print(f"  Expert sizes match: {sizes_match}")
+
+    # Compare mappings: verify kern_dst_src is a valid permutation
+    # that groups tokens by expert
+    kern_sorted_eids = flat_ids[kern_dst_src.long()]
+    ref_sorted_eids = flat_ids[ref_dst_src.long()]
+    # Both should be sorted by expert
+    kern_sorted = torch.all(kern_sorted_eids[:-1] <= kern_sorted_eids[1:]).item()
+    ref_sorted = torch.all(ref_sorted_eids[:-1] <= ref_sorted_eids[1:]).item()
+    print(f"  Kernel produces sorted expert order: {kern_sorted}")
+    print(f"  Ref produces sorted expert order: {ref_sorted}")
+
+    # ---- Test 2: moe_combine_result ----
+    print("\n--- Test 2: moe_combine_result (32 tokens, top_k=8, H=2048) ---")
+    H = 2048
+    expert_outputs = torch.randn(num_tokens * topk, H,
+                                 device="cuda", dtype=torch.float16)
+    weights = torch.rand(num_tokens, topk,
+                         device="cuda", dtype=torch.float32)
+    weights = weights / weights.sum(dim=-1, keepdim=True)  # normalize
+
+    kern_out = mod.moe_combine_result(expert_outputs, weights, num_tokens, topk)
+    ref_out = pytorch_combine(expert_outputs, weights, topk, num_tokens, H)
+
+    max_diff = (kern_out.float() - ref_out.float()).abs().max().item()
+    print(f"  Max diff: {max_diff:.8f}")
+    print(f"  Match (tol=1e-3): {max_diff < 1e-3}")
+
+    # ---- Test 3: Performance ----
+    print("\n--- Performance: moe_compute_index ---")
+    flat_ids = torch.randint(0, 256, (256,), device="cuda", dtype=torch.int64)
+
+    # Warmup
+    for _ in range(10):
+        mod.moe_compute_index(flat_ids, 256)
+        pytorch_compute_index(flat_ids, 256)
+    torch.cuda.synchronize()
+
+    N = 200
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(N):
+        mod.moe_compute_index(flat_ids, 256)
+    torch.cuda.synchronize()
+    kern_ms = (time.perf_counter() - t0) / N * 1000
+
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(N):
+        pytorch_compute_index(flat_ids, 256)
+    torch.cuda.synchronize()
+    pt_ms = (time.perf_counter() - t0) / N * 1000
+
+    print(f"  Kernel: {kern_ms:.3f} ms")
+    print(f"  PyTorch: {pt_ms:.3f} ms")
+    print(f"  Speedup: {pt_ms/kern_ms:.2f}x")
+
+    print("\n--- Performance: moe_combine_result ---")
+    expert_outputs = torch.randn(32 * 8, 2048, device="cuda", dtype=torch.float16)
+    weights = torch.rand(32, 8, device="cuda", dtype=torch.float32)
+
+    for _ in range(10):
+        mod.moe_combine_result(expert_outputs, weights, 32, 8)
+        pytorch_combine(expert_outputs, weights, 8, 32, 2048)
+    torch.cuda.synchronize()
+
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(N):
+        mod.moe_combine_result(expert_outputs, weights, 32, 8)
+    torch.cuda.synchronize()
+    kern_ms = (time.perf_counter() - t0) / N * 1000
+
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(N):
+        pytorch_combine(expert_outputs, weights, 8, 32, 2048)
+    torch.cuda.synchronize()
+    pt_ms = (time.perf_counter() - t0) / N * 1000
+
+    print(f"  Kernel: {kern_ms:.3f} ms")
+    print(f"  PyTorch: {pt_ms:.3f} ms")
+    print(f"  Speedup: {pt_ms/kern_ms:.2f}x")
+
+    print("\n" + "=" * 60)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())