Enable FA4 for context-parallel attention

[sudhakarsingh27]

Description

Enable FlashAttention 4 for context-parallel attention on the already-supported CP communication shapes: p2p, all_gather, and a2a.

FA4 can represent padded THD layouts with physical padded offsets plus seqused_* visible lengths. This lets the CP path handle fa_pad_between_seqs=True with the same layout distinction used by the non-CP reference path. The hierarchical a2a+p2p path remains disabled because it has not been validated for FA4.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Route FA4 raw forward/backward calls through the CP p2p, all_gather, and a2a FlashAttention paths.
• Pass padded THD cu_seqlens_* together with seqused_* for FA4 when padding exists between sequences.
• Keep FA4 disabled for a2a+p2p until that path is separately validated.
• Update FlashAttention CP tests so FA4-capable environments exercise the CP matrix and padded THD guards.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

Validation

• git diff --check
• python3.12 -m py_compile on the touched attention/test files
• Static forward/backward arity check for the three touched CP autograd functions
• Padded THD FlashAttention CP pytest subset on H100-class and B200-class GPUs: 6 passed, 10 skipped, 48 deselected on each platform. The skips are existing unsupported combinations: p2p+SWA, MLA with non-P2P CP, and padded THD with a2a+p2p.

[greptile-apps]

Greptile Summary

Enables FlashAttention 4 for supported context-parallel attention paths.

• Routes FA4 raw forward and backward calls through p2p, all-gather, and all-to-all communication modes.
• Adds padded THD sequence metadata handling with physical offsets and visible sequence lengths.
• Keeps hierarchical a2a+p2p excluded and expands CP test eligibility for FA4 environments.

Confidence Score: 5/5

The PR appears safe to merge with no concrete changed-code defect identified.

The FA4 flag is consistently propagated through backend selection and all three supported CP autograd paths, padded THD metadata is supplied separately from visible lengths, and the unsupported hierarchical communication mode remains excluded.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Threads FA4 through the three supported CP autograd paths and adapts raw forward/backward arguments, outputs, THD metadata, and stream serialization.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Imports FA4 raw kernels, forwards FA4 selection into CP, and distinguishes padded physical offsets from visible THD lengths.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Enables FA4 for p2p, all-gather, and a2a CP while retaining the exclusion for unsupported hierarchical a2a+p2p.
tests/pytorch/attention/test_attention_with_cp.py	Extends the existing CP matrix and padded-THD eligibility checks to FA4-capable environments.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[DotProductAttention backend selection] --> B{FA4 and context parallel?}
  B -->|p2p| C[FA4 CP p2p]
  B -->|all_gather| D[FA4 CP all-gather]
  B -->|a2a| E[FA4 CP all-to-all]
  B -->|a2a+p2p| F[Disable FA4 and use fallback]
  C --> G{THD padded layout?}
  D --> G
  E --> G
  G -->|Yes| H[Padded cu_seqlens plus seqused lengths]
  G -->|No| I[Standard layout metadata]
  H --> J[Raw FA4 forward and backward]
  I --> J

Loading

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}