NVFP4 fused router kernel: full rewrite with proper block-scaled GEMM setup
Major fixes: - Added tiled_mma_sfb creation (always CtaGroup.ONE, rounded N) - Added mma_tiler_sfb, cta_tile_shape_mnk_sfb, cluster_layout_sfb_vmnk - Use blockscaled_utils.make_smem_layout_sfa/sfb (with sf_vec_size) instead of sm100_utils (which doesn't support block-scaled SF layouts) - Proper TMEM column accounting for SFA + SFB + accumulator - Fixed make_blockscaled_trivial_tiled_mma argument order (a_dtype, b_dtype, a_major, b_major, sf_dtype, sf_vec_size, cta_group, mma_inst_shape) - Fixed SFB TMA atom to use tiled_mma_sfb and cluster_layout_sfb_vmnk - Fixed SFB partition_SFB to use tiled_mma_sfb.get_slice - Fixed SFB global tile partitioning to use mma_tiler_sfb - Fixed mainloop_s2t_copy_and_partition to use TMEM fragments (make_fragment_SFA/SFB) as the tSF parameter - Updated run_nvfp4_fused_router wrapper to accept processed weight tensors from Nvfp4Linear._mat_b and _scale_b - Updated test to properly build Nvfp4Linear and use processed weights The old code was a rough sketch that never worked — it was missing the entire tiled_mma_sfb infrastructure, used wrong SMEM layout functions, and had broken TMA atom setup for scale factors.
This commit is contained in:
@@ -1,27 +1,57 @@
|
||||
"""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).
|
||||
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 # tokens
|
||||
K = 7168 # hidden size
|
||||
N = 384 # num experts
|
||||
M = 4
|
||||
K = 7168
|
||||
N = 384
|
||||
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)
|
||||
@@ -33,53 +63,195 @@ def test_reference_router():
|
||||
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
|
||||
# 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 (approximately)
|
||||
# Verify weights sum to routed_scaling_factor
|
||||
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})")
|
||||
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})")
|
||||
|
||||
# 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}) ✓")
|
||||
print(f" All expert IDs in [0, {N}) OK")
|
||||
|
||||
# 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 ✓")
|
||||
print(f" No duplicate IDs OK")
|
||||
|
||||
# Verify: weights are non-negative
|
||||
assert (out_w >= 0).all(), "Negative weights"
|
||||
print(f" All weights non-negative ✓")
|
||||
print(f" All weights non-negative OK")
|
||||
|
||||
print("Reference router test PASSED ✓")
|
||||
# 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_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.")
|
||||
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()
|
||||
test_nvfp4_fused_router_import()
|
||||
# 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}")
|
||||
|
||||
Reference in New Issue
Block a user