Reduce prepared SDPA execution overhead - #1151
YangXu1990uiuc wants to merge 13 commits into
Conversation
YangXu1990uiuc
added
orig-nv-eng
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: NVIDIA/cudnn-frontend/.coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (11)
💤 Files with no reviewable changes (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe changes route dense SDPA through prepared launches, add runtime-bound pointer ABIs for SM100 decode kernels, project native operand facts, preserve 64-bit strides, support strided split-combine outputs, and harden CUDA resource cleanup and backend error handling. ChangesPrepared SDPA execution
Priority: ⬇️ Low Estimated code review effort: 5 (Critical) | ~90 minutes Change: Refactor Sequence Diagram(s)sequenceDiagram
participant Caller
participant FrostSdpaFwdPlan
participant PreparedBinder
participant SDPAKernel
participant SplitCombine
Caller->>FrostSdpaFwdPlan: execute runtime tensors and metadata
FrostSdpaFwdPlan->>PreparedBinder: bind validated facts, workspace, and outputs
PreparedBinder->>SDPAKernel: launch prepared dense or decode kernel
SDPAKernel->>SplitCombine: provide split partials when split-KV is enabled
SplitCombine-->>Caller: write final output and optional LSE
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 50.31% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 161 functions across 30 files. (1 skipped: 1 unsupported.)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
@cudnn-ci-bot run frost python_tests |
|
🏁 Pipeline finished SHA: 26 passed, 7 failed, 5 manual
|
|
@cudnn-ci-bot run frost python_tests |
|
🏁 Pipeline finished SHA: 26 passed, 7 failed, 5 manual
|
YangXu1990uiuc
left a comment
YangXu1990uiuc
left a comment
There was a problem hiding this comment.
Codex bot review · model: gpt-6-astra
Reviewed 50c2e5494512 against base 667fe4ce8ce4.
The design makes sense: one explicit decode entry, bounded caching of pure geometry, and native construction of immutable buffer facts. I found no actionable correctness or approval-blocking issue in this head. The cache hits retain per-call buffer/device/span checks and separate launch frames. Independent B200 measurements confirm the paged CPU-overhead improvement without a material GPU regression in the sampled cases. Comment-only because this PR belongs to the account running the bot.
Validation: Reviewed the complete diff and the final tensor-entry removal, traced graph/standalone/direct-kernel callers, and built the changed native binding in isolation. The final head's native sources match that binding. The exact-head SM100 targeted gate passed 253 tests, with one marker-filtered case: decode, prepared THD/dense execution, geometry caching, native facts, variant-pack normalization and dispatch. On the supported CuTe DSL 4.7 floor, 18 focused tests passed on the final head: migrated direct pointer launches, NHD/HND page sizes, forced splits with empty ranges, Int64 page-table stride, graph decode, buffer rebinding and CUDA Graph capture/replay. An additional independent probe reused one decode plan across B8/Q4/KV1024 -> B3/Q2/KV512 with padded strides -> B8/Q4/KV1024. O and a nontrivial H/B/S Stats layout matched FP64 reference math, output padding guards remained intact, and guarded compiler entry points were never called during execute. Production B200 head/base/base/head comparison against develop 667fe4c, with each arm's matching native extension, cuDNN 9.25.1, CUDA 13.2 and pinned CPU affinity. BF16, B8, Hq/Hkv=32/8, D128, KV1024, Stats enabled. All 16 case/runs passed route assertions, full FP64 O/LSE comparisons and changed-input, poisoned-output CUDA Graph replay. B200 paged Q4/page16 CPU enqueue medians (us): NHD 55.47/61.21 -> 44.21/42.17; HND 60.44/58.66 -> 40.14/42.93. GPU graph time remained about 14.36 us. Dense Q1 GPU time was 12.33 -> 12.39 us and Q4 stayed about 12.32 us. Dense CPU improvements were small or inconclusive; one Q4 head run was noisier. These are component timings, not end-to-end latency. Current-head pre-commit passed. The author already started pipeline 68830210 for 50c2e54; it is still running, so no duplicate CI was triggered.
Limitations: The independent B200 probe covers the stated B8 decode cases, not the author's B64 FlashInfer wrapper workload or a full performance sweep. A dense GPU slowdown initially seen on an engineering board did not reproduce on B200 and is not being reported as a PR defect. Full feature/wheel and architecture CI remains pending on pipeline 68830210. The focused local tests used cuDNN 9.25.1; they do not substitute for that matrix.
Comment only: GitHub does not allow reviewing one's own PR with approval/request-changes.
|
CI follow-up for I inspected all seven failing job summaries. The six Python jobs fail the same four cuTile linear-attention correlated-key cases. SM107 has 12 ragged-decode native backend NVRTC plan-construction failures, plus a backward worker crash; those decode stacks do not run the changed FROST decode kernel. The CI bot labels these jobs as existing nightly failures. I have not independently compared the nightly failure signatures, so this is not a clean-CI claim. No new PR-specific defect was identified in the inspected traces; my review conclusion is unchanged. Codex bot review — model |
|
@cudnn-ci-bot run frost python_tests |
|
🏁 Pipeline finished SHA: 25 passed, 8 failed, 5 manual
|
YangXu1990uiuc
left a comment
YangXu1990uiuc
left a comment
There was a problem hiding this comment.
Codex bot review · model: gpt-6-astra
Reviewed f4189279273d against base 667fe4ce8ce4.
The expanded prepared split path and D256 pointer ABI make sense. I found no correctness or approval blocker in the tested domain, with one nonblocking D256 unsplit GPU regression below. Split execution shows substantial host savings, and the newly prepared path honors the tested KV shape overrides. Comment-only because this is the account owner's PR.
[P2] Keep the small-grid D256 unsplit case in the performance gate (source)
The final pointer host still regresses a supported small-grid unsplit case beyond the approximately 2% range reported for the PR's matrix. On B200, BF16, B=2/Hq=8/Hkv=2/Q=4/KV=1024/D=256, BSHD Q/K/V/O, Stats enabled, and explicitly selected split_kv=1/pack_gqa=True, the dense kernel goes from 14.59 to 16.00 us (+9.7%); paged NHD with page_size=16 goes from 16.36 to 17.10 us (+4.6%). These numbers use ordinary contiguous Stats. A separate padded HBS Stats run shows the same effect. Please add this Q=4, multi-KV-head case to the matrix and investigate preserving a common-layout specialization in the pointer host, analogous to the page-table fast path. This is a localized GPU performance issue, not a numerical failure or an approval blocker; the split-path host savings remain substantial.
Evidence: Independent final f418927 versus previously reviewed 50c2e54 on one allocated B200, ABBA order, identical concrete knobs and dependencies (cuDNN 9.25.1.1, PyTorch 2.13.0, DSL 4.8). Each GPU sample averages 640 executions from captured 32-call graphs; seven samples per process and two processes per source. Every arm passes full FP64 O/LSE references and changed-input/poisoned-output replay. Dense contiguous-Stats process medians: previous 14.588/14.591 us, final 15.998/15.997 us. Paged: previous 16.351/16.359 us, final 17.103/17.103 us. Host enqueue was measured separately; this is not an end-to-end latency claim. Attribution: Introduced between the previously reviewed 50c2e54 and final f418927 by the D256 host/ABI changes. The unchanged device-function AST does not preserve its former static tensor-layout specialization. The precise generated-code cause has not been profiled.
Validation: Latest f418927: 49 D256/pointer-ABI tests passed on SM100, with 3 architecture-specific skips. CuTe DSL 4.7: all 28 selected dynamic-host and pointer-combine cases passed. Native extension reused only after confirming native-source identity with the freshly built bc0054e binding. The immediately preceding bc0054e revision passed 213 focused native-facts/prepared/ABI/combine/decode tests plus 211 split/paged regressions, including other prefill flavors and quantized combine callers. Its three SM107 Int64 ABI cases also passed independently on Rubin. The only subsequent production change is the D256 host branch revalidated above; unchanged coverage is reused explicitly. Final-head B200 ABBA: 8 dense/paged NHD/HND cases, D128/D256, selected split2 and D256 unsplit, all passed complete FP64 O/LSE checks, output-padding canaries and changed-input/poisoned-output CUDA Graph replay. Six padded-output split cases reduced host enqueue from 67–138 us to 50–57 us and captured GPU time by 7–18%; these include eliminating the old output copy. A second ABBA control with contiguous Stats confirms the unsplit finding. Final-head analytic shape-override probe: D128 and D256 prepared split plans, declared KV=1024 overridden to 512 with full-capacity storage retained, produced exact expected O=1 and LSE=log(512). This validates this prepared domain, not every fallback path discussed in #1152. Fresh-process persistent-cache check: compilation forbidden, both unit/nonunit page-table paths execute correctly across four batch/split cases using two reloaded artifacts; zero misses, bypasses, invalid entries or export failures.
Limitations: Author-triggered pipeline 68850105 is still running; its mirror commit was verified against the full f418927 head. The inspected SM107 failure summary contains 12 native backend NVRTC plan-construction failures; two inspected Python jobs each have the same four cuTile correlated-key failures seen in earlier pipelines. This is a prior-pipeline category comparison, not an independent nightly-baseline rerun or a clean-CI claim. No duplicate CI was requested. Performance is a bounded B200 component comparison, not a full shape sweep or FlashInfer/model benchmark. Independent local/remote probes used cuDNN 9.25.1.1; the author's separate 9.26 GA numbers were not substituted for these measurements.
Comment only: GitHub does not allow reviewing one's own PR with approval/request-changes.
|
CI follow-up for Review update: P1, blocking approval — collecting a retired prepared plan can invalidate another graph's capture. This updates my earlier no-blocker conclusion with new CI-driven evidence; please address it before merging. The new D256 prepared path owns raw Independent SM100 controls, using two D256 split2 plans (BF16, B2/Hq8/Hkv2/Q4/KV1024), retaining the retired plan in a Python reference cycle and collecting it inside the other plan's capture:
I also instrumented the new 18-case test with collection at the prepared execution boundary: 13 fail / 5 pass; draining garbage before capture gives 18 pass, with no device frees during capture. These controlled probes establish the lifecycle failure. The CI's two failures have the matching launch/capture-end stack, but that CI run itself was not instrumented to prove the first invalidating call. Please make retired-plan cleanup safe around capture, for example by retaining/reusing the immutable dummy allocation owners or deferring their frees to a safe point. Pre-collecting garbage can stabilize the test, but a test-only workaround would leave the newly exposed runtime lifetime issue. This is approval-blocking because multi-plan capture can fail nondeterministically depending on collection timing; it is not an attention numerical error or a P0. The remaining six Python jobs have the previously observed four cuTile correlated-key failures; SM107 has the previously observed 12 native NVRTC failures. SM103 additionally has 35 native FP8 launch failures that the CI triage labels pre-existing; I have not independently established their nightly baseline. No additional broad CI run was requested. Codex bot review — model |
|
Updated to @cudnn-ci-bot run frost python_tests |
|
@cudnn-ci-bot run frost python_tests |
|
🏁 Pipeline finished SHA: 26 passed, 7 failed, 5 manual
|
YangXu1990uiuc
left a comment
YangXu1990uiuc
left a comment
There was a problem hiding this comment.
Codex bot review · model: gpt-6-astra
Reviewed a2814ba54be7 against base b100e34239d4.
The capture-lifetime fix addresses my earlier approval-blocking P1: the original retired-plan reproducer now passes, and forced collection at the prepared execution boundary passes all 18 cases. Removing owned dummy buffers, detaching saved exception chains, and scoping native destruction under a restored capture mode make sense together. I found no remaining approval blocker in this follow-up, with one nonblocking test-portability issue below. The earlier small-grid D256 performance concern is now represented in the matrix and addressed by the common-stride specialization; the description appropriately retains the remaining original-baseline tradeoff. Comment-only; the existing approval remains in place.
[P2] Use a portable native layout in the new lifetime regression (source)
The unconditional L0 regression uses native ReLU with dim=[1,1,32,32] and stride=[1024,1024,32,1]. On both A100 and H100 with cuDNN 9.26.0.51, neither resource case reaches collection: graph.build() rejects the C=1 layout, while the additional pointwise engine is SM10x-only. Please use a supported native layout for this cross-architecture test. I independently verified dim=[1,32,1,32], stride=[1024,1,1024,32] for both X and Y, with matching tensor views, on A100: collecting either resource during global capture, replay, subsequent native execution, and actual GLOBAL/THREAD_LOCAL mode restoration all pass. This preserves the lifetime detector rather than losing the case to an unconditional architecture skip.
Evidence: Exact a2814ba source and freshly rebuilt native extension, PyTorch 2.13.0+cu130, cuDNN 9.26.0.51: the new module reports 2 failed / 2 passed on H100; both resource cases also fail at the same graph construction on A100. The unchanged module passes 4/4 on SM100. A separate A100 native-layout control passes both resource kinds under both initial thread modes. The failure is before capture, not a recurrence of the cleanup defect. Attribution: Introduced by the new cross-architecture lifetime test in this revision; the underlying backend layout restriction is pre-existing.
Validation: Rebuilt the native binding from a2814ba and verified the loaded Python/native paths. On SM100 with cuDNN 9.26.0.51 and the minimum CuTe DSL 4.7.0, prepared THD/dense/split, geometry, pointer-ABI and D256 decode suites: 147 passed, 3 architecture skips. These include nonunit page-table and output strides through the general branch. Original two-plan D256 split2 cyclic-GC reproducer: passes without pre-draining garbage, with full FP64 O/LSE and changed-input/poisoned-output replay checks. All 18 prepared split capture cases also pass with gc.collect() injected at the prepared execution boundary; no failed device frees during capture. This closes the earlier P1. Native lifetime module on SM100: 4 passed. A100 independent native-layout control: device-buffer and owned-backend-graph collection both preserve capture/replay and later native execution, and restore the real thread capture mode from both GLOBAL and THREAD_LOCAL starting states. The scoped exchange follows CUDA's documented push/pop model. Default-path dispatch suite on A100 with FROST opt-in disabled: 61 passed, including chained backend-exception lifetime and repeated-build diagnostics. A bounded local SM100 ABBA comparison against f418927, using each revision's corresponding native binding, checked dense and paged-NHD D256 unsplit execution. Both arms passed full O/LSE references and changed-input replay; captured GPU time improved in both cases, with no consistent material host-enqueue regression. GC was excluded equally from timing only; the separate forced-GC tests above validate lifetime behavior.
Limitations: Owner-triggered pipeline 68921996 is still running. Its mirrored commit explicitly identifies the full a2814ba head. No duplicate CI was requested; pending architecture lanes are not counted as passes. The timing spot check covers two local SM100 component cases. The broader final-head B200 matrix in the PR description remains author-reported; I did not repeat that allocation or establish a uniform improvement over 50c2e54.
Comment only: GitHub does not allow reviewing one's own PR with approval/request-changes.
|
Pipeline 68921996 has finished for a2814ba. The SM80/SM100/SM103/SM120 FROST SDPA lanes, SM100 linear/GEMM, builds/analysis and C++ sanitizer lanes passed; SM103 reported 5,921 passed / 1,904 skipped. The remaining reds need to be distinguished despite the summary's blanket “Old Failure” labels:
The capture-lifetime blocker remains resolved by the focused tests and the now-passing FROST lanes. The new lifetime-test layout issue remains actionable; this is not an all-green Python run. Codex bot review — model |
3 participants
Before submitting
pre-commit runand committed any formatting changes.AGENTS.mdfor each directory this PR touches and my changes comply.cat-*, one or morearea:*/op:*, and oneorig-*.Affected area
Python API or bindings; FE OSS kernels or CuTeDSL.
Summary
Reduce the host work in an already built FROST SDPA plan, including small-query split-KV execution. Previously split decode stayed on the tensor adapter and rebuilt buffer facts and workspace views on every call. The prepared path now binds the caller's partial workspace and final outputs directly, then launches attention and combine on the same stream.
Migrate SM100 D128 and D256 f16/bf16 decode to the existing explicit pointer ABI. Delete the obsolete f16 tensor host, shape/stride-specialized fake-tensor compile branches, and adapter compile/execute fallback. Keep each kernel's host and pipeline explicit.
Specialize D256 contiguous page-table columns and common O/Stats inner strides in the compiled host, with per-call branches retaining the general Int64 stride path. The device pipeline and public ABI remain unchanged.
Prepare split-KV launches with independent per-call frames. Partial O uses the kernel's required dtype; partial LSE stays natural-log even without user Stats. The combine writes the actual final O/LSE strides, including padded layouts, without an O scratch/copy-back.
Preserve Int64 stride leaves through annotations and compile placeholders, including all seven sibling prefill page-table slots. Page-table indices remain int32.
Cache pure THD/dense geometry and project immutable facts directly from the native variant pack. Remove the unused Python
facts_of_packand nativefacts()projection; tests use an independent oracle.Reject mixed K/V pool layout kinds during standalone support checking. Validate output span/device/alignment/overlap and compiled shapes before launching; metadata inputs also check the producer storage span after shape overrides.
Make FE-owned resource destruction safe during unrelated CUDA Graph capture: restore the thread capture mode after cleanup, avoid retaining backend exception tracebacks, and use existing aligned buffers for unused prepared ABI slots. Zero-KV THD descriptor bindings require no device allocation or initialization.
No public API or attention pipeline redesign. Runtime buffers, workspace and streams are rebound per call; no execute-time JIT or tensor views/copies in the prepared graph path. Split plans retain their declared
(B, S_q)workspace envelope. Quantized callers retain their used tensor ABI and existing scaling/amax semantics. Capability rows and heuristic selection are unchanged.Follows #1119 and #1132; completes the page-table stride coverage left outside #1146.
Validation
Current revision:
a2814ba54.Performance
NVIDIA B200, BF16, cuDNN 9.26 GA. Current
a2814ba54versus precedingf41892792, ABBA order. Each value is the mean of two process medians. Both arms load the native extension corresponding to their revision; this compares the complete revisions.These small cases use B2/Hq8/Hkv2/Q4/KV1024; paged cases use page16 and noncontiguous block IDs. The matrix also includes D128/D256 B64 controls and a long-KV paged D256 control. Small D256 Stats/split GPU time improves 3.48–6.28%; dense split1 O is essentially unchanged. Mean host enqueue increases by 0.005–1.213 µs across the measured settings, with mixed directions in some repeated pairs.
The larger host reductions from the tensor-to-prepared migration were measured earlier (
50c2e5494→f41892792): for example, paged D128 split2 131.22 → 48.57 µs and paged D256 split2 90.80 → 48.30 µs, both B1/Hkv1/Q1/KV32768/page16 without Stats. That earlier matched-source ABBA passed 72/72 settings and held the compatible native extension fixed. Those numbers belong to the earlier heads, not the current revision.A same-input comparison with a separate earlier
50c2e5494batch still leaves up to 4.37% GPU overhead for the current small dense D256 case; it is cross-batch context, not a direct final-versus-50c2 ABBA. The current change recovers much of the small D256 GPU loss while retaining the prepared host path; it does not establish a uniform GPU improvement over the original tensor path.Host timing measures idle single-call public FE
graph.execute()enqueue, with synchronization outside the timer. GPU timing measures warm captured replay on one ordered non-default stream. Compilation, reference computation and capture are excluded. These are component results, not FlashInfer/vLLM or model throughput.Summary by CodeRabbit
Performance
Bug Fixes
Testing