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

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

Reference path: Nvfp4Linear (NVFP4 GEMM) -> activation_topk CUDA kernel
Fused path:     Nvfp4FusedRouterKernel (single kernel, no intermediate logits)

Tests:
1. Reference path correctness (BF16 GEMM + activation_topk)
2. NVFP4 fused kernel matches reference (cosine similarity)
3. NVFP4 fused kernel matches at various (M, N, K) shapes
"""

import sys
import os
import torch
import math

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


def _sqrt_softplus(x: torch.Tensor) -> torch.Tensor:
    """Reference: sqrt(softplus(x))."""
    abs_x = x.abs()
    sp = x.clamp(min=0) + torch.log1p(torch.exp(-abs_x))
    return sp.sqrt()


def _reference_topk(
    logits: torch.Tensor,    # [M, N] FP32
    e_bias: torch.Tensor,    # [N] FP32
    routed_scaling_factor: float,
    top_k: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Pure PyTorch reference: sqrt(softplus) + bias + topk + renorm."""
    act = _sqrt_softplus(logits)
    scores = act + e_bias.unsqueeze(0)
    topk_vals, topk_ids = scores.topk(top_k, dim=1)
    # Gather activations for the selected experts
    selected_act = act.gather(1, topk_ids)
    # Renormalize
    act_sum = selected_act.sum(dim=1, keepdim=True)
    weights = selected_act / act_sum * routed_scaling_factor
    return weights, topk_ids


def test_reference_router():
    """Test the reference BF16 linear + activation_topk path."""
    torch.manual_seed(42)
    device = "cuda"
    M = 4
    K = 7168
    N = 384
    top_k = 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)

    # Reference: BF16 linear + activation_topk
    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)

    # Also compute pure PyTorch reference
    ref_w, ref_ids = _reference_topk(logits, e_bias, routed_scaling_factor, top_k)

    print(f"Reference router test (M={M}, K={K}, N={N}, top_k={top_k}):")
    print(f"  Top-k IDs (row 0): {out_ids[0].tolist()}")
    print(f"  Top-k weights (row 0): {[f'{w:.4f}' for w in out_w[0].tolist()]}")

    # Verify weights sum to routed_scaling_factor
    w_sum = out_w.sum(dim=1)
    for row in range(M):
        diff = abs(w_sum[row].item() - routed_scaling_factor)
        assert diff < 0.01, f"Row {row}: weight sum {w_sum[row].item():.4f} != {routed_scaling_factor:.4f}"
    print(f"  Weight sums: {[f'{s:.4f}' for s in w_sum.tolist()]} (expected {routed_scaling_factor})")

    assert (out_ids >= 0).all() and (out_ids < N).all(), "Invalid expert IDs"
    print(f"  All expert IDs in [0, {N}) OK")

    for row in range(M):
        row_ids = out_ids[row].tolist()
        assert len(set(row_ids)) == len(row_ids), f"Row {row} has duplicate IDs: {row_ids}"
    print(f"  No duplicate IDs OK")

    assert (out_w >= 0).all(), "Negative weights"
    print(f"  All weights non-negative OK")

    # Cross-check with pure PyTorch reference
    ids_match = (out_ids == ref_ids).all().item()
    if ids_match:
        print(f"  IDs match PyTorch reference OK")
    else:
        mismatches = (out_ids != ref_ids).sum().item()
        print(f"  WARNING: {mismatches} ID mismatches vs PyTorch reference")

    print("Reference router test PASSED")


def test_nvfp4_fused_router():
    """Test the NVFP4 fused router kernel against reference path.

    This test requires B200 (CuTeDSL compilation + NVFP4 tensor cores).
    Must be run via fire_b200_test.
    """
    torch.manual_seed(42)
    device = "cuda"
    M = 1        # decode: single token
    K = 7168     # hidden size
    N = 384      # num experts
    top_k = 6
    routed_scaling_factor = 0.5

    print(f"NVFP4 fused router test (M={M}, K={K}, N={N}, top_k={top_k}):")

    # Create BF16 hidden states
    hidden_states = torch.randn(M, K, dtype=torch.bfloat16, device=device)
    # Create BF16 gate weight and quantize to NVFP4
    W_gate_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device)
    e_bias = torch.randn(N, dtype=torch.float32, device=device)

    # Quantize weight to NVFP4 (matching Nvfp4Linear's quantization)
    from dsv4.ops.quantize import quantize_weight_nvfp4, quantize_activation_nvfp4
    w_fp4, w_sf, ws2_scalar = quantize_weight_nvfp4(W_gate_bf16)

    # Reference path: NVFP4 GEMM + activation_topk
    # Build Nvfp4Linear the same way single_shot_inference does
    from dsv4.layers.linear import Nvfp4Linear
    gate_lin = Nvfp4Linear(K, N, max_num_tokens=8, device=device)
    gate_lin.fp4 = [w_fp4]
    gate_lin.sf = [w_sf]
    gate_lin.gs = [1.0]
    gate_lin.ws2 = [ws2_scalar.to(device) if ws2_scalar is not None else None]
    gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0)
    gate_lin.finalize_weights()

    logits_ref = gate_lin(hidden_states).float()
    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_ref, e_bias, routed_scaling_factor, top_k, ref_w, ref_ids)

    print(f"  Reference: IDs={ref_ids[0].tolist()}, weights={[f'{w:.4f}' for w in ref_w[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()}, weights={[f'{w:.4f}' for w in fused_w[0].tolist()]}")

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

    # Compare weights
    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.999:
            print(f"  Weight match: EXCELLENT")
        elif cos > 0.99:
            print(f"  Weight match: GOOD")
        else:
            print(f"  Weight match: POOR - needs investigation")

    # Verify weight normalization
    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")


def test_nvfp4_fused_router_multitoken():
    """Test with M=4 tokens (batched decode or small prefill)."""
    torch.manual_seed(123)
    device = "cuda"
    M = 4
    K = 7168
    N = 384
    top_k = 6
    routed_scaling_factor = 0.5

    print(f"NVFP4 fused router multi-token test (M={M}, K={K}, N={N}):")

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

    from dsv4.ops.quantize import quantize_weight_nvfp4
    w_fp4, w_sf, ws2_scalar = quantize_weight_nvfp4(W_gate_bf16)

    # Reference
    from dsv4.layers.linear import Nvfp4Linear
    gate_lin = Nvfp4Linear(K, N, max_num_tokens=8, device=device)
    gate_lin.fp4 = [w_fp4]
    gate_lin.sf = [w_sf]
    gate_lin.gs = [1.0]
    gate_lin.ws2 = [ws2_scalar.to(device) if ws2_scalar is not None else None]
    gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0)
    gate_lin.finalize_weights()
    logits_ref = gate_lin(hidden_states).float()
    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_ref, e_bias, routed_scaling_factor, top_k, ref_w, ref_ids)

    # Fused
    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,
    )

    # Compare
    ids_match = (fused_ids == ref_ids).all().item()
    mismatches = (fused_ids != ref_ids).sum().item() if not ids_match else 0
    cos = torch.nn.functional.cosine_similarity(
        fused_w.flatten().unsqueeze(0), ref_w.flatten().unsqueeze(0)).item()

    print(f"  IDs match: {ids_match} ({mismatches} mismatches)")
    print(f"  Weight cosine: {cos:.6f}")
    print("NVFP4 fused router multi-token test DONE")


if __name__ == "__main__":
    test_reference_router()
    print()
    # NVFP4 fused tests require B200 — they'll fail on non-SM100 hardware
    try:
        test_nvfp4_fused_router()
    except Exception as e:
        print(f"NVFP4 fused router test skipped (requires B200): {e}")
    print()
    try:
        test_nvfp4_fused_router_multitoken()
    except Exception as e:
        print(f"NVFP4 fused router multi-token test skipped (requires B200): {e}")