Add GQA XQA and Flash workspace recipes

[Ti-Tai Wang (titaiwangms)]

Summary

Stacked on #32446, add graph-free checked workspace recipes for the CUDA GroupQueryAttention decode backends:

• reproduce XQA semaphore and multi-block scratch sizing from plain device/shape scalars;
• retain XQA RoPE Q/K and dynamic FP32 head-sink conversion allocation parity;
• reproduce Flash regular and fast-decode LSE/split-accumulator sizing;
• preserve the GQA fast-decode distinction between KV heads used for split selection and query heads used for accumulator allocation;
• fail closed for unsupported XQA cache/storage configurations.

Flash bounded-estimation constraint

Flash workspace is not monotonic in KV length. For B=1, S_q=1, two heads, head size 64, and 108 SMs, increasing KV length from 13,824 to 13,825 changes the selected split count from 54 to 28 and reduces workspace. Future Level-1/Level-2 aggregation must compute a conservative envelope or report unavailable rather than evaluating only the maximum shape.

The split heuristic matches runtime double-threshold semantics, including a regression boundary where using 0.85f would select 17 splits instead of the runtime 20.

Scope

This PR adds concrete selected-backend recipes only. It does not add MEA/unfused recipes, complete-route composition, dynamic-bound route reachability, L1/L2 adapters, cuDNN workspace queries, runtime allocation changes, or #32071 planner integration.

Validation

• CUDA provider test target built successfully.
• 201 CUDA internal tests passed; 2 unrelated Lean Attention tests skipped.
• 44 GQA preparation/XQA/Flash tests passed, including runtime-helper parity and Flash discontinuity regressions.

Tracking: #29775
Dependency: #32446

This draft temporarily targets main because GitHub cannot use a fork-only branch as the base of an upstream PR. Its diff will reduce to this commit after #32446 merges.

[Copilot]

🔵 Needs a closer look

Low-level CUDA workspace parity and the stacked dependency warrant final human review despite comprehensive tests.

Pull request overview

Adds graph-free, checked workspace recipes for CUDA GroupQueryAttention preparation, XQA, and Flash backends. No actionable issues were identified.

Changes:

• Models preparation, XQA, and Flash workspace layouts with checked arithmetic.
• Adds runtime-parity, validation, overflow, and discontinuity tests.
• Documents non-monotonic Flash workspace sizing and validates plugin-safe header inclusion.

File summaries

File	Description
group_query_attention_workspace.h	Defines recipe APIs and data structures.
group_query_attention_workspace.cc	Implements preparation workspace sizing.
group_query_attention_workspace_xqa_flash.cc	Implements XQA and Flash recipes.
group_query_attention_workspace_test.cc	Tests preparation recipes.
group_query_attention_workspace_xqa_flash_test.cc	Tests backend parity and boundaries.
group_query_attention_workspace_header_test.cc	Checks isolated header compatibility.
attention_workspace_estimation.md	Documents recipes and non-monotonic sizing.
onnxruntime_unittests.cmake	Adds the plugin header test target source.

Review details

• Files reviewed: 8/8 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)]

XQA launches dim3{nbSubSeqPerSeq, nbKHeads, batchSize} , so  kv_num_heads  and  batch_size  become  grid.y  and  grid.z , respectively. CUDA limits both axes to 65,535, but this currently checks only their product against  UINT32_MAX . Please reject either dimension above 65,535 and add boundary tests

[Copilot]

🟡 Changes recommended

Add the FP8 build-capability gate and validate Flash fast-decode route eligibility.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_xqa_flash.cc:449

• fast_decode changes the split heuristic here without validating that the runtime can actually select Flash fast decode. For example, is_windowed_kv_cache, quantized K/V, QK-Norm, or is_first_prompt makes data.use_flash_attention_fast_decode false in group_query_attention.cc:717, while this function still returns an OK fast-decode recipe. Mirror the existing preparation-route validation (or require an already validated route) so the recipe cannot describe the wrong backend allocation for these inputs.

  GQAFlashWorkspaceRecipe recipe;
  recipe.split_heuristic_head_count = static_cast<size_t>(
      config.fast_decode ? problem.kv_num_heads : problem.num_heads);
  recipe.split_heuristic_kv_length = static_cast<size_t>(config.total_sequence_length);
  if (config.fast_decode && config.local_window_size > 0) {
    recipe.split_heuristic_kv_length = std::min(
        recipe.split_heuristic_kv_length,
        static_cast<size_t>(config.local_window_size));
  }

• Files reviewed: 7/7 changed files
• Comments generated: 1
• Review effort level: Lite (auto)

Note

Copilot is running an experiment and ran this review at Lite.

[Ti-Tai Wang (titaiwangms)]

Addressed the current review and lint findings in 01d6082:

• enforce the CUDA XQA grid-y/grid-z 65,535 limits independently, with exact boundary tests;
• gate FP8 XQA recipes on USE_FP8_KV_CACHE as well as SM capability, with FP8-on/off compile coverage;
• reject available facts that contradict a claimed final Flash fast-decode route while preserving regular Flash;
• apply the clang-format ordering required by Python format CI.

The XQA/Flash source and tests compile with FP8 enabled and disabled, lintrunner is clean, and fresh linked smoke tests exercised the new status paths. The prior Suggest fixes failure was secondary to the lint suggestion API permission error.