"""test_moe_bridge.py — Integration test for ix_moe_bridge on real device. Run after build_moe_bridge.sh. No model weights needed — uses random tensors. Tests each of the 5 MoE functions + the fused pipeline. Usage: python3 test_moe_bridge.py """ import sys import os import torch import time # Qwen3.5-27B MoE params NUM_EXPERTS = 128 TOPK = 8 HIDDEN_SIZE = 3584 INTERMEDIATE_SIZE = 18944 # per-partition (full=18944*2 for gate+up, /TP if sharded) NUM_TOKENS = 4 def load_bridge(): """Try to load ix_moe_bridge.""" # Try prebuilt script_dir = os.path.dirname(os.path.abspath(__file__)) for p in [ os.path.join(script_dir, "prebuilt", "ix_moe_bridge.so"), os.path.join(script_dir, "ix_moe_bridge.so"), ]: if os.path.isfile(p): import importlib.util spec = importlib.util.spec_from_file_location("ix_moe_bridge", p) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod # Try import import ix_moe_bridge return ix_moe_bridge def test_topk_softmax(bridge, device): print("\n--- topk_softmax ---") gating = torch.randn(NUM_TOKENS, NUM_EXPERTS, device=device, dtype=torch.float32) topk_w, topk_ids, token_expert_ids = bridge.topk_softmax(gating, TOPK, True) assert topk_w.shape == (NUM_TOKENS, TOPK), f"weights shape: {topk_w.shape}" assert topk_ids.shape == (NUM_TOKENS, TOPK), f"ids shape: {topk_ids.shape}" assert topk_w.dtype == torch.float32 assert topk_ids.dtype == torch.int32 assert (topk_ids >= 0).all() and (topk_ids < NUM_EXPERTS).all(), "ids out of range" assert torch.allclose(topk_w.sum(-1), torch.ones(NUM_TOKENS, device=device), atol=1e-5), \ f"weights don't sum to 1: {topk_w.sum(-1)}" print(f" ✓ shape={topk_w.shape}, sum={topk_w.sum(-1).tolist()}") print(f" ✓ top expert ids (row 0): {topk_ids[0].tolist()}") def test_moe_gen_idx(bridge, device): print("\n--- moe_gen_idx ---") expert_ids = torch.randint(0, NUM_EXPERTS, (NUM_TOKENS * TOPK,), device=device, dtype=torch.int32) results = bridge.moe_gen_idx(expert_ids, NUM_EXPERTS) src_dst, dst_src, expert_sizes, expert_cumsum = results assert src_dst.shape == (NUM_TOKENS * TOPK,), f"src_dst shape: {src_dst.shape}" assert dst_src.shape == (NUM_TOKENS * TOPK,), f"dst_src shape: {dst_src.shape}" assert expert_sizes.shape[0] == NUM_EXPERTS, f"expert_sizes shape: {expert_sizes.shape}" assert expert_sizes.sum().item() == NUM_TOKENS * TOPK, \ f"expert_sizes sum: {expert_sizes.sum().item()} != {NUM_TOKENS * TOPK}" print(f" ✓ src_dst={src_dst.shape}, expert_sizes sum={expert_sizes.sum().item()}") def test_moe_expand_input(bridge, device): print("\n--- moe_expand_input ---") hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, device=device, dtype=torch.float16) # Create simple gather index: [0,1,2,...,NUM_TOKENS*TOPK-1] mod NUM_TOKENS gather_idx = torch.arange(NUM_TOKENS * TOPK, device=device, dtype=torch.int32) % NUM_TOKENS combine_idx = torch.arange(NUM_TOKENS * TOPK, device=device, dtype=torch.int32) expanded = bridge.moe_expand_input(hidden, gather_idx, combine_idx, TOPK) assert expanded.shape == (NUM_TOKENS * TOPK, HIDDEN_SIZE), f"shape: {expanded.shape}" print(f" ✓ shape={expanded.shape}, dtype={expanded.dtype}") def test_group_gemm(bridge, device): print("\n--- group_gemm ---") # Simulate: expanded tokens × expert weights total_tokens = NUM_TOKENS * TOPK # 32 inputs = torch.randn(total_tokens, HIDDEN_SIZE, device=device, dtype=torch.float16) # weights: [NUM_EXPERTS, 2*INTERMEDIATE, HIDDEN] — 3D weights = torch.randn(NUM_EXPERTS, INTERMEDIATE_SIZE * 2, HIDDEN_SIZE, device=device, dtype=torch.float16) * 0.01 # tokens_per_expert: distribute evenly tpe = torch.zeros(NUM_EXPERTS, device=device, dtype=torch.int32) for i in range(total_tokens): tpe[i % NUM_EXPERTS] += 1 output_n = INTERMEDIATE_SIZE * 2 result = bridge.group_gemm(inputs, weights, tpe, output_n) assert result.shape == (total_tokens, output_n), f"shape: {result.shape}" assert not torch.isnan(result).any(), "NaN in group_gemm output" print(f" ✓ shape={result.shape}, max={result.abs().max().item():.4f}") def test_silu_and_mul(bridge, device): print("\n--- silu_and_mul ---") gate_up = torch.randn(NUM_TOKENS, INTERMEDIATE_SIZE * 2, device=device, dtype=torch.float16) activated = bridge.silu_and_mul(gate_up) assert activated.shape == (NUM_TOKENS, INTERMEDIATE_SIZE), f"shape: {activated.shape}" print(f" ✓ shape={activated.shape}") def test_moe_combine_result(bridge, device): print("\n--- moe_combine_result ---") expert_out = torch.randn(NUM_TOKENS * TOPK, HIDDEN_SIZE, device=device, dtype=torch.float16) weights = torch.randn(NUM_TOKENS, TOPK, device=device, dtype=torch.float32) weights = torch.softmax(weights, dim=-1) combined = bridge.moe_combine_result(expert_out, weights) assert combined.shape == (NUM_TOKENS, HIDDEN_SIZE), f"shape: {combined.shape}" assert not torch.isnan(combined).any(), "NaN in combine output" print(f" ✓ shape={combined.shape}") def test_fused_pipeline(bridge, device): print("\n--- fused_moe_forward (7-step pipeline) ---") hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, device=device, dtype=torch.float16) router = torch.randn(NUM_TOKENS, NUM_EXPERTS, device=device, dtype=torch.float16) w13 = torch.randn(NUM_EXPERTS, INTERMEDIATE_SIZE * 2, HIDDEN_SIZE, device=device, dtype=torch.float16) * 0.01 w2 = torch.randn(NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE, device=device, dtype=torch.float16) * 0.01 t0 = time.time() output = bridge.fused_moe_forward(hidden, router, w13, w2, TOPK, NUM_EXPERTS, True) torch.cuda.synchronize() elapsed = time.time() - t0 assert output.shape == (NUM_TOKENS, HIDDEN_SIZE), f"shape: {output.shape}" assert not torch.isnan(output).any(), "NaN in fused output" print(f" ✓ shape={output.shape}, time={elapsed*1000:.1f}ms") def main(): if not torch.cuda.is_available(): print("CUDA not available, skipping GPU tests") sys.exit(0) device = torch.device("cuda:0") print(f"Device: {torch.cuda.get_device_name(0)}") print(f"Params: {NUM_EXPERTS} experts, topk={TOPK}, hidden={HIDDEN_SIZE}, " f"inter={INTERMEDIATE_SIZE}, tokens={NUM_TOKENS}") bridge = load_bridge() funcs = [f for f in dir(bridge) if not f.startswith('_')] print(f"Bridge loaded: {len(funcs)} functions: {funcs}") passed = 0 failed = 0 for test_fn in [ test_topk_softmax, test_moe_gen_idx, test_moe_expand_input, test_group_gemm, test_silu_and_mul, test_moe_combine_result, test_fused_pipeline, ]: try: test_fn(bridge, device) passed += 1 except Exception as e: print(f" ✗ FAILED: {e}") import traceback; traceback.print_exc() failed += 1 print(f"\n{'='*40}") print(f"Results: {passed} passed, {failed} failed") if failed > 0: sys.exit(1) if __name__ == "__main__": main()