Add a cuDNN SDPA decode tier to the ONNX standard Attention CUDA kernel (Phase 1) - #29715
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
There was a problem hiding this comment.
Pull request overview
Adds a cuDNN SDPA-backed decode-specialized tier to the ONNX-domain Attention CUDA kernel (opset 24 external-KV-cache, single-token causal decode), inserting it ahead of the existing Flash → MEA → Unfused cascade, and introduces a Python parity test intended to exercise the new tier via sdpa_kernel provider options.
Changes:
- Add
RunCudnnSdpaAttentionand a tightly gated cuDNN SDPA eligibility block so the effective cascade becomes cuDNN → Flash → MEA → Unfused for the targeted decode case. - Add AttentionKernelDebugInfo wiring to record which tier ran (when debug info is enabled).
- Extend the TensorScatter+Attention Python test harness to accept CUDA provider options and add cuDNN-decode parity cases (with unfused fallback).
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxruntime/core/providers/cuda/llm/attention.cc | Implements cuDNN SDPA decode tier, adds gating in ComputeInternal, and emits tier debug-info prior to early returns. |
| onnxruntime/core/providers/cuda/llm/attention.h | Declares the cuDNN SDPA entrypoint and stores cuDNN enable/auto-enable flags on the kernel instance. |
| onnxruntime/test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py | Adds provider-option plumbing to the test runner and introduces cuDNN-decode parity test cases using sdpa_kernel selection with fallback. |
Consolidated review — Phase-1 cuDNN SDPA decode tierThanks for the Phase-1 implementation. Five reviewers went over the diff against the design (#29714). The core implementation is correct: the spec-level mappings were verified faithful — the decode frontier ( The findings below are about gating, test coverage, and graph-capture safety — not wrong outputs. Major (please address before merge)1. Drop 2. Tests never assert cuDNN was actually selected (can silently pass on MATH fallback). 3. The 4-D BNSH mixed-layout branch is completely uncovered. 4. The CUDA-graph invariant is untested and unenforced. Minor / Nits
Verified — no action neededpresent-population (no double-write, correct length, GQA-identical), fully-masked guard (layout-safe, batch==row for decode), layout/stride branch, includes/layering, |
Copilot Can you fix them, if these are valid? |
…sal=0 decode cases
Addressed the valid gating/coverage findings in 0490bdc: Major 1 — Removed Major 2 — The decode test now enables Minors — Added Deferred (flagged as follow-ups): Major 3 (4-D BNSH branch) and the BF16 / asymmetric |
Re-review (round 2) — after the updateThanks for the revisions. The two headline items are in good shape; four Major items remain, all in the test layer (they don't affect the kernel logic, which was verified correct, but they block trustworthy verification of this path). ✅ Resolved & verified faithful
🔴 Major — please address1. (New) The cuDNN-capability probe reads torch's cuDNN version, not ORT's.
2. The 4-D BNSH mixed-layout path is still uncovered. 3. CUDA-graph capture is still untested. 4. BF16 is still untested. 🟡 Nits (non-blocking)
Verified — no action needed
|
|
Related to #29717? |
|
Xavier Dupré (@xadupre) Yes, they are related through the cuDNN SDPA integration, but they target different Attention implementations and use cases. #29717 extends the contrib |
Scope-B minimal: extends test_tensorscatter_attention.py with opset-24 external-KV-cache q_seq==1 decode fast-path coverage. Pre-existing tests and torch usage are unchanged. - M1: replace the torch.backends.cudnn.version() probe with an observe-dispatch cudnn_decode_supported() that runs a minimal decode with ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO and asserts ORT actually routes to CUDNN_FLASH_ATTENTION (fixes the torch-vs-ORT trustworthiness bug). - M2: 4-D BNSH q/k/v coverage via use_4d in the graph builder and runner. - M3: CUDA graph capture/replay test for the decode tier. - M4: bfloat16 coverage using the torch data_ptr() IO-binding pattern (test_gqa style). - Canary: ORT_TEST_REQUIRE_CUDNN_SDPA makes every decode dispatch assertion non-skippable and fails loud on MATH fallback, closing the all-skips-green hole. - Minors: bit-exact present_k/present_v parity (pure D2D copy), explicit session/io-binding release (del + gc.collect()), bf16 tolerance comment, and stale bf16 comment fixes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Fills the multi-batch validation-loop-index gap: a valid batch-0 must not mask a
bad batch-1. Feeds nonpad_kv_seqlen={0,-3} (batch=2, kv=4) on the CPU EP and asserts
the host-validation rejection (core/providers/cpu/llm/attention_helper.h). Mirrors
GQA's MultiBatchOneBadSeqlensK_OOB. TEST-ONLY; no kernel change.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Ti-Tai Wang (@titaiwangms) — measured the Phase-1 tier against the design's premise (#29714: cuDNN takes external-cache decode from the Flash floor to GQA-class latency), using the #28352/#29684 benchmark plus a tier-on arm. Setup: PR head Decode latency, µs (fp16 / bf16)
ᵃ residual box jitter. Routing asserted per timed session via this PR's debug-info wiring (
Repro: |
…patch verification) Documents the hard-won GPU transformers pytest gotchas: the neutral + private mktemp-cwd rule (avoids the repo-root source-shadow ModuleNotFoundError and the shared /tmp sitecustomize injection vector), LD_PRELOAD lib-pinning when torch's bundled CUDA/cuDNN shadow ORT's, and how to prove real cuDNN SDPA dispatch via ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO + the ORT_TEST_REQUIRE_CUDNN_SDPA non-skippable canary. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…nSdpaAttention
RunCudnnSdpaAttention called context->GetComputeStream() directly, which
does not exist on the plugin-EP adapter OpKernelContext
(onnxruntime::ep::adapter::OpKernelContext exposes GetGPUComputeStream),
breaking all three CUDA Plugin EP builds (Linux, Windows with/without cuDNN).
Switch to the established dual-build-safe helper GetOrtStream(context) and
pass ort_stream.get() to cudnn_sdpa::run(), matching the pattern used by the
other CUDA attention kernels (contrib_ops/cuda/bert/{attention,multihead_attention,
group_query_attention,decoder_attention,paged_attention}.cc).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: f38b9775-a6b6-4422-8f42-9e48e626a358
275719e to
df87b1a
Compare
Description
Adds a cuDNN SDPA tier to the ONNX
Attention(opset 23/24) CUDA kernel, making the effective dispatch cascade cuDNN → Flash → MEA → Unfused. The tier is narrowly gated to the opset-24 external-KV-cache single-token decode path (its actual value proposition) and takes the valid KV length as a deviceint*, so it stays CUDA-graph-capturable — unlike the host-readback approach that sank #29689.This is a draft landing Phase 1 only; correctness-sensitive follow-ups are gated out and flagged inline.
attention.ccComputeInternal) — new cuDNN eligibility block before the Flash block. Hard gate (all required):nonpad_kv_seqlen != nullptr,past_key == nullptr,is_causal,q_sequence_length == 1, noattn_mask/output_qk/softcap, fp16/bf16,cudnn_sdpa::is_supported. Not built-guarded (relies on the wrapper'sCUDNN_MAJORstubs +is_stable()).RunCudnnSdpaAttention— honors existing ONNX-Attention semantics: rank-branchedqkv_format(Q_K_V_BSNHfor 3-D, mixedQ_K_V_BSNH_BNSH_BNSHfor 4-D with Q transposed BNSH→BSNH), BSNH output scratch → transpose to BNSH, device int32 seqlens via the existingLaunchConvertNonpadKvSeqlenToFlashSeqlensK,sequence_length_kv = total_sequence_length(capacity) + per-batch mask,mask_sequence_lengths_q = nullptr, separatepresent_key/present_valuepopulation.LaunchZeroOutputForFullyMaskedBatchesafterrun(). cuDNN emits NaN for anonpad==0row where every other tier defines0; this is a spec-equivalence requirement, not defense-in-depth (can't be a host-side gate without a D2H sync).UseCudnnFlashAttention()/AllowCudnnFlashAttentionAuto()(no new option key); addsenable_/auto_enable_cudnn_flash_attention_members mirroring GQA.AttentionKernelDebugInfodispatch block the op previously lacked, recording which tier ran.sdpa_kernel, covering MHA/MQA/GQA, fully-masked and heterogeneousnonpad; falls back to unfused (still parity-correct) where cuDNN is unavailable.Deferred / needs validation (inline TODOs): Phase 2 (internal
past_key/past_value) and Phase 3 (prefills_q>1) are gated out. Cold cuDNN plan-cache miss under graph capture relies on the warmup-before-capture invariant (documented in the function header); wrapper-API enhancement + a C++ cold-cache capture test are follow-ups. CUDA was not buildable in this environment — compilation and GPU tests need validation on cuDNN 9.3+/SM≥90.Motivation and Context
The standard ONNX
AttentionCUDA kernel has no decode-specialized tier, capping decode latency at the Flash floor (~13 µs) versus GQA's cuDNN/XQA tiers (~10 µs). Fixing Flash split-sizing (#29686, wontfix) can't close this structural gap. cuDNN SDPA is the widest-eligibility, already-linked option and reads valid length device-side, so it reaches GQA-class decode latency without breaking CUDA-graph capture.