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

D1: Use cutlass.range for k_sub loops (CuTeDSL immutable handle)

Browse Source
This commit is contained in:
biondizzle
2026-05-24 06:43:30 +00:00
parent 2bf3ee40aa
commit dd39c2ebdf
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -216,8 +216,8 @@ class FmhaKernel:
# K sub-tiling path (hd=512): load Q and K per k_sub
qp.reset()
kvp.reset()
for kt in range(self.n_kv_tiles):
for k_sub in range(self.n_k_sub_tiles):
for kt in cutlass.range(0, self.n_kv_tiles, 1, unroll=1):
for k_sub in cutlass.range(0, self.n_k_sub_tiles, 1, unroll=1):
# Load Q[k_sub] → sQ
qh = qp.acquire_and_advance()
cute.copy(tma_q, tAgQ[(None, Int32(k_sub))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
@@ -248,8 +248,8 @@ class FmhaKernel:
if const_expr(self.n_k_sub_tiles > 1):
# K sub-tiling path (hd=512): pipeline sync for k_sub iterations
kvh = kvc.wait_and_advance() # initialize kvh before loops
for kt in range(self.n_kv_tiles):
for k_sub in range(self.n_k_sub_tiles):
for kt in cutlass.range(0, self.n_kv_tiles, 1, unroll=1):
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()
qk_mma.set(tcgen05.Field.ACCUMULATE, k_sub != 0)