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

TMA: use self.n_kv_tiles + kv_coord pattern from working diag test

Browse Source
This commit is contained in:
biondizzle
2026-05-22 20:08:07 +00:00
parent 74380ffc21
commit 6a82036ebd
1 changed files with 9 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
tests/fmha_v3_stage_c_example7.py
View File

@@ -56,6 +56,7 @@ class FmhaV3StageCMulti:
def __init__(self, s_k=128, scale_softmax=None):
# s_k MUST equal actual sequence length n.
self.s_k = s_k
self.n_kv_tiles = s_k // 128
self.acc_dtype = Float32; self.qk_acc_dtype = Float32
self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16
self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1
@@ -203,22 +204,20 @@ class FmhaV3StageCMulti:
pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)
# ===== TMA LOAD warp =====
# GMEM tile coordinate: use the cutlass.range induction variable kt
# directly. CuTeDSL's `cutlass.range` doesn't auto-detect a Python `+=`
# rebinding as a loop-carried iter_args update — the JIT traces the
# body once and captures whatever value `kv_coord` had at trace time,
# so an outer `kv_coord = Int32(0)` plus a `kv_coord += 1` inside the
# loop bakes 0 into every iteration's TMA descriptor at runtime.
# The induction variable IS the loop-carried state, properly tracked.
# Use the EXACT pattern from test_fmha_v3_diag.py which works for n=256.
# kv_coord = Int32(0+0) + kv_coord += 1 in cutlass.range(self.n_kv_tiles, unroll=1)
# with tBgK = tBgK[(None,None,0,0)] (GMEM tile dim free)
if warp_idx == self.tma_warp_id:
qp.reset(); qh = qp.acquire_and_advance()
cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
qp.tail()
kvp.reset(); pk = kvp.try_acquire()
for kt in cutlass.range(0, n_kv_tiles, 1, unroll=1):
kv_coord = Int32(0 + 0)
for kt in cutlass.range(self.n_kv_tiles, unroll=1):
kvh = kvp.acquire_and_advance(pk)
cute.copy(tma_k, tBgK[(None, kt)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_v, tVgV[(None, kt)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_k, tBgK[(None, kv_coord)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_v, tVgV[(None, kv_coord)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
kv_coord += 1
pk = cutlass.Boolean(1)
kvp.tail()