Skip to content

Enable bounded non-windowed GQA workspace estimation - #32696

Open
Chi Lo (chilo-ms) wants to merge 1 commit into
mainfrom
chilo/bounded-non-windowed-gqa-workspace-estimation
Open

Chi Lo (chilo-ms) wants to merge 1 commit into
mainfrom
chilo/bounded-non-windowed-gqa-workspace-estimation

Conversation

@chilo-ms

Copy link
Copy Markdown
Contributor

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.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot AI balanced review requested due to automatic review settings September 18, 2026 23:32

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 High severity · 1 Low severity

Open (2)
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.

GTEST_SKIP() << "A CUDA device is required to construct the CUDA kernel.";
}

ScopedEnvironmentVariables scoped_env_vars{{{"ORT_ENABLE_XQA", "0"}}};
Comment on lines +512 to +514
/// Optional positive upper bound for the total_sequence_length scalar of non-windowed CUDA
/// GroupQueryAttention nodes. The scalar value is unavailable during workspace declaration,
/// so callers must provide a sound bound when it may exceed the past-cache capacity.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants