Repository navigation
[Framework/CUDA] Add generic kernel dispatch policy reader - #33042
Conversation
Chi Lo (chilo-ms)
changed the title
Ti-Tai Wang (titaiwangms)
left a comment
Ti-Tai Wang (titaiwangms)
left a comment
There was a problem hiding this comment.
Reviewed the reader-only delta at 691c089f0963307500a82cc690a8bfb5973e98de, against its stacked base a17ecc6de4ee533fbdf58ee42edfaf4fca226210 (#33016).
No actionable finding. Parsing and forwarding are consistent with the documented reader-only contract: recognized lowercase policy tokens are parsed, unrecognized values resolve to Auto, the raw configuration value is retained, and the policy does not change route eligibility, workspace estimates, partitioning enablement, or runtime dispatch. The direct tests cover parsing, retention, estimate invariance, and the missing-envelope case.
This is not approval of future policy enforcement or a reduced/no-OOM admission contract. No local build or tests were run; the parent implementation was reviewed separately.
a17ecc6 to
1f175d3
Compare
## 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
1f175d3 to
3633cea
Compare
There was a problem hiding this comment.
🟡 Changes recommended
The stack is missing #33016’s latest ABI validation and fail-closed GQA placement hardening.
2 open findings
What changed in this PR
Adds a framework-owned, hardware-neutral kernel dispatch policy reader and forwards the intent to CUDA GQA without changing dispatch or sizing.
Changes:
- Adds shared policy tokens and exact parser.
- Retains raw session configuration and forwards parsed intent to GQA.
- Extends stacked GQA envelope accounting and regression coverage.
| File | Description |
|---|---|
include/onnxruntime/core/framework/kernel_dispatch_policy.h |
Defines policy enum and parser. |
include/onnxruntime/core/framework/resource_accountant.h |
Stores raw policy configuration. |
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h |
Documents the session option. |
onnxruntime/core/framework/resource_accountant.cc |
Reads the policy option. |
onnxruntime/core/providers/cuda/cuda_execution_provider.cc |
Forwards policy intent to GQA. |
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.h |
Adds policy to GQA estimate configuration. |
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_estimate.cc |
Handles bounded non-windowed estimation. |
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.h |
Models partial-alias preservation. |
onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_bounds.cc |
Accounts for preservation scratch. |
onnxruntime/test/framework/resource_accountant_test.cc |
Tests parsing and raw retention. |
onnxruntime/test/providers/cuda/test_cases/group_query_attention_workspace_estimate_test.cc |
Tests policy neutrality and GQA bounds. |
🧠 Review effort: Balanced
💡 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 the isolated reader-only change from 2cff2bafad0fa370c61075b94e74ff53ff6867e1 to f148b46dc8c60bbd6d3008b6c99975a5cde3e4da with the five-role review team. No introduced Major/Critical correctness blocker was established.
The shared policy representation remains graph/CUDA-free, and this delta forwards policy metadata without changing workspace sizing, dispatch, or partition behavior. Unknown tokens falling back to Auto are an explicit documented/tested parser contract, not an inferred strict-token requirement. Setting the policy alone does not enable resource accounting.
Stack-readiness caveat: this recorded head does not include the latest #33016 tip 83019a940215caaff5694eaec613df1ad8237e3c. Its private-field test access is inherited from the recorded base, not introduced by this nine-file delta. The reader does not modify that test file, so an ordinary three-way merge retains the updated target's fix; no new compile regression or merge conflict is claimed. Please refresh the stack as already planned. This approval covers the pinned isolated delta, not a hypothetical refreshed stack or the subsequent enforcement in #33316.
Validation limits: no local build/tests were run, and there is no qualifying exact-head CUDA CI evidence for this PR. Parent CI is not substituted for child validation; shared-provider link/load, exhaustive build-matrix coverage, runtime outputs and physical memory were not verified.
The base branch was changed.
… estimator (reader-only) Adds the hardware-neutral session option session.attention_dispatch_policy and threads its value through to the CUDA GroupQueryAttention Level-1 workspace estimator. This is the Tier-1 intent knob from the policy-based workspace estimation design: "latency" favors high-workspace fast routes, "memory" favors low-workspace routes (so more nodes fit on GPU), and "safe" assumes the bounded fallback everywhere. Unset (default) or any unrecognized value means "auto" and keeps the current route-aware behavior. This is a pure no-functional-change reader. The raw string is parsed (WorkspaceEstimatorConfig::attention_dispatch_policy -> AttentionDispatchPolicy enum via ParseAttentionDispatchPolicy) and stored on GQAWorkspaceEstimateConfig::dispatch_policy, but the estimator does not yet consume it, so every policy value produces the identical estimate. A follow-up PR biases route selection on the committed policy (coverage validation and the bounded fallback land in their own PRs). Plumbing mirrors the ep.cuda.gqa_workspace_max_total_sequence_length reader. The enum/parser live in the CUDA GQA estimator header because it is the only consumer today; hoist to a neutral header when a second op consumes it. Tests: extends the FactoryRetainsNarrowWorkspaceEstimatorConfig framework test (config string -> WorkspaceEstimatorConfig field, CPU CI) and adds a provider-test asserting the token parsing (exact-match, unknown -> Auto) and the reader-only invariant (no policy value changes the estimate). 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, #33016). Base retargets toward main as the stack merges. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b95559c4-5cbe-4787-8dd7-ffa45a2157f9
Rename the unreleased option to session.kernel_dispatch_policy. Move the shared enum and parser into a CUDA-independent framework header and retain GQA as the first reader-only consumer. Move token tests into the framework suite and preserve estimator behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b95559c4-5cbe-4787-8dd7-ffa45a2157f9
f148b46 to
9f1c163
Compare
Default the optional reader policy explicitly to std::nullopt so partial WorkspaceEstimatorConfig aggregate initializers do not trigger -Wmissing-field-initializers under GCC and Clang. Add a regression for the omitted-policy initialization pattern while preserving reader-only behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b95559c4-5cbe-4787-8dd7-ffa45a2157f9
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
Clean reader-only plumbing for session.kernel_dispatch_policy across WorkspaceEstimatorConfig, CUDAExecutionProvider::GetCapability, and EstimateGroupQueryAttentionWorkspace. Left one suggestion on validating unrecognized session.kernel_dispatch_policy values (instead of silently falling back to Auto) and one minor nit on GetGroupQueryAttentionWorkspaceEstimateConfig.
Reject unknown non-empty session.kernel_dispatch_policy values with INVALID_ARGUMENT before creating accountants, including when partitioning is disabled. Keep empty and auto defaults, retain validated raw tokens, and enforce the CUDA consumer invariant. Forward policy through the GQA config helper and cover parser rejection, factory validation, and node-adapter defaults/forwarding. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 24b20465-fd4d-4a64-8174-aa0622c5164d
Ti-Tai Wang (titaiwangms)
left a comment
Ti-Tai Wang (titaiwangms)
left a comment
There was a problem hiding this comment.
Re-reviewed the isolated a8fda2fe6892f9bd0f318d20dd7cb137f44ae6ea -> ce93b4849f438c1e45061e317e1bc376cf70c38b delta with the full five-role team. No actionable correctness blocker was established.
The revised strict-token contract is consistent: empty/auto, latency, memory and safe are accepted; other tokens fail without overwriting the parser output. In the inspected non-minimal path, validation occurs before accountant enablement, and the CUDA consumer forwards the policy through the shared GQA config adapter. Valid policies remain reader-only here: they do not change route selection, workspace sizing or partitioning. The shared policy header remains graph/provider independent.
Validation is stronger than a green-check inference: exact-head Linux CPU job 117241749368 and CUDA job 117253800039 contain actual OK results for the parser/factory cases, including FactoryRejectsInvalidKernelDispatchPolicy and KernelDispatchPolicyDoesNotEnablePartitioning. The CUDA log also contains actual OK results for NodeAdapterForwardsDispatchPolicy, DispatchPolicyIsReaderOnlyForBothCacheModes and DispatchPolicyDoesNotEnableUnspecifiedEnvelope.
Non-blocking follow-ups: clarify why raw policy text is retained and revalidated on the provider side; clarify the documented invalid-token behavior for minimal sessions that skip CreateAccountants. No minimal-session test was executed in this review. No Critical/Major finding was rejected or deferred for this PR in the new round.
Scope limits: no local build/tests, exhaustive transitive-consumer audit, numerical/performance/concurrency or physical-memory validation. This approval covers the pinned reader delta, not the enforcing child #33316, and does not imply a workload-admission or whole-session no-OOM guarantee.
Description
Add reader-only plumbing for the hardware-neutral session option
session.kernel_dispatch_policy.The intent is generic across operators and EPs.
KernelDispatchPolicyand the shared, header-onlyTryParseKernelDispatchPolicylive incore/framework/kernel_dispatch_policy.h, with no CUDA or graph dependencies. CUDA GQA is the first consumer, not the owner of the policy.For valid values, no policy changes runtime dispatch, workspace estimates, or partitioning in this PR.
latency,memory, andsaferemain intent metadata for future policy resolution. In particular,memorydoes not reduce the reservation, andsafedoes not select a bounded fallback or provide a no-OOM guarantee.Reader contract
autoresolve toAuto. The other valid tokens are exactlylatency,memory, andsafe.INVALID_ARGUMENT; they never silently resolve toAuto. Validation runs even when partitioning is disabled.GetGroupQueryAttentionWorkspaceEstimateConfigacceptsdispatch_policy = KernelDispatchPolicy::Autoalongside the envelope and initializes both fields in one place. The node-estimator overload and CUDA test bridge forward the policy through that helper.Integration and regression coverage
Preserve #33014/#33016's envelope validation against
[0, INT32_MAX], invalid-value diagnostics, and fail-closed GQA placement when no proven estimate is available. Reuse the exported graph-config helper instead of duplicating node parsing. Keep envelopes, eligibility, and sizing local to GQA; other operators and EPs can adopt the shared intent independently.Framework regressions cover valid-token parsing, rejected tokens (including
Memory,SAFE,memroy, andmemory), unchanged parser output on failure, validation with partitioning enabled and disabled, raw-token retention, unset defaults, and policy-only configuration not enabling partitioning. Existing KV-envelope boundary tests remain intact.CUDA metadata regressions cover default and explicit policy forwarding through the node-config helper/test bridge, identical bytes/backend sets for windowed and non-windowed estimation, and unchanged decline without a non-windowed envelope.
Validation for this review update
git diff --check: passed.Motivation and Context
Provide a generic policy representation and reader without coupling other consumers to GQA or CUDA. Rejecting typos now prevents future enforcement (including #33316's committed policies) from silently falling back to
Autoand bypassing the requested restrictions.This PR builds on #33016, which is now merged, and targets
main. Committed policies, operator/EP-specific resolution, route restrictions, runtime coverage/envelope enforcement, bounded fallback, and reduced policy-bound estimates remain separate follow-ups.