Enable bounded non-windowed GQA workspace estimation

[Chi Lo (chilo-ms)]

Description

Follow-up to #32617 for non-windowed CUDA GroupQueryAttention workspace estimation.

• add ep.cuda.gqa_workspace_max_total_sequence_length as an explicit positive bound for the runtime total_sequence_length scalar used by Level-2 declaration;
• size non-windowed present-cache/backend workspace from the larger of the scalar bound and past-cache capacity;
• model MayInplace conservatively by including the valid one-sided past/present alias case and its full past-tensor preservation buffer; and
• keep XQA and Flash fast decode out of the partial-alias case because those runtime routes require both cache pairs to alias.

The Level-1 node adapter remains unavailable for non-windowed inputs because it cannot access session configuration. This PR changes estimation and declaration only; it does not consume a planned workspace root or change runtime allocation topology.

Validation

• Release CUDA provider target built successfully on RTX 5090 / CUDA 13.3.
• Changed-file lintrunner and git diff --check passed.
• Added direct estimator, aggregate, and kernel-declaration coverage in group_query_attention_workspace_estimate_test.cc.
• The updated CUDA-internal test source compiled locally. Running the module was blocked by the existing Windows onnxruntime_providers_cuda_ut CMake dependency cycle/module-loading setup when internal tests are enabled.

Tracks #29775.

[Copilot]

Copilot review overview

🟡 Changes recommended

The new kernel-declaration test disables every backend route compatible with its model and will fail.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 1

Open (2)

• Test expectation is unreachable with XQA disabled · New
• Clarify option is required for non-windowed declarations · New

What changed in this PR

Adds bounded Level-2 workspace estimation for non-windowed CUDA GroupQueryAttention.

Changes:

• Adds a session-configured total-sequence-length bound.
• Models one-sided cache alias preservation.
• Adds estimator, aggregation, declaration tests, and documentation.

File	Description
group_query_attention.h	Stores the configured bound.
group_query_attention.cc	Parses and applies the bound.
group_query_attention_workspace_estimate.h	Extends estimator configuration.
group_query_attention_workspace_estimate.cc	Builds non-windowed bounds.
group_query_attention_workspace_bounds.h	Adds past-capacity and alias metadata.
group_query_attention_workspace_bounds.cc	Aggregates partial-alias workspace routes.
group_query_attention_workspace_estimate_test.cc	Tests estimation and declaration behavior.
onnxruntime_session_options_config_keys.h	Defines the new session option.
attention_workspace_estimation.md	Documents non-windowed estimation.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Ti-Tai Wang (titaiwangms)]

Full review

Verdict: Request changes. The partial-alias and workspace-sizing implementation is generally sound, but there are three confirmed Major issues. The first is a functional integration gap rather than merely a test problem.

1. Major: the new session bound does not reach Level 1, so constrained-memory planning still cannot use it

WorkspaceEstimatorConfig already carries session configuration into CUDA GetCapability() for options such as FP-A/INT-B (resource_accountant.h:66-69, resource_accountant.cc:588-595). However, ep.cuda.gqa_workspace_max_total_sequence_length is only read during GQA kernel construction and is not passed to the Node-based Level-1 estimator.

As a result, every non-windowed GQA node still uses the generic workspace fallback during partitioning. After kernel construction, Level 2 can declare the larger bounded root. session_state.cc:2124-2184 then either:

• logs that the declaration exceeds the partitioning reservation; or
• fails session initialization when session.strict_workspace_verification=1.

The documentation statement that Level 1 “cannot access the session option” is also inaccurate: the framework already provides a session-config channel; this key has simply not been plumbed through it.

Suggested fix: add this key to WorkspaceEstimatorConfig, parse it once through the resource-accountant path, and pass the same value to the Node-based GQA estimator. Please add tests proving:

• Level 1 and Level 2 produce the same root size for bounded non-windowed GQA;
• strict workspace verification initializes successfully; and
• a CUDA budget between the generic fallback and the real estimate rejects assignment.

2. Major: KernelDeclaresBoundedNonWindowedRoot has no reachable backend

The model builder always includes head_sink (group_query_attention_workspace_estimate_test.cc:183-222). The new test then combines:

• ORT_ENABLE_XQA=0;
• sdpa_kernel = kMath, disabling Flash and MEA; and
• head_sink, which excludes the unfused and cuDNN routes.

The estimator therefore correctly returns nullopt, and ASSERT_TRUE(expected.has_value()) at line 972 deterministically fails. This matches the current CUDA and TensorRT CI failures.

Suggested fix: remove head_sink from this test and validate the kMath/unfused route, or retain the sink, enable XQA, and gate the test on the required GPU architecture. The assertion should not simply be weakened because the purpose of this test is to prove that the kernel and estimator declare identical bytes.

3. Major: the configured bound is not checked against the runtime scalar

The option is parsed in the constructor and used for Level-2 declaration (group_query_attention.cc:56-65, 149-150, 238), but execution never verifies that the actual non-windowed total_sequence_length is within max_total_sequence_length_.

Today this does not directly cause an out-of-bounds workspace access because runtime still allocates with GetScratchBuffer() using the live dimensions. It does, however, permit an understated declaration and invalid memory accounting. Once #32071 consumes the planner-owned root, the same unchecked mismatch becomes a workspace safety issue.

Suggested fix: before runtime allocation, reject a non-windowed invocation whose actual total/present sequence length exceeds the configured bound. At minimum, this validation must be tracked as a blocking prerequisite for #32071. The parser should also reject values above INT32_MAX, since the runtime scalar and workspace geometry are int32_t-bounded; currently such a value is accepted but later makes estimation silently unavailable.

Additional issues

• Public contract: onnxruntime_session_options_config_keys.h:512-516 suggests the option is needed only when the scalar may exceed past-cache capacity. The implementation requires a positive value for every non-windowed Level-2 estimate. Please document that unset means non-windowed Level 2 is disabled, that this is one session-wide bound covering all non-windowed GQA nodes, and what happens when runtime exceeds it.
• Stale comments: the Flash fast-decode comment still describes non-windowed adapter support as future work, although this PR makes that route adapter-reachable.
• Attention bias: estimation now rejects non-windowed bias as well, but the documentation only describes sliding-window bias as unavailable.
• API clarity: all callers explicitly pass requires_separate_past_buffer; removing its default value from MakeProblem would keep alias assumptions explicit.

What looks correct

• max(past_capacity, configured_bound) is the correct present-cache/backend KV upper bound.
• The one-sided-alias preservation region matches the runtime full-past-tensor copy.
• Regular Flash, MEA, and unfused envelopes include the partial-alias variant.
• XQA and Flash fast decode correctly exclude partial alias because their runtime routes require both cache pairs to alias.
• The graph-free estimator remains free of graph-type leakage.

This was a static full-team review. I did not run additional local CUDA tests; the impossible-route kernel test failure is already reproduced by CUDA and TensorRT CI.