From fae9f6fbb5c10235f23e71cd1b9dc23efb6ff2bd Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 00:27:41 +0000 Subject: [PATCH] Reset to working_softmax_maybe.py + TMA fix only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previous O rescale attempt broke n=128 (0.464773). Revert to known-good softmax code, only apply TMA fix: tBgK[(None,None,0,0)] → tBgK[(None,0,None,0)] Expected: n=128 cos 0.999998 (same as working), n=256 cos 0.71 (TMA fix loads 2 tiles but no O rescale) --- tests/unit/test_fmha_v3_stage_c.py | 129 ++++++++++++++--------------- 1 file changed, 61 insertions(+), 68 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index e8dfc384..c507ad72 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -180,13 +180,7 @@ class FmhaV3StageCMulti: b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape) tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) - # After tma_partition, tBgK has 4 modes: (V_grouped, ?, KV_tiles, ?). - # Mode 4 is the GMEM tile iteration axis (size = n_kv_tiles). - # 8-None no-op pre-slice opens the full TMA coord space. - # tAgQ is fine with 4-mode slice (Q has only 1 tile). - tAgQ = tAgQ[(None,0,None,0)] - tBgK = tBgK[(None,0,None,0)] - tVgV = tVgV[(None,0,None,0)] + tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) tCrV = pv_mma.make_fragment_B(sV) @@ -210,13 +204,9 @@ class FmhaV3StageCMulti: pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk) # ===== TMA LOAD warp ===== - # GMEM tile coordinate: use the cutlass.range induction variable kt - # directly. CuTeDSL's `cutlass.range` doesn't auto-detect a Python `+=` - # rebinding as a loop-carried iter_args update — the JIT traces the - # body once and captures whatever value `kv_coord` had at trace time, - # so an outer `kv_coord = Int32(0)` plus a `kv_coord += 1` inside the - # loop bakes 0 into every iteration's TMA descriptor at runtime. - # The induction variable IS the loop-carried state, properly tracked. + # NOTE: using kt from cutlass.range works for n=128 (single tile). + # Multi-tile (n>128) loads from tile 0 only — the JIT constant-folds kt. + # TODO: fix multi-tile TMA indexing (kv_coord pattern from diag test). if warp_idx == self.tma_warp_id: qp.reset(); qh = qp.acquire_and_advance() cute.copy(tma_q, tAgQ[(None, Int32(0))], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier) @@ -224,8 +214,8 @@ class FmhaV3StageCMulti: kvp.reset(); pk = kvp.try_acquire() for kt in cutlass.range(0, n_kv_tiles, 1, unroll=1): kvh = kvp.acquire_and_advance(pk) - cute.copy(tma_k, tBgK[None, kt], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) - cute.copy(tma_v, tVgV[None, kt], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) + cute.copy(tma_k, tBgK[(None, kt)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier) + cute.copy(tma_v, tVgV[(None, kt)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier) pk = cutlass.Boolean(1) kvp.tail() @@ -238,7 +228,7 @@ class FmhaV3StageCMulti: kvc.reset(); pk = kvc.try_wait() acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) acc_pipe.producer_acquire(acc_st) - for kt in range(self.n_kv_tiles): + for kt in range(n_kv_tiles): kvh = kvc.wait_and_advance(pk); pk = cutlass.Boolean(1) sh = s_prod.acquire_and_advance() qk_mma.set(tcgen05.Field.ACCUMULATE, False) @@ -296,31 +286,7 @@ class FmhaV3StageCMulti: row_sum = Float32(0.0) scale_log2 = Float32(self.scale_softmax_log2) - # === O rescale / correction_rescale setup (needed inside softmax loop) === - corr_tile_size = 16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - 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.acc_dtype, - ) - tmem_store_o_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), - self.acc_dtype, - ) - tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) - thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i) - tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i) - tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i) - - # Per-tile softmax loop with online O rescale. + # Per-tile softmax loop. # Online softmax row_max/row_sum tracking is maintained, but the # in-place TMEM O rescale (which would multiply existing O by # exp2(old_max - new_max) before PV[kt]) is DISABLED — this is the @@ -330,7 +296,7 @@ class FmhaV3StageCMulti: # For now: kernel is correct when row_max growth across tiles is # mild (typical for short n with random data); for very long n # the missing rescale shows as accuracy drift. - for kt in range(self.n_kv_tiles): + for kt in range(n_kv_tiles): si_handle = s_cons.wait_and_advance() # Load S[kt] @@ -378,41 +344,68 @@ class FmhaV3StageCMulti: cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) cute.arch.fence_view_async_tmem_store() - # === Per-tile O rescale: O *= acc_scale for kt > 0 === - # Must run BEFORE softmax_done_bar.arrive() so MMA's PV[kt] - # reads the rescaled O. Uses same paired atoms as final normalize. - if kt > 0: - for i in range(HEAD_DIM // corr_tile_size): - 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, - ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - si_handle.release() softmax_done_bar.arrive() - # Wait for MMA's PV[N-1] to commit before reading O for normalize. + # === Reference-style scaled epilogue (no TMEM round-trip) === + # + # Pattern (mirrors CUTLASS Blackwell FMHA reference's + # correction_epilog): for each column sub-tile, + # 1. TMEM -> registers via PAIRED tmem_load atom + # 2. scale in registers (1/row_sum) + # 3. FP32 -> BF16 conversion in registers + # 4. registers -> SMEM via PAIRED smem_store atom + # Then TMA SMEM -> GMEM as a separate step. + # + # Critical: the load and store atoms MUST be a matched pair. + # Independently constructed Ld32x32bOp + St32x32bOp atoms (the + # previous code) don't preserve the register tile shape, so even a + # no-op load+store corrupts data. Using utils.blackwell_helpers + # (sm100_utils) gives a paired set keyed to the same epi_subtile. + + # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # === Final O normalization: O *= 1/row_sum === - # Uses the working 2D register tensor pattern from working_softmax_maybe.py. - inv_row_sum = Float32(1.0) / row_sum + # === O normalization via TMEM load → scale → TMEM store === + # Matches CUTLASS reference's correction_rescale pattern exactly. + corr_tile_size = 16 + + cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + tOcO = pv_thr.partition_C(cO) + + tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) + tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) + + tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) + tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.acc_dtype, + ) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.acc_dtype, + ) + + tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_i) + tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_i) + + thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx) + thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx) + + tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i) + tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i) + tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i) + + # 2D register tensor: (frg_shape, n_corr_tiles) tTMrO = cute.make_rmem_tensor( (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype ) + inv_row_sum = Float32(1.0) / row_sum + for i in range(HEAD_DIM // corr_tile_size): tTMrO_i_ = tTMrO[None, i] tTMrO_i_layout = cute.composition(