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.
This commit is contained in:
2026-05-26 20:31:28 +00:00
parent 34d64137ec
commit 20ed6d5114

View File

@@ -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