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

D1: Use same pipeline API as working code (acquire_and_advance) for k_sub path

Browse Source
This commit is contained in:
biondizzle
2026-05-24 06:36:19 +00:00
parent 85af7f4cf3
commit 3bf1e62b58
1 changed files with 18 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -213,25 +213,25 @@ 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=512): load Q and K per k_sub using pipeline barriers
qp_stage = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.q_stage)
kvp_stage = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.kv_stage)
# 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):
# Load Q[k_sub] → sQ
qp.producer_acquire(qp_stage)
cute.copy(tma_q, tAgQ[(None, Int32(k_sub))], tAsQ[(None, qp_stage.index)], tma_bar_ptr=qp_stage.barrier)
qp.producer_commit(qp_stage); qp_stage.advance()
qh = qp.acquire_and_advance()
cute.copy(tma_q, tAgQ[(None, Int32(k_sub))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
# Load K[k_sub] → sK
kvp.producer_acquire(kvp_stage)
cute.copy(tma_k, tBgK[(None, Int32(k_sub))], tBsK[(None, kvp_stage.index)], tma_bar_ptr=kvp_stage.barrier)
kvp.producer_commit(kvp_stage); kvp_stage.advance()
# Load V[kt] → sV (uses same K/V pipeline)
kvp.producer_acquire(kvp_stage)
cute.copy(tma_v, tVgV[(None, Int32(kt))], tVsV[(None, kvp_stage.index)], tma_bar_ptr=kvp_stage.barrier)
kvp.producer_commit(kvp_stage); kvp_stage.advance()
qp.producer_tail(qp_stage)
kvp.producer_tail(kvp_stage)
kvh = kvp.acquire_and_advance()
cute.copy(tma_k, tBgK[(None, Int32(k_sub))], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
qh = qp.wait_and_advance(); qh.release()
kvh = kvp.wait_and_advance(); pk = cutlass.Boolean(1)
# Load V[kt] → sV
kvh = kvp.acquire_and_advance(pk)
cute.copy(tma_v, tVgV[(None, Int32(kt))], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
kvh = kvp.wait_and_advance(); pk = cutlass.Boolean(1)
qp.tail()
kvp.tail()
else:
# Original pipeline path (hd≤256)
qp.reset(); qh = qp.acquire_and_advance()
@@ -250,8 +250,6 @@ class FmhaKernel:
tmem.wait_for_alloc()
if const_expr(self.n_k_sub_tiles > 1):
# K sub-tiling path (hd=512): pipeline sync for k_sub iterations
qp_stage = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.q_stage)
kvp_stage = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.kv_stage)
for kt in range(self.n_kv_tiles):
for k_sub in range(self.n_k_sub_tiles):
qh = qc.wait_and_advance(); qh.release()
@@ -260,6 +258,7 @@ class FmhaKernel:
for kb in cutlass.range(cute.size(tCrQ, mode=[2]), unroll_full=True):
cute.gemm(qk_mma, tStS0, tCrQ[(None,None,kb,0)], tCrK[(None,None,kb,kvh.index)], tStS0)
qk_mma.set(tcgen05.Field.ACCUMULATE, True)
kvh.release()
# After all k_sub: S has full QK for this kt
cute.arch.fence_view_async_tmem_store()
softmax_done_bar.arrive()
@@ -267,14 +266,13 @@ class FmhaKernel:
pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0)
if not self.use_smem_p:
for kb in cutlass.range(cute.size(tOrP0, mode=[2]), unroll_full=True):
cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0)
cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kvb.index)], tOtO0)
pv_mma.set(tcgen05.Field.ACCUMULATE, True)
else:
for kb in cutlass.range(cute.size(tCrP, mode=[2]), unroll_full=True):
cute.gemm(pv_mma, tOtO0, tCrP[(None,None,kb,0)], tCrV[(None,None,kb,kvh.index)], tOtO0)
cute.gemm(pv_mma, tOtO0, tCrP[(None,None,kb,0)], tCrV[(None,None,kvb.index)], tOtO0)
pv_mma.set(tcgen05.Field.ACCUMULATE, True)
cute.arch.fence_view_async_tmem_store()
kvh.release()
final_o_bar.arrive()
else:
# Original pipeline path (hd≤256)