[CUDA] Bound GQA non-windowed workspace via the KV-length envelope (route-aware L1 estimate)

[Chi Lo (chilo-ms)]

Summary

Adds opt-in, route-aware Level-1 workspace estimation for non-windowed CUDA GroupQueryAttention using the declared total-KV envelope from ep.cuda.gqa_workspace_max_total_sequence_length.

The present-cache capacity bound is max(allocated past-cache capacity, declared total-KV envelope). Static caches may be larger than the active sequence: GQA permits total_sequence_length <= past_sequence_length and retains that allocated capacity. Bounding only the envelope would underestimate MEA expansion and partial-alias preservation.

The existing graph-free aggregate sizes reachable routes and admitted prompt/decode phase corners; this PR does not introduce a second sizing engine.

Envelope validation and caller contract

• Explicit values are parsed once during session initialization and must be decimal integers in [0, INT32_MAX], matching the runtime ABI. Malformed, negative, overflowing, and greater-than-INT32_MAX values fail initialization.
• Zero remains the unspecified value. With the envelope unset or zero, non-windowed estimation is unavailable; windowed estimation continues to use its cache-shape bound.
• The envelope must cover the query-length bound, and the resulting cache-capacity bound must fit the runtime int32 ABI.
• The query/max-shape and total-KV declarations are estimation inputs, not runtime-enforced admission limits or a no-OOM guarantee. The proof applies only within the caller-declared workload envelope; runtime enforcement remains follow-up work.

Workspace accounting and proof policy

• Capacity-sensitive preparation and MEA expansion cover the full static-cache extent, even when it exceeds the active total-KV envelope.
• Partial aliasing copies the full past tensor. Its preservation scratch coexists with the chosen route and is added above the per-route maximum, bounded by the full cache-capacity bound.
• Preparation recipes already account for sequence-length vectors; there is no new missing-vector gap.
• Reachable cuDNN without a workspace oracle, unsupported attention-bias cases, and otherwise unprovable route domains conservatively make the estimator unavailable.
• For capacity-aware CUDA placement, every GQA candidate must have a successful analytical estimate. If proof is unavailable, CUDA assignment for that node is declined.
• The generic accountant fallback is an ad hoc safety margin, not a proven GQA bound. Profile data is empirical and likewise cannot substitute for the missing analytical proof. Unrelated operators retain their existing generic fallback behavior.

Regression coverage

• Accepted boundary value INT32_MAX; rejected INT32_MAX + 1, INT64_MAX, decimal overflow, malformed text, and negative values.
• Unset-envelope behavior and monotonicity when the envelope increases.
• Static past capacity 1024 with total-KV envelope 512, checking full preservation and MEA expansion against complete concrete runtime recipes.
• An envelope below the query-length bound still declines, even when the static cache is larger.
• Static capacity outside the runtime int32 ABI declines.
• Parsed-envelope plumbing and unchanged windowed estimates.
• Additive full-past preservation above the per-route maximum.
• Capacity-aware CUDA placement declines GQA when the analytical estimator is unavailable, including when profile data exists; unrelated operators retain generic fallback.

Dependencies

Targets main, which now includes #33118's shared preparation-sizing implementation. The duplicated sizing commits were dropped during the post-merge rebase; this PR contains only its envelope-estimation, hardening, test, and documentation changes.

Part of the policy-based workspace-estimation series described in #32944. Committed per-(node, phase) route policy and runtime workload-envelope enforcement remain follow-up work.

[Copilot]

🟡 Changes recommended

Valid static-cache executions can exceed the estimate when allocated past capacity is larger than the declared total-length envelope.

1 open finding

• Static-cache capacity can exceed envelope, underestimating L1 allocation · New

What changed in this PR

Adds opt-in, envelope-bounded workspace estimation for non-windowed CUDA GroupQueryAttention.

Changes:

• Uses the configured KV-length envelope for non-windowed estimates.
• Accounts for partial-alias past-cache preservation scratch.
• Adds envelope and monotonicity tests.

File	Description
group_query_attention_workspace_estimate_test.cc	Tests non-windowed envelope behavior and preservation accounting.
cuda_execution_provider.cc	Updates envelope-consumption commentary.
group_query_attention_workspace_estimate.h	Documents the active envelope configuration.
group_query_attention_workspace_estimate.cc	Builds non-windowed bounds from the envelope.
group_query_attention_workspace_bounds.h	Adds preservation-scratch accounting state.
group_query_attention_workspace_bounds.cc	Adds preservation bytes to aggregate workspace.
onnxruntime_session_options_config_keys.h	Documents the session option’s behavior.
resource_accountant.h	Removes obsolete reader-only wording.

🧠 Review effort: Balanced

---

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

[Ti-Tai Wang (titaiwangms)]

Reviewed at a17ecc6de4ee533fbdf58ee42edfaf4fca226210.

No blocking source finding established. The bound now accounts for allocated cache capacity independently of the active KV envelope, and adds partial-alias preservation on top of the backend maximum. The explicit L1-only scope and the absence of runtime envelope enforcement are appropriate for this slice; missing L2 forwarding is not a blocker for this PR.

Two small documentation updates would help:

• The bounds header still says callers must separately account for partial-alias preservation, although the new flag makes the aggregate account for it. See the inline comment.
• Please update docs/annotated_partitioning/attention_workspace_estimation.md, which still describes successful graph-adapter estimation as sliding-window-only. Describe positive-envelope non-windowed eligibility, remaining decline conditions such as a reachable cuDNN route without an oracle, and the unchanged generic accountant fallback when estimation is unavailable.

Validation evidence: In the existing exact-head Linux CUDA CI, the static-cache preservation and MEA-expansion tests, plus the envelope tests, executed and passed. I did not run a local ORT build or tests. This review does not establish a runtime-enforced capacity or no-OOM guarantee, nor does it validate a combined tree with #32071.

[Ti-Tai Wang (titaiwangms)]

Minor: update the partial-alias contract in this comment block. The preceding sentences still require callers to charge preservation separately, but account_partial_alias_preservation now makes the aggregate add the full past-copy charge itself. A future caller following the old wording could double-charge or misunderstand the flag. Please say that both pairs are assumed to alias unless this flag is enabled, in which case the aggregate accounts for preservation.

[Chi Lo (chilo-ms)]

The updated wording is already included in the pushed head. It now says that both past/present K/V pairs are assumed to alias unless account_partial_alias_preservation is enabled, in which case the aggregate itself adds the worst-case preservation scratch. This removes the old instruction to charge that scratch separately and avoids double-counting.

[Chi Lo (chilo-ms)]

Thanks for the exact-head review. Both documentation findings are addressed in the prepared local hardening stack at 1f175d33b3:

• the bounds contract now says the aggregate assumes both K/V pairs alias unless account_partial_alias_preservation is enabled, in which case it adds the worst-case preservation copy;
• the Attention roadmap now describes positive-envelope non-windowed eligibility, max(allocated cache capacity, envelope), remaining unavailable cases (including reachable cuDNN without an oracle), and the final proof-unavailable policy.

The final policy intentionally requires a proven GQA bound for capacity-aware CUDA placement: unavailable GQA estimation declines assignment, while generic fallback remains unchanged for unrelated operators and profile data cannot substitute for the missing proof. The update is committed locally but will be published after #33118 merges so the old duplicated sizing stack can be dropped cleanly. The public head remains a17ecc6de4ee533fbdf58ee42edfaf4fca226210 until then; no re-review is requested yet.

[Ti-Tai Wang (titaiwangms)]

Reviewed the four-commit delta from the approved #33118 head (80292cc) to 1f175d3. The branch contains that exact preparation baseline; the divergent main comparison results from the squash merge, not missing preparation changes.

No blocking source findings. The non-windowed estimate covers the greater of allocated cache capacity and the declared envelope, plus worst-case partial-alias preservation. The shared parser preserves zero/unset semantics while rejecting values outside the runtime int32 ABI. In the in-tree capacity-aware CUDA path, an unavailable GQA estimate now declines assignment before generic fallback or profile accounting; ordinary placement without an accountant is unchanged. The corresponding placement and parser boundary tests are present, and the previously noted stale envelope documentation has been updated.

Nonblocking documentation follow-up: please qualify the new fail-closed guarantee as applying to the in-tree CUDA estimation path. BUILD_CUDA_EP_AS_PLUGIN still excludes the GQA Level-1 estimation block and its integration tests, so this PR does not establish the same guarantee for CUDA-plugin placement.

Validation caveat: at the latest check snapshot for this exact head, 27 checks succeeded, one was neutral, 57 were pending, and none had failed. The CUDA jobs were still running, so I have not verified execution of the new placement/static-cache boundary tests on this revision. This is a source-review approval, not an assertion that CI is complete or that the declared envelope is runtime-enforced. No local build or tests were run.

[Copilot]

🟡 Changes recommended

CUDA plugin accounting bypasses the new fail-closed GQA estimation requirement and still uses the generic fallback.

1 open finding

• CUDA plugin GQA bypasses fail-closed capacity estimation · New

1 resolved since last review

• Static-cache capacity can exceed envelope, underestimating L1 allocation

🧠 Review effort: Balanced

[Chi Lo (chilo-ms)]

Fixed in 2cff2bafad.

The host-side CUDA plugin accounting path now declines new com.microsoft::GroupQueryAttention assignments whenever a resource accountant is present, before calling ComputeResourceCount. Because that path cannot receive analytical estimates yet, this deliberately remains conservative even when a KV envelope is declared. A fused group containing an unproven GQA node is declined as a whole, so neither generic nor profiled costs can bypass the proof requirement. Placement without an accountant, previously accounted assignments, unrelated operators, and other plugin EPs retain their existing behavior.

Added regression coverage for single-node/fused groups, ad-hoc/profile accounting, unset/positive envelopes, and the unchanged cases above, and documented the plugin limitation. Tests were not run locally; the local build was stopped before completion.