From 79be9cb8da469bd63bd3f68668fd90129020aca1 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 07:20:23 +0000 Subject: [PATCH] Fix: hardcode mma_inst_shape_k=32 for NVFP4 (avoids MLIR unpack error in JIT) --- dsv4/kernels/router/nvfp4_fused_router_kernel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index 2cd47635..6f945207 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -125,15 +125,15 @@ class Nvfp4FusedRouterKernel: self._tiled_mma = tiled_mma self._tiled_mma_sfb = tiled_mma_sfb - mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_shape_k = 32 # NVFP4 sf_vec_size=16: MMA inst K = sf_vec_size * 2 (FP4 packing) mma_inst_tile_k = 4 self.mma_tiler = ( self.mma_inst_shape_mn[0], self.mma_inst_shape_mn[1], - mma_inst_shape_k * mma_inst_tile_k, + cutlass.Int32(mma_inst_shape_k * mma_inst_tile_k), ) self.mma_tiler_sfb = ( self.mma_inst_shape_mn_sfb[0], self.mma_inst_shape_mn_sfb[1], - mma_inst_shape_k * mma_inst_tile_k, + cutlass.Int32(mma_inst_shape_k * mma_inst_tile_k), ) self.cta_tile_shape_mnk = ( self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),