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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user