Add cuDNN OSS backward dispatch for context parallel attention

[YangXu1990uiuc]

Description

Add an opt-in cuDNN OSS (FROST) backward provider to TE's fused-attention path so context-parallel attention can execute the open-source D128 and D256 kernels through the common C++ dispatcher.

The first call prepares a cuDNN Python graph and registers its compiled TVM-FFI backward chain. Subsequent calls bind current pointers, stream and caller-owned workspace directly in C++. Forward and context-parallel communication retain their existing implementations.

Type of change

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

Implementation and supported scope

• Reuse TE's optional CuTeDSL/TVM-FFI registration infrastructure. Cache immutable plans per declaration and device, including mask alignment; validate binding alignment, workspace capacity and overlap before dispatch.
• Support SM100 BF16, separate Q/K/V, D128/D256, nondeterministic backward, no dropout/bias/sinks, and unmasked or top-left causal attention. D256 also supports rectangular bottom-right causal attention when Q length is no greater than KV length; Q-longer-than-KV uniform THD falls back in mode 1 and declines in mode 2. Existing forward rejects that orientation for dense BSHD. D128 bottom-right remains unsupported. Dense BSHD lengths are multiples of 128.
• D128 uses native packed THD, including unequal sequence lengths. D256 THD requires an explicit uniform-layout promise, checked on the GPU on every execution, and converts packed FP32 Stats on the caller's stream.
• A separate default-off flag admits causal THD AllGather for CP2/4/8. It preserves physical batch pitches, selects the current Q half by pointer origin, and retains existing gradient initialization for unused rows. No new Q/K/V repack kernel is added.
• Add a default-off D128 output-initialization policy for strict FROST, BF16 GQA and square compact equal-length THD. The PyTorch binding and common dispatcher share the eligibility predicate; an every-execution GPU check validates lengths/origins before the complete backward chain overwrites dQ/dK/dV. Native/fallback, unequal and rectangular cases retain their fills.
• Retain the complete backward helper chain, scratch initialization and grouped-gradient reductions. Add off, fallback and strict-routing modes, plus tests for changed inputs/addresses, nondefault streams, mask-cache separation and CUDA Graph replay. Replay removes the Python preparation callback after warmup to verify native dispatch.

This PR adds the execution path; it does not change FROST kernel algorithms. TE #3683 separately optimizes layouts and CP copies.

Dependencies and use

This draft uses the prepared-plan contract in cuDNN Frontend #1688, now merged. The original measurements below retain their frozen FE 333278627eb2ccffc5f47da26d733f87a2d0a8aa attribution. D256 THD AllGather additionally requires the physical-stride support in cuDNN Frontend #1690. D128 validation uses the frozen #1688 source with the same D256 fixes; D256 AllGather measurements below use the published #1690 source. Installing #1690 alone does not supply Roman's D128 engine.

A compatibility rerun of the same 62 directed tests passes with zero skips on production B200 using the merged Roman source c98264fc9222d85dcac747fa9295ed9d8d4b2c47 plus the complete published #1690 patch, unchanged TE native binaries, and the CUDA 13 CuTeDSL flavor. It includes all five isolated metadata guards. This validates the newer dependency, without reattributing earlier performance measurements. A separate default-CUDA-12 probe is rejected by the existing TE compiler gate before kernel compilation; CUDA 12 support is not claimed or enabled. Merged-dependency qualification.

The adapter is confined to one module and rejects unfamiliar prepared-plan metadata. The metadata is currently private; a stable FE export contract remains a prerequisite for graduating this integration.

Install compatible FE source in TE's Python environment with nvidia-cutlass-dsl[cu13]>=4.8.0 and apache-tvm-ffi. Build TE with NVTE_WITH_CUTEDSL=1 (-DNVTE_WITH_CUTEDSL=ON for direct CMake). A newer system toolkit alone does not replace the required CUDA 13 CuTeDSL compiler flavor.

export NVTE_ENABLE_CUTEDSL_BACKEND=1
export CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1
export NVTE_FUSED_ATTN_FROST=2  # strict OSS backward; 1 permits backend fallback
# Optional: omit fully overwritten output fills for compact D128 GQA THD:
export NVTE_FUSED_ATTN_FROST_D128_NOZERO=1
# For equal-length D256 THD:
export NVTE_FUSED_ATTN_FROST_UNIFORM=1
# Additionally enable the qualified causal THD AllGather layout:
export NVTE_FUSED_ATTN_FROST_ALLGATHER=1
python -m pytest -v tests/pytorch/attention/test_frost_attn.py

Warm up each declaration before CUDA Graph capture. Leave the uniform-layout promise disabled for unequal THD and gapped layouts outside the explicit AllGather contract. NVTE_FUSED_ATTN_FROST=0 retains the existing route.

Review follow-up at a7e91a56

Add qa/L1_pytorch_frost_unittest/test.sh for a provisioned SM100 target, with an opt-in hook from PyTorch L0 via NVTE_FROST_TEST_REQUIRED=1. It requires the compatible FE/CUDA 13 CuTeDSL stack and a bridge-enabled TE build, checks actual capability, and rejects skipped tests or missing isolated metadata-guard children. The CI owner still needs to provision and enable the qualified lane; ordinary L0 lanes do not inherit the experimental dependencies.

The actual checked-in QA launcher passes 62/62 tests with zero skips on B200, including all five isolated guard children. The added D256 bottom-right Q=256, KV=128 test verifies strict decline, mode-1 fallback against CPU FP64 gradients, and exactly zero dQ on leading keyless rows. It also returns to fallback after strict decline to exercise the cached unsupported declaration. Native source and binaries are unchanged from cf04d4e1; no new CP performance claim is attached to this QA/documentation commit.

TE_PATH="$PWD" XML_LOG_DIR=/tmp/te-frost-qa \
  bash qa/L1_pytorch_frost_unittest/test.sh

Changed-file pre-commit, both shell syntax checks, license checks and the exact staged-source comparison pass. QA qualification. The comment-only review found no new correctness blocker in the supported scope at cf04d4e1; this does not replace maintainer approval or the remaining FE export and broader qualification work.

Validation

Native revision cf04d4e1 rebuilds the changed common and PyTorch native units with CUDA 13.2, cuDNN 9.27 and PyTorch 2.14.1+cu130. Changed-file formatting, C++ lint, license checks and exact-runtime import checks pass. The existing unused custom_types_converters compiler warning is unchanged. All eleven reported GitHub checks passed at cf04d4e1, including both minimum-version jobs. All eleven reported checks also pass at the QA/documentation follow-up a7e91a56. Full GPU L0/JAX qualification remains separate.

61/61 directed repository tests, zero skips, pass through the rebuilt TE C++ entry on a production B200. They retain the 46 existing FP64-gradient, D128 unequal-THD, D256 uniform/bottom-right/pitched-layout, rebinding, stream, fallback and graph-replay cases. Ten added initialization tests check changed inputs, exact OFF/ON agreement, NaN-poisoned graph replay without Python preparation, actual output-fill counts, and native/fallback/unequal/rectangular controls. Five further tests reject changed device lengths/origins in isolated processes, including mutation after graph capture. Exact-source qualification.

The following distributed correctness results belong to the preceding 4152a201 build. The newly rebuilt head's full CP16K/A2A comparison is reported separately below:

• Actual CP2: B2/S512/Hq8/Hkv2/D256 BF16 causal, all six BSHD/THD × A2A/P2P/AllGather pairs. Native controls and strict FROST pass 96 complete O/dQ/dK/dV comparisons against CPU FP64, maximum relative RMS 0.003870.
• Actual CP4 and CP8 THD AllGather: B2/S1024 and B2/S2048 respectively, Hq8/Hkv2/D256. Native and strict FROST pass 96 additional complete CPU-FP64 comparisons across 12 ranks, maximum relative RMS 0.005176.
• The companion FE Support computing zero-centered gamma in compute dtype for CuDNN #1690 draft passes its 99-test matrix, including physical-stride correctness and padding canaries, on a develop base.

Full GPU L0 suites, JAX execution, broader layout coverage and clean sanitizer qualification remain pending. The earlier integration's sanitizer run retained a Selective device code recompilation in progress API warning in an existing native-forward launch; this is not a clean sanitizer pass, and the later bottom-right/batch-pitch changes are not sanitizer-qualified.

Complete CP performance at the preceding 4152a201 integration

TE 4152a201 + FE #1690 5ba587f1, D256 BF16 causal equal-length THD, Hq32/Hkv8, CP4/AllGather, 524288 total tokens. Each node uses four production B200 GPUs: 148 SMs, default 1000 W limit, NV18 links; cuDNN 9.27.0.42, CUDA 13.2-built TE, CuTeDSL 4.8 CUDA 13 flavor, FA4 4.0.0b35.dev1+gbb25ba3.

Every shape/node run uses all six provider orders with fresh four-rank process sets: 72 workers, 288 complete output/gradient comparisons and 18 qualified telemetry windows. Five warmups and 50 timed iterations use the default allocator and dynamic clocks, without CUDA Graphs or profiling. Latency includes attention forward + backward + communication and reports the slowest rank. Percentages are medians of paired round ratios.

Shape / node	Native cuDNN	FROST backward	FA4	FROST/native change	FROST/FA4 change
64×8K / B200-250	119.82 ms	102.67 ms	106.15 ms	-14.31%	-3.35%
64×8K / B200-048	119.19 ms	101.58 ms	105.42 ms	-14.46%	-3.59%
32×16K / B200-018	188.00 ms	157.06 ms	170.50 ms	-16.54%	-7.92%
16×32K / B200-069	333.57 ms	272.50 ms	313.43 ms	-18.24%	-13.06%
16×32K / B200-079	335.76 ms	269.79 ms	307.66 ms	-19.63%	-12.36%
8×64K / B200-018	606.35 ms	482.19 ms	558.33 ms	-20.45%	-13.66%
8×64K / B200-073	606.86 ms	487.23 ms	563.12 ms	-19.72%	-13.46%

FROST is faster in 6/6 rounds for both comparisons in every row. These are four tested coordinates: 8K, 32K and 64K have independent-node confirmation with disjoint GPU UUIDs; 16K has one-node evidence. The two 64K nodes pass 576 complete output/gradient comparisons and 36 telemetry windows across 144 workers. They do not establish an original benchmark-grid win rate or extrapolate to D128 or other communication methods.

Peak allocated memory is 61.05 GiB native, 53.05 GiB FROST, 31.03 GiB FA4. Only the independent native reference is batch-chunked to avoid a separately diagnosed helper-offset limitation; actual CP calls retain the complete shape. Separate vector-load and first-dK/dV reuse changes are excluded. The D128 initialization policy does not apply to these D256 shapes.

8K two-node qualification · Initial 16K/32K qualification · 32K two-node qualification · 64K two-node qualification

Separate full-shape trace validation at 4152a201 passes 48 complete comparisons each for D256 AllGather and D128 A2A. All providers execute matching collective signatures/message sizes. D128 backward main-kernel sums are 20.06 ms native, 18.15 ms FROST and 18.06 ms FA4; helpers, forward and copies remain in the full path. These instrumented, separate-process durations are diagnostic, not paired CP latency or a D128 win. D256 attribution · D128 attribution.

D128 initialization policy: two-node evidence before packaging

Before this commit, the same native policy was measured as a private patch over 4152a201, with unchanged Roman 33327862 FROST kernels. The published implementation preserves those native statements; the new build is separately covered by the 61 tests above. These timings belong to the frozen pre-packaging binaries.

Workload: B64/S8192/Hq32/Hkv8/D128, BF16 causal compact THD, 524288 tokens, CP4/A2A. Two independent four-B200 nodes use the same six-order/fresh-process protocol, software, hardware and complete-CP timing definition above. All 576 full O/dQ/dK/dV comparisons and 36 telemetry windows pass across 144 workers with disjoint GPU UUIDs. Untimed audits confirm three versus zero output-fill calls in every FROST arm/rank.

Node	FROST: initialize	FROST: omit fills	FA4	Paired omit/initialize	Paired omit/FA4
B200-049	45.919 ms	45.608 ms	44.430 ms	-0.75%	+2.54%
B200-249	45.453 ms	45.064 ms	43.981 ms	-0.98%	+2.45%

Omitting eligible fills is faster in 6/6 rounds on both nodes, but remains slower than FA4 in every round. Allocated output sizes and internal FROST scratch initialization are unchanged. This is an incremental benefit for one layout, not an additional original-grid win; percentages from other experiments are not additive. Two-node qualification and raw evidence.

D128 16K complete CP at the published head

Rebuilt cf04d4e1 binaries with unchanged Roman 33327862 kernels complete B32/S16384/Hq32/Hkv8/D128, BF16 causal compact THD, CP4/A2A, 524288 tokens on one production B200 node. The same six-order protocol passes 288 full O/dQ/dK/dV comparisons, 18 telemetry windows and 72 fresh workers; maximum relative RMS is 0.003663. Untimed audits confirm three versus zero output fills in every FROST arm/rank.

FROST: initialize	FROST: omit fills	FA4	Paired omit/initialize	Paired omit/FA4
73.021 ms	72.806 ms	69.168 ms	-0.47%	+5.14%

The initialization option is faster in 5/6 paired rounds and remains slower than FA4 in all six. This is one shape/node result from the published build, not a new original-grid win or evidence for other communication methods. Only the native reference uses B8 chunks; actual CP uses the complete shape. Published-head 16K qualification.

The rolling explanation and evidence retains earlier-head measurements and separate composition experiments with their source identities and limitations.

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented the nontrivial implementation details
• I have made corresponding documentation changes
• My changes generate no new warnings
• I have added focused tests
• New and existing unit tests pass locally with my changes

Draft pending human author review, stable FE export ABI and the remaining integration qualification.

Follow-up D128 kernel synchronization qualification

cuDNN Frontend #1698 repairs terminal dQ barrier-phase reuse in the dense and THD FROST D128 backward kernels, including the grouped-query routes. Its exact-source regression reproduces pre-fix Missing wait reports and passes 70 ordinary tests plus all eight cases under each of racecheck and synccheck after the repair. Unfiltered memcheck qualification remains pending because of a retained CUDA runtime-recompilation API warning.

The complete-CP timings above remain attributed to their stated FE/kernel snapshots, before this repair. Repaired 8K and 2K CP4/A2A remeasurements are reported separately below; the full current-FE package remains outside that qualification. Repair validation and limitations.

Repaired D128 8K complete CP remeasurement

A new complete CP4/A2A run uses the qualified packed terminal-credit repair, AST-equivalent to FE #1698, on the frozen FE197 integration with #1696 and #1690. Only that packed kernel changes; TE native binaries, all other FE inputs and communication remain identical across the four providers. This is a new comparison with the repaired dependency, not a measurement of repair overhead against an unrepaired arm.

Workload: B64/S8192/Hq32/Hkv8/D128, BF16 causal compact THD, 524288 tokens, CP4/A2A. Four production B200 GPUs use 148 SMs, default 1000 W limits and NV18 links; cuDNN 9.27.0.42, CUDA 13.2-built TE, CuTeDSL 4.8 cu13, FA4 4.0.0b35.dev1+gbb25ba3. Warm eager timing includes forward, backward, helpers and communication and reports the slowest rank.

Provider	Median complete-CP latency
cuDNN backend	48.132 ms
FROST backward; cuDNN forward	45.110 ms
FROST forward and backward	44.823 ms
FlashAttention 4	43.759 ms

Paired latency changes: backward-only/FA4 +3.03%; combined/FA4 +2.30%; combined/backward-only -0.81%; combined/cuDNN -6.88%. Negative means lower latency. Percentages are medians of four paired round ratios. The combined path is faster than FA4 in 0/4 orders.

All 256 complete output/gradient comparisons, 64 fresh workers and 48 timing windows pass, with maximum relative RMS 0.003660. Actual packed-GQA backward route, forward engine, source/binary hashes, output-fill audits and telemetry/stability gates are checked. A separately versioned collector corrects only an expected worker-status spelling; the frozen workloads, measurements and all admission gates are unchanged.

This is one coordinate on one node, not an original-grid win rate or full qualification of the current FE package. The earlier results retain their own frozen-source attribution; the repaired 2K remeasurement follows below. Public current-upstream memcheck and the remaining integration graduation requirements above remain pending. Repaired CP qualification · Collector correction.

Repaired D128 2K complete CP remeasurement

The same repaired FE dependency and unchanged TE binaries now complete the four-provider comparison at B256/S2048/Hq32/Hkv8/D128, BF16 causal compact THD, 524288 tokens, CP4/A2A. The dependency is frozen FE197 + #1696/#1690 with the qualified packed terminal-credit repair, AST-equivalent to FE #1698. All providers use the same TE communication implementation.

Four production B200 GPUs use 148 SMs, default 1000 W and NV18 links; cuDNN 9.27.0.42, CUDA 13.2-built TE, CuTeDSL 4.8 cu13, FA4 4.0.0b35.dev1+gbb25ba3. Warm eager latency includes forward, backward, helpers and communication and reports the slowest rank. This coordinate runs on a different node from the repaired 8K result above.

Provider	Median complete-CP latency
cuDNN backend	26.017 ms
FROST backward; cuDNN forward	24.493 ms
FROST forward and backward	24.287 ms
FlashAttention 4	25.559 ms

Paired latency changes: backward-only/FA4 -4.16%; combined/FA4 -5.01%; combined/backward-only -0.83%; combined/cuDNN -6.69%. Negative means lower latency. These are medians of four paired round ratios; all four orders favor the combined path for every comparison.

All 256 complete output/gradient comparisons, 64 fresh workers and 48 timing windows pass (maximum relative RMS 0.003662). Actual packed-GQA backward and forward routes, source/binary hashes, output-fill audits, hardware telemetry and timing stability are checked.

The repaired 2K and 8K results establish gains at these two bounded coordinates; 8K remains 2.30% slower than FA4. They do not establish an original-324-grid win rate, measure repair overhead against an unrepaired arm, or qualify the entire current FE package. Earlier D256 measurements retain their stated source versions. Stable FE export/packaging, public current-upstream memcheck and the remaining integration qualification remain pending. Repaired 2K qualification and frozen evidence.

D256 CP8 AllGather complete attention comparison

At B64/S8192/Hq32/Hkv8/D256, BF16 causal THD, 524288 tokens, CP8/AllGather, the integrated OSS backward reduces latency relative to the cuDNN backend, but remains slower than FA4. Forward stays on cuDNN; #3705 forward is not enabled. All providers use the same TE communication implementation.

Eight production B200 GPUs use 148 SMs each, default 1000 W limits and NV18 links. The frozen environment is TE forward-v195 (equivalent native backward from this head), FE197 + #1696/#1690 + the packed D128 terminal-credit repair (D256 kernel unchanged), 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.

Provider	Median complete-CP latency	Peak allocated memory per rank
cuDNN backend	77.112 ms	54.52 GiB
FROST backward; cuDNN forward	70.483 ms	46.52 GiB
FlashAttention 4	67.895 ms	24.52 GiB

Median paired latency change: FROST/cuDNN -8.86%; FROST/FA4 +3.79%. Negative means lower latency. FROST is faster than cuDNN in 6/6 orders and faster than FA4 in 0/6. The prior CP4 advantage does not extend to this CP8 coordinate; the memory gap also remains.

All 576 full output/gradient comparisons pass (maximum relative RMS 0.005357), along with 96 independent CPU FP64 comparisons through actual small-shape CP8 (maximum relative RMS 0.005143). Six provider permutations use 18 fresh eight-rank process sets, 144 workers and 18 admitted timing windows. Source/binary/route checks, 1000 W/NV18 hardware checks, sampled HBM-clock/temperature/throttle checks, and stability gates pass. This is a one-node complete-CP screen, not a cross-node result or original-324-grid win. Earlier CP4 results keep their own frozen source and protocol attribution. CP8 qualification and frozen evidence.

[YangXu1990uiuc]

Codex bot review — gpt-6-astra

Reviewed cf04d4e11821b3b93b76915a03637bc2b572b460. Thank you for your contribution! Comment-only review, as requested.

I did not find a new correctness blocker in the supported scope. The separation between cached declarations and per-call pointers/stream/workspace looks sound. Keeping the complete native backward chain, checking the uniform-layout promise on device on every execution, and restricting output-fill removal to a separately guarded strict-mode contract are the right boundaries.

Two follow-ups:

1. [P2] Connect the new regression suite to a qualified QA target before graduating the draft. qa/L0_pytorch_unittest/test.sh explicitly enumerates attention test files, and no checked-in QA launcher currently invokes test_frost_attn.py. Simply adding this file does not give the bridge or the no-zero policy ongoing CI coverage. A dedicated SM100 target with the bridge enabled and compatible FE/DSL dependencies would be sufficient; it should verify that the native-route tests actually execute, including the isolated metadata-guard children. This need not impose the experimental dependency stack on every ordinary L0 lane.

2. [P3] Qualify the D256 bottom-right rectangular support statement and retain a decline/fallback test. With S_q=256, S_kv=128, uniform THD forward works, but the equivalent dense FE backward declaration rejects max_s_q > max_s_kv. I verified that mode 1 correctly falls back with correct gradients (including exactly zero dQ on the leading keyless rows), while mode 2 reports the unsupported declaration. Dense BSHD is rejected by the existing forward path before this bridge. This is a support/documentation boundary, not a silent numerical failure; the current wording should make clear that rectangular support does not include both orientations.

Validation performed independently on one SM100 GPU:

• All 61 committed tests passed, including changed bindings/streams, all tested AllGather step geometries, poisoned CUDA-graph replay, actual output-fill counts, and all five isolated invalid-metadata checks.
• 8 additional GPU checks passed for D128 MHA/GQA with physical gaps and an empty sequence, dense MHA, and repeated replay through supported/fallback routes. Gradients were checked against CPU FP64 references.
• Separately verified the D256 Q-longer-than-KV strict decline and numerically correct THD fallback described above.
• A small alternating no-zero off/on measurement at B=2, Hq/Hkv=4/1, D=128, S=128/1024 showed the intended host-submission reduction and no observed replay regression. Gradients were checked again after timing and poisoned replay. This is not a distributed-CP performance qualification.

I verified exact-head Python/test imports and reused the author's rebuilt native binaries after checking their hashes and matching native source snapshots; this was not an independent clean build. I did not independently rerun multi-GPU CP, JAX, the full L0 matrix, or Compute Sanitizer. The draft's stated stable-FE-export and broader-qualification work therefore remains relevant.