Stabilize FlashAttention CI environments

[sudhakarsingh27]

Summary

Stabilize attention CI when moving PyTorch images install multiple FlashAttention generations and CUTLASS packages.

Changes

• Keep PyTorch L0 on FA2 while retaining fused and unfused attention coverage.
• Run H100 L3 with FA3 and context parallelism; run B200 L3 with non-CP FA4.
• Isolate each L3 FlashAttention distribution and backend flag.
• Keep L1 context-parallel comparisons on FA2/FA3 until separate FA4 CP integration lands.
• Give repository-local test utilities precedence over installed top-level utils packages.
• Gate FA4 selection and dedicated tests on backend enablement and supported architectures/shapes.
• Reject FA4 b24 with an older CUTLASS DSL and treat a broken optional FA4 interface import as unavailable instead of failing Transformer Engine import.

Why

The FlashAttention distributions share a Python namespace, and the moving image can make an unintended generation available to otherwise generic tests. Explicit suite ownership keeps L0 stable, makes L3 labels truthful, and prevents unsupported FA4 paths from obscuring other CI results.

An installed FA4 distribution is not necessarily usable when a later package installation changes its transitive CUTLASS stack. Checking the known b24 minimum and guarding the public FA4 interface import prevents that optional dependency failure from breaking unrelated test collection.

Validation

• Bash syntax checks for the changed QA scripts.
• Python compile check and clean diff check.
• Focused FA4 import-guard checks cover absent FA4, an older CUTLASS DSL, a valid stack, and a nested dependency import failure.
• Directed GPU validation on the preceding revision passed H100 FA2, FA3, and FA3 CP, and B200 FA4; follow-up validation is pending for this revision.
• Full pre-commit bootstrap is blocked by the configured package mirror lacking setuptools>=80; relevant syntax and behavior checks pass independently.

[greptile-apps]

Greptile Summary

The PR stabilizes FlashAttention CI by assigning backend generations to specific suites and architectures, isolating conflicting distributions, and making optional FA4 loading resilient to incompatible CUTLASS installations.

• Pins L0 attention coverage to FA2 and keeps context-parallel L1 coverage on FA2/FA3.
• Assigns H100 L3 to FA3 with context parallelism and newer architectures to FA4 without context parallelism.
• Adds FA4 architecture, shape, enablement, and dependency compatibility guards.
• Ensures repository-local test utilities take precedence over installed top-level packages.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
qa/L3_pytorch_FA_versions_test/test.sh	Isolates FlashAttention distributions and assigns one backend generation and context-parallel policy to each supported architecture.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Guards FA4 loading against an incompatible CUTLASS DSL version and nested optional-interface import failures.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Refines FA4 architecture and unsupported head-dimension filtering before backend dispatch.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Recognizes FA3 as supporting sliding-window attention in all-gather and all-to-all context-parallel paths.
tests/pytorch/attention/test_attention.py	Gates dedicated FA4 tests on backend enablement, installation status, architecture, and supported shapes.
qa/L0_pytorch_unittest/test.sh	Pins the general L0 PyTorch suite to FA2 while preserving fused and unfused attention coverage.
qa/L1_pytorch_distributed_unittest/test.sh	Explicitly disables FA4 for context-parallel comparisons in both parallel and sequential CI execution.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  CI[Attention CI suite] --> L0[L0 PyTorch]
  CI --> L1[L1 distributed]
  CI --> L3[L3 FA versions]
  L0 --> FA2[FlashAttention 2]
  L1 --> CP[Context parallel tests]
  CP --> FA23[FlashAttention 2 or 3]
  L3 --> Arch{GPU architecture}
  Arch -->|SM90| FA3[FlashAttention 3 + CP]
  Arch -->|Above SM90| FA4[FlashAttention 4, non-CP]
  FA4 --> Guard{FA4 interface and CUTLASS compatible?}
  Guard -->|Yes| Enabled[Enable FA4 backend]
  Guard -->|No| Fallback[Mark FA4 unavailable]

Loading

_{Reviews (9): Last reviewed commit: "Guard FA4 against incompatible CUTLASS i..." | Re-trigger Greptile}

[sudhakarsingh27]

/te-ci pytorch L3