From 65669596d4c28fb50450ae727354f00f0acb770c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 09:30:15 +0000 Subject: [PATCH] Fix: all CuTe shape values must be cutlass.Int32 for MLIR compatibility 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(). --- .../router/nvfp4_fused_router_kernel.py | 30 ++++++++++++++----- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/dsv4/kernels/router/nvfp4_fused_router_kernel.py b/dsv4/kernels/router/nvfp4_fused_router_kernel.py index 3c730206..d0b7d29b 100644 --- a/dsv4/kernels/router/nvfp4_fused_router_kernel.py +++ b/dsv4/kernels/router/nvfp4_fused_router_kernel.py @@ -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])