From f5003a793dfa353fa19e67a223458b212a356641 Mon Sep 17 00:00:00 2001
From: Sudhakar Singh <sudhakars@nvidia.com>
Date: Mon, 10 Aug 2026 17:53:25 -0700
Subject: [PATCH 1/8] Gate FA4 and stabilize attention test imports

FA4 can be installed on SM8x even though its current implementation rejects those GPUs. Disable selection and skip dedicated FA4 tests there so A100 and L40 use supported attention backends.

FA4 and CUTLASS can also expose a generic utils package on sys.path. Prepend the Transformer Engine test helper directory in the context-parallel test so collection resolves the intended utilities.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---
 tests/pytorch/attention/test_attention.py     | 28 ++++++++-----------
 .../attention/test_attention_with_cp.py       |  2 +-
 .../attention/dot_product_attention/utils.py  |  6 ++--
 3 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py
index 82791084d8..2e8a713080 100644
--- a/tests/pytorch/attention/test_attention.py
+++ b/tests/pytorch/attention/test_attention.py
@@ -341,9 +341,13 @@ 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."
+requires_fa4 = pytest.mark.skipif(
+    not FlashAttentionUtils.v4_is_installed or device_compute_capability < (9, 0),
+    reason="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 +366,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 +444,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 +466,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 +487,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 +510,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..00760db237 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)
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 6eb3ce54f1..58a031b8b7 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

From 0f6c71ebcdd94849ce1321e59af93648f020c756 Mon Sep 17 00:00:00 2001
From: Sudhakar Singh <sudhakars@nvidia.com>
Date: Tue, 11 Aug 2026 16:03:03 -0700
Subject: [PATCH 2/8] Isolate FlashAttention CI backends

Moving images can install FA4 alongside older FlashAttention generations, which mixes a shared Python namespace and can make context-parallel reference runs compile an unsupported backend. Isolate the L3 version matrix, keep current CP comparisons on FA2/FA3, and temporarily reject symmetric D512 FA4 on Blackwell until upstream kernel support is complete.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---
 qa/L3_pytorch_FA_versions_test/test.sh        | 30 +++++++++++++++----
 .../attention/run_attention_with_cp.py        | 21 ++++++++++++-
 .../attention/test_attention_with_cp.py       |  5 +++-
 .../attention/dot_product_attention/utils.py  | 14 +++++++++
 4 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh
index 468ef04d76..ec49322886 100644
--- a/qa/L3_pytorch_FA_versions_test/test.sh
+++ b/qa/L3_pytorch_FA_versions_test/test.sh
@@ -33,12 +33,12 @@ sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); pri
 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
+#   - sm>90 (B200): FA2 (2.8.3) - FA4 CP remains disabled until its CP path is supported
 # Non-CP tests still run for every FA version in the array.
 if [ $sm_arch -gt 90 ]
 then
   FA_versions=(2.8.3 4.0.0b11)
-  CP_FA_VERSION="${FA_versions[-1]}"
+  CP_FA_VERSION="${FA_versions[0]}"
 elif [ $sm_arch -eq 90 ]
 then
   FA_versions=(2.8.3 3.0.0b1 4.0.0b11)
@@ -48,20 +48,40 @@ 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
diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py
index 7c6cdefd15..c2dce3792b 100644
--- a/tests/pytorch/attention/run_attention_with_cp.py
+++ b/tests/pytorch/attention/run_attention_with_cp.py
@@ -6,7 +6,7 @@
 import os
 import sys
 import logging
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
 import torch
 import torch.distributed as dist
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
@@ -46,6 +46,24 @@
 dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp8": torch.bfloat16}
 
 
+@contextmanager
+def _use_supported_flash_attention_for_cp():
+    """Keep the reference and CP runs on the same supported FA generation."""
+    # The CP implementation currently supports FA2/FA3, but the reference run
+    # happens before the CP group is attached and can otherwise select FA4.
+    # Restore the caller's setting because pool workers reuse this process.
+    env_var = "NVTE_FLASH_ATTN_V4"
+    previous = os.environ.get(env_var)
+    os.environ[env_var] = "0"
+    try:
+        yield
+    finally:
+        if previous is None:
+            os.environ.pop(env_var, None)
+        else:
+            os.environ[env_var] = previous
+
+
 def generate_input_shapes(
     qkv_format: str,
     config: ModelConfig,
@@ -192,6 +210,7 @@ def get_tols(config, dtype):
     return atol, rtol, rmse_tol
 
 
+@_use_supported_flash_attention_for_cp()
 def run_dpa_with_cp(
     dtype="bf16",
     model=None,
diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py
index 00760db237..91a8b576a5 100644
--- a/tests/pytorch/attention/test_attention_with_cp.py
+++ b/tests/pytorch/attention/test_attention_with_cp.py
@@ -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/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 58a031b8b7..1fa3daa388 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -996,6 +996,20 @@ 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

From 41ac022e422a1fe06ba85a90d6072175b6f6871b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 11 Aug 2026 23:04:44 +0000
Subject: [PATCH 3/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../pytorch/attention/dot_product_attention/utils.py           | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 1fa3daa388..e9fc03b174 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -1006,8 +1006,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             and head_dim_qk == head_dim_v == 512
         ):
             logger.debug(
-                "Disabling FlashAttention 4 for unsupported symmetric head_dim=512 "
-                "on SM100/SM110."
+                "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

From ae40007ef7e5d1528ae246b60ab661f03dc8c76b Mon Sep 17 00:00:00 2001
From: Sudhakar Singh <sudhakars@nvidia.com>
Date: Tue, 11 Aug 2026 17:27:41 -0700
Subject: [PATCH 4/8] Prepend PyTorch test utility imports

FA4 and its CUTLASS dependency expose a top-level utils module after Transformer Engine imports. Appending the test root can therefore bind these late imports to the installed module and fail collection. Give the repository helper precedence in the four test files that exhibited this ordering.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---
 tests/pytorch/attention/test_kv_cache.py                       | 3 ++-
 tests/pytorch/distributed/test_fusible_ops.py                  | 3 ++-
 tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py | 3 ++-
 tests/pytorch/distributed/test_sanity.py                       | 3 ++-
 4 files changed, 8 insertions(+), 4 deletions(-)

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 = {

From c88114db57d552811aefd59f072a2cfb5d114b69 Mon Sep 17 00:00:00 2001
From: Sudhakar Singh <sudhakars@nvidia.com>
Date: Tue, 11 Aug 2026 22:50:35 -0700
Subject: [PATCH 5/8] Recognize FA3 sliding-window CP support

The all-gather and a2a guards use the FA2 package version check to recognize FlashAttention support, so an isolated FA3 run is rejected even though FA3 implements sliding-window attention. Accept the explicit FA3 backend in both guards.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---
 .../attention/dot_product_attention/context_parallel.py    | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

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 ="

From eb9a249f33308f5e88bb448b152fa100e903a41b Mon Sep 17 00:00:00 2001
From: Sudhakar Singh <sudhakars@nvidia.com>
Date: Tue, 11 Aug 2026 23:58:36 -0700
Subject: [PATCH 6/8] Keep CP backend selection at CI boundaries

The CP runner must honor an explicit generation selected by its caller, particularly the existing B200 L3 FA4 lane. Remove its internal V4 override, restore the L3 SM100 selection changed in 0f6c71eb, and disable V4 only for the L1 suite that still targets FA2/FA3. This keeps per-generation L3 isolation intact without making the shared runner silently override directed coverage.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---
 qa/L1_pytorch_distributed_unittest/test.sh    | 10 +++++----
 qa/L3_pytorch_FA_versions_test/test.sh        |  4 ++--
 .../attention/run_attention_with_cp.py        | 21 +------------------
 3 files changed, 9 insertions(+), 26 deletions(-)

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 ec49322886..a6f3275fb7 100644
--- a/qa/L3_pytorch_FA_versions_test/test.sh
+++ b/qa/L3_pytorch_FA_versions_test/test.sh
@@ -33,12 +33,12 @@ sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); pri
 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): FA2 (2.8.3) - FA4 CP remains disabled until its CP path is supported
+#   - 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.
 if [ $sm_arch -gt 90 ]
 then
   FA_versions=(2.8.3 4.0.0b11)
-  CP_FA_VERSION="${FA_versions[0]}"
+  CP_FA_VERSION="${FA_versions[-1]}"
 elif [ $sm_arch -eq 90 ]
 then
   FA_versions=(2.8.3 3.0.0b1 4.0.0b11)
diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py
index c2dce3792b..7c6cdefd15 100644
--- a/tests/pytorch/attention/run_attention_with_cp.py
+++ b/tests/pytorch/attention/run_attention_with_cp.py
@@ -6,7 +6,7 @@
 import os
 import sys
 import logging
-from contextlib import contextmanager, nullcontext
+from contextlib import nullcontext
 import torch
 import torch.distributed as dist
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
@@ -46,24 +46,6 @@
 dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp8": torch.bfloat16}
 
 
-@contextmanager
-def _use_supported_flash_attention_for_cp():
-    """Keep the reference and CP runs on the same supported FA generation."""
-    # The CP implementation currently supports FA2/FA3, but the reference run
-    # happens before the CP group is attached and can otherwise select FA4.
-    # Restore the caller's setting because pool workers reuse this process.
-    env_var = "NVTE_FLASH_ATTN_V4"
-    previous = os.environ.get(env_var)
-    os.environ[env_var] = "0"
-    try:
-        yield
-    finally:
-        if previous is None:
-            os.environ.pop(env_var, None)
-        else:
-            os.environ[env_var] = previous
-
-
 def generate_input_shapes(
     qkv_format: str,
     config: ModelConfig,
@@ -210,7 +192,6 @@ def get_tols(config, dtype):
     return atol, rtol, rmse_tol
 
 
-@_use_supported_flash_attention_for_cp()
 def run_dpa_with_cp(
     dtype="bf16",
     model=None,

From a620611fde30d298ea05431fc1a2e99fca0828bf Mon Sep 17 00:00:00 2001
From: Sudhakar Singh <sudhakars@nvidia.com>
Date: Wed, 12 Aug 2026 13:39:19 -0700
Subject: [PATCH 7/8] Align FlashAttention CI coverage by architecture

Keep L0 on the mature FA2 path while L3 owns newer-generation coverage. Restrict H100 L3 to FA3 and B200 L3 to non-CP FA4 so unsupported H100 FA4 kernels and mislabeled Blackwell CP results do not obscure the intended signal. Make FA4-specific tests honor backend enablement to prevent silent fallback under an FA4 label.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---
 qa/L0_pytorch_debug_unittest/test.sh      |  5 +++++
 qa/L0_pytorch_unittest/test.sh            |  5 +++++
 qa/L3_pytorch_FA_versions_test/test.sh    | 22 +++++++++++++---------
 tests/pytorch/attention/test_attention.py |  9 +++++++--
 4 files changed, 30 insertions(+), 11 deletions(-)

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/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh
index a6f3275fb7..047358b301 100644
--- a/qa/L3_pytorch_FA_versions_test/test.sh
+++ b/qa/L3_pytorch_FA_versions_test/test.sh
@@ -31,18 +31,18 @@ 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[@]}"
@@ -102,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 ))
@@ -127,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 2e8a713080..c4f852f112 100644
--- a/tests/pytorch/attention/test_attention.py
+++ b/tests/pytorch/attention/test_attention.py
@@ -341,9 +341,14 @@ def test_dpa_num_splits(dtype, model_configs, model):
 }
 
 
+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 FlashAttentionUtils.v4_is_installed or device_compute_capability < (9, 0),
-    reason="Flash-attn v4 and compute capability >= SM90 are required.",
+    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.",
 )
 
 

From c0486820fce2e57474d847072e4c9c3c4d29e467 Mon Sep 17 00:00:00 2001
From: Sudhakar Singh <sudhakars@nvidia.com>
Date: Wed, 12 Aug 2026 15:11:47 -0700
Subject: [PATCH 8/8] Guard FA4 against incompatible CUTLASS installs

Package metadata can report FA4 present even when a later dependency install leaves its transitive CUTLASS stack unusable. Reject the known b24/CUTLASS combination below the stable 4.6.2 release and treat a nested interface ImportError as an unavailable optional backend so unrelated Transformer Engine imports can continue.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
---
 .../dot_product_attention/backends.py         | 34 +++++++++++++++----
 1 file changed, 27 insertions(+), 7 deletions(-)

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"