Skip to content

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

Description

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:

  • nonpad_kv_seqlen != nullptr (opset-24 external cache)
  • 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.

Non-goals

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.

4.3 RunCudnnSdpaAttention contract (layout, output, present, masking)

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:

  1. 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).
  2. 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.
  3. 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.
  4. 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.
  5. 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.)
  6. 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.
  7. 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 RunCudnnSdpaAttention must 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.
  8. scale, handle, allocator, stream. scale = 1/sqrt(head_size) unless
    overridden (group_query_attention_impl.cu:1377); handle = GetCudnnHandle(context) (available: Attention<T> inherits CudaKernel,
    attention.h:12; cuda_kernel.h:132); temp allocator + CUDA stream via
    GetOrtStream (mirror group_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_cache in place; nonpad_kv_seqlen [B] int64 device gives valid
length.

  1. K/V are already the full padded cache after TensorScatter — no append/concat.
  2. Convert nonpad_kv_seqlen → int32 device seqlens (§4.3 step 4).
  3. Invoke cudnn_sdpa::run with the §4.3 contract (mixed format, BSNH output
    scratch, mask_sequence_lengths_kv from step 2, mask_sequence_lengths_q = nullptr, is_causal = parameters.is_causal).
  4. Transpose output BSNH→BNSH (§4.3 step 2); populate present outputs (§4.3 step 3).
  5. Apply the fully-masked-batch zero-fill guard (§4.3 step 7).

4.5 Path 2 data flow (internal cache, past_key/past_value) — Phase 2

Fast-follow. past_key/past_value are separate from present; new K/V appended via
LaunchConcatNewToPastKV before 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:

  • 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 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 existing kernel_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
int64 → int32 device seqlens (mask_sequence_lengths_kv) LaunchConvertNonpadKvSeqlenToFlashSeqlensK, attention_mask_impl.cu:72 (emits count, :56-70)
Fully-masked zero-fill guard LaunchZeroOutputForFullyMaskedBatches / LaunchZeroFullyMaskedRows, attention.cc:885-921; kernel attention_mask_impl.cu:98
Fill uniform int32 device buffer (Path 2) LaunchFillInt32, attention_mask_impl.cu:279
Q BNSH→BSNH / output BSNH→BNSH transpose TransposeBNSHtoBSNH / TransposeBSNHtoBNSH, attention.cc:127-146, used :321-334, :529-531
present_key/value population (Path 1) attention.cc:538-560
Concat new→past for Path 2 LaunchConcatNewToPastKV, attention.cc:391-460
Debug-info dispatch wiring pattern contrib_ops/cuda/bert/attention.cc:196-206
Method signature/style + verbose routing logs attention.h:18-32, attention.cc:1442-1500
Scale / stream handling group_query_attention_impl.cu:1376-1377
cuDNN build guard (CUDNN_MAJOR stubs) cudnn_fmha/cudnn_flash_attention.cc:10-58

Test harness reuse:

Component file:line
Kernel-routing assertion via VERBOSE logs test_cuda_graph_capture.py:52-57,138-160,221-230
Debug-info kernel parsing helper / context mgr test_gqa.py:113-127
SdpaKernel.CUDNN_FLASH_ATTENTION enum benchmark_mha.py:70
C++ kernel-options/env tests test/providers/cuda/test_cases/attention_kernel_options_test.cc:171-197
unittest + @parameterized.expand idiom test_mha.py:1305,1381,1390
Existing external-cache decode test (extend) test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py (e.g. TestTensorScatterAttentionCUDAFP16:530)
Existing CUDA-graph capture harness (extend) test/python/transformers/test_onnx_attention/test_cuda_graph_capture.py:236

6. Affected files

  • onnxruntime/core/providers/cuda/llm/attention.cc — cuDNN eligibility block in
    the cascade + RunCudnnSdpaAttention + debug-info wiring; include cuDNN wrapper
    header.
  • onnxruntime/core/providers/cuda/llm/attention.hRunCudnnSdpaAttention decl
    (:18-32 style) + disable_cudnn_flash_attention_ member.
  • (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.
  • Tests (§8).

7. Implementation checklist + Phase-1 acceptance criteria

Ordered steps:

  1. Add disable_cudnn_flash_attention_ member, initialized from the existing
    kernel_options->UseCudnnFlashAttention() — no new option key (§4.7).
  2. Hoist/widen has_output_qk visibility (§4.1).
  3. Add the cuDNN eligibility block before #if USE_FLASH_ATTENTION with the full
    Phase-1 hard gate (§3, §4.2).
  4. Implement RunCudnnSdpaAttention honoring the §4.3 contract (rank-branched
    qkv_format, BSNH output scratch + transpose, separate present population, device
    seqlens, mask_sequence_lengths_q=nullptr, fully-masked zero-fill).
  5. Wire AttentionKernelDebugInfo dispatch recording (§4.6).
  6. Add build guard per GQA (§4.1); ensure cold-miss/graph invariant (§4.8).
  7. 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):

Correctness parity (vs Unfused/Flash reference, fp16 + bf16 tolerance):

  • 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.
  • batch = 1 and > 1.
  • valid < buffer: nonpad_kv_seqlen strictly < total_sequence_length (padded).
  • 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: availableAttention<T> inherits CudaKernel
    (attention.h:12), GetCudnnHandle at cuda_kernel.h:132.
  • Q2 Layout: mixed Q-BSNH / KV-BNSH → Q_K_V_BSNH_BNSH_BNSH (§4.3 step 1).
  • Q3 sequence_length_kv: pass capacity + per-batch mask_sequence_lengths_kv;
    frontier equivalence confirmed for decode, mechanism verified (§4.2).
  • FP32: fp16/bf16 only, FP32 keeps existing cascade (matches GQA). Acceptable.

Genuinely open (need a run or product decision):

  • [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.

Metadata

Metadata

Labels

ep:CUDAissues related to the CUDA execution providermodel:transformerissues related to a transformer model: BERT, GPT2, Hugging Face, Longformer, T5, etc.performanceissues related to performance regressions

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions