diff --git a/dsv4/kernels/attention/fmha_smem_acc.py b/dsv4/kernels/attention/fmha_smem_acc.py index 18007057..a127f078 100644 --- a/dsv4/kernels/attention/fmha_smem_acc.py +++ b/dsv4/kernels/attention/fmha_smem_acc.py @@ -458,25 +458,28 @@ class FmhaKernel: # ======================================================== # SMEM ACCUMULATOR: load O_kt from TMEM, accumulate in SMEM # ======================================================== - # O_kt is in TMEM (PV with ACCUMULATE=False → fresh output) - # Load via one-way Ld32x32bOp (read-only, NO write-back to TMEM) - # Then: sO_acc = acc_scale * sO_acc + O_kt - # Using coordinate-indexed writes to sO_acc + # Only needed for n_kv_tiles > 1. + # For n_kv_tiles=1, epilogue_tma_store reads directly from TMEM. + # + # WARNING: Ld32x32bOp may be a destructive TMEM read! + # If so, the O load here would consume TMEM data, + # making epilogue_tma_store read garbage. + # Guard with const_expr to avoid this for n_kv_tiles=1. # ======================================================== - rO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.qk_acc_dtype) - cute.copy(tiled_o_load, tTMEM_LOADtO, rO) - cute.arch.fence_view_async_tmem_load() + if const_expr(self.n_kv_tiles > 1): + rO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.qk_acc_dtype) + cute.copy(tiled_o_load, tTMEM_LOADtO, rO) + cute.arch.fence_view_async_tmem_load() - # Rescale existing sO_acc and add O_kt - # Use coordinate tensor to map each register to (row, col) in sO_acc - for j0 in range(32): - for j1 in range(4): - coord = tTMEM_LOADcO[(j0, 0), j1, 0, 0] - row = coord[0] - col = coord[1] - old_val = sO_acc[row, col] - new_val = acc_scale * old_val + rO[(j0, 0), j1, 0, 0] - sO_acc[row, col] = new_val + # Rescale existing sO_acc and add O_kt + for j0 in range(32): + for j1 in range(4): + coord = tTMEM_LOADcO[(j0, 0), j1, 0, 0] + row = coord[0] + col = coord[1] + old_val = sO_acc[row, col] + new_val = acc_scale * old_val + rO[(j0, 0), j1, 0, 0] + sO_acc[row, col] = new_val # Wait for MMA's final signal final_o_bar.arrive_and_wait()