Fix: mma_tiler K=1 placeholder in __init__, refined in _setup_attributes
Same pattern as fused_swiglu.py: - __init__ sets mma_tiler = (M, N, 1) with K=1 placeholder - _setup_attributes refines K to the actual value from cute.size(tiled_mma.shape_mnk) - cute.slice_ and cute.local_tile work correctly with the K=1 initial value - mma_tiler_sfb also gets K=1 placeholder This fixes the MLIR crash on cute.slice_(self.mma_tiler, (None, 0, None)) which couldn't handle the full (128, 128, 64) tuple.
This commit is contained in:
@@ -108,8 +108,10 @@ class Nvfp4FusedRouterKernel:
|
||||
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
|
||||
mma_inst_tile_k = self.mma_tiler_mnk[2] // mma_inst_shape_k
|
||||
|
||||
self.mma_tiler = (self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], 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])
|
||||
# ── MMA tiler — K is refined in _setup_attributes ──
|
||||
# ── MMA tiler — K is refined in _setup_attributes ──
|
||||
self.mma_tiler = (mma_tiler_mnk[0], mma_tiler_mnk[1], 1)
|
||||
self.mma_tiler_sfb = (mma_tiler_mnk[0] // (2 if self.use_2cta_instrs else 1), cute.round_up(mma_tiler_mnk[1], 128), 1)
|
||||
self.cta_tile_shape_mnk = (
|
||||
self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
|
||||
self.mma_tiler[1],
|
||||
|
||||
Reference in New Issue
Block a user