From 18ab5078969023cc6b44dc9fcfc6660b267adb3c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 03:40:12 +0000 Subject: [PATCH] revert: fmha.py to D1 (0f52c34, before D1.5 changes) --- dsv4/kernels/attention/fmha.py | 73 ++++++++++++---------------------- 1 file changed, 26 insertions(+), 47 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 46d336c6..37ed7a68 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -319,6 +319,12 @@ class FmhaKernel: tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i) n_corr_tiles = self.pv_n_tile // corr_tile_size + # tTMrO register tensor (defined unconditionally for CuTeDSL scoping). + # Used for O rescale (kt > 0) and O normalization (after loop). + tTMrO = cute.make_rmem_tensor( + (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype + ) + for kt in range(self.n_kv_tiles): si_handle = s_cons.wait_and_advance() @@ -363,9 +369,6 @@ class FmhaKernel: cute.arch.fence_view_async_tmem_store() else: # SMEM-P: write P to sP using coordinate-indexed store. - # Uses tTMEM_LOADcS identity tensor to get (m, k) coordinates. - # DEBUG: Write a known pattern to sP to verify the coordinate mapping. - # Pattern: sP[m, k] = (m + k) % 256 as BF16 (unique per position) for j0 in range(32): for j1 in range(4): coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0] @@ -374,13 +377,9 @@ class FmhaKernel: k0 = k_coord % 16 k1 = (k_coord // 16) % 4 k2 = k_coord // 64 - # Debug: write (m + k) mod 256 instead of actual P value _sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16[(j0, 0), j1, 0, 0] cute.arch.fence_proxy("async.shared", space="cta") if kt > 0: - tTMrO = cute.make_rmem_tensor( - (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype - ) for i in range(n_corr_tiles): tTMrO_i_ = tTMrO[None, i] tTMrO_i_layout = cute.composition( @@ -408,39 +407,20 @@ class FmhaKernel: final_o_bar.arrive_and_wait() # ============================================================ - # EPILOGUE: Normalize O + TMA store to GMEM + # EPILOGUE: TMA store O to GMEM + compute LSE # ============================================================ - # Step 1: Normalize O in TMEM via round-trip (3% error from hand-constructed - # atoms — D1.5 tracks the paired-atom fix). - # Step 2: Use CUTLASS epilogue_tma_store for TMEM→SMEM→GMEM write. + # The raw un-normalized O in TMEM is perfect (cos 0.999998). + # TMEM round-trip normalization with hand-constructed atoms causes + # severe data corruption (53% error) due to layout mismatch with + # epilogue_tma_store's paired-atom addressing. + # Solution: always write raw O via epilogue_tma_store, compute LSE, + # and let the caller normalize externally using LSE. + # This is the D5a path — production-quality with zero precision loss. + # The TMEM round-trip normalization (normalize=True) is tracked as D1.5. # ============================================================ - # D5a: When normalize=False, skip 1/row_sum (emit un-normalized O + LSE). - if const_expr(self.normalize): - inv_row_sum = Float32(1.0) / row_sum - # Normalize O: TMEM round-trip O *= inv_row_sum - for i in range(n_corr_tiles): - 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) - for k in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[k] = tTMrO_i[k] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # TMA store via CUTLASS epilogue_tma_store - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tOtO.layout) + # TMA store via CUTLASS epilogue_tma_store (reads raw O from TMEM) + tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) acc_cons_st = pipeline.make_pipeline_state( @@ -453,17 +433,16 @@ class FmhaKernel: ) c_pipe.producer_tail() - # D5a: Write LSE (log-softmax) when normalize=False - # lse = ln(row_sum) + row_max * ln(2) + # Compute LSE: lse = ln(row_sum) + row_max * ln(2) + # Always compute LSE (needed for external normalization). # row_max is in scale_log2 domain, multiply by ln(2) to convert. - if const_expr(not self.normalize): - _row_max_safe = row_max - if row_max == -cutlass.Float32.inf: - _row_max_safe = Float32(0.0) - if sfw_idx == 0: - _ln2 = Float32(0.6931471805599453) # ln(2) - lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2 - mLSE[0] = lse_val + _row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + _row_max_safe = Float32(0.0) + if sfw_idx == 0: + _ln2 = Float32(0.6931471805599453) # ln(2) + lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2 + mLSE[0] = lse_val tmem.relinquish_alloc_permit() tmem.free(tmem_ptr)