From f631ff16d683309ac7cca83ef2f1a9baa527707d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 22 May 2026 09:28:32 +0000 Subject: [PATCH] fix: use cute.arch.fmax instead of if-else in vectorized loop --- tests/unit/test_fmha_v3_stage_c_full.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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)