diff --git a/cmake/deps.txt b/cmake/deps.txt index 5a835e0c6fe9f..d8cd7713bfe65 100644 --- a/cmake/deps.txt +++ b/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, diff --git a/cmake/external/cudnn_frontend.cmake b/cmake/external/cudnn_frontend.cmake index d89ab0f669f35..9af915a272b07 100644 --- a/cmake/external/cudnn_frontend.cmake +++ b/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() diff --git a/cmake/onnxruntime_providers_cuda_plugin.cmake b/cmake/onnxruntime_providers_cuda_plugin.cmake index 32856eda42400..1db0c67b6b64d 100644 --- a/cmake/onnxruntime_providers_cuda_plugin.cmake +++ b/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 diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index 60f2d05446da1..7d5c9bc1e221e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/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 diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu index 22862f1ba7b4c..fd1a5b26fb049 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_impl.cu +++ b/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::value; T* attention_bias = const_cast(data.attention_bias); int* mask_sequence_lengths_kv = const_cast(data.mask_index); diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.cc b/onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.cc index 1e0c9cb8baffd..0723a48528dfe 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.cc +++ b/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(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(kDisableMemoryEfficientAttention, false); use_trt_fused_attention_ = !ParseEnvironmentVariableWithDefault(kDisableFusedSelfAttention, false); - use_cudnn_flash_attention_ = ParseEnvironmentVariableWithDefault(kEnableCudnnFlashAttention, false); + const auto cudnn_flash_attention_env = ParseEnvironmentVariable(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(kDisableTrtFlashAttention, false); diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.h b/onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.h index fd3b90387a235..d94eb5ec06e5c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.h +++ b/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}; diff --git a/onnxruntime/contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.cc b/onnxruntime/contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.cc index ba786931bb39a..76a556b23f7df 100644 --- a/onnxruntime/contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/cudnn_fmha/cudnn_flash_attention.cc @@ -59,6 +59,7 @@ void run( #include #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 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, BytesHash > 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 CreateConstantSeqLenBuffer(AllocatorPtr allocator, + Stream* stream, + int batch_size, + int value) { + IAllocatorUniquePtr buffer = + IAllocator::MakeUniquePtr(allocator, static_cast(batch_size), false, stream); + cudaStream_t cuda_stream = stream ? static_cast(stream->GetHandle()) : nullptr; + onnxruntime::cuda::Fill(cuda_stream, buffer.get(), value, static_cast(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 synthesized_seq_len_q; + IAllocatorUniquePtr 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; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 44408e6ce4af9..3caa5d4ad5270 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/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::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::value || std::is_same::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(kZerosCount, nullptr); CUDA_CALL_THROW(cudaMemset(zeros_.get(), 0, kZerosCount * sizeof(int))); @@ -392,11 +398,34 @@ Status GroupQueryAttention::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::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(device_prop, parameters.head_size, @@ -475,7 +504,7 @@ Status GroupQueryAttention::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::ComputeInternal(OpKernelContext* context) cons // See LaunchUnfusedAttention in contrib_ops/cuda/bert/unfused_attention.h. // --------------------------------------------------------------------- IAllocatorUniquePtr 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::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::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(GetCudnnHandle(context)); + } + ORT_RETURN_IF_ERROR((QkvToContext( device_prop, cublas, ort_stream.get(), parameters, data))); return Status::OK(); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h index 75a613724e746..34847983ad7de 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h @@ -36,6 +36,8 @@ class GroupQueryAttention final : public CudaKernel { bool disable_memory_efficient_attention_; bool disable_flash_decode_; bool enable_xqa_; + bool enable_cudnn_flash_attention_; // cuDNN SDPA explicitly enabled (env / sdpa_kernel) + bool auto_enable_cudnn_flash_attention_; // auto-prefer cuDNN SDPA on SM>=90 when no explicit kernel pinned KVQuantizationType k_quant_type_; KVQuantizationType v_quant_type_; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 4b365cb304a43..1c39de01fef66 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -37,6 +37,7 @@ limitations under the License. #include "contrib_ops/cuda/bert/attention_softmax.h" #include "contrib_ops/cuda/bert/bert_padding.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/flash_attention/flash_api.h" #include "contrib_ops/cuda/bert/unfused_attention.h" #include "contrib_ops/cuda/bert/group_query_attention_impl.h" @@ -1183,6 +1184,81 @@ Status UnfusedGqaAttention( ////////// API Functions +// ============================================================================ +// CudnnSdpaAttention: cuDNN scaled-dot-product-attention (cudnn_frontend) path. +// +// Preferred on SM>=90 (Hopper/Blackwell) for non-quantized FP16/BF16 GQA. cuDNN handles +// grouped-query attention natively (no KV head expansion) and applies causal masking through the +// diagonal-band API. The present KV cache (BNSH, padded to seqlen_present_kv_cache) is passed with +// the capacity as the KV sequence length so strides match the physical buffer; a per-batch padding +// mask (total_seq_lens) ignores the unused tail. RoPE / packed-QKV unpack and the new-token append +// are performed by PrepareQKV. +// +// Not handled here (excluded by op-level eligibility, caller falls through elsewhere): +// - Quantized KV cache (U != T), softcap, smooth softmax / head sink, sliding window. +// ============================================================================ +template +Status CudnnSdpaAttention( + const cudaDeviceProp& device_prop, + cudaStream_t stream, + Stream* ort_stream, + GroupQueryAttentionParameters& parameters, + GroupQueryAttentionData& data, + float scale) { + if constexpr (!std::is_same::value) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "cuDNN SDPA GQA path requires a non-quantized KV cache (T == U)."); + } else { + ORT_GQA_TRACE("CudnnSdpaAttention"); + + const int max_threads_per_block = device_prop.maxThreadsPerBlock; + + // Prepare Q (optional RoPE) and append the new K/V into the present cache (BNSH, padded to + // seqlen_present_kv_cache). After this, present_key/value hold the full KV history. + const T* q_prep = nullptr; + ORT_RETURN_IF_ERROR((PrepareQKV(stream, max_threads_per_block, parameters, data, q_prep))); + + ORT_RETURN_IF(data.cudnn_handle == nullptr || data.allocator == nullptr, + "cuDNN SDPA GQA path is missing the cuDNN handle or temp-space allocator."); + + constexpr bool is_bf16 = std::is_same::value; + + // First prompt may right-pad the query, so its valid query length equals total_seq_lens. + // Otherwise every one of the sequence_length query tokens is valid and the wrapper synthesizes a + // full-length (no-op) query padding mask. + int* seq_len_q = parameters.is_first_prompt ? data.total_seq_lens : nullptr; + int* seq_len_kv = data.total_seq_lens; + + ::onnxruntime::cudnn_sdpa::run( + reinterpret_cast(data.output), + const_cast(reinterpret_cast(q_prep)), + reinterpret_cast(data.present_key), + reinterpret_cast(data.present_value), + /*attn_bias=*/nullptr, + seq_len_q, + seq_len_kv, + parameters.batch_size, + parameters.num_heads, // num_heads_q + parameters.kv_num_heads, // num_heads_kv + parameters.head_size, // head_size_qk + parameters.head_size, // head_size_v (GQA: v_head_size == head_size) + parameters.sequence_length, // sequence_length_q + parameters.seqlen_present_kv_cache, // sequence_length_kv (capacity, matches buffer strides) + scale, + /*is_causal=*/true, + is_bf16, + /*broadcast_attn_bias_dim_0=*/false, + /*broadcast_attn_bias_dim_1=*/false, + /*sliding_window=*/0, + AttentionQkvFormat::Q_K_V_BSNH_BNSH_BNSH, + static_cast(data.cudnn_handle), + ort_stream, + data.allocator); + + return Status::OK(); + } +} + template Status QkvToContext( const cudaDeviceProp& device_prop, @@ -1196,6 +1272,10 @@ Status QkvToContext( return ExtremeDecoding(device_prop, stream, parameters, data, scale); } + if (data.use_cudnn_sdpa) { + return CudnnSdpaAttention(device_prop, stream, ort_stream, parameters, data, scale); + } + #if USE_FLASH_ATTENTION if (data.use_flash_attention_fast_decode) { return FlashDecoding(device_prop, stream, parameters, data, scale); diff --git a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc index a2af4831a3a00..fe4c4fc784c4a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.cc @@ -60,6 +60,7 @@ MultiHeadAttention::MultiHeadAttention(const OpKernelInfo& info) kernel_options_ = this->GetAttentionKernelOptions(); constexpr bool kIsFp16 = std::is_same::value; + constexpr bool kIsFp16OrBf16 = kIsFp16 || std::is_same::value; disable_fused_self_attention_ = !kIsFp16 || !kernel_options_->UseTrtFusedAttention(); enable_trt_flash_attention_ = kIsFp16 && kernel_options_->UseTrtFlashAttention(); @@ -74,7 +75,13 @@ MultiHeadAttention::MultiHeadAttention(const OpKernelInfo& info) disable_fused_cross_attention_ = !kIsFp16 || !kernel_options_->UseTrtCrossAttention(); - enable_cudnn_flash_attention_ = kIsFp16 && kernel_options_->UseCudnnFlashAttention(); + // cuDNN SDPA (cudnn_frontend) supports both FP16 and BF16. + enable_cudnn_flash_attention_ = kIsFp16OrBf16 && kernel_options_->UseCudnnFlashAttention(); + + // On SM>=90 (Hopper/Blackwell) cuDNN SDPA is generally the fastest backend, so it is auto-preferred + // ahead of flash / cutlass attention unless the user explicitly pinned a different SDPA kernel via + // the sdpa_kernel provider option. The SM check itself is done at compute time. + auto_enable_cudnn_flash_attention_ = kIsFp16OrBf16 && kernel_options_->AllowCudnnFlashAttentionAuto(); disable_decoder_attention_ = !kernel_options_->UseDecoderAttention(); @@ -298,8 +305,35 @@ Status MultiHeadAttention::ComputeInternal(OpKernelContext* context) cons constexpr bool use_lean_attention = false; #endif + // === cuDNN SDPA eligibility (computed before flash so it can take priority on SM>=90) === + // cuDNN SDPA only supports no mask, or a 1D key sequence length (padding) mask. + bool is_mask_none_or_1d_k_len = parameters.mask_type == AttentionMaskType::MASK_NONE || + parameters.mask_type == AttentionMaskType::MASK_1D_KEY_SEQ_LEN; + // cuDNN SDPA is enabled when explicitly requested, or auto-preferred on SM>=90 (unless the user + // pinned a different SDPA kernel through the sdpa_kernel provider option). + bool cudnn_sdpa_enabled = enable_cudnn_flash_attention_ || + (auto_enable_cudnn_flash_attention_ && sm >= 90); + // Bottom-right causal masking (used by cuDNN when s_q != s_kv) does not support attention bias. + bool cudnn_sdpa_bias_ok = attention_bias == nullptr || + !is_unidirectional_ || + parameters.sequence_length == parameters.total_sequence_length; + bool cudnn_sdpa_supported = cudnn_sdpa_enabled && + cudnn_sdpa_bias_ok && + is_mask_none_or_1d_k_len && + onnxruntime::cudnn_sdpa::is_stable() && + onnxruntime::cudnn_sdpa::is_supported(device_prop, + parameters.num_heads, // num_heads_q + parameters.num_heads, // num_heads_kv + parameters.head_size, // head_size_qk + parameters.v_head_size, // head_size_v + parameters.sequence_length, // seq_len_q + parameters.total_sequence_length, // seq_len_kv + is_unidirectional_); + #if USE_FLASH_ATTENTION + // On SM>=90 (Hopper/Blackwell) prefer cuDNN SDPA ahead of flash attention. bool use_flash_attention = kernel_type == AttentionKernelType::AttentionKernel_Default && + !(cudnn_sdpa_supported && sm >= 90) && !disable_flash_attention_ && nullptr == attention_bias && nullptr == key_padding_mask && @@ -348,19 +382,8 @@ Status MultiHeadAttention::ComputeInternal(OpKernelContext* context) cons } #endif - bool is_mask_none_or_1d_k_len = parameters.mask_type == AttentionMaskType::MASK_NONE || - parameters.mask_type == AttentionMaskType::MASK_1D_KEY_SEQ_LEN; bool use_cudnn_sdpa = kernel_type == AttentionKernelType::AttentionKernel_Default && - enable_cudnn_flash_attention_ && - is_mask_none_or_1d_k_len && - onnxruntime::cudnn_sdpa::is_supported(device_prop, - parameters.num_heads, // num_heads_q - parameters.num_heads, // num_heads_kv - parameters.head_size, // head_size_qk - parameters.v_head_size, // head_size_v - parameters.sequence_length, // seq_len_q - parameters.total_sequence_length, // seq_len_kv - is_unidirectional_); + cudnn_sdpa_supported; DUMP_STRING("Use cuDNN SDPA = ", (use_cudnn_sdpa == true)); if (use_cudnn_sdpa) { kernel_type = AttentionKernelType::AttentionKernel_CudnnFlashAttention; diff --git a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h index 7de27c24fb601..67f88aada9132 100644 --- a/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/multihead_attention.h @@ -37,6 +37,7 @@ class MultiHeadAttention final : public CudaKernel { #endif bool disable_memory_efficient_attention_; bool enable_cudnn_flash_attention_; + bool auto_enable_cudnn_flash_attention_; // auto-prefer cuDNN SDPA on SM>=90 when no explicit kernel pinned bool disable_decoder_attention_; // These mutable members are readonly after they are initialized so that they can be shared among multiple threads. diff --git a/onnxruntime/core/providers/cuda/cuda_utils.cu b/onnxruntime/core/providers/cuda/cuda_utils.cu index f7b0bc2b91b2b..a4d1eb6b30d9c 100644 --- a/onnxruntime/core/providers/cuda/cuda_utils.cu +++ b/onnxruntime/core/providers/cuda/cuda_utils.cu @@ -33,6 +33,7 @@ void Fill(cudaStream_t stream, T* output, T value, int64_t count) { _Fill <<>>(output, value, N); } + template class ConstantBufferImpl : public IConstantBuffer { public: