Skip to content

[PyT] [common] Add cuDNN OSS forward dispatch for context parallel attention - #3705

Draft
YangXu1990uiuc wants to merge 8 commits into
NVIDIA:mainfrom
YangXu1990uiuc:yanxu/cp-frost-forward-20261017
Draft

YangXu1990uiuc wants to merge 8 commits into
NVIDIA:mainfrom
YangXu1990uiuc:yanxu/cp-frost-forward-20261017

Conversation

@YangXu1990uiuc

@YangXu1990uiuc YangXu1990uiuc commented Oct 17, 2026 •

Copy link
Copy Markdown

Description

Add an opt-in cuDNN OSS (FROST) forward provider for context-parallel attention. Supported D128 BF16 compact-THD causal training calls execute the complete compiled FROST forward chain through TE's common C++ dispatcher and return the packed FP32 Stats consumed by backward.

This draft is stacked on #3699 at a7e91a56. The current forward head is a708a5a8: incremental diff. It includes compatibility with the known FE prepared host signatures described below. The backward branch is unchanged. The combined GitHub diff includes that dependency until it lands.

Execution and supported scope

  • Python participates in cold preparation: construct the normal FE graph, select its supported forward engine, and retain the complete compiled host entry and immutable arguments. Steady-state C++ binds the current pointers, workspace and stream and calls TVM-FFI directly. The compiled entry includes metadata/remap, descriptor setup and the main SDPA kernel.
  • Support SM100, D128, BF16 training, separate Q/K/V, compact THD, top-left causal padding, and no dropout, bias, sinks, paging or sliding window. Stats remain packed token-major FP32 natural logarithms.
  • Compactness is an explicit promise checked on the GPU every execution, including graph replay: physical origins must equal zero-based cumulative lengths, with valid bounds and nonempty KV for live Q. Unequal sequence lengths are supported; there is no dense conversion. Physical storage holes, AllGather forward subcalls with those holes, D256 forward and noncausal forward are outside this initial scope.
  • Cache plans per declaration/device without holding a C++ cache lock across Python. Validate current operand/workspace alignment, capacity and overlap. Retain immutable FFI values in the cache without adding FFI array-type registration dependencies to the common library.
  • Add default-off, fallback and strict modes, documentation, seventeen directed forward tests and 25 CPU contract tests. The dedicated FROST QA target requires all 104 tests and all twelve isolated metadata-guard children.
  • Accept only the known 27-, 31- and 37-slot FE host signatures. Extra score, bias and KV-mask slots must contain their exact disabled defaults. Unknown signatures and enabled unsupported features are rejected. C++ uses bounded stack storage and passes the actual argument count, retaining native steady-state execution.

The integration selects existing FE kernels; it does not change their arithmetic or communication implementation. The separate private backward exp2 experiment is not part of this PR.

Use and dependencies

Build TE with CuTeDSL support and install compatible FE source, nvidia-cutlass-dsl[cu13]>=4.8.0, and apache-tvm-ffi. This draft validates and consumes FE's private prepared forward metadata. A stable FE export contract and compatible dependency packaging remain prerequisites for graduating the integration.

Current-FE directed requalification used develop 7701e232 plus FE #1690, the terminal dQ-credit synchronization repair in FE #1698, and an isolated FP32 padding-helper correction. This is a composed qualification package, not unmodified develop or the unmodified published FE heads. Ordinary numerical tests do not detect the unrepaired D128 dense/THD barrier-phase violation. The current package passes the directed TE suite and unfiltered terminal-credit synccheck; the FE performance holds remain in force. The current-head 2K and 8K CP comparisons below use this package and explicitly exclude #1696 and #1702. Older comparisons retain their earlier frozen dependency versions.

export NVTE_ENABLE_CUTEDSL_BACKEND=1
export CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1
export NVTE_FUSED_ATTN=1
export NVTE_FLASH_ATTN=0
export NVTE_FUSED_ATTN_FROST_FWD=2          # strict; 1 permits backend fallback
export NVTE_FUSED_ATTN_FROST_FWD_COMPACT=1  # only for the documented layout
export NVTE_FUSED_ATTN_FROST=2             # optional: pair with OSS backward

Warm up each declaration before CUDA Graph capture. Leave the compactness promise disabled for physical holes or unsupported declarations. NVTE_FUSED_ATTN_FROST_FWD=0 retains the existing forward route.

The historical CP measurements below used FE develop 2bcdd1bb plus #1696 for the forward lookup optimization and the complete #1690 patch for the existing D256 pitched-backward regressions. The D256 pitch dependency is not required by the D128 forward implementation itself.

Current compatibility validation

The original current-develop run passed 62 backward cases, then rejected the first forward declaration because FE's host signature had grown beyond the fixed 27-slot bridge. The repair updates both Python export and the native frame while keeping unsupported optional features disabled.

  • 25 CPU contract tests pass: legacy/current signatures, neutral optional slots, malformed or unknown signatures, and retained feature guards.
  • 79/79 original TE GPU tests pass, zero failures/errors/skips, on an exclusive production B200 (148 SM, 1000 W) with cuDNN 9.27.0.42, CUDA 13.2 and the dependency package above. This includes all twelve isolated metadata guards, independent CPU-FP64 O/Stats/gradient comparisons, rebinding, streams, changed-length replay and cached execution with Python preparation callbacks disabled. The native forward plans use the current 37-slot host signature.
  • Eight terminal-credit FE regressions pass ordinarily and all eight pass under unfiltered synccheck with zero reported errors. No unexpected GPU process was observed. This is directed synchronization coverage, not a complete sanitizer matrix.
  • Recompiled the changed common CUDA unit with the existing CUDA 13.2 configuration, linked with 113 hash-verified unchanged parent objects, and loaded that library through the normal current-head TE/PyTorch path. The submitted native code matches the tested code; Python preparation differs only by an AST-equivalent trailing-comma formatting fix.
  • Changed-file pre-commit, Pylint 3.3.1, C++ lint and diff checks pass. The QA launcher now includes the 25 contract tests; the reported CPU and GPU suites were run separately.

The current-head complete-CP comparison below exposes a forward-path slowdown. Isolated host-submission and device-runtime attribution for this combination remains pending. The PR remains a draft pending that investigation, current-head review/CI, remaining integration qualification and the dependency holds.

Current-software original-method CP results at 8K, 32K and 512K

Three additional endpoints complete the original benchmark timing protocol on production B200 GPUs (148 SMs, default 1000 W, NV18 links): BF16 causal THD, Hq32/Hkv8, D128, CP4/A2A, and 524,288 global tokens. The original clipped CPU inputs, padding, per-step cloning/synchronization, ten warmups, five timed iterations, and mean of per-iteration rank maxima are preserved. All providers use the same TE communication implementation. Timing includes forward, backward, helpers and communication.

Global B × S Native cuDNN, ms cuDNN forward + FROST backward, ms Early-exit FROST forward + backward, ms FA4, ms
64 × 8192 48.448 45.353 45.980 44.761
16 × 32768 134.125 126.900 130.422 119.282
1 × 524288 1855.774 1699.237 1698.782 1585.922

The table gives medians of four endpoint means. Within-round paired median time changes for the hybrid are −6.70%, −5.39%, and −8.50% versus native cuDNN, but +0.99%, +6.25%, and +7.12% versus FA4, respectively. All four rounds agree with each comparison's direction. Combined FROST forward/backward remains 2.52%, 9.28%, and 7.13% slower than FA4. Percentages are medians of paired ratios, not ratios of table medians.

Four balanced provider orders per shape use 48 fresh four-rank process sets. All 768 complete O/dQ/dK/dV comparisons pass (maximum relative RMS 0.003417), together with source/route, hardware, process-isolation and timing-stability checks. Independent offline recomputation verifies all 1,599 raw-file hashes and reproduces the timings and comparisons. The software is cuDNN 9.27, CUDA 13.2-built TE, CuTeDSL 4.8 cu13 and the previously documented FA4 build. TE is the published C++ integration at a708a5a8; the FE dependency stack remains the frozen experimental package documented in the preceding results. Only the FROST-forward arm adds the unpublished early-exit decoder; B=1 retains the original decoder. A newer work-prefix experiment is not included in these measurements.

Including the earlier original-method 2K result, the hybrid is faster than FA4 in every round at one of four sampled current-software endpoints, namely 2K. These samples do not reconstruct the historical software, missing CSV or all 81 B200 configuration pairs, and establish no historical-grid win rate.

Peak memory remains a separate limitation. Both FROST routes allocate 15,587,086,848 bytes versus FA4's 15,587,086,336 bytes (+512 bytes), and reserve 18,274,582,528 bytes versus approximately 16.65–16.70 GB for FA4, 1.47–1.51 GiB more reserved memory. Native cuDNN allocates approximately 17.21 GB and reserves 20.42 GB. These are allocator peaks, not a leak diagnosis.

No merge gate is cleared by this update. Public-head regressions, all nine private component precision holds, exact-head GPU CI, dependency and integration requirements remain. Host submission, eager device/wall time and graph replay require their own nonregression evidence; complete-CP gains cannot substitute for those checks. The public source heads are unchanged.

Original-method CP4 A2A follow-up on the current software stack

One endpoint now uses the original published benchmark's timing loop: production B200 (148 SMs, 1000 W), BF16 causal THD, B256×S2048, Hq32/Hkv8, D128, 524,288 tokens, CP4/A2A. The original clipped CPU inputs, padding flags, per-step cloning/synchronization, ten warmups, five timed iterations, and mean of per-iteration rank maxima are preserved. Full O/dQ/dK/dV validation runs outside timing. Four balanced provider orders use 16 fresh four-rank process sets; all 256 full-tensor comparisons pass, with maximum relative RMS 0.003263.

Provider Median endpoint time
Native cuDNN forward/backward 25.850 ms
cuDNN forward + FROST backward 24.446 ms
Early-exit FROST forward + FROST backward 25.538 ms
FA4 25.581 ms

Within-round paired medians: the hybrid reduces time by 5.45% versus native cuDNN and 4.62% versus FA4, faster in all four rounds. FROST forward/backward is approximately tied with FA4 (−0.21%, faster in three of four rounds); this is not a robust forward win. These measurements use the published TE C++ integration at a708a5a8 and the documented FE dependency stack. Only the FROST-forward arm adds the unpublished early-exit decoder, which has not replaced the public FE lookup PR.

Peak memory is retained as a separate result: both FROST routes allocate 15,587,088,896 bytes versus FA4's 15,587,088,384 bytes, a 512-byte difference. FROST reserves 18,274,582,528 bytes versus FA4's 16,846,422,016 bytes, 1.33 GiB more reserved memory. Native cuDNN peaks at 17,214,608,384 allocated and 20,422,066,176 reserved bytes. Reserved allocator memory is not a leak diagnosis; the difference is not dismissed. Temporary preparation callbacks are restored, with no corresponding nanobind shutdown warning across all 64 processes.

Scope: this is one original-method endpoint on current software, not reconstruction of the historical FA4 commit or missing endpoint CSV, and not a claim about the other 80 B200 pairs. The first offline collection rejected TE's "0"/"1" flag strings; a strict four-form boolean parser was verified against all 128 retained route records and negative controls, then the same immutable GPU data were reread. No benchmark rerun or numerical/timing threshold change was made. Existing public-head regressions, all nine private component precision holds, dependency/CI requirements, and merge holds remain.

Current-head D128 2K complete CP comparison — hold

Published TE head a708a5a8, through the actual native C++ execution route, now completes the comparison with the current composed FE package described above: develop 7701e232 + #1690 + the full #1698 synchronization repair + the isolated FP32 padding-helper correction. This package excludes the held #1696 forward lookup optimization and #1702. The separate terminal-credit wait-overlap candidate is also absent.

Workload: B256/S2048/Hq32/Hkv8/D128, BF16 causal compact THD, 524288 tokens, CP4/A2A. Four production B200 GPUs (148 SM, default 1000 W) use cuDNN 9.27.0.42 and CUDA 13.2-built TE. All arms share the TE source and communication implementation. Warm eager latency includes forward, backward, helpers and communication, reporting the slowest rank; prepared Python callbacks are disabled during timing.

Provider Median complete-CP latency
cuDNN backend 25.829 ms
FROST backward; cuDNN forward 24.329 ms
FROST forward and backward 28.411 ms
FlashAttention 4 25.462 ms

Median paired latency changes across four balanced orders:

  • Backward-only / FA4: -4.44%, faster in 4/4 orders.
  • Combined / backward-only: +16.77%, slower in 4/4 orders.
  • Combined / cuDNN backend: +10.00%.
  • Combined / FA4: +11.59%.

All 256 complete O/dQ/dK/dV comparisons, 64 fresh workers and 48 timing windows pass (maximum relative RMS 0.003662). Telemetry/stability gates pass and no foreign GPU process is observed. Peak allocated memory is 15,604,025,856 bytes for native and 13,976,506,368 bytes for either FROST arm; preserved nanobind shutdown leak warnings remain undiagnosed, so this does not qualify lifetime or leak behavior.

Hold for forward performance attribution. The previous favorable 2K results below included #1696 and used different frozen FE versions. This comparison does not isolate the TE ABI repair as the cause. The missing bounded THD lookup is a candidate explanation requiring a controlled comparison; #1696 itself remains held for its separate small-batch regression. This result does not qualify a merge, establish an original-324-grid win rate, or measure the private wait-overlap candidate.

Current-head D128 8K complete CP comparison — hold

The same published TE a708a5a8 and frozen FE package used for the 2K result above also complete B64/S8192/Hq32/Hkv8/D128, BF16 causal compact THD, 524288 tokens, CP4/A2A. This still excludes #1696, #1702 and the private terminal-credit wait-overlap candidate. Hardware, native C++ execution, communication, callback-free timing and measurement definitions match the 2K comparison.

Provider Median complete-CP latency
cuDNN backend 48.569 ms
FROST backward; cuDNN forward 45.464 ms
FROST forward and backward 46.296 ms
FlashAttention 4 44.451 ms

Across the four balanced orders, the median paired latency changes are:

  • Backward-only / FA4: +2.35%, slower in 4/4 orders.
  • Combined / backward-only: +1.73%, slower in 4/4 orders.
  • Combined / cuDNN backend: -4.64%, faster in 4/4 orders.
  • Combined / FA4: +4.05%, slower in 4/4 orders.

All 256 complete O/dQ/dK/dV comparisons, 64 fresh workers and 48 timing windows pass; maximum relative RMS is 0.003660. Telemetry/stability and full-node isolation gates pass. Peak allocated memory is 15,604,026,880 bytes for native and 13,976,603,648 bytes for combined FROST; this is not a lifetime/leak qualification.

Forward remains on hold. Enabling it adds 1.73% here, compared with 16.77% at B256/S2048. Batch and sequence length both changed, so this is evidence of workload dependence, not isolated attribution to the lookup algorithm. A bounded-lookup candidate that preserves the B=1 specialization is undergoing separate production qualification; it is absent from both current CP measurements. Neither measurement qualifies a merge or establishes an original-324-grid win rate.

Historical dependency-specific validation and CP measurements

Historical validation and CP measurements

On production B200 (148 SM), cuDNN 9.27.0.42, CUDA 13.2-built TE, PyTorch 2.14.1+cu130 and CuTeDSL 4.8 with its CUDA 13 compiler:

  • Recompiled the changed common CUDA units and linked them with hash-verified unchanged build objects; no undefined link symbols were added or removed. A fresh FE binding is used. Native source snapshots match the historical 020ed08b implementation after whitespace-only formatting.
  • 79/79 directed tests pass, zero skips: seventeen new forward cases and 62 backward regressions, including all twelve isolated device-metadata guards. Coverage includes independent CPU FP64 O/Stats/gradient checks, MHA/GQA, unequal lengths, changed bindings and streams, changed-input/length and empty-to-live graph replay, and cached execution with Python preparation callbacks disabled.
  • Changed-file pre-commit and diff checks pass at 020ed08b.
  • Follow-up 11aab95c resolves the new Python lint findings: dict literal spelling, redundant parentheses, an unused parameter name, and scoped annotations for explicit ABI checks/compiled CUDA bindings. Pylint 3.3.1 and changed-file hooks pass locally. A normalized AST comparison verifies equivalent preparation code; the native binaries, kernel and GPU qualification remain attributed to 020ed08b. Equivalence evidence. GitHub checks subsequently passed for that historical head.

Native qualification and exact source/binary evidence · Native link audit.

A complete CP4/A2A comparison now passes on four B200 GPUs: B256 × S2048, Hq32/Hkv8, D128 BF16 causal compact THD, 524288 total tokens. All arms use the same TE source and communication implementation. Warm eager timing includes forward, backward and communication and reports the slowest rank's host elapsed time.

Provider Median total latency
cuDNN backend 25.762 ms
FROST backward; cuDNN forward 24.241 ms
FROST forward and backward 24.063 ms
FlashAttention 4 25.452 ms

Adding forward to the FROST backward path reduces complete-CP latency by 0.73%. The combined path reduces latency by 6.60% versus cuDNN and 5.45% versus FA4. Percentages are medians of four per-round paired ratios; all four orders favor the combined path. This is one coordinate on one node, using the FE dependencies above, and is not an original-324-grid win-rate claim or a projection to longer sequences/other communication types.

Four balanced process orders comprise sixteen fresh four-rank sets, 64 workers and 48 timing windows. All 256 full O/dQ/dK/dV comparisons pass (maximum relative RMS 0.003662); power, HBM-clock, temperature and stability checks pass. Untimed route/fill audits verify native backward in both FROST arms, engine 20511 only in the combined arm, and no redundant gradient-output fills. GPU evidence is attributed to 020ed08b; 11aab95c has the equivalent Python preparation spelling described above. Complete CP measurements, route checks and raw hashes.

The same four-arm protocol also completes at B64 × S8192, retaining 524288 tokens, CP4/A2A, Hq32/Hkv8, D128 BF16 causal compact THD and the same FE dependencies. All four arms run on one B200 node; the 2K experiment above ran on a different node.

Provider Median total latency at 8K
cuDNN backend 48.143 ms
FROST backward; cuDNN forward 45.230 ms
FROST forward and backward 45.043 ms
FlashAttention 4 44.143 ms

At 8K, adding forward changes complete-CP latency by -0.33% relative to backward alone. The combined path changes latency by -6.44% versus cuDNN and +1.90% versus FA4; backward alone is +2.51% versus FA4. Negative means lower latency. Ratios are paired within each of four balanced orders. The combined path is faster than FA4 in 0/4 orders, and adding forward improves on backward alone in 4/4 orders.

The 2K advantage does not extend to this 8K coordinate: the integrated path still trails FA4. This PR provides native execution of the supported OSS forward chain; it does not establish a universal performance advantage.

The 8K run passes another 256 full O/dQ/dK/dV comparisons (maximum relative RMS 0.003660), with 64 fresh workers and all 48 timing windows passing the predeclared telemetry/stability checks. Native route, source/binary identity and output-fill checks pass. No private backward kernel experiment is used in either CP comparison, and no original-324-grid wins are added. Complete 8K CP evidence.

An earlier attempt stopped at native-control TCPStore initialization with EADDRINUSE, before that worker executed attention. Its partial timings remain excluded. The successful repeat uses a unique file rendezvous with an otherwise unchanged attention/timing body. Retained startup diagnostic.

The first native QA run also exposed a dependency-snapshot assembly error: the D256 pitch patch had been skipped. That failure is retained; after verifying the applied dependency files, the complete 79-case repeat passed. Dependency diagnostic.

Full GPU L0, JAX execution and a clean sanitizer qualification for this new bridge remain pending. The dedicated CI target still requires owner provisioning/enrollment. This remains a draft pending those integration decisions and maintainer review.

Rolling explanation and evidence.

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.

Follow-up causal THD forward kernel optimization

cuDNN Frontend #1702, head 41be3b14, builds on #1696 and moves one quarter of softmax exp2 work on eligible SM100 BF16 D128-flavor causal THD interior tiles to the existing FMA implementation. Masked/diagonal tiles retain MUFU; paged, split-KV, sliding-window and FP16 causal paths are excluded. This is a separate FROST kernel change; the TE dispatch code and this PR's head are unchanged.

At B32/S16384/Hq32/Hkv8/D128, BF16 causal THD, 524288 tokens, CP4/A2A, two independent B200 nodes show 1.46–1.72% lower complete-CP latency against the repaired FROST forward+backward baseline. Timing includes forward, backward, helpers and communication, with identical TE binaries and communication, four balanced orders, 256 full output/gradient checks and 48 telemetry/stability windows per node. Baseline/candidate/FA4 medians are 74.617/73.342/70.940 ms on one node and 74.517/73.367/70.716 ms on the other. Paired ratios are analyzed separately per node; the optimized path still takes 3.34–3.78% longer than FA4.

The published forward source independently passes 22 numerical/replay cases plus one public-layout rejection on B200, with two cases under each of memcheck, racecheck and synccheck, zero errors/hazards. Tests include both packed Stats layouts, changed-length replay and smaller/asymmetric D128-envelope widths. The original complete-CP measurements retain their frozen experimental-package attribution; they do not remeasure the entire new public FE package or add an original-grid win rate. Existing TE integration/export/packaging requirements remain unchanged.

Two-node complete-CP comparison · Published kernel qualification.

Share the strict compact-layout predicate between PyTorch output allocation and common dispatch, and validate device metadata on every execution. Keep the policy default-off and retain native/fallback initialization.

Add FP64, changed-input, poisoned-replay, actual-fill and isolated metadata-guard tests; all 61 directed tests pass on B200.

Signed-off-by: Yang Xu <[email protected]>
@github-actions github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Oct 17, 2026
@YangXu1990uiuc
YangXu1990uiuc force-pushed the yanxu/cp-frost-forward-20261017 branch from 6b0570b to a708a5a Compare October 19, 2026 13:53
@KshitijLakhani KshitijLakhani changed the title Add cuDNN OSS forward dispatch for context parallel attention [PyT] [common] Add cuDNN OSS forward dispatch for context parallel attention Oct 20, 2026
@YangXu1990uiuc

Copy link
Copy Markdown
Author

Posted by Claude Code, an AI agent acting for @YangXu1990uiuc.

A note for when this FROST path grows an FP8 variant (it is BF16-only today, so nothing to change now). The FROST FP8 rows do not produce Amax_S (cudnn-frontend #602), and a graph that declares it is routed to backend engines only. cudnn-frontend #1684 documents Amax_S as deprecated and measured the cost on Rubin: d256 GQA FP8 prefill runs 1.7-2.3x slower through the backend. TE's fused_attn_fp8.cu declares it today. Leaving it virtual on its own is not enough under DelayedScaling: the S amax history would stay at 0 and scale_s would stay at its initial 1.0. So the FP8 variant should also set the S scale from the fixed bound: 448 for E4M3 without dropout, 448(1-p) with dropout probability p (docs/operations/Attention.md, "Amax S").

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

attention community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants