Skip to content

Add a cuDNN SDPA decode tier to the ONNX standard Attention CUDA kernel (Phase 1) - #29715

Merged
Ti-Tai Wang (titaiwangms) merged 12 commits into
mainfrom
copilot/add-cudnn-sdpa-decode-tier
Jul 29, 2026
Merged

Add a cuDNN SDPA decode tier to the ONNX standard Attention CUDA kernel (Phase 1)#29715
Ti-Tai Wang (titaiwangms) merged 12 commits into
mainfrom
copilot/add-cudnn-sdpa-decode-tier

Conversation

Copilot AI commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

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 device int*, 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.

  • Cascade (attention.cc ComputeInternal) — 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, no attn_mask/output_qk/softcap, fp16/bf16, cudnn_sdpa::is_supported. Not built-guarded (relies on the wrapper's CUDNN_MAJOR stubs + is_stable()).
  • RunCudnnSdpaAttention — honors existing ONNX-Attention semantics: rank-branched qkv_format (Q_K_V_BSNH for 3-D, mixed Q_K_V_BSNH_BNSH_BNSH for 4-D with Q transposed BNSH→BSNH), BSNH output scratch → transpose to BNSH, device int32 seqlens via the existing LaunchConvertNonpadKvSeqlenToFlashSeqlensK, sequence_length_kv = total_sequence_length (capacity) + per-batch mask, mask_sequence_lengths_q = nullptr, separate present_key/present_value population.
  • Fully-masked-batch guardLaunchZeroOutputForFullyMaskedBatches after run(). cuDNN emits NaN for a nonpad==0 row where every other tier defines 0; this is a spec-equivalence requirement, not defense-in-depth (can't be a host-side gate without a D2H sync).
  • Option plumbing — reuses the shared UseCudnnFlashAttention() / AllowCudnnFlashAttentionAuto() (no new option key); adds enable_/auto_enable_cudnn_flash_attention_ members mirroring GQA.
  • Debug-info wiring — adds the AttentionKernelDebugInfo dispatch block the op previously lacked, recording which tier ran.
  • Tests — cuDNN decode parity test forcing cuDNN via sdpa_kernel, covering MHA/MQA/GQA, fully-masked and heterogeneous nonpad; 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 (prefill s_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 Attention CUDA 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.

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

Copilot AI changed the title [WIP] Add a cuDNN SDPA decode tier to ONNX Attention CUDA kernel Add a cuDNN SDPA decode tier to the ONNX standard Attention CUDA kernel (Phase 1) Jul 14, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 RunCudnnSdpaAttention and 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.

@titaiwangms

Copy link
Copy Markdown
Contributor

Consolidated review — Phase-1 cuDNN SDPA decode tier

Thanks 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 (sequence_length_kv = total_sequence_length + per-batch padding mask, causal dropped for s_q==1), the present-population length (kv_sequence_length == total_sequence_length when past==0, byte-identical to RunFlashAttention's Path-1), the fully-masked zeroing (batch-level == row-level for decode), and the BSNH/BNSH layout/stride branch. Includes/layering, the unconditional has_output_qk, and the emit_debug_info wiring (all four return paths) are also fine. No Critical issues.

The findings below are about gating, test coverage, and graph-capture safety — not wrong outputs.

Major (please address before merge)

1. Drop is_causal from the eligibility gate for the q_sequence_length == 1 external-cache case (dead-code risk).
For s_q == 1, cuDNN collapses causal masking to a no-op: params.is_causal = is_causal && (sequence_length_q > 1) (cudnn_flash_attention.cc:430) forces false for both attribute values, so is_causal=0 and is_causal=1 enter the identical code path and both reduce to padding-only attention j ∈ [0, nonpad[b]−1]. Keeping parameters.is_causal && in the gate therefore buys zero correctness margin, while the repo's own decode contract documents is_causal=0 as the correct TensorScatter-decode setting (test_tensorscatter_attention.py:464-468, and every pre-existing decode test uses is_causal=0). If the production graph emits is_causal=0, this tier is never selected even though the new test (the only is_causal=1 decode case) passes. q_sequence_length == 1 already does the only job is_causal was gating for (excluding prefill). Fix: remove is_causal from the q_seq==1 external-cache gate and add an is_causal=0 row to the decode test matrix.

2. Tests never assert cuDNN was actually selected (can silently pass on MATH fallback).
The new tests force sdpa_kernel = CUDNN_FLASH_ATTENTION | MATH and assert only numeric parity / no-NaN. On any box that isn't SM≥80 + cuDNN≥9.3, routing falls back to the unfused MATH kernel and the test still reports green — so the cuDNN path, the fully-masked guard, and the debug-info wiring can all be broken with every new test passing. Fix: use the debug-info hook this PR adds (AttentionKernelDebugInfo / ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1, the get_sdpa_kernel_from_debug_info pattern in test_gqa.py:113-127) to assert SdpaKernel=CUDNN_FLASH_ATTENTION on configs expected to support it; keep the parity tests for the fallback path. Do not treat fallback as success.

3. The 4-D BNSH mixed-layout branch is completely uncovered.
run_tensorscatter_attention builds 3-D (BSNH) inputs, so the tests only exercise the Q_K_V_BSNH branch. The 4-D input path — which does the Q-transpose, the BSNH→BNSH output transpose, and the mixed Q_K_V_BSNH_BNSH_BNSH cuDNN layout — has never executed once. Fix: add at least one 4-D [batch, num_heads, seq, head_size] case to _CUDNN_DECODE_CASES.

4. The CUDA-graph invariant is untested and unenforced.
The PR comments that the cuDNN plan must be built before cudaStreamBeginCapture, but adds no graph-capture test and no cold-miss detection — a thread-local cache miss still calls build_graph() inside capture. The existing test_cuda_graph_capture.py uses q_sequence_length=16, so the new q_seq==1 gate excludes it. This is the design's own Phase-1 acceptance bar. Fix: add a cuDNN decode graph-capture test (changing nonpad_kv_seqlen across replay, with a dispatch assertion) plus the direct C++ cold-cache-during-capture test, or add capture detection before plan construction.

Minor / Nits

  • Auto-enable on SM90 with no Flash fallback (a cuDNN plan-build failure aborts inference rather than falling through to Flash): this is inherited GQA behavior, mirrored verbatim (group_query_attention.cc:561-576, where use_flash_attention = !use_cudnn_sdpa), not a defect this PR introduces. Changing it belongs in shared cudnn_sdpa resilience work across both ops — not a blocker here.
  • Missing defensive guard: add ORT_ENFORCE(parameters.q_sequence_length == 1, ...) alongside the existing nonpad_kv_seqlen/past_sequence_length enforces — q_seq==1 is the actual correctness boundary but is the one gate condition with no runtime guard if the cascade is loosened later.
  • Naming: the class now mixes disable_flash_attention_ / disable_memory_efficient_attention_ with enable_cudnn_flash_attention_ / auto_enable_cudnn_flash_attention_. This mirrors GQA (fine as a rationale) but reverses the reviewed design decision without a note — please either match the local disable_* convention or add a comment stating the GQA-parity reason.
  • Coverage gaps: only fp16 and head_size=64 are tested. Add BF16 and an asymmetric v_head_size case (both are admitted by the gate and have distinct type/layout handling).
  • Docstring accuracy: the test docstring says cuDNN requires "SM≥90 (Hopper/Blackwell)", but the forced path only needs SM≥80 + cuDNN≥9.3 (cudnn_flash_attention.cc:113); SM≥90 is only the auto-enable heuristic. Please correct.
  • Forward-looking (not a Phase-1 blocker): the bundled numpy_attention_ref uses a capacity-anchored causal offset (kv_seq − q_seq); harmless for s_q==1, but if reused for Phase-3 prefill with s_q>1 heterogeneous batches it diverges from the per-batch (nonpad[b] − q_seq) onnx#8068 frontier. Flag it in the reference for whoever picks up prefill.

Verified — no action needed

present-population (no double-write, correct length, GQA-identical), fully-masked guard (layout-safe, batch==row for decode), layout/stride branch, includes/layering, has_output_qk scope, and emit_debug_info completeness were all checked and are correct.

@titaiwangms

Copy link
Copy Markdown
Contributor

Consolidated review — Phase-1 cuDNN SDPA decode tier

Thanks 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 (sequence_length_kv = total_sequence_length + per-batch padding mask, causal dropped for s_q==1), the present-population length (kv_sequence_length == total_sequence_length when past==0, byte-identical to RunFlashAttention's Path-1), the fully-masked zeroing (batch-level == row-level for decode), and the BSNH/BNSH layout/stride branch. Includes/layering, the unconditional has_output_qk, and the emit_debug_info wiring (all four return paths) are also fine. No Critical issues.

The findings below are about gating, test coverage, and graph-capture safety — not wrong outputs.

Major (please address before merge)

1. Drop is_causal from the eligibility gate for the q_sequence_length == 1 external-cache case (dead-code risk). For s_q == 1, cuDNN collapses causal masking to a no-op: params.is_causal = is_causal && (sequence_length_q > 1) (cudnn_flash_attention.cc:430) forces false for both attribute values, so is_causal=0 and is_causal=1 enter the identical code path and both reduce to padding-only attention j ∈ [0, nonpad[b]−1]. Keeping parameters.is_causal && in the gate therefore buys zero correctness margin, while the repo's own decode contract documents is_causal=0 as the correct TensorScatter-decode setting (test_tensorscatter_attention.py:464-468, and every pre-existing decode test uses is_causal=0). If the production graph emits is_causal=0, this tier is never selected even though the new test (the only is_causal=1 decode case) passes. q_sequence_length == 1 already does the only job is_causal was gating for (excluding prefill). Fix: remove is_causal from the q_seq==1 external-cache gate and add an is_causal=0 row to the decode test matrix.

2. Tests never assert cuDNN was actually selected (can silently pass on MATH fallback). The new tests force sdpa_kernel = CUDNN_FLASH_ATTENTION | MATH and assert only numeric parity / no-NaN. On any box that isn't SM≥80 + cuDNN≥9.3, routing falls back to the unfused MATH kernel and the test still reports green — so the cuDNN path, the fully-masked guard, and the debug-info wiring can all be broken with every new test passing. Fix: use the debug-info hook this PR adds (AttentionKernelDebugInfo / ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1, the get_sdpa_kernel_from_debug_info pattern in test_gqa.py:113-127) to assert SdpaKernel=CUDNN_FLASH_ATTENTION on configs expected to support it; keep the parity tests for the fallback path. Do not treat fallback as success.

3. The 4-D BNSH mixed-layout branch is completely uncovered. run_tensorscatter_attention builds 3-D (BSNH) inputs, so the tests only exercise the Q_K_V_BSNH branch. The 4-D input path — which does the Q-transpose, the BSNH→BNSH output transpose, and the mixed Q_K_V_BSNH_BNSH_BNSH cuDNN layout — has never executed once. Fix: add at least one 4-D [batch, num_heads, seq, head_size] case to _CUDNN_DECODE_CASES.

4. The CUDA-graph invariant is untested and unenforced. The PR comments that the cuDNN plan must be built before cudaStreamBeginCapture, but adds no graph-capture test and no cold-miss detection — a thread-local cache miss still calls build_graph() inside capture. The existing test_cuda_graph_capture.py uses q_sequence_length=16, so the new q_seq==1 gate excludes it. This is the design's own Phase-1 acceptance bar. Fix: add a cuDNN decode graph-capture test (changing nonpad_kv_seqlen across replay, with a dispatch assertion) plus the direct C++ cold-cache-during-capture test, or add capture detection before plan construction.

Minor / Nits

  • Auto-enable on SM90 with no Flash fallback (a cuDNN plan-build failure aborts inference rather than falling through to Flash): this is inherited GQA behavior, mirrored verbatim (group_query_attention.cc:561-576, where use_flash_attention = !use_cudnn_sdpa), not a defect this PR introduces. Changing it belongs in shared cudnn_sdpa resilience work across both ops — not a blocker here.
  • Missing defensive guard: add ORT_ENFORCE(parameters.q_sequence_length == 1, ...) alongside the existing nonpad_kv_seqlen/past_sequence_length enforces — q_seq==1 is the actual correctness boundary but is the one gate condition with no runtime guard if the cascade is loosened later.
  • Naming: the class now mixes disable_flash_attention_ / disable_memory_efficient_attention_ with enable_cudnn_flash_attention_ / auto_enable_cudnn_flash_attention_. This mirrors GQA (fine as a rationale) but reverses the reviewed design decision without a note — please either match the local disable_* convention or add a comment stating the GQA-parity reason.
  • Coverage gaps: only fp16 and head_size=64 are tested. Add BF16 and an asymmetric v_head_size case (both are admitted by the gate and have distinct type/layout handling).
  • Docstring accuracy: the test docstring says cuDNN requires "SM≥90 (Hopper/Blackwell)", but the forced path only needs SM≥80 + cuDNN≥9.3 (cudnn_flash_attention.cc:113); SM≥90 is only the auto-enable heuristic. Please correct.
  • Forward-looking (not a Phase-1 blocker): the bundled numpy_attention_ref uses a capacity-anchored causal offset (kv_seq − q_seq); harmless for s_q==1, but if reused for Phase-3 prefill with s_q>1 heterogeneous batches it diverges from the per-batch (nonpad[b] − q_seq) onnx#8068 frontier. Flag it in the reference for whoever picks up prefill.

Verified — no action needed

present-population (no double-write, correct length, GQA-identical), fully-masked guard (layout-safe, batch==row for decode), layout/stride branch, includes/layering, has_output_qk scope, and emit_debug_info completeness were all checked and are correct.

Copilot Can you fix them, if these are valid?

Copilot AI commented Jul 14, 2026

Copy link
Copy Markdown
Contributor Author

Consolidated review — Phase-1 cuDNN SDPA decode tier

Thanks for the Phase-1 implementation. Five reviewers went over the diff against the design (#29714). **The core implementation is correc...

Addressed the valid gating/coverage findings in 0490bdc:

Major 1 — Removed is_causal from the q_seq==1 eligibility gate (and updated the three related comment blocks + the RunCudnnSdpaAttention header). Confirmed against cudnn_flash_attention.cc:430 that s_q==1 collapses both is_causal values to the same padding-only frontier, so gating on it only left the tier dead for the repo's is_causal=0 decode contract. Added is_causal=0 rows by parametrizing the decode matrix over both values.

Major 2 — The decode test now enables ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO, captures fd-1 output (mirroring test_gqa.py), and asserts SdpaKernel=CUDNN_FLASH_ATTENTION when the platform is cuDNN-capable (cudnn_decode_supported: SM≥80 + cuDNN≥9.3). Fallback configs keep parity-only assertions; a MATH fallback on a capable box now fails instead of silently passing.

Minors — Added ORT_ENFORCE(q_sequence_length == 1, ...); corrected the docstring to SM≥80 + cuDNN≥9.3; added the GQA-parity rationale note for the enable_/auto_enable_ naming; documented provider_options; flagged the capacity-anchored causal offset in numpy_attention_ref as decode-only (Phase-3 caveat).

Deferred (flagged as follow-ups): Major 3 (4-D BNSH branch) and the BF16 / asymmetric v_head_size cases require extending run_tensorscatter_attention, which currently builds only 3-D fp16/fp32 inputs (BF16 also needs ml_dtypes since numpy has no native bf16); Major 4 (graph-capture test + C++ cold-cache-during-capture test) is a larger standalone addition. All three need GPU validation this environment can't provide, so I left them for a targeted follow-up rather than land unvalidated harness changes.

@titaiwangms

Copy link
Copy Markdown
Contributor

Re-review (round 2) — after the update

Thanks 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

  • is_causal dropped from the eligibility gate. Confirmed spec-faithful, and actually stronger than the comment claims: for q_sequence_length == 1, params.is_causal = is_causal && (sequence_length_q > 1) collapses to false for both values (cudnn_flash_attention.cc:430), and since that collapsed value is the plan-cache key, is_causal=0 and is_causal=1 resolve to the same cached graph — bit-identical, not merely equivalent. The critical check also passes: is_supported() gates every causal-specific restriction behind is_causal && sequence_length_q > 1 (cudnn_flash_attention.cc:123), so at this q_seq==1 call site the is_causal argument is provably inert — there is no config where one value is eligible and the other silently falls back. The test now iterating both values is correct.
  • Dispatch assertion via the debug-info hook. The fd-1 capture lifecycle is sound: AttentionKernelDebugInfo::Print uses std::endl (synchronous flush), the reader thread starts before execution and drains concurrently (output well below pipe capacity), and both the normal and exception paths restore the env var and close descriptors.
  • Minors all addressed: ORT_ENFORCE(q_sequence_length == 1) guard, docstring corrected to SM≥80 + cuDNN≥9.3, the numpy_attention_ref Phase-3 caveat, and the naming comment acknowledging the deliberate GQA-parity polarity.

🔴 Major — please address

1. (New) The cuDNN-capability probe reads torch's cuDNN version, not ORT's.
cudnn_decode_supported() gates the routing assertion on torch.backends.cudnn.version() >= 90300, but ORT dlopens its own cuDNN and is_stable() additionally rejects 9.10.0 / 9.10.1 (cudnn_flash_attention.cc:95-100, return version >= 90300 with an explicit 91000/91001 exclusion). Two failure modes:

  • False positive: torch reports 9.10.x (or a different cuDNN than ORT loaded) → ORT falls back to MATH → the assertEqual("CUDNN_FLASH_ATTENTION", ...) fails spuriously.
  • False negative: torch reports old/no cuDNN while ORT loaded a supported one → a real routing regression passes silently through MATH.
    Fix: base the assertion predicate on ORT's own loaded cuDNN version (or expose it), and at minimum mirror the 9.10.0/9.10.1 exclusion.

2. The 4-D BNSH mixed-layout path is still uncovered.
run_tensorscatter_attention builds only 3-D inputs (query=[batch, q_seq, q_hidden], key/value=[batch, total_kv, kv_hidden]), so every _CUDNN_DECODE_CASES case takes is_bsnh=trueQ_K_V_BSNH. The !is_bsnh branch — Q transpose BNSH→BSNH, output transpose BSNH→BNSH, the mixed Q_K_V_BSNH_BNSH_BNSH cuDNN layout, and the D2D present copies — is never executed. Fix: add a use_4d path to the harness that reshapes inputs to [batch, num_heads, seq, head_size] and at least one 4-D case with the routing + output/present assertions.

3. CUDA-graph capture is still untested.
The code documents a "plan must be built before cudaStreamBeginCapture" invariant, but the new tests are eager IO-binding runs only; the existing test_cuda_graph_capture.py covers Flash at q_seq=16, which the q_seq==1 gate excludes. Fix: add a cuDNN decode capture + replay test that mutates the device nonpad_kv_seqlen across replays and first asserts cuDNN routing. (This is the design's own Phase-1 acceptance bar.)

4. BF16 is still untested.
Every new case fixes torch.float16 / TensorProto.FLOAT16, so the distinct is_bf16 cuDNN path has no coverage. Fix: add BF16 routing + parity cases.

🟡 Nits (non-blocking)

  • The Phase-3 caveat's divergence trigger is "heterogeneous nonpad lengths," but the actual condition is q_seq > 1 and any nonpad[b] < capacity (a homogeneous padded batch diverges too). Suggest: "prefill (q_seq > 1) with padded KV (any nonpad[b] < capacity)."
  • The test helpers (_CaptureStdout, _parse_sdpa_kernel, _run_capturing_sdpa_kernel) are a near-copy of the test_gqa.py idiom, and the SDPA bitmask constants (8, 16) duplicate AttentionBackend. Consider a small shared test utility to avoid future drift (nit, not a blocker).
  • emit_debug_info is called with three positional bools at four sites, kept readable only by /*use_cudnn=*/ comments — a future 4th tier risks an argument-order slip with no compiler safety net.

Verified — no action needed

is_causal pass-through vs GQA's constant true (safe — inert at q_seq==1), option plumbing (mirrors GQA, no new key), present-population length/handling, and the debug-info wiring completeness were all re-checked and are consistent.

@xadupre

Copy link
Copy Markdown
Member

Related to #29717?

@mastryukov1990

Copy link
Copy Markdown
Contributor

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 com.microsoft::Attention kernel in contrib_ops/cuda/bert for BERT/BGE-style self-attention with compact 1D sequence-length masks and no KV cache. This PR targets the standard ONNX Attention opset 24 kernel in core/providers/cuda/llm for single-token decode with an external KV cache. They reuse the same underlying cuDNN SDPA runner, but the dispatch paths and tests are separate, so the PRs are complementary rather than replacements.

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>
@titaiwangms
Ti-Tai Wang (titaiwangms) marked this pull request as ready for review July 16, 2026 00:34
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>
@namgyu-youn

namgyu-youn commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

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 34767c6bb1, one Release wheel — the A/B toggles ORT_ENABLE_CUDNN_FLASH_ATTENTION (=0 also disables auto-preference). RTX 6000 Ada (SM89, engages via the explicit env; auto is SM90+ — absolute
numbers not comparable to the #28352 table), CUDA 12.8, cuDNN 9.24, ORT d4ac37e. B=1, H=32/8, D=128, buffer 8192, S_q=1. CUDA-event timer, 100 reps, L2 flushed; cells are medians of 3 runs (shared box, unlockable clocks: single runs jitter up to ±50%, medians ~±10%).

Decode latency, µs (fp16 / bf16)

past attn_scatter (Flash) attn_scatter_cudnn gqa_cudnn gqa_xqa
128 35.7 / 35.7 34.6 / 34.3 31.3 / 31.1 31.0 / 30.6
256 67.2 / 72.4 34.7 / 35.5 31.6 / 31.6 68.2ᵃ / 31.9
512 45.9 / 57.7 35.9 / 35.6 33.3 / 32.8 34.1 / 32.8
1024 91.4 / 53.8 41.4 / 38.9 78.0ᵃ / 37.5 51.8 / 36.5
2048 109.0 / 113.5 48.0 / 48.4 44.4 / 44.0 51.1 / 44.5
4096 79.0 / 77.8 61.0 / 60.6 57.6 / 56.4 72.7 / 53.4

ᵃ residual box jitter. Routing asserted per timed session via this PR's debug-info wiring (SdpaKernel=CUDNN_FLASH_ATTENTION).

  1. The tier does what the design claims: external-cache decode goes from Flash's erratic 36–113 µs (the ONNX Attention CUDA: Performance Optimization Opportunities vs Contrib GQA #28352 split-KV pathology) to a smooth 35–61 µs curve.
  2. Attention-only, the gap to GQA is closed: with TensorScatter removed (--attention-only) the cuDNN arm matches gqa_cudnn within noise at every length (fp16 29.5→55.7 µs); the ~3–5 µs left in the full-graph numbers is the scatter nodes, not the kernel.
  3. 4-D BNSH inputs work at ~3–4 µs transpose cost (sweep_4d.csv), and parity + routing passed in both layouts and both dtypes — data points for review Majors 2/4, though not a substitute for the requested tests.

Repro: bench/README.md on namgyu-youn/onnxruntime@explore/29715-bench, or add the attn_scatter_cudnn env pins to any build containing this PR.

Copilot AI and others added 2 commits July 16, 2026 20:11
…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
@titaiwangms
Ti-Tai Wang (titaiwangms) force-pushed the copilot/add-cudnn-sdpa-decode-tier branch from 275719e to df87b1a Compare July 17, 2026 21:12
@titaiwangms
Ti-Tai Wang (titaiwangms) enabled auto-merge (squash) July 29, 2026 16:19
@titaiwangms
Ti-Tai Wang (titaiwangms) merged commit cac326b into main Jul 29, 2026
90 of 91 checks passed
@titaiwangms
Ti-Tai Wang (titaiwangms) deleted the copilot/add-cudnn-sdpa-decode-tier branch July 29, 2026 16:19
This was referenced Aug 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Perf] Add a cuDNN SDPA decode tier to the ONNX standard Attention CUDA kernel

7 participants