Fix: pass self.mma_tiler_mnk (full K) to _compute_stages, not self.mma_tiler (K=1 placeholder)

2026-06-01 10:55:43 +00:00
parent be2bb2fe84
commit 44fb9b6c00
1 changed files with 1 additions and 1 deletions
--- 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)