From 2bf3ee40aa9aa4b823f927ba0c18ac507f41c11e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 06:42:26 +0000 Subject: [PATCH] D1: Fix kvh scoping - define before loops, consume V via pipeline --- dsv4/kernels/attention/fmha.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 5d8ca0b2..29fdf6c7 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -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)