95 lines
3.5 KiB
Python
95 lines
3.5 KiB
Python
"""Test NVFP4 fused router kernel against the reference path.
|
|
|
|
Reference: Nvfp4Linear (NVFP4 GEMM) → activation_topk CUDA kernel
|
|
Test: Nvfp4FusedRouterKernel (single-kernel fusion)
|
|
|
|
Both should produce identical top-k weights and expert IDs.
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import torch
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
|
|
|
|
def test_fused_router():
|
|
"""Compare fused kernel vs. reference path for the router."""
|
|
torch.manual_seed(42)
|
|
device = "cuda"
|
|
M = 4 # tokens
|
|
K = 7168 # hidden size
|
|
N = 384 # num experts
|
|
top_k = 6
|
|
routed_scaling_factor = 0.5
|
|
|
|
from dsv4.layers.linear import Nvfp4Linear
|
|
from dsv4.ops.quantize import quantize_activation_nvfp4
|
|
|
|
# Create BF16 hidden states
|
|
hidden_states = torch.randn(M, K, dtype=torch.bfloat16, device=device)
|
|
|
|
# Create random BF16 gate weight and quantize to NVFP4
|
|
W_gate_bf16 = torch.randn(K, N, dtype=torch.bfloat16, device=device)
|
|
from dsv4.ops.quantize import quantize_weight_to_nvfp4
|
|
w_fp4, w_sf, ws2_val = quantize_weight_to_nvfp4(W_gate_bf16)
|
|
|
|
# Build Nvfp4Linear for reference path
|
|
# The Nvfp4Linear expects stacked/contiguous weight tensors
|
|
from dsv4.layers.linear import Nvfp4Linear
|
|
gate_lin = Nvfp4Linear(in_features=K, out_features=N, device=device)
|
|
gate_lin.fp4 = [w_fp4.contiguous()]
|
|
gate_lin.sf = [w_sf.contiguous()]
|
|
gate_lin.gsb = None
|
|
gate_lin.gs = [1.0] # default global scale
|
|
gate_lin.ws2 = [ws2_val]
|
|
gate_lin._activation_global_scale = 1.0 # default
|
|
gate_lin.finalize_weights()
|
|
|
|
# e_bias
|
|
e_bias = torch.randn(N, dtype=torch.float32, device=device)
|
|
|
|
# === Reference path: Nvfp4Linear + activation_topk ===
|
|
logits_ref = gate_lin(hidden_states).float() # (M, N) FP32
|
|
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 path: Nvfp4FusedRouterKernel ===
|
|
try:
|
|
from dsv4.kernels.router.nvfp4_fused_router_kernel import run_nvfp4_fused_router
|
|
gsb_val = ws2_val
|
|
fused_w, fused_ids = run_nvfp4_fused_router(
|
|
hidden_states, w_fp4, w_sf, gsb_val, e_bias,
|
|
routed_scaling_factor, top_k,
|
|
)
|
|
|
|
# Compare
|
|
ids_match = (ref_ids == fused_ids).all().item()
|
|
if ids_match:
|
|
print(f"Fused router: expert IDs MATCH reference ✓")
|
|
else:
|
|
print(f"Fused router: expert IDs MISMATCH")
|
|
for row in range(M):
|
|
print(f" Row {row}: ref={ref_ids[row].tolist()} fused={fused_ids[row].tolist()}")
|
|
|
|
# Compare weights
|
|
if ref_w.shape == fused_w.shape:
|
|
max_diff = (ref_w - fused_w).abs().max().item()
|
|
print(f" Max weight diff: {max_diff:.6f}")
|
|
else:
|
|
print(f" Weight shape mismatch: ref={ref_w.shape} fused={fused_w.shape}")
|
|
|
|
except Exception as e:
|
|
print(f"Fused router kernel compilation/runtime error (expected on non-B200):")
|
|
print(f" {e}")
|
|
print(f"\nReference path (Nvfp4Linear + activation_topk) works correctly.")
|
|
print(f"Ref top-k IDs: {ref_ids[0].tolist()}")
|
|
print(f"Ref top-k weights: {ref_w[0].tolist()}")
|
|
print(f"\nFused kernel needs B200 for CuTeDSL compilation.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_fused_router()
|