Add GQA workspace estimation

[Ti-Tai Wang (titaiwangms)]

Description

Adds conservative Level-1 and Level-2 workspace estimation for CUDA contrib GroupQueryAttention, tracking #29775.

This is PR4, the final GQA workspace-estimation PR in the current stack. It depends on #32602, which bounds windowed runtime workspace by the resident/staged KV extent.

The estimator:

• inventories all reachable GQA CUDA backend routes over dynamic shape bounds;
• composes checked preparation and backend workspace recipes;
• uses C for single-token windowed decode and checked C + S for multi-token staging;
• bounds Flash split workspace independently of the non-monotonic split heuristic;
• reports non-windowed and reachable cuDNN cases as unavailable when no sound graph-free bound exists;
• reports Level-1 runtime_workspace_bytes; and
• declares one operator-owned slot-0 Level-2 root with 256-byte alignment.

This PR does not change runtime allocation topology, opt into planner-owned allocation, or slice a planned root. Those remain follow-up integration work.

Validation

• Fresh Release CUDA build on A100 (SM80)
• CUDA internal module: 264 passed, 2 unrelated LeanAttention skips
• GQA workspace estimate/bounds tests: 18/18 passed
• GQA workspace recipe/arithmetic tests: 19/19 passed
• Lintrunner and git diff --check passed
• Deep route-dominance review swept 2.7M complete routes and 25.6M Flash boundary points with no underestimation

Stack

• Depends on Bound windowed GQA workspace allocations #32602
• Tracks [Tracking] Constrained-environment memory and layer-offload roadmap #29775

[Copilot]

🟡 Changes recommended

The new Level-2 kernel override lacks direct tests of its configuration translation and virtual dispatch.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds conservative CUDA GroupQueryAttention workspace estimation and exposes it through Level-1 accounting and Level-2 declarations.

Changes:

• Adds checked preparation/backend recipes and bounded route aggregation.
• Integrates estimates into CUDA partition accounting and kernel workspace declarations.
• Expands workspace and windowed-cache regression coverage.

File summaries

File	Description
onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_test.cc	Expands preparation-recipe tests.
onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_mea_unfused_test.cc	Tests MEA, unfused, and composed recipes.
onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_header_test.cc	Extends shared-header checks.
onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc	Tests bounds and estimator adapters.
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc	Adds windowed CUDA regressions.
onnxruntime/core/providers/cuda/cuda_execution_provider.cc	Adds Level-1 GQA accounting.
onnxruntime/contrib_ops/cuda/bert/group_query_attention.h	Declares Level-2 workspace support.
onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc	Implements declaration and bounded runtime sizing.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace.h	Defines workspace recipe types.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace.cc	Models preparation allocations.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_xqa_flash.cc	Records Flash route metadata.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_routes.cc	Composes complete route recipes.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_mea_unfused.cc	Implements MEA and unfused recipes.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h	Declares estimator adapters.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc	Parses graph/kernel facts into bounds.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h	Defines bounded estimation inputs.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc	Aggregates reachable backend envelopes.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu	Applies effective windowed KV length.
docs/annotated_partitioning/attention_workspace_estimation.md	Documents GQA estimation behavior.

Review details

• Files reviewed: 19/19 changed files
• Comments generated: 1
• Review effort level: Balanced

---

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

[Copilot]

🔵 Needs a closer look

Backend dispatch and memory-bound calculations directly affect partition acceptance across a large CUDA route matrix.

Review details

• Files reviewed: 19/19 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Valid partial-RoPE models are rejected by the estimator, and the XQA kernel test lacks an SM80 capability gate.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 19/19 changed files
• Comments generated: 3
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The estimator accepts cache geometry that the CUDA runtime rejects, and its primary fixtures exercise that invalid case.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 19/19 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The Level-1 accounting integration lacks end-to-end coverage, and one test relies on a transitive standard-library include.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc:8

• This file uses std::max at line 645 but does not include its declaring header, so it currently relies on transitive test/framework includes. Include <algorithm> directly to keep the test portable across standard-library configurations.

• Files reviewed: 19/19 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

Dynamic CUDA route envelopes directly affect partition memory accounting and warrant final human validation.

Review details

• Files reviewed: 19/19 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Chi Lo (chilo-ms)]

1 actionable finding (P2) —  onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc:504 

 SetGroupQueryAttentionLevel1MemoryEstimate()  only populates  runtime_workspace_bytes , leaving  persistent_prepack_bytes  at zero. For a constant FP16/BF16  head_sink ,  PrePack()  retains the original initializer and allocates an additional session-lived FP32 buffer ( group_query_attention.cc:170-201 ). The Level-1 contract explicitly requires such kernel-owned prepacked buffers to be reported separately.

Treating the sink as dynamic only adds its storage to the XQA route. Because route aggregation takes the maximum across XQA, Flash, and MEA, that storage can disappear when another route dominates, even though the prepacked buffer remains live. This can make CUDA partition accounting underestimate memory by  num_heads * sizeof(float)  and accept a partition above its configured limit. Detect a constant initializer at Level 1—or conservatively charge the FP32 copy whenever  head_sink  is present—and populate  persistent_prepack_bytes ; add a regression where Flash/MEA workspace exceeds XQA workspace

[Ti-Tai Wang (titaiwangms)]

Fixed in fd85715. Level-1 now detects a constant head_sink initializer in CUDA GetCapability, models its session-lived FP32 copy separately as persistent_prepack_bytes = num_heads * sizeof(float), and keeps it outside the mutually exclusive route workspace maximum. The regression also verifies that this persistent charge remains when unfused workspace dominates XQA. The branch is rebased onto merged #32602/current main; the CUDA internal suite passed 269 tests with only the 2 existing LeanAttention skips.

[Copilot]

🟡 Changes recommended

Level-1 accounting can underestimate XQA workspace when session prepacking is disabled.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 10/10 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

Correctness depends on conservative bounds across numerous hardware-dependent CUDA routes and warrants final human validation.

Review details

• Files reviewed: 10/10 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Chi Lo (chilo-ms)]

Low-priority test gap:  group_query_attention_workspace_estimate_test.cc:805  exercises  SetGroupQueryAttentionLevel1MemoryEstimate()  only with an aggregate whose persistent and initialization fields are zero. The end-to-end budget test would not detect those fields being dropped. Add a direct positive setter assertion with nonzero values or inspect the resource accountant’s pending breakdown. Required CI is still mostly pending; completed checks currently pass.

[Ti-Tai Wang (titaiwangms)]

Addressed in 5bdd479. The Level-1 setter test now supplies nonzero persistent_prepack_bytes and initialization_scratch_bytes and directly verifies that both propagate. CUDA internal tests still pass 270 tests with only the 2 existing LeanAttention skips.

[Copilot]

🔵 Needs a closer look

The estimator spans dynamic CUDA dispatch, non-monotonic workspace formulas, and partition-budget behavior requiring final expert review.

Review details

• Files reviewed: 10/10 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

CUDA backend route coverage and partition-budget effects warrant final human validation despite strong tests.

Review details

• Files reviewed: 10/10 changed files
• Comments generated: 0 new
• Review effort level: Balanced