Enable bounded non-windowed GQA workspace estimation

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

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)

onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include <algorithm>
+#include <charconv>
 #include <cmath>
 #include <cstdint>
 #include <limits>
@@ -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<T, U>::GroupQueryAttention(const OpKernelInfo& info)
   k_quant_type_ = StringToKVQuantizationType(info.GetAttrOrDefault<std::string>("k_quant_type", "NONE"));
   v_quant_type_ = StringToKVQuantizationType(info.GetAttrOrDefault<std::string>("v_quant_type", "NONE"));
   kv_cache_bit_width_ = static_cast<int>(info.GetAttrOrDefault<int64_t>("kv_cache_bit_width", 0));
+  max_total_sequence_length_ =
+      ParsePositiveInt64Config(info, kOrtSessionOptionsCudaGqaWorkspaceMaxTotalSequenceLength);
 
   constexpr bool kIsFp16OrBf16 = std::is_same_v<T, MLFloat16> || std::is_same_v<T, BFloat16>;
   // XQA defaults on for fp16/bf16; ORT_ENABLE_XQA=0 disables it explicitly.
@@ -218,6 +235,7 @@ Status GroupQueryAttention<T, U>::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_;

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<int> zeros_;

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<int64_t, 2> 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<GQAWorkspaceStatus, size_t>{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);

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;

onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc

@@ -198,12 +198,12 @@ std::optional<GQAWorkspaceBounds> 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,15 +248,19 @@ std::optional<GQAWorkspaceBounds> 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<int64_t>::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 ||
+      (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<GQAWorkspaceBounds> 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<GQAWorkspaceBounds> 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;

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;