[CUDA] Add cuDNN paged SDPA decode tier for PagedAttention

[Hariharan Seshadri (hariharans29)]

Description

Adds a cuDNN paged SDPA tier to the CUDA PagedAttention dispatch cascade,
sitting between XQA and FlashAttention 2. Decode-only (max_query_len_bound == 1, token_count == batch_size), unquantized KV cache, standard causal
SDPA — no softcap, sliding window, head sink, or bias.

The gate mirrors GroupQueryAttention's cuDNN tier:

• UseCudnnFlashAttention() (ORT_ENABLE_CUDNN_FLASH_ATTENTION=1 or the
  sdpa_kernel bit) enables it explicitly.
• AllowCudnnFlashAttentionAuto() enables it automatically on sm >= 90.
• ORT_ENABLE_CUDNN_FLASH_ATTENTION=0 is the shared kill switch across
  every cuDNN attention path in the CUDA EP.

No new env var; cuDNN attention configuration stays uniform across
Attention / MultiHeadAttention / GroupQueryAttention / PagedAttention.

Motivation and Context

FlashAttention 2's paged split-KV (Flash Decoding) leaves substantial
throughput on the table for GQA at long context. cuDNN 9.5+ ships a paged
SDPA path that is dramatically faster in that regime.

Benchmarks

RTX 5060 Ti, CUDA 12.8, cuDNN 9.12.0.46, sm_120, 15-rep median.
onnxruntime_benchmark BM_PagedAttentionDecode_Cuda/*, cuDNN paged
(ORT_ENABLE_CUDNN_FLASH_ATTENTION=1) vs. FA2 paged
(ORT_ENABLE_CUDNN_FLASH_ATTENTION=0).

GQA (num_heads > kv_num_heads), H=128:

Shape	past=512	past=2048	past=8192	past=16384
B=1 nH:14 kv:2	0.76×	1.34×	3.42×	5.77×
B=1 nH:32 kv:4	0.83×	1.30×	2.95×	4.80×
B=2 nH:14 kv:2	0.76×	1.24×	2.96×	4.77×
B=2 nH:32 kv:4	1.03×	2.04×	4.89×	4.16×

MHA (num_heads == kv_num_heads):

Shape	past=512	past=2048	past=8192	past=16384
B=1 nH:16 kv:16 H:64	0.71×	0.97×	1.83×	1.28×
B=1 nH:16 kv:16 H:128	0.77×	1.28×	1.24×	1.33×
B=2 nH:16 kv:16 H:64	0.73×	0.91×	0.93×	0.95×
B=2 nH:16 kv:16 H:128	0.78×	1.12×	0.93×	0.96×

GQA at past >= 2048 is the intended target (1.3–5.8× win). MHA long
context is roughly parity. Short-context decode (past=512) shows a
bounded 0.71–0.83× regression (20–25 µs absolute) that only hits very
early in a generation before context grows past 2048. Users that need the
short-context FA2 path there can opt out with
ORT_ENABLE_CUDNN_FLASH_ATTENTION=0.

Changes

Kernel + API surface

• cudnn_fmha/cudnn_flash_attention.{h,cc}: new is_supported_paged(...)
  and run_paged(...) entry points, wrapping cuDNN's frontend
  SDPA_attributes with a page-table graph.
• attention_data.h: three new fields on PagedAttentionData for the
  cuDNN handle, per-batch KV-length scratch, and allocator.
• paged_attention.{h,cc}: constructor reads the new gates
  (enable_cudnn_paged_, auto_enable_cudnn_paged_); ComputeInternal
  adds the cuDNN paged tier with metadata-gated eligibility so no D→H
  readback is triggered.
• paged_attention_impl.{h,cu}: RunPagedAttentionCudnn(...) builds the
  per-batch KV-length scratch and calls cudnn_sdpa::run_paged.

Build

• cmake/onnxruntime_providers_cuda.cmake: suppress nvcc warning
  #20303-D (CUTLASS subbyte_reference.h uses __nv_atomic_load_n with
  a memory-order arg, which nvcc flags as sm_70+ on the sm_60 arch list;
  -Werror all-warnings otherwise promotes it to a hard error).
• cmake/onnxruntime_unittests.cmake: include the microbenchmark in the
  CUDA build (was WebGPU-only).

Tests + benchmarks

• test/contrib_ops/paged_attention_op_test.cc: new cuDNN paged tests and
  eligibility-gate assertions.
• Four pre-existing tests (Cuda_AttentionMetadataShape2Compatibility...,
  Cuda_FlashSplitKvLongContext, Cuda_FlashSplitKvCudaGraphReplay,
  Cuda_FlashSplitKvSkipsShortReplayRange) now set
  ORT_ENABLE_CUDNN_FLASH_ATTENTION=0. These tests force non-cuDNN
  backends by disabling FA/MEA/decoder; after the auto-on-sm≥90 gate they
  would otherwise dispatch to cuDNN instead of the intended tier.
• test/onnx/microbenchmark/paged_attention.cc: parallel _Cuda
  benchmark family alongside the existing WebGPU family.

Testing

• onnxruntime_provider_test.exe --gtest_filter=PagedAttention.* on
  Windows / CUDA 12.8 / cuDNN 9.12.0.46 / RTX 5060 Ti (sm_120):
  22 pass, 16 skipped (WebGPU + EndToEnd requiring absent runtime),
  0 fail.
• 15-rep A/B benchmark sweep as tabulated above.

[Tianlei Wu (tianleiwu)]

Reviewed the cuDNN paged SDPA tier. The design is sound: the stride trick that presents ORT's [num_blocks, block_size, num_heads_kv, head_size] cache as cuDNN's [num_blocks, num_heads_kv, block_size, head_size] container is correct and zero-copy, dropping the causal mask for s_q == 1 matches the reasoning already documented in run(), the gate is genuinely metadata-only (no new D->H readback), and the mutual-exclusion bookkeeping across use_paged_decode / use_flash_attention / use_memory_efficient_attention / the XQA-restore block / the "no backend" error is complete. Reusing UseCudnnFlashAttention() / AllowCudnnFlashAttentionAuto() instead of inventing a new env var is the right call, and the max_seq_len_kv == max_num_blocks_per_seq * block_size alignment note plus its regression test capture a hard-won debugging result.

Two things I would like resolved before merge, plus a set of suggestions. Inline comments carry the details.

Blocking-ish

1. Cuda_CudnnPagedBuildsWhenPageTableOverAllocated hard-asserts the cuDNN backend with only a GetCudaArchitecture() < 800 guard. On a CUDA build without cuDNN flash attention, on cuDNN < 9.5, or on an sm_8x device where the paged planner declines, this is a red test rather than a skip. Its sibling Cuda_CudnnPagedDispatchWhenEnabled already has the correct shape.
2. is_supported_paged() admits dprops.major >= 8, but a planner rejection inside build_paged_graph() is an ORT_THROW, not a fallback. Unlike the dense run() path, PagedAttention always has FA2/MEA available for the same shape, so an unsupported (arch, shape) under explicit opt-in kills the user's inference instead of degrading.

Suggestions

3. PagedGraphParams params; leaves 7 padding bytes indeterminate and BytesHash hashes all of them — a miss means the graph is rebuilt every decode step, silently erasing the benefit of this PR.
4. LaunchGetSeqlensKVDecode hardcodes past_seqlens[i] + 1 rather than deriving the query count from cumulative_seqlens_q.
5. build_paged_graph() calls cudaGetDeviceProperties on every cache miss purely to build a string used only in the two throw branches.
6. No CUDA-graph-capture guard on the first-call graph build, in an op that has capture in its contract (Cuda_FlashSplitKvCudaGraphReplay) and where the XQA path does guard with isCapturing().
7. docs/contrib_ops/cuda/paged_attention.md is the normative description of the backend cascade, the sdpa_kernel bits and the env-var matrix; it should gain the new tier, the short-context regression and the ORT_ENABLE_CUDNN_FLASH_ATTENTION=0 opt-out.

Test coverage gap

Nice that RunIoBindingCase checks against a CPU reference, so these are real numerical tests. But every case where cuDNN actually fires uses batch_size == 1 — precisely the configuration where the two batch-indexed pieces of this change cannot fail (the page-table leading stride max_num_blocks_per_seq, and the per-batch seqlens_kv padding mask). A batch >= 2 GQA case with unequal past_seqlens would exercise both. is_bf16 in run_paged also has no coverage; one c.bf16_query = true case closes that.

Scope

moe_kernels.cu, moe_quantization.cc, QgemmU8X8KernelAvx2.asm, manifest_parser.cc and the deep_gemm/WIN32 cmake edits belong to #32485 — noted that you plan to rebase. The --diag-suppress=20303 change is a CUTLASS/sm_60 arch-list interaction unrelated to this feature; splitting it out would let it be reverted independently when sm_60 leaves the arch list.

Also worth verifying before merge: that the graph cache actually hits across decode steps (count build_paged_graph entries over a ~100-token generation), and the behaviour on an sm_80/sm_86 device with ORT_ENABLE_CUDNN_FLASH_ATTENTION=1.

[Copilot]

🟡 Changes recommended

The node-wide readiness probe can disagree with the shape- and thread-specific graph cache, and one regression test can silently skip its target failure.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds a cuDNN paged-SDPA decode tier to CUDA PagedAttention, with dispatch gating, graph caching, tests, documentation, and benchmarks.

Changes:

• Implements cuDNN paged-attention graph construction and execution.
• Integrates capability-based dispatch and fallback handling.
• Adds CUDA tests and benchmark coverage.

File summaries

File	Description
onnxruntime/contrib_ops/cuda/bert/attention_data.h	Adds cuDNN execution state.
onnxruntime/contrib_ops/cuda/bert/paged_attention.h	Adds cuDNN dispatch configuration and probe state.
onnxruntime/contrib_ops/cuda/bert/paged_attention.cc	Integrates eligibility, probing, and dispatch.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h	Declares KV-length generation.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu	Implements cuDNN decode execution.
onnxruntime/contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.h	Exposes paged-SDPA APIs.
onnxruntime/contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.cc	Builds, caches, and executes cuDNN graphs.
onnxruntime/test/contrib_ops/paged_attention_op_test.cc	Adds dispatch and correctness tests.
onnxruntime/test/onnx/microbenchmark/paged_attention.cc	Adds CUDA benchmark variants.
docs/contrib_ops/cuda/paged_attention.md	Documents the new dispatch tier.
cmake/onnxruntime_unittests.cmake	Includes the benchmark for CUDA builds.

Review details

• Files reviewed: 11/11 changed files
• Comments generated: 3
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Planner failures are repeatedly rebuilt, XQA failure can bypass cuDNN, and capture behavior lacks coverage.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

onnxruntime/contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.cc:836

• Planner failures are not actually cached here: a null graph is returned without inserting this key. Because ComputeInternal probes on every decode step, any statically eligible shape rejected by cuDNN will repeat the expensive validate()/build() attempt for every generated token. Cache the negative result as intended, and make run_paged return false when it finds a cached null entry instead of dereferencing it.

  auto mha_graph = build_paged_graph(params);
  if (mha_graph == nullptr) {
    return false;

• Files reviewed: 11/11 changed files
• Comments generated: 3
• Review effort level: Balanced

[Tianlei Wu (tianleiwu)]

Second review round on head 63c119b5. All eight threads from the previous round are addressed and resolved — the SafeInt narrowing, the FillPagedGraphParams memset, the major >= 9 static gate, the non-throwing build_paged_graph + try_build_paged_graph probe, the capture guard, the cumulative_seqlens_q-derived KV mask, the cudnn_allocator rename, and dropping the unrelated cmake suppression. Thanks for the thorough turnaround.

Two new findings this round (inline), plus notes below.

What I re-verified as correct

• Q / K / V / O / page-table dims and strides match ORT's physical [num_blocks, block_size, kv_num_heads, head_size] cache layout and the packed decode Q/O layout.
• head_sink is covered transitively by !parameters.use_smooth_softmax (the helper sets use_smooth_softmax = head_sink != nullptr), and v_head_size != head_size is only reachable in LATENT mode, which !use_latent_attention excludes.
• CreateConstantSeqLenBuffer uses a stream-ordered Fill, so the synthesized seq_len_q buffer is capture-safe.
• needs_dense_kv cannot go stale when the late XQA-fallback retry flips use_cudnn_paged, because mea_eligible requires !flash_eligible and the retry requires fp16_xqa_eligible.
• num_heads_kv == 0 short-circuits before the modulo in is_supported_paged.

Minor, not worth a thread

• paged_attention.cc:756 — the "no backend available" return now also guards on !use_cudnn_paged, but the message still lists only FlashAttention / MemoryEfficientAttention / paged decode and the three ORT_DISABLE_* env vars. Worth mentioning the cuDNN paged tier and ORT_ENABLE_CUDNN_FLASH_ATTENTION so a user whose only viable backend was cuDNN (and whose planner probe failed) gets an actionable message.
• The PR description still documents a cmake/onnxruntime_providers_cuda.cmake #20303-D suppression that was removed from the branch. Please drop that bullet before merge.

CI coverage
is_supported_paged now requires dprops.major >= 9, and every new test that actually exercises the kernel either gates on GetCudaArchitecture() >= 900 or skips when the SdpaKernel=CUDNN_FLASH_ATTENTION marker is absent. On pre-Hopper CI agents only the negative gate tests (RespectsGlobalCudnnDisable, FallsBackWithoutMetadata, FallsBackForMultiTokenBound, FallsBackForQuantizedCache) will execute. That is the right conservative choice, but it does mean the kernel path itself has no automated coverage in CI today — worth calling out in the design doc so a future regression is understood to be caught only by manual sm_90+ runs.

[Tianlei Wu (tianleiwu)]

Will create a follow up PR