diff --git a/docs/annotated_partitioning/attention_workspace_estimation.md b/docs/annotated_partitioning/attention_workspace_estimation.md index d7120d989f616..04899befead52 100644 --- a/docs/annotated_partitioning/attention_workspace_estimation.md +++ b/docs/annotated_partitioning/attention_workspace_estimation.md @@ -207,6 +207,62 @@ alias fact and the past-cache capacity, and the preparation recipe includes the checked preservation region. Windowed execution, XQA, and Flash fast decode require both K/V pairs to alias and reject this state. +### GroupQueryAttention bounded estimator and framework adapters + +The GQA framework follow-up translates positional inputs and immutable kernel +attributes into a graph-free bounded domain. It covers packed and separate +Q/K/V, optional past-cache pairs, cache quantization, prompt and decode states, +windowed staging and compaction, XQA (including unresolved shared-memory +fallback), regular and fast-decode Flash, MEA, unfused, and cuDNN reachability. +Head geometry derived from `WorkspaceInputShape` is an upper bound: route +reachability searches supported smaller head sizes instead of testing only the +componentwise maximum. + +Flash split selection is not monotonic in KV length. The bounded estimator uses +the runtime block size (`256` for head size at most 64, `128` through 128, and +`64` above 128) and bounds the selected split count by: + +```text +min(128, SM count, ceil(KV length bound / block size)) +``` + +It then sizes the split accumulators at that envelope. Complete routes are +mutually exclusive and are aggregated with `max`, not sum. The resulting +nonzero estimate is one operator-owned slot-0 root with 256-byte alignment. +Level 1 reports it as `runtime_workspace_bytes`; Level 2 declares one root from +the same estimator. Level 1 conservatively includes dynamic head-sink conversion +because prepack state is unavailable during capability analysis. For a constant +head sink, it separately charges the possible session-lived FP32 copy and the +initialization-only device staging copy, even when session configuration later +disables prepacking. Level 2 can omit the transient conversion region when the +constructed kernel has a prepacked head sink, so its root can be smaller. +Neither adapter changes runtime +`GetScratchBuffer()` calls or allocation topology. + +The CPU `total_sequence_length` scalar and past/present aliasing are not +available through `WorkspaceInputShape`. For non-windowed GQA, the scalar can +exceed the past-cache sequence dimension and directly scale Flash, MEA, and +unfused workspace. Non-windowed execution can also preserve one full past +tensor when exactly one past/present pair aliases. The current adapter therefore +reports non-windowed GQA as unavailable rather than treating the past shape as +a total-KV bound or omitting alias-preservation scratch. + +Successful estimates are currently limited to sliding-window cache nodes. +Their final cache capacity remains `C`, while a multi-token step uses a +transient staged/effective attention extent of `C + S`; a single-token step +uses `C`. This shape-only bound remains sound even when the scalar +`total_sequence_length` is much larger. Runtime requires both past/present +pairs to alias before staging or compaction, excluding the partial-alias +preservation path. Sliding-window attention bias remains +unsupported by runtime and is also unavailable to the estimator. If cuDNN can +be reached anywhere in a graph-free bounded domain, aggregation remains +unavailable because cuDNN's allocator-based workspace has no sound graph-free +oracle. Level 1 then uses the generic fallback and Level 2 emits no requirement. + +Unlike the log-only PA/PMHA Level-1 probe, a successful GQA Level-1 estimate is +passed to the #31962 resource accountant as `runtime_workspace_bytes` and can +affect CUDA partition acceptance. + MHA and GQA are high-value coverage targets and have high estimation-drift risk. Their runtime behavior can include dynamic internal backend dispatch, cache lifecycle and aliasing, optional inputs, non-monotonic fallback paths, and unfused workspace governed by `S_q * S_kv_total`. GQA additionally has different Q and KV head counts. MHA can have diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index bfb2b97a65112..b14d8509a15dd 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -13,6 +13,9 @@ #include "contrib_ops/cuda/bert/group_query_attention_impl.h" #include "contrib_ops/cuda/bert/group_query_attention.h" #include "contrib_ops/cuda/bert/group_query_attention_workspace.h" +#if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) +#include "contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h" +#endif #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" @@ -201,6 +204,51 @@ Status GroupQueryAttention::PrePack(const Tensor& tensor, int input_idx, A return Status::OK(); } +#if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) +template +Status GroupQueryAttention::DeclareWorkspaceRequirements( + gsl::span input_shapes, + InlinedVector& requirements) const { + requirements.clear(); + GQAWorkspaceEstimateConfig config; + config.qkv_element_size = sizeof(T); + config.cache_element_size = sizeof(U); + config.num_heads = num_heads_; + config.kv_num_heads = kv_num_heads_; + config.causal = is_unidirectional_ ? 1 : 0; + config.local_window_size = local_window_size_; + config.sliding_window_cache = sliding_window_cache_; + config.do_rotary = do_rotary_; + config.smooth_softmax = use_smooth_softmax_; + config.softcap = softcap_; + config.k_quantization = + k_quant_type_ == KVQuantizationType::PER_TENSOR + ? GQAKvQuantizationType::PerTensor + : (k_quant_type_ == KVQuantizationType::PER_CHANNEL + ? GQAKvQuantizationType::PerChannel + : GQAKvQuantizationType::None); + config.v_quantization = + v_quant_type_ == KVQuantizationType::PER_TENSOR + ? GQAKvQuantizationType::PerTensor + : (v_quant_type_ == KVQuantizationType::PER_CHANNEL + ? GQAKvQuantizationType::PerChannel + : GQAKvQuantizationType::None); + config.kv_cache_bit_width = kv_cache_bit_width_; + config.is_bf16 = std::is_same_v; + config.cache_is_fp8 = std::is_same_v; + config.enable_xqa = enable_xqa_; + config.disable_flash_decode = disable_flash_decode_; + config.head_sink_is_prepacked = xqa_head_sink_count_ == num_heads_; + + const auto estimate = EstimateGroupQueryAttentionWorkspace( + config, input_shapes, GetDeviceProp(), *kernel_options_); + if (estimate.has_value()) { + SetGroupQueryAttentionWorkspaceRequirements(*estimate, requirements); + } + return Status::OK(); +} +#endif + // ComputeInternal executes the GQA kernel. // // Inputs: diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h index 8cdcedc3726f5..acaba2593db44 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h @@ -25,6 +25,12 @@ class GroupQueryAttention final : public CudaKernel { Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, bool& is_packed, PrePackedWeights* prepacked_weights) override; +#if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) + Status DeclareWorkspaceRequirements( + gsl::span input_shapes, + InlinedVector& requirements) const override; +#endif + protected: int num_heads_; // number of attention heads int kv_num_heads_; // different for k and v for group query attention diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc new file mode 100644 index 0000000000000..66b74ec2f19df --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc @@ -0,0 +1,389 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h" + +#include +#include +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace { + +constexpr GQAWorkspaceStatus Invalid(const char* message) noexcept { + return {GQAWorkspaceError::InvalidArgument, message}; +} + +constexpr GQAWorkspaceStatus Unavailable(const char* message) noexcept { + return {GQAWorkspaceError::Unavailable, message}; +} + +GQAWorkspaceStatus Mul(size_t a, size_t b, size_t& out) noexcept { + return CheckedGQAWorkspaceMultiply(a, b, out); +} + +GQAWorkspaceStatus Add(size_t a, size_t b, size_t& out) noexcept { + return CheckedGQAWorkspaceAdd(a, b, out); +} + +GQAWorkspaceStatus Mul4(size_t a, size_t b, size_t c, size_t d, size_t& out) noexcept { + auto status = Mul(a, b, out); + if (!status.IsOK()) return status; + status = Mul(out, c, out); + if (!status.IsOK()) return status; + return Mul(out, d, out); +} + +GQAWorkspaceStatus Append(size_t bytes, size_t& cursor) noexcept { + if (bytes == 0) return {}; + size_t offset = 0; + auto status = CheckedGQAWorkspaceAlign(cursor, kGQAWorkspaceAlignment, offset); + if (!status.IsOK()) return status; + return Add(offset, bytes, cursor); +} + +GQAWorkspaceStatus Compose(size_t preparation, size_t backend, size_t& total) noexcept { + if (backend == 0) { + total = preparation; + return {}; + } + size_t offset = 0; + auto status = CheckedGQAWorkspaceAlign(preparation, kGQAWorkspaceAlignment, offset); + if (!status.IsOK()) return status; + return Add(offset, backend, total); +} + +GQAWorkspaceStatus EffectiveKvLengthBound(const GQAWorkspaceBounds& bounds, + int64_t sequence_length, + int64_t& effective_kv_length) noexcept { + effective_kv_length = bounds.present_kv_cache_capacity_bound; + if (!bounds.is_windowed_kv_cache || sequence_length == 1) return {}; + + size_t staged_capacity = 0; + auto status = Add(static_cast(bounds.present_kv_cache_capacity_bound), + static_cast(sequence_length), staged_capacity); + if (!status.IsOK()) return status; + if (staged_capacity > static_cast(std::numeric_limits::max())) { + return Invalid("Windowed GQA effective KV length bound must fit int32."); + } + + effective_kv_length = static_cast(staged_capacity); + return {}; +} + +GQAWorkspaceProblem MakeProblem(const GQAWorkspaceBounds& bounds, + int64_t sequence_length, + int64_t head_size, + bool first_prompt) noexcept { + GQAWorkspaceProblem problem; + problem.qkv_element_size = bounds.qkv_element_size; + problem.cache_element_size = bounds.cache_element_size; + problem.batch_size = bounds.batch_size_bound; + problem.sequence_length = sequence_length; + problem.num_heads = bounds.num_heads; + problem.kv_num_heads = bounds.kv_num_heads; + problem.head_size = head_size; + problem.present_kv_cache_capacity = bounds.present_kv_cache_capacity_bound; + problem.kv_cache_bit_width = bounds.kv_cache_bit_width; + problem.k_quantization = bounds.k_quantization; + problem.v_quantization = bounds.v_quantization; + problem.is_windowed_kv_cache = bounds.is_windowed_kv_cache; + problem.is_first_prompt = first_prompt; + problem.do_rotary = bounds.do_rotary; + problem.is_packed_qkv = bounds.is_packed_qkv; + problem.use_qk_norm = bounds.use_qk_norm; + return problem; +} + +int64_t LargestMultipleOfEight(int64_t bound, int64_t limit) noexcept { + return (std::min(bound, limit) / 8) * 8; +} + +GQAWorkspaceStatus FlashBackendEnvelope(const GQAWorkspaceProblem& problem, + const GQAWorkspaceBounds& bounds, + int64_t effective_kv_length, + bool fast_decode, + size_t& bytes) noexcept { + // flash::num_splits_heuristic never selects more than this. Unlike evaluating + // the heuristic at max KV length, this remains sound across its discontinuities. + const size_t head_size = static_cast(problem.head_size); + const size_t block_n = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64); + size_t kv_length = static_cast(effective_kv_length); + if (fast_decode && bounds.local_window_size > 0) { + kv_length = std::min(kv_length, static_cast(bounds.local_window_size)); + } + size_t numerator = 0; + auto status = Add(kv_length, block_n - 1, numerator); + if (!status.IsOK()) return status; + const size_t max_splits = std::min( + {size_t{128}, static_cast(bounds.multi_processor_count), numerator / block_n}); + + size_t lse_bytes = 0; + status = Mul4(static_cast(problem.batch_size), + static_cast(problem.num_heads), + static_cast(problem.sequence_length), sizeof(float), lse_bytes); + if (!status.IsOK()) return status; + bytes = 0; + status = Append(lse_bytes, bytes); + if (!status.IsOK() || max_splits <= 1) return status; + + size_t split_elements = 0; + status = Mul4(max_splits, static_cast(problem.batch_size), + static_cast(problem.num_heads), + static_cast(problem.sequence_length), split_elements); + if (!status.IsOK()) return status; + size_t split_lse = 0; + status = Mul(split_elements, sizeof(float), split_lse); + if (!status.IsOK()) return status; + status = Append(split_lse, bytes); + if (!status.IsOK()) return status; + size_t rounded_head_size = 0; + status = CheckedGQAWorkspaceAlign(head_size, 32, rounded_head_size); + if (!status.IsOK()) return status; + size_t output = 0; + status = Mul4(split_elements, rounded_head_size, sizeof(float), 1, output); + if (!status.IsOK()) return status; + return Append(output, bytes); +} + +GQAWorkspaceStatus XqaBackendEnvelope(const GQAWorkspaceProblem& problem, + const GQAWorkspaceBounds& bounds, + size_t& bytes) noexcept { + // Route eligibility is established by the graph adapter. This envelope + // intentionally bounds GetXQAScratchSize without reapplying recipe guards. + size_t sequence_count = 0; + auto status = Mul(static_cast(problem.batch_size), + static_cast(problem.kv_num_heads), sequence_count); + if (!status.IsOK()) return status; + const size_t subsequences = + std::max(sequence_count, static_cast(bounds.multi_processor_count)); + size_t semaphore = 0; + status = Mul(sequence_count, sizeof(uint32_t), semaphore); + if (!status.IsOK()) return status; + status = CheckedGQAWorkspaceAlign(semaphore, 128, bytes); + if (!status.IsOK()) return status; + size_t rows = 0; + status = Mul(128, subsequences, rows); + if (!status.IsOK()) return status; + status = Add(bytes, rows, bytes); + if (!status.IsOK()) return status; + // Alignment phase varies with sequence_count. Add the maximum possible + // padding instead of evaluating one endpoint. + status = Add(bytes, 127, bytes); + if (!status.IsOK()) return status; + status = Add(bytes, rows, bytes); + if (!status.IsOK()) return status; + + const int64_t group_size = problem.num_heads / problem.kv_num_heads; + const size_t tile = group_size <= 8 ? 8 : (group_size <= 16 ? 16 : 32); + size_t vector_bytes = 0; + status = Mul4(static_cast(problem.head_size), tile, 2, 1, vector_bytes); + if (!status.IsOK()) return status; + status = Add(bytes, vector_bytes - 1, bytes); + if (!status.IsOK()) return status; + size_t output = 0; + status = Mul(vector_bytes, subsequences, output); + if (!status.IsOK()) return status; + status = Add(bytes, output, bytes); + if (!status.IsOK()) return status; + + if (problem.do_rotary) { + size_t q = 0; + status = Mul4(static_cast(problem.batch_size), + static_cast(problem.num_heads), + static_cast(problem.head_size), + problem.qkv_element_size, q); + if (!status.IsOK()) return status; + status = CheckedGQAWorkspaceAlign(q, kGQAWorkspaceAlignment, q); + if (!status.IsOK()) return status; + status = Add(bytes, q, bytes); + if (!status.IsOK()) return status; + size_t k = 0; + status = Mul4(static_cast(problem.batch_size), + static_cast(problem.kv_num_heads), + static_cast(problem.head_size), + problem.qkv_element_size, k); + if (!status.IsOK()) return status; + status = CheckedGQAWorkspaceAlign(k, kGQAWorkspaceAlignment, k); + if (!status.IsOK()) return status; + status = Add(bytes, k, bytes); + if (!status.IsOK()) return status; + } + if (bounds.xqa_head_sink_storage == GQAXqaHeadSinkStorage::DynamicConversion) { + size_t sink = 0; + status = Mul(static_cast(problem.num_heads), sizeof(float), sink); + if (!status.IsOK()) return status; + status = CheckedGQAWorkspaceAlign(sink, kGQAWorkspaceAlignment, sink); + if (!status.IsOK()) return status; + status = Add(bytes, sink, bytes); + } + return status; +} + +} // namespace + +GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( + const GQAWorkspaceBounds& bounds) noexcept { + GQAWorkspaceAggregate aggregate; + if (bounds.qkv_element_size != 2 || + (bounds.cache_element_size != 1 && bounds.cache_element_size != 2) || + bounds.batch_size_bound <= 0 || bounds.sequence_length_bound <= 0 || + bounds.num_heads <= 0 || bounds.kv_num_heads <= 0 || + bounds.num_heads % bounds.kv_num_heads != 0 || + bounds.head_size_bound < 8 || bounds.present_kv_cache_capacity_bound <= 0 || + bounds.multi_processor_count <= 0 || + (!bounds.prompt_reachable && !bounds.decode_reachable)) { + aggregate.status = Invalid("GQA workspace bounds are incomplete or invalid."); + return aggregate; + } + if (HasGQAReachableBackend(bounds.reachable_backends, GQAReachableBackend::Cudnn)) { + aggregate.status = Unavailable( + "A reachable cuDNN GQA route has no graph-free workspace oracle."); + return aggregate; + } + const bool head_sink_is_or_may_be_prepacked = + bounds.head_sink_may_be_prepacked || + bounds.xqa_head_sink_storage == GQAXqaHeadSinkStorage::PrepackedFp32; + if (head_sink_is_or_may_be_prepacked) { + aggregate.status = Mul( + static_cast(bounds.num_heads), sizeof(float), + aggregate.persistent_prepack_bytes); + if (!aggregate.status.IsOK()) return aggregate; + aggregate.status = Mul( + static_cast(bounds.num_heads), bounds.qkv_element_size, + aggregate.initialization_scratch_bytes); + if (!aggregate.status.IsOK()) return aggregate; + } + + const auto retain = [&aggregate](GQAReachableBackend backend, + GQAWorkspaceStatus status, + size_t bytes) { + if (!aggregate.status.IsOK()) return; + if (!status.IsOK()) { + aggregate.status = status; + return; + } + aggregate.total_workspace_bytes = std::max(aggregate.total_workspace_bytes, bytes); + aggregate.sized_backends = aggregate.sized_backends | backend; + }; + + const std::array sequence_candidates{ + 1, bounds.sequence_length_bound}; + const auto size_complete = [&](GQABackend backend, int64_t head_size, + int64_t sequence_length, bool first_prompt) { + int64_t effective_kv_length = 0; + const auto effective_status = + EffectiveKvLengthBound(bounds, sequence_length, effective_kv_length); + if (!effective_status.IsOK()) { + GQACompleteWorkspaceResult result; + result.status = effective_status; + return result; + } + GQAConcreteRoute route; + route.backend = backend; + route.preparation.preprocess_mode = + backend == GQABackend::MemoryEfficient ? GQAPreprocessMode::MemoryEfficient + : GQAPreprocessMode::Unfused; + route.unfused.total_sequence_length = effective_kv_length; + return GetGQACompleteWorkspaceRecipe( + MakeProblem(bounds, sequence_length, head_size, first_prompt), route); + }; + + const int64_t general_head = LargestMultipleOfEight( + bounds.head_size_bound, std::numeric_limits::max()); + for (GQAReachableBackend backend : + {GQAReachableBackend::MemoryEfficient, GQAReachableBackend::Unfused}) { + if (!HasGQAReachableBackend(bounds.reachable_backends, backend)) continue; + const int64_t head = backend == GQAReachableBackend::MemoryEfficient + ? LargestMultipleOfEight(bounds.head_size_bound, 1024) + : general_head; + for (int64_t sequence : sequence_candidates) { + if (sequence == 1 && !bounds.decode_reachable && bounds.sequence_length_bound != 1) continue; + if (sequence > 1 && !bounds.prompt_reachable) continue; + for (bool first : {false, true}) { + if (first && !bounds.prompt_reachable) continue; + const auto result = size_complete( + backend == GQAReachableBackend::MemoryEfficient + ? GQABackend::MemoryEfficient + : GQABackend::Unfused, + head, sequence, first); + retain(backend, result.status, result.recipe.total_workspace_bytes); + } + } + } + + const auto size_flash = [&](bool fast, int64_t sequence, bool first) { + int64_t effective_kv_length = 0; + auto status = EffectiveKvLengthBound(bounds, sequence, effective_kv_length); + if (!status.IsOK()) { + return std::pair{status, 0}; + } + const int64_t head = LargestMultipleOfEight(bounds.head_size_bound, 256); + const auto problem = MakeProblem(bounds, sequence, head, first); + GQAPreparationRoute prep_route; + prep_route.preprocess_mode = GQAPreprocessMode::Flash; + prep_route.use_flash_attention_fast_decode = fast; + const auto prep = GetGQAPreparationRecipe(problem, prep_route); + if (!prep.status.IsOK()) return std::pair{prep.status, 0}; + size_t backend = 0; + status = FlashBackendEnvelope(problem, bounds, effective_kv_length, fast, backend); + size_t total = 0; + if (status.IsOK()) status = Compose(prep.recipe.total_preparation_bytes, backend, total); + return std::pair{status, total}; + }; + if (HasGQAReachableBackend(bounds.reachable_backends, GQAReachableBackend::Flash)) { + for (int64_t sequence : sequence_candidates) { + if (sequence == 1 && !bounds.decode_reachable && bounds.sequence_length_bound != 1) continue; + if (sequence > 1 && !bounds.prompt_reachable) continue; + for (bool first : {false, true}) { + if (first && !bounds.prompt_reachable) continue; + const auto result = size_flash(false, sequence, first); + retain(GQAReachableBackend::Flash, result.first, result.second); + } + } + } + if (HasGQAReachableBackend(bounds.reachable_backends, + GQAReachableBackend::FlashFastDecode)) { + // The current graph adapter rejects non-windowed inputs, so this route is + // retained for graph-free callers and a future non-windowed adapter. + // FlashBackendEnvelope reserves the maximum reachable split count instead + // of evaluating the non-monotonic split heuristic. Its storage and the + // preparation storage are nondecreasing in sequence length, so Smax covers + // every dynamic S in the bounded domain. + const auto result = size_flash(true, bounds.sequence_length_bound, false); + retain(GQAReachableBackend::FlashFastDecode, result.first, result.second); + } + + if (HasGQAReachableBackend(bounds.reachable_backends, GQAReachableBackend::Xqa)) { + bool found_head = false; + for (int64_t head : {int64_t{64}, int64_t{128}, int64_t{256}}) { + if (head > bounds.head_size_bound) continue; + found_head = true; + const auto problem = MakeProblem(bounds, 1, head, false); + GQAPreparationRoute route; + route.preprocess_mode = GQAPreprocessMode::Xqa; + const auto prep = GetGQAPreparationRecipe(problem, route); + size_t backend = 0; + auto status = prep.status; + if (status.IsOK()) status = XqaBackendEnvelope(problem, bounds, backend); + size_t total = 0; + if (status.IsOK()) status = Compose(prep.recipe.total_preparation_bytes, backend, total); + retain(GQAReachableBackend::Xqa, status, total); + } + if (!found_head) { + aggregate.status = Unavailable("No XQA head geometry is reachable in the supplied bounds."); + } + } + + if (aggregate.status.IsOK() && + aggregate.sized_backends == GQAReachableBackend::None) { + aggregate.status = Unavailable("No GQA backend route can be sized."); + } + return aggregate; +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h new file mode 100644 index 0000000000000..0783f90172a0c --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h @@ -0,0 +1,90 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "contrib_ops/cuda/bert/group_query_attention_workspace.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +enum class GQAReachableBackend : uint32_t { + None = 0, + Xqa = 1u << 0, + Flash = 1u << 1, + FlashFastDecode = 1u << 2, + MemoryEfficient = 1u << 3, + Unfused = 1u << 4, + Cudnn = 1u << 5, +}; + +constexpr GQAReachableBackend operator|(GQAReachableBackend left, + GQAReachableBackend right) noexcept { + return static_cast( + static_cast(left) | static_cast(right)); +} + +constexpr bool HasGQAReachableBackend(GQAReachableBackend mask, + GQAReachableBackend backend) noexcept { + return (static_cast(mask) & static_cast(backend)) != 0; +} + +// Graph-free upper bounds and immutable feature facts. Positive dimensions are +// componentwise bounds, not a claim that their combination is one runtime shape. +// The aggregate assumes both past/present K/V pairs alias. A caller that permits +// partial aliasing must separately account for preservation scratch. For a +// non-windowed cache, present_kv_cache_capacity_bound must bound the runtime +// present KV length, not only the past input tensor's sequence dimension. +struct GQAWorkspaceBounds { + size_t qkv_element_size = 0; + size_t cache_element_size = 0; + int64_t batch_size_bound = 0; + int64_t sequence_length_bound = 0; + int64_t num_heads = 0; + int64_t kv_num_heads = 0; + int64_t head_size_bound = 0; + int64_t present_kv_cache_capacity_bound = 0; + int64_t kv_cache_bit_width = 0; + GQAKvQuantizationType k_quantization = GQAKvQuantizationType::None; + GQAKvQuantizationType v_quantization = GQAKvQuantizationType::None; + bool is_windowed_kv_cache = false; + bool do_rotary = false; + bool is_packed_qkv = false; + bool use_qk_norm = false; + bool prompt_reachable = false; + bool decode_reachable = false; + + GQAReachableBackend reachable_backends = GQAReachableBackend::None; + int64_t device_major = 0; + int64_t device_minor = 0; + int64_t multi_processor_count = 0; + bool is_bf16 = false; + GQAXqaKvType xqa_kv_type = GQAXqaKvType::None; + // Exact state controls per-run XQA conversion scratch and implies prepack + // lifetime accounting. Possible state extends that charge to Level 1. + GQAXqaHeadSinkStorage xqa_head_sink_storage = GQAXqaHeadSinkStorage::None; + bool head_sink_may_be_prepacked = false; + int64_t local_window_size = -1; +}; + +struct GQAWorkspaceAggregate { + GQAWorkspaceStatus status; + size_t total_workspace_bytes = 0; + size_t persistent_prepack_bytes = 0; + size_t initialization_scratch_bytes = 0; + GQAReachableBackend sized_backends = GQAReachableBackend::None; +}; + +// Computes a sound maximum over all complete routes reachable anywhere in the +// bounded domain. cuDNN reachability makes the result unavailable because its +// allocator-based workspace has no graph-free oracle. +GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( + const GQAWorkspaceBounds& bounds) noexcept; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc new file mode 100644 index 0000000000000..e11b666169c86 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc @@ -0,0 +1,521 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) + +// Must be first so this translation unit uses the shared-provider Node world. +#include "core/providers/shared_library/provider_api.h" + +#include "contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h" + +#include +#include +#include +#include +#include + +#include "core/common/float16.h" +#include "core/platform/env_var_utils.h" +#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h" +#include "contrib_ops/cuda/bert/flash_attention/flash_api.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace { + +// Positional indices from the Microsoft-domain GroupQueryAttention v1 schema. +constexpr size_t kQuery = 0; +constexpr size_t kKey = 1; +constexpr size_t kValue = 2; +constexpr size_t kPastKey = 3; +constexpr size_t kPastValue = 4; +constexpr size_t kSequenceLengths = 5; +constexpr size_t kTotalSequenceLength = 6; +constexpr size_t kCosCache = 7; +constexpr size_t kSinCache = 8; +constexpr size_t kPositionIds = 9; +constexpr size_t kAttentionBias = 10; +constexpr size_t kHeadSink = 11; +constexpr size_t kKScale = 12; +constexpr size_t kVScale = 13; +constexpr size_t kQNorm = 14; +constexpr size_t kKNorm = 15; + +bool Present(gsl::span shapes, size_t index) { + return GetWorkspaceInputShape(shapes, index).GetState() != WorkspaceInputShapeState::Missing; +} + +const TensorShape* Shape(gsl::span shapes, size_t index) { + return GetWorkspaceInputShape(shapes, index).GetShape(); +} + +bool PositiveDims(const TensorShape& shape) { + return std::all_of(shape.GetDims().begin(), shape.GetDims().end(), + [](int64_t dim) { return dim > 0; }); +} + +bool PairPresent(gsl::span shapes, size_t left, size_t right) { + return Present(shapes, left) == Present(shapes, right); +} + +std::optional ParseQuantization(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), + [](unsigned char c) { return static_cast(std::toupper(c)); }); + if (value == "NONE") return GQAKvQuantizationType::None; + if (value == "PER_TENSOR") return GQAKvQuantizationType::PerTensor; + if (value == "PER_CHANNEL") return GQAKvQuantizationType::PerChannel; + return std::nullopt; +} + +bool HasZeroDimension(gsl::span shapes) { + for (const auto& input : shapes) { + const auto* shape = input.GetShape(); + if (shape != nullptr && + std::find(shape->GetDims().begin(), shape->GetDims().end(), int64_t{0}) != + shape->GetDims().end()) { + return true; + } + } + return false; +} + +bool ValidateAuxiliaryShapes(const GQAWorkspaceEstimateConfig& config, + gsl::span shapes, + int64_t batch_bound, + int64_t sequence_bound, + int64_t head_bound, + int64_t total_kv_bound) { + const TensorShape* sequence_lengths = Shape(shapes, kSequenceLengths); + const TensorShape* total = Shape(shapes, kTotalSequenceLength); + if (sequence_lengths == nullptr || total == nullptr || + sequence_lengths->NumDimensions() == 0 || + !(total->NumDimensions() == 0 || + (total->NumDimensions() == 1 && (*total)[0] == 1))) { + return false; + } + size_t batch_dimension_count = 0; + for (int64_t dim : sequence_lengths->GetDims()) { + if (dim != 1 && dim != batch_bound) return false; + if (dim == batch_bound && batch_bound != 1) ++batch_dimension_count; + } + if (batch_bound != 1 && batch_dimension_count != 1) return false; + + if (!PairPresent(shapes, kCosCache, kSinCache) || + !PairPresent(shapes, kQNorm, kKNorm)) { + return false; + } + if (config.do_rotary && !Present(shapes, kCosCache)) return false; + if (Present(shapes, kCosCache)) { + const auto* cos = Shape(shapes, kCosCache); + const auto* sin = Shape(shapes, kSinCache); + if (cos == nullptr || sin == nullptr || cos->NumDimensions() != 2 || + sin->NumDimensions() != 2 || !PositiveDims(*cos) || !PositiveDims(*sin)) { + return false; + } + const int64_t rotary_width = (*cos)[1]; + if (rotary_width != (*sin)[1] || rotary_width % 8 != 0 || + rotary_width > head_bound / 2) { + return false; + } + } + if (Present(shapes, kQNorm)) { + const auto* q = Shape(shapes, kQNorm); + const auto* k = Shape(shapes, kKNorm); + if (q == nullptr || k == nullptr || q->NumDimensions() != 1 || + k->NumDimensions() != 1 || (*q)[0] <= 0 || (*k)[0] <= 0 || + std::min((*q)[0], (*k)[0]) < head_bound) { + return false; + } + } + if (Present(shapes, kHeadSink)) { + const auto* sink = Shape(shapes, kHeadSink); + if (sink == nullptr || sink->NumDimensions() != 1 || + (*sink)[0] != config.num_heads) { + return false; + } + } + if (Present(shapes, kAttentionBias)) { + const auto* bias = Shape(shapes, kAttentionBias); + if (bias == nullptr || bias->NumDimensions() != 4 || !PositiveDims(*bias) || + ((*bias)[0] != 1 && (*bias)[0] != batch_bound) || + ((*bias)[1] != 1 && (*bias)[1] != config.num_heads) || + (*bias)[2] < sequence_bound || (*bias)[3] < total_kv_bound) { + return false; + } + } + if (Present(shapes, kPositionIds)) { + const auto* positions = Shape(shapes, kPositionIds); + if (positions == nullptr || positions->NumDimensions() != 2 || + !PositiveDims(*positions) || (*positions)[0] != batch_bound || + (*positions)[1] < sequence_bound) { + return false; + } + } + return true; +} + +bool ValidateScaleShape(GQAKvQuantizationType type, const TensorShape* shape, + int64_t kv_heads, int64_t head_bound) { + if (type == GQAKvQuantizationType::None) return true; + if (shape == nullptr || !PositiveDims(*shape)) return false; + if (type == GQAKvQuantizationType::PerTensor) { + return std::all_of(shape->GetDims().begin(), shape->GetDims().end(), + [](int64_t dim) { return dim == 1; }); + } + if (shape->NumDimensions() > 4) return false; + const auto& dims = shape->GetDims(); + if (dims.size() >= 2 && dims[dims.size() - 2] != 1 && + dims[dims.size() - 2] > kv_heads) { + return false; + } + return dims.back() == 1 || dims.back() <= head_bound; +} + +std::optional BuildBounds( + const GQAWorkspaceEstimateConfig& config, + gsl::span shapes, + const cudaDeviceProp& device_prop, + const AttentionKernelOptions& kernel_options) { + if (config.qkv_element_size != 2 || + (config.cache_element_size != 1 && config.cache_element_size != 2) || + config.num_heads <= 0 || config.kv_num_heads <= 0 || + config.num_heads % config.kv_num_heads != 0 || + config.num_heads > std::numeric_limits::max() || + config.kv_num_heads > std::numeric_limits::max() || + (config.causal != 0 && config.causal != 1) || + (config.local_window_size != -1 && config.local_window_size <= 0) || + (config.sliding_window_cache && config.local_window_size <= 0) || + (config.causal == 0 && config.local_window_size != -1) || + !std::isfinite(config.softcap) || + !Present(shapes, kQuery) || !Present(shapes, kSequenceLengths) || + !Present(shapes, kTotalSequenceLength) || + !PairPresent(shapes, kKey, kValue) || + !PairPresent(shapes, kPastKey, kPastValue) || + !PairPresent(shapes, kKScale, kVScale) || + HasZeroDimension(shapes)) { + return std::nullopt; + } + + // WorkspaceInputShape exposes the total_sequence_length scalar's shape, not + // its value. For a non-windowed cache that value can exceed past dim 2 and + // directly scales backend workspace. Non-windowed execution can also copy a + // full past tensor when only one past/present pair aliases, but alias state is + // unavailable here. Windowed execution is bounded by the cache shape and the + // runtime requires both past/present pairs to alias before allocating scratch. + if (!config.sliding_window_cache) return std::nullopt; + + const bool packed = !Present(shapes, kKey); + const TensorShape* query = Shape(shapes, kQuery); + if (query == nullptr || query->NumDimensions() != 3 || !PositiveDims(*query)) { + return std::nullopt; + } + int64_t batch_bound = (*query)[0]; + int64_t sequence_bound = (*query)[1]; + int64_t head_bound = 0; + if (packed) { + const int64_t head_factor = config.num_heads + 2 * config.kv_num_heads; + if ((*query)[2] < head_factor) return std::nullopt; + head_bound = (*query)[2] / head_factor; + } else { + const auto* key = Shape(shapes, kKey); + const auto* value = Shape(shapes, kValue); + if (key == nullptr || value == nullptr || key->NumDimensions() != 3 || + value->NumDimensions() != 3 || !PositiveDims(*key) || !PositiveDims(*value)) { + return std::nullopt; + } + batch_bound = std::min({batch_bound, (*key)[0], (*value)[0]}); + sequence_bound = std::min({sequence_bound, (*key)[1], (*value)[1]}); + head_bound = std::min({(*query)[2] / config.num_heads, + (*key)[2] / config.kv_num_heads, + (*value)[2] / config.kv_num_heads}); + } + if (head_bound < 8 || batch_bound <= 0 || sequence_bound <= 0) return std::nullopt; + + // Cache capacity is the only sound bound for total KV length available at + // this API. Without past tensors, the CPU scalar value is unavailable and a + // single-token invocation can name an arbitrarily larger total length. + if (!Present(shapes, kPastKey)) return std::nullopt; + const auto* past_key = Shape(shapes, kPastKey); + const auto* past_value = Shape(shapes, kPastValue); + if (past_key == nullptr || past_value == nullptr || + past_key->NumDimensions() != 4 || past_value->NumDimensions() != 4 || + !PositiveDims(*past_key) || !PositiveDims(*past_value)) { + return std::nullopt; + } + batch_bound = std::min({batch_bound, (*past_key)[0], (*past_value)[0]}); + if ((*past_key)[1] < config.kv_num_heads || (*past_value)[1] < config.kv_num_heads) { + return std::nullopt; + } + if ((*past_key)[2] != (*past_value)[2]) return std::nullopt; + const int64_t capacity_bound = (*past_key)[2]; + int64_t cache_head_bound = std::min((*past_key)[3], (*past_value)[3]); + if (config.kv_cache_bit_width == 4) { + if (cache_head_bound > std::numeric_limits::max() / 2) return std::nullopt; + cache_head_bound *= 2; + } + head_bound = std::min(head_bound, cache_head_bound); + if (batch_bound <= 0 || capacity_bound <= 0 || head_bound < 8 || + capacity_bound != config.local_window_size || + !ValidateAuxiliaryShapes(config, shapes, batch_bound, sequence_bound, + head_bound, capacity_bound)) { + return std::nullopt; + } + + // Runtime rejects attention bias for sliding-window GQA. + if (Present(shapes, kAttentionBias)) return std::nullopt; + + const bool k_quantized = config.k_quantization != GQAKvQuantizationType::None; + const bool v_quantized = config.v_quantization != GQAKvQuantizationType::None; + if (k_quantized != v_quantized || + (k_quantized && (!Present(shapes, kKScale) || config.cache_element_size != 1 || + (config.kv_cache_bit_width != 8 && config.kv_cache_bit_width != 4))) || + (!k_quantized && (config.cache_element_size != 2 || + config.kv_cache_bit_width != 0))) { + return std::nullopt; + } + if (Present(shapes, kKScale) && + (!ValidateScaleShape(config.k_quantization, Shape(shapes, kKScale), + config.kv_num_heads, head_bound) || + !ValidateScaleShape(config.v_quantization, Shape(shapes, kVScale), + config.kv_num_heads, head_bound))) { + return std::nullopt; + } + + GQAWorkspaceBounds bounds; + bounds.qkv_element_size = config.qkv_element_size; + bounds.cache_element_size = config.cache_element_size; + bounds.batch_size_bound = batch_bound; + bounds.sequence_length_bound = sequence_bound; + bounds.num_heads = config.num_heads; + bounds.kv_num_heads = config.kv_num_heads; + bounds.head_size_bound = head_bound; + bounds.present_kv_cache_capacity_bound = capacity_bound; + bounds.kv_cache_bit_width = config.kv_cache_bit_width; + bounds.k_quantization = config.k_quantization; + bounds.v_quantization = config.v_quantization; + bounds.is_windowed_kv_cache = config.sliding_window_cache; + bounds.do_rotary = config.do_rotary; + bounds.is_packed_qkv = packed; + bounds.use_qk_norm = Present(shapes, kQNorm); + bounds.prompt_reachable = true; + bounds.decode_reachable = true; + bounds.device_major = device_prop.major; + bounds.device_minor = device_prop.minor; + bounds.multi_processor_count = device_prop.multiProcessorCount; + bounds.is_bf16 = config.is_bf16; + bounds.local_window_size = config.local_window_size; + bounds.xqa_head_sink_storage = + Present(shapes, kHeadSink) + ? (config.head_sink_is_prepacked ? GQAXqaHeadSinkStorage::PrepackedFp32 + : GQAXqaHeadSinkStorage::DynamicConversion) + : GQAXqaHeadSinkStorage::None; + bounds.head_sink_may_be_prepacked = + Present(shapes, kHeadSink) && config.head_sink_may_be_prepacked; + + const bool has_bias = Present(shapes, kAttentionBias); + const bool has_sink = Present(shapes, kHeadSink); + const bool smooth_supported_by_xqa = !config.smooth_softmax || has_sink; + const bool quantized_xqa = k_quantized && config.kv_cache_bit_width == 8; + const int64_t group_size = config.num_heads / config.kv_num_heads; + const bool xqa_head_reachable = + (head_bound >= 64) && + ((head_bound >= 64 && IsSupportedGQAXqaHeadSize(64)) || + (head_bound >= 128 && IsSupportedGQAXqaHeadSize(128)) || + (head_bound >= 256 && IsSupportedGQAXqaHeadSize(256))); + const bool fp8_cache = k_quantized && config.cache_is_fp8; + const bool xqa_reachable = + config.enable_xqa && config.causal == 1 && !has_bias && + device_prop.major >= 8 && config.softcap == 0.0f && + smooth_supported_by_xqa && (!bounds.use_qk_norm || !k_quantized) && + (!k_quantized || config.kv_cache_bit_width == 8) && + xqa_head_reachable && + IsSupportedGQAXqaGroupSize(group_size, k_quantized) && + (!fp8_cache || device_prop.major >= 9 || + (device_prop.major == 8 && device_prop.minor == 9)); + if (xqa_reachable) { + bounds.reachable_backends = + bounds.reachable_backends | GQAReachableBackend::Xqa; + bounds.xqa_kv_type = quantized_xqa + ? (config.cache_is_fp8 ? GQAXqaKvType::Fp8 + : GQAXqaKvType::Int8) + : GQAXqaKvType::None; + } + + bool flash_reachable = false; +#if USE_FLASH_ATTENTION + flash_reachable = + !has_bias && kernel_options.UseFlashAttention() && + onnxruntime::flash::is_any_supported_head_size( + device_prop, static_cast(head_bound), + static_cast(config.num_heads), + static_cast(config.kv_num_heads)); +#endif + if (flash_reachable) { + bounds.reachable_backends = + bounds.reachable_backends | GQAReachableBackend::Flash; + if (!config.disable_flash_decode && !k_quantized && !bounds.use_qk_norm && + !config.sliding_window_cache) { + bounds.reachable_backends = + bounds.reachable_backends | GQAReachableBackend::FlashFastDecode; + } + } + + bool mea_reachable = false; +#if USE_MEMORY_EFFICIENT_ATTENTION + mea_reachable = + !k_quantized && !has_bias && kernel_options.UseEfficientAttention() && + has_memory_efficient_attention_for_head_size_bounds( + device_prop.major * 10 + device_prop.minor, + /*is_half=*/!config.is_bf16, config.is_bf16, + static_cast(std::min(head_bound, std::numeric_limits::max())), + static_cast(std::min(head_bound, std::numeric_limits::max()))); +#endif + if (mea_reachable) { + bounds.reachable_backends = + bounds.reachable_backends | GQAReachableBackend::MemoryEfficient; + } + if (!k_quantized && !config.smooth_softmax && !has_sink) { + bounds.reachable_backends = + bounds.reachable_backends | GQAReachableBackend::Unfused; + } + + const bool cudnn_reachable = + !has_bias && !k_quantized && config.softcap == 0.0f && + !config.smooth_softmax && !has_sink && config.local_window_size == -1 && + (kernel_options.UseCudnnFlashAttention() || + (kernel_options.AllowCudnnFlashAttentionAuto() && device_prop.major >= 9)); + if (cudnn_reachable) { + bounds.reachable_backends = + bounds.reachable_backends | GQAReachableBackend::Cudnn; + } + return bounds; +} + +std::optional ConfigFromNode(const Node& node) { + if (node.OpType() != "GroupQueryAttention") return std::nullopt; + GQAWorkspaceEstimateConfig config; + bool found_heads = false; + bool found_kv_heads = false; + for (const auto& [name, attr] : node.GetAttributes()) { + if (name == "num_heads") { + config.num_heads = attr.i(); + found_heads = true; + } else if (name == "kv_num_heads") { + config.kv_num_heads = attr.i(); + found_kv_heads = true; + } else if (name == "causal") { + config.causal = attr.i(); + } else if (name == "local_window_size") { + config.local_window_size = attr.i(); + } else if (name == "sliding_window_cache") { + config.sliding_window_cache = attr.i() == 1; + } else if (name == "do_rotary") { + config.do_rotary = attr.i() == 1; + } else if (name == "smooth_softmax") { + config.smooth_softmax = attr.i() == 1; + } else if (name == "softcap") { + config.softcap = attr.f(); + } else if (name == "kv_cache_bit_width") { + config.kv_cache_bit_width = attr.i(); + } else if (name == "k_quant_type") { + const auto type = ParseQuantization(attr.s()); + if (!type.has_value()) return std::nullopt; + config.k_quantization = *type; + } else if (name == "v_quant_type") { + const auto type = ParseQuantization(attr.s()); + if (!type.has_value()) return std::nullopt; + config.v_quantization = *type; + } + } + if (!found_heads || !found_kv_heads) return std::nullopt; + const auto& inputs = node.InputDefs(); + if (inputs.empty() || inputs[0] == nullptr) return std::nullopt; + const auto* query_type = inputs[0]->TypeAsProto(); + if (query_type == nullptr || !query_type->has_tensor_type()) return std::nullopt; + const int query_elem = query_type->tensor_type().elem_type(); + if (query_elem != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 && + query_elem != ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16) { + return std::nullopt; + } + config.qkv_element_size = 2; + config.is_bf16 = query_elem == ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16; + + const NodeArg* cache_arg = + inputs.size() > kPastKey && inputs[kPastKey] != nullptr && inputs[kPastKey]->Exists() + ? inputs[kPastKey] + : nullptr; + if (cache_arg == nullptr || cache_arg->TypeAsProto() == nullptr || + !cache_arg->TypeAsProto()->has_tensor_type()) { + return std::nullopt; + } + const int cache_elem = cache_arg->TypeAsProto()->tensor_type().elem_type(); + config.cache_is_fp8 = + cache_elem == ONNX_NAMESPACE::TensorProto_DataType_FLOAT8E4M3FN; + config.cache_element_size = + cache_elem == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 || + cache_elem == ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16 + ? 2 + : 1; + config.enable_xqa = + ParseEnvironmentVariableWithDefault("ORT_ENABLE_XQA", 1) != 0; + config.disable_flash_decode = + ParseEnvironmentVariableWithDefault("ORT_DISABLE_FLASH_DECODE", false); + return config; +} + +} // namespace + +std::optional EstimateGroupQueryAttentionWorkspace( + const GQAWorkspaceEstimateConfig& config, + gsl::span input_shapes, + const cudaDeviceProp& device_prop, + const AttentionKernelOptions& kernel_options) { + const auto bounds = BuildBounds(config, input_shapes, device_prop, kernel_options); + if (!bounds.has_value()) return std::nullopt; + const auto aggregate = GetGQAWorkspaceAggregateForBounds(*bounds); + return aggregate.status.IsOK() + ? std::optional{aggregate} + : std::nullopt; +} + +std::optional EstimateGroupQueryAttentionWorkspace( + const Node& node, + gsl::span input_shapes, + const cudaDeviceProp& device_prop, + const AttentionKernelOptions& kernel_options, + bool head_sink_is_constant_initializer) { + auto config = ConfigFromNode(node); + if (!config.has_value()) return std::nullopt; + config->head_sink_may_be_prepacked = head_sink_is_constant_initializer; + return EstimateGroupQueryAttentionWorkspace( + *config, input_shapes, device_prop, kernel_options); +} + +void SetGroupQueryAttentionWorkspaceRequirements( + const GQAWorkspaceAggregate& estimate, + InlinedVector& requirements) { + requirements.clear(); + if (!estimate.status.IsOK() || estimate.total_workspace_bytes == 0) return; + requirements.push_back( + {estimate.total_workspace_bytes, 0, kGQAWorkspaceAlignment}); +} + +void SetGroupQueryAttentionLevel1MemoryEstimate( + const GQAWorkspaceAggregate& workspace, + Level1MemoryEstimate& estimate) { + if (workspace.status.IsOK() && workspace.total_workspace_bytes != 0) { + estimate.runtime_workspace_bytes = workspace.total_workspace_bytes; + } + if (workspace.status.IsOK()) { + estimate.persistent_prepack_bytes = workspace.persistent_prepack_bytes; + estimate.initialization_scratch_bytes = + workspace.initialization_scratch_bytes; + } +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime + +#endif diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h new file mode 100644 index 0000000000000..02c6acdf308ac --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) + +#include +#include +#include + +#include +#include + +#include "core/common/inlined_containers.h" +#include "core/framework/level1_memory_estimate.h" +#include "core/framework/workspace_input_shape.h" +#include "core/framework/workspace_requirement.h" +#include "contrib_ops/cuda/bert/attention_kernel_options.h" +#include "contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h" + +namespace onnxruntime { +// Do not forward-declare Node. In-tree and shared-provider headers define +// different Node class keys; translation units must supply their own world. +namespace contrib { +namespace cuda { + +struct GQAWorkspaceEstimateConfig { + size_t qkv_element_size = 0; + size_t cache_element_size = 0; + int64_t num_heads = 0; + int64_t kv_num_heads = 0; + int64_t causal = 1; + int64_t local_window_size = -1; + bool sliding_window_cache = false; + bool do_rotary = false; + bool smooth_softmax = false; + float softcap = 0.0f; + GQAKvQuantizationType k_quantization = GQAKvQuantizationType::None; + GQAKvQuantizationType v_quantization = GQAKvQuantizationType::None; + int64_t kv_cache_bit_width = 0; + bool is_bf16 = false; + bool cache_is_fp8 = false; + bool enable_xqa = true; + bool disable_flash_decode = false; + // Exact kernel state; true omits per-run conversion scratch. + bool head_sink_is_prepacked = false; + // Level-1 possibility; true charges persistent and initialization lifetimes. + bool head_sink_may_be_prepacked = false; +}; + +std::optional EstimateGroupQueryAttentionWorkspace( + const GQAWorkspaceEstimateConfig& config, + gsl::span input_shapes, + const cudaDeviceProp& device_prop, + const AttentionKernelOptions& kernel_options); + +std::optional EstimateGroupQueryAttentionWorkspace( + const Node& node, + gsl::span input_shapes, + const cudaDeviceProp& device_prop, + const AttentionKernelOptions& kernel_options, + bool head_sink_is_constant_initializer = false); + +void SetGroupQueryAttentionWorkspaceRequirements( + const GQAWorkspaceAggregate& estimate, + InlinedVector& requirements); + +void SetGroupQueryAttentionLevel1MemoryEstimate( + const GQAWorkspaceAggregate& workspace, + Level1MemoryEstimate& estimate); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime + +#endif diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index e007c27d1d6fc..441966194fc1d 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -51,6 +51,7 @@ #if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) #include "contrib_ops/cuda/bert/packed_attention_workspace_estimate.h" +#include "contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h" #endif using namespace onnxruntime::common; @@ -3594,9 +3595,8 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, #endif #if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) - // PackedAttention and PackedMultiHeadAttention use the same Level-1 - // log-only contract as MatMulNBits. Route-aware workspace is not added to - // the partition budget until the planner integration is available. + // PackedAttention and PackedMultiHeadAttention remain log-only. Their + // route-aware workspace is not added to the partition budget yet. if (node != nullptr && (node->OpType() == "PackedAttention" || node->OpType() == "PackedMultiHeadAttention") && @@ -3612,6 +3612,32 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, << ": " << ws->total_workspace_bytes << " bytes"; } } + + if (node != nullptr && node->OpType() == "GroupQueryAttention" && + node->Domain() == kMSDomain) { + // Unlike PA/PMHA above, GQA participates in #31962 accounting. A + // successful estimate is supplied to ComputeResourceCount and can + // affect the CUDA partition acceptance decision. + const auto input_shapes = ResolveNodeInputShapes( + *node, &graph.GetGraph(), + resource_accountant->GetMaxShapeInferenceResult()); + const auto& input_defs = node->InputDefs(); + const bool head_sink_is_constant_initializer = + input_defs.size() > 11 && input_defs[11] != nullptr && + input_defs[11]->Exists() && + graph.IsConstantInitializer(input_defs[11]->Name(), true); + const auto ws = contrib::cuda::EstimateGroupQueryAttentionWorkspace( + *node, gsl::make_span(input_shapes), GetDeviceProp(), + *GetAttentionKernelOptions(), head_sink_is_constant_initializer); + if (ws.has_value()) { + Level1MemoryEstimate estimate; + contrib::cuda::SetGroupQueryAttentionLevel1MemoryEstimate(*ws, estimate); + level1_memory_estimate = estimate; + LOGS(logger, VERBOSE) << "Level-1 memory estimate for " << node->Name() + << ": runtime workspace=" + << ws->total_workspace_bytes << " bytes"; + } + } #endif const auto resource_count_variant = diff --git a/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc b/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc new file mode 100644 index 0000000000000..f1efe514cf322 --- /dev/null +++ b/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc @@ -0,0 +1,913 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" + +#if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/framework/op_kernel.h" +#include "core/framework/session_state.h" +#include "core/graph/graph.h" +#include "core/providers/cuda/cuda_execution_provider.h" +#include "core/providers/cuda/cuda_execution_provider_info.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "contrib_ops/cpu/bert/attention_common.h" +#include "contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h" +#include "test/test_environment.h" +#include "test/util/include/asserts.h" +#include "test/util/include/inference_session_wrapper.h" +#include "test/util/include/scoped_env_vars.h" + +namespace onnxruntime { +namespace test { +namespace { + +using contrib::attention::AttentionBackend; +using contrib::cuda::EstimateGroupQueryAttentionWorkspace; +using contrib::cuda::GetGQACompleteWorkspaceRecipe; +using contrib::cuda::GetGQAEffectiveWorkspaceKvLength; +using contrib::cuda::GetGQAFlashWorkspaceRecipe; +using contrib::cuda::GetGQAWorkspaceAggregateForBounds; +using contrib::cuda::GQABackend; +using contrib::cuda::GQAConcreteRoute; +using contrib::cuda::GQAFlashConfig; +using contrib::cuda::GQAKvQuantizationType; +using contrib::cuda::GQAPreprocessMode; +using contrib::cuda::GQAReachableBackend; +using contrib::cuda::GQAWorkspaceAggregate; +using contrib::cuda::GQAWorkspaceBounds; +using contrib::cuda::GQAWorkspaceEstimateConfig; +using contrib::cuda::GQAWorkspaceProblem; +using contrib::cuda::GQAXqaHeadSinkStorage; +using contrib::cuda::HasGQAReachableBackend; +using contrib::cuda::SetGroupQueryAttentionLevel1MemoryEstimate; +using contrib::cuda::SetGroupQueryAttentionWorkspaceRequirements; + +constexpr int kMath = static_cast(AttentionBackend::MATH); +constexpr int kFlash = static_cast(AttentionBackend::FLASH_ATTENTION); + +WorkspaceInputShape Known(std::initializer_list dims) { + return WorkspaceInputShape::PresentWithShape(TensorShape{TensorShapeVector{dims}}); +} + +std::array SeparateShapes(int64_t sequence = 4, + int64_t head = 64, + int64_t capacity = 256) { + std::array shapes; + shapes[0] = Known({2, sequence, 8 * head}); + shapes[1] = Known({2, sequence, 2 * head}); + shapes[2] = Known({2, sequence, 2 * head}); + shapes[3] = Known({2, 2, capacity, head}); + shapes[4] = Known({2, 2, capacity, head}); + shapes[5] = Known({2}); + shapes[6] = Known({}); + return shapes; +} + +std::array PackedShapes() { + auto shapes = SeparateShapes(); + shapes[0] = Known({2, 4, 12 * 64}); + shapes[1] = {}; + shapes[2] = {}; + return shapes; +} + +GQAWorkspaceEstimateConfig Config() { + GQAWorkspaceEstimateConfig config; + config.qkv_element_size = 2; + config.cache_element_size = 2; + config.num_heads = 8; + config.kv_num_heads = 2; + config.local_window_size = 256; + config.sliding_window_cache = true; + return config; +} + +cudaDeviceProp Device(int major = 8, int minor = 0) { + cudaDeviceProp device{}; + device.major = major; + device.minor = minor; + device.multiProcessorCount = 108; + return device; +} + +GQAWorkspaceBounds Bounds() { + GQAWorkspaceBounds bounds; + bounds.qkv_element_size = 2; + bounds.cache_element_size = 2; + bounds.batch_size_bound = 2; + bounds.sequence_length_bound = 4; + bounds.num_heads = 8; + bounds.kv_num_heads = 2; + bounds.head_size_bound = 64; + bounds.present_kv_cache_capacity_bound = 256; + bounds.prompt_reachable = true; + bounds.decode_reachable = true; + bounds.multi_processor_count = 108; + return bounds; +} + +std::optional EstimateFromNode( + gsl::span input_shapes, + const AttentionKernelOptions& options) { + ONNX_NAMESPACE::TypeProto query_type; + query_type.mutable_tensor_type()->set_elem_type( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); + ONNX_NAMESPACE::TypeProto cache_type; + cache_type.mutable_tensor_type()->set_elem_type( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); + NodeArg query{"query", &query_type}; + NodeArg key{"key", &query_type}; + NodeArg value{"value", &query_type}; + NodeArg past_key{"past_key", &cache_type}; + + NodeAttributes attributes; + const auto add_int_attribute = [&attributes](const char* name, int64_t value) { + ONNX_NAMESPACE::AttributeProto attribute; + attribute.set_name(name); + attribute.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + attribute.set_i(value); + attributes.emplace(name, std::move(attribute)); + }; + add_int_attribute("num_heads", 8); + add_int_attribute("kv_num_heads", 2); + add_int_attribute("causal", 1); + add_int_attribute("local_window_size", 256); + add_int_attribute("sliding_window_cache", 1); + + const std::vector inputs{&query, &key, &value, &past_key}; + const std::vector outputs; + Node node{"gqa", "GroupQueryAttention", "", inputs, outputs, &attributes, kMSDomain}; + return EstimateGroupQueryAttentionWorkspace( + node, input_shapes, Device(), options); +} + +void SetValueInfo(ONNX_NAMESPACE::ValueInfoProto& value_info, + const char* name, + int32_t element_type, + std::initializer_list dimensions) { + value_info.set_name(name); + auto* tensor_type = value_info.mutable_type()->mutable_tensor_type(); + tensor_type->set_elem_type(element_type); + auto* shape = tensor_type->mutable_shape(); + for (int64_t dimension : dimensions) { + shape->add_dim()->set_dim_value(dimension); + } +} + +std::string BuildGroupQueryAttentionKernelModel() { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(ONNX_NAMESPACE::IR_VERSION); + auto* onnx_opset = model.add_opset_import(); + onnx_opset->set_domain(""); + onnx_opset->set_version(17); + auto* ms_opset = model.add_opset_import(); + ms_opset->set_domain(kMSDomain); + ms_opset->set_version(1); + + auto* graph = model.mutable_graph(); + graph->set_name("group_query_attention_workspace_level2"); + auto* node = graph->add_node(); + node->set_domain(kMSDomain); + node->set_name("gqa"); + node->set_op_type("GroupQueryAttention"); + for (const char* input_name : + {"query", "key", "value", "past_key", "past_value", "seqlens_k", + "total_sequence_length", "", "", "", "", "head_sink"}) { + node->add_input(input_name); + } + node->add_output("output"); + node->add_output("present_key"); + node->add_output("present_value"); + + const auto add_int_attribute = [node](const char* name, int64_t value) { + auto* attribute = node->add_attribute(); + attribute->set_name(name); + attribute->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + attribute->set_i(value); + }; + add_int_attribute("num_heads", 8); + add_int_attribute("kv_num_heads", 2); + add_int_attribute("local_window_size", 256); + add_int_attribute("sliding_window_cache", 1); + + constexpr int32_t kFloat16 = ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; + constexpr int32_t kInt32 = ONNX_NAMESPACE::TensorProto_DataType_INT32; + SetValueInfo(*graph->add_input(), "query", kFloat16, {2, 4, 512}); + SetValueInfo(*graph->add_input(), "key", kFloat16, {2, 4, 128}); + SetValueInfo(*graph->add_input(), "value", kFloat16, {2, 4, 128}); + SetValueInfo(*graph->add_input(), "past_key", kFloat16, {2, 2, 256, 64}); + SetValueInfo(*graph->add_input(), "past_value", kFloat16, {2, 2, 256, 64}); + SetValueInfo(*graph->add_input(), "seqlens_k", kInt32, {2}); + SetValueInfo(*graph->add_input(), "total_sequence_length", kInt32, {}); + SetValueInfo(*graph->add_output(), "output", kFloat16, {2, 4, 512}); + SetValueInfo(*graph->add_output(), "present_key", kFloat16, {2, 2, 256, 64}); + SetValueInfo(*graph->add_output(), "present_value", kFloat16, {2, 2, 256, 64}); + + auto* head_sink = graph->add_initializer(); + head_sink->set_name("head_sink"); + head_sink->set_data_type(kFloat16); + head_sink->add_dims(8); + head_sink->mutable_raw_data()->assign(8 * sizeof(uint16_t), '\0'); + + std::string bytes; + model.SerializeToString(&bytes); + return bytes; +} + +const Node* FindNodeByOpType(const Graph& graph, const char* op_type) { + for (const auto& node : graph.Nodes()) { + if (node.OpType() == op_type) { + return &node; + } + } + return nullptr; +} + +bool HasCudaDevice() { + int device_count = 0; + return cudaGetDeviceCount(&device_count) == cudaSuccess && device_count > 0; +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, ParsesPackedAndSeparateLayouts) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + const auto separate = EstimateGroupQueryAttentionWorkspace( + Config(), SeparateShapes(), Device(), options); + const auto packed = EstimateGroupQueryAttentionWorkspace( + Config(), PackedShapes(), Device(), options); + ASSERT_TRUE(separate.has_value()); + ASSERT_TRUE(packed.has_value()); + EXPECT_GT(separate->total_workspace_bytes, 0u); + EXPECT_GT(packed->total_workspace_bytes, separate->total_workspace_bytes); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, NodeAdapterParsesAttributesAndTypes) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + const auto estimate = EstimateFromNode(SeparateShapes(), options); + ASSERT_TRUE(estimate.has_value()); + EXPECT_GT(estimate->total_workspace_bytes, 0u); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Estimate) { + if (!HasCudaDevice()) { + GTEST_SKIP() << "A CUDA device is required for the budget integration test."; + } + + ScopedEnvironmentVariables scoped_env_vars{{{"ORT_ENABLE_XQA", "1"}}}; + CUDAExecutionProviderInfo provider_info; + provider_info.sdpa_kernel = kMath; + const std::string model_bytes = BuildGroupQueryAttentionKernelModel(); + + auto make_session_options = [](size_t memory_limit_kb) { + SessionOptions session_options; + session_options.graph_optimization_level = TransformerLevel::Default; + const std::string partitioning_settings = + std::to_string(memory_limit_kb) + ","; + ORT_THROW_IF_ERROR(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsResourceCudaPartitioningSettings, + partitioning_settings.c_str())); + return session_options; + }; + + std::optional estimate; + { + InferenceSessionWrapper session(make_session_options(1024), + GetEnvironment()); + auto cuda_ep = std::make_shared(provider_info); + if (cuda_ep->GetDeviceProp().major < 8) { + GTEST_SKIP() << "XQA requires compute capability 8.0 or newer."; + } + ASSERT_STATUS_OK(session.RegisterExecutionProvider(cuda_ep)); + ASSERT_STATUS_OK(session.Load(model_bytes.data(), static_cast(model_bytes.size()))); + ASSERT_STATUS_OK(session.Initialize()); + + const Node* node = FindNodeByOpType(session.GetGraph(), "GroupQueryAttention"); + ASSERT_NE(node, nullptr); + ASSERT_EQ(node->GetExecutionProviderType(), kCudaExecutionProvider); + auto shapes = SeparateShapes(); + shapes[11] = Known({8}); + estimate = EstimateGroupQueryAttentionWorkspace( + *node, gsl::make_span(shapes), cuda_ep->GetDeviceProp(), + *cuda_ep->GetAttentionKernelOptions(), + /*head_sink_is_constant_initializer=*/true); + ASSERT_TRUE(estimate.has_value()); + + auto prepacked_config = Config(); + prepacked_config.head_sink_is_prepacked = true; + const auto prepacked = EstimateGroupQueryAttentionWorkspace( + prepacked_config, gsl::make_span(shapes), cuda_ep->GetDeviceProp(), + *cuda_ep->GetAttentionKernelOptions()); + ASSERT_TRUE(prepacked.has_value()); + EXPECT_GT(estimate->total_workspace_bytes, prepacked->total_workspace_bytes); + EXPECT_EQ(estimate->persistent_prepack_bytes, + prepacked->persistent_prepack_bytes); + EXPECT_EQ(estimate->initialization_scratch_bytes, + prepacked->initialization_scratch_bytes); + } + + constexpr size_t kHeadSinkInitializerBytes = 8 * sizeof(MLFloat16); + constexpr size_t kOutputBytes = 2 * 4 * 512 * sizeof(MLFloat16); + constexpr size_t kPresentCacheBytes = 2 * 2 * 256 * 64 * sizeof(MLFloat16); + constexpr size_t kAccountedTensorBytes = + kHeadSinkInitializerBytes + kOutputBytes + 2 * kPresentCacheBytes; + constexpr size_t kFallbackWorkspaceBytes = kAccountedTensorBytes / 2; + ASSERT_GT(estimate->total_workspace_bytes, kFallbackWorkspaceBytes); + EXPECT_EQ(estimate->persistent_prepack_bytes, 8 * sizeof(float)); + EXPECT_EQ(estimate->initialization_scratch_bytes, + 8 * sizeof(MLFloat16)); + const auto kilobytes_above = [](size_t bytes) { + return bytes / 1024 + 1; + }; + + { + const size_t accepted_limit_kb = kilobytes_above( + kAccountedTensorBytes + estimate->total_workspace_bytes + + estimate->persistent_prepack_bytes); + InferenceSessionWrapper session(make_session_options(accepted_limit_kb), + GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider( + std::make_shared(provider_info))); + ASSERT_STATUS_OK(session.Load(model_bytes.data(), static_cast(model_bytes.size()))); + ASSERT_STATUS_OK(session.Initialize()); + + const Node* node = FindNodeByOpType(session.GetGraph(), "GroupQueryAttention"); + ASSERT_NE(node, nullptr); + EXPECT_EQ(node->GetExecutionProviderType(), kCudaExecutionProvider); + } + + { + // Fallback accounting would admit this node. The larger Level-1 estimate + // must instead make CUDA reject it. + const size_t rejected_limit_kb = kilobytes_above( + kAccountedTensorBytes + kFallbackWorkspaceBytes); + ASSERT_LT(rejected_limit_kb * 1024, + kAccountedTensorBytes + estimate->total_workspace_bytes + + estimate->persistent_prepack_bytes); + InferenceSessionWrapper session(make_session_options(rejected_limit_kb), + GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider( + std::make_shared(provider_info))); + ASSERT_STATUS_OK(session.Load(model_bytes.data(), static_cast(model_bytes.size()))); + ASSERT_STATUS_OK(session.Initialize()); + + const Node* node = FindNodeByOpType(session.GetGraph(), "GroupQueryAttention"); + ASSERT_NE(node, nullptr); + EXPECT_NE(node->GetExecutionProviderType(), kCudaExecutionProvider); + } +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedTotalKvAndAliasingAreUnavailable) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto config = Config(); + config.sliding_window_cache = false; + config.local_window_size = -1; + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(/*sequence=*/1, /*head=*/64, /*capacity=*/128), + Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsCacheCapacityDifferentFromWindow) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto config = Config(); + config.local_window_size = 128; + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(), Device(), options) + .has_value()); + EXPECT_TRUE(EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(/*sequence=*/4, /*head=*/64, /*capacity=*/128), + Device(), options) + .has_value()); + + auto shapes = SeparateShapes(/*sequence=*/4, /*head=*/64, /*capacity=*/128); + shapes[4] = Known({2, 2, 256, 64}); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, shapes, Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsOptionalHolesAndMalformedGeometry) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto shapes = SeparateShapes(); + shapes[2] = {}; + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + Config(), shapes, Device(), options) + .has_value()); + + shapes = SeparateShapes(); + shapes[3] = Known({2, 2, 256}); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + Config(), shapes, Device(), options) + .has_value()); + + shapes = SeparateShapes(); + shapes[0] = WorkspaceInputShape::PresentWithoutShape(); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + Config(), shapes, Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsBiasWithoutHeadSink) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto shapes = SeparateShapes(); + shapes[10] = Known({2, 8, 4}); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + Config(), shapes, Device(), options) + .has_value()); + + shapes[10] = Known({2, 8, 4, 256}); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + Config(), shapes, Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsUndersizedPositionIds) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto shapes = SeparateShapes(); + shapes[9] = Known({1, 1}); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + Config(), shapes, Device(), options) + .has_value()); + + shapes[9] = Known({2, 4}); + EXPECT_TRUE(EstimateGroupQueryAttentionWorkspace( + Config(), shapes, Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, ValidatesPartialRotaryCaches) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto config = Config(); + config.do_rotary = true; + auto shapes = SeparateShapes(); + constexpr int64_t kMaxDimension = std::numeric_limits::max(); + shapes[7] = Known({kMaxDimension, 8}); + shapes[8] = Known({kMaxDimension, 8}); + EXPECT_TRUE(EstimateGroupQueryAttentionWorkspace( + config, shapes, Device(), options) + .has_value()); + + shapes[8] = Known({kMaxDimension, 16}); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, shapes, Device(), options) + .has_value()); + + shapes[7] = Known({kMaxDimension, 40}); + shapes[8] = Known({kMaxDimension, 40}); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, shapes, Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsNoncausalLocalWindow) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto config = Config(); + config.causal = 0; + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(), Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, ValidatesQuantizedCacheMetadata) { + AttentionKernelOptions options; + options.InitializeOnce(kFlash | kMath, true); + for (int bit_width : {8, 4}) { + auto config = Config(); + config.cache_element_size = 1; + config.kv_cache_bit_width = bit_width; + config.k_quantization = GQAKvQuantizationType::PerTensor; + config.v_quantization = GQAKvQuantizationType::PerChannel; + auto shapes = SeparateShapes(); + const int64_t stored_head = bit_width == 4 ? 32 : 64; + shapes[3] = Known({2, 2, 256, stored_head}); + shapes[4] = Known({2, 2, 256, stored_head}); + shapes[12] = Known({1}); + shapes[13] = Known({1, 2, 1, 64}); + EXPECT_TRUE(EstimateGroupQueryAttentionWorkspace( + config, shapes, Device(), options) + .has_value()); + } + + auto config = Config(); + config.cache_element_size = 1; + config.kv_cache_bit_width = 8; + config.k_quantization = GQAKvQuantizationType::PerTensor; + config.v_quantization = GQAKvQuantizationType::PerTensor; + auto shapes = SeparateShapes(); + shapes[12] = Known({1}); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, shapes, Device(), options) + .has_value()); +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, FindsSmallerSupportedHeadRoutes) { + auto bounds = Bounds(); + bounds.head_size_bound = 263; + bounds.reachable_backends = + GQAReachableBackend::Xqa | GQAReachableBackend::Flash | + GQAReachableBackend::MemoryEfficient; + const auto estimate = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(estimate.status.IsOK()) << estimate.status.message; + EXPECT_TRUE(HasGQAReachableBackend(estimate.sized_backends, + GQAReachableBackend::Xqa)); + EXPECT_TRUE(HasGQAReachableBackend(estimate.sized_backends, + GQAReachableBackend::Flash)); +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, PromptDecodeAndWindowRoutesAreBounded) { + auto bounds = Bounds(); + bounds.reachable_backends = GQAReachableBackend::Unfused; + bounds.is_windowed_kv_cache = true; + const auto mixed = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(mixed.status.IsOK()); + + bounds.prompt_reachable = false; + bounds.sequence_length_bound = 1; + const auto decode = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(decode.status.IsOK()); + EXPECT_GT(mixed.total_workspace_bytes, decode.total_workspace_bytes); + + bounds.decode_reachable = false; + bounds.prompt_reachable = true; + bounds.sequence_length_bound = 4; + const auto prompt = GetGQAWorkspaceAggregateForBounds(bounds); + EXPECT_TRUE(prompt.status.IsOK()); +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, WindowedBoundsUseFinalCapacityAndTransientExtent) { + constexpr int64_t kCapacity = 256; + constexpr int64_t kSequence = 4; + constexpr int64_t kLargeTotalSequenceLength = 1'000'000; + auto bounds = Bounds(); + bounds.present_kv_cache_capacity_bound = kCapacity; + bounds.sequence_length_bound = kSequence; + bounds.is_windowed_kv_cache = true; + + for (GQAReachableBackend backend : + {GQAReachableBackend::Flash, GQAReachableBackend::MemoryEfficient, + GQAReachableBackend::Unfused}) { + bounds.reachable_backends = backend; + const auto estimate = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(estimate.status.IsOK()) << estimate.status.message; + + GQAWorkspaceProblem problem; + problem.qkv_element_size = bounds.qkv_element_size; + problem.cache_element_size = bounds.cache_element_size; + problem.batch_size = bounds.batch_size_bound; + problem.sequence_length = kSequence; + problem.num_heads = bounds.num_heads; + problem.kv_num_heads = bounds.kv_num_heads; + problem.head_size = bounds.head_size_bound; + problem.present_kv_cache_capacity = kCapacity; + problem.is_windowed_kv_cache = true; + + GQAConcreteRoute route; + route.backend = backend == GQAReachableBackend::Flash + ? GQABackend::Flash + : (backend == GQAReachableBackend::MemoryEfficient + ? GQABackend::MemoryEfficient + : GQABackend::Unfused); + route.preparation.preprocess_mode = + backend == GQAReachableBackend::Flash + ? GQAPreprocessMode::Flash + : (backend == GQAReachableBackend::MemoryEfficient + ? GQAPreprocessMode::MemoryEfficient + : GQAPreprocessMode::Unfused); + const int64_t effective_kv_length = GetGQAEffectiveWorkspaceKvLength( + kLargeTotalSequenceLength, kCapacity + kSequence, true); + route.flash.total_sequence_length = effective_kv_length; + route.flash.multi_processor_count = bounds.multi_processor_count; + route.unfused.total_sequence_length = effective_kv_length; + + const auto concrete = GetGQACompleteWorkspaceRecipe(problem, route); + ASSERT_TRUE(concrete.status.IsOK()) << concrete.status.message; + EXPECT_EQ(concrete.recipe.preparation.effective_kv_cache_capacity, + kCapacity + kSequence); + EXPECT_EQ(concrete.recipe.preparation.staged_key_bytes, + static_cast(bounds.batch_size_bound * bounds.kv_num_heads * + (kCapacity + kSequence) * bounds.head_size_bound * + bounds.cache_element_size)); + if (backend == GQAReachableBackend::MemoryEfficient) { + EXPECT_EQ(concrete.recipe.memory_efficient.effective_kv_cache_capacity, + kCapacity + kSequence); + } else if (backend == GQAReachableBackend::Flash) { + EXPECT_EQ(concrete.recipe.flash.split_heuristic_kv_length, + static_cast(kCapacity + kSequence)); + } else { + const size_t qk_bytes = + static_cast(bounds.batch_size_bound * bounds.num_heads * + kSequence * (kCapacity + kSequence) * sizeof(float)); + EXPECT_EQ(concrete.recipe.unfused.qk_bytes, qk_bytes); + } + EXPECT_GE(estimate.total_workspace_bytes, + concrete.recipe.total_workspace_bytes); + } + + bounds.sequence_length_bound = 1; + bounds.prompt_reachable = false; + bounds.reachable_backends = GQAReachableBackend::Unfused; + const auto decode = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(decode.status.IsOK()); + + GQAWorkspaceProblem decode_problem; + decode_problem.qkv_element_size = bounds.qkv_element_size; + decode_problem.cache_element_size = bounds.cache_element_size; + decode_problem.batch_size = bounds.batch_size_bound; + decode_problem.sequence_length = 1; + decode_problem.num_heads = bounds.num_heads; + decode_problem.kv_num_heads = bounds.kv_num_heads; + decode_problem.head_size = bounds.head_size_bound; + decode_problem.present_kv_cache_capacity = kCapacity; + decode_problem.is_windowed_kv_cache = true; + GQAConcreteRoute decode_route; + decode_route.backend = GQABackend::Unfused; + decode_route.preparation.preprocess_mode = GQAPreprocessMode::Unfused; + decode_route.unfused.total_sequence_length = GetGQAEffectiveWorkspaceKvLength( + kLargeTotalSequenceLength, kCapacity, true); + const auto concrete_decode = + GetGQACompleteWorkspaceRecipe(decode_problem, decode_route); + ASSERT_TRUE(concrete_decode.status.IsOK()); + EXPECT_EQ(concrete_decode.recipe.preparation.effective_kv_cache_capacity, + kCapacity); + EXPECT_GE(decode.total_workspace_bytes, + concrete_decode.recipe.total_workspace_bytes); +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, FlashEnvelopeDominatesDiscontinuousEndpoints) { + auto bounds = Bounds(); + bounds.batch_size_bound = 1; + bounds.sequence_length_bound = 1; + bounds.num_heads = 2; + bounds.kv_num_heads = 2; + bounds.present_kv_cache_capacity_bound = 13830; + bounds.prompt_reachable = false; + bounds.reachable_backends = GQAReachableBackend::Flash; + const auto envelope = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(envelope.status.IsOK()); + + for (int64_t kv_length = 13820; kv_length <= 13830; ++kv_length) { + GQAWorkspaceProblem problem; + problem.qkv_element_size = 2; + problem.cache_element_size = 2; + problem.batch_size = 1; + problem.sequence_length = 1; + problem.num_heads = 2; + problem.kv_num_heads = 2; + problem.head_size = 64; + problem.present_kv_cache_capacity = kv_length; + + GQAFlashConfig flash; + flash.total_sequence_length = kv_length; + flash.multi_processor_count = 108; + const auto flash_recipe = GetGQAFlashWorkspaceRecipe(problem, flash); + ASSERT_TRUE(flash_recipe.status.IsOK()) << kv_length; + + GQAConcreteRoute route; + route.backend = GQABackend::Flash; + route.preparation.preprocess_mode = GQAPreprocessMode::Flash; + route.flash = flash; + const auto concrete = GetGQACompleteWorkspaceRecipe(problem, route); + ASSERT_TRUE(concrete.status.IsOK()) << kv_length; + EXPECT_GE(envelope.total_workspace_bytes, + concrete.recipe.total_workspace_bytes) + << kv_length << " selected splits " + << flash_recipe.recipe.selected_split_count; + } +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, FlashFastDecodeEnvelopeCoversEverySequenceLength) { + auto bounds = Bounds(); + bounds.batch_size_bound = 1; + bounds.sequence_length_bound = 130; + bounds.num_heads = 8; + bounds.kv_num_heads = 2; + bounds.head_size_bound = 128; + bounds.present_kv_cache_capacity_bound = 1024; + bounds.is_windowed_kv_cache = false; + bounds.local_window_size = -1; + bounds.reachable_backends = GQAReachableBackend::FlashFastDecode; + const auto envelope = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(envelope.status.IsOK()) << envelope.status.message; + + for (int64_t sequence_length = 1; + sequence_length <= bounds.sequence_length_bound; ++sequence_length) { + GQAWorkspaceProblem problem; + problem.qkv_element_size = bounds.qkv_element_size; + problem.cache_element_size = bounds.cache_element_size; + problem.batch_size = bounds.batch_size_bound; + problem.sequence_length = sequence_length; + problem.num_heads = bounds.num_heads; + problem.kv_num_heads = bounds.kv_num_heads; + problem.head_size = bounds.head_size_bound; + problem.present_kv_cache_capacity = bounds.present_kv_cache_capacity_bound; + + GQAConcreteRoute route; + route.backend = GQABackend::Flash; + route.preparation.preprocess_mode = GQAPreprocessMode::Flash; + route.preparation.use_flash_attention_fast_decode = true; + route.flash.total_sequence_length = bounds.present_kv_cache_capacity_bound; + route.flash.multi_processor_count = bounds.multi_processor_count; + route.flash.fast_decode = true; + const auto concrete = GetGQACompleteWorkspaceRecipe(problem, route); + ASSERT_TRUE(concrete.status.IsOK()) + << sequence_length << ": " << concrete.status.message; + EXPECT_GE(envelope.total_workspace_bytes, + concrete.recipe.total_workspace_bytes) + << sequence_length << " selected splits " + << concrete.recipe.flash.selected_split_count; + } +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, AggregatesExclusiveRoutesWithMax) { + auto xqa = Bounds(); + xqa.reachable_backends = GQAReachableBackend::Xqa; + xqa.xqa_head_sink_storage = GQAXqaHeadSinkStorage::DynamicConversion; + xqa.head_sink_may_be_prepacked = true; + const auto xqa_only = GetGQAWorkspaceAggregateForBounds(xqa); + ASSERT_TRUE(xqa_only.status.IsOK()); + auto unfused = xqa; + unfused.reachable_backends = GQAReachableBackend::Unfused; + const auto unfused_only = GetGQAWorkspaceAggregateForBounds(unfused); + ASSERT_TRUE(unfused_only.status.IsOK()); + ASSERT_GT(unfused_only.total_workspace_bytes, + xqa_only.total_workspace_bytes); + auto both = xqa; + both.reachable_backends = + GQAReachableBackend::Xqa | GQAReachableBackend::Unfused; + const auto aggregate = GetGQAWorkspaceAggregateForBounds(both); + ASSERT_TRUE(aggregate.status.IsOK()); + EXPECT_EQ(aggregate.total_workspace_bytes, + std::max(xqa_only.total_workspace_bytes, + unfused_only.total_workspace_bytes)); + EXPECT_EQ(aggregate.persistent_prepack_bytes, + static_cast(xqa.num_heads) * sizeof(float)); + EXPECT_EQ(aggregate.initialization_scratch_bytes, + static_cast(xqa.num_heads) * xqa.qkv_element_size); +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, + ExactPrepackedStateImpliesPrepackLifetimes) { + auto bounds = Bounds(); + bounds.reachable_backends = GQAReachableBackend::Xqa; + bounds.xqa_head_sink_storage = GQAXqaHeadSinkStorage::PrepackedFp32; + const auto aggregate = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(aggregate.status.IsOK()); + EXPECT_EQ(aggregate.persistent_prepack_bytes, + static_cast(bounds.num_heads) * sizeof(float)); + EXPECT_EQ(aggregate.initialization_scratch_bytes, + static_cast(bounds.num_heads) * bounds.qkv_element_size); +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, CudnnReachabilityIsUnavailable) { + auto bounds = Bounds(); + bounds.reachable_backends = + GQAReachableBackend::Unfused | GQAReachableBackend::Cudnn; + EXPECT_FALSE(GetGQAWorkspaceAggregateForBounds(bounds).status.IsOK()); + bounds.reachable_backends = GQAReachableBackend::Unfused; + EXPECT_TRUE(GetGQAWorkspaceAggregateForBounds(bounds).status.IsOK()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, DeclaresOneAlignedSlotAndOnlyWorkspaceAtLevel1) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + const auto estimate = EstimateGroupQueryAttentionWorkspace( + Config(), SeparateShapes(), Device(), options); + ASSERT_TRUE(estimate.has_value()); + InlinedVector requirements; + SetGroupQueryAttentionWorkspaceRequirements(*estimate, requirements); + ASSERT_EQ(requirements.size(), 1u); + EXPECT_EQ(requirements[0].size_bytes, estimate->total_workspace_bytes); + EXPECT_EQ(requirements[0].slot_id, 0); + EXPECT_EQ(requirements[0].alignment_bytes, 256u); + + Level1MemoryEstimate level1; + level1.runtime_transient_bytes = 17; + level1.persistent_prepack_bytes = 19; + level1.initialization_scratch_bytes = 23; + SetGroupQueryAttentionLevel1MemoryEstimate(*estimate, level1); + EXPECT_EQ(level1.runtime_workspace_bytes, estimate->total_workspace_bytes); + EXPECT_EQ(level1.runtime_transient_bytes, 17u); + EXPECT_EQ(level1.persistent_prepack_bytes, 0u); + EXPECT_EQ(level1.initialization_scratch_bytes, 0u); + + auto lifetime_estimate = *estimate; + lifetime_estimate.persistent_prepack_bytes = 29; + lifetime_estimate.initialization_scratch_bytes = 31; + SetGroupQueryAttentionLevel1MemoryEstimate(lifetime_estimate, level1); + EXPECT_EQ(level1.persistent_prepack_bytes, 29u); + EXPECT_EQ(level1.initialization_scratch_bytes, 31u); + + auto unavailable = *estimate; + unavailable.status.error = contrib::cuda::GQAWorkspaceError::Unavailable; + SetGroupQueryAttentionWorkspaceRequirements(unavailable, requirements); + EXPECT_TRUE(requirements.empty()); + + GQAWorkspaceAggregate zero; + zero.status = {}; + Level1MemoryEstimate zero_level1; + SetGroupQueryAttentionLevel1MemoryEstimate(zero, zero_level1); + EXPECT_FALSE(zero_level1.runtime_workspace_bytes.has_value()); +} + +TEST(GroupQueryAttentionWorkspaceEstimateTest, KernelDeclaresPrepackedHeadSinkRoot) { + if (!HasCudaDevice()) { + GTEST_SKIP() << "A CUDA device is required to construct the CUDA kernel."; + } + + ScopedEnvironmentVariables scoped_env_vars{{{"ORT_ENABLE_XQA", "1"}}}; + SessionOptions session_options; + session_options.graph_optimization_level = TransformerLevel::Default; + session_options.session_logid = "GroupQueryAttentionWorkspaceLevel2"; + InferenceSessionWrapper session(session_options, GetEnvironment()); + + CUDAExecutionProviderInfo provider_info; + provider_info.sdpa_kernel = kMath; + auto cuda_ep = std::make_shared(provider_info); + if (cuda_ep->GetDeviceProp().major < 8) { + GTEST_SKIP() << "XQA requires compute capability 8.0 or newer."; + } + ASSERT_STATUS_OK(session.RegisterExecutionProvider(cuda_ep)); + const std::string model_bytes = BuildGroupQueryAttentionKernelModel(); + ASSERT_STATUS_OK(session.Load(model_bytes.data(), static_cast(model_bytes.size()))); + ASSERT_STATUS_OK(session.Initialize()); + + const Graph& graph = session.GetGraph(); + const Node* node = FindNodeByOpType(graph, "GroupQueryAttention"); + ASSERT_NE(node, nullptr); + ASSERT_EQ(node->GetExecutionProviderType(), kCudaExecutionProvider); + const OpKernel* kernel = session.GetSessionState().GetKernel(node->Index()); + ASSERT_NE(kernel, nullptr); + + auto shapes = SeparateShapes(); + shapes[11] = Known({8}); + auto expected_config = Config(); + expected_config.head_sink_is_prepacked = true; + const auto expected = EstimateGroupQueryAttentionWorkspace( + expected_config, gsl::make_span(shapes), cuda_ep->GetDeviceProp(), + *cuda_ep->GetAttentionKernelOptions()); + ASSERT_TRUE(expected.has_value()); + + auto dynamic_config = expected_config; + dynamic_config.head_sink_is_prepacked = false; + const auto dynamic = EstimateGroupQueryAttentionWorkspace( + dynamic_config, gsl::make_span(shapes), cuda_ep->GetDeviceProp(), + *cuda_ep->GetAttentionKernelOptions()); + ASSERT_TRUE(dynamic.has_value()); + EXPECT_LT(expected->total_workspace_bytes, dynamic->total_workspace_bytes); + + InlinedVector requirements; + ASSERT_STATUS_OK(kernel->DeclareWorkspaceRequirements( + gsl::make_span(shapes), requirements)); + ASSERT_EQ(requirements.size(), 1U); + EXPECT_EQ(requirements[0].slot_id, 0); + EXPECT_EQ(requirements[0].size_bytes, expected->total_workspace_bytes); + EXPECT_EQ(requirements[0].alignment_bytes, 256U); + + shapes[0] = WorkspaceInputShape::PresentWithoutShape(); + ASSERT_STATUS_OK(kernel->DeclareWorkspaceRequirements( + gsl::make_span(shapes), requirements)); + EXPECT_TRUE(requirements.empty()); +} + +TEST(GroupQueryAttentionWorkspaceBoundsTest, CheckedOverflowIsUnavailable) { + auto bounds = Bounds(); + bounds.batch_size_bound = std::numeric_limits::max(); + bounds.sequence_length_bound = std::numeric_limits::max(); + bounds.num_heads = std::numeric_limits::max(); + bounds.kv_num_heads = 1; + bounds.head_size_bound = std::numeric_limits::max(); + bounds.present_kv_cache_capacity_bound = std::numeric_limits::max(); + bounds.reachable_backends = GQAReachableBackend::Unfused; + EXPECT_FALSE(GetGQAWorkspaceAggregateForBounds(bounds).status.IsOK()); + + bounds = Bounds(); + bounds.is_windowed_kv_cache = true; + bounds.present_kv_cache_capacity_bound = std::numeric_limits::max(); + bounds.sequence_length_bound = 2; + bounds.reachable_backends = GQAReachableBackend::Unfused; + EXPECT_FALSE(GetGQAWorkspaceAggregateForBounds(bounds).status.IsOK()); +} + +} // namespace +} // namespace test +} // namespace onnxruntime + +#endif diff --git a/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_header_test.cc b/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_header_test.cc index 5ed0525822904..4e842ab36d725 100644 --- a/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_header_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_header_test.cc @@ -5,6 +5,7 @@ // test-framework headers. Both the in-tree CUDA target and the plugin-internal // target compile it. #include "contrib_ops/cuda/bert/group_query_attention_workspace.h" +#include "contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h" #include @@ -19,6 +20,8 @@ using contrib::cuda::GQAMemoryEfficientWorkspaceRecipe; using contrib::cuda::GQAPreparationRecipe; using contrib::cuda::GQAPreparationRoute; using contrib::cuda::GQAUnfusedWorkspaceRecipe; +using contrib::cuda::GQAWorkspaceAggregate; +using contrib::cuda::GQAWorkspaceBounds; using contrib::cuda::GQAWorkspaceProblem; using contrib::cuda::GQAWorkspaceStatus; using contrib::cuda::GQAXqaWorkspaceRecipe; @@ -26,6 +29,8 @@ using contrib::cuda::IsSupportedGQAXqaGroupSize; using contrib::cuda::IsSupportedGQAXqaHeadSize; static_assert(std::is_trivially_copyable_v); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_trivially_copyable_v); static_assert(std::is_trivially_copyable_v); static_assert(std::is_trivially_copyable_v); static_assert(std::is_trivially_copyable_v); @@ -45,11 +50,15 @@ static_assert(GetGQAEffectiveWorkspaceKvLength(257, 8, false) == 257); void CompileGroupQueryAttentionWorkspaceHeaderInIsolation() { GQAWorkspaceProblem problem; + problem.requires_separate_past_buffer = true; + problem.past_kv_cache_capacity = 1; + GQAWorkspaceBounds bounds; GQAPreparationRoute route; GQAPreparationRecipe recipe; GQAConcreteRoute concrete_route; GQACompleteWorkspaceRecipe complete_recipe; (void)problem; + (void)bounds; (void)route; (void)recipe; (void)concrete_route;