nvfp4-megamoe-kernel/tests/unit/test_fused_router.py

"""Test NVFP4 fused router kernel against the reference path.

Phase 1: Reference path (BF16 linear + activation_topk)
Phase 2: NVFP4 fused kernel vs BF16 reference
Phase 3: NVFP4 fused kernel vs NVFP4 2-kernel path
"""

import sys
import os
import torch

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))


def test_reference_router():
    """Test the reference BF16 linear + activation_topk path."""
    torch.manual_seed(42)
    device = "cuda"
    M, K, N, top_k = 4, 7168, 384, 6
    routed_scaling_factor = 0.5

    hidden_states = torch.randn(M, K, dtype=torch.bfloat16, device=device)
    W_gate = torch.randn(K, N, dtype=torch.bfloat16, device=device)
    e_bias = torch.randn(N, dtype=torch.float32, device=device)

    logits = torch.nn.functional.linear(hidden_states.float(), W_gate.T.float())
    from dsv4.kernels.router._activation_topk import run_fused_activation_topk
    out_w = torch.empty(M, top_k, dtype=torch.float32, device=device)
    out_ids = torch.empty(M, top_k, dtype=torch.int32, device=device)
    run_fused_activation_topk(logits, e_bias, routed_scaling_factor, top_k, out_w, out_ids)

    w_sum = out_w.sum(dim=1)
    assert all(abs(w_sum[r].item() - routed_scaling_factor) < 0.01 for r in range(M))
    assert (out_ids >= 0).all() and (out_ids < N).all()
    for r in range(M):
        assert len(set(out_ids[r].tolist())) == top_k
    assert (out_w >= 0).all()

    print(f"Reference router (M={M}, K={K}, N={N}): PASSED")
    print(f"  IDs row0: {out_ids[0].tolist()}")
    print(f"  Weights row0: {[f'{w:.4f}' for w in out_w[0].tolist()]}")


def test_nvfp4_fused_router():
    """Test NVFP4 fused router: compare fused kernel vs 2-kernel path.

    Both use the same Nvfp4Linear (same quantized weights), so they should
    match exactly (same GEMM, same activation_topk math).
    """
    torch.manual_seed(42)
    device = "cuda"
    M, K, N, top_k = 1, 7168, 384, 6
    routed_scaling_factor = 0.5

    print(f"\nNVFP4 fused router (M={M}, K={K}, N={N}):")

    hidden_states = torch.randn(M, K, dtype=torch.bfloat16, device=device)
    e_bias = torch.randn(N, dtype=torch.float32, device=device)

    # Build Nvfp4Linear from checkpoint-style quantized weights
    # The checkpoint stores: weight (N_packed, K_packed) uint8, weight_scale (N_packed, K_sf)
    # For random BF16 weights, we need to quantize ourselves.
    # quantize_weight_to_nvfp4 expects (K, N) BF16, returns (K//2, N) FP4
    # But Nvfp4Linear expects (N_packed, K_packed) = (N, K//2) after view
    # We need to transpose the output of quantize_weight_to_nvfp4

    from dsv4.ops.quantize import quantize_weight_to_nvfp4
    W_gate_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device)
    w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(W_gate_bf16)
    # w_fp4: (K//2, N) = (3584, 384) float4_e2m1fn_x2
    # w_sf: (K//16, N) float8_e4m3fn
    # w_gs: scalar

    # Nvfp4Linear expects fp4 in (N, K//2) format — transpose
    w_fp4_nk = w_fp4.T.contiguous()  # (N, K//2) = (384, 3584)
    w_sf_nk = w_sf.T.contiguous()    # (N, K//16)

    from dsv4.layers.linear import Nvfp4Linear
    gate_lin = Nvfp4Linear(K, N, max_num_tokens=8, device=device)
    gate_lin.fp4 = [w_fp4_nk]
    gate_lin.sf = [w_sf_nk]
    gate_lin.gs = [w_gs.item()]
    gate_lin.ws2 = [None]
    gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0)
    gate_lin.finalize_weights()

    # 2-kernel NVFP4 reference path
    logits_nvfp4 = gate_lin(hidden_states).float()
    print(f"  NVFP4 GEMM output shape: {logits_nvfp4.shape}")
    print(f"  NVFP4 GEMM output[0,:5]: {logits_nvfp4[0,:5].tolist()}")

    from dsv4.kernels.router._activation_topk import run_fused_activation_topk
    ref_w = torch.empty(M, top_k, dtype=torch.float32, device=device)
    ref_ids = torch.empty(M, top_k, dtype=torch.int32, device=device)
    run_fused_activation_topk(logits_nvfp4, e_bias, routed_scaling_factor, top_k, ref_w, ref_ids)
    print(f"  2-kernel: IDs={ref_ids[0].tolist()}")

    # Fused kernel path — use Nvfp4Linear's processed weight tensors
    from dsv4.kernels.router.nvfp4_fused_router_kernel import run_nvfp4_fused_router
    gsb_val = gate_lin._gsb.item()
    gsa = gate_lin._activation_global_scale
    fused_w, fused_ids = run_nvfp4_fused_router(
        hidden_states, gate_lin._mat_b, gate_lin._scale_b,
        gsa, gsb_val, e_bias, routed_scaling_factor, top_k,
    )
    print(f"  Fused:    IDs={fused_ids[0].tolist()}")

    # Compare
    ids_match = (fused_ids == ref_ids).all().item()
    if ids_match:
        print(f"  IDs match: OK")
    else:
        mismatches = (fused_ids != ref_ids).sum().item()
        print(f"  ID mismatches: {mismatches}")

    if fused_w.shape == ref_w.shape:
        cos = torch.nn.functional.cosine_similarity(
            fused_w.flatten().unsqueeze(0), ref_w.flatten().unsqueeze(0)).item()
        max_diff = (fused_w - ref_w).abs().max().item()
        print(f"  Weight cosine: {cos:.6f}, max_diff: {max_diff:.6f}")
        if cos < 0.99:
            print(f"  WARNING: Poor weight match — needs investigation")

    w_sum = fused_w.sum(dim=1)
    for row in range(M):
        diff = abs(w_sum[row].item() - routed_scaling_factor)
        if diff > 0.01:
            print(f"  WARNING: Row {row} weight sum {w_sum[row].item():.4f} != {routed_scaling_factor:.4f}")

    print("NVFP4 fused router test DONE")


if __name__ == "__main__":
    test_reference_router()
    print()
    try:
        test_nvfp4_fused_router()
    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"NVFP4 fused router test failed: {e}")