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