From 44fb9b6c00fc3ccdd5c0500122e4754303bf4550 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 10:55:43 +0000 Subject: [PATCH] Fix: pass self.mma_tiler_mnk (full K) to _compute_stages, not self.mma_tiler (K=1 placeholder) --- dsv4/kernels/router/nvfp4_fused_router_kernel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index cb22ab06..d59caeb7 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -148,7 +148,7 @@ class Nvfp4FusedRouterKernel: # Stage counts (same as MoE) self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( - tiled_mma, self.mma_tiler, a_dtype, b_dtype, + tiled_mma, self.mma_tiler_mnk, a_dtype, b_dtype, self.epi_tile, c_dtype, c_layout, sf_dtype, self.sf_vec_size, self.smem_capacity, self.occupancy)