[PyTorch] Add head-parallel FA4 backward for cuDNN CP attention

[bzantium]

Description

Add an opt-in path that keeps cuDNN's P2P context-parallel forward and uses head-parallel FlashAttention-4 for backward. The saved attention tensors are packed into one all-to-all before backward, and the three gradients return in a second exchange. Packing preserves the FP32 log-sum-exp values and reuses the send allocation for the unpacked tensors.

The path targets causal FP16/BF16 attention on SM100 with CP2/4/8. Enable it with NVTE_FUSED_ATTN_CP_USE_FAv4_BWD=1; unsupported configurations retain the existing backward implementation.

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

• Add head-parallel FA4 backward and fused packing for balanced CP partitions.
• Document the opt-in switch and supported configurations.
• Add distributed correctness, graph replay and fallback tests to the FA4 CI runner, plus a reproducible attention benchmark.

Validation

B200, 32K sequence, BF16, 32 heads, QK/V dimensions 192/128. Native–mixed–mixed–native order, 50 synchronized graph replays per run; maximum across ranks per sample, then median of run medians. These are attention forward/backward times and incremental allocated peaks during graph capture/replay, excluding projections, MoE, optimizer and checkpoint I/O.

CP	Batch	Native (ms)	Mixed (ms)	Time change	Native / mixed peak (MiB)
2	1	20.527	22.085	+7.6%	2821.0 / 2818.0
4	1	12.915	12.164	-5.8%	1411.5 / 1409.0
8	1	12.213	9.720	-20.4%	932.5 / 932.5
4	2	24.554	23.736	-3.3%	2821.0 / 2818.0

The benchmark uses CUDA 13.1, PyTorch 2.11.0a0 (NGC 26.02), cuDNN 9.20, FA4 commit 71bf77c and CUTLASS DSL 4.5.1. CP2 is slower in this configuration, so this remains opt-in.

• Attention checks pass across CP2/4/8 and CP2/4 with attention heads partitioned across two CP groups. FP16/BF16, batch 1/2 and both supported head dimensions are covered.
• Forward outputs match native cuDNN exactly. Gradients pass elementwise tolerances and relative L2 < 0.008 against native CP and full-sequence FP32 math attention. Bit-preserving packing is checked with both 32-bit and 64-bit indices on small tensors. Changed-input graph replay, deterministic mode and native fallbacks also pass. The real large-stride regression is opt-in with NVTE_TEST_CP_LARGE_STRIDE=1 because it needs 6 GiB per rank.
• The same five configurations pass with CI-pinned FA4 4.0.0b11 / CUTLASS DSL 4.4.2. Existing CP utilities (23 tests and 4 subtests) and five existing CP cases covering GQA, MLA, bias and multiple communication modes also pass.
• pre-commit run --all-files, production-file pylint and the license check pass.

Reproduce a benchmark row from the repository root with the dependencies above:

NVTE_FLASH_ATTN=0 NVTE_FUSED_ATTN=1 NVTE_FUSED_ATTN_BACKEND=1 \
CUDA_DEVICE_MAX_CONNECTIONS=8 NVTE_BATCH_MHA_P2P_COMM=0 \
python -m torch.distributed.run --standalone --nproc-per-node=4 \
  benchmarks/benchmark_mixed_cp.py --sequence 32768 --batch 1 --output cp4.json

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

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the previous memory-pressure concern is fixed and no new actionable failures remain.

Summary

• Packs saved attention tensors into one all-to-all exchange and returns gradients through a second exchange.
• Restricts dispatch to supported SM100, causal FP16/BF16 configurations and preserves the native fallback elsewhere.
• Adds distributed correctness, CUDA graph replay, fallback, bit-preservation, wide-index, and benchmark coverage.
• The changes since the previous review avoid the default multi-GiB sparse allocation while retaining explicit large-offset regression coverage.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[cuDNN P2P CP forward] --> B[Saved Q, K, V, O, dO, LSE]
    B --> C[Pack sequence partitions]
    C --> D[All-to-all: sequence to heads]
    D --> E[Head-parallel FA4 backward]
    E --> F[Pack dQ, dK, dV]
    F --> G[All-to-all: heads to sequence]
    G --> H[Original CP gradient layout]
    U[Unsupported configuration] --> N[Existing native backward]

Loading

_{Reviews (3) · Last reviewed commit: "Keep wide-index CP tests small by defaul..."}

[bzantium]

I checked the failed Build jobs on 37813c6:

• PyTorch and All report that the hosted runner lost communication with the server; the PyTorch job log is unavailable.
• Core hit the six-hour execution limit.
• JAX failed compiling the unchanged fused_attn/flash_attn.cu through sccache/CUDA 13.3: fatbinary fatal: Could not open input file 'flash_attn.compute_75.cubin'.

These failures occurred before the attention tests. Lint, license, documentation, and minimum-version checks passed. The new workflows on b517bbf are awaiting maintainer approval (action_required), including Build. Could a maintainer approve them? I do not have permission to approve or rerun upstream workflows. If the JAX failure repeats, a run without sccache would help isolate it.