diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index a0a4e30e..4ad7338a 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -410,22 +410,46 @@ class FmhaKernel: scale_log2 = Float32(self.scale_softmax_log2) # ============================================================ - # D1.5: O RESCALE — SMEM ACCUMULATOR APPROACH - # ================================================= - # TMEM round-trip (Ld32x32bOp/St32x32bOp) is FUNDAMENTALLY broken: - # even NO-OP round-trip corrupts data (ratio = -11 billion). - # Instead, we use one-way TMEM→REGS→SMEM after each PV, - # accumulate in SMEM with acc_scale multiplication, and - # TMA store SMEM→GMEM after all kt iterations. - # - # For n_kv_tiles=1 (s_k=128), the existing epilogue_tma_store - # path works perfectly (cos=0.999998). The SMEM accumulator - # is only needed for n_kv_tiles > 1. + # D1.5: O RESCALE ATOMS (using tCtO_base — proven-correct TMEM layout) # ============================================================ + # Use tCtO_base (from epilogue, proven to correctly read O from TMEM) + # for BOTH load and store atoms. Both copies built from same tensor + # so register layouts are compatible. + # ============================================================ + tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - # NOTE: The code below is the BROKEN TMEM round-trip approach. - # It's kept as reference but should NOT be used. - # The SMEM accumulator implementation is TODO. + # Build TMEM load+store atoms from tCtO_base via composition + corr_tile_size = 16 + tCtO_i_layout = cute.composition( + tCtO_base.layout, cute.make_layout((128, corr_tile_size)) + ) + tCtO_i = cute.make_tensor(tCtO_base.iterator, tCtO_i_layout) + + # Coordinate tensor for partition_D of load + cO = cute.make_identity_tensor((128, self.head_dim)) + # Use pv_mma (not pv_thr) for partition_C — matches tCtO_base's layout + tOcO = pv_mma.get_slice(0).partition_C(cO) + tOcO_i_layout = cute.composition( + tOcO.layout, cute.make_layout((128, corr_tile_size)) + ) + tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) + + tmem_load_o_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.qk_acc_dtype, + ) + tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tCtO_i) + thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx) + tTMEM_LOADtO = thr_tmem_load_o.partition_S(tCtO_i) + tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i) + + tmem_store_o_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.qk_acc_dtype, + ) + tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tCtO_i) + thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx) + tTMEM_STOREtO = thr_tmem_store_o.partition_D(tCtO_i) # prev_acc_scale: unused, kept for clarity. acc_scale at kt is used # to rescale O from kt=0..kt-1 before PV[kt]. @@ -529,12 +553,33 @@ class FmhaKernel: k2 = k_coord // 64 _sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16[(j0, 0), j1, 0, 0] cute.arch.fence_proxy("async.shared", space="cta") - # D1.5: O rescale for kt > 0 — NOT YET IMPLEMENTED. - # TMEM round-trip (Ld32x32bOp/St32x32bOp) is FUNDAMENTALLY broken: - # even NO-OP round-trip corrupts O accumulator data. - # Production path for multi-KV-tile: Python KV merge (cos 0.999998). - # Future: SMEM accumulator approach (one-way TMEM→REGS→SMEM per kt). - # n_kv_tiles=1 is the only supported path for in-kernel processing. + # D1.5: O rescale for kt > 0 — using tCtO_base (proven-correct TMEM layout). + # Both load and store atoms built from same tCtO_i tensor via composition. + if const_expr(self.n_kv_tiles > 1): + if kt > 0: + pv_done_bar.arrive_and_wait() # Wait for PV[kt-1] + n_slices = self.head_dim // corr_tile_size + tTMrO = cute.make_rmem_tensor( + (tTMEM_LOADcO.shape, n_slices), self.qk_acc_dtype + ) + for i in range(n_slices): + tTMrO_i_ = tTMrO[None, i] + tTMrO_i_layout = cute.composition( + tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0]) + ) + tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) + tTMEM_LOADtO_i = cute.make_tensor( + tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout + ) + tTMEM_STOREtO_i = cute.make_tensor( + tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout + ) + cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i) + cute.arch.fence_view_async_tmem_load() + for k in cutlass.range(cute.size(tTMrO_i), vectorize=True): + tTMrO_i[k] = tTMrO_i[k] * acc_scale + cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i) + cute.arch.fence_view_async_tmem_store() si_handle.release() softmax_done_bar.arrive()