diff --git a/tests/unit/test_fmha_v3_stage_c_full.py b/tests/unit/test_fmha_v3_stage_c_full.py
index 6f086071..91a9f5c8 100644
--- 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)