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

Fix: mma_tiler and cluster_layout must use MLIR values for cute.slice_

Browse Source 
cute.slice_ on Python int tuples fails. All values in mma_tiler and
cluster_layout need to be cutlass.Int32() since they flow into
cute.slice_ and cute.local_tile inside @cute.kernel.

Now consistent: mma_inst_shape_mn, mma_tiler, cluster_layout_vmnk all
use MLIR-typed values created inside @cute.jit context.
This commit is contained in:
biondizzle
2026-06-01 09:42:17 +00:00
parent ec79f30709
commit d9d3ca42b0
1 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -70,8 +70,12 @@ class Nvfp4FusedRouterKernel:
self.arch = "sm_100"
# Set up MMA instruction shapes
# These are now used inside @cute.jit context, so cute.round_up is fine
self.mma_inst_shape_mn = (mma_tiler_mnk[0], mma_tiler_mnk[1])
# All values must be CuTe/MLIR values for cute.slice_ compatibility
# since mma_tiler is used inside @cute.kernel with cute.slice_
self.mma_inst_shape_mn = (
cutlass.Int32(mma_tiler_mnk[0]),
cutlass.Int32(mma_tiler_mnk[1]),
)
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),
@@ -125,12 +129,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),
@@ -144,10 +148,10 @@ 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])