diff --git a/qa/L0_pytorch_debug_unittest/test.sh b/qa/L0_pytorch_debug_unittest/test.sh index 3efa462628..36efe485f5 100644 --- a/qa/L0_pytorch_debug_unittest/test.sh +++ b/qa/L0_pytorch_debug_unittest/test.sh @@ -22,6 +22,11 @@ FAILED_CASES="" : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" +# L0 keeps one mature FlashAttention generation; L3 owns newer-generation coverage. +export NVTE_FLASH_ATTN_V2=1 +export NVTE_FLASH_ATTN_V3=0 +export NVTE_FLASH_ATTN_V4=0 + # Config with the dummy feature which prevents nvinspect from being disabled. # Nvinspect will be disabled if no feature is active. : ${NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE:=$TE_PATH/tests/pytorch/debug/test_configs/dummy_feature.yaml} diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 759432857a..ee3ef89e66 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -22,6 +22,11 @@ set -x : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" +# L0 keeps one mature FlashAttention generation; L3 owns newer-generation coverage. +export NVTE_FLASH_ATTN_V2=1 +export NVTE_FLASH_ATTN_V3=0 +export NVTE_FLASH_ATTN_V4=0 + pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 90ff6ba2fd..ec19492ee7 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -24,20 +24,22 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" # Run CP tests (deterministic + non-deterministic) first so they can be parallelized. # Each needs 4 GPUs, so >=8 GPUs allows them to run concurrently on disjoint GPU sets. +# Main's CP implementation supports FA2/FA3. Keep FA4 disabled at the suite +# boundary so both the reference and CP halves use the same backend generation. NUM_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())") echo "Detected $NUM_GPUS GPU(s)" if [ "$NUM_GPUS" -ge 8 ]; then echo "Running CP tests in parallel: non-deterministic on GPUs 0-3, deterministic on GPUs 4-7" - CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & + CUDA_VISIBLE_DEVICES=0,1,2,3 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & PID_CP_NONDET=$! - CUDA_VISIBLE_DEVICES=4,5,6,7 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & + CUDA_VISIBLE_DEVICES=4,5,6,7 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & PID_CP_DET=$! wait $PID_CP_NONDET || test_fail "test_attention_with_cp.py" wait $PID_CP_DET || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" else echo "Running CP tests sequentially: need >=8 GPUs for parallel execution" - python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py" - NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" + NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py" + NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" fi python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py" diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index 468ef04d76..047358b301 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -31,37 +31,57 @@ export NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 # Iterate over Flash Attention versions sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); print(sm[0]*10+sm[1])"` export FLASH_ATTN_CUDA_ARCHS=$sm_arch -# CP tests are expensive and run only once per arch: -# - sm90 (H100): FA3 (3.0.0b1) - context_parallel.py only supports FA3 on Hopper -# - sm>90 (B200): latest FA4 - FA3 is not built/installed for sm>90 -# Non-CP tests still run for every FA version in the array. +# Run one architecture-owned FlashAttention generation. CP remains FA3-only +# until the production selector and runner support FA4 CP end to end. +CP_FA_VERSION="" if [ $sm_arch -gt 90 ] then - FA_versions=(2.8.3 4.0.0b11) - CP_FA_VERSION="${FA_versions[-1]}" + FA_versions=(4.0.0b11) elif [ $sm_arch -eq 90 ] then - FA_versions=(2.8.3 3.0.0b1 4.0.0b11) + FA_versions=(3.0.0b1) CP_FA_VERSION="3.0.0b1" +else + error_exit "No L3 FlashAttention generation is defined for sm${sm_arch}" fi for fa_version in "${FA_versions[@]}" do + # The FA distributions share the flash_attn namespace. Keep exactly one + # installed so import-time discovery and the iteration label cannot disagree. + pip3 uninstall -y flash-attn flash-attn-3 flash-attn-4 \ + || error_exit "Failed to isolate Flash Attention $fa_version" + export NVTE_FLASH_ATTN_V2=0 + export NVTE_FLASH_ATTN_V3=0 + export NVTE_FLASH_ATTN_V4=0 + # Build Flash Attention if [ "${fa_version}" \< "3.0.0" ] then - pip3 install flash-attn==${fa_version} --no-build-isolation + export NVTE_FLASH_ATTN_V2=1 + pip3 install flash-attn==${fa_version} --no-build-isolation \ + || error_exit "Failed to install Flash Attention $fa_version" elif [[ "${fa_version}" == 4.* ]] then - pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 --no-build-isolation + export NVTE_FLASH_ATTN_V4=1 + # FA4 is intentionally last in every version array. Its b11 test pin needs + # CUTLASS DSL 4.4.2, so replace the image-matched stack only for this final + # iteration; later iterations would otherwise need that stack restored. + pip3 uninstall -y nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base \ + nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 \ + || error_exit "Failed to isolate CUTLASS DSL for Flash Attention $fa_version" + pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 \ + --no-build-isolation || error_exit "Failed to install Flash Attention $fa_version" else + export NVTE_FLASH_ATTN_V3=1 # FA3 source build (~20 min). Skip if FA3 is already installed. if python3 -c "import flash_attn_3" 2>/dev/null; then echo "FA3 already installed (from base image); skipping source build" else git clone https://github.com/Dao-AILab/flash-attention.git - cd flash-attention/hopper && python setup.py install + cd flash-attention/hopper && python setup.py install \ + || error_exit "Failed to install Flash Attention $fa_version" cd ../../ fi fi @@ -82,7 +102,7 @@ do # test_attention.py reloads its own trusted delayed-scaling FP8 checkpoint, # whose legacy extra state requires an explicit pickle opt-in. - if [ "$fa_version" = "$CP_FA_VERSION" ]; then + if [ -n "$CP_FA_VERSION" ] && [ "$fa_version" = "$CP_FA_VERSION" ]; then echo "Running CP tests with FA $fa_version (CP version for sm$sm_arch)" if [ "$NUM_GPUS" -ge 5 ]; then CP_NUM_GPUS=$(( NUM_GPUS - 1 > 4 ? 4 : NUM_GPUS - 1 )) @@ -107,7 +127,11 @@ do NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_CP $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py (FA $fa_version)" fi else - echo "Skipping CP tests for FA $fa_version (CP only runs with FA $CP_FA_VERSION on sm$sm_arch)" + if [ -n "$CP_FA_VERSION" ]; then + echo "Skipping CP tests for FA $fa_version (CP uses FA $CP_FA_VERSION on sm$sm_arch)" + else + echo "CP tests are not scheduled for the FA generation on sm$sm_arch" + fi NVTE_TORCH_COMPILE=0 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)" fi done diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 82791084d8..c4f852f112 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -341,9 +341,18 @@ def test_dpa_num_splits(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." +fa4_enabled = bool(int(os.getenv("NVTE_FLASH_ATTN", "1"))) and bool( + int(os.getenv("NVTE_FLASH_ATTN_V4", "1")) +) +requires_fa4 = pytest.mark.skipif( + not fa4_enabled + or not FlashAttentionUtils.v4_is_installed + or device_compute_capability < (9, 0), + reason="Enabled Flash-attn v4 and compute capability >= SM90 are required.", ) + + +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_base]) @pytest.mark.parametrize("model", model_configs_fa4_base.keys()) @@ -362,9 +371,7 @@ def test_dpa_fa4_base(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.skipif( device_compute_capability not in ((10, 0), (10, 3)), reason="FA4 head_dim=256 dedicated kernel is SM100/103-only.", @@ -442,9 +449,7 @@ def test_dpa_d256(dtype, model_configs, model, qkv_layout): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_mla]) @pytest.mark.parametrize("model", model_configs_fa4_mla.keys()) @@ -466,9 +471,7 @@ def test_dpa_fa4_mla(dtype, model_configs, model): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_swa]) @pytest.mark.parametrize("model", model_configs_fa4_swa.keys()) @@ -489,9 +492,7 @@ def test_dpa_fa4_sliding_window(dtype, model_configs, model, qkv_layout): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_varlen]) @pytest.mark.parametrize("model", model_configs_fa4_varlen.keys()) @@ -514,9 +515,7 @@ def test_dpa_fa4_varlen(dtype, model_configs, model, qkv_layout): } -@pytest.mark.skipif( - not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." -) +@requires_fa4 @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_mask]) @pytest.mark.parametrize("model", model_configs_fa4_mask.keys()) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 2c0a5d9217..91a8b576a5 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -29,7 +29,7 @@ from transformer_engine.pytorch.attention.dot_product_attention.utils import FlashAttentionUtils _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ModelConfig, get_available_attention_backends pytest_logging_level = logging.getLevelName(logging.root.level) @@ -300,7 +300,10 @@ def _submit(pool: PoolWorker, **kwargs) -> None: qkv_formats = ["sbhd", "thd"] -@pytest.mark.skipif(not FlashAttentionUtils.v2_plus, reason="Flash-attn 2.0+ is required.") +@pytest.mark.skipif( + not (FlashAttentionUtils.v2_plus or FlashAttentionUtils.v3_is_installed), + reason="Flash-attn v2 or v3 is required.", +) @pytest.mark.skipif(get_device_compute_capability() < (8, 0), reason="CP tests require sm80+.") @pytest.mark.parametrize("dtype", dtypes) @pytest.mark.parametrize("model", model_configs_flash_attn.keys()) diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index cdd98d2445..2a857a10dc 100644 --- a/tests/pytorch/attention/test_kv_cache.py +++ b/tests/pytorch/attention/test_kv_cache.py @@ -33,7 +33,8 @@ ) _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ( ModelConfig, reset_rng_states, diff --git a/tests/pytorch/distributed/test_fusible_ops.py b/tests/pytorch/distributed/test_fusible_ops.py index 9d08a81a8f..d733286093 100644 --- a/tests/pytorch/distributed/test_fusible_ops.py +++ b/tests/pytorch/distributed/test_fusible_ops.py @@ -32,7 +32,8 @@ # Import utility functions _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import dtype_tols, make_recipe, quantization_tols diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py index eb43ba7e75..07dffebf5f 100644 --- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py +++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py @@ -37,7 +37,8 @@ # Import utility functions _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import dtype_tols, make_recipe, run_distributed, str_to_dtype # Check if FP8 is supported diff --git a/tests/pytorch/distributed/test_sanity.py b/tests/pytorch/distributed/test_sanity.py index 2e7a63e0a2..1b8bf5890f 100644 --- a/tests/pytorch/distributed/test_sanity.py +++ b/tests/pytorch/distributed/test_sanity.py @@ -19,7 +19,8 @@ from transformer_engine.common import recipe _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +# Prepend so installed packages with a top-level utils module cannot shadow the test helpers. +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ModelConfig model_configs = { diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index aeac501ff8..aa3bc15e84 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -167,14 +167,34 @@ flash_attn_func_v4 = None flash_attn_varlen_func_v4 = None else: - from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module - flash_attn_func as flash_attn_func_v4, - flash_attn_varlen_func as flash_attn_varlen_func_v4, - _validate_head_dims as _fa4_validate_head_dims, - ) + try: + cutlass_dsl_version = PkgVersion(get_pkg_version("nvidia-cutlass-dsl")) + + # FA4 4.0.0b24 requires CUTLASS DSL 4.6.2 or newer. + if fa_utils.fa4_version == PkgVersion("4.0.0b24") and cutlass_dsl_version < PkgVersion( + "4.6.2" + ): + raise ImportError( + "flash-attn-4 4.0.0b24 requires nvidia-cutlass-dsl>=4.6.2; " + f"found {cutlass_dsl_version}" + ) - fa_utils.v4_validate_head_dims = _fa4_validate_head_dims - fa_utils.set_flash_attention_4_params() + from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module + flash_attn_func as flash_attn_func_v4, + flash_attn_varlen_func as flash_attn_varlen_func_v4, + _validate_head_dims as _fa4_validate_head_dims, + ) + except ImportError as exc: + flash_attn_func_v4 = None + flash_attn_varlen_func_v4 = None + warnings.warn( + f"FlashAttention 4 is installed but cannot be loaded: {exc}", + RuntimeWarning, + stacklevel=2, + ) + else: + fa_utils.v4_validate_head_dims = _fa4_validate_head_dims + fa_utils.set_flash_attention_4_params() # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16 _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index b444f034c8..ea89ca97eb 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -3087,10 +3087,12 @@ def forward( window_size == (-1, 0) or window_size == (-1, -1) or use_fused_attention + or use_flash_attn_3 or fa_utils.v2_3_plus ), ( "cp_comm_type='all_gather' only supports SWA through FusedAttention or FlashAttention" - f" >= 2.3. Found {use_fused_attention=} and {fa_utils.v2_3_plus=}." + f" >= 2.3. Found {use_fused_attention=}, {use_flash_attn_3=}, " + f"and {fa_utils.v2_3_plus=}." ) assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( "cp_comm_type='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" @@ -4231,10 +4233,11 @@ def forward( window_size == (-1, 0) or window_size == (-1, -1) or use_fused_attention + or use_flash_attn_3 or fa_utils.v2_3_plus ), ( "cp_comm_type='a2a' only supports SWA through FusedAttention or FlashAttention >= 2.3." - f" Found {use_fused_attention=} and {fa_utils.v2_3_plus=}." + f" Found {use_fused_attention=}, {use_flash_attn_3=}, and {fa_utils.v2_3_plus=}." ) assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( "cp_comm_type='a2a' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6eb3ce54f1..e9fc03b174 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -563,10 +563,10 @@ def _disable_all_flash_attention() -> None: if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: logger.debug("Disabling FlashAttention 3 for compute capability != sm90") use_flash_attention_3 = False - # FA4 supports SM80, SM90, SM100, SM120 - if device_compute_capability < (8, 0): + # FA4 does not currently support SM8x. + if device_compute_capability < (9, 0): if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: - logger.debug("Disabling FlashAttention 4 for compute capability < sm80") + logger.debug("Disabling FlashAttention 4 for compute capability < sm90") use_flash_attention_4 = False # On SM90, prefer FA3 over FA4 when FA3 is available. # FA3 is more mature on Hopper; FA4's SM90 backward has limitations @@ -996,6 +996,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt device_compute_capability[0] * 10 + device_compute_capability[1], ) use_flash_attention_4 = False + # FA4's validator currently accepts symmetric (512, 512) on SM100/SM110, + # but the generic forward kernel exceeds its TMEM allocation for that shape. + # Preserve the supported asymmetric (64, 512) MLA path while D512 support + # is completed upstream. + if ( + use_flash_attention_4 + and (10, 0) <= device_compute_capability < (12, 0) + and head_dim_qk == head_dim_v == 512 + ): + logger.debug( + "Disabling FlashAttention 4 for unsupported symmetric head_dim=512 on SM100/SM110." + ) + use_flash_attention_4 = False # flash-attn-4 4.0.0b11 validates (256, 256) on SM100, but its dedicated # hd256 kernel diverges from the reference for cross-attention/decode-like # shapes such as sq=1, skv=2048. Keep FA4 enabled for the self-attention