Repository navigation
[CUDA] Plumb GQA workspace KV-length envelope knob into the L1 estimator (reader-only) - #33014
Conversation
…tor (reader-only) Adds the session option ep.cuda.gqa_workspace_max_total_sequence_length and threads its value through to the CUDA GroupQueryAttention Level-1 workspace estimator. total_sequence_length (accumulated past + current KV tokens) is a runtime scalar input whose value cannot be recovered from graph shapes, so the non-windowed workspace estimate cannot currently bound it. This knob hands the estimator that KV-length envelope directly for capacity-aware partitioning. This is a pure no-functional-change reader: the value is parsed (WorkspaceEstimatorConfig -> GQAWorkspaceEstimateConfig::max_total_sequence_length) and stored, but not yet consumed. BuildBounds still returns std::nullopt for the non-windowed case, so the estimate output is unchanged. A follow-up PR consumes the field to produce the non-windowed bound. Plumbing mirrors the existing ep.cuda.fpa_intb_* knobs. Extends the existing FactoryRetainsNarrowWorkspaceEstimatorConfig framework test to validate the reader end-to-end (config string -> WorkspaceEstimatorConfig field) in CPU CI. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The public option documentation incorrectly promises a bound that the current estimator still ignores.
Review effort: Balanced
Findings: 1
Open (1)
What changed in this PR
Adds reader-only plumbing for a CUDA GQA workspace KV-length envelope without changing estimates.
Changes:
- Adds and retains the new session option.
- Parses and forwards the envelope into the GQA estimator configuration.
- Extends framework coverage for option retention.
| File | Description |
|---|---|
onnxruntime/test/framework/resource_accountant_test.cc |
Tests option retention. |
onnxruntime/core/providers/cuda/cuda_execution_provider.cc |
Parses and forwards the envelope. |
onnxruntime/core/framework/resource_accountant.cc |
Copies the session option into estimator configuration. |
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h |
Adds the envelope field and parameter. |
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc |
Stores the forwarded value without consuming it. |
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h |
Declares and documents the session option. |
include/onnxruntime/core/framework/resource_accountant.h |
Extends workspace estimator configuration. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Ti-Tai Wang (titaiwangms)
left a comment
Ti-Tai Wang (titaiwangms)
left a comment
There was a problem hiding this comment.
Reviewed head 249e090a859709e9c9133e5ad230561351110913 against base 4bc564ea41b8a1e7e9d47accacf2c4d233ef6e7f, in the context of #32944.
The implementation is genuinely reader-only: it retains, parses, forwards, and stores the option, but BuildBounds() does not consume it and non-windowed estimation remains unavailable. I found no newly lowered admission bound or runtime envelope enforcement in this diff.
1. Minor — public documentation promises a bound that is currently ignored
Location: include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h:468-475.
The public comment says a positive integer "sets the bound," but group_query_attention_workspace_estimate.cc:206-209 explicitly leaves the field unconsumed and still returns std::nullopt for non-windowed caches. Setting "4096" therefore does not produce the advertised estimate or change partitioning.
Please describe this version as reserved/reader-only with no effect on estimates, or publish the behavioral promise with the consuming change. Distinguish an estimation envelope from a runtime-enforced input limit; this PR establishes neither a hard KV limit nor a no-OOM guarantee.
2. Minor — malformed explicit values silently become unspecified, and the test stops before parsing
Locations: onnxruntime/core/providers/cuda/cuda_execution_provider.cc:3645-3657 and onnxruntime/test/framework/resource_accountant_test.cc:810-827.
Malformed, negative, or overflowing values leave the effective value at zero with no diagnostic. The added test only checks retention of the raw "4096" string; it would still pass if parsing or forwarding regressed.
Please define whether invalid explicit values fail or produce a warning, and cover positive, zero, unset, negative, malformed, and overflowing values. Add coverage that the parsed value reaches GQAWorkspaceEstimateConfig, not just WorkspaceEstimatorConfig.
Follow-up integration question
Please reconcile this option with #32696's L2 use and #33016's L1 consumer so one key has a consistent parsed meaning across the stages. Before a future consumer trusts a reduced admission bound, the query/KV workload domain must be proven or enforced, and every reachable route must be safely bounded. Merely parsing this knob is not that prerequisite.
These are documentation/parsing/test-coverage findings, not a claim that the current reader-only change underestimates memory. No build, end-to-end parser/estimator execution, or GPU validation was performed for this PR.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 37ad7066-fcd3-45df-891e-44bee9e1cf04
|
Ti-Tai Wang (@titaiwangms) Addressed the documentation, parsing, and test-coverage findings from your review in Reader-only contract: The public option documentation now states that the value is reserved/forwarded but does not change estimates or partitioning. Non-windowed estimation remains unavailable; this PR does not enforce a runtime input limit or provide a no-OOM guarantee. Parsing and coverage: Parsing now happens once in Integration follow-up: Alignment with #32696 and #33016 is not claimed as completed here. Their consumers need to share this parsed meaning, including zero-as-unspecified. Before a consumer trusts a reduced admission bound, the query/KV workload domain must be proven or enforced and every reachable route safely bounded; parsing this option alone does not establish that prerequisite. This PR does not modify either consumer. Validation: Six factory tests passed in a focused native run. The forwarding regression has been added but was not executed locally; full ORT/GPU validation is not claimed. |
Ti-Tai Wang (titaiwangms)
left a comment
Ti-Tai Wang (titaiwangms)
left a comment
There was a problem hiding this comment.
Follow-up review at 97bf527: the concerns from my previous review are addressed.
The public option is now accurately documented as reader-only: it does not change estimates or partitioning, non-windowed estimation remains unavailable, and it is neither a runtime-enforced limit nor a no-OOM guarantee. Validation happens once in the accountant factory, including when partitioning is disabled, and explicit invalid values return INVALID_ARGUMENT rather than silently becoming unspecified.
The new tests cover valid/invalid parsing and forwarding to the node configuration while asserting that estimates remain unchanged. I also checked the pinned parsing helper: signed int64 uses entire-string decimal std::from_chars, matching the rejection tests for leading plus, whitespace, suffixes, and overflow.
No additional actionable source findings. This reader-only slice looks ready from the code-review perspective, subject to normal CI completion. The failed logs inspected show setup/tooling problems: the Java manifest cannot be loaded, Rust toolchain setup retries fail, the Python-format job subsequently lacks lintrunner_adapters, and the iOS vcpkg download times out. I did not find a changed-source lint diagnostic in those logs; that is not equivalent to a completed lint pass.
Validation scope: source/test review, pinned parser inspection, and failed CI logs; no new local build or test execution at this head. This COMMENT is not a formal approval.
Ti-Tai Wang (titaiwangms)
left a comment
Ti-Tai Wang (titaiwangms)
left a comment
There was a problem hiding this comment.
Reviewed 97bf527. The reader-only contract, strict factory validation, and forwarding/unchanged-estimate coverage address the previous concerns; no remaining actionable source findings. Approving the code changes; the setup/tooling CI failures still need normal CI resolution and this does not waive required checks. No new local build/test execution was performed.
…tor (reader-only) (#33014) ## What Adds the session option `ep.cuda.gqa_workspace_max_total_sequence_length` and threads its value into the CUDA GroupQueryAttention (GQA) Level-1 workspace estimator. `total_sequence_length` (accumulated past + current KV tokens) is a **runtime scalar input** whose value cannot be recovered from graph shapes. For the non-windowed KV cache, the backend workspace scales linearly with it, so the L1 estimator currently returns `std::nullopt` for that case (it cannot bound the value). This knob lets the user hand the estimator that **KV-length envelope** directly, for capacity-aware partitioning (#31962). ## No functional change (reader-only) This is a **pure NFC reader**: - The value is parsed (`WorkspaceEstimatorConfig` → `GQAWorkspaceEstimateConfig::max_total_sequence_length`) and **stored**, but **not consumed**. - `BuildBounds` still returns `std::nullopt` for the non-windowed case, so **the estimate output is unchanged** for all inputs. - A follow-up PR (PR 4) consumes the field to produce the non-windowed bound. The plumbing mirrors the existing `ep.cuda.fpa_intb_*` knobs (config key → `WorkspaceEstimatorConfig` → CUDA EP parse → estimator param). ## Test Extends the existing `RealAccountantTest.FactoryRetainsNarrowWorkspaceEstimatorConfig` framework test to validate the reader end-to-end (config string → `WorkspaceEstimatorConfig` field) in CPU CI. ## Context Part of the policy-based workspace-estimation PR chain (design doc: `docs/annotated_partitioning/policy_based_workspace_estimation.md`). This is **PR 3 (envelope plumbing)** in the roadmap: - PR 1 — GQA scratch sizing consolidation (#32945) - PR 2 — route-eligibility extraction (#33001) - **PR 3 — envelope plumbing (this PR)** - PR 4 — route-aware, envelope-bounded estimate (consumes this field) --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf Copilot-Session: 37ad7066-fcd3-45df-891e-44bee9e1cf04
…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
## Summary Complete the preparation-sizing half of Stage 1 in the policy-based, route-aware workspace-estimation design (#32944). This combines the agreed QKV-preprocessing refactor and the remaining cache/sequence preparation sizing into one PR. Builds on #32945, now merged into `main`, which consolidates MEA, unfused, XQA and Flash backend scratch sizing. - Share checked cache-row, past-preservation, staging, compaction, sequence-vector and QKV-preprocessing calculations between `ComputeInternal` and `GetGQAPreparationRecipe`. - Remove the duplicate `GQABufferRequirements` implementation. - Keep the existing separate allocations, ownership, pointer wiring, dispatch priorities and lifetimes. - Preserve the preparation recipe's aligned layout and full validation; the live kernel calls narrow sizing helpers rather than inheriting unrelated recipe-admission restrictions. This is sizing-code consolidation, not a missing-buffer fix. It does not activate dispatch policies, exclude routes from L1, coalesce allocations, or change static-preallocation strategy. ## Compatibility details - Cache sizing runs before backend selection and retains the original present-cache capacity. Windowed staging still changes the runtime extent from `C` to `C + S`. - Past preservation uses the aliased tensor's physical row width; staging uses the physical output row width; compaction retains its packed-row interpretation. INT4 rows are not packed a second time. - QKV sizing runs after the existing backend selection and receives the effective post-staging capacity. cuDNN retains the existing generic Q-materialization behavior. - Flash fast decode still reuses the device sequence-length input instead of allocating sequence vectors, and still suppresses the QKV-preprocessing buffer. - Existing healthy-input byte counts are retained. Arithmetic overflow and invalid sizing inputs produce explicit sizing statuses at the runtime boundary. ## Coverage and validation - Fresh standalone build of the existing preparation Google Test source: **50 tests passed**, including the legacy QKV feature matrix, dense/byte/INT4 cache geometry, physical packed rows, sequence-vector suppression, invalid dimensions and overflow. - MSVC `/W4 /WX` compilation of the actual CUDA GQA runtime translation unit in both shared-provider and plugin-adapter configurations, with Flash, MEA, FP8 and INT4 enabled. - MSVC `/W4 /WX` compilation of the graph-time estimator and XQA/Flash sizing-test translation units. - MSVC `/W4 /WX` compilation of the CUDA GQA operator-test translation unit. - Four new CUDA regressions compare packed and separate QKV on unfused/Flash paths, including windowed staging. They reuse the aliasing helper, which disables CPU fallback and verifies actual backend dispatch. - `clang-format --dry-run --Werror` and `git diff --check` passed. **Full CUDA inference execution is not validated locally against a freshly linked provider.** The new operator regressions and existing GQA tests need full CUDA build/test CI. ## Stack Base: `main`. The conflict-resolution merge retains the early preparation-sizing initialization, the upstream seq-free eligibility helpers, and the INT8/FP8 storage-width normalization. The envelope and policy-reader PRs (#33014, #33016 and #33042) remain independent of this refactor. Committed-policy resolution, coverage/envelope enforcement, bounded fallback and policy-bound estimates remain follow-up work. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf
…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
What
Adds the session option
ep.cuda.gqa_workspace_max_total_sequence_lengthand threads its value into the CUDA GroupQueryAttention (GQA) Level-1 workspace estimator.total_sequence_length(accumulated past + current KV tokens) is a runtime scalar input whose value cannot be recovered from graph shapes. For the non-windowed KV cache, the backend workspace scales linearly with it, so the L1 estimator currently returnsstd::nulloptfor that case (it cannot bound the value). This knob lets the user hand the estimator that KV-length envelope directly, for capacity-aware partitioning (#31962).No functional change (reader-only)
This is a pure NFC reader:
WorkspaceEstimatorConfig→GQAWorkspaceEstimateConfig::max_total_sequence_length) and stored, but not consumed.BuildBoundsstill returnsstd::nulloptfor the non-windowed case, so the estimate output is unchanged for all inputs.The plumbing mirrors the existing
ep.cuda.fpa_intb_*knobs (config key →WorkspaceEstimatorConfig→ CUDA EP parse → estimator param).Test
Extends the existing
RealAccountantTest.FactoryRetainsNarrowWorkspaceEstimatorConfigframework test to validate the reader end-to-end (config string →WorkspaceEstimatorConfigfield) in CPU CI.Context
Part of the policy-based workspace-estimation PR chain (design doc:
docs/annotated_partitioning/policy_based_workspace_estimation.md). This is PR 3 (envelope plumbing) in the roadmap: