diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index e59405db74..a87dd12a6b 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -1569,12 +1569,15 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
         # was observed to grow super-linearly (B=4 took ~4x the B=2 amount, not 2x) —
         # revisit if a config uses B>2.
         SM90_DET_FUSED_THD_BWD_MAX_BHSS = 1 << 30
-        if (
-            use_fused_attention
-            and fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value
+        is_sm90_f16_thd_training = (
+            fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value
             and is_training
             and qkv_format == "thd"
             and device_compute_capability == (9, 0)
+        )
+        if (
+            use_fused_attention
+            and is_sm90_f16_thd_training
             and batch_size * num_heads * max_seqlen_q * max_seqlen_kv
             >= SM90_DET_FUSED_THD_BWD_MAX_BHSS
         ):