fix: use cute.arch.fmax instead of if-else in vectorized loop

2026-05-22 09:28:32 +00:00
parent 941bcae8e1
commit f631ff16d6
1 changed files with 1 additions and 1 deletions
--- a/tests/unit/test_fmha_v3_stage_c_full.py
+++ b/tests/unit/test_fmha_v3_stage_c_full.py
@@ -258,7 +258,7 @@ class FmhaV3StageC:
                for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True):
                    for j in range(frg_cnt):
                        s_val = tTMEM_LOADrS_frg[k, j] * scale_log2
-                        if s_val > row_max: row_max = s_val
+                        row_max = cute.arch.fmax(row_max, s_val)

                row_max_safe = row_max
                if row_max == -cutlass.Float32.inf: row_max_safe = Float32(0.0)