From 60429e5378555b1fa95ccad7faf55c9d474b7544 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:12:15 +0000 Subject: [PATCH 01/12] Add GQA workspace estimation Add conservative graph-free route bounds, Level-1 accounting, and a single aligned Level-2 workspace declaration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../attention_workspace_estimation.md | 50 ++ .../cuda/bert/group_query_attention.cc | 48 ++ .../cuda/bert/group_query_attention.h | 6 + .../group_query_attention_workspace_bounds.cc | 368 ++++++++++++ .../group_query_attention_workspace_bounds.h | 81 +++ ...roup_query_attention_workspace_estimate.cc | 507 +++++++++++++++++ ...group_query_attention_workspace_estimate.h | 74 +++ .../providers/cuda/cuda_execution_provider.cc | 27 +- ...query_attention_workspace_estimate_test.cc | 536 ++++++++++++++++++ ...p_query_attention_workspace_header_test.cc | 9 + 10 files changed, 1703 insertions(+), 3 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc create mode 100644 onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h create mode 100644 onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc create mode 100644 onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h create mode 100644 onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc diff --git a/docs/annotated_partitioning/attention_workspace_estimation.md b/docs/annotated_partitioning/attention_workspace_estimation.md index d7120d989f616..586849448ed5b 100644 --- a/docs/annotated_partitioning/attention_workspace_estimation.md +++ b/docs/annotated_partitioning/attention_workspace_estimation.md @@ -207,6 +207,56 @@ 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 the same +root. 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..04d150a264b9c --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc @@ -0,0 +1,368 @@ +// 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 { + 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 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)) { + 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..aa0780b2dbb0a --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h @@ -0,0 +1,81 @@ +// 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. +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; + GQAXqaHeadSinkStorage xqa_head_sink_storage = GQAXqaHeadSinkStorage::None; + int64_t local_window_size = -1; +}; + +struct GQAWorkspaceAggregate { + GQAWorkspaceStatus status; + size_t total_workspace_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..7b93621f45b62 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc @@ -0,0 +1,507 @@ +// 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 "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) || + std::min((*cos)[1], (*sin)[1]) * 2 < head_bound) { + 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; + } + const int64_t capacity_bound = std::min((*past_key)[2], (*past_value)[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; + + 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) { + const auto config = ConfigFromNode(node); + return config.has_value() + ? EstimateGroupQueryAttentionWorkspace( + *config, input_shapes, device_prop, kernel_options) + : std::nullopt; +} + +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; + } +} + +} // 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..11621cfc1423b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h @@ -0,0 +1,74 @@ +// 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; + // A present sink may be prepacked at runtime, but Level 1 cannot prove that. + bool head_sink_is_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); + +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..60635f72454d6 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,27 @@ 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 ws = contrib::cuda::EstimateGroupQueryAttentionWorkspace( + *node, gsl::make_span(input_shapes), GetDeviceProp(), + *GetAttentionKernelOptions()); + 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..05f30c1fc31f1 --- /dev/null +++ b/onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc @@ -0,0 +1,536 @@ +// 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 "core/graph/graph.h" +#include "contrib_ops/cpu/bert/attention_common.h" +#include "contrib_ops/cuda/bert/group_query_attention_workspace_estimate.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::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 = 128; + 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", 128); + 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); +} + +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, 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, 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, RejectsNoncausalLocalWindow) { + AttentionKernelOptions options; + options.InitializeOnce(kMath, true); + auto config = Config(); + config.causal = 0; + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(), Device(), options) + .has_value()); + + config = Config(); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( + config, SeparateShapes(/*sequence=*/4, /*head=*/64, /*capacity=*/64), + 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, AggregatesExclusiveRoutesWithMax) { + auto xqa = Bounds(); + xqa.sequence_length_bound = 1; + xqa.prompt_reachable = false; + xqa.reachable_backends = GQAReachableBackend::Xqa; + 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()); + 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)); +} + +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, 19u); + EXPECT_EQ(level1.initialization_scratch_bytes, 23u); + + 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(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; From 1e6add3aa32bcbb8f52d4a493168e03b9738291d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:31:23 +0000 Subject: [PATCH 02/12] Harden GQA workspace bounds Avoid overflow while validating rotary cache dimensions and document and test the heuristic-independent Flash fast-decode envelope across the full dynamic sequence domain. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../group_query_attention_workspace_bounds.cc | 4 ++ ...roup_query_attention_workspace_estimate.cc | 7 ++- ...query_attention_workspace_estimate_test.cc | 57 +++++++++++++++++++ 3 files changed, 66 insertions(+), 2 deletions(-) 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 index 04d150a264b9c..b03927e08dd34 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc @@ -331,6 +331,10 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( } if (HasGQAReachableBackend(bounds.reachable_backends, GQAReachableBackend::FlashFastDecode)) { + // 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); } 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 index 7b93621f45b62..eedff22ad6aec 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc @@ -109,8 +109,11 @@ bool ValidateAuxiliaryShapes(const GQAWorkspaceEstimateConfig& config, 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) || - std::min((*cos)[1], (*sin)[1]) * 2 < head_bound) { + sin->NumDimensions() != 2 || !PositiveDims(*cos) || !PositiveDims(*sin)) { + return false; + } + const int64_t rotary_width = std::min((*cos)[1], (*sin)[1]); + if (rotary_width < head_bound / 2 + head_bound % 2) { return false; } } 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 index 05f30c1fc31f1..10a741350756f 100644 --- 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 @@ -222,6 +222,20 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsUndersizedPositionIds) { .has_value()); } +TEST(GroupQueryAttentionWorkspaceEstimateTest, AcceptsLargeRotaryCacheDimensionsWithoutOverflow) { + 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, kMaxDimension}); + shapes[8] = Known({kMaxDimension, kMaxDimension}); + EXPECT_TRUE(EstimateGroupQueryAttentionWorkspace( + config, shapes, Device(), options) + .has_value()); +} + TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsNoncausalLocalWindow) { AttentionKernelOptions options; options.InitializeOnce(kMath, true); @@ -445,6 +459,49 @@ TEST(GroupQueryAttentionWorkspaceBoundsTest, FlashEnvelopeDominatesDiscontinuous } } +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.sequence_length_bound = 1; From fd5a2fe635ff092c376e4f1ba86ed9c78c7e9852 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:41:57 +0000 Subject: [PATCH 03/12] Clarify GQA workspace bound contracts Document aliasing and non-windowed capacity preconditions, explain the graph-free route envelopes, and qualify Level 1 versus Level 2 head-sink sizing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../attention_workspace_estimation.md | 9 ++++++--- .../cuda/bert/group_query_attention_workspace_bounds.cc | 4 ++++ .../cuda/bert/group_query_attention_workspace_bounds.h | 4 ++++ 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/docs/annotated_partitioning/attention_workspace_estimation.md b/docs/annotated_partitioning/attention_workspace_estimation.md index 586849448ed5b..6c21b2e33528b 100644 --- a/docs/annotated_partitioning/attention_workspace_estimation.md +++ b/docs/annotated_partitioning/attention_workspace_estimation.md @@ -229,9 +229,12 @@ 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 the same -root. Neither adapter changes runtime `GetScratchBuffer()` calls or allocation -topology. +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. Level 2 can +omit that transient 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 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 index b03927e08dd34..cdfb0989c6cb1 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc @@ -151,6 +151,8 @@ GQAWorkspaceStatus FlashBackendEnvelope(const GQAWorkspaceProblem& problem, 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); @@ -331,6 +333,8 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( } 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 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 index aa0780b2dbb0a..c6731d2db22a8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h @@ -35,6 +35,10 @@ constexpr bool HasGQAReachableBackend(GQAReachableBackend mask, // 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; From ff30fd6400306b4ee1674cde05512cbe9c1b5861 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:23:01 +0000 Subject: [PATCH 04/12] Test GQA Level-2 workspace declaration Construct the CUDA GQA kernel and verify virtual workspace declaration, prepacked head-sink translation, and unavailable-shape behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...query_attention_workspace_estimate_test.cc | 155 ++++++++++++++++++ 1 file changed, 155 insertions(+) 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 index 10a741350756f..7f2c67ed4b7d9 100644 --- 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 @@ -6,13 +6,24 @@ #if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) #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 "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 { @@ -137,6 +148,93 @@ std::optional EstimateFromNode( 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", 128); + 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); @@ -567,6 +665,63 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, DeclaresOneAlignedSlotAndOnlyWork 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); + 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(); From a7271ad2a19f5a22c24713d0504bde710ea731c6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:03:47 +0000 Subject: [PATCH 05/12] Fix GQA partial rotary estimation Match runtime rotary-cache validation, add the missing cctype include, and gate the XQA kernel declaration test on SM80. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...roup_query_attention_workspace_estimate.cc | 6 ++++-- ...query_attention_workspace_estimate_test.cc | 20 ++++++++++++++++--- 2 files changed, 21 insertions(+), 5 deletions(-) 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 index eedff22ad6aec..a955dee638323 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc @@ -9,6 +9,7 @@ #include "contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h" #include +#include #include #include #include @@ -112,8 +113,9 @@ bool ValidateAuxiliaryShapes(const GQAWorkspaceEstimateConfig& config, sin->NumDimensions() != 2 || !PositiveDims(*cos) || !PositiveDims(*sin)) { return false; } - const int64_t rotary_width = std::min((*cos)[1], (*sin)[1]); - if (rotary_width < head_bound / 2 + head_bound % 2) { + const int64_t rotary_width = (*cos)[1]; + if (rotary_width != (*sin)[1] || rotary_width % 8 != 0 || + rotary_width > head_bound / 2) { return false; } } 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 index 7f2c67ed4b7d9..d417f93cd85f4 100644 --- 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 @@ -320,18 +320,29 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsUndersizedPositionIds) { .has_value()); } -TEST(GroupQueryAttentionWorkspaceEstimateTest, AcceptsLargeRotaryCacheDimensionsWithoutOverflow) { +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, kMaxDimension}); - shapes[8] = Known({kMaxDimension, kMaxDimension}); + 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) { @@ -679,6 +690,9 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, KernelDeclaresPrepackedHeadSinkRo 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()))); From 44ee6fd9d5cd9b6e1a83c921c7e95e7c1a9535d5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:16:11 +0000 Subject: [PATCH 06/12] Match GQA estimator cache geometry Require the bounded past cache capacity to match the sliding-window size, mirror the runtime past key/value capacity check, and make estimator fixtures use executable CUDA geometry. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...roup_query_attention_workspace_estimate.cc | 5 +-- ...query_attention_workspace_estimate_test.cc | 32 +++++++++++++------ 2 files changed, 26 insertions(+), 11 deletions(-) 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 index a955dee638323..54abb3cb69284 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc @@ -247,7 +247,8 @@ std::optional BuildBounds( if ((*past_key)[1] < config.kv_num_heads || (*past_value)[1] < config.kv_num_heads) { return std::nullopt; } - const int64_t capacity_bound = std::min((*past_key)[2], (*past_value)[2]); + 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; @@ -255,7 +256,7 @@ std::optional BuildBounds( } 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 || + capacity_bound != config.local_window_size || !ValidateAuxiliaryShapes(config, shapes, batch_bound, sequence_bound, head_bound, capacity_bound)) { return std::nullopt; 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 index d417f93cd85f4..2f66807da2ca2 100644 --- 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 @@ -84,7 +84,7 @@ GQAWorkspaceEstimateConfig Config() { config.cache_element_size = 2; config.num_heads = 8; config.kv_num_heads = 2; - config.local_window_size = 128; + config.local_window_size = 256; config.sliding_window_cache = true; return config; } @@ -138,7 +138,7 @@ std::optional EstimateFromNode( add_int_attribute("num_heads", 8); add_int_attribute("kv_num_heads", 2); add_int_attribute("causal", 1); - add_int_attribute("local_window_size", 128); + add_int_attribute("local_window_size", 256); add_int_attribute("sliding_window_cache", 1); const std::vector inputs{&query, &key, &value, &past_key}; @@ -194,7 +194,7 @@ std::string BuildGroupQueryAttentionKernelModel() { }; add_int_attribute("num_heads", 8); add_int_attribute("kv_num_heads", 2); - add_int_attribute("local_window_size", 128); + add_int_attribute("local_window_size", 256); add_int_attribute("sliding_window_cache", 1); constexpr int32_t kFloat16 = ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; @@ -268,6 +268,26 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedTotalKvAndAliasingAreU .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); @@ -353,12 +373,6 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsNoncausalLocalWindow) { EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( config, SeparateShapes(), Device(), options) .has_value()); - - config = Config(); - EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace( - config, SeparateShapes(/*sequence=*/4, /*head=*/64, /*capacity=*/64), - Device(), options) - .has_value()); } TEST(GroupQueryAttentionWorkspaceEstimateTest, ValidatesQuantizedCacheMetadata) { From 2924e21015ab998ff71490105a2c1215793c22c1 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:57:30 +0000 Subject: [PATCH 07/12] Test GQA Level-1 budget integration Drive CUDA GetCapability across thresholds that distinguish the structured GQA workspace estimate from fallback accounting and verify strict rejection at the estimated resource count. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...query_attention_workspace_estimate_test.cc | 84 +++++++++++++++++++ 1 file changed, 84 insertions(+) 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 index 2f66807da2ca2..c458cd720c8c8 100644 --- 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 @@ -5,6 +5,7 @@ #if !defined(USE_CUDA_MINIMAL) && !defined(DISABLE_CONTRIB_OPS) && !defined(BUILD_CUDA_EP_AS_PLUGIN) +#include #include #include #include @@ -256,6 +257,89 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, NodeAdapterParsesAttributesAndTyp 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_threshold) { + SessionOptions session_options; + session_options.graph_optimization_level = TransformerLevel::Default; + const std::string partitioning_settings = + std::to_string(memory_threshold) + ","; + 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 * 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()); + ASSERT_TRUE(estimate.has_value()); + } + + 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_LT(estimate->total_workspace_bytes, kFallbackWorkspaceBytes); + + { + const size_t threshold = + kAccountedTensorBytes + estimate->total_workspace_bytes + 1; + InferenceSessionWrapper session(make_session_options(threshold), + 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); + } + + { + const size_t threshold = + kAccountedTensorBytes + estimate->total_workspace_bytes; + InferenceSessionWrapper session(make_session_options(threshold), + 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); From 1379b86b7f481abf225f5cc937e250c3a474a738 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:05:17 +0000 Subject: [PATCH 08/12] Fix GQA budget test thresholds Declare the test dependencies directly and convert computed byte boundaries to the KiB units expected by capacity-aware partitioning. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...query_attention_workspace_estimate_test.cc | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) 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 index c458cd720c8c8..208357b6ce478 100644 --- 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 @@ -19,6 +19,7 @@ #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" @@ -267,11 +268,11 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti provider_info.sdpa_kernel = kMath; const std::string model_bytes = BuildGroupQueryAttentionKernelModel(); - auto make_session_options = [](size_t memory_threshold) { + 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_threshold) + ","; + std::to_string(memory_limit_kb) + ","; ORT_THROW_IF_ERROR(session_options.config_options.AddConfigEntry( kOrtSessionOptionsResourceCudaPartitioningSettings, partitioning_settings.c_str())); @@ -280,7 +281,7 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti std::optional estimate; { - InferenceSessionWrapper session(make_session_options(1024 * 1024), + InferenceSessionWrapper session(make_session_options(1024), GetEnvironment()); auto cuda_ep = std::make_shared(provider_info); if (cuda_ep->GetDeviceProp().major < 8) { @@ -307,12 +308,15 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti constexpr size_t kAccountedTensorBytes = kHeadSinkInitializerBytes + kOutputBytes + 2 * kPresentCacheBytes; constexpr size_t kFallbackWorkspaceBytes = kAccountedTensorBytes / 2; - ASSERT_LT(estimate->total_workspace_bytes, kFallbackWorkspaceBytes); + ASSERT_GT(estimate->total_workspace_bytes, kFallbackWorkspaceBytes); + const auto kilobytes_above = [](size_t bytes) { + return bytes / 1024 + 1; + }; { - const size_t threshold = - kAccountedTensorBytes + estimate->total_workspace_bytes + 1; - InferenceSessionWrapper session(make_session_options(threshold), + const size_t accepted_limit_kb = kilobytes_above( + kAccountedTensorBytes + estimate->total_workspace_bytes); + InferenceSessionWrapper session(make_session_options(accepted_limit_kb), GetEnvironment()); ASSERT_STATUS_OK(session.RegisterExecutionProvider( std::make_shared(provider_info))); @@ -325,9 +329,13 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti } { - const size_t threshold = - kAccountedTensorBytes + estimate->total_workspace_bytes; - InferenceSessionWrapper session(make_session_options(threshold), + // 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); + InferenceSessionWrapper session(make_session_options(rejected_limit_kb), GetEnvironment()); ASSERT_STATUS_OK(session.RegisterExecutionProvider( std::make_shared(provider_info))); From 72f42a06c924e696be74fa42a0fe1fb03a0eaa0e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:50:21 +0000 Subject: [PATCH 09/12] Account for prepacked GQA head sink Report the session-lived FP32 head-sink copy separately from route workspace, detect constant initializers during Level-1 estimation, and cover the case where another backend dominates the workspace maximum. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../group_query_attention_workspace_bounds.cc | 6 ++++++ .../group_query_attention_workspace_bounds.h | 1 + ...roup_query_attention_workspace_estimate.cc | 16 +++++++++------ ...group_query_attention_workspace_estimate.h | 3 ++- .../providers/cuda/cuda_execution_provider.cc | 7 ++++++- ...query_attention_workspace_estimate_test.cc | 20 +++++++++++++------ 6 files changed, 39 insertions(+), 14 deletions(-) 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 index cdfb0989c6cb1..f4d5f604b56ae 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc @@ -243,6 +243,12 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( "A reachable cuDNN GQA route has no graph-free workspace oracle."); return aggregate; } + if (bounds.xqa_head_sink_storage == GQAXqaHeadSinkStorage::PrepackedFp32) { + aggregate.status = Mul( + static_cast(bounds.num_heads), sizeof(float), + aggregate.persistent_prepack_bytes); + if (!aggregate.status.IsOK()) return aggregate; + } const auto retain = [&aggregate](GQAReachableBackend backend, GQAWorkspaceStatus status, 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 index c6731d2db22a8..9cee6c855db5d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h @@ -71,6 +71,7 @@ struct GQAWorkspaceBounds { struct GQAWorkspaceAggregate { GQAWorkspaceStatus status; size_t total_workspace_bytes = 0; + size_t persistent_prepack_bytes = 0; GQAReachableBackend sized_backends = GQAReachableBackend::None; }; 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 index 54abb3cb69284..ec101e80de7d1 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc @@ -481,12 +481,13 @@ std::optional EstimateGroupQueryAttentionWorkspace( const Node& node, gsl::span input_shapes, const cudaDeviceProp& device_prop, - const AttentionKernelOptions& kernel_options) { - const auto config = ConfigFromNode(node); - return config.has_value() - ? EstimateGroupQueryAttentionWorkspace( - *config, input_shapes, device_prop, kernel_options) - : std::nullopt; + const AttentionKernelOptions& kernel_options, + bool head_sink_is_constant_initializer) { + auto config = ConfigFromNode(node); + if (!config.has_value()) return std::nullopt; + config->head_sink_is_prepacked = head_sink_is_constant_initializer; + return EstimateGroupQueryAttentionWorkspace( + *config, input_shapes, device_prop, kernel_options); } void SetGroupQueryAttentionWorkspaceRequirements( @@ -504,6 +505,9 @@ void SetGroupQueryAttentionLevel1MemoryEstimate( 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; + } } } // namespace cuda 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 index 11621cfc1423b..5230ffbaa7c58 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h @@ -57,7 +57,8 @@ std::optional EstimateGroupQueryAttentionWorkspace( const Node& node, gsl::span input_shapes, const cudaDeviceProp& device_prop, - const AttentionKernelOptions& kernel_options); + const AttentionKernelOptions& kernel_options, + bool head_sink_is_constant_initializer = false); void SetGroupQueryAttentionWorkspaceRequirements( const GQAWorkspaceAggregate& estimate, diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 60635f72454d6..441966194fc1d 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -3621,9 +3621,14 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, 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()); + *GetAttentionKernelOptions(), head_sink_is_constant_initializer); if (ws.has_value()) { Level1MemoryEstimate estimate; contrib::cuda::SetGroupQueryAttentionLevel1MemoryEstimate(*ws, estimate); 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 index 208357b6ce478..5e3e8a1e8204c 100644 --- 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 @@ -47,6 +47,7 @@ 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; @@ -298,7 +299,8 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti shapes[11] = Known({8}); estimate = EstimateGroupQueryAttentionWorkspace( *node, gsl::make_span(shapes), cuda_ep->GetDeviceProp(), - *cuda_ep->GetAttentionKernelOptions()); + *cuda_ep->GetAttentionKernelOptions(), + /*head_sink_is_constant_initializer=*/true); ASSERT_TRUE(estimate.has_value()); } @@ -309,13 +311,15 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti 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)); const auto kilobytes_above = [](size_t bytes) { return bytes / 1024 + 1; }; { const size_t accepted_limit_kb = kilobytes_above( - kAccountedTensorBytes + estimate->total_workspace_bytes); + kAccountedTensorBytes + estimate->total_workspace_bytes + + estimate->persistent_prepack_bytes); InferenceSessionWrapper session(make_session_options(accepted_limit_kb), GetEnvironment()); ASSERT_STATUS_OK(session.RegisterExecutionProvider( @@ -334,7 +338,8 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti const size_t rejected_limit_kb = kilobytes_above( kAccountedTensorBytes + kFallbackWorkspaceBytes); ASSERT_LT(rejected_limit_kb * 1024, - kAccountedTensorBytes + estimate->total_workspace_bytes); + kAccountedTensorBytes + estimate->total_workspace_bytes + + estimate->persistent_prepack_bytes); InferenceSessionWrapper session(make_session_options(rejected_limit_kb), GetEnvironment()); ASSERT_STATUS_OK(session.RegisterExecutionProvider( @@ -719,15 +724,16 @@ TEST(GroupQueryAttentionWorkspaceBoundsTest, FlashFastDecodeEnvelopeCoversEveryS TEST(GroupQueryAttentionWorkspaceBoundsTest, AggregatesExclusiveRoutesWithMax) { auto xqa = Bounds(); - xqa.sequence_length_bound = 1; - xqa.prompt_reachable = false; xqa.reachable_backends = GQAReachableBackend::Xqa; + xqa.xqa_head_sink_storage = GQAXqaHeadSinkStorage::PrepackedFp32; 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; @@ -736,6 +742,8 @@ TEST(GroupQueryAttentionWorkspaceBoundsTest, AggregatesExclusiveRoutesWithMax) { 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)); } TEST(GroupQueryAttentionWorkspaceBoundsTest, CudnnReachabilityIsUnavailable) { @@ -767,7 +775,7 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, DeclaresOneAlignedSlotAndOnlyWork 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, 19u); + EXPECT_EQ(level1.persistent_prepack_bytes, 0u); EXPECT_EQ(level1.initialization_scratch_bytes, 23u); auto unavailable = *estimate; From 07a61dfa539cd7040c8ea895ab2cb94f3329569b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:24:41 +0000 Subject: [PATCH 10/12] Conservatively account for disabled GQA prepacking Keep Level-1 dynamic head-sink conversion workspace even for constant initializers while separately charging the possible session-lived FP32 prepack copy. Preserve exact constructed-kernel behavior for Level 2. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/group_query_attention.cc | 1 + .../group_query_attention_workspace_bounds.cc | 2 +- .../bert/group_query_attention_workspace_bounds.h | 1 + .../group_query_attention_workspace_estimate.cc | 4 +++- .../group_query_attention_workspace_estimate.h | 2 +- ...oup_query_attention_workspace_estimate_test.cc | 15 ++++++++++++++- 6 files changed, 21 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index b14d8509a15dd..70af22eaf14df 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -239,6 +239,7 @@ Status GroupQueryAttention::DeclareWorkspaceRequirements( config.enable_xqa = enable_xqa_; config.disable_flash_decode = disable_flash_decode_; config.head_sink_is_prepacked = xqa_head_sink_count_ == num_heads_; + config.head_sink_may_be_prepacked = config.head_sink_is_prepacked; const auto estimate = EstimateGroupQueryAttentionWorkspace( config, input_shapes, GetDeviceProp(), *kernel_options_); 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 index f4d5f604b56ae..ec5a878529ed5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc @@ -243,7 +243,7 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( "A reachable cuDNN GQA route has no graph-free workspace oracle."); return aggregate; } - if (bounds.xqa_head_sink_storage == GQAXqaHeadSinkStorage::PrepackedFp32) { + if (bounds.head_sink_may_be_prepacked) { aggregate.status = Mul( static_cast(bounds.num_heads), sizeof(float), aggregate.persistent_prepack_bytes); 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 index 9cee6c855db5d..a379ea4ef8802 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h @@ -65,6 +65,7 @@ struct GQAWorkspaceBounds { bool is_bf16 = false; GQAXqaKvType xqa_kv_type = GQAXqaKvType::None; GQAXqaHeadSinkStorage xqa_head_sink_storage = GQAXqaHeadSinkStorage::None; + bool head_sink_may_be_prepacked = false; int64_t local_window_size = -1; }; 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 index ec101e80de7d1..84c37ed04f184 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc @@ -310,6 +310,8 @@ std::optional BuildBounds( ? (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); @@ -485,7 +487,7 @@ std::optional EstimateGroupQueryAttentionWorkspace( bool head_sink_is_constant_initializer) { auto config = ConfigFromNode(node); if (!config.has_value()) return std::nullopt; - config->head_sink_is_prepacked = head_sink_is_constant_initializer; + config->head_sink_may_be_prepacked = head_sink_is_constant_initializer; return EstimateGroupQueryAttentionWorkspace( *config, input_shapes, device_prop, kernel_options); } 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 index 5230ffbaa7c58..15f6b39108015 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h @@ -43,8 +43,8 @@ struct GQAWorkspaceEstimateConfig { bool cache_is_fp8 = false; bool enable_xqa = true; bool disable_flash_decode = false; - // A present sink may be prepacked at runtime, but Level 1 cannot prove that. bool head_sink_is_prepacked = false; + bool head_sink_may_be_prepacked = false; }; std::optional EstimateGroupQueryAttentionWorkspace( 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 index 5e3e8a1e8204c..6e1af69b8a0e3 100644 --- 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 @@ -302,6 +302,17 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti *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; + prepacked_config.head_sink_may_be_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); } constexpr size_t kHeadSinkInitializerBytes = 8 * sizeof(MLFloat16); @@ -725,7 +736,8 @@ TEST(GroupQueryAttentionWorkspaceBoundsTest, FlashFastDecodeEnvelopeCoversEveryS TEST(GroupQueryAttentionWorkspaceBoundsTest, AggregatesExclusiveRoutesWithMax) { auto xqa = Bounds(); xqa.reachable_backends = GQAReachableBackend::Xqa; - xqa.xqa_head_sink_storage = GQAXqaHeadSinkStorage::PrepackedFp32; + 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; @@ -823,6 +835,7 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, KernelDeclaresPrepackedHeadSinkRo shapes[11] = Known({8}); auto expected_config = Config(); expected_config.head_sink_is_prepacked = true; + expected_config.head_sink_may_be_prepacked = true; const auto expected = EstimateGroupQueryAttentionWorkspace( expected_config, gsl::make_span(shapes), cuda_ep->GetDeviceProp(), *cuda_ep->GetAttentionKernelOptions()); From f5716ad40c2cb01b0dad9814ec2bb5eb9ce54b1d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:31:40 +0000 Subject: [PATCH 11/12] Harden GQA prepack lifetime accounting Make exact prepack imply persistent accounting, include conservative initialization staging, document the lifetime split, and add adversarial coverage for inconsistent caller facts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../attention_workspace_estimation.md | 9 +++++--- .../cuda/bert/group_query_attention.cc | 1 - .../group_query_attention_workspace_bounds.cc | 9 +++++++- .../group_query_attention_workspace_bounds.h | 3 +++ ...roup_query_attention_workspace_estimate.cc | 2 ++ ...group_query_attention_workspace_estimate.h | 2 ++ ...query_attention_workspace_estimate_test.cc | 23 ++++++++++++++++--- 7 files changed, 41 insertions(+), 8 deletions(-) diff --git a/docs/annotated_partitioning/attention_workspace_estimation.md b/docs/annotated_partitioning/attention_workspace_estimation.md index 6c21b2e33528b..04899befead52 100644 --- a/docs/annotated_partitioning/attention_workspace_estimation.md +++ b/docs/annotated_partitioning/attention_workspace_estimation.md @@ -231,9 +231,12 @@ 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. Level 2 can -omit that transient region when the constructed kernel has a prepacked head -sink, so its root can be smaller. Neither adapter changes runtime +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 diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 70af22eaf14df..b14d8509a15dd 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -239,7 +239,6 @@ Status GroupQueryAttention::DeclareWorkspaceRequirements( config.enable_xqa = enable_xqa_; config.disable_flash_decode = disable_flash_decode_; config.head_sink_is_prepacked = xqa_head_sink_count_ == num_heads_; - config.head_sink_may_be_prepacked = config.head_sink_is_prepacked; const auto estimate = EstimateGroupQueryAttentionWorkspace( config, input_shapes, GetDeviceProp(), *kernel_options_); 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 index ec5a878529ed5..66b74ec2f19df 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc @@ -243,11 +243,18 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( "A reachable cuDNN GQA route has no graph-free workspace oracle."); return aggregate; } - if (bounds.head_sink_may_be_prepacked) { + 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, 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 index a379ea4ef8802..0783f90172a0c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h @@ -64,6 +64,8 @@ struct GQAWorkspaceBounds { 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; @@ -73,6 +75,7 @@ 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; }; 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 index 84c37ed04f184..e11b666169c86 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc @@ -509,6 +509,8 @@ void SetGroupQueryAttentionLevel1MemoryEstimate( } if (workspace.status.IsOK()) { estimate.persistent_prepack_bytes = workspace.persistent_prepack_bytes; + estimate.initialization_scratch_bytes = + workspace.initialization_scratch_bytes; } } 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 index 15f6b39108015..02c6acdf308ac 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h @@ -43,7 +43,9 @@ struct GQAWorkspaceEstimateConfig { 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; }; 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 index 6e1af69b8a0e3..9020a180d82a2 100644 --- 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 @@ -305,7 +305,6 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti auto prepacked_config = Config(); prepacked_config.head_sink_is_prepacked = true; - prepacked_config.head_sink_may_be_prepacked = true; const auto prepacked = EstimateGroupQueryAttentionWorkspace( prepacked_config, gsl::make_span(shapes), cuda_ep->GetDeviceProp(), *cuda_ep->GetAttentionKernelOptions()); @@ -313,6 +312,8 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti 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); @@ -323,6 +324,8 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti 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; }; @@ -756,6 +759,21 @@ TEST(GroupQueryAttentionWorkspaceBoundsTest, AggregatesExclusiveRoutesWithMax) { 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) { @@ -788,7 +806,7 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, DeclaresOneAlignedSlotAndOnlyWork 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, 23u); + EXPECT_EQ(level1.initialization_scratch_bytes, 0u); auto unavailable = *estimate; unavailable.status.error = contrib::cuda::GQAWorkspaceError::Unavailable; @@ -835,7 +853,6 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, KernelDeclaresPrepackedHeadSinkRo shapes[11] = Known({8}); auto expected_config = Config(); expected_config.head_sink_is_prepacked = true; - expected_config.head_sink_may_be_prepacked = true; const auto expected = EstimateGroupQueryAttentionWorkspace( expected_config, gsl::make_span(shapes), cuda_ep->GetDeviceProp(), *cuda_ep->GetAttentionKernelOptions()); From 9bd7af2d432584ff07f01b955b1d7e5d7f1b6b8a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:16:19 +0000 Subject: [PATCH 12/12] Test GQA Level-1 lifetime propagation Add direct nonzero assertions for persistent prepack and initialization scratch fields in the Level-1 setter. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../group_query_attention_workspace_estimate_test.cc | 7 +++++++ 1 file changed, 7 insertions(+) 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 index 9020a180d82a2..f1efe514cf322 100644 --- 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 @@ -808,6 +808,13 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, DeclaresOneAlignedSlotAndOnlyWork 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);