[Add] Add FlashAttention 4 context parallel support.

[Baibaifan]

Summary

This PR adds FlashAttention 4 support for TransformerEngine context parallel attention.

The change enables FA4 in the CP backend selection path and wires the FA4 raw CuTe forward/backward APIs into the existing P2P context-parallel implementation. It also adds a lightweight distributed comparison script that validates FA4 CP outputs and gradients against non-CP full-sequence attention.

Changes

• Enable FlashAttention 4 when context parallelism is used.
• Import optional FA4 raw _flash_attn_fwd and _flash_attn_bwd APIs.
• Add FA4 raw forward/backward wrappers for CP P2P attention.
• Pass use_flash_attn_4 through the FlashAttention CP path.
• Add run_fa4_cp_vs_non_cp.py to compare FA4 CP with non-CP attention.

Validation

• python3 -m py_compile passed for modified Python files.
• git diff --check passed.
• Autograd forward/backward return arity was statically checked.
• New test script supports:

torchrun --nproc_per_node=2 tests/pytorch/attention/run_fa4_cp_vs_non_cp.py

[greptile-apps]

Greptile Summary

Adds FlashAttention 4 support to the P2P context-parallel attention path.

• Imports the optional FA4 raw CuTe forward and backward APIs.
• Adds FA4-specific CP kernel wrappers, window arguments, LSE handling, and autograd propagation.
• Enables FA4 during context-parallel backend selection for P2P and hybrid a2a+p2p modes.
• Adds a two-process numerical comparison script for zero-dropout causal and unmasked attention.

Confidence Score: 2/5

This PR should not merge until FA4 CP preserves configured attention dropout and accepts supported hybrid group layouts.

The new FA4 raw path omits dropout state for an admitted training configuration, causing incorrect attention behavior, and its mode guard rejects a reachable a2a+p2p subgroup layout before execution.

Files Needing Attention: transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py and transformer_engine/pytorch/attention/dot_product_attention/utils.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Adds the FA4 P2P implementation, but drops configured attention dropout and rejects a valid degenerate hybrid-group configuration.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Optionally imports FA4 private raw APIs and propagates the FA4 selection flag into context parallelism.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Enables FA4 for CP broadly, allowing the nonzero-dropout configuration mishandled by the new P2P wrapper.
tests/pytorch/attention/run_fa4_cp_vs_non_cp.py	Adds a useful distributed numerical comparison, but its fixed zero-dropout setup does not exercise the affected training configuration.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[DotProductAttention] --> B[Backend selection]
  B -->|FA4 + context parallel| C[attn_forward_func_with_cp]
  C -->|p2p or a2a+p2p| D[AttnFuncWithCPAndKVP2P]
  D --> E[FA4 raw forward]
  E --> F[Ring communication and output merge]
  F --> G[FA4 raw backward]
  G --> H[Q/K/V gradients]

Loading

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

[greptile-apps]

FA4 drops attention dropout

When FA4 context-parallel training uses nonzero attention_dropout, this branch omits dropout_p from the raw forward call, while the backward path also omits the dropout probability and corresponding RNG state, causing the configured dropout to be ignored or gradients to use an inconsistent dropout mask.

Knowledge Base Used: PyTorch Attention Stack

[greptile-apps]

Hybrid mode normalization rejects FA4

When cp_comm_type="a2a+p2p" uses a P2P subgroup of size one, the preceding normalization rewrites the mode to "a2a", so this guard rejects an advertised supported hybrid configuration and raises ValueError before attention executes.

Knowledge Base Used: PyTorch Attention Stack