D1: Fix kvh scoping - define before loops, consume V via pipeline
This commit is contained in:
@@ -247,6 +247,7 @@ 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
|
||||
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):
|
||||
qh = qc.wait_and_advance(); qh.release()
|
||||
@@ -261,6 +262,8 @@ class FmhaKernel:
|
||||
softmax_done_bar.arrive()
|
||||
softmax_done_bar.arrive_and_wait()
|
||||
pv_mma.set(tcgen05.Field.ACCUMULATE, kt != 0)
|
||||
# Load V: consume from K/V pipeline
|
||||
kvh = kvc.wait_and_advance()
|
||||
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)
|
||||
@@ -270,6 +273,7 @@ class FmhaKernel:
|
||||
cute.gemm(pv_mma, tOtO0, tCrP[(None,None,kb,0)], tCrV[(None,None,kb,kvh.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)
|
||||
|
||||
Reference in New Issue
Block a user