Repository navigation
[CUDA] Single-source remaining GQA preparation scratch sizing - #33118
Conversation
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
CUDA allocation sizing is high-risk and still requires full freshly linked provider CI validation.
Review effort: Balanced
Findings: None
What changed in this PR
Consolidates CUDA GQA preparation scratch sizing into shared, overflow-checked workspace helpers.
Changes:
- Shares cache, sequence-vector, and QKV preprocessing size calculations.
- Removes legacy
GQABufferRequirements. - Adds sizing and CUDA packed-QKV regression coverage.
| File | Description |
|---|---|
group_query_attention.cc |
Uses shared preparation-sizing helpers. |
group_query_attention_workspace.h |
Declares cache and preparation sizing APIs. |
group_query_attention_workspace.cc |
Implements checked shared sizing logic. |
group_query_attention_impl.h |
Removes duplicate sizing implementation. |
group_query_attention_impl.cu |
Updates QKV sizing terminology. |
group_query_attention_workspace_xqa_flash.cc |
Updates preparation terminology. |
group_query_attention_workspace_test.cc |
Adds compatibility, validation, and overflow tests. |
group_query_attention_workspace_xqa_flash_test.cc |
Updates XQA allocation documentation. |
group_query_attention_op_test.cc |
Adds packed-QKV CUDA regressions. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
The base branch was changed.
…ace recipes ComputeInternal sized the memory-efficient (MEA) head-expansion / FMHA-accumulator scratch and the unfused-fallback Q/Y/QK/softmax scratch with inline math that duplicated the Level-1 workspace estimator's recipe builders. Route the runtime allocations through GetGQAMemoryEfficientWorkspaceRecipe and GetGQAUnfusedWorkspaceRecipe so a single source of truth backs both the estimate and the allocation. This is a no-functional-change refactor: the recipe byte counts are identical to the previous inline sizing. MEA expanded-K/V equals sizeof(T)*batch*num_heads*seqlen_present_kv_cache*head_size and the FP32 accumulator equals sizeof(float)*batch*sequence_length*num_heads*head_size, gated exactly as MemoryEfficientAttentionParams::need_workspace(head_size, false) (head_size > 128). The unfused recipe reproduces the 256-byte-aligned Q and Y regions plus the two GetUnfusedAttentionWorkspaceSize-equivalent QK/softmax regions, so total bytes and sub-buffer offsets are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf
…pace recipes Migrate the runtime XQA and FlashAttention scratch sizing in GroupQueryAttention::ComputeInternal onto the shared GQA workspace recipes (GetGQAXqaWorkspaceRecipe / GetGQAFlashWorkspaceRecipe) that already back the Level-1 workspace estimator. This is a no-functional-change refactor: the recipes are byte-for-byte equivalent to the previous inline math, as proven by the existing equivalence tests (MatchesLegacyHelperAcrossSupportedFiniteDomain for XQA, MatchesLegacySplitAndBufferHelpers / DoubleThresholdBoundaryMatchesRuntime for Flash), so runtime allocations are identical. - Build one fully-populated GQAWorkspaceProblem before the XQA block and share it across the XQA, Flash, memory-efficient, and unfused recipes (the previous 6-field problem, defined near the MEA block, is moved up and expanded). - XQA: source internal_scratch_bytes (== GetXQAScratchSize), the optional RoPE Q/K regions, and the dynamic head-sink region from the recipe, preserving the single-allocation layout and pointer wiring (data.xqa_buffer / xqa_buffer_bytes, data.qkv_buffer at rotary_q_offset, data.xqa_head_sink at dynamic_head_sink_offset). - Flash: source softmax LSE and split-accumulator sizes from the recipe (which reproduces get_num_splits_and_buffer_sizes and the fast-decode accumulator correction), keep the three separate GetScratchBuffer calls and the accumulator memsets, and preserve the parameters.num_splits side effect used by the kernel launch and debug info. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf
Share checked cache, sequence-vector, and QKV preparation sizing between runtime allocation and the workspace recipe builders. Keep separate allocations, pointer wiring, dispatch order, and scratch lifetimes. Preserve physical packed rows and original versus staged cache capacities; remove the duplicate GQABufferRequirements implementation. Add legacy-size matrix coverage and packed CUDA regressions for aliasing and staging. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf
caccdc2 to
80292cc
Compare
|
Rebased onto current |
Ti-Tai Wang (titaiwangms)
left a comment
Ti-Tai Wang (titaiwangms)
left a comment
There was a problem hiding this comment.
Reviewed at 80292cc2b4c2cd646fcd9fc84f01058500361232.
No confirmed blocking source finding. The shared helpers preserve the separate runtime allocation strategy and expose checked preparation sizing without imposing the full recipe's admission rules on every runtime call.
I specifically checked the logical-versus-physical cache-row concern: for supported packing, these are equivalent representations, not different byte counts. For example, logical head 128 with INT4/uint8 storage and physical head 64 with packing disabled both give a 64-byte row. With B=1, KV heads=2, C=8 and S=4, staging is 1536 bytes per K/V buffer and preservation is 1024 bytes. I am not reporting that representation difference as a defect.
Updated CI evidence: The completed exact-head Test Linux CUDA x64 Release job 114032128652 in run 37988295839 shows all four new tests actually executing and passing, not skipping:
GroupQueryAttentionTest.CudaPackedPreparationUnfusedGroupQueryAttentionTest.CudaWindowedPackedPreparationUnfusedGroupQueryAttentionTest.CudaPackedPreparationFlashGroupQueryAttentionTest.CudaWindowedPackedPreparationFlash
No local ORT build or tests were run. These FP16 fixtures do not establish exhaustive INT4/INT8/FP8 or formerly accepted odd-head compatibility, and this review does not validate the combined preparation/planner layout after stacking with #32071.
|
Thanks for the exact-head review and CUDA test confirmation. #32071 is now closed as a broad, outdated combined-tree integration and will be reconstructed as focused follow-up work after this sizing/estimation stack settles. Accordingly, this PR does not claim or depend on validation of the old #32071 combined layout. No source change is needed for the findings in this review; #33118 remains the standalone preparation-sizing consolidation described in its scope. |
Ti-Tai Wang (titaiwangms)
left a comment
Ti-Tai Wang (titaiwangms)
left a comment
There was a problem hiding this comment.
Approved the source changes at 80292cc2b4c2cd646fcd9fc84f01058500361232.
The head is unchanged from my previous review, which established no blocking source finding. CUDA build/test coverage is now complete in the current check snapshot, including the Linux and Windows CUDA/plugin legs. The exact-head Linux CUDA test log also confirms that all four new packed/separate preparation regressions executed and passed, rather than skipping.
This approval is for the preparation-sizing refactor, not policy enforcement or the combined layout after stacking with #32071. No local ORT build, sanitizer, or exhaustive quantized/odd-head runtime matrix was run.
CI caveat: webgpu_plugin_build_x64_RelWithDebInfo is still failing. Its log shows D3D12 command-queue/device initialization failing with DXGI_ERROR_DEVICE_REMOVED in SliceTest.Slice2D_WithPositiveSteps_2, following the preceding Slice test. Those WebGPU/Slice paths are not modified by this PR. I have not established the device-removal root cause or a successful retry; this outstanding check still requires separate triage or rerun and is not being represented as green by this approval.
…oute-aware L1 estimate) (#33016) ## Summary Adds opt-in, route-aware Level-1 workspace estimation for non-windowed CUDA GroupQueryAttention using the declared total-KV envelope from `ep.cuda.gqa_workspace_max_total_sequence_length`. The present-cache capacity bound is **`max(allocated past-cache capacity, declared total-KV envelope)`**. Static caches may be larger than the active sequence: GQA permits `total_sequence_length <= past_sequence_length` and retains that allocated capacity. Bounding only the envelope would underestimate MEA expansion and partial-alias preservation. The existing graph-free aggregate sizes reachable routes and admitted prompt/decode phase corners; this PR does not introduce a second sizing engine. ## Envelope validation and caller contract - Explicit values are parsed once during session initialization and must be decimal integers in **`[0, INT32_MAX]`**, matching the runtime ABI. Malformed, negative, overflowing, and greater-than-`INT32_MAX` values fail initialization. - Zero remains the unspecified value. With the envelope unset or zero, non-windowed estimation is unavailable; windowed estimation continues to use its cache-shape bound. - The envelope must cover the query-length bound, and the resulting cache-capacity bound must fit the runtime int32 ABI. - The query/max-shape and total-KV declarations are estimation inputs, **not runtime-enforced admission limits or a no-OOM guarantee**. The proof applies only within the caller-declared workload envelope; runtime enforcement remains follow-up work. ## Workspace accounting and proof policy - Capacity-sensitive preparation and MEA expansion cover the full static-cache extent, even when it exceeds the active total-KV envelope. - Partial aliasing copies the full past tensor. Its preservation scratch coexists with the chosen route and is added above the per-route maximum, bounded by the full cache-capacity bound. - Preparation recipes already account for sequence-length vectors; there is no new missing-vector gap. - Reachable cuDNN without a workspace oracle, unsupported attention-bias cases, and otherwise unprovable route domains conservatively make the estimator unavailable. - For capacity-aware CUDA placement, every GQA candidate must have a successful analytical estimate. If proof is unavailable, CUDA assignment for that node is declined. - The generic accountant fallback is an ad hoc safety margin, not a proven GQA bound. Profile data is empirical and likewise cannot substitute for the missing analytical proof. Unrelated operators retain their existing generic fallback behavior. ## Regression coverage - Accepted boundary value `INT32_MAX`; rejected `INT32_MAX + 1`, `INT64_MAX`, decimal overflow, malformed text, and negative values. - Unset-envelope behavior and monotonicity when the envelope increases. - Static past capacity **1024** with total-KV envelope **512**, checking full preservation and MEA expansion against complete concrete runtime recipes. - An envelope below the query-length bound still declines, even when the static cache is larger. - Static capacity outside the runtime int32 ABI declines. - Parsed-envelope plumbing and unchanged windowed estimates. - Additive full-past preservation above the per-route maximum. - Capacity-aware CUDA placement declines GQA when the analytical estimator is unavailable, including when profile data exists; unrelated operators retain generic fallback. ## Dependencies Targets `main`, which now includes #33118's shared preparation-sizing implementation. The duplicated sizing commits were dropped during the post-merge rebase; this PR contains only its envelope-estimation, hardening, test, and documentation changes. Part of the policy-based workspace-estimation series described in #32944. Committed per-`(node, phase)` route policy and runtime workload-envelope enforcement remain follow-up work. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Chi Lo <lochi@microsoft.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf Copilot-Session: c132130a-b0c7-422c-9ff7-75826b6ea3fc 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.
Solid consolidation of the remaining GQA preparation scratch calculations into shared GetGQACacheRowBytes, GetGQACachePreparationSizes, GetGQASequenceLengthsSize, and GetGQAQkvPreprocessBytes helpers, with thorough parity tests against the legacy formulas. Left two minor suggestions on GetGQACacheRowBytes packing validation and cache_sizes.effective_kv_cache_capacity reuse.
| // Preservation and staging size physical tensor rows. Compaction uses the | ||
| // packed-row width from the cache attributes, matching its copy kernel. | ||
| auto cache_row_problem = workspace_problem; | ||
| if (workspace_problem.requires_separate_past_buffer) { | ||
| cache_row_problem.head_size = shared_past->Shape()[3]; | ||
| cache_row_problem.kv_cache_bit_width = 0; | ||
| } else if (parameters.is_windowed_kv_cache && parameters.sequence_length > 1) { | ||
| cache_row_problem.head_size = dense_head_size; | ||
| cache_row_problem.kv_cache_bit_width = 0; | ||
| } | ||
| size_t cache_row_bytes = 0; | ||
| const auto cache_row_status = GetGQACacheRowBytes(cache_row_problem, cache_row_bytes); | ||
| ORT_RETURN_IF_NOT(cache_row_status.IsOK(), |
There was a problem hiding this comment.
Suggestion: enforce problem.kv_cache_bit_width == 0 || problem.cache_element_size == 1 in GetGQACacheRowBytes to eliminate the cache_row_problem overrides.
ValidateProblem (group_query_attention_workspace.cc:151) and BuildBounds both reject kv_cache_bit_width != 0 && cache_element_size != 1, whereas GetGQACacheRowBytes allows kv_cache_bit_width of 4 or 8 with cache_element_size == 2. Whenever kv_cache_bit_width == 0 || cache_element_size == 1 holds (and head_size % 8 == 0), shared_past->Shape()[3] * sizeof(U), dense_head_size * sizeof(U), and GetGQACacheRowBytes(workspace_problem, ...) produce the exact same row byte count. Enforcing problem.kv_cache_bit_width == 0 || problem.cache_element_size == 1 inside GetGQACacheRowBytes would catch invalid 2-byte packed configs at runtime as well and let ComputeInternal pass workspace_problem directly without mutating cache_row_problem.
| size_t qkv_preprocess_bytes = 0; | ||
| const auto qkv_status = GetGQAQkvPreprocessBytes( | ||
| workspace_problem, preparation_route, parameters.seqlen_present_kv_cache, qkv_preprocess_bytes); | ||
| ORT_RETURN_IF_NOT(qkv_status.IsOK(), | ||
| "GQA QKV preparation sizing failed: ", qkv_status.message); |
There was a problem hiding this comment.
Nit: pass cache_sizes.effective_kv_cache_capacity to GetGQAQkvPreprocessBytes.
GetGQAPreparationRecipe (group_query_attention_workspace.cc:746) passes recipe.effective_kv_cache_capacity (computed by GetGQACachePreparationSizes) into GetGQAQkvPreprocessBytes. Passing cache_sizes.effective_kv_cache_capacity here instead of the mutated parameters.seqlen_present_kv_cache makes the runtime dataflow identical to GetGQAPreparationRecipe.
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.ComputeInternalandGetGQAPreparationRecipe.GQABufferRequirementsimplementation.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
CtoC + S.Coverage and validation
/W4 /WXcompilation of the actual CUDA GQA runtime translation unit in both shared-provider and plugin-adapter configurations, with Flash, MEA, FP8 and INT4 enabled./W4 /WXcompilation of the graph-time estimator and XQA/Flash sizing-test translation units./W4 /WXcompilation of the CUDA GQA operator-test translation unit.clang-format --dry-run --Werrorandgit diff --checkpassed.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.