86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
"""Test NVFP4 fused router kernel against the reference path.
|
|
|
|
Phase 1: Verify reference path (BF16 linear + activation_topk) works.
|
|
Phase 2: Test CuTeDSL fused kernel (needs B200 for CuTeDSL compilation).
|
|
"""
|
|
|
|
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 = 4 # tokens
|
|
K = 7168 # hidden size
|
|
N = 384 # num experts
|
|
top_k = 6
|
|
routed_scaling_factor = 0.5
|
|
|
|
# Create BF16 hidden states and weight
|
|
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)
|
|
|
|
# Verify results
|
|
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 (approximately)
|
|
w_sum = out_w.sum(dim=1)
|
|
expected = routed_scaling_factor
|
|
for row in range(M):
|
|
diff = abs(w_sum[row].item() - expected)
|
|
assert diff < 0.01, f"Row {row}: weight sum {w_sum[row].item():.4f} != {expected:.4f}"
|
|
print(f" Weight sums: {[f'{s:.4f}' for s in w_sum.tolist()]} (expected {expected})")
|
|
|
|
# Verify: IDs are valid (0 <= id < N)
|
|
assert (out_ids >= 0).all() and (out_ids < N).all(), "Invalid expert IDs"
|
|
print(f" All expert IDs in [0, {N}) ✓")
|
|
|
|
# Verify: no duplicate IDs within each row
|
|
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 ✓")
|
|
|
|
# Verify: weights are non-negative
|
|
assert (out_w >= 0).all(), "Negative weights"
|
|
print(f" All weights non-negative ✓")
|
|
|
|
print("Reference router test PASSED ✓")
|
|
|
|
|
|
def test_nvfp4_fused_router_import():
|
|
"""Test that the fused router kernel module imports without error."""
|
|
try:
|
|
from dsv4.kernels.router.nvfp4_fused_router_kernel import Nvfp4FusedRouterKernel
|
|
print(f"Nvfp4FusedRouterKernel imported successfully")
|
|
kernel = Nvfp4FusedRouterKernel(top_k=6)
|
|
print(f" mma_tiler_mn: {kernel.mma_tiler_mn}")
|
|
print(f" threads_per_cta: {kernel.threads_per_cta}")
|
|
print(f" sf_vec_size: {kernel.sf_vec_size}")
|
|
print(f" top_k: {kernel.top_k}")
|
|
print("Fused router kernel class construction PASSED ✓")
|
|
except Exception as e:
|
|
print(f"Fused router kernel import failed: {e}")
|
|
print("This is expected if CuTeDSL is not available.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_reference_router()
|
|
print()
|
|
test_nvfp4_fused_router_import()
|