From 20ed6d51146620b5cd0950a745e9bf24c100e2b0 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 26 May 2026 20:31:28 +0000 Subject: [PATCH] D1.5: Add TMEM load fence before PV with ACCUMULATE, revert debug rescale factor The MMA warp needs fence_view_async_tmem_load() before PV[kt>0] to ensure the rescaled O values are visible. NamedBarrier synchronizes warps but may not guarantee TMEM visibility without an explicit fence. --- dsv4/kernels/attention/fmha.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index a0e81c94..1513254d 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -338,6 +338,8 @@ class FmhaKernel: cute.arch.fence_view_async_tmem_store() sh.commit() softmax_done_bar.arrive_and_wait() + if const_expr(self.n_kv_tiles > 1): + cute.arch.fence_view_async_tmem_load() # Ensure rescaled O visible before PV[kt] 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): @@ -558,7 +560,7 @@ class FmhaKernel: pv_done_bar.arrive_and_wait() # Wait for PV[kt-1] # Rescale O: load, multiply by acc_scale, store back to TMEM. # CUTLASS pattern: both copies use same tOtO_i (composition-tiled). - rescale_factor = Float32(0.5) # DEBUG: force known value + rescale_factor = acc_scale if const_expr(self.debug_noop_rescale): rescale_factor = Float32(1.0) n_slices = self.head_dim // corr_tile_size