[CUDA] Add cuDNN paged SDPA decode tier for PagedAttention

cmake/onnxruntime_unittests.cmake

@@ -1566,7 +1566,7 @@ if (NOT onnxruntime_ENABLE_TRAINING_TORCH_INTEROP)
       ${BENCHMARK_DIR}/quantize.cc
       ${BENCHMARK_DIR}/reduceminmax.cc
       ${BENCHMARK_DIR}/layer_normalization.cc
-      $<$<BOOL:${onnxruntime_USE_WEBGPU}>:${BENCHMARK_DIR}/paged_attention.cc>)
+      $<$<OR:$<BOOL:${onnxruntime_USE_WEBGPU}>,$<BOOL:${onnxruntime_USE_CUDA}>>:${BENCHMARK_DIR}/paged_attention.cc>)
     target_include_directories(onnxruntime_benchmark PRIVATE ${ONNXRUNTIME_ROOT} ${onnxruntime_graph_header} ${ONNXRUNTIME_ROOT}/core/mlas/inc)
     target_compile_definitions(onnxruntime_benchmark PRIVATE BENCHMARK_STATIC_DEFINE)
     target_compile_definitions(onnxruntime_benchmark PRIVATE ${mlas_private_compile_definitions})

docs/contrib_ops/cuda/paged_attention.md

@@ -1260,6 +1260,7 @@ Target dispatch order once all phases land, first eligible wins:
 |---|---|---|
 | 0 | **MLA backend** (new, §12.7) | `kv_cache_layout == "LATENT"` — FlashMLA / FlashInfer MLA / TRT-LLM MLA where available, unfused MLA reference otherwise |
 | 1 | **Paged decode kernel** (new, Phase 4) | Decode-shaped batch per the static shape test in §4.7 (`token_count <= batch_size`); non-quantized or `PER_TENSOR`/`PER_CHANNEL` INT8/FP8; sliding window and `head_sink` supported |
+| 1.5 | **cuDNN paged SDPA** (new) | Decode-shaped (`max_query_len_bound == 1` **and** `token_count == batch_size`); FP16/BF16; non-quantized cache; `is_causal`; no softcap / sliding window / smooth softmax (QK-Norm rides the shared prologue — §7); `sm >= 90` and cuDNN `>= 9.5`; `is_supported_paged` shape gate accepts `(num_heads, kv_num_heads, head_size, block_size)`. Selected ahead of FlashAttention/MEA when eligible and XQA is not the winner. Reports `SdpaKernel=CUDNN_FLASH_ATTENTION`. Kill switch: `ORT_ENABLE_CUDNN_FLASH_ATTENTION=0` and the `sdpa_kernel` bit both disable it (they share the state with the CUDA EP's other cuDNN attention paths). Also gated by a one-shot per-node build probe: if `try_build_paged_graph` fails (planner rejection for that node's shape signature), the node falls back to FlashAttention/MEA for the rest of its lifetime instead of throwing. During CUDA graph capture the probe is skipped and cuDNN paged is suppressed for that capture — run at least one warm-up Compute before capture. **Short-context regression**: the planner requires `max_seq_len_kv == max_num_blocks_per_seq * block_size`, so `run_paged` aligns the metadata bound up to a page multiple when the caller over-allocates the page table (e.g. `past+1` reported but reserving an extra page). |
 | 2 | **FlashAttention varlen** | FP16/BF16, SM80+, non-quantized cache (or quantized via dequant-gather); supports sliding window, softcap, packed QKV, `head_sink` via LSE epilogue |
 | 3 | **Memory-Efficient Attention (CUTLASS fMHA)** | Fallback for supported combinations and pre-SM80 |
 | 4 | **Unfused** | Last resort — arbitrary `head_size`, `attention_bias`, `output_qk` |

onnxruntime/contrib_ops/cuda/bert/attention_data.h

@@ -311,6 +311,17 @@ struct PagedAttentionData {
   // use_paged_decode when set.
   bool use_xqa_decode = false;
   bool use_xqa_spec_dec = false;
+  // cuDNN paged SDPA decode kernel. Opt-in on H100+ for one-token-per-sequence decode when the
+  // cache is unquantized and none of the fused options (softcap / head sink / sliding window /
+  // bias) are requested. Takes precedence over use_flash_attention when set.
+  bool use_cudnn_paged = false;
+
+  // cuDNN paged SDPA path: temp-space allocator, cuDNN handle (stored as void* to avoid pulling the
+  // cuDNN headers into this file; cast to cudnnHandle_t in the .cu runner) and per-batch KV length
+  // scratch (int32, [batch_size]) filled from past_seqlens on device before dispatch.
+  AllocatorPtr cudnn_allocator = nullptr;
+  void* cudnn_handle = nullptr;
+  int* cudnn_seqlens_kv = nullptr;
 };
 
 }  // namespace cuda

onnxruntime/contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.h

@@ -47,4 +47,83 @@ void run(
     Stream* stream,
     AllocatorPtr allocator);
 
+// Paged-KV variant of the SDPA above. The K/V "cache" tensors point into a shared block pool with
+// physical layout [cache_num_blocks, block_size, num_heads_kv, head_size]; a per-batch page_table
+// selects which blocks belong to which sequence. Same masking surface as run() with these first-cut
+// restrictions (matching cuDNN 9.5+ paged SDPA support):
+//   * s_q == 1 (decode, one query token per batch entry). Q is BNSH_1: [batch_size, num_heads_q, 1,
+//     head_size_qk].
+//   * No attention bias, no softcap, no sliding window, no head sink.
+//   * Causal masking is a no-op for s_q == 1 (the token attends to all keys up to its own
+//     position, which the padding mask bounds), so this entry point never asks for it.
+//   * mask_sequence_lengths_kv is required: [batch_size] int32 device buffer of per-sequence KV
+//     lengths.
+// Shape gate only; is_supported_paged does not probe the cuDNN planner. The paged planner is
+// validated on H100+ (sm>=90), so this shape gate refuses sm_8x uniformly -- callers that opt in
+// on sm_8x fall through to another PagedAttention backend rather than reaching a graph->build()
+// planner-rejection path with no in-op fallback.
+bool is_supported_paged(const cudaDeviceProp& dprops,
+                        int num_heads_q,
+                        int num_heads_kv,
+                        int head_size_qk,
+                        int head_size_v,
+                        int sequence_length_q,       // must be 1
+                        int max_sequence_length_kv,  // upper bound for graph build
+                        int block_size);
+
+// Pre-dispatch buildability probe. Returns true iff a cuDNN paged graph for these shape/type
+// parameters is present in the thread-local cache OR was just built successfully. On planner
+// rejection it returns false without throwing, so the PagedAttention cascade can clear
+// use_cudnn_paged for the current Run and fall back to FlashAttention / MemoryEfficientAttention.
+// The lookup key is byte-identical to the one run_paged uses, so the caller must invoke this
+// every Run rather than caching the answer in a node-scalar: buildability is a per-(thread,
+// shape) property (the cache is thread_local and keyed on the full PagedGraphParams).
+//
+// Capture-safe: on a cache miss with `stream` capturing the returned graph, this returns false
+// rather than attempting a non-capturable build. A well-behaved producer therefore warms the
+// cache with at least one non-capturing Compute per (thread, shape) before capture.
+//
+// The graph is compiled at max_seq_len_kv == max_num_blocks_per_seq * block_size (the natural
+// page-table capacity) because cuDNN 9.12's planner rejects any smaller value; per-sequence
+// lengths in mask_sequence_lengths_kv still bound actual attention range, so this coarsening is
+// semantically safe.
+bool try_build_paged_graph(
+    int batch_size,
+    int num_heads_q,
+    int num_heads_kv,
+    int head_size_qk,
+    int head_size_v,
+    int cache_num_blocks,
+    int block_size,
+    int max_num_blocks_per_seq,
+    float scale,
+    bool is_bf16,
+    cudnnHandle_t handle,
+    Stream* stream);
+
+// Executes the paged SDPA. Returns true on success, false if the cuDNN graph is unavailable
+// (planner rejection cached from a prior try_build_paged_graph, cache miss on a capturing stream,
+// or a build attempted here that the planner rejected). On false the output is not written, so the
+// caller can propagate a Status without leaving partial state.
+bool run_paged(
+    void* output,                   // [batch_size, num_heads_q, 1, head_size_v]
+    void* q,                        // [batch_size, num_heads_q, 1, head_size_qk]
+    void* k_cache,                  // [cache_num_blocks, block_size, num_heads_kv, head_size_qk]
+    void* v_cache,                  // [cache_num_blocks, block_size, num_heads_kv, head_size_v]
+    int* block_table,               // [batch_size, max_num_blocks_per_seq] int32
+    int* mask_sequence_lengths_kv,  // [batch_size] int32 per-sequence KV length
+    int batch_size,
+    int num_heads_q,
+    int num_heads_kv,
+    int head_size_qk,
+    int head_size_v,
+    int cache_num_blocks,  // leading dimension of the block pool
+    int block_size,
+    int max_num_blocks_per_seq,  // stride of block_table's leading dim
+    float scale,
+    bool is_bf16,  // True if bfloat16, otherwise float16
+    cudnnHandle_t handle,
+    Stream* stream,
+    AllocatorPtr allocator);
+
 }  // namespace onnxruntime::cudnn_sdpa

onnxruntime/contrib_ops/cuda/bert/paged_attention.cc

@@ -15,6 +15,7 @@
 #include "contrib_ops/cpu/bert/paged_attention_helper.h"
 #include "contrib_ops/cuda/bert/flash_attention/flash_api.h"
 #include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h"
+#include "contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.h"
 #include "contrib_ops/cuda/bert/xqa/xqa_paged_loader.h"
 #include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h"
 
@@ -147,6 +148,10 @@ PagedAttention<T, TCACHE>::PagedAttention(const OpKernelInfo& info)
   enable_xqa_ = sizeof(T) == 2 && (ParseEnvironmentVariableWithDefault<int>("ORT_ENABLE_XQA", 1) != 0);
   enable_native_xqa_ =
       enable_xqa_ && (ParseEnvironmentVariableWithDefault<int>("ORT_ENABLE_XQA_NATIVE_KV", 0) != 0);
+  // cuDNN paged SDPA follows the same opt-in / auto-on pattern as GroupQueryAttention's cuDNN tier.
+  constexpr bool kIsFp16OrBf16 = std::is_same_v<T, MLFloat16> || std::is_same_v<T, BFloat16>;
+  enable_cudnn_paged_ = kIsFp16OrBf16 && kernel_options_->UseCudnnFlashAttention();
+  auto_enable_cudnn_paged_ = kIsFp16OrBf16 && kernel_options_->AllowCudnnFlashAttentionAuto();
 }
 
 template <typename T, typename TCACHE>
@@ -475,13 +480,79 @@ Status PagedAttention<T, TCACHE>::ComputeInternal(OpKernelContext* context) cons
       ((quantized_xqa_eligible && std::is_same<T, MLFloat16>::value) || native_spec_xqa_eligible) &&
       parameters.head_size == 256 && group_size == 6 &&
       max_query_len_bound > 1 && max_query_len_bound <= 8;
+
+  // cuDNN paged SDPA (decode-only, unquantized cache). Preferred over FlashAttention when eligible;
+  // XQA still wins its target case (fp16, group_size 6, head_size 256, native page size). The
+  // eligibility is intentionally metadata-gated so the selection never triggers a new D->H readback
+  // -- max_query_len_bound == 1 with has_metadata_bounds is the same signal XQA uses. The
+  // xqa_spec_dec_candidate path is orthogonal because it requires max_query_len_bound > 1.
+  // Enablement mirrors GroupQueryAttention's cuDNN tier: explicit opt-in via the sdpa_kernel bit /
+  // ORT_ENABLE_CUDNN_FLASH_ATTENTION=1, plus auto-on for sm>=90. The env var set to 0 kills every
+  // cuDNN attention path uniformly.
+  const bool cudnn_paged_enabled =
+      enable_cudnn_paged_ || (auto_enable_cudnn_paged_ && device_prop.major >= 9);
+  const bool cudnn_paged_eligible =
+      cudnn_paged_enabled &&
+      has_metadata_bounds &&
+      max_query_len_bound == 1 &&
+      parameters.token_count == parameters.batch_size &&
+      !use_latent_attention &&
+      !kIsQuantizedCache &&
+      parameters.is_causal &&
+      parameters.softcap == 0.0f &&
+      parameters.local_window_size <= 0 &&
+      !parameters.use_smooth_softmax &&
+      onnxruntime::cudnn_sdpa::is_stable() &&
+      onnxruntime::cudnn_sdpa::is_supported_paged(
+          device_prop,
+          parameters.num_heads, parameters.kv_num_heads,
+          parameters.head_size, parameters.head_size,
+          /*sequence_length_q=*/1,
+          /*max_sequence_length_kv=*/max_kv_len_bound,
+          parameters.block_size);
+  bool use_cudnn_paged = cudnn_paged_eligible && !fp16_xqa_eligible;
+
+  // Pre-dispatch buildability probe for cuDNN paged. is_supported_paged only checks static shapes;
+  // the planner may still reject the compiled graph, and we cannot fall back once dispatch has
+  // committed. The graph cache in run_paged is thread_local and keyed on the full PagedGraphParams
+  // (shape + handle), so a node-wide scalar latch would be wrong on two axes: (a) after a
+  // successful probe for one shape, a later Compute with a different batch_size /
+  // max_num_blocks_per_seq / num_blocks would skip the check and hit an unbuildable dispatch, and
+  // (b) a second worker thread has its own empty thread_local cache. So probe the *current* shape
+  // and handle every Run using the same key run_paged will use. try_build_paged_graph is a
+  // cache-first read (returns true on hit without touching cuDNN) and folds isCapturing()
+  // internally: on a cache miss during graph capture it returns false rather than attempting a
+  // non-capturable build. A false result clears use_cudnn_paged for this Run only, so the cascade
+  // drops to FlashAttention / MemoryEfficientAttention.
+  if (use_cudnn_paged) {
+    const float probe_scale = parameters.scale == 0.0f
+                                  ? 1.f / std::sqrt(static_cast<float>(parameters.head_size))
+                                  : parameters.scale;
+    const bool ok = onnxruntime::cudnn_sdpa::try_build_paged_graph(
+        parameters.batch_size,
+        parameters.num_heads, parameters.kv_num_heads,
+        parameters.head_size, parameters.head_size,
+        parameters.num_blocks,
+        parameters.block_size,
+        parameters.max_num_blocks_per_seq,
+        probe_scale,
+        std::is_same<T, BFloat16>::value,
+        GetCudnnHandle(context),
+        ort_stream.get());
+    if (!ok) {
+      use_cudnn_paged = false;
+    }
+  }
+
   // Only the FlashAttention backend takes a causality flag; the paged decode and CUTLASS kernels
   // both hard-code a bottom-right causal mask.
   bool use_paged_decode =
+      !use_cudnn_paged &&
       decode_eligible && parameters.is_causal &&
       ((decode_shaped && (kIsQuantizedCache || fp16_xqa_eligible || !flash_eligible)) || xqa_spec_dec_candidate);
-  bool use_flash_attention = flash_eligible && !use_paged_decode;
-  const bool use_memory_efficient_attention = mea_eligible && !use_paged_decode && parameters.is_causal;
+  bool use_flash_attention = flash_eligible && !use_paged_decode && !use_cudnn_paged;
+  const bool use_memory_efficient_attention =
+      mea_eligible && !use_paged_decode && !use_cudnn_paged && parameters.is_causal;
 
   if (!parameters.is_causal && !use_flash_attention) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
@@ -604,19 +675,23 @@ Status PagedAttention<T, TCACHE>::ComputeInternal(OpKernelContext* context) cons
   // Native-cache XQA promotion is speculative until the one-token-per-sequence and shared-memory
   // checks pass. Restore Flash for ragged decode steps and unsupported devices instead of leaving
   // them on the portable scalar paged-decode fallback. This is safe without dense KV staging
-  // because the native FP16/BF16 cache is already a Flash-supported dtype.
-  if (!use_xqa_decode && (fp16_xqa_eligible || native_spec_xqa_eligible) && !kIsQuantizedCache &&
-      flash_eligible) {
+  // because the native FP16/BF16 cache is already a Flash-supported dtype. Skip when cuDNN paged
+  // already won this step; the two selections would otherwise coexist for the corner case where
+  // native XQA is opted in on a bf16 cache with head_size=256 / group_size=6 and max_query_len=1
+  // (native_spec_xqa_eligible is a shape gate and stays true regardless of query length).
+  if (!use_xqa_decode && !use_cudnn_paged && (fp16_xqa_eligible || native_spec_xqa_eligible) &&
+      !kIsQuantizedCache && flash_eligible) {
     use_paged_decode = false;
     use_flash_attention = true;
   }
   DUMP_STRING("Backend = ", use_latent_attention  ? "latent"
                             : use_xqa_decode      ? "paged decode (XQA)"
+                            : use_cudnn_paged     ? "cuDNN paged SDPA"
                             : use_paged_decode    ? "paged decode"
                             : use_flash_attention ? "flash attention"
                                                   : "memory efficient attention");
 
-  if (!use_latent_attention && !use_paged_decode && !use_flash_attention && !use_memory_efficient_attention) {
+  if (!use_latent_attention && !use_paged_decode && !use_flash_attention && !use_memory_efficient_attention && !use_cudnn_paged) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                            "PagedAttention requires FlashAttention (sm>=80, fp16/bf16, block_size a multiple of ",
                            flash_min_block_size, " for head_size ", parameters.head_size,
@@ -761,6 +836,14 @@ Status PagedAttention<T, TCACHE>::ComputeInternal(OpKernelContext* context) cons
   }
 #endif
 
+  // cuDNN paged SDPA needs a small [batch_size] int32 buffer holding per-sequence KV lengths,
+  // filled on device by LaunchGetSeqlensKVDecode right before dispatch.
+  IAllocatorUniquePtr<void> cudnn_seqlens_kv_buffer;
+  if (use_cudnn_paged) {
+    cudnn_seqlens_kv_buffer = GetScratchBuffer<void>(
+        sizeof(int) * static_cast<size_t>(parameters.batch_size), GetComputeStream(context));
+  }
+
   // Print debug info
   if (kernel_options_->AllowDebugInfo()) {
     AttentionKernelDebugInfo debug_info;
@@ -769,6 +852,7 @@ Status PagedAttention<T, TCACHE>::ComputeInternal(OpKernelContext* context) cons
     debug_info.use_flash_attention = use_flash_attention;
     debug_info.use_efficient_attention = use_memory_efficient_attention;
     debug_info.use_decoder_attention = use_paged_decode && !use_xqa_decode;
+    debug_info.use_cudnn_flash_attention = use_cudnn_paged;
     if (use_flash_attention) {
       debug_info.num_splits = std::max(1, flash_num_splits);
     } else if (use_paged_decode && !use_xqa_decode) {
@@ -812,6 +896,7 @@ Status PagedAttention<T, TCACHE>::ComputeInternal(OpKernelContext* context) cons
   data.use_paged_decode = use_paged_decode;
   data.use_xqa_decode = use_xqa_decode;
   data.use_xqa_spec_dec = use_xqa_spec_dec;
+  data.use_cudnn_paged = use_cudnn_paged;
   if (softmax_lse_buffer != nullptr) {
     // FlashAttention always writes fp32 log-sum-exp, independent of T.
     data.softmax_lse = reinterpret_cast<float*>(softmax_lse_buffer.get());
@@ -851,6 +936,12 @@ Status PagedAttention<T, TCACHE>::ComputeInternal(OpKernelContext* context) cons
     data.fmha_buffer = reinterpret_cast<CudaT*>(fmha_buffer.get());
   }
 
+  if (use_cudnn_paged) {
+    ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&data.cudnn_allocator));
+    data.cudnn_handle = static_cast<void*>(GetCudnnHandle(context));
+    data.cudnn_seqlens_kv = reinterpret_cast<int*>(cudnn_seqlens_kv_buffer.get());
+  }
+
   cublasHandle_t cublas = GetCublasHandle(context);
 
   return QkvToContext<CudaT, CudaTCache>(

onnxruntime/contrib_ops/cuda/bert/paged_attention.h

@@ -54,10 +54,22 @@ class PagedAttention final : public CudaKernel {
   bool enable_xqa_;
   // Native FP16/BF16 cache specializations are opt-in because FlashAttention is competitive.
   bool enable_native_xqa_;
+  // cuDNN paged SDPA (decode-only tier). Mirrors GroupQueryAttention: the standard sdpa_kernel bit
+  // and ORT_ENABLE_CUDNN_FLASH_ATTENTION opt users in explicitly, and sm>=90 gets it automatically
+  // via AllowCudnnFlashAttentionAuto(). ORT_ENABLE_CUDNN_FLASH_ATTENTION=0 is the shared kill switch
+  // for every cuDNN attention path in the CUDA EP.
+  bool enable_cudnn_paged_;
+  bool auto_enable_cudnn_paged_;
   // -1 = not yet resolved, 0 = the kernel needs more shared memory than this device allows,
   // 1 = it fits. Resolved once per node because it only depends on head_size / group size.
   mutable std::atomic<int> xqa_shared_memory_ok_{-1};
   mutable std::atomic<int> xqa_spec_dec_shared_memory_ok_{-1};
+  // No node-scalar latch for cuDNN paged: the graph cache in cudnn_flash_attention.cc is
+  // thread_local and keyed on the full PagedGraphParams (shape + handle), so buildability is a
+  // per-(thread, shape) property, not a per-node one. ComputeInternal calls try_build_paged_graph
+  // every Run with the current shape; that call is a cache-first read (~10 ns steady state) and
+  // folds the capture-mode check internally, so a shape change or a second worker thread cannot
+  // stale a global latch into a hard dispatch failure.
   const AttentionKernelOptions* kernel_options_;
 };