diff --git a/docs/annotated_partitioning/attention_workspace_estimation.md b/docs/annotated_partitioning/attention_workspace_estimation.md index 04899befead52..977d726edad47 100644 --- a/docs/annotated_partitioning/attention_workspace_estimation.md +++ b/docs/annotated_partitioning/attention_workspace_estimation.md @@ -240,20 +240,27 @@ 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 +available through `WorkspaceInputShape`. Non-windowed Level-2 estimation is +enabled when the caller supplies a positive +`ep.cuda.gqa_workspace_max_total_sequence_length` session option. The estimator +uses the larger of that value and the past-cache capacity as the present-cache +and backend KV-length bound. Callers must provide a sound scalar bound whenever +the runtime value may exceed the past-cache capacity. + +`MayInplace(3, 1)` and `MayInplace(4, 2)` permit but do not guarantee cache +aliasing. For non-windowed GQA, the aggregate therefore also evaluates the +valid one-sided-alias case, which preserves one complete past tensor before +preprocessing overwrites the aliased present buffer. XQA and Flash fast decode +are excluded from that case because their runtime routes require both +past/present pairs to alias. The Level-1 node adapter remains unavailable for +non-windowed inputs because it cannot access the session option. + +For sliding-window nodes, the 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 diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 079aaa5d80aca..6fc9b724ba9fc 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -509,6 +509,12 @@ static const char* const kOrtSessionOptionsMaxShapeOverride = "session.max_shape static const char* const kOrtSessionOptionsStrictWorkspaceVerification = "session.strict_workspace_verification"; +/// Optional positive upper bound for the total_sequence_length scalar of non-windowed CUDA +/// GroupQueryAttention nodes. The scalar value is unavailable during workspace declaration, +/// so callers must provide a sound bound when it may exceed the past-cache capacity. +static const char* const kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength = + "ep.cuda.gqa_workspace_max_total_sequence_length"; + // Enable EP context feature to dump the partitioned graph which includes the EP context into Onnx file. // The dumped Onnx model with EP context can be used for future inference to avoid the EP graph partitioning/compile overhead. // "0": disable. (default) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index b14d8509a15dd..c04cc1e234084 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include +#include #include #include #include @@ -10,6 +11,7 @@ #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" #include "core/platform/env_var_utils.h" +#include "core/session/onnxruntime_session_options_config_keys.h" #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" @@ -50,6 +52,19 @@ KVQuantizationType StringToKVQuantizationType(std::string s) { } return KVQuantizationType::NONE; } + +int64_t ParsePositiveInt64Config(const OpKernelInfo& info, const char* key) { + const std::string value = info.GetConfigOptions().GetConfigOrDefault(key, ""); + if (value.empty()) { + return 0; + } + + int64_t parsed = 0; + const auto [end, error] = std::from_chars(value.data(), value.data() + value.size(), parsed); + ORT_ENFORCE(error == std::errc{} && end == value.data() + value.size() && parsed > 0, + key, " must be a positive int64 value, but got: ", value); + return parsed; +} } // namespace #define REGISTER_KERNEL_TYPED(T, U) \ @@ -131,6 +146,8 @@ GroupQueryAttention::GroupQueryAttention(const OpKernelInfo& info) k_quant_type_ = StringToKVQuantizationType(info.GetAttrOrDefault("k_quant_type", "NONE")); v_quant_type_ = StringToKVQuantizationType(info.GetAttrOrDefault("v_quant_type", "NONE")); kv_cache_bit_width_ = static_cast(info.GetAttrOrDefault("kv_cache_bit_width", 0)); + max_total_sequence_length_ = + ParsePositiveInt64Config(info, kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength); constexpr bool kIsFp16OrBf16 = std::is_same_v || std::is_same_v; // XQA defaults on for fp16/bf16; ORT_ENABLE_XQA=0 disables it explicitly. @@ -218,6 +235,7 @@ Status GroupQueryAttention::DeclareWorkspaceRequirements( config.causal = is_unidirectional_ ? 1 : 0; config.local_window_size = local_window_size_; config.sliding_window_cache = sliding_window_cache_; + config.max_total_sequence_length = max_total_sequence_length_; config.do_rotary = do_rotary_; config.smooth_softmax = use_smooth_softmax_; config.softcap = softcap_; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h index acaba2593db44..5bbf296cf87d6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h @@ -56,6 +56,7 @@ class GroupQueryAttention final : public CudaKernel { KVQuantizationType k_quant_type_; KVQuantizationType v_quant_type_; int kv_cache_bit_width_; + int64_t max_total_sequence_length_ = 0; static constexpr int kZerosCount = 256; // In prompt case we create a zero buffer of size 256 for seqlen (assume batch_size <= 256) IAllocatorUniquePtr zeros_; 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 66b74ec2f19df..189b5b8d3fc48 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 @@ -76,7 +76,8 @@ GQAWorkspaceStatus EffectiveKvLengthBound(const GQAWorkspaceBounds& bounds, GQAWorkspaceProblem MakeProblem(const GQAWorkspaceBounds& bounds, int64_t sequence_length, int64_t head_size, - bool first_prompt) noexcept { + bool first_prompt, + bool requires_separate_past_buffer = false) noexcept { GQAWorkspaceProblem problem; problem.qkv_element_size = bounds.qkv_element_size; problem.cache_element_size = bounds.cache_element_size; @@ -86,6 +87,8 @@ GQAWorkspaceProblem MakeProblem(const GQAWorkspaceBounds& bounds, 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.past_kv_cache_capacity = + requires_separate_past_buffer ? bounds.past_kv_cache_capacity_bound : 0; problem.kv_cache_bit_width = bounds.kv_cache_bit_width; problem.k_quantization = bounds.k_quantization; problem.v_quantization = bounds.v_quantization; @@ -94,6 +97,7 @@ GQAWorkspaceProblem MakeProblem(const GQAWorkspaceBounds& bounds, problem.do_rotary = bounds.do_rotary; problem.is_packed_qkv = bounds.is_packed_qkv; problem.use_qk_norm = bounds.use_qk_norm; + problem.requires_separate_past_buffer = requires_separate_past_buffer; return problem; } @@ -233,6 +237,7 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( 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.partial_alias_reachable && bounds.past_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."); @@ -272,7 +277,8 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( 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 sequence_length, bool first_prompt, + bool requires_separate_past_buffer) { int64_t effective_kv_length = 0; const auto effective_status = EffectiveKvLengthBound(bounds, sequence_length, effective_kv_length); @@ -288,7 +294,9 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( : GQAPreprocessMode::Unfused; route.unfused.total_sequence_length = effective_kv_length; return GetGQACompleteWorkspaceRecipe( - MakeProblem(bounds, sequence_length, head_size, first_prompt), route); + MakeProblem(bounds, sequence_length, head_size, first_prompt, + requires_separate_past_buffer), + route); }; const int64_t general_head = LargestMultipleOfEight( @@ -304,24 +312,29 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( 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); + for (bool partial_alias : {false, true}) { + if (partial_alias && !bounds.partial_alias_reachable) continue; + const auto result = size_complete( + backend == GQAReachableBackend::MemoryEfficient + ? GQABackend::MemoryEfficient + : GQABackend::Unfused, + head, sequence, first, partial_alias); + retain(backend, result.status, result.recipe.total_workspace_bytes); + } } } } - const auto size_flash = [&](bool fast, int64_t sequence, bool first) { + const auto size_flash = [&](bool fast, int64_t sequence, bool first, + bool requires_separate_past_buffer) { 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); + const auto problem = MakeProblem( + bounds, sequence, head, first, requires_separate_past_buffer); GQAPreparationRoute prep_route; prep_route.preprocess_mode = GQAPreprocessMode::Flash; prep_route.use_flash_attention_fast_decode = fast; @@ -339,8 +352,11 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( 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); + for (bool partial_alias : {false, true}) { + if (partial_alias && !bounds.partial_alias_reachable) continue; + const auto result = size_flash(false, sequence, first, partial_alias); + retain(GQAReachableBackend::Flash, result.first, result.second); + } } } } @@ -352,7 +368,9 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( // 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); + const auto result = size_flash( + true, bounds.sequence_length_bound, false, + /*requires_separate_past_buffer=*/false); retain(GQAReachableBackend::FlashFastDecode, result.first, result.second); } @@ -361,7 +379,8 @@ GQAWorkspaceAggregate GetGQAWorkspaceAggregateForBounds( 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); + const auto problem = MakeProblem( + bounds, 1, head, false, /*requires_separate_past_buffer=*/false); GQAPreparationRoute route; route.preprocess_mode = GQAPreprocessMode::Xqa; const auto prep = GetGQAPreparationRecipe(problem, route); 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 0783f90172a0c..3e8099a74a535 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,9 +35,7 @@ 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 +// 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; @@ -48,6 +46,7 @@ struct GQAWorkspaceBounds { int64_t kv_num_heads = 0; int64_t head_size_bound = 0; int64_t present_kv_cache_capacity_bound = 0; + int64_t past_kv_cache_capacity_bound = 0; int64_t kv_cache_bit_width = 0; GQAKvQuantizationType k_quantization = GQAKvQuantizationType::None; GQAKvQuantizationType v_quantization = GQAKvQuantizationType::None; @@ -57,6 +56,8 @@ struct GQAWorkspaceBounds { bool use_qk_norm = false; bool prompt_reachable = false; bool decode_reachable = false; + // Non-windowed execution may bind exactly one past/present pair in-place. + bool partial_alias_reachable = false; GQAReachableBackend reachable_backends = GQAReachableBackend::None; int64_t device_major = 0; 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 e11b666169c86..3898b702452c0 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 @@ -198,12 +198,12 @@ std::optional BuildBounds( } // 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; + // its value. Non-windowed estimation therefore requires an explicit bound. + // Aliasing is also unknown, so the aggregate includes the valid one-sided + // past/present alias case. + if (!config.sliding_window_cache && config.max_total_sequence_length <= 0) { + return std::nullopt; + } const bool packed = !Present(shapes, kKey); const TensorShape* query = Shape(shapes, kQuery); @@ -248,7 +248,11 @@ std::optional BuildBounds( return std::nullopt; } if ((*past_key)[2] != (*past_value)[2]) return std::nullopt; - const int64_t capacity_bound = (*past_key)[2]; + const int64_t past_capacity_bound = (*past_key)[2]; + const int64_t capacity_bound = + config.sliding_window_cache + ? past_capacity_bound + : std::max(past_capacity_bound, config.max_total_sequence_length); 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; @@ -256,7 +260,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 || + (config.sliding_window_cache && capacity_bound != config.local_window_size) || !ValidateAuxiliaryShapes(config, shapes, batch_bound, sequence_bound, head_bound, capacity_bound)) { return std::nullopt; @@ -291,6 +295,7 @@ std::optional BuildBounds( bounds.kv_num_heads = config.kv_num_heads; bounds.head_size_bound = head_bound; bounds.present_kv_cache_capacity_bound = capacity_bound; + bounds.past_kv_cache_capacity_bound = past_capacity_bound; bounds.kv_cache_bit_width = config.kv_cache_bit_width; bounds.k_quantization = config.k_quantization; bounds.v_quantization = config.v_quantization; @@ -300,6 +305,7 @@ std::optional BuildBounds( bounds.use_qk_norm = Present(shapes, kQNorm); bounds.prompt_reachable = true; bounds.decode_reachable = true; + bounds.partial_alias_reachable = !config.sliding_window_cache; bounds.device_major = device_prop.major; bounds.device_minor = device_prop.minor; bounds.multi_processor_count = device_prop.multiProcessorCount; 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 02c6acdf308ac..870792b9e9a75 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 @@ -33,6 +33,9 @@ struct GQAWorkspaceEstimateConfig { int64_t causal = 1; int64_t local_window_size = -1; bool sliding_window_cache = false; + // Required for non-windowed estimation because WorkspaceInputShape does not + // carry the total_sequence_length scalar value. + int64_t max_total_sequence_length = 0; bool do_rotary = false; bool smooth_softmax = false; float softcap = 0.0f; 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 95d3ef0304477..9fe4103853e94 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 @@ -164,7 +164,7 @@ void SetValueInfo(ONNX_NAMESPACE::ValueInfoProto& value_info, } } -std::string BuildGroupQueryAttentionKernelModel() { +std::string BuildGroupQueryAttentionKernelModel(bool sliding_window_cache = true) { ONNX_NAMESPACE::ModelProto model; model.set_ir_version(ONNX_NAMESPACE::IR_VERSION); auto* onnx_opset = model.add_opset_import(); @@ -197,8 +197,10 @@ std::string BuildGroupQueryAttentionKernelModel() { }; add_int_attribute("num_heads", 8); add_int_attribute("kv_num_heads", 2); - add_int_attribute("local_window_size", 256); - add_int_attribute("sliding_window_cache", 1); + if (sliding_window_cache) { + add_int_attribute("local_window_size", 256); + add_int_attribute("sliding_window_cache", 1); + } constexpr int32_t kFloat16 = ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; constexpr int32_t kInt32 = ONNX_NAMESPACE::TensorProto_DataType_INT32; @@ -367,16 +369,34 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, GetCapabilityBudgetUsesLevel1Esti } } -TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedTotalKvAndAliasingAreUnavailable) { +TEST(GroupQueryAttentionWorkspaceEstimateTest, NonWindowedRequiresExplicitTotalKvBound) { 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()); + config.enable_xqa = false; + const auto shapes = SeparateShapes(/*sequence=*/4, /*head=*/64, /*capacity=*/128); + EXPECT_FALSE(EstimateGroupQueryAttentionWorkspace(config, shapes, Device(), options).has_value()); + + config.max_total_sequence_length = 64; + const auto past_capacity_bound = + EstimateGroupQueryAttentionWorkspace(config, shapes, Device(), options); + ASSERT_TRUE(past_capacity_bound.has_value()); + + config.max_total_sequence_length = 128; + const auto matching_bound = + EstimateGroupQueryAttentionWorkspace(config, shapes, Device(), options); + ASSERT_TRUE(matching_bound.has_value()); + EXPECT_EQ(matching_bound->total_workspace_bytes, + past_capacity_bound->total_workspace_bytes); + + config.max_total_sequence_length = 256; + const auto larger_bound = + EstimateGroupQueryAttentionWorkspace(config, shapes, Device(), options); + ASSERT_TRUE(larger_bound.has_value()); + EXPECT_GT(larger_bound->total_workspace_bytes, + matching_bound->total_workspace_bytes); } TEST(GroupQueryAttentionWorkspaceEstimateTest, RejectsCacheCapacityDifferentFromWindow) { @@ -552,6 +572,31 @@ TEST(GroupQueryAttentionWorkspaceBoundsTest, PromptDecodeAndWindowRoutesAreBound EXPECT_TRUE(prompt.status.IsOK()); } +TEST(GroupQueryAttentionWorkspaceBoundsTest, PartialAliasAddsPastPreservationBuffer) { + auto bounds = Bounds(); + bounds.reachable_backends = GQAReachableBackend::Unfused; + bounds.present_kv_cache_capacity_bound = 256; + bounds.past_kv_cache_capacity_bound = 128; + const auto shared_or_separate = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(shared_or_separate.status.IsOK()) << shared_or_separate.status.message; + + bounds.partial_alias_reachable = true; + const auto partial_alias = GetGQAWorkspaceAggregateForBounds(bounds); + ASSERT_TRUE(partial_alias.status.IsOK()) << partial_alias.status.message; + ASSERT_GE(partial_alias.total_workspace_bytes, + shared_or_separate.total_workspace_bytes); + const size_t past_tensor_bytes = + static_cast(bounds.batch_size_bound * bounds.kv_num_heads * + bounds.past_kv_cache_capacity_bound * bounds.head_size_bound) * + bounds.cache_element_size; + EXPECT_GE(partial_alias.total_workspace_bytes - + shared_or_separate.total_workspace_bytes, + past_tensor_bytes); + + bounds.past_kv_cache_capacity_bound = 0; + EXPECT_FALSE(GetGQAWorkspaceAggregateForBounds(bounds).status.IsOK()); +} + TEST(GroupQueryAttentionWorkspaceBoundsTest, WindowedBoundsUseFinalCapacityAndTransientExtent) { constexpr int64_t kCapacity = 256; constexpr int64_t kSequence = 4; @@ -887,6 +932,54 @@ TEST(GroupQueryAttentionWorkspaceEstimateTest, KernelDeclaresPrepackedHeadSinkRo EXPECT_TRUE(requirements.empty()); } +TEST(GroupQueryAttentionWorkspaceEstimateTest, KernelDeclaresBoundedNonWindowedRoot) { + if (!HasCudaDevice()) { + GTEST_SKIP() << "A CUDA device is required to construct the CUDA kernel."; + } + + ScopedEnvironmentVariables scoped_env_vars{{{"ORT_ENABLE_XQA", "0"}}}; + SessionOptions session_options; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength, "512")); + 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(/*sliding_window_cache=*/false); + 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); + 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.sliding_window_cache = false; + expected_config.local_window_size = -1; + expected_config.max_total_sequence_length = 512; + expected_config.enable_xqa = false; + 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()); + + 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); +} + TEST(GroupQueryAttentionWorkspaceBoundsTest, CheckedOverflowIsUnavailable) { auto bounds = Bounds(); bounds.batch_size_bound = std::numeric_limits::max();