From 7b3f6cb13c01e1531df1761ac5bc523a364ef353 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 09:19:48 +0000 Subject: [PATCH] Fix fused router: use run_nvfp4_fused_router wrapper, correct CuTe tensor API - kernel wrapper converts torch tensors to CuTe tensors with mark_layout_dynamic - test uses the wrapper instead of calling kernel.run() directly - mat_b/scale_b are now torch tensors (converted inside wrapper) --- .../router/nvfp4_fused_router_kernel.py | 31 +++++------ tests/unit/test_fused_router.py | 53 +++++-------------- 2 files changed, 26 insertions(+), 58 deletions(-) diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index 576df42d..2973403d 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -1024,8 +1024,8 @@ class Nvfp4FusedRouterKernel: # ================================================================ def run_nvfp4_fused_router( hidden_states: torch.Tensor, # [M, hidden_size] BF16 - mat_b, # CuTe tensor: gate weight (NVFP4, blockscaled layout) - scale_b, # CuTe tensor: gate weight scale factors (FP8 E4M3) + mat_b: torch.Tensor, # [K_packed, E_packed] NVFP4 gate weight (K-major, torch tensor) + scale_b: torch.Tensor, # [K_sf, E_sf] FP8 E4M3 weight scales (torch tensor) gsa, # Activation global scale (scalar or 1-elem tensor) gsb_val: float, # Weight global scale value e_bias: torch.Tensor, # [E] FP32 @@ -1052,18 +1052,15 @@ def run_nvfp4_fused_router( act_gsa = gsa if gsa is not None else _gsa act_gsb = gsb_val - # Create CuTe tensors for activation + scales - # A tensor: [K_packed, M, L] where K_packed = K/2 (2 elements per byte for FP4) - K_packed = K // 2 - mat_a = cutlass_torch.from_dlpack(act_nvfp4) - mat_a = mat_a.mark_layout_dynamic() - # SFA tensor: [K_sf, M, L] - scale_a = cutlass_torch.from_dlpack(act_sf) - scale_a = scale_a.mark_layout_dynamic() + def to_cute(t): + ct = cutlass_torch.from_dlpack(t) + return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t)) - # e_bias must be a CuTe tensor - e_bias_cute = cutlass_torch.from_dlpack(e_bias) - e_bias_cute = e_bias_cute.mark_layout_dynamic() + mat_a_cute = to_cute(act_nvfp4) + mat_b_cute = to_cute(mat_b) + scale_a_cute = to_cute(act_sf) + scale_b_cute = to_cute(scale_b) + e_bias_cute = to_cute(e_bias) # Number of experts from e_bias E = e_bias.shape[0] @@ -1071,10 +1068,8 @@ def run_nvfp4_fused_router( # Output buffers out_weights = torch.zeros(M, top_k, dtype=torch.float32, device=device) out_ids = torch.zeros(M, top_k, dtype=torch.int32, device=device) - out_w_cute = cutlass_torch.from_dlpack(out_weights) - out_w_cute = out_w_cute.mark_layout_dynamic() - out_id_cute = cutlass_torch.from_dlpack(out_ids) - out_id_cute = out_id_cute.mark_layout_dynamic() + out_w_cute = to_cute(out_weights) + out_id_cute = to_cute(out_ids) # MMA tiler: (128, 128, 64) for decode mma_tiler_mnk = (128, 128, 64) @@ -1086,7 +1081,7 @@ def run_nvfp4_fused_router( top_k=top_k, ) kernel.run( - mat_a, mat_b, scale_a, scale_b, + mat_a_cute, mat_b_cute, scale_a_cute, scale_b_cute, e_bias_cute, out_w_cute, out_id_cute, M, E, K, routed_scaling_factor, top_k, ) diff --git a/tests/unit/test_fused_router.py b/tests/unit/test_fused_router.py index efd3b722..973f3205 100644 --- a/tests/unit/test_fused_router.py +++ b/tests/unit/test_fused_router.py @@ -92,50 +92,23 @@ def test_fused_router(): # ---- Fused kernel ---- print("\n[3] Running fused NVFP4 GEMM + router epilogue...") - from dsv4.kernels.router.nvfp4_fused_router_kernel import Nvfp4FusedRouterKernel - import cutlass.cute as cute - import cutlass.torch as cutlass_torch - - # Quantize activation - act_gs = float(hidden_states.float().abs().max()) / (6.0 * 448.0) - act_nvfp4, act_sf = quantize_activation_nvfp4(hidden_states, act_gs) - - # CuTe tensors for A (activation) - mat_a = cutlass_torch.from_dlpack(act_nvfp4) - mat_a = mat_a.mark_layout_dynamic() - scale_a = cutlass_torch.from_dlpack(act_sf) - scale_a = scale_a.mark_layout_dynamic() - - # CuTe tensors for B (weight) — from gate_lin - mat_b = gate_lin._mat_b - scale_b = gate_lin._scale_b - - # e_bias CuTe tensor - e_bias_cute = cutlass_torch.from_dlpack(e_bias) - e_bias_cute = e_bias_cute.mark_layout_dynamic() - - # Output buffers - out_weights = torch.zeros(M, top_k, dtype=torch.float32, device=device) - out_ids = torch.zeros(M, top_k, dtype=torch.int32, device=device) - out_w_cute = cutlass_torch.from_dlpack(out_weights) - out_w_cute = out_w_cute.mark_layout_dynamic() - out_id_cute = cutlass_torch.from_dlpack(out_ids) - out_id_cute = out_id_cute.mark_layout_dynamic() - - kernel = Nvfp4FusedRouterKernel( - sf_vec_size=sf_vec_size, - mma_tiler_mnk=(128, 128, 64), - cluster_shape_mnk=(1, 1, 1), - top_k=top_k, - ) + from dsv4.kernels.router.nvfp4_fused_router_kernel import run_nvfp4_fused_router try: - kernel.run( - mat_a, mat_b, scale_a, scale_b, - e_bias_cute, out_w_cute, out_id_cute, - M, E, K, routed_scaling_factor, top_k, + fused_weights, fused_ids = run_nvfp4_fused_router( + hidden_states=hidden_states, + mat_b=gate_lin._mat_b, + scale_b=gate_lin._scale_b, + gsa=gate_lin._gsa_buf, + gsb_val=float(gate_lin._gsb), + e_bias=e_bias, + routed_scaling_factor=routed_scaling_factor, + top_k=top_k, + sf_vec_size=sf_vec_size, ) print(" Fused kernel compilation and execution succeeded!") + print(f" Fused topk_ids: {fused_ids[0].tolist()}") + print(f" Fused topk_weights: {fused_weights[0].tolist()}") except Exception as ex: print(f" FUSED KERNEL FAILED: {ex}") import traceback