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

D1.4: Use cutlass.range(unroll=1) for k_sub loops in both TMA and MMA warps

Browse Source
This commit is contained in:
biondizzle
2026-05-24 17:55:33 +00:00
parent f21defe6dd
commit 828f81ca8c
1 changed files with 8 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
dsv4/kernels/attention/fmha.py
View File

@@ -219,17 +219,16 @@ class FmhaKernel:
# ===== TMA LOAD warp =====
if warp_idx == self.tma_warp_id:
if const_expr(self.n_k_sub_tiles > 1):
# K sub-tiling path (hd>256): Python for loop (unrolled at trace time).
# cutlass.range(unroll=1) creates runtime loops that the MLIR optimizer
# struggles with (45+ min compilation). Python range unrolls at trace time
# but produces simpler IR since each iteration is a flat copy.
# K sub-tiling path (hd>256): use cutlass.range loop to avoid IR explosion
# from Python range unrolling. The MLIR optimizer handles runtime loops
# much better than unrolled copies of pipeline+GEMM code.
qp.reset()
kvp.reset()
for k_sub in range(self.n_k_sub_tiles):
for k_sub in cutlass.range(0, self.n_k_sub_tiles, 1, unroll=1):
qh = qp.acquire_and_advance()
cute.copy(tma_q, tAgQ[(None, Int32(k_sub))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
cute.copy(tma_q, tAgQ[(None, k_sub)], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
kvh = kvp.acquire_and_advance()
cute.copy(tma_k, tBgK[(None, Int32(k_sub))], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_k, tBgK[(None, k_sub)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
# Load V[0]
kvh_v = kvp.acquire_and_advance()
cute.copy(tma_v, tVgV[(None, Int32(0))], tVsV[(None, kvh_v.index)], tma_bar_ptr=kvh_v.barrier)
@@ -252,11 +251,11 @@ class FmhaKernel:
if warp_idx == self.mma_warp_id:
tmem.wait_for_alloc()
if const_expr(self.n_k_sub_tiles > 1):
# K sub-tiling path (hd>256): Python for loop (unrolled at trace time)
# K sub-tiling path (hd>256): cutlass.range loop (runtime, not unrolled)
qc.reset()
kvc.reset()
qk_mma.set(tcgen05.Field.ACCUMULATE, False)
for k_sub in range(self.n_k_sub_tiles):
for k_sub in cutlass.range(0, self.n_k_sub_tiles, 1, unroll=1):
qh = qc.wait_and_advance(); qh.release()
kvh = kvc.wait_and_advance()
for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True):