You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Design: Add a cuDNN SDPA decode tier to the ONNX Attention CUDA kernel
Status: Design draft, ready for an execution team to pick up Op:Attention (kOnnxDomain, opset 23/24), CUDA EP Primary file:onnxruntime/core/providers/cuda/llm/attention.cc Blueprint op:GroupQueryAttention (com.microsoft, CUDA) — copy its proven cuDNN integration Context: follow-up to the closed issue #29686 / PR #29689 (Flash split-sizing; wontfix — host D2H sync breaks CUDA graph capture)
1. Problem & motivation
The standard ONNX Attention CUDA kernel dispatches only through Flash → Memory-Efficient (MEA) → Unfused (attention.cc:1378-1379, === KERNEL SELECTION CASCADE ===). It has no decode-specialized tier.
The contrib GroupQueryAttention (GQA) op has an XQA → cuDNN SDPA → Flash → MEA → Unfused cascade
(group_query_attention.cc:452-664). Its decode-latency lead comes from the XQA
(~11 µs) and cuDNN SDPA (~10 µs) tiers; the Flash tier (~13–22 µs) is the ceiling
ONNX Attention is stuck at today.
Measured gap (issue #28352, single-token decode, eager): ONNX Attention's Flash
path is ~34 µs (over-partitioned) / ~13 µs (best case) vs GQA ~10–11 µs. Fixing
the Flash split sizing (#29686, wontfix) cannot close this gap — it only reaches
the Flash-tier floor (~13 µs). The gap is structural: a missing kernel tier.
Goal: add a cuDNN SDPA tier so the opset-24 external-KV-cache decode path
(TensorScatter + nonpad_kv_seqlen, the mobius path) reaches GQA-class decode
latency, without the host-side valid-length readback that broke PR #29689.
2. Why cuDNN SDPA first (not XQA)
cuDNN SDPA
XQA
Eligibility breadth
Wide (any clean causal SDPA)
Narrow (single-token, shared buffer, head∈{64,128,256}, group whitelist)
seq_len_q restriction
None (this design still gates to s_q==1 for v1 — see §3)
== 1 only
Buffer requirement
None (reads a full present KV; not past==present)
Hard past_present_share_buffer
Interface
Clean standalone cudnn_sdpa::run(...)
TRT-LLM kernel, strict shape/layout assumptions
Existing ORT dependency
Yes (cuDNN already linked)
Vendored kernels only
Recommendation
Do this first
Follow-up (§9)
Crucially, cuDNN SDPA takes the valid KV length as a device int* array
(mask_sequence_lengths_kv, cudnn_flash_attention.h), read device-side.
This sidesteps the entire host-side num_splits / valid-vs-buffer / placement /
H2D problem Flash hit — no host round-trip, capturable under CUDA graph (see §4.8).
3. Scope & phasing (tightened after review)
The v1 landing must be narrowly gated to its actual value proposition so it
cannot steal traffic from currently-correct paths.
Phase 1 (this work item) — external-cache decode only. Hard eligibility gate,
ALL required:
past_key == nullptr (external cache, not internal past/present)
parameters.is_causal == true
parameters.q_sequence_length == 1 (decode) — see §4.2 for why this is the
only unconditionally-safe causal case
!has_output_qk (cuDNN cannot produce the optional output_qk output)
attn_mask == nullptr (explicit mask routes to MEA/Unfused, unchanged)
no softcap / smooth-softmax / head-sink (ONNX Attention has no head-sink; keep
the check defensive)
dtype ∈ {fp16, bf16}; head_size % 8 == 0 and ≤ 256 (qk and v checked
independently — cudnn_flash_attention.cc:113-116)
Phase 2 (fast-follow, same design) — internal-cache decode (Path 2, past_key/past_value). cuDNN does not require a shared buffer, so it applies
after LaunchConcatNewToPastKV. Actually simpler than Phase 1: the present
buffer equals valid length (no over-allocation), so pass sequence_length_kv = total_sequence_length and mask_sequence_lengths_kv = nullptr (uniform, no per-batch mask). See §4.5.
Not changing the ONNX Attention schema (no new inputs, no opset bump).
Not enabling prefill (s_q>1) through cuDNN in v1. Prefill is correct only
with additional per-batch anchor + query-padding handling;
deferred to a Phase 3 with its own tests. See §4.2.
4. Design
4.1 Cascade placement + build guard
Insert a new cuDNN eligibility block in ComputeInternal's cascade before the #if USE_FLASH_ATTENTION block (~attention.cc:1424), so the priority becomes cuDNN → Flash → MEA → Unfused.
Build guard: do not gate the cuDNN
branch under #if USE_FLASH_ATTENTION. GQA's cuDNN eligibility lives outside
that macro (group_query_attention.cc:552-571), and the wrapper itself stubs out
when cuDNN is too old via CUDNN_MAJOR (cudnn_flash_attention.cc:10-58). Gate
the call site on the same condition GQA uses; rely on cudnn_sdpa::is_stable() / is_supported() + the wrapper's CUDNN_MAJOR handling for availability.
has_output_qk visibility:has_output_qk is
currently defined inside #if USE_FLASH_ATTENTION || USE_MEMORY_EFFICIENT_ATTENTION
(attention.cc:1393). The new cuDNN block sits before it and needs it — widen the
macro to include the cuDNN build condition, or hoist has_output_qk above all
three blocks.
4.2 Why s_q == 1 is the v1 gate (causal-frontier correctness)
This is the highest-risk correctness area.
Grounded conclusion:
Decode (s_q == 1): exact ONNX equivalence, unconditional. cuDNN drops
causal masking when s_q == 1 — the assignment params.is_causal = is_causal && (sequence_length_q > 1) runs at cudnn_flash_attention.cc:430 (the comment at :120-124 is the analogous
eligibility note in is_supported(), not the executed line); the causal bound is
vacuous for a single
query row, and the per-batch padding mask (mask_sequence_lengths_kv) alone
yields j ∈ [0, nonpad[b]-1] — exactly the ONNX frontier. Holds for any
capacity and any heterogeneous nonpad. ✅
Prefill (s_q > 1): conditional, deferred. The ONNX bottom-right frontier
(onnx#8068, offset[b] = nonpad_kv_seqlen[b] - q_sequence_length) is reproduced
by cuDNN's BOTTOM_RIGHT anchor only because the per-batch mask_sequence_lengths_kv drives it (sliding_window_mask() in cuDNN-frontend attn_score_modifiers.h; scaled_dot_product_flash_attention.h:639-653). If it
were ever omitted, cuDNN falls back to the static capacity s_kv → wrong global
offset → causal leak. Additionally, right-padded prefill queries require mask_sequence_lengths_q.
Both are real correctness requirements, so prefill is out of scope for v1 and
gated out by q_sequence_length == 1.
Preserve the existing causal_cross_no_past exclusion (attention.cc:1409).
With nonpad_kv_seqlen != nullptr the required frontier IS bottom-right, so the
external-cache decode case is correctly eligible; pure cross-attention without an
external cache must still be excluded.
Add a RunCudnnSdpaAttention(...) method on Attention<T> parallel to RunFlashAttention / RunMemoryEfficientAttention (attention.h:18-32 for
signature/style). It must reproduce these existing ONNX-Attention semantics —
missing any of them corrupts results:
Q layout / qkv_format — branch on is_bsnh (= parameters.transpose_output, attention.cc:252). There are two legal input ranks and they need different
handling; do not hardcode one format:
3-D inputs (is_bsnh == true, e.g. [B,S,hidden]): Q/K/V are already
physical BSNH. No Q transpose. Pass Q_K_V_BSNH.
4-D inputs (is_bsnh == false, [B,N,S,H]): transpose Q BNSH→BSNH (as
Flash/MEA do, attention.cc:321-334) while the K/V cache stays BNSH → this is
the mixed case, pass Q_K_V_BSNH_BNSH_BNSH, exactly as GQA
(group_query_attention_impl.cu:~1360).
Mirror RunFlashAttention's own if (!is_bsnh) transpose gate
(attention.cc:324). Add parity tests for BOTH ranks (§8).
Output layout — BSNH scratch then transpose to BNSH. cuDNN always writes O
as BSNH regardless of QKV format (cudnn_flash_attention.cc:322-328). The
standard Attention 4D output is BNSH. Mirror RunFlashAttention: write cuDNN
output to a BSNH scratch buffer, then TransposeBSNHtoBNSH when !is_bsnh
(attention.cc:336, :529-531). Passing output = Y directly corrupts layout.
present_key/present_value are separate outputs, not aliases of K/V. They
are independently requested outputs; cuDNN must read the input K/V cache
buffers, and present outputs are populated separately, matching current Path-1
behavior (attention.cc:538-560). Passing present_key into cuDNN before it is
filled reads uninitialized memory.
Valid KV length → device int32. Reuse LaunchConvertNonpadKvSeqlenToFlashSeqlensK
(attention_mask_impl.cu:72), which emits the valid token count (not index,
clamped to [0, total_sequence_length], attention_mask_impl.cu:56-70) — this
is exactly cuDNN's mask_sequence_lengths_kv. No new conversion kernel.
sequence_length_kv = parameters.total_sequence_length (buffer capacity;
physical strides derive from it — cudnn_flash_attention.cc:145-156). GQA passes seqlen_present_kv_cache for the same reason. The per-batch mask bounds the
valid region. (Numerically equal to kv_sequence_length in Path 1 only because past_sequence_length == 0 there — anchor for the reviewer note.)
mask_sequence_lengths_q = nullptr for decode: all q tokens valid. Do not
infer q validity from nonpad_kv_seqlen (that describes KV only) — GQA's is_first_prompt/query-valid-length concept does not transfer. If prefill is ever enabled, synthesize a full/real seq_len_q device buffer.
Fully-masked batch guard (REQUIRED). nonpad_kv_seqlen[b] is arbitrary user input and the converter permits 0.
cuDNN's softmax over an all--inf row produces an unspecified result
(likely NaN — treat as unverified pending the §10 needs-run); every existing
tier defines output = 0 here (Flash early-exits; MEA calls LaunchZeroOutputForFullyMaskedBatches / LaunchZeroFullyMaskedRows, attention.cc:885-921). GQA's blueprint omits this guard only because GQA
guarantees total_seq_lens[b] >= 1 — standard Attention has no such
invariant. Therefore RunCudnnSdpaAttentionmust call the device-side LaunchZeroOutputForFullyMaskedBatches after cudnn_sdpa::run. This is the sole mechanism — do not "exclude nonpad == 0 batches from eligibility",
because eligibility is a host-side gate and nonpad_kv_seqlen is a device
buffer, so excluding on its value would require a D2H copy + sync — exactly the
CUDA-graph-breaking pattern that sank PR Size Flash split-KV from valid KV length in ONNX Attention opset-24 external-cache path (CUDA) #29689 and that §4.8 exists to avoid.
The zero-fill is a spec-equivalence requirement, not defense-in-depth.
4.4 Path 1 data flow (opset-24 external cache — the primary target)
mobius's path: TensorScatter writes new K/V into the pre-allocated key_cache/value_cache in place; nonpad_kv_seqlen [B] int64 device gives valid
length.
K/V are already the full padded cache after TensorScatter — no append/concat.
Fast-follow. past_key/past_value are separate from present; new K/V appended via LaunchConcatNewToPastKVbefore cuDNN (as in RunFlashAttention Path 2, attention.cc:391-460). cuDNN needs no shared
buffer. Simpler than Path 1: the present buffer equals valid length (no
over-allocation), so:
mask_sequence_lengths_kv = nullptr (uniform valid length, no per-batch mask) — or,
if a per-batch device buffer is preferred for uniformity with Path 1, materialize
it explicitly with LaunchFillInt32 (attention_mask_impl.cu:279); never pass a
host scalar.
Recommend landing Phase 1 first; Phase 2 reuses the same RunCudnnSdpaAttention.
4.6 Dispatch debug-info wiring
onnxruntime/core/providers/cuda/llm/attention.{cc,h} contain zero AttentionKernelDebugInfo / AllowDebugInfo() wiring today (verified: grep returns
no matches) — only LOGS_DEFAULT(VERBOSE) routing logs. §8's "assert the cuDNN tier
was selected" test therefore has nothing deterministic to assert on, and parity tests
could silently pass through Flash/MEA/Unfused.
Deliverable: wire the same AttentionKernelDebugInfo block the contrib op uses
(contrib_ops/cuda/bert/attention.cc:196-206: if (kernel_options_->AllowDebugInfo()) { AttentionKernelDebugInfo debug_info; ...; debug_info.Print(...); }) into ComputeInternal, recording which tier ran. AttentionKernelOptions/AttentionKernelDebugInfo are already reachable from this
kernel (GetAttentionKernelOptions() is used at attention.cc:114; attention_kernel_options.h:10-24).
4.7 Enable-flag plumbing & naming
Reuse the EXISTING shared cuDNN option — do NOT add a new option key. AttentionKernelOptions already exposes UseCudnnFlashAttention()
(attention_kernel_options.h:32, backed by the shared ORT_ENABLE_CUDNN_FLASH_ATTENTION
env / provider option) and AllowCudnnFlashAttentionAuto() (:45-46) — GQA reads
exactly these (group_query_attention.cc:135-136). Adding an Attention-specific
option key would create redundant, conflicting env vars across the Attention family.
Deliverable: in Attention<T>, add a local member disable_cudnn_flash_attention_ (matching the local disable_X_ style at attention.h:66-67) and initialize it in the constructor at attention.cc:114-117 from the existingkernel_options->UseCudnnFlashAttention() (and AllowCudnnFlashAttentionAuto() if you want GQA-style auto-enable). Do NOT modify attention_kernel_options.{h,cc}, cuda_execution_provider.h, or attention_common.h — the option is already surfaced.
Failure-safety:is_supported() is a coarse shape/device check;
the actual cuDNN graph build inside run() can still throw. Moving cuDNN ahead of a
working Flash path must not convert success into a hard failure. For the initial
landing, keep cuDNN behind the explicit opt (default-safe), and/or wrap the plan
build so a build failure falls back to the next tier rather than aborting inference
(but see §4.8 — the current wrapper API conflates build vs execution failure).
4.8 CUDA-graph safety — grounded, with a cold-cache caveat
Warm-path (steady-state decode): safe by construction.
No host read of valid length; the converter is a device kernel; the synthesized
q-length buffer (when used) is a stream-ordered Fill explicitly for capture safety
(cudnn_flash_attention.cc:363-368); the cuDNN graph-plan cache is keyed on capacity
(sequence_length_kv), stable across decode steps. This is the real structural
advantage over the closed host-readback PR #29689.
Cold-cache caveat. On a thread-local plan-cache miss, run() builds the cuDNN graph + execution plan and allocates workspace inside the
call — that is not a single capturable launch. The current wrapper API cannot help
the call site here:run() returns void, the plan cache is private
(cudnn_flash_attention.cc:356-358), and it conflates plan-build vs execution failure
through exceptions (:441-474) — so the call site cannot reliably detect an in-capture
cold miss nor "wrap the plan build" to fall back. Realistic v1 stance:
Required invariant: the plan must be built (warmup) before cudaStreamBeginCapture. This matches how the existing GQA/graph flow already
works — ORT's graph regression does two warmup runs before capture
(test_cuda_graph_capture.py:255-268), so steady-state decode is safe. Document
this as a hard invariant.
Optional enhancement (own follow-up): extend the wrapper API to expose
prepared/cache state and a distinguishable build failure, so the call site can
reject or fall back on an in-capture cold miss. If pursued, add cudnn_fmha/cudnn_flash_attention.{h,cc} to §6 Affected files.
Test cold-cache with a direct C++ capture test, not the Python harness (which
always warms twice and therefore cannot exercise a cold miss during capture).
5. Reuse table (copy-adapt from here — user emphasis: don't reinvent GQA/MHA)
Component
Existing implementation (file:line)
cudnn_sdpa::run invocation + arg mapping
group_query_attention_impl.cu:1339-1363
cuDNN handle + temp allocator acquisition
group_query_attention.cc:807-810
Read existing cuDNN option (do NOT add a new key)
AttentionKernelOptions::UseCudnnFlashAttention()attention_kernel_options.h:32, AllowCudnnFlashAttentionAuto():45; consumed by GQA group_query_attention.cc:135-136
(reuse, no change) cudnn_fmha/cudnn_flash_attention.{h,cc} (unless the §4.8
optional wrapper-API enhancement is pursued), attention_mask_impl.cu, attention_kernel_options.{h,cc} — the cuDNN option is already surfaced via UseCudnnFlashAttention(); do NOT add a new key.
Add build guard per GQA (§4.1); ensure cold-miss/graph invariant (§4.8).
Tests (§8).
Phase-1 done when: cuDNN tier is selected for the gated decode shapes (asserted
via debug-info), output matches the Unfused/MEA reference within fp16/bf16 tolerance
(including nonpad==0 and heterogeneous nonpad batches), present outputs match, and
a decode loop under enable_cuda_graph=True captures + replays with matching output.
8. Testing plan (strengthened matrix)
Add/extend under test/python/transformers/test_onnx_attention/
(test_tensorscatter_attention.py, test_cuda_graph_capture.py):
Input rank: BOTH 3-D (is_bsnh==true, Q_K_V_BSNH) and 4-D
(is_bsnh==false, mixed Q_K_V_BSNH_BNSH_BNSH) — guards §4.3 step 1.
head_size ∈ {64, 128, 256} plus at least one non-canonical multiple of 8
(e.g. 72, 248) since cuDNN accepts any %8==0 ≤256 (cudnn_flash_attention.cc:113-116).
group sizes: MHA (kv_heads==q_heads), MQA (kv_heads=1), GQA (kv_heads=2/4).
asymmetric (head_size, v_head_size) e.g. (64,32), (128,64) — cuDNN allows
it and Flash rejects it (attention.cc:1432), so cuDNN newly serves these shapes.
fully-masked batch:nonpad_kv_seqlen[b] == 0 in a mixed batch, e.g. nonpad=[0, 5], capacity=16 → assert output row = 0, not NaN (guards §4.3 step 7).
heterogeneous nonpad across the batch (v1 s_q==1 drops the causal mask, so
this guards the per-batch padding mask col_idx < seq_len_kv[b]; the causal
anchor is exercised only when prefill/Phase 3 lands).
Negative / routing:
FP32 input → assert fallback to MEA/Unfused (cuDNN excluded).
output_qk requested → assert cuDNN NOT selected (guards §3 gate).
attn_mask != nullptr → assert MEA/Unfused path.
causal_cross_no_past with nonpad_kv_seqlen == nullptr → assert NOT cuDNN.
Dispatch assertion: using the §4.6 debug-info hook + test_gqa.py:113-127
helper, assert the selected tier is cuDNN — fail (don't silently pass) if it falls
back. Force cuDNN in CI where supported.
CUDA-graph: decode loop under enable_cuda_graph=True; assert capture succeeds
and outputs match eager (the regression PR #29689 could not pass). Ensure nonpad_kv_seqlen is updated across replay, not only Q/K/V. Add a cold-cache case
(§4.8) distinct from the warmed-capture case.
Availability guard: skip gracefully when cuDNN SDPA unsupported on the test GPU
(mirror GQA test guards).
C++ cold-cache capture test: a direct C++ test that invokes cudnn_sdpa::run
between cudaStreamBeginCapture/EndCapture (§4.8), since the Python harness always
warms twice and cannot exercise a cold miss during capture. (No new attention_kernel_options_test.cc case is needed — the shared ORT_ENABLE_CUDNN_FLASH_ATTENTION behavior is already covered at attention_kernel_options_test.cc:171-197, and the new op member is private; assert
routing via the §4.6 dispatch hook instead.)
9. Follow-up: XQA tier (separate work item)
XQA is the fastest decode kernel but the narrowest. Reconcile with the in-file
comment first:attention.cc:221-237 states XQA is "fundamentally incompatible
with this op's spec design" — that is about Path 2 internal past/present (which
cannot share a buffer, XQA's hard requirement). Path 1 (opset-24 TensorScatter
external cache) is fixed-size/in-place and does plausibly satisfy XQA's sequence_length == 1 + shared-buffer preconditions. Any XQA follow-up must scope
itself to Path 1 and say so, or it reads as contradicting the file.
Blockers: head_size ∈ {64,128,256} + group whitelist
(group_query_attention.cc:485-489), and confirming the XQA-internal KV layout
matches the TensorScatter-written cache. Land cuDNN SDPA first.
10. Risks & questions
Resolved during review (documented, not open):
Q1 Handle/allocator: available — Attention<T> inherits CudaKernel
(attention.h:12), GetCudnnHandle at cuda_kernel.h:132.
[needs-run: cuDNN SDPA tier emits NaN on a nonpad_kv_seqlen[b]==0 batch where
Flash/MEA emit 0; repro=opset24 Attention external-cache decode, B=2, nonpad=[0,5],
capacity=16, fp16, is_causal=1; expect=cuDNN out[batch0]=NaN vs reference 0;
cost=cheap] — confirms §4.3 step 7 is load-bearing before landing.
[needs-run: a cold cudnn_sdpa graph-plan cache miss is CUDA-capturable;
repro=invoke cudnn_sdpa::run once between cudaStreamBeginCapture/EndCapture for
fp16 B=1,Hq=8,Hkv=1,Sq=1,Skv=128,D=64; expect=end-capture succeeds and replay
matches eager; cost=cheap] — validates the §4.8 warmup invariant.
attn_mask-present product behavior for opset-24 padding-mask usage (Q6): confirm
routing to MEA/Unfused is the desired behavior.
Graph-plan cache is unbounded (TODO at cudnn_flash_attention.cc:356-358);
Phase-1 fixed-shape decode keeps it bounded, but note if scope widens.
Design: Add a cuDNN SDPA decode tier to the ONNX
AttentionCUDA kernelStatus: Design draft, ready for an execution team to pick up
Op:
Attention(kOnnxDomain, opset 23/24), CUDA EPPrimary file:
onnxruntime/core/providers/cuda/llm/attention.ccBlueprint op:
GroupQueryAttention(com.microsoft, CUDA) — copy its proven cuDNN integrationContext: follow-up to the closed issue #29686 / PR #29689 (Flash split-sizing; wontfix — host D2H sync breaks CUDA graph capture)
1. Problem & motivation
The standard ONNX
AttentionCUDA kernel dispatches only throughFlash → Memory-Efficient (MEA) → Unfused (
attention.cc:1378-1379,=== KERNEL SELECTION CASCADE ===). It has no decode-specialized tier.The contrib
GroupQueryAttention(GQA) op has anXQA → cuDNN SDPA → Flash → MEA → Unfused cascade
(
group_query_attention.cc:452-664). Its decode-latency lead comes from the XQA(~11 µs) and cuDNN SDPA (~10 µs) tiers; the Flash tier (~13–22 µs) is the ceiling
ONNX
Attentionis stuck at today.Measured gap (issue #28352, single-token decode, eager): ONNX Attention's Flash
path is ~34 µs (over-partitioned) / ~13 µs (best case) vs GQA ~10–11 µs. Fixing
the Flash split sizing (#29686, wontfix) cannot close this gap — it only reaches
the Flash-tier floor (~13 µs). The gap is structural: a missing kernel tier.
Goal: add a cuDNN SDPA tier so the opset-24 external-KV-cache decode path
(TensorScatter +
nonpad_kv_seqlen, the mobius path) reaches GQA-class decodelatency, without the host-side valid-length readback that broke PR #29689.
2. Why cuDNN SDPA first (not XQA)
seq_len_qrestrictions_q==1for v1 — see §3)== 1onlypresentKV; notpast==present)past_present_share_buffercudnn_sdpa::run(...)Crucially, cuDNN SDPA takes the valid KV length as a device
int*array(
mask_sequence_lengths_kv,cudnn_flash_attention.h), read device-side.This sidesteps the entire host-side
num_splits/ valid-vs-buffer / placement /H2D problem Flash hit — no host round-trip, capturable under CUDA graph (see §4.8).
3. Scope & phasing (tightened after review)
The v1 landing must be narrowly gated to its actual value proposition so it
cannot steal traffic from currently-correct paths.
Phase 1 (this work item) — external-cache decode only. Hard eligibility gate,
ALL required:
nonpad_kv_seqlen != nullptr(opset-24 external cache)past_key == nullptr(external cache, not internal past/present)parameters.is_causal == trueparameters.q_sequence_length == 1(decode) — see §4.2 for why this is theonly unconditionally-safe causal case
!has_output_qk(cuDNN cannot produce the optionaloutput_qkoutput)attn_mask == nullptr(explicit mask routes to MEA/Unfused, unchanged)Attentionhas no head-sink; keepthe check defensive)
head_size % 8 == 0and≤ 256(qk and v checkedindependently —
cudnn_flash_attention.cc:113-116)Phase 2 (fast-follow, same design) — internal-cache decode (Path 2,
past_key/past_value). cuDNN does not require a shared buffer, so it appliesafter
LaunchConcatNewToPastKV. Actually simpler than Phase 1: the presentbuffer equals valid length (no over-allocation), so pass
sequence_length_kv = total_sequence_lengthandmask_sequence_lengths_kv = nullptr(uniform, no per-batch mask). See §4.5.Non-goals
Attentionschema (no new inputs, no opset bump).s_q>1) through cuDNN in v1. Prefill is correct onlywith additional per-batch anchor + query-padding handling;
deferred to a Phase 3 with its own tests. See §4.2.
4. Design
4.1 Cascade placement + build guard
Insert a new cuDNN eligibility block in
ComputeInternal's cascade before the#if USE_FLASH_ATTENTIONblock (~attention.cc:1424), so the priority becomescuDNN → Flash → MEA → Unfused.
Build guard: do not gate the cuDNN
branch under
#if USE_FLASH_ATTENTION. GQA's cuDNN eligibility lives outsidethat macro (
group_query_attention.cc:552-571), and the wrapper itself stubs outwhen cuDNN is too old via
CUDNN_MAJOR(cudnn_flash_attention.cc:10-58). Gatethe call site on the same condition GQA uses; rely on
cudnn_sdpa::is_stable()/is_supported()+ the wrapper'sCUDNN_MAJORhandling for availability.has_output_qkvisibility:has_output_qkiscurrently defined inside
#if USE_FLASH_ATTENTION || USE_MEMORY_EFFICIENT_ATTENTION(
attention.cc:1393). The new cuDNN block sits before it and needs it — widen themacro to include the cuDNN build condition, or hoist
has_output_qkabove allthree blocks.
4.2 Why
s_q == 1is the v1 gate (causal-frontier correctness)This is the highest-risk correctness area.
Grounded conclusion:
s_q == 1): exact ONNX equivalence, unconditional. cuDNN dropscausal masking when
s_q == 1— the assignmentparams.is_causal = is_causal && (sequence_length_q > 1)runs atcudnn_flash_attention.cc:430(the comment at:120-124is the analogouseligibility note in
is_supported(), not the executed line); the causal bound isvacuous for a single
query row, and the per-batch padding mask (
mask_sequence_lengths_kv) aloneyields
j ∈ [0, nonpad[b]-1]— exactly the ONNX frontier. Holds for anycapacity and any heterogeneous
nonpad. ✅s_q > 1): conditional, deferred. The ONNX bottom-right frontier(onnx#8068,
offset[b] = nonpad_kv_seqlen[b] - q_sequence_length) is reproducedby cuDNN's BOTTOM_RIGHT anchor only because the per-batch
mask_sequence_lengths_kvdrives it (sliding_window_mask()in cuDNN-frontendattn_score_modifiers.h;scaled_dot_product_flash_attention.h:639-653). If itwere ever omitted, cuDNN falls back to the static capacity
s_kv→ wrong globaloffset → causal leak. Additionally, right-padded prefill queries require
mask_sequence_lengths_q.Both are real correctness requirements, so prefill is out of scope for v1 and
gated out by
q_sequence_length == 1.causal_cross_no_pastexclusion (attention.cc:1409).With
nonpad_kv_seqlen != nullptrthe required frontier IS bottom-right, so theexternal-cache decode case is correctly eligible; pure cross-attention without an
external cache must still be excluded.
4.3
RunCudnnSdpaAttentioncontract (layout, output, present, masking)Add a
RunCudnnSdpaAttention(...)method onAttention<T>parallel toRunFlashAttention/RunMemoryEfficientAttention(attention.h:18-32forsignature/style). It must reproduce these existing ONNX-Attention semantics —
missing any of them corrupts results:
qkv_format— branch onis_bsnh(= parameters.transpose_output,attention.cc:252). There are two legal input ranks and they need differenthandling; do not hardcode one format:
is_bsnh == true, e.g.[B,S,hidden]): Q/K/V are alreadyphysical BSNH. No Q transpose. Pass
Q_K_V_BSNH.is_bsnh == false,[B,N,S,H]): transpose Q BNSH→BSNH (asFlash/MEA do,
attention.cc:321-334) while the K/V cache stays BNSH → this isthe mixed case, pass
Q_K_V_BSNH_BNSH_BNSH, exactly as GQA(
group_query_attention_impl.cu:~1360).Mirror
RunFlashAttention's ownif (!is_bsnh)transpose gate(
attention.cc:324). Add parity tests for BOTH ranks (§8).Oas BSNH regardless of QKV format (
cudnn_flash_attention.cc:322-328). Thestandard
Attention4D output is BNSH. MirrorRunFlashAttention: write cuDNNoutput to a BSNH scratch buffer, then
TransposeBSNHtoBNSHwhen!is_bsnh(
attention.cc:336,:529-531). Passingoutput = Ydirectly corrupts layout.present_key/present_valueare separate outputs, not aliases of K/V. Theyare independently requested outputs; cuDNN must read the input K/V cache
buffers, and present outputs are populated separately, matching current Path-1
behavior (
attention.cc:538-560). Passingpresent_keyinto cuDNN before it isfilled reads uninitialized memory.
LaunchConvertNonpadKvSeqlenToFlashSeqlensK(
attention_mask_impl.cu:72), which emits the valid token count (not index,clamped to
[0, total_sequence_length],attention_mask_impl.cu:56-70) — thisis exactly cuDNN's
mask_sequence_lengths_kv. No new conversion kernel.sequence_length_kv = parameters.total_sequence_length(buffer capacity;physical strides derive from it —
cudnn_flash_attention.cc:145-156). GQA passesseqlen_present_kv_cachefor the same reason. The per-batch mask bounds thevalid region. (Numerically equal to
kv_sequence_lengthin Path 1 only becausepast_sequence_length == 0there — anchor for the reviewer note.)mask_sequence_lengths_q = nullptrfor decode: all q tokens valid. Do notinfer q validity from
nonpad_kv_seqlen(that describes KV only) — GQA'sis_first_prompt/query-valid-length concept does not transfer. If prefill is ever enabled, synthesize a full/realseq_len_qdevice buffer.nonpad_kv_seqlen[b]is arbitrary user input and the converter permits0.cuDNN's softmax over an all-
-infrow produces an unspecified result(likely NaN — treat as unverified pending the §10 needs-run); every existing
tier defines output = 0 here (Flash early-exits; MEA calls
LaunchZeroOutputForFullyMaskedBatches/LaunchZeroFullyMaskedRows,attention.cc:885-921). GQA's blueprint omits this guard only because GQAguarantees
total_seq_lens[b] >= 1— standardAttentionhas no suchinvariant. Therefore
RunCudnnSdpaAttentionmust call the device-sideLaunchZeroOutputForFullyMaskedBatchesaftercudnn_sdpa::run. This is thesole mechanism — do not "exclude
nonpad == 0batches from eligibility",because eligibility is a host-side gate and
nonpad_kv_seqlenis a devicebuffer, so excluding on its value would require a D2H copy + sync — exactly the
CUDA-graph-breaking pattern that sank PR Size Flash split-KV from valid KV length in ONNX Attention opset-24 external-cache path (CUDA) #29689 and that §4.8 exists to avoid.
The zero-fill is a spec-equivalence requirement, not defense-in-depth.
scale = 1/sqrt(head_size)unlessoverridden (
group_query_attention_impl.cu:1377);handle = GetCudnnHandle(context)(available:Attention<T>inheritsCudaKernel,attention.h:12;cuda_kernel.h:132); temp allocator + CUDA stream viaGetOrtStream(mirrorgroup_query_attention.cc:807-810,group_query_attention_impl.cu:1328-1376).4.4 Path 1 data flow (opset-24 external cache — the primary target)
mobius's path: TensorScatter writes new K/V into the pre-allocated
key_cache/value_cachein place;nonpad_kv_seqlen[B] int64 device gives validlength.
nonpad_kv_seqlen→ int32 device seqlens (§4.3 step 4).cudnn_sdpa::runwith the §4.3 contract (mixed format, BSNH outputscratch,
mask_sequence_lengths_kvfrom step 2,mask_sequence_lengths_q = nullptr,is_causal = parameters.is_causal).4.5 Path 2 data flow (internal cache,
past_key/past_value) — Phase 2Fast-follow.
past_key/past_valueare separate from present; new K/V appended viaLaunchConcatNewToPastKVbefore cuDNN (as inRunFlashAttentionPath 2,attention.cc:391-460). cuDNN needs no sharedbuffer. Simpler than Path 1: the present buffer equals valid length (no
over-allocation), so:
sequence_length_kv = parameters.total_sequence_length(=past_sequence_length + kv_sequence_length, NOT+ q_sequence_length;ONNX permits
s_q != s_kv).mask_sequence_lengths_kv = nullptr(uniform valid length, no per-batch mask) — or,if a per-batch device buffer is preferred for uniformity with Path 1, materialize
it explicitly with
LaunchFillInt32(attention_mask_impl.cu:279); never pass ahost scalar.
Recommend landing Phase 1 first; Phase 2 reuses the same
RunCudnnSdpaAttention.4.6 Dispatch debug-info wiring
onnxruntime/core/providers/cuda/llm/attention.{cc,h}contain zeroAttentionKernelDebugInfo/AllowDebugInfo()wiring today (verified: grep returnsno matches) — only
LOGS_DEFAULT(VERBOSE)routing logs. §8's "assert the cuDNN tierwas selected" test therefore has nothing deterministic to assert on, and parity tests
could silently pass through Flash/MEA/Unfused.
Deliverable: wire the same
AttentionKernelDebugInfoblock the contrib op uses(
contrib_ops/cuda/bert/attention.cc:196-206:if (kernel_options_->AllowDebugInfo()) { AttentionKernelDebugInfo debug_info; ...; debug_info.Print(...); }) intoComputeInternal, recording which tier ran.AttentionKernelOptions/AttentionKernelDebugInfoare already reachable from thiskernel (
GetAttentionKernelOptions()is used atattention.cc:114;attention_kernel_options.h:10-24).4.7 Enable-flag plumbing & naming
Reuse the EXISTING shared cuDNN option — do NOT add a new option key.
AttentionKernelOptionsalready exposesUseCudnnFlashAttention()(
attention_kernel_options.h:32, backed by the sharedORT_ENABLE_CUDNN_FLASH_ATTENTIONenv / provider option) and
AllowCudnnFlashAttentionAuto()(:45-46) — GQA readsexactly these (
group_query_attention.cc:135-136). Adding anAttention-specificoption key would create redundant, conflicting env vars across the Attention family.
Deliverable: in
Attention<T>, add a local memberdisable_cudnn_flash_attention_(matching the localdisable_X_style atattention.h:66-67) and initialize it in the constructor atattention.cc:114-117from the existing
kernel_options->UseCudnnFlashAttention()(andAllowCudnnFlashAttentionAuto()if you want GQA-style auto-enable). Do NOT modifyattention_kernel_options.{h,cc},cuda_execution_provider.h, orattention_common.h— the option is already surfaced.Failure-safety:
is_supported()is a coarse shape/device check;the actual cuDNN graph build inside
run()can still throw. Moving cuDNN ahead of aworking Flash path must not convert success into a hard failure. For the initial
landing, keep cuDNN behind the explicit opt (default-safe), and/or wrap the plan
build so a build failure falls back to the next tier rather than aborting inference
(but see §4.8 — the current wrapper API conflates build vs execution failure).
4.8 CUDA-graph safety — grounded, with a cold-cache caveat
Warm-path (steady-state decode): safe by construction.
No host read of valid length; the converter is a device kernel; the synthesized
q-length buffer (when used) is a stream-ordered
Fillexplicitly for capture safety(
cudnn_flash_attention.cc:363-368); the cuDNN graph-plan cache is keyed on capacity(
sequence_length_kv), stable across decode steps. This is the real structuraladvantage over the closed host-readback PR #29689.
Cold-cache caveat. On a thread-local plan-cache miss,
run()builds the cuDNN graph + execution plan and allocates workspace inside thecall — that is not a single capturable launch. The current wrapper API cannot help
the call site here:
run()returnsvoid, the plan cache is private(
cudnn_flash_attention.cc:356-358), and it conflates plan-build vs execution failurethrough exceptions (
:441-474) — so the call site cannot reliably detect an in-capturecold miss nor "wrap the plan build" to fall back. Realistic v1 stance:
cudaStreamBeginCapture. This matches how the existing GQA/graph flow alreadyworks — ORT's graph regression does two warmup runs before capture
(
test_cuda_graph_capture.py:255-268), so steady-state decode is safe. Documentthis as a hard invariant.
prepared/cache state and a distinguishable build failure, so the call site can
reject or fall back on an in-capture cold miss. If pursued, add
cudnn_fmha/cudnn_flash_attention.{h,cc}to §6 Affected files.always warms twice and therefore cannot exercise a cold miss during capture).
5. Reuse table (copy-adapt from here — user emphasis: don't reinvent GQA/MHA)
cudnn_sdpa::runinvocation + arg mappinggroup_query_attention_impl.cu:1339-1363group_query_attention.cc:807-810AttentionKernelOptions::UseCudnnFlashAttention()attention_kernel_options.h:32,AllowCudnnFlashAttentionAuto():45; consumed by GQAgroup_query_attention.cc:135-136mask_sequence_lengths_kv)LaunchConvertNonpadKvSeqlenToFlashSeqlensK,attention_mask_impl.cu:72(emits count,:56-70)LaunchZeroOutputForFullyMaskedBatches/LaunchZeroFullyMaskedRows,attention.cc:885-921; kernelattention_mask_impl.cu:98LaunchFillInt32,attention_mask_impl.cu:279TransposeBNSHtoBSNH/TransposeBSNHtoBNSH,attention.cc:127-146, used:321-334,:529-531attention.cc:538-560LaunchConcatNewToPastKV,attention.cc:391-460contrib_ops/cuda/bert/attention.cc:196-206attention.h:18-32,attention.cc:1442-1500group_query_attention_impl.cu:1376-1377CUDNN_MAJORstubs)cudnn_fmha/cudnn_flash_attention.cc:10-58Test harness reuse:
test_cuda_graph_capture.py:52-57,138-160,221-230test_gqa.py:113-127SdpaKernel.CUDNN_FLASH_ATTENTIONenumbenchmark_mha.py:70test/providers/cuda/test_cases/attention_kernel_options_test.cc:171-197@parameterized.expandidiomtest_mha.py:1305,1381,1390test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py(e.g.TestTensorScatterAttentionCUDAFP16:530)test/python/transformers/test_onnx_attention/test_cuda_graph_capture.py:2366. Affected files
onnxruntime/core/providers/cuda/llm/attention.cc— cuDNN eligibility block inthe cascade +
RunCudnnSdpaAttention+ debug-info wiring; include cuDNN wrapperheader.
onnxruntime/core/providers/cuda/llm/attention.h—RunCudnnSdpaAttentiondecl(
:18-32style) +disable_cudnn_flash_attention_member.cudnn_fmha/cudnn_flash_attention.{h,cc}(unless the §4.8optional wrapper-API enhancement is pursued),
attention_mask_impl.cu,attention_kernel_options.{h,cc}— the cuDNN option is already surfaced viaUseCudnnFlashAttention(); do NOT add a new key.7. Implementation checklist + Phase-1 acceptance criteria
Ordered steps:
disable_cudnn_flash_attention_member, initialized from the existingkernel_options->UseCudnnFlashAttention()— no new option key (§4.7).has_output_qkvisibility (§4.1).#if USE_FLASH_ATTENTIONwith the fullPhase-1 hard gate (§3, §4.2).
RunCudnnSdpaAttentionhonoring the §4.3 contract (rank-branchedqkv_format, BSNH output scratch + transpose, separate present population, deviceseqlens,
mask_sequence_lengths_q=nullptr, fully-masked zero-fill).AttentionKernelDebugInfodispatch recording (§4.6).Phase-1 done when: cuDNN tier is selected for the gated decode shapes (asserted
via debug-info), output matches the Unfused/MEA reference within fp16/bf16 tolerance
(including
nonpad==0and heterogeneousnonpadbatches), present outputs match, anda decode loop under
enable_cuda_graph=Truecaptures + replays with matching output.8. Testing plan (strengthened matrix)
Add/extend under
test/python/transformers/test_onnx_attention/(
test_tensorscatter_attention.py,test_cuda_graph_capture.py):Correctness parity (vs Unfused/Flash reference, fp16 + bf16 tolerance):
is_bsnh==true,Q_K_V_BSNH) and 4-D(
is_bsnh==false, mixedQ_K_V_BSNH_BNSH_BNSH) — guards §4.3 step 1.(e.g. 72, 248) since cuDNN accepts any
%8==0 ≤256(cudnn_flash_attention.cc:113-116).kv_heads==q_heads), MQA (kv_heads=1), GQA (kv_heads=2/4).(head_size, v_head_size)e.g.(64,32),(128,64)— cuDNN allowsit and Flash rejects it (
attention.cc:1432), so cuDNN newly serves these shapes.nonpad_kv_seqlenstrictly <total_sequence_length(padded).nonpad_kv_seqlen[b] == 0in a mixed batch, e.g.nonpad=[0, 5], capacity=16 → assert output row = 0, not NaN (guards §4.3 step 7).nonpadacross the batch (v1s_q==1drops the causal mask, sothis guards the per-batch padding mask
col_idx < seq_len_kv[b]; the causalanchor is exercised only when prefill/Phase 3 lands).
Negative / routing:
output_qkrequested → assert cuDNN NOT selected (guards §3 gate).attn_mask != nullptr→ assert MEA/Unfused path.causal_cross_no_pastwithnonpad_kv_seqlen == nullptr→ assert NOT cuDNN.Dispatch assertion: using the §4.6 debug-info hook +
test_gqa.py:113-127helper, assert the selected tier is cuDNN — fail (don't silently pass) if it falls
back. Force cuDNN in CI where supported.
CUDA-graph: decode loop under
enable_cuda_graph=True; assert capture succeedsand outputs match eager (the regression PR #29689 could not pass). Ensure
nonpad_kv_seqlenis updated across replay, not only Q/K/V. Add a cold-cache case(§4.8) distinct from the warmed-capture case.
Availability guard: skip gracefully when cuDNN SDPA unsupported on the test GPU
(mirror GQA test guards).
C++ cold-cache capture test: a direct C++ test that invokes
cudnn_sdpa::runbetween
cudaStreamBeginCapture/EndCapture(§4.8), since the Python harness alwayswarms twice and cannot exercise a cold miss during capture. (No new
attention_kernel_options_test.cccase is needed — the sharedORT_ENABLE_CUDNN_FLASH_ATTENTIONbehavior is already covered atattention_kernel_options_test.cc:171-197, and the new op member is private; assertrouting via the §4.6 dispatch hook instead.)
9. Follow-up: XQA tier (separate work item)
XQA is the fastest decode kernel but the narrowest. Reconcile with the in-file
comment first:
attention.cc:221-237states XQA is "fundamentally incompatiblewith this op's spec design" — that is about Path 2 internal past/present (which
cannot share a buffer, XQA's hard requirement). Path 1 (opset-24 TensorScatter
external cache) is fixed-size/in-place and does plausibly satisfy XQA's
sequence_length == 1+ shared-buffer preconditions. Any XQA follow-up must scopeitself to Path 1 and say so, or it reads as contradicting the file.
Blockers: head_size ∈ {64,128,256} + group whitelist
(
group_query_attention.cc:485-489), and confirming the XQA-internal KV layoutmatches the TensorScatter-written cache. Land cuDNN SDPA first.
10. Risks & questions
Resolved during review (documented, not open):
Attention<T>inheritsCudaKernel(
attention.h:12),GetCudnnHandleatcuda_kernel.h:132.Q_K_V_BSNH_BNSH_BNSH(§4.3 step 1).sequence_length_kv: pass capacity + per-batchmask_sequence_lengths_kv;frontier equivalence confirmed for decode, mechanism verified (§4.2).
Genuinely open (need a run or product decision):
nonpad_kv_seqlen[b]==0batch whereFlash/MEA emit 0; repro=opset24 Attention external-cache decode, B=2, nonpad=[0,5],
capacity=16, fp16, is_causal=1; expect=cuDNN out[batch0]=NaN vs reference 0;
cost=cheap] — confirms §4.3 step 7 is load-bearing before landing.
repro=invoke cudnn_sdpa::run once between cudaStreamBeginCapture/EndCapture for
fp16 B=1,Hq=8,Hkv=1,Sq=1,Skv=128,D=64; expect=end-capture succeeds and replay
matches eager; cost=cheap] — validates the §4.8 warmup invariant.
routing to MEA/Unfused is the desired behavior.
cudnn_flash_attention.cc:356-358);Phase-1 fixed-shape decode keeps it bounded, but note if scope widens.