fix: add missing old_row_max = row_max before softmax max computation

2026-05-22 09:30:32 +00:00
parent ce06478e56
commit 3aa597460e
1 changed files with 1 additions and 0 deletions
--- a/tests/unit/test_fmha_v3_stage_c_full.py
+++ b/tests/unit/test_fmha_v3_stage_c_full.py
@@ -248,6 +248,7 @@ class FmhaV3StageC:
                cute.arch.fence_view_async_tmem_load()

                # Compute row_max from S (element-wise, before exp2)
+                old_row_max = row_max
                frg_cnt = 4
                frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt
                tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile))