Repository navigation
[CUDA] Plumb GQA workspace envelope into Level-2 declaration - #33240
Draft
Chi Lo (chilo-ms) wants to merge 5 commits into
Draft
Chi Lo (chilo-ms) wants to merge 5 commits into
Chi Lo (chilo-ms) wants to merge 5 commits into
Conversation
…oute-aware L1 estimate) Make the GQA Level-1 workspace estimator produce a sound bound for the non-windowed KV-cache case instead of declining. Previously BuildBounds returned nullopt for every non-windowed node because total_sequence_length is a runtime scalar not recoverable from shapes. It now uses the declared KV-length envelope (config.max_total_sequence_length, from session option ep.cuda.gqa_workspace_max_total_sequence_length, plumbed in the prior PR) as present_kv_cache_capacity_bound and evaluates the existing route-aware, phase-cornered aggregate over it. Opt-in and default-off: when the envelope is unset (== 0) the non-windowed path still returns nullopt, so behavior is byte-identical to today. The windowed path is untouched -- every change is guarded by the windowed flag. Soundness: - Backend workspace is non-decreasing in KV length and the runtime sizes non-windowed scratch from total_sequence_length <= envelope (GetGQAEffectiveWorkspaceKvLength), so capacity_bound = envelope is an upper bound. - Non-windowed partial aliasing (exactly one past/present K/V pair shared) copies the full past cache into preservation scratch that coexists with the backend workspace. This is now charged on top of the per-route maximum via GQAWorkspaceBounds::account_partial_alias_preservation, bounded by the envelope and head_size_bound. Windowed GQA rejects non-shared buffers, so it never incurs this copy. - cuDNN (when reachable) and attention bias conservatively decline (nullopt), keeping the estimate an upper bound. The small seq_lens buffer remains an unmodeled pre-existing gap shared with the windowed path. Tests (onnxruntime_provider_test): non-windowed without an envelope still declines; with an envelope it is bounded and monotonic in the envelope; the preservation copy is additive on top of the selected route. Part of the policy-based workspace estimation effort; see docs/annotated_partitioning/policy_based_workspace_estimation.md and the PR chain (#32944, #32945, #33001, #33014). Base retargets to main after #33014 merges. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf
Use the larger of allocated past capacity and the total-KV envelope to cover MEA expansion and full partial-alias preservation. Keep the query-length envelope check independent and add static-cache and runtime-ABI regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf
Reject out-of-ABI envelopes during initialization, require a proven route-aware GQA estimate for capacity-aware CUDA placement, and keep the caller-declared envelope contract explicit. Share the validated parser for Level-2 reuse and cover invalid values plus fail-closed placement. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Reuse the validated workspace-envelope parser during CUDA GQA kernel construction and pass the stored value into the shared estimator for workspace declaration. Add direct L1/L2 agreement and strict-verification coverage for bounded non-windowed GQA. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This was referenced Oct 9, 2026
Base automatically changed from
chilo/gqa-workspace-nonwindowed-estimate
to
main
October 15, 2026 22:53
This branch has not been deployed
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
1 participant
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.
Summary
Extract the unique Level-2 CUDA GroupQueryAttention declaration plumbing from the closed #32696 integration and stack it narrowly on #33016.
ParseCudaGqaWorkspaceMaxTotalSequenceLengthparser during GQA kernel construction; no private integer parser is introduced.GQAWorkspaceEstimateConfigfromDeclareWorkspaceRequirements.[0, INT32_MAX]validation and caller-contract semantics.Coverage
Scope
This is constructor/member/declaration plumbing only: three files, one commit. It does not enable planned workspace allocation, enforce runtime query/KV bounds, restrict runtime routes, or revive #32696's private parser.
Draft while #33016 completes review and CI. Base:
chilo/gqa-workspace-nonwindowed-estimate.