[Chi Lo (chilo-ms)]

Description

Resolve session.kernel_dispatch_policy for CUDA GroupQueryAttention into runtime-enforced route masks for prefill, decode, and chunked prefill.

• auto preserves existing phase-applicable dispatch behavior.
• latency preserves the existing per-phase priority: XQA -> cuDNN -> Flash/fast-decode -> MEA -> unfused.
• memory independently selects the minimum-workspace covering route set for each phase, including a phase-local covering fallback when needed. Decode does not inherit higher-workspace prefill selections.
• safe restricts dispatch to routes with analytical workspace recipes and excludes cuDNN. It remains usable without a declared shape/envelope workspace bound; this route restriction is not itself a workspace-bound commitment.

Every runtime dispatch branch checks the current phase's permitted mask. A non-auto policy returns an explicit phase-specific error when no permitted route can serve the run, rather than escaping to an unaccounted route.

Accounting and eligibility

• L1 and L2 resolve the same policy and budget the maximum across admitted phases and their permitted routes, not the sum of phase workspaces. L2 commits the masks after kernel creation and PrePack. The code now documents initialization-before-Run publication and the absence of concurrent re-declaration support; the sized set is the union of committed routes, not every existentially reachable route.
• Exact head geometry from executable-graph metadata is kept separate from max-shape hints. Flash reachability and coverage use the matching FP16/BF16 predicate; exact-static heads are checked directly, while unknown heads retain the conservative admitted-domain sweep, including ORT_QUICK_BUILD restrictions.
• Runtime-only rejection, such as non-aliasing caches disabling XQA, may use only a fallback explicitly committed for that phase. Decode-only routes do not leak into prefill or chunked prefill.
• memory requires successful Level-2 commitment and is rejected at GQA initialization in CUDA-plugin, CUDA-minimal, and ORT-minimal configurations where that commitment cannot run. safe retains its constructor-installed route restrictions if estimation is unavailable; capacity-aware placement still requires a proven estimate.
• cuDNN has no graph-free workspace oracle: auto/latency preserve its dispatch behavior, while strict capacity-aware accounting declines an estimate when it is reachable.

Regression coverage

Added scalar, node-adapter, Level-2, plugin-initialization, and production InferenceSession::Run() regressions for phase-local decode selection, exact-static head-128 Flash coverage across FP16/BF16 and quantized/nonquantized caches, quick-build versus unknown-head behavior, and Safe with missing envelopes or dynamic sequence shapes. Runtime cases assert CUDA assignment, permitted masks, selected routes, numerical output, and explicit no-permitted-route failures. Safe/cuDNN cases include an Auto control and explicit unsupported-device/build/version skips. A focused initialization regression also verifies that non-windowed Memory without an envelope identifies the required session option explicitly.

Latest review fixture corrections

The five reported CI failures are addressed without weakening all-admitted-phase accounting. H512 rotary/sink sizing now uses explicit decode-only bounds, and a separate regression preserves estimate rejection when prefill has no route under Auto, Latency, Memory, or Safe. The Level-1 budget and Level-2 prepacking fixtures use static head128 with Flash enabled for general phases (including quick builds) and a window64 that keeps XQA conversion savings observable. The single-token Unfused-only fixture disables XQA explicitly. The exclusive-route maximum fixture uses KV2048 so its Unfused decode workspace exceeds XQA's at query length one.

Latest revision: repository lintrunner and whitespace checks passed on the complete 15-file C++ PR delta. Static integer-formula checks also confirmed the corrected workspace inequalities and prepacking sensitivity; these are not C++/CUDA execution results. No local builds or C++/GPU test execution were performed, per request; new-head CI must confirm the corrections.

Validation

Latest base merge ea7f7524fa: lintrunner on all 15 PR C++ files and git diff --check origin/main -- include onnxruntime / git diff --cached --check origin/main -- include onnxruntime passed. Source checks confirm both conflicting test bodies are unchanged, all framework regressions from both parents are retained (accounting for the PR's existing Reader-to-Value test rename), and the prior GQA fixture corrections remain unchanged. Main's MatMulNBits/packed-attention placement block, strict dispatch-policy parsing, and chunk forwarding are preserved. No local build or C++/GPU test execution was performed; the merged head requires CI.

Passed on the refreshed 15-file C++ PR delta, using the existing isolated environment and CI-pinned clang-format 20.1.8 / lintrunner 0.13.1:

. .\.venv\Scripts\Activate.ps1
$files = @(git diff --name-only origin/main -- include onnxruntime | Where-Object { $_ -match '\.(cc|h)$' } | ForEach-Object { $_.Replace('/', '\') })
& .\.venv\Scripts\lintrunner.exe @files

Result: passed, ok No lint issues.

• git diff --check origin/main -- include onnxruntime - passed.
• git diff --cached --check origin/main -- include onnxruntime before the merge commit - passed.
• Source audit of the policy delta - passed: no conflict markers or duplicate test definitions; strict parser matches merged main; the raw shared parser has no graph/config facade dependency; main's missing-profile Level-1 accounting behavior and the review regressions are retained.

Not run, per request: local ORT/CUDA builds or C++/GPU test execution. Compilation and CUDA/quick-build regression execution require new-head CI confirmation; lint/source checks are not substitutes for those results.

Stack and scope

#33042 has merged. This PR targets main directly and includes main at 6188bc0900, integrated in ea7f7524fa. The merge preserves strict dispatch-policy validation and GQA policy enforcement alongside main's MatMulNBits/packed-attention Level-1 accounting and separate Q/K normalization epsilon updates. Both tests at the conflicting framework-test insertion point are retained. The Level-2 GQA envelope plumbing from #33240 is included. The coordinated design is documented in #32944.

GQA only; other kernels are unchanged. latency is an ordered heuristic, not a calibrated timing model. Shape/KV envelopes remain estimation declarations rather than runtime-enforced input limits. Whole-graph lifetime accounting and allocator/runtime reserves remain separate work; this does not provide a whole-session no-OOM guarantee.

Motivation and Context

Make the selected phase-aware policy enforceable at runtime without admitting an unaccounted fallback. Address review findings that could inflate decode workspace through prefill-route inheritance, reject statically supported Flash-only nodes under quick builds, or inconsistently require a workspace envelope for Safe in ordinary versus plugin/minimal builds.

[Copilot]

🟡 Changes recommended

The memory policy leaves CUDA plugin GQA kernels with permanently empty route masks, causing every run to fail.

1 open finding

• Plugin kernels fail GQA runs with the memory dispatch policy · New

What changed in this PR

Adds phase-aware CUDA GQA dispatch policies tied to workspace accounting and runtime enforcement.

Changes:

• Resolves auto, latency, memory, and safe into phase-specific backend masks.
• Enforces committed masks at runtime and validates GQA KV envelopes.
• Expands policy, phase, accounting, and Level-2 commitment tests.

File	Description
group_query_attention_workspace_header_test.cc	Adds phase and ABI assertions.
group_query_attention_workspace_estimate_test.cc	Tests policy resolution and accounting.
resource_accountant_test.cc	Tests INT32 envelope validation.
cuda_execution_provider.cc	Passes policy into L1 estimation.
resource_accountant.cc	Centralizes envelope parsing.
group_query_attention.h	Stores committed phase masks.
group_query_attention.cc	Enforces masks during dispatch.
group_query_attention_workspace_estimate.h	Documents resolved policy semantics.
group_query_attention_workspace_estimate.cc	Builds phase-aware bounds.
group_query_attention_workspace_bounds.h	Defines phases and committed masks.
group_query_attention_workspace_bounds.cc	Resolves policies and aggregates phase bounds.
onnxruntime_session_options_config_keys.h	Documents policy and envelope behavior.
resource_accountant.h	Exposes the validated envelope parser.
kernel_dispatch_policy.h	Adds policy-name formatting.

🧠 Review effort: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Ti-Tai Wang (titaiwangms)]

Reviewed f148b46dc8c60bbd6d3008b6c99975a5cde3e4da -> 3526b5a8329c2ae8b763bb0acaeaea51bc4f6c7c with the full five-role team. Requesting changes for two implementation defects and one important regression-test gap; inline comments provide triggering conditions and minimal corrections.

ID	Severity	Raised by	Finding	Disposition
F1	Major	Critical, Integration	Host-only parser dependency crosses the ordinary CUDA shared-provider boundary	Confirmed by source/link/bridge inspection; no executed link/load repro
F2	Major	Code	Added tests do not execute production runtime policy gates/fallback/no-route behavior	Confirmed missing-test finding, not a claimed route escape
F5	Major	Deep	Existence-based head reachability becomes an invalid sole-route Memory commitment	Confirmed by pinned integer recipes and runtime predicates; no GPU execution

The F5 head512/envelope64 example is an availability regression (no permitted route despite supported alternatives), not a demonstrated OOM. The declared envelope is still caller-trusted and is not enforced as a runtime input limit.

Minority/disposition notes: Integration's F1 Critical rating is normalized to Major because this is a blocking unmerged build/load regression, not a demonstrated production incident. Readability's proposed Major for unknown policy tokens is rejected as a correctness blocker: fallback to Auto is the explicit parser contract; a typo warning would be a non-blocking improvement. Readability's generic-diagnostics Major is downgraded to Minor; please preserve the missing envelope/shape/oracle reason where practical. Integration retracted the proposed stack-tip merge-conflict/new-regression finding after confirming the bad test access was inherited and normal merging retains the parent fix.

Non-blocking open questions: Critical/Integration ask whether normal/plugin Safe initialization prerequisites intentionally differ when proof is unavailable, and whether a supported ORT_MINIMAL_BUILD configuration can retain GQA while excluding L2 without constructor rejection. Deep/Integration request clarification of L1-to-L2 selection consistency when shape bounds and candidate feasibility change; no additional counterexample beyond F5 was established. Deep also leaves Unfused's BNSH feasibility capture and an inherited head512 XQA recipe question unverified, neither used as a new blocker. CUDA decline/CPU fallback semantics need intent clarification, not universal CPU policy enforcement. Readability's custom-accountant policy drift was not established on the normal factory path; the Latency no-accountant fail-closed question was resolved by the final non-Auto error guard.

Other non-blocking improvements: document constructor/L2/Run mask lifecycle and relaxed-atomic assumptions; distinguish feasible, committed and sized backend sets; explain per-phase aggregation and Auto versus Latency. Failed-XQA-scratch/fallback overlap was not established: the inspected rejection precedes scratch allocation, and cuDNN returns instead of falling through.

Validation/exclusions: no local build, provider-load repro, Run tests, sanitizer, benchmark, concurrency stress or physical-memory/OOM measurement. There is no qualifying exact-head CUDA CI evidence for this child; #33016's framework/plugin results do not validate these new runtime paths. Readability's root coverage was partial; this is not an exhaustive whole-graph accounting or refreshed-stack review. Threat-model coverage concerns availability/resource admission and provider boundaries; no new privilege or information-disclosure issue was established.

[Tianlei Wu (tianleiwu)]

Strong implementation of phase-aware GQA dispatch policy commitment and runtime enforcement, including the separation of existential reachable_backends from full-domain covering_backends and the header-only ParseCudaGqaWorkspaceMaxTotalSequenceLength validator. Left three inline comments on decode_fallbacks unioning in Memory policy decode selection, ORT_QUICK_BUILD / BFloat16 handling in BuildBounds's coverage loop, and Safe policy initialization consistency across normal vs. plugin/minimal builds.

[Ti-Tai Wang (titaiwangms)]

Re-reviewed the actual stacked delta ce93b4849f438c1e45061e317e1bc376cf70c38b -> 8c746850cea0130b98c66d563992ca0118d1db2a with the full five-role team. No new substantiated implementation blocker was established in the bounded review. The previous findings are addressed in the inspected source, but I am leaving a COMMENT rather than approval because the exact-head CUDA build/Run evidence remains outstanding. This PR still targets the #33042 reader branch, not main, and is currently draft.

Prior-finding dispositions:

• Shared-provider parser dependency: addressed by the raw-string/scalar header helper. Host/provider retrieve config locally; the normal CUDA constructor no longer calls the native host ConfigOptions overload. No new framework linkage/bridge is used.
• Existential Flash eligibility used as sole-route coverage: addressed by separating reachable from covering backends, retaining a proven phase-local fallback, and using executable-graph static-head provenance and the correct FP16/BF16 predicates. The previous unknown-head-bound-512 counterexample now retains Flash|Unfused with the 213,248-byte bound; exact static geometry is checked directly rather than inferred from a smaller eligible head.
• Missing production enforcement tests: coverage added. Tests now call InferenceSession::Run and check CUDA assignment, phase/mask/selected-route diagnostics, independent output-cache alias rejection, output or intentional no-route failure, and Auto/Safe controls.
• Decode fallback and initialization: decode commits its own selected.backends, not prefill selections. Safe may retain route restrictions without a proven workspace bound; capacity-aware placement still requires proof. Memory rejects unavailable commitment, including the inspected plugin/minimal cases.

Validation limitation: this exact head has only seven lint/docs/policy checks, with no qualifying CUDA build or test result. QA verified that the requested selectors exist, but none of the six available local binaries has proven reviewed-head provenance. No stale binary was executed as evidence and no full rebuild was started. The needed production batch is:

GTEST_FILTER='GroupQueryAttentionDispatchPolicyTest.*:GroupQueryAttentionWorkspaceEstimateTest.KernelCommits*' \
GTEST_COLOR=no timeout 300s \
./onnxruntime_provider_test --gtest_filter=CUDA_EP_Unittest.All

This requires an exact-head nonminimal/non-plugin CUDA internal-test build; Flash/MEA/cuDNN controls require their corresponding build/device support. Please retain actual OK versus SKIPPED versus compile-excluded distinctions. Source assertions do not establish numerical output, actual backend dispatch, or link/load success.

Non-blocking questions/follow-ups: consider nonuniform inputs and returned-cache assertions in a follow-up numerical/cache-sensitive case; clarify Auto versus Latency and the constructor/L2/Run mask lifecycle/commitment side effect; document reachable/covering/committed/permitted terminology. The explicit-empty numeric-envelope versus empty-policy difference is inherited contract behavior, not a new blocker. Cross-build initialization and token-validation behavior still need relevant execution coverage.

No Critical/Major finding was overruled for this PR in the new round. The Critical lane covered availability/resource boundaries and opt-in diagnostic metadata, not a broader exploit audit. Excluded: local build/Run/sanitizer, measured VRAM/OOM, numerical/performance/concurrency certification and exhaustive transitive consumers. Level-2 is still verification/commitment, not an automatic repartition loop; the caller-declared envelope remains an estimation input rather than runtime admission enforcement or a whole-session no-OOM guarantee. This comment records source-level closure of the previous issues, not runtime clearance or dismissal of the previous review.

[Chi Lo (chilo-ms)]

Ready for re-review at 8fc52ed. Merged current main after #33042 landed and resolved the conflicts while preserving strict token validation, main's missing-profile Level-1 accounting, and all seven inline-comment fixes/regressions.

Also documented initialization-before-Run / relaxed-atomic mask assumptions and max-of-phase workspace aggregation, and added a focused initialization regression for the explicit missing non-windowed Memory-envelope diagnostic.

Repository lintrunner passed on the final 15-file C++ delta, as did policy-delta whitespace and source-invariant checks. No local builds or C++/GPU test execution were performed, per request; CUDA/quick-build execution still requires new-head CI. The four original threads remain open for reviewer confirmation; their existing fix explanations are retained. The three later threads were already resolved by the reviewer.

[Chi Lo (chilo-ms)]

CI fixes pushed in e1ec319.

• Ordinary CUDA and TensorRT jobs on both Linux and Windows diagnosed native/provider logging type and macro redefinitions in group_query_attention.cc. Removed the direct host logging.h include; cuda_common.h already supplies the appropriate provider logging surface.
• Windows CUDA also diagnosed the new dispatch test's braced output-name list failing conversion to gsl::span<const std::string>. The test now passes a typed std::array<std::string, 3> with unchanged output order and assertions.
• The separate Linux no-cuDNN failure was an HTTP 504 downloading cpp_client_telemetry during configure, before compilation. No unrelated dependency/workflow workaround was added.

Passed: . .\.venv\Scripts\Activate.ps1; & .\.venv\Scripts\lintrunner.exe onnxruntime\contrib_ops\cuda\bert\group_query_attention.cc onnxruntime\test\providers\cuda\test_cases\group_query_attention_workspace_estimate_test.cc and staged/working diff whitespace checks. No local builds or C++/GPU execution were performed, per request. The push starts new-head CI, including a fresh attempt of the no-cuDNN dependency download; no stale-head reruns were queued. Compilation and regression execution still need CI confirmation.

[Tianlei Wu (tianleiwu)]

Thanks for addressing the three findings from the previous round (Memory decode fallback independence, static_head_size / BFloat16 Flash coverage under ORT_QUICK_BUILD, and Safe policy retention without a workspace bound).

In the current CUDA CI runs (Test Linux CUDA x64 Release, Windows GPU CUDA CI Pipeline Test Job, Test Linux TensorRT x64 Release), 5 unit tests in CUDA_EP_Unittest.All are failing:

1. GroupQueryAttentionWorkspaceEstimateTest.H512EstimatorDominatesCompleteRuntimeWorkspace (line 542), GroupQueryAttentionWorkspaceEstimateTest.GetCapabilityBudgetUsesLevel1Estimate (line 730), and GroupQueryAttentionWorkspaceEstimateTest.KernelDeclaresPrepackedHeadSinkRoot (line 1686):
   • See inline comment on AggregateCommittedPhases in group_query_attention_workspace_bounds.cc:476: when head_sink is present (has_sink == true, excluding Unfused/MEA) and FlashAttention is either disabled (kMath) or does not cover the head size (head = 512, or head = 64 in CI's ORT_QUICK_BUILD=ON where Flash only compiles head = 128), bounds.reachable_backends is {GQAReachableBackend::Xqa} (decode-only), so committed_backends.prefill is None and AggregateCommittedPhases fails with Unavailable.
2. GroupQueryAttentionWorkspaceEstimateTest.SingleTokenShapeDoesNotAdmitChunkedPrefill (line 1001):
   • See inline comment on group_query_attention_workspace_estimate_test.cc:1000: Config() defaults to config.enable_xqa = true, so committed_backends.decode under Safe is GQAReachableBackend::Xqa | GQAReachableBackend::Unfused (0x11) rather than GQAReachableBackend::Unfused (0x10).
3. GroupQueryAttentionWorkspaceBoundsTest.AggregatesExclusiveRoutesWithMax (group_query_attention_workspace_estimate_test.cc:1376):
   • Because GetGQAWorkspaceAggregateForPermittedPhase (group_query_attention_workspace_bounds.cc:282) now sizes GQAExecutionPhase::Decode with sequence_length = 1 instead of bounds.sequence_length_bound (4), unfused_only.total_workspace_bytes at S = 1 (37,120 bytes) is now smaller than xqa_only.total_workspace_bytes (140,030 bytes for 108 SMs), failing ASSERT_GT(unfused_only.total_workspace_bytes, xqa_only.total_workspace_bytes). Increasing xqa.present_kv_cache_capacity_bound (e.g., to 2048) or xqa.batch_size_bound in that test will make unfused_only exceed xqa_only at S = 1.

[Chi Lo (chilo-ms)]

Review-fixture corrections pushed in f0d95174ac, addressing all five failures in the latest review without changing production estimator behavior.

The two inline threads document the H512 decode-only bounds, quick-build-compatible head128 sink fixtures, and explicit XQA disablement in the single-token case.

The remaining review-summary failure, GroupQueryAttentionWorkspaceBoundsTest.AggregatesExclusiveRoutesWithMax, now uses KV capacity2048. At the production Decode phase's query length one, the pinned formulas give Unfused 266,496 bytes versus XQA 140,030 bytes, restoring the test's intended strict inequality while retaining its maximum-not-sum and prepack-lifetime assertions. These byte values are static calculations, not observed GPU allocations.

Validation:

$env:PATH="$PWD\.venv\Scripts;$env:PATH"
$files = @(git diff 314ff8d4e87816bafa1b8e0fefa476c896c4dedd --name-only --diff-filter=ACMR -- include onnxruntime | Where-Object { $_ -match '\.(cc|h)$' } | ForEach-Object { $_.Replace('/', '\') })
& .\.venv\Scripts\lintrunner.exe @files

Result: passed, all 15 C++ files, ok No lint issues.

git diff --check 314ff8d4e87816bafa1b8e0fefa476c896c4dedd -- include onnxruntime and the staged whitespace check: passed.

Static integer-formula checks: passed for the corrected decode inequality and for observable XQA prepack savings / the generous discovery budget over 1-1024 SMs. These checks do not execute production C++ or CUDA.

Not run, per request: local ORT/CUDA builds or C++/GPU tests. The exact new-head CUDA/TensorRT CI results still need to confirm the reported failures are cleared. The four older review threads remain open for reviewer confirmation.

[Chi Lo (chilo-ms)]

Resolved the conflict by merging current main (6188bc0900) in ea7f752. GitHub now reports no merge conflicts.

Kept both additions in resource_accountant_test.cc: the GQA raw-envelope parser regression and main's default chunk-policy regression. The merge also preserves main's MatMulNBits/packed-attention accounting and separate Q/K normalization epsilon updates alongside the GQA policies; the earlier GQA fixture corrections are unchanged.

Formatting for all 15 PR C++ files, staged/unstaged PR whitespace checks, and targeted source merge checks passed. No local build or C++/GPU test execution was performed, per request; new-head CI still needs to confirm the merge.

[Tianlei Wu (tianleiwu)]

Re-reviewed ea7f7524fa: all five GroupQueryAttentionWorkspace* unit test failures from the previous round are cleanly resolved in f0d95174ac while keeping admitted-phase workspace accounting fail-closed, and CUDA CI is green.

[Tianlei Wu (tianleiwu)]

Reviewed the current PR against its current base: no new actionable findings. Rechecked the five prior tianleiwu threads and author replies; all are addressed and already resolved. Current-head Linux CUDA logs show 56 relevant GQA tests passing and the plugin initialization regression passing; rebuilt graph-free helpers passed 4,109 UBSan assertions. Documentation drift, WebAssembly-proxy ArrayBuffer failures, and the canceled React Native job do not establish an introduced CUDA-policy defect. No full local ORT build or production CUDA inference was run.

[Tianlei Wu (tianleiwu)]

Reviewed all 22 changed files at ae324eae42bf305c45a2a32601ac6ca1590fbda5. Requesting changes for one newly identified numerical-correctness issue: Memory can commit MEA for a learned head_sink, replacing its value with zero-logit smoothing. The inline comment gives a concrete admitted geometry and correction.

The five previously owned threads are addressed and already resolved; no duplicate thread actions are needed. Exact-head CI is SUCCESS, including native Linux/Windows CUDA and plugin test jobs, but the added production cases do not cover nonzero sinks with MEA enabled. No local build or CUDA execution was performed; this finding follows from the pinned dispatch, workspace, and softmax implementations.

[Tianlei Wu (tianleiwu)]

[P1] MEA is not a mathematically covering route for an explicit head_sink. IsGQAMemoryEfficientEligible only uses has_sink in the H512 preference; it does not reject the tensor. EfficientAttention never forwards data.head_sink and only enables zero-logit smooth softmax.

For Flash+MEA enabled, FP16 SM80/108 SMs, B1/S4/Hq=Hkv8/head128/cache64/envelope64, separate Q/K/V, no bias/rotary/QK norm, and head_sink=1, Memory ranks MEA (131,084 bytes) below Flash (131,456 bytes) and commits it for prefill. Auto executes Flash. With Q=K=0 and V=1, the final row should be 4/(4+exp(1)) ≈ 0.5954; Memory instead computes 4/5 = 0.8.

Please exclude sink-bearing MEA candidates and enforce that eligibility guard at dispatch, or implement actual per-head sink support. Add a nonzero-sink Memory-versus-Auto Run() regression with both backends enabled. This is a source-derived counterexample, not a local GPU execution result.