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

Fix: all CuTe shape values must be cutlass.Int32 for MLIR compatibility

Browse Source 
Python ints cause 'Expected an MLIR object (got None)' in _pack_shape.
This is the same fix we applied to the FMHA kernel mma_tiler.
All mma_inst_shape, mma_tiler, cluster_shape values now use cutlass.Int32().
This commit is contained in:
biondizzle
2026-06-01 09:30:15 +00:00
parent df48dacc2b
commit 65669596d4
1 changed files with 23 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -70,10 +70,18 @@ class Nvfp4FusedRouterKernel:
self.arch = "sm_100"
# Set up MMA instruction shapes before any MMA creation calls
self.mma_inst_shape_mn = (mma_tiler_mnk[0], mma_tiler_mnk[1])
# ALL shape values must be cutlass.Int32 so CuTe layout construction
# produces static (not dynamic) dimensions. Python ints cause
# "Expected an MLIR object (got None)" in _pack_shape.
self.mma_inst_shape_mn = (
cutlass.Int32(mma_tiler_mnk[0]),
cutlass.Int32(mma_tiler_mnk[1]),
)
# round_up with Python math, then wrap in cutlass.Int32
sfb_n = ((mma_tiler_mnk[1] + 128 - 1) // 128) * 128
self.mma_inst_shape_mn_sfb = (
self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
cute.round_up(self.mma_inst_shape_mn[1], 128),
cutlass.Int32(mma_tiler_mnk[0] // (2 if self.use_2cta_instrs else 1)),
cutlass.Int32(sfb_n),
)
# 6-warp specialization (no scheduler warp for dense GEMM)
@@ -124,12 +132,12 @@ class Nvfp4FusedRouterKernel:
self.mma_tiler = (
self.mma_inst_shape_mn[0],
self.mma_inst_shape_mn[1],
self.mma_tiler_mnk[2],
cutlass.Int32(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],
cutlass.Int32(self.mma_tiler_mnk[2]),
)
self.cta_tile_shape_mnk = (
self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
@@ -143,10 +151,18 @@ class Nvfp4FusedRouterKernel:
)
self.cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
cute.make_layout((
cutlass.Int32(self.cluster_shape_mn[0]),
cutlass.Int32(self.cluster_shape_mn[1]),
cutlass.Int32(1),
)),
(tiled_mma.thr_id.shape,))
self.cluster_layout_sfb_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
cute.make_layout((
cutlass.Int32(self.cluster_shape_mn[0]),
cutlass.Int32(self.cluster_shape_mn[1]),
cutlass.Int32(1),
)),
(tiled_mma_sfb.thr_id.shape,))
self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])