Enable FA4 for context-parallel attention

[sudhakarsingh27]

Description

Enable FlashAttention 4 for context-parallel attention on the supported p2p, all_gather, and a2a paths; a2a+p2p remains disabled pending validation.

FA4 can represent padded THD layouts with physical padded offsets plus seqused_* visible lengths. This lets the CP path preserve the layout distinction used by the non-CP reference path.

This PR is limited to FA4 context-parallel enablement. FlashAttention package isolation, unsupported-configuration gating, and test-helper import fixes remain in #3341.

Type of change

• Documentation change
• Bug fix
• New feature
• Breaking change
• Infra/Build change
• Code refactoring

Changes

• Route FA4 forward/backward calls through the CP p2p, all_gather, and a2a paths.
• Pass padded THD cu_seqlens_* together with seqused_* for FA4.
• Keep FA4 disabled for a2a+p2p.
• Extend CP test eligibility so FA4-only environments reach the FA4-specific guards.

Validation

• Existing H100- and B200-class padded THD CP canaries: 6 passed, 10 skipped, 48 deselected per platform.
• Python compileall, Black, and git diff --check.
• Updated GPU CI pending after restoring the PR to the FA4-only change set.

Checklist

• I have read and followed the contributing guidelines.
• I have commented non-obvious behavior.
• My changes generate no new warnings.
• I have added tests or CI coverage for the behavior.

[greptile-apps]

Greptile Summary

Enables FlashAttention 4 for supported context-parallel attention paths.

• Routes FA4 forward and backward kernels through the p2p, all_gather, and a2a implementations while leaving a2a+p2p disabled.
• Preserves padded THD physical offsets using cu_seqlens_* and supplies visible token lengths through seqused_*.
• Extends backend selection and context-parallel tests to recognize supported FA4 configurations.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Adds FA4 kernel dispatch, argument adaptation, result handling, and autograd propagation across the supported context-parallel algorithms.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Imports FA4 private forward/backward entry points and propagates FA4 selection and padded THD metadata into context-parallel execution.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Enables FA4 backend selection for p2p, all-gather, and A2A context parallelism while retaining the combined A2A/P2P exclusion.
tests/pytorch/attention/test_attention_with_cp.py	Extends context-parallel test eligibility and padded THD coverage conditions to FlashAttention 4.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[DotProductAttention] --> B{Context parallel?}
  B -->|No| C[Standard attention backend]
  B -->|Yes| D{CP communication type}
  D -->|p2p| E[FA4 P2P forward/backward]
  D -->|all_gather| F[FA4 all-gather forward/backward]
  D -->|a2a| G[FA4 A2A forward/backward]
  D -->|a2a+p2p| H[Disable FA4]
  E --> I{Padded THD?}
  F --> I
  G --> I
  I -->|Yes| J[Padded cu_seqlens plus seqused lengths]
  I -->|No| K[Regular FA4 arguments]

Loading

_{Reviews (7): Last reviewed commit: "Merge current main into FA4 CP support" | Re-trigger Greptile}

[sudhakarsingh27]

Shouldn't 2.0+ also apply to flash-attn 4?

[sudhakarsingh27]

is it possible to extend get_fa_args instead of creating a new get_fa4_thd_kwargs?

[cyanguwa]

Please make sure the CI is running FA4 and for all the relevant tests.

[cyanguwa]

Same here, as the comment above about fwd.