From 54f463b8a1f06ea0d6d28ea43bb8d8db87467ffc Mon Sep 17 00:00:00 2001 From: Minho Ryu Date: Tue, 4 Aug 2026 19:39:46 +0900 Subject: [PATCH] Prefer Flash Attn 2 over Fused Attn for THD dropout on Blackwell cuDNN's dropout kernels for THD attention are much slower than FlashAttention 2's on SM100/103, and the cuDNN team confirmed a fix is not on their roadmap. The generic Hopper+ rule prefers FusedAttention, so this combination silently took the slow path with nothing reporting it. Prefer FA2 for THD training with dropout on SM100/103 when FA2 is confirmed usable, and leave every other configuration on the existing rule. The earlier advisory is dropped: it ran before backend selection was final and could name FusedAttention when a later filter selected something else. Signed-off-by: Minho Ryu --- tests/pytorch/attention/test_attention.py | 58 +++++++++++++++++++ .../attention/dot_product_attention/utils.py | 20 ++++++- 2 files changed, 76 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 82791084d8..f76a4d90b3 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -28,8 +28,10 @@ _attention_backends, ) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( + AttentionParams, FlashAttentionUtils, check_set_window_size, + get_attention_backend, ) from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext @@ -322,6 +324,62 @@ def test_dpa_num_splits(dtype, model_configs, model): ) +@pytest.mark.skipif( + device_compute_capability not in ((10, 0), (10, 3)), + reason="This backend preference applies only to SM100/SM103.", +) +@pytest.mark.skipif( + not FlashAttentionUtils.is_installed, + reason="A supported FlashAttention 2 installation is required.", +) +@pytest.mark.skipif( + get_cudnn_version() < (9, 0, 0), + reason="cuDNN 9.0.0+ is required for THD FusedAttention.", +) +def test_thd_dropout_prefers_flash_attention_2(monkeypatch): + """THD training with dropout must select FA2 over FusedAttention on SM100/103. + + cuDNN's dropout kernels for this combination are far slower than FA2's and are not on + the roadmap, so the choice is made here rather than left to the generic Hopper+ rule + that prefers FusedAttention. + """ + monkeypatch.setenv("NVTE_FLASH_ATTN", "1") + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + + attention_params = AttentionParams( + qkv_dtype=torch.bfloat16, + qkv_layout="thd_thd_thd", + batch_size=4, + num_heads=16, + num_gqa_groups=16, + max_seqlen_q=1024, + max_seqlen_kv=1024, + head_dim_qk=128, + head_dim_v=128, + attn_mask_type="padding_causal", + core_attention_bias_shape=None, + attention_dropout=0.1, + is_training=True, + ) + ( + use_flash_attention, + flash_attention_backend, + use_fused_attention, + fused_attention_backend, + _, + available_backends, + ) = get_attention_backend(attention_params) + + # Assert availability first: without both candidates the preference is not exercised. + assert available_backends[0], "FlashAttention is unavailable for this configuration" + assert available_backends[1], "FusedAttention is unavailable for this configuration" + + assert use_flash_attention + assert flash_attention_backend == FlashAttentionUtils.version + assert not use_fused_attention + assert fused_attention_backend is None + + # ============================== # Flash Attention 4 (FA4) tests # ============================== diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6eb3ce54f1..61338f6660 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1615,8 +1615,24 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt bool(available_backends[2]), ) - # Select FusedAttention for performance - if use_flash_attention and use_fused_attention and device_compute_capability >= (9, 0): + # Prefer FA2 for THD training with dropout on SM100/103, where FusedAttention has a known + # performance issue. At this point use_flash_attention_2 confirms a usable installation. + if ( + is_training + and qkv_format == "thd" + and attention_dropout != 0.0 + and device_compute_capability in ((10, 0), (10, 3)) + and use_flash_attention_2 + and use_fused_attention + ): + logger.debug( + "Disabling FusedAttention to give FlashAttention 2 preference for THD with dropout" + " on SM100/103" + ) + use_fused_attention = False + fused_attention_backend = None + # Select FusedAttention for performance in all other Hopper+ configurations. + elif use_flash_attention and use_fused_attention and device_compute_capability >= (9, 0): logger.debug( "Disabling FlashAttention to give FusedAttention preference on Hopper+ " "for performance reasons"