biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix: hardcode mma_inst_shape_k=32 for NVFP4 (avoids MLIR unpack error in JIT)

Browse Source
This commit is contained in:
biondizzle
2026-06-01 07:20:23 +00:00
parent c3a64ceed7
commit 79be9cb8da
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
dsv4/kernels/router/nvfp4_fused_router_kernel.py
View File

@@ -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),