Add GQA workspace preparation recipes

[Ti-Tai Wang (titaiwangms)]

Summary

Add the graph-free workspace preparation foundation for CUDA GroupQueryAttention, following the checked-recipe architecture established by #32321.

• model windowed KV-cache staging and compaction without changing runtime allocations;
• model 3-vector/6-vector sequence metadata and Flash fast-decode suppression;
• model route-selected QKV preprocessing requirements;
• use checked arithmetic and one internally 256-byte-aligned preparation layout;
• distinguish transient staging/compaction from output KV cache and persistent/prepacked state.

Scope

This is the first GQA workspace-estimation PR in the #29775 rollout. It intentionally does not add backend-internal workspace recipes, route aggregation, Level-1 estimation, Level-2 declaration, cuDNN workspace queries, or #32071 planner integration. Existing GetScratchBuffer() allocation count, size, lifetime, and pointer layout are unchanged.

Follow-up PRs will add:

1. XQA and Flash regular/fast-decode workspace recipes.
2. CUTLASS MEA and unfused fallback recipes plus the complete route aggregate.
3. Optional-aware Level-1/Level-2 adapters with one 256-byte-aligned operator root.
4. cuDNN support only if an exact runtime-parity workspace query can be established.

Validation

• 30 GQA workspace recipe and validation tests passed.
• 73 related GroupQueryAttention runtime tests passed; 12 WebGPU-only tests skipped because WebGPU EP was unavailable.
• CUDA provider test target built successfully.

Tracking: #29775

[Copilot]

🟢 Approval recommended

The implementation matches runtime allocation formulas and has comprehensive boundary and parity coverage.

Pull request overview

Adds graph-free CUDA GroupQueryAttention workspace recipes with checked sizing and validation.

Changes:

• Models staging, compaction, sequence metadata, and QKV preprocessing.
• Adds overflow, validation, layout, and runtime-parity tests.
• Verifies plugin-safe standalone header compilation.

No actionable issues identified.

File summaries

File	Description
group_query_attention_workspace.h	Defines workspace recipe interfaces.
group_query_attention_workspace.cc	Implements checked sizing and validation.
group_query_attention_workspace_test.cc	Tests recipes and edge cases.
group_query_attention_workspace_header_test.cc	Verifies isolated header compilation.
cmake/onnxruntime_unittests.cmake	Adds the plugin header test.

Review details

• Files reviewed: 5/5 changed files
• Comments generated: 0
• Review effort level: Balanced

---

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

[Chi Lo (chilo-ms)]

Found one blocking issue:  GetGQAPreparationRecipe  accepts impossible XQA routes and may underestimate workspace ( group_query_attention_workspace.cc:127-166, 254-264 ). Runtime XQA requires decode with  sequence_length == 1 , not first prompt, and either both caches unquantized or both K/V quantization types supported ( group_query_attention.cc:489-552 ). The recipe currently permits multi-token/first-prompt XQA and even tests asymmetric quantization ( K=PerTensor ,  V=None ) as valid. For packed multi-token input it returns Q-only storage, while the actual fallback route can require Q+K+V. I recommend validating those available XQA eligibility facts or representing this as preprocessing behavior rather than an XQA route.

[Ti-Tai Wang (titaiwangms)]

Fixed the blocking XQA-route validation issue in 02064f7. GQAPreprocessMode::Xqa now rejects first-prompt and multi-token inputs, asymmetric K/V quantization, packed unquantized caches, non-8-bit quantized caches, and quantized QK-Norm. The former asymmetric/multi-token “valid” fixtures now use runtime-reachable single-token configurations, and parameterized regression tests cover each rejected case. The same fix is also pushed to stacked PRs #32453 and #32454.

[Ti-Tai Wang (titaiwangms)]

Chi Lo (@chilo-ms) PTAL

[Chi Lo (chilo-ms)]

ValidateProblem()  now rejects first-prompt, multi-token, asymmetric quantization, INT4, and quantized QK-Norm XQA routes. However, it still accepts XQA geometries the runtime cannot select. Runtime XQA requires:

•  head_size  ∈  {64, 128, 256} 
• Quantized group size ∈  {4, 8, 16, 32} 
• Unquantized group size ∈  {1, 2, 4, 5, 8, 16, 32} 

The recipe has both  head_size  and head counts available but does not validate these constraints ( group_query_attention_workspace.cc:167-190  versus  group_query_attention.cc:547-568 ). The updated XQA tests even use  head_size = 8 , so they validate routes that runtime always rejects. This can still underestimate preprocessing when actual dispatch falls through to a route such as MEA, particularly for packed QKV.

[Ti-Tai Wang (titaiwangms)]

Thanks — this was correct. Fixed in c0dadf6.

The fix is structural rather than another estimator-only copy of the constraints:

• runtime dispatch and workspace validation now share graph-free predicates for the GQA XQA head-size and quantized/unquantized group-size domains;
• selected-route validation also requires the matching one-byte/8-bit quantized or two-byte/unpacked cache representation;
• XQA tests now start from a runtime-reachable H=64, group-4 decode fixture and cover the complete supported/unsupported finite domains.

This prevents the preparation estimator from accepting a claimed final XQA route that runtime would fall back from and therefore underestimating fallback preprocessing. The same shared-predicate change is synchronized to #32453 (4ac4ce9) and #32454 (8fac283).