diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index e3aeaafe..53436c6d 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -108,8 +108,10 @@ class Nvfp4FusedRouterKernel: mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) mma_inst_tile_k = self.mma_tiler_mnk[2] // mma_inst_shape_k - self.mma_tiler = (self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], self.mma_tiler_mnk[2]) - self.mma_tiler_sfb = (self.mma_inst_shape_mn_sfb[0], self.mma_inst_shape_mn_sfb[1], self.mma_tiler_mnk[2]) + # ── MMA tiler — K is refined in _setup_attributes ── + # ── MMA tiler — K is refined in _setup_attributes ── + self.mma_tiler = (mma_tiler_mnk[0], mma_tiler_mnk[1], 1) + self.mma_tiler_sfb = (mma_tiler_mnk[0] // (2 if self.use_2cta_instrs else 1), cute.round_up(mma_tiler_mnk[1], 128), 1) self.cta_tile_shape_mnk = ( self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), self.mma_tiler[1],