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

Test: kv_coord = warp_idx() * 0 — force SSA from runtime value

Browse Source
This commit is contained in:
biondizzle
2026-05-22 20:33:40 +00:00
parent 1de848c5ca
commit b3778896b9
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -224,9 +224,12 @@ class FmhaV3StageCMulti:
cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
qp.tail()
kvp.reset()
kv_coord = n_kv_tiles - n_kv_tiles # SSA runtime zero
# Force kv_coord to be a runtime SSA value.
# The JIT folds n_kv_tiles - n_kv_tiles and Int32(0) to compile-time constants.
# cute.arch.make_warp_uniform(cute.arch.warp_idx()) is a runtime value.
# Multiplying by 0 gives a runtime zero the JIT can't constant-fold.
kv_coord = cute.arch.make_warp_uniform(cute.arch.warp_idx()) * 0
for kt in range(n_kv_tiles):
cute.printf("TMA kt=%d kv_coord=%d\n", kt, kv_coord)
kvh = kvp.acquire_and_advance()
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)