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

[Copilot]

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 guard — LaunchZeroOutputForFullyMaskedBatches 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]

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

[Copilot]

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]

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]

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]

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]

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=true → Q_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]

Related to #29717?

[mastryukov1990]

@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.

[namgyu-youn]

@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.

[tianleiwu]

Review: cuDNN SDPA decode tier (Phase 1)

Verdict: COMMENT (not approving yet — the author notes CUDA was not buildable in their environment, so compilation + GPU tests are unvalidated, and this is an explicitly phased draft).

Overall

This is a well-scoped and unusually well-documented change. The new RunCudnnSdpaAttention faithfully mirrors RunFlashAttention's Path-1/prompt population, the cudnn_sdpa::run(...) and is_supported(...) argument order matches the wrapper header exactly, and the whole design correctly stays CUDA-graph-capturable by reading the valid KV length device-side (LaunchConvertNonpadKvSeqlenToFlashSeqlensK) with no D2H sync — the exact failure mode that sank #29689.

Things that are done well

• Cascade placement — cuDNN is gated in front of Flash/MEA/Unfused and is intentionally not wrapped in USE_FLASH_ATTENTION/USE_MEMORY_EFFICIENT_ATTENTION, matching GQA; hoisting has_output_qk to an unconditional [[maybe_unused]] is the correct consequence.
• Fully-masked-batch guard — zeroing nonpad==0 rows after run() is a genuine spec-equivalence requirement (cuDNN emits NaN over an all -inf row), applied on the BSNH buffer before the BSNH→BNSH transpose with the correct per-batch element stride. It cannot be a host-side gate without breaking capture; that reasoning is sound.
• is_causal not gated — correct for s_q==1: is_supported skips the causal branch and run() drops causal masking, so is_causal={0,1} collapse to the identical padding-only frontier. Confirmed against cudnn_flash_attention.cc.
• Parameter consistency — present_key/value population uses kv_sequence_length, which equals total_sequence_length (used for cuDNN strides) precisely because past_sequence_length == 0 is enforced. Correct, though implicit.
• Tests — the debug-info routing assertion (SdpaKernel=CUDNN_FLASH_ATTENTION), the ORT_TEST_REQUIRE_CUDNN_SDPA non-skippable CI gate + canary, and the capture/replay-with-mutation test (proving the length is re-read from the device buffer, not baked in) close the exact holes that make a decode-tier test silently pass on MATH fallback.

Cross-cutting notes

• Scope — .agents/skills/ort-transformers-gpu-pytest/SKILL.md (+218 lines) is developer tooling unrelated to the runtime feature. Consider splitting it into a separate PR so the functional change reviews cleanly on its own.
• Merge gating — since the code path could not be compiled/run in the author's environment, please attach a green GPU CI leg that actually shows CUDNN_FLASH_ATTENTION dispatch (e.g. with ORT_TEST_REQUIRE_CUDNN_SDPA=1) before this leaves draft.
• The two concerns from the earlier automated review round (undocumented provider_options; missing routing assertion) are already addressed on this head.

One minor test-robustness nit inline.

[tianleiwu]

Review round 2 — head 76f8cad4

Verdict: COMMENT (still not approving — the PR states CUDA was not buildable in the author's environment, so compilation and GPU tests remain unvalidated).

My previous inline nit (cudnn_decode_supported() probing a hardcoded config) is addressed on this head and that thread is resolved. Re-checking the kernel on this head, the following are confirmed correct and need no action:

• sequence_length_kv = total_sequence_length (cuDNN strides) vs kv_sequence_length (present-cache population) — equal, because past_key == nullptr forces past_sequence_length == 0.
• Fully-masked zero-fill elements_per_batch = q_seq * q_num_heads * v_head_size matches the BSNH scratch allocation in both the 3-D and 4-D branches, and runs before the BSNH→BNSH transpose.
• is_causal intentionally ungated for s_q == 1: is_supported skips the causal branch and run() folds is_causal &&= (s_q > 1), so both values share one cached plan and the same padding-only frontier.
• The unguarded cuDNN block still builds under --disable_contrib_ops (the cmake contrib_ops/cuda/bert/.* filter covers cudnn_fmha/) and is excluded under CUDA_MINIMAL (non-recursive glob over core/providers/cuda/).
• AttentionKernelOptions is a per-provider member, so the per-session sdpa_kernel / ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO manipulation in the tests behaves as intended (no cross-test std::once_flag leakage).

The remaining concerns are about how this lands, not the math. Details inline; two are worth calling out here:

1. The tier becomes default-ON on SM≥90, and that path has no test coverage. AllowCudnnFlashAttentionAuto() is true in the default configuration, so auto_enable_cudnn_flash_attention_ is true and every opset-24 external-cache decode on H100/H200/B200 silently moves off Flash onto this tier. But every test here passes sdpa_kernel=24, which sets has_explicit_kernel_selection_ and therefore disables the auto path — so the route production will actually take is never exercised. Given the code could not be compiled or run by the author, please either make Phase 1 opt-in (enable_ only) until a green GPU leg exists, or add a case that omits sdpa_kernel.

2. No latency measurements. The tier is justified entirely by a ~3 µs decode delta, but the PR contains no before/after numbers, and the implementation adds several per-step launches/allocations (details inline). Please attach a measurement, and note that the ~10 µs target is only reachable when present_key/present_value are not requested — otherwise the full-KV-cache copy dominates any kernel-level win.

Merge gating (repeating from the previous round): please attach a green GPU CI leg that actually shows SdpaKernel=CUDNN_FLASH_ATTENTION dispatch before this leaves draft. And .agents/skills/ort-transformers-gpu-pytest/SKILL.md (+218 lines of developer tooling) would still review more cleanly as a separate PR.