Add GQA workspace estimation

docs/annotated_partitioning/attention_workspace_estimation.md

@@ -207,6 +207,62 @@ alias fact and the past-cache capacity, and the preparation recipe includes the
 checked preservation region. Windowed execution, XQA, and Flash fast decode
 require both K/V pairs to alias and reject this state.
 
+### GroupQueryAttention bounded estimator and framework adapters
+
+The GQA framework follow-up translates positional inputs and immutable kernel
+attributes into a graph-free bounded domain. It covers packed and separate
+Q/K/V, optional past-cache pairs, cache quantization, prompt and decode states,
+windowed staging and compaction, XQA (including unresolved shared-memory
+fallback), regular and fast-decode Flash, MEA, unfused, and cuDNN reachability.
+Head geometry derived from `WorkspaceInputShape` is an upper bound: route
+reachability searches supported smaller head sizes instead of testing only the
+componentwise maximum.
+
+Flash split selection is not monotonic in KV length. The bounded estimator uses
+the runtime block size (`256` for head size at most 64, `128` through 128, and
+`64` above 128) and bounds the selected split count by:
+
+```text
+min(128, SM count, ceil(KV length bound / block size))
+```
+
+It then sizes the split accumulators at that envelope. Complete routes are
+mutually exclusive and are aggregated with `max`, not sum. The resulting
+nonzero estimate is one operator-owned slot-0 root with 256-byte alignment.
+Level 1 reports it as `runtime_workspace_bytes`; Level 2 declares one root from
+the same estimator. Level 1 conservatively includes dynamic head-sink conversion
+because prepack state is unavailable during capability analysis. For a constant
+head sink, it separately charges the possible session-lived FP32 copy and the
+initialization-only device staging copy, even when session configuration later
+disables prepacking. Level 2 can omit the transient conversion region when the
+constructed kernel has a prepacked head sink, so its root can be smaller.
+Neither adapter changes runtime
+`GetScratchBuffer()` calls or allocation topology.
+
+The CPU `total_sequence_length` scalar and past/present aliasing are not
+available through `WorkspaceInputShape`. For non-windowed GQA, the scalar can
+exceed the past-cache sequence dimension and directly scale Flash, MEA, and
+unfused workspace. Non-windowed execution can also preserve one full past
+tensor when exactly one past/present pair aliases. The current adapter therefore
+reports non-windowed GQA as unavailable rather than treating the past shape as
+a total-KV bound or omitting alias-preservation scratch.
+
+Successful estimates are currently limited to sliding-window cache nodes.
+Their final cache capacity remains `C`, while a multi-token step uses a
+transient staged/effective attention extent of `C + S`; a single-token step
+uses `C`. This shape-only bound remains sound even when the scalar
+`total_sequence_length` is much larger. Runtime requires both past/present
+pairs to alias before staging or compaction, excluding the partial-alias
+preservation path. Sliding-window attention bias remains
+unsupported by runtime and is also unavailable to the estimator. If cuDNN can
+be reached anywhere in a graph-free bounded domain, aggregation remains
+unavailable because cuDNN's allocator-based workspace has no sound graph-free
+oracle. Level 1 then uses the generic fallback and Level 2 emits no requirement.
+
+Unlike the log-only PA/PMHA Level-1 probe, a successful GQA Level-1 estimate is
+passed to the #31962 resource accountant as `runtime_workspace_bytes` and can
+affect CUDA partition acceptance.
+
 MHA and GQA are high-value coverage targets and have high estimation-drift risk. Their runtime behavior can include
 dynamic internal backend dispatch, cache lifecycle and aliasing, optional inputs, non-monotonic fallback paths, and
 unfused workspace governed by `S_q * S_kv_total`. GQA additionally has different Q and KV head counts. MHA can have

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<T, U>::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 <typename T, typename U>
+Status GroupQueryAttention<T, U>::DeclareWorkspaceRequirements(
+    gsl::span<const WorkspaceInputShape> input_shapes,
+    InlinedVector<WorkspaceRequirement>& 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<T, BFloat16>;
+  config.cache_is_fp8 = std::is_same_v<U, Float8E4M3FN>;
+  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:

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<const WorkspaceInputShape> input_shapes,
+      InlinedVector<WorkspaceRequirement>& 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