diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 4ad7338a..a0a4e30e 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -410,46 +410,22 @@ class FmhaKernel: scale_log2 = Float32(self.scale_softmax_log2) # ============================================================ - # D1.5: O RESCALE ATOMS (using tCtO_base — proven-correct TMEM layout) + # 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. # ============================================================ - # 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) - # 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) + # 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. # prev_acc_scale: unused, kept for clarity. acc_scale at kt is used # to rescale O from kt=0..kt-1 before PV[kt]. @@ -553,33 +529,12 @@ 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 — 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() + # 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. si_handle.release() softmax_done_bar.arrive()