Optimize context parallel attention layouts and gradient accumulation

[YangXu1990uiuc]

Description

Optimize context parallel attention with eligible dense cuDNN descriptors for equal-length BF16 THD inputs, bounded initialization removal, vectorized A2A layout copies, reuse of the first AllGather backward gradient contribution, and bounded native head partitioning to reduce workspace.

This is an opt-in experimental draft. All new switches are disabled by default. It keeps the existing communication flow and preserves the externally packed Stats layout through explicit stream-ordered conversions.

Type of change

• New feature (non-breaking change which adds functionality)

Changes

• Add uniform THD forward/backward descriptor paths for D128/D256.
• Check the promised sequence lengths and physical offsets on device on each execution; an invalid opt-in promise reports a CUDA device error.
• Preserve physical pitches and sequence origins in the bounded AllGather backward experiment and its graph cache.
• Skip redundant output/gradient zeroing only for the checked BF16 D128, separate-Q/K/V, GQA4 subset.
• Add a bounded, opt-in SM100 BF16 A2A pack/unpack fast path inside the existing swap API, used by both attention backends.
• Add focused PyTorch and C++ copy regressions and fix test-helper dtype dispatch with older CUDA headers.
• Reuse the first dK/dV contribution as the accumulator for opt-in BF16 THD CP4 or CP8/D256 AllGather, preserving the second addition, stream events and communication.
• Add sixteen repository regressions covering both backends, unequal padded THD at CP4/CP8 and CP2/SBHD fallbacks.
• Add default-off native H16/H8/H4 partitioning for qualified AllGather backward shapes; preserve physical strides and reuse workspace.
• Add a C++ address-contract regression and sixteen fresh-worker CP activation/fallback cases.
• Document the seven experimental controls and their limitations in docs/envvars.rst.

CP8 first-gradient accumulator reuse

Extend the existing default-off NVTE_CP_EXPERIMENTAL_AG_FIRST_DKV eligibility from CP4 to CP4 or CP8 for BF16 THD D256 AllGather. The first contribution already contains the accumulator's initial value; reusing it avoids a separate zero-filled allocation and the first add into zero. The accumulation algorithm, later addition, stream events, reorder and collective calls are unchanged. This is a three-file Python/test/documentation extension; no backend or FROST kernel changes.

The exact candidate passes 16/16 repository tests, zero skips, across 72 fresh rank workers on B200 with cuDNN 9.27 and FA4. Coverage includes OFF/ON, unequal padded THD CP4/CP8, GQA8 and CP2/SBHD fallback controls. Full outputs and gradients are checked. Loaded Python paths and native binaries are verified for every worker. All 318 common/Torch-native/build inputs match the existing build byte-for-byte; the corresponding binaries are reused. Changed-file hooks and the license check pass. Whole-repository PyTorch lint retains the same four diagnostic lines as the isolated unmodified parent, all in byte-identical untouched files; this is not a clean global lint pass. Full GPU L0/JAX coverage remains outstanding.

Independent-node complete-CP evidence

B64/S8192/Hq32/Hkv8/D256, BF16 causal uniform THD, 524288 tokens, CP8/AllGather. Each node uses eight 148-SM B200 GPUs at default 1000 W with NV18 links; cuDNN 9.27.0.42, CUDA 13.2-built TE, CuTeDSL 4.8 cu13 and FA4 4.0.0b35.dev1+gbb25ba3. Warm eager latency includes forward, backward, helpers and communication and reports the slowest rank. Each provider uses its own OFF/ON control; both FROST and FA4 receive the shared optimization.

Node	Provider	OFF	ON	Paired ON/OFF latency change	Faster orders
A	FROST	68.380 ms	64.871 ms	-5.05%	4/4
A	FA4	67.940 ms	65.289 ms	-3.94%	4/4
B	FROST	68.623 ms	65.333 ms	-4.75%	4/4
B	FA4	68.533 ms	65.879 ms	-3.83%	4/4

With reuse enabled on both providers, the paired FROST/FA4 latency changes are -0.64% (4/4 orders) / -0.79% (4/4 orders). Negative means lower latency. Percentages are medians of within-order paired ratios; nodes are analyzed separately.

Peak allocated memory decreases from 46.52 to 42.52 GiB per rank for FROST and from 24.52 to 22.77 GiB for FA4. The substantial remaining FROST memory gap is unchanged by the scope of this patch.

Each node passes 512 complete output/gradient checks and 128 independent CPU FP64 checks, with 128 fresh workers and 48 telemetry/stability windows. Four balanced provider/flag orders use three blocks of five warmups and 50 timed iterations. The repeat preserves source, software and protocol and uses eight disjoint GPU UUIDs.

Performance retains its frozen TE333 + FE323 composition attribution: it includes the native FROST bridge from #3699 and aligned gradient reduction, but excludes the separate pitched-dot change. The public candidate's gate/reuse body is AST-identical to that experiment; the current public package itself is qualified by the repository tests above, not remeasured for performance. Do not add these percentages to other experiments. This existing coordinate adds no original-grid win count.

Current-source repository qualification · First-node CP8 evidence · Independent-node repeat · Lint control.

Validation and performance

The descriptor/initialization path at initial main-based commit 36b0c8c builds with CUDA 13.2, cuDNN 9.27 and PyTorch 2.14.1+cu130. Its pinned binaries pass 48 directed forward/backward, tail, fallback and changed-input/poisoned-graph checks plus 18 existing repository tests on preproduction SM100. This is focused correctness coverage, separate from current-head measurements below.

At TE head 90d6b49, the changed common CUDA code and Python execution logic match the frozen runtime qualification (byte-identical CUDA; identical Python AST). On preproduction SM100, eight native positive cases and seven fallback cases pass bitwise, changed-input, poisoned graph, storage/stream-rebinding and route checks, with zero memcheck errors; all ten added repository PyTorch cases pass without skips. The focused C++ test binary passes 21 tests without skips, including the new CP pack/unpack case, with zero memcheck errors. A test-only if constexpr correction prevents an irrelevant unsupported dtype from being instantiated with older CUDA headers. Changed-file pre-commit, PyTorch/JAX C++ lint and changed Python-file lint pass. These are focused checks, not full GPU L0 CI. The four-B200 run of this revised integration is now fully admitted at the same B64/S8192/Hq32/Hkv8/D128 BF16 causal THD CP4/A2A shape, with cuDNN 9.29, dense descriptors already enabled and backend experiments disabled. Median CP latency is 43.62 -> 40.07 ms for cuDNN (8.08% paired reduction) and 43.75 -> 40.03 ms for FA4 (8.58%), both in 8/8 alternating rounds. Twelve small and twelve full-shape comparisons, all-rank native routes/library hashes, stability, hardware and telemetry checks pass; each round uses ten warmup and ten timed iterations, taking the slowest rank per iteration. The fair enabled/enabled cuDNN/FA4 ratio is 1.0010, with cuDNN leading only 3/8 rounds: the shared CP optimization adds no reliable backend winning cell. This is a power-managed 1000 W partial-node allocation, not full-model training. All 11 reported GitHub checks on TE head 90d6b49 now pass, including minimum supported framework versions, documentation, C++/Python lint, DCO and pre-commit; full GPU L0 coverage remains pending.

A prior cuDNN 9.29 CP4/A2A 8K experiment at TE 36b0c8c measured 48.61 ms with controls disabled, 43.91 ms with TE descriptor controls and 44.32 ms with TE plus backend controls. Paired reductions were 9.48% and 8.89%, both 8/8 rounds. This precedes the current copy optimization and shows a non-additive interaction; the gains should not be summed.

First-gradient buffer reuse

Commit 7a555de adds the disabled-by-default NVTE_CP_EXPERIMENTAL_AG_FIRST_DKV control. The exact Python candidate with the 90d6b49 native C++ build passes 96 OFF/ON rank pairs and 384 complete output/gradient comparisons with changed Q/dO and allocator churn before synchronization. Twelve additional four-B200 repository cases pass without skips: cuDNN/FA4, OFF/ON, unequal D256/GQA8 THD CP4, plus THD CP2 and SBHD CP4 fallbacks. Relevant Python QA lint, license and changed-file pre-commit pass. Native C++ is unchanged. This is focused coverage; the old head's CI status above does not establish CI on this new commit, and no new sanitizer coverage or general stream-lifetime guarantee is claimed.

The paired latency and peak-memory measurements linked below use the separately frozen private composed runtime. They establish an incremental shared TE optimization on an existing AllGather coordinate; they are not measurements of all current draft heads combined.

Current public integration: shared buffer-reuse increment

An isolated composition of TE #3699 (4152a201) and only the two first-dK/dV reuse blocks from TE #3683 (d379b933) passes full CP OFF/ON timing with FE #1690 (5ba587f1). The native libraries are unchanged; the FE vector-load composition is not included. This is a composition test, not either TE PR alone.

The shape is B64/S8192/Hq32/Hkv8/D256, BF16 causal equal-length THD, CP4/AllGather, 524288 tokens, on four production B200 GPUs (148 SMs, default 1000 W, NV18), cuDNN 9.27 and FA4 4.0.0b35.dev1+gbb25ba3. All six provider orders alternate OFF/ON order. Each provider/flag arm uses a fresh four-rank process set: 36 process sets, 144 workers, 576 complete output/gradient comparisons and 36 qualified telemetry windows. Each arm uses 5 warmups and 50 timed iterations with the default allocator and dynamic clocks, without CUDA Graphs or profiling. Latency covers attention forward, backward and communication and uses the slowest rank; percentages are medians of paired round ratios.

Provider	OFF → ON latency	Paired ON/OFF change	Faster rounds	Peak allocated OFF → ON
native	116.92 → 114.65 ms	-1.95%	6/6	61.05 → 57.05 GiB
frost	100.00 → 98.73 ms	-1.37%	6/6	53.05 → 49.05 GiB
fa4	104.40 → 101.49 ms	-2.77%	6/6	31.03 → 29.53 GiB

With both paths OFF, paired FROST/FA4 latency changes by -4.15% (6/6 faster rounds); with both paths ON, it changes by -2.70% (6/6). The shared optimization benefits must not be added to percentages from another run. This is one-node incremental evidence at an existing coordinate; different-node confirmation is pending and no original-grid win count is added. Qualified measurements and raw evidence.

Separate published-FE artifact qualification

The artifact exported from FE e2424cdb now passes a separate four-B200 CP4/A2A run at global B4/S131072/Hq32/Hkv8/D256, BF16 causal THD, cuDNN 9.27. Median training CP latency is 1105.86 ms stock, 970.69 ms TE dense, 910.16 ms FROST C++ and 1068.39 ms FA4. Median paired FROST latency is 14.87% lower than FA4 and 6.26% lower than TE dense, both in 8/8 interleaved process rounds; stock loses to FA4 in every round. Twelve small and twelve full-shape rank/variant comparisons, capture routes, binaries, NV18 topology, power/clocks and stability checks pass. Each round uses ten warm timed iterations and the slowest rank per iteration. Qualified CP evidence.

This uses the separate experimental TE/bridge binaries and the older A2A copy, with the default allocator. It is not a combined test of the three current draft heads, and the whole FROST algorithm's gain is not attributable to the FE workspace patch alone. This artifact has one-node qualification; the older private-hook artifact's two-node 12.19–15.07% result remains separately attributed. Requalifying the same coordinate adds no cell.

Native head partitioning

Commit 7fa9e1d extends process-static NVTE_CP_EXPERIMENTAL_NATIVE_AG_HEADS to 0/4/8/16 (default 0). Sixteen repository cases pass; H4 saves another 6 GiB versus H8 without a reliable extra latency benefit. Full GPU L0/JAX and composition checks remain pending.

Related optimization work

Current English report · 中文版. These same URLs contain the detailed protocols, raw measurements, rejected experiments and revision history. The implementation and validation sections above describe this draft's actual source. Results below span separate runtime families; combined qualification of all three current draft heads remains outstanding.

• Confirmed coverage. D256 BF16 causal THD CP4/A2A has nine measured coordinates (2K–512K), nine all-round advantages and eight strict flips, confirmed across independent B200 nodes. The 64K point is not a strict flip. The separate historical D128 selected set has seven strict flips in ten configurations, including two from backend-only changes. Exact original-record matching remains 0/324; these are selected datasets, not a full-grid win rate. Scope and evidence ledger.

• One additional AllGather strict flip, confirmed across two nodes. At B64/S8192/Hq32/Hkv8/D256, BF16 causal equal-length THD CP4/AllGather, native-forward plus the private FROST C++ backward lowers paired CP latency versus FA4 by 6.42% / 6.17%, all eight rounds on each node. Native THD/dense/FROST/FA4 medians are 115.74/105.06/99.13/106.16 ms and 118.09/107.12/100.35/106.99 ms. Dense versus FA4 is near parity. Peak PyTorch allocations are 64.05/80.05/56.05/34.03 GiB/GPU. Both runs use identical frozen inputs, cuDNN 9.29/NVRTC 13.5, FA4 bb25ba3, four 148-SM/1000-W B200s and NV18; ten warmup plus ten timed iterations use the slowest rank per iteration. Full outputs/gradients, runtime routes, hardware and stability pass. This includes communication and helpers; it is not pure SDPA timing, full-model memory or combined-three-head qualification. Two-node results.

• Shared TE copy optimization remains separately attributed. The current-copy 2K controls measure native/dense/FROST/FA4 at 46.10/42.71/39.79/43.02 ms. FROST versus FA4 is 7.61% lower in all eight rounds, with 8.05% on the other node's matching FROST/FA4 arms. This is an existing coordinate. Reorder is neutral and remains unpromoted. Current-copy attribution.

• First dK/dV buffer reuse now has complete CP timing. At the same B64/S8192/D256 AllGather coordinate, native THD/dense/FROST/FA4 OFF → ON medians are 117.48 → 115.24 / 107.11 → 103.99 / 100.33 → 98.71 / 107.68 → 104.80 ms. Median paired latency reductions are 2.06% / 2.98% / 1.70% / 2.74%, all eight rounds per path. Peak PyTorch allocations fall by 4 / 4 / 4 / 1.5 GiB per GPU. With both sides ON, FROST is 5.75% below FA4 in all eight rounds; dense is near parity at five of eight. This is a shared TE optimization on an existing coordinate, not an additional cell or an additive extension of the older FROST/FA4 percentages. It is one-node incremental evidence in the private composed runtime. Timing, memory, route, hardware, source and stability checks pass; an SGR-only topology parser repair preserves the original failed admission and every sample. Paired results and scope.

• First dK/dV reuse is in TE draft Optimize context parallel attention layouts and gradient accumulation #3683 at 7a555de, default off. Exact candidate Python with the TE90d6 native build passes 96 rank pairs / 384 complete comparisons, changed Q/dO and allocator churn on both streams. 12/12 repository cases, zero skips cover cuDNN/FA4, unequal D256/GQA8 THD CP4 and CP2/SBHD fallbacks; changed-file lint/license/pre-commit pass. This is focused coverage, not full GPU CI or a new alias sanitizer proof. Source and validation.

• Native head partition is now in the TE draft; 30% / 45% lower allocation reproduces on two B200 nodes. d54e149d adds default-off H16/H8 chunks in TE C++, preserving physical strides and reusing workspace; backend compute kernels are unchanged. 12/12 repository cases pass, including CP2/SBHD/D128 fallbacks. At B64×S8192, H32/Hkv8, D256, causal BF16 equal THD, CP4 AllGather, nodes B200-037/-017 use eight balanced rounds each: four 148-SM/1000-W/NV18 B200s, 64 fresh processes total. H32/H16/H8/FA4 median CP latency is 107.863 / 107.270 / 106.698 / 108.225 ms and 108.001 / 107.840 / 107.752 / 109.226 ms. Both nodes reproduce peak allocation 80.05 / 56.04 / 44.04 / 34.03 GiB; FA4 still uses less memory. H8 paired reductions are 1.16% / 0.05% versus H32 (7/8 / 4/8 rounds) and 1.30% / 1.40% versus FA4 (8/8 / 7/8 rounds). H8 is not an all-round cross-node win. Second-node H8/H32 is near zero: incremental latency benefit is unconfirmed. Percentages are medians of per-round ratios. The original 10-warmup/10-timed rank-max host-wall loop is unchanged. First-dKV reuse is off in all arms; dense/strided backward is on, forward/no-zero and FROST/backend experiments are off. 176/176 native full-CP controls are byte-identical; each allocation passes 80 full-tensor comparisons for FA4 (unchanged 0.025/0.025, relative-RMS<0.02). Allocation covers setup/warmup/timing, not a full model. Clocks are dynamic. One existing coordinate, no new winning cell or combined-three-head qualification. Every round and source evidence · Source and repository tests.

Private C++ FROST P2P: two-node timing confirmed. 8K/16K/32K/64K/128K confirmed. At 128K, CP passes 320 output/gradient comparisons; the 80 GiB reference is archived. Two B200 nodes measure FROST 886.91 / 890.81 ms versus FA4 1053.72 / 1053.99 ms: paired latency −15.82% / −15.51%, 8/8 faster rounds on each. Scope: BF16 causal THD, B4×131072, H32/Hkv8, D256, CP4/P2P, cuDNN 9.29; rank-max forward+backward+communication, fresh-process runs. Native already wins this point, so no strict flip is added. Component timing remains rejected (3.154 s / 2.081 s gaps). Separate public integration: TE FROST integration #3699, with separate full-CP qualification. Native dK trial 5059128 has no useful gain, not promoted. A separate 64-bit dot fix passes 912 helper cases and 128 full-SDK checks; the baseline reproduces 8 failures. Outside these draft heads; no performance claim. Original matching remains 0/324, full GPU L0/JAX remains separate. Evidence and limits.

Remaining before ready

• Agree on an explicit uniform-layout promise and a production eligibility/fallback contract instead of private environment switches.
• Complete PyTorch/JAX GPU L0 coverage, add further descriptor-path repository regressions, and broaden multi-GPU performance coverage beyond selected shapes.
• Review the initialization and graph-cache eligibility independently before enabling any default route.

The FROST C++ bridge used in separate experiments is outside this PR. This draft calls the existing cuDNN backend.

Checklist

• I have read the contributing guidelines.
• The functionality is complete for production use.
• I have commented the non-obvious contract and documented the experimental controls.
• A fresh build generates no new warnings.
• Focused copy tests have been added to the upstream repository.
• New and existing L0 tests pass on this commit.

Please keep this PR in draft until the author has reviewed it and the remaining checks are complete.

CP8 AllGather: composition with first-gradient reuse

The first-dK/dV buffer reuse in TE #3683 and the pitched O/dO dot loads in FE #1629 now have a complete-CP composition measurement on two independent production B200 nodes. OSS backward executes through the separate TE #3699 bridge; forward remains cuDNN. This addendum changes no code or default routing.

B64/S8192/Hq32/Hkv8/D256, BF16 causal equal-length THD, 524288 tokens, CP8/AllGather. Every provider, including FA4, enables the same first-dK/dV reuse. The two FROST arms retain identical aligned reductions and compact vector loads and differ only in the pitched-dot branch. The measured candidate's shared-helper file is byte-identical to FE #1629 at 232cba22; the frozen runtime is TE333 with FE362/FE365 and the previously qualified D128, physical-pitch and terminal-credit dependencies. This is qualification of that composition, not of the standalone FE branch or the entire current develop package.

Node	FROST + reuse	FROST + reuse + pitched dot	FA4 + reuse	Paired dot increment	Faster orders	Paired composed/FA4
B200-013	65.353 ms	65.135 ms	65.551 ms	-0.33%	6/6	-0.66%
B200-081	64.459 ms	64.175 ms	65.154 ms	-0.36%	5/6	-1.36%

The additional pitched-dot branch lowers complete-CP latency by 0.33% / 0.36% on the two nodes. The complete composition is 0.66% / 1.36% lower than FA4, with all six provider orders on each node favoring the composition in that comparison. The second node has one dot-versus-baseline order at +0.062%; the incremental benefit is not an all-order improvement there. Both nodes independently meet the predeclared 0.25% median and 5/6-direction confirmation gates. Percentages are medians of within-order ratios; node latencies are not pooled and earlier independent optimization percentages are not added.

Each node uses eight 148-SM B200 GPUs at default 1000 W with NV18 links, cuDNN 9.27.0.42, CUDA 13.2-built TE, CuTeDSL 4.8 cu13 and FA4 4.0.0b35.dev1+gbb25ba3. Six provider permutations use fresh eight-rank process sets, five warmups and 50 timed iterations per set. Warm eager latency includes forward, backward, helpers and communication and reports the slowest rank. Pointer audits occur outside timing; no observer or profiler is present in the measured interval.

Across both nodes, 1152 complete output/gradient checks, 192 independent CPU FP64 checks, 288 pointer-alias audits and all 36 hardware/stability windows pass. Node names and all eight GPU UUIDs differ between runs. Peak allocated memory remains 42.52 GiB for FROST versus 22.77 GiB for FA4. This is one verified coordinate, with no new original-324-grid win claim. The earlier reuse-off and reuse-only measurements retain their own sources and scope.

Two-node confirmation and frozen inputs · First node · Independent repeat.