Add GQA workspace estimation - #32617
Open
Ti-Tai Wang (titaiwangms) wants to merge 17 commits into
Open
Ti-Tai Wang (titaiwangms) wants to merge 17 commits into
Ti-Tai Wang (titaiwangms) wants to merge 17 commits into
Conversation
Ti-Tai Wang (titaiwangms)
marked this pull request as ready for review
September 15, 2026 22:06
Ti-Tai Wang (titaiwangms)
added
ep:CUDA
Copilot started reviewing on behalf of
Ti-Tai Wang (titaiwangms)
September 15, 2026 22:07
View session
Contributor
There was a problem hiding this comment.
🟡 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 started reviewing on behalf of
Ti-Tai Wang (titaiwangms)
September 15, 2026 22:26
View session
Add MEA and unfused workspace recipes and compose one aligned concrete route root. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b24d04bb-12ab-41a0-8941-76eb9769005d
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Model the separate past-cache preservation allocation in complete workspace recipes and strengthen cross-recipe validation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Persist the Flash fast-decode witness and reject impossible windowed alias preparation states. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Cross-check separate-past preparation against XQA and Flash fast-decode backends. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Use the resident or staged cache extent when sizing Flash and unfused workspace for sliding-window GQA, while preserving non-windowed behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Expose the selected Flash split count in debug output and assert that staged C+S sizing avoids raw-length split workspace. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Assert the zero runtime encoding used when staged windowed KV sizing avoids split-KV workspace. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Apply the runtime windowed KV-length bound when composing Flash and unfused workspace recipes, and expose the effective allocation length in debug output. Cover staged, decode, below-capacity, and non-windowed route semantics. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add conservative graph-free route bounds, Level-1 accounting, and a single aligned Level-2 workspace declaration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Avoid overflow while validating rotary cache dimensions and document and test the heuristic-independent Flash fast-decode envelope across the full dynamic sequence domain. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Document aliasing and non-windowed capacity preconditions, explain the graph-free route envelopes, and qualify Level 1 versus Level 2 head-sink sizing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Construct the CUDA GQA kernel and verify virtual workspace declaration, prepacked head-sink translation, and unavailable-shape behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Ti-Tai Wang (titaiwangms)
force-pushed
the
titaiwangms/gqa-workspace-estimation
branch
from
September 16, 2026 20:53
39eb9f9 to
e248002
Compare
Copilot started reviewing on behalf of
Ti-Tai Wang (titaiwangms)
September 16, 2026 20:54
View session
Contributor
There was a problem hiding this comment.
🟡 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
Match runtime rotary-cache validation, add the missing cctype include, and gate the XQA kernel declaration test on SM80. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot started reviewing on behalf of
Ti-Tai Wang (titaiwangms)
September 16, 2026 21:07
View session
Contributor
There was a problem hiding this comment.
🟡 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
Require the bounded past cache capacity to match the sliding-window size, mirror the runtime past key/value capacity check, and make estimator fixtures use executable CUDA geometry. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot started reviewing on behalf of
Ti-Tai Wang (titaiwangms)
September 16, 2026 21:19
View session
Contributor
There was a problem hiding this comment.
🟡 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::maxat 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
Drive CUDA GetCapability across thresholds that distinguish the structured GQA workspace estimate from fallback accounting and verify strict rejection at the estimated resource count. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Declare the test dependencies directly and convert computed byte boundaries to the KiB units expected by capacity-aware partitioning. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot started reviewing on behalf of
Ti-Tai Wang (titaiwangms)
September 16, 2026 22:06
View session
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
Cfor single-token windowed decode and checkedC + Sfor multi-token staging;runtime_workspace_bytes; andThis PR does not change runtime allocation topology, opt into planner-owned allocation, or slice a planned root. Those remain follow-up integration work.
Validation
git diff --checkpassedStack