diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 1b934af323..d4640ac0b2 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -33,11 +33,13 @@ _attention_backends, ) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( + AttentionParams, FlashAttentionUtils, FusedAttentionParams, _get_supported_versions, check_set_window_size, get_fused_attn_spec, + get_attention_backend, ) from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext @@ -212,6 +214,49 @@ def test_fused_attn_backend_message(): assert message != "" +@pytest.mark.parametrize("head_dim_qk,head_dim_v", [(192, 128), (64, 512)]) +@pytest.mark.parametrize("is_training,expected_fa3", [(True, False), (False, True)]) +def test_fa3_mismatched_head_dims_mode_selection( + monkeypatch, head_dim_qk, head_dim_v, is_training, expected_fa3 +): + """FA3 supports mismatched head dimensions only for forward-only execution.""" + monkeypatch.setattr( + "transformer_engine.pytorch.attention.dot_product_attention.utils.get_device_compute_capability", + lambda: (9, 0), + ) + monkeypatch.setattr(FlashAttentionUtils, "v3_is_installed", True) + monkeypatch.setattr(FlashAttentionUtils, "fa3_version", PkgVersion("3.0.0b1")) + monkeypatch.setenv("NVTE_FLASH_ATTN", "1") + monkeypatch.setenv("NVTE_FLASH_ATTN_V2", "0") + monkeypatch.setenv("NVTE_FLASH_ATTN_V3", "1") + monkeypatch.setenv("NVTE_FLASH_ATTN_V4", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN", "0") + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "1") + + ( + use_flash_attention, + flash_attention_backend, + use_fused_attention, + _, + use_unfused_attention, + available_backends, + ) = get_attention_backend( + AttentionParams( + qkv_dtype=torch.bfloat16, + qkv_layout="bshd_bshd_bshd", + head_dim_qk=head_dim_qk, + head_dim_v=head_dim_v, + is_training=is_training, + ) + ) + + assert bool(use_flash_attention) == expected_fa3 + assert bool(available_backends[0]) == expected_fa3 + assert not use_fused_attention + assert bool(use_unfused_attention) != expected_fa3 + assert flash_attention_backend == (FlashAttentionUtils.fa3_version if expected_fa3 else None) + + # Define F16 data types to test param_types = [torch.float16] if is_bf16_available(): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 23b89287df..c094027c5e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1084,6 +1084,17 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) use_flash_attention_3 = False + if use_flash_attention_3 and is_training and head_dim_qk != head_dim_v: + logger.debug( + "Disabling FlashAttention 3 for training with head_dim_qk != head_dim_v, " + "as its backward pass does not support it " + "(Dao-AILab/flash-attention#1487). " + "Found: head_dim_qk = %s, head_dim_v = %s.", + head_dim_qk, + head_dim_v, + ) + use_flash_attention_3 = False + if ( use_flash_attention_4 and FlashAttentionUtils.v4_is_installed