Upgrade cudnn_frontend to 1.24 and enable cuDNN SDPA for MHA/GQA

cmake/deps.txt

@@ -54,7 +54,7 @@ tensorboard;https://github.com/tensorflow/tensorboard/archive/373eb09e4c5d2b3cc2
 cutlass;https://github.com/NVIDIA/cutlass/archive/refs/tags/v4.4.2.zip;4b0bae4428b84370407c0a71778b13dc2eee5be1
 extensions;https://github.com/microsoft/onnxruntime-extensions/archive/c24b7bab0c12f53da76d0c31b03b9f0f8ec8f3b4.zip;239063aee4946a9af147b473a4c3da78ba7413b4
 directx_headers;https://github.com/microsoft/DirectX-Headers/archive/refs/tags/v1.613.1.zip;47653509a3371eabb156360f42faf582f314bf2e
-cudnn_frontend;https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.12.0.zip;7e733cfdc410d777b76122d64232499205589a96
+cudnn_frontend;https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.24.0.zip;a55a1980bf5c57692d66ae7bc3b39798f5535e1f
 dawn;https://github.com/google/dawn/archive/ec7b457e5bb1fcec6f59733c4f3dd84d2f885a38.zip;d4d64d1729104b61e654073566f1a376e16cad92
 kleidiai;https://github.com/ARM-software/kleidiai/archive/refs/tags/v1.20.0.tar.gz;6895e72b3d5cf1173358164cb3d64c9d7d33cc84
 # kleidiai-qmx is pinned to a specific commit as there are no tagged releases. When an appropriate tagged release becomes available,

cmake/external/cudnn_frontend.cmake

@@ -13,3 +13,14 @@ set(CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS OFF CACHE BOOL "" FORCE)
 set(CUDNN_PATH ${onnxruntime_CUDNN_HOME})
 
 onnxruntime_fetchcontent_makeavailable(cudnn_frontend)
+
+# Treat the cudnn_frontend headers as system headers so that warnings originating from them (e.g.
+# unused static helper functions introduced in newer releases) are not promoted to errors by
+# ONNX Runtime's -Werror flags. This applies to every consumer that links the cudnn_frontend target.
+if(TARGET cudnn_frontend)
+  get_target_property(_cudnn_frontend_include_dirs cudnn_frontend INTERFACE_INCLUDE_DIRECTORIES)
+  if(_cudnn_frontend_include_dirs)
+    set_target_properties(cudnn_frontend PROPERTIES
+      INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "${_cudnn_frontend_include_dirs}")
+  endif()
+endif()

cmake/onnxruntime_providers_cuda_plugin.cmake

@@ -385,6 +385,8 @@ target_link_libraries(onnxruntime_providers_cuda_plugin PRIVATE
     CUDA::cublas
     CUDA::cublasLt
     CUDA::cufft
+    CUDA::nvrtc
+    CUDA::cuda_driver
     CUDNN::cudnn_all
     cudnn_frontend
     Boost::mp11

onnxruntime/contrib_ops/cuda/bert/attention_data.h

@@ -197,6 +197,8 @@ struct GroupQueryAttentionData {
   bool use_memory_efficient_attention = false;
   bool use_flash_attention_fast_decode = false;
   bool use_xqa = false;
+  // cuDNN SDPA (cudnn_frontend) path: preferred on SM>=90 for non-quantized FP16/BF16 GQA.
+  bool use_cudnn_sdpa = false;
   // GQA-capable unfused fallback (issue #28195): used when Flash/MEA/XQA are all ineligible,
   // e.g. fp16 head_size > 256 with past_key, or GQA on old GPUs without MEA/Flash support.
   bool use_unfused = false;
@@ -213,6 +215,11 @@ struct GroupQueryAttentionData {
   T* unfused_q_bnsh = nullptr;
   T* unfused_y_bnsh = nullptr;
   void* unfused_workspace = nullptr;
+
+  // cuDNN SDPA path: temp-space allocator and cuDNN handle (stored as void* to avoid pulling the
+  // cuDNN headers into this file; cast to cudnnHandle_t in the .cu runner).
+  AllocatorPtr allocator = nullptr;
+  void* cudnn_handle = nullptr;
 };
 
 template <typename T>

onnxruntime/contrib_ops/cuda/bert/attention_impl.cu

@@ -426,7 +426,7 @@ Status CudnnFlashAttention(
          data.qkv_format == AttentionQkvFormat::Q_K_V_BNSH);
   assert(parameters.mask_type == AttentionMaskType::MASK_NONE ||
          parameters.mask_type == AttentionMaskType::MASK_1D_KEY_SEQ_LEN);
-  constexpr bool is_bf16 = false;
+  constexpr bool is_bf16 = std::is_same<T, BFloat16>::value;
 
   T* attention_bias = const_cast<T*>(data.attention_bias);
   int* mask_sequence_lengths_kv = const_cast<int*>(data.mask_index);

onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.cc

@@ -15,6 +15,8 @@ using namespace onnxruntime::contrib::attention;
 
 namespace onnxruntime {
 void AttentionKernelOptions::Initialize(int value, bool use_build_flag, bool check_cudnn_version) {
+  has_explicit_kernel_selection_ = (value > 0);
+  disable_auto_cudnn_flash_attention_ = false;
   if (value > 0) {
     use_flash_attention_ = (value & static_cast<int>(AttentionBackend::FLASH_ATTENTION)) > 0;
 #if USE_LEAN_ATTENTION
@@ -37,7 +39,13 @@ void AttentionKernelOptions::Initialize(int value, bool use_build_flag, bool che
 #endif
     use_efficient_attention_ = !ParseEnvironmentVariableWithDefault<bool>(kDisableMemoryEfficientAttention, false);
     use_trt_fused_attention_ = !ParseEnvironmentVariableWithDefault<bool>(kDisableFusedSelfAttention, false);
-    use_cudnn_flash_attention_ = ParseEnvironmentVariableWithDefault<bool>(kEnableCudnnFlashAttention, false);
+    const auto cudnn_flash_attention_env = ParseEnvironmentVariable<bool>(kEnableCudnnFlashAttention);
+    if (cudnn_flash_attention_env.has_value()) {
+      use_cudnn_flash_attention_ = *cudnn_flash_attention_env;
+      disable_auto_cudnn_flash_attention_ = !use_cudnn_flash_attention_;
+    } else {
+      use_cudnn_flash_attention_ = false;
+    }
 
     use_unfused_ = true;
     use_trt_flash_attention_ = !ParseEnvironmentVariableWithDefault<bool>(kDisableTrtFlashAttention, false);

onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.h

@@ -36,6 +36,17 @@ class AttentionKernelOptions {
   bool UseTrtCausalAttention() const { return use_trt_causal_attention_; }
   bool UseDecoderAttention() const { return use_decoder_attention_; }
 
+  // True when the SDPA kernel was explicitly selected via the sdpa_kernel provider option
+  // (a positive bitmask). When false, the kernel selection follows defaults / environment
+  // variables, which allows operators to auto-prefer cuDNN SDPA on SM>=90.
+  bool HasExplicitKernelSelection() const { return has_explicit_kernel_selection_; }
+
+  // True when operators may auto-prefer cuDNN SDPA on SM>=90. This is disabled when
+  // the user explicitly pins kernels via sdpa_kernel or sets ORT_ENABLE_CUDNN_FLASH_ATTENTION=0.
+  bool AllowCudnnFlashAttentionAuto() const {
+    return !has_explicit_kernel_selection_ && !disable_auto_cudnn_flash_attention_;
+  }
+
   bool AllowDebugInfo() const { return enable_kernel_debug_info_; }
 
   int MinSeqLenForFlashAttentionPackedQkv() const { return min_seq_len_for_flash_attention_packed_qkv_; }
@@ -61,6 +72,9 @@ class AttentionKernelOptions {
 
   bool use_decoder_attention_{true};
 
+  bool has_explicit_kernel_selection_{false};
+  bool disable_auto_cudnn_flash_attention_{false};
+
   bool enable_kernel_debug_info_{false};
 
   int min_seq_len_for_flash_attention_packed_qkv_{0};

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

@@ -59,6 +59,7 @@ void run(
 
 #include <cudnn_frontend.h>
 #include "core/providers/cuda/shared_inc/cudnn_fe_call.h"
+#include "core/providers/cuda/shared_inc/cuda_utils.h"
 #include "core/providers/cuda/cuda_stream_handle.h"
 
 namespace onnxruntime::cudnn_sdpa {
@@ -82,12 +83,21 @@ bool is_stable() {
   //    Sliding window attention
   // version 90300 (9.3.0)
   //    Bug fixes; Variable sequence length supports zero-sequence-length values
+  // version 90600 (9.6.0)
+  //    Bottom-right causal mask no longer requires sequence lengths to be multiples of 64
   // For more information, please refer to cuDNN release notes, and the following links:
   //    https://docs.nvidia.com/deeplearning/cudnn/latest/developer/graph-api.html#fused-flash-attention-fprop
-  //    https://github.com/NVIDIA/cudnn-frontend/blob/v1.5.2/docs/operations/Attention.md
+  //    https://github.com/NVIDIA/cudnn-frontend/blob/v1.24.0/docs/operations/Attention.md
 
   // For cuDNN version < 9.3, we will disable it by default.
-  return cudnnGetVersion() >= 90300;
+  const size_t version = cudnnGetVersion();
+
+  // cuDNN 9.10.0 and 9.10.1 have known bugs in the FP16/BF16 SDPA forward kernels, so skip them.
+  if (version == 91000 || version == 91001) {
+    return false;
+  }
+
+  return version >= 90300;
 }
 
 namespace fe = cudnn_frontend;
@@ -100,14 +110,36 @@ bool is_supported(const cudaDeviceProp& dprops,
                   int sequence_length_q,
                   int sequence_length_kv,
                   bool is_causal) {
-  return (dprops.major >= 8) &&
-         (head_size_qk % 8 == 0) && (head_size_qk <= 256) &&
-         (head_size_v % 8 == 0) && (head_size_v <= 256) &&
-         (num_heads_q % num_heads_kv == 0) &&
-         // Bottom right causal mask is only supported with s_q multiple of 64 and s_kv multiple of 64.
-         (!is_causal || (sequence_length_q != sequence_length_kv &&
-                         sequence_length_q % 64 == 0 &&
-                         sequence_length_kv % 64 == 0));
+  if (dprops.major < 8 ||
+      (head_size_qk % 8 != 0) || (head_size_qk > 256) ||
+      (head_size_v % 8 != 0) || (head_size_v > 256) ||
+      (num_heads_kv == 0) || (num_heads_q % num_heads_kv != 0)) {
+    return false;
+  }
+
+  // For a single query token (s_q == 1, e.g. decode) causal masking is a no-op: the token attends to
+  // every key up to its own position, which the padding / kv sequence length already bounds. cuDNN is
+  // therefore called without a causal mask in that case (see run()), so the causal-specific support
+  // restrictions below do not apply.
+  if (is_causal && sequence_length_q > 1) {
+    // cuDNN expresses causal masking through diagonal alignment (cudnn_frontend >= 1.24):
+    //   * s_q == s_kv : top-left aligned mask (standard self-attention causal, e.g. prefill).
+    //   * s_q  < s_kv : bottom-right aligned mask (decode / cross attention with past KV).
+    //   * s_q  > s_kv : not supported by cuDNN bottom-right causal masking.
+    if (sequence_length_q > sequence_length_kv) {
+      return false;
+    }
+
+    // Bottom-right causal masking requires s_q and s_kv to be multiples of 64 on cuDNN < 9.6.0.
+    // Top-left causal masking (s_q == s_kv) has no such restriction.
+    if (sequence_length_q != sequence_length_kv &&
+        cudnnGetVersion() < 90600 &&
+        (sequence_length_q % 64 != 0 || sequence_length_kv % 64 != 0)) {
+      return false;
+    }
+  }
+
+  return true;
 }
 
 // A helper function to set stride for q, k, v or output tensor.
@@ -228,13 +260,23 @@ std::shared_ptr<fe::graph::Graph> build_graph(GraphParams& params) {
 
   auto attributes = fe::graph::SDPA_attributes()
                         .set_name("SDPA")
-                        .set_is_inference(true)
-                        .set_causal_mask(is_causal)
-                        .set_causal_mask_bottom_right(is_causal && sequence_length_q != sequence_length_kv)
+                        .set_generate_stats(false)
                         .set_attn_scale(scale);
 
+  if (is_causal) {
+    // Use diagonal-alignment based causal masking (cudnn_frontend >= 1.24). A right bound of 0 keeps
+    // only the lower-triangular region. Standard self-attention (s_q == s_kv) is top-left aligned;
+    // decode / cross attention (s_q < s_kv) is bottom-right aligned so the query rows line up with the
+    // most recent keys.
+    attributes.set_diagonal_alignment(sequence_length_q != sequence_length_kv
+                                          ? fe::DiagonalAlignment_t::BOTTOM_RIGHT
+                                          : fe::DiagonalAlignment_t::TOP_LEFT)
+        .set_diagonal_band_right_bound(0);
+  }
+
   if (params.sliding_window > 0) {
-    attributes.set_sliding_window_length(params.sliding_window);
+    // Sliding window length maps to the left bound of the attention diagonal band.
+    attributes.set_diagonal_band_left_bound(params.sliding_window);
   }
 
   if (params.has_bias) {
@@ -315,6 +357,22 @@ struct BytesHash {
 // TODO(tianleiwu): since the key includes sequence lengths, we may want to limit the cache size.
 thread_local std::unordered_map<GraphParams, std::shared_ptr<fe::graph::Graph>, BytesHash<GraphParams> > mha_graph_cache;
 
+// Allocate a device buffer of shape [batch_size] filled with a constant sequence length.
+// Used to synthesize a no-op padding mask for one side when cuDNN requires both seq_len_q and
+// seq_len_kv to be set (cudnn_frontend validates that both are present when padding mask is on).
+// The buffer is filled with a stream-ordered Fill kernel (rather than a synchronous cudaMemcpy)
+// so this path is safe to capture into a CUDA graph.
+static IAllocatorUniquePtr<int> CreateConstantSeqLenBuffer(AllocatorPtr allocator,
+                                                           Stream* stream,
+                                                           int batch_size,
+                                                           int value) {
+  IAllocatorUniquePtr<int> buffer =
+      IAllocator::MakeUniquePtr<int>(allocator, static_cast<size_t>(batch_size), false, stream);
+  cudaStream_t cuda_stream = stream ? static_cast<cudaStream_t>(stream->GetHandle()) : nullptr;
+  onnxruntime::cuda::Fill<int>(cuda_stream, buffer.get(), value, static_cast<int64_t>(batch_size));
+  return buffer;
+}
+
 void run(
     void* output,
     void* q,
@@ -340,6 +398,22 @@ void run(
     cudnnHandle_t handle,
     Stream* stream,
     AllocatorPtr allocator) {
+  // cuDNN requires both seq_len_q and seq_len_kv to be present when a padding mask is used. When the
+  // caller provides only one side, synthesize the other as the full (unpadded) sequence length so it
+  // behaves as a no-op padding mask on that side.
+  IAllocatorUniquePtr<int> synthesized_seq_len_q;
+  IAllocatorUniquePtr<int> synthesized_seq_len_kv;
+  if (mask_sequence_lengths_q != nullptr || mask_sequence_lengths_kv != nullptr) {
+    if (mask_sequence_lengths_q == nullptr) {
+      synthesized_seq_len_q = CreateConstantSeqLenBuffer(allocator, stream, batch_size, sequence_length_q);
+      mask_sequence_lengths_q = synthesized_seq_len_q.get();
+    }
+    if (mask_sequence_lengths_kv == nullptr) {
+      synthesized_seq_len_kv = CreateConstantSeqLenBuffer(allocator, stream, batch_size, sequence_length_kv);
+      mask_sequence_lengths_kv = synthesized_seq_len_kv.get();
+    }
+  }
+
   GraphParams params;
   params.batch_size = batch_size;
   params.num_heads_q = num_heads_q;
@@ -349,7 +423,11 @@ void run(
   params.sequence_length_q = sequence_length_q;
   params.sequence_length_kv = sequence_length_kv;
   params.scale = scale;
-  params.is_causal = is_causal;
+  // A single query token (s_q == 1, e.g. decode) attends to all keys up to its own position, so causal
+  // masking is a no-op and the padding / kv sequence length bounds the valid keys. Dropping the causal
+  // mask here also avoids a cuDNN limitation where decode-only graphs (s_q == 1) with a causal
+  // right-bound fail to build on cuDNN backend versions <= 9.9.0.
+  params.is_causal = is_causal && (sequence_length_q > 1);
   params.is_bf16 = is_bf16;
   params.qkv_format = qkv_format;
   params.handle = handle;

onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc

@@ -11,6 +11,7 @@
 #include "contrib_ops/cuda/bert/group_query_attention_impl.h"
 #include "contrib_ops/cuda/bert/group_query_attention.h"
 #include "contrib_ops/cpu/bert/group_query_attention_helper.h"
+#include "contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.h"
 #include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h"
 #include "contrib_ops/cuda/bert/flash_attention/flash_api.h"
 #include "contrib_ops/cuda/bert/xqa/xqa_loader.h"
@@ -119,6 +120,11 @@ GroupQueryAttention<T, U>::GroupQueryAttention(const OpKernelInfo& info)
   // Memory efficient attention supports float and float16. BFloat16 support added for SM80+.
   disable_memory_efficient_attention_ = !kernel_options_->UseEfficientAttention();
 
+  // cuDNN SDPA (cudnn_frontend) supports FP16 and BF16 and is auto-preferred on SM>=90.
+  constexpr bool kIsFp16OrBf16 = std::is_same<T, MLFloat16>::value || std::is_same<T, BFloat16>::value;
+  enable_cudnn_flash_attention_ = kIsFp16OrBf16 && kernel_options_->UseCudnnFlashAttention();
+  auto_enable_cudnn_flash_attention_ = kIsFp16OrBf16 && kernel_options_->AllowCudnnFlashAttentionAuto();
+
   if (!disable_flash_attention_) {
     zeros_ = this->GetScratchBuffer<int>(kZerosCount, nullptr);
     CUDA_CALL_THROW(cudaMemset(zeros_.get(), 0, kZerosCount * sizeof(int)));
@@ -392,11 +398,34 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
     }
   }
 
-  // Compute past_present_share_buffer early since it's needed for flash attention path selection.
-  // This compares the final pointer values after quantization handling.
+  // === cuDNN SDPA eligibility (preferred on SM>=90, Hopper/Blackwell) ===
+  // Constrained to the well-supported causal path: non-quantized FP16/BF16 KV cache, no softcap,
+  // no smooth-softmax / head sink, and no sliding window. Rotary and packed QKV are handled by
+  // PrepareQKV before the kernel runs; cuDNN handles grouped-query attention natively.
+  bool use_cudnn_sdpa = !data.use_xqa &&
+                        !is_inputs_quantized &&
+                        std::is_same<T, U>::value &&
+                        parameters.softcap == 0.0f &&
+                        !parameters.use_smooth_softmax &&
+                        head_sink == nullptr &&
+                        parameters.local_window_size == -1 &&
+                        parameters.past_kv_format == AttentionQkvFormat::Q_K_V_BNSH &&
+                        (enable_cudnn_flash_attention_ ||
+                         (auto_enable_cudnn_flash_attention_ && device_prop.major >= 9)) &&
+                        onnxruntime::cudnn_sdpa::is_stable() &&
+                        onnxruntime::cudnn_sdpa::is_supported(device_prop,
+                                                              parameters.num_heads,
+                                                              parameters.kv_num_heads,
+                                                              parameters.head_size,
+                                                              parameters.head_size,
+                                                              parameters.sequence_length,          // seq_len_q
+                                                              parameters.seqlen_present_kv_cache,  // seq_len_kv (capacity)
+                                                              /*is_causal=*/true);
+  data.use_cudnn_sdpa = use_cudnn_sdpa;
 
 #if USE_FLASH_ATTENTION
   bool use_flash_attention = !data.use_xqa &&
+                             !data.use_cudnn_sdpa &&
                              !disable_flash_attention_ &&
                              onnxruntime::flash::is_supported<T>(device_prop,
                                                                  parameters.head_size,
@@ -475,7 +504,7 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
   }
 
 #if USE_MEMORY_EFFICIENT_ATTENTION
-  if (!data.use_xqa && !data.use_flash_attention) {
+  if (!data.use_xqa && !data.use_cudnn_sdpa && !data.use_flash_attention) {
     // Fall back to memory efficient attention.
     int sm = (device_prop.major * 10) + device_prop.minor;
     bool use_memory_efficient_attention =
@@ -524,7 +553,7 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
   // See LaunchUnfusedAttention in contrib_ops/cuda/bert/unfused_attention.h.
   // ---------------------------------------------------------------------
   IAllocatorUniquePtr<void> unfused_scratch;
-  if (!data.use_xqa && !data.use_flash_attention && !data.use_memory_efficient_attention &&
+  if (!data.use_xqa && !data.use_cudnn_sdpa && !data.use_flash_attention && !data.use_memory_efficient_attention &&
       !is_inputs_quantized && !parameters.use_smooth_softmax && head_sink == nullptr &&
       parameters.past_kv_format == AttentionQkvFormat::Q_K_V_BNSH) {
     data.use_unfused = true;
@@ -562,6 +591,7 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
     AttentionKernelDebugInfo debug_info;
     debug_info.use_flash_attention = data.use_flash_attention;
     debug_info.use_efficient_attention = data.use_memory_efficient_attention;
+    debug_info.use_cudnn_flash_attention = data.use_cudnn_sdpa;
 
     debug_info.Print("GroupQueryAttention",
                      this->Node().Name(),
@@ -601,6 +631,11 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
 
   cublasHandle_t cublas = GetCublasHandle(context);
 
+  if (data.use_cudnn_sdpa) {
+    ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&data.allocator));
+    data.cudnn_handle = static_cast<void*>(GetCudnnHandle(context));
+  }
+
   ORT_RETURN_IF_ERROR((QkvToContext<CudaT, CudaU>(
       device_prop, cublas, ort_stream.get(), parameters, data)));
   return Status::OK();