diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index eaa06fe6..920600ac 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -73,7 +73,7 @@ class Nvfp4FusedRouterKernel: self.use_2cta_instrs = mma_tiler_mn[0] == 256 self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE - self.mma_kind = tcgen05.mma.Kind.BLOCK_SCALE + # No mma_kind needed — make_blockscaled_trivial_tiled_mma infers from dtypes # Warp layout (6 warps: 4 epi + 1 MMA + 1 TMA) self.epilog_warp_id = (0, 1, 2, 3) @@ -95,10 +95,10 @@ class Nvfp4FusedRouterKernel: # ---------------------------------------------------------------- def _create_tiled_mma(self): """Create the tiled MMA for NVFP4 block-scaled GEMM.""" - return utils.sm100.make_trivial_tiled_mma( + return sm100_utils.make_blockscaled_trivial_tiled_mma( self.a_dtype, self.a_major_mode, self.b_major_mode, - self.acc_dtype, self.cta_group, self.mma_tiler_mn, - self.mma_kind, self.sf_dtype, + self.sf_dtype, self.cta_group, self.mma_tiler_mn, + self.sf_vec_size, ) def _setup_attributes(self):