From 4248589c614ce21b3da8f4f36e258378b2aa5261 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 01:27:03 +0000 Subject: [PATCH] diag: skip final normalize, test raw PV output via epilogue_tma_store --- tests/unit/test_fmha_v3_stage_c.py | 31 ++---------------------------- 1 file changed, 2 insertions(+), 29 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 62f7f997..899e036a 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -407,35 +407,8 @@ class FmhaV3StageCMulti: # 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 === - inv_row_sum = Float32(1.0) / row_sum - - 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( - 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 j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i) - - cute.arch.fence_view_async_tmem_store() - - # Standard epilogue: TMEM → SMEM → GMEM via TMA store. - # O in TMEM is now scaled by 1/row_sum. + # DIAG: skip final normalize, just use epilogue_tma_store directly + # to test raw PV output tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) acc_cons_st = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_acc_stage