Skip to content

Add GQA XQA and Flash workspace recipes - #32453

Draft
Ti-Tai Wang (titaiwangms) wants to merge 2 commits into
microsoft:mainfrom
titaiwangms:titaiwangms/gqa-workspace-xqa-flash
Draft

Add GQA XQA and Flash workspace recipes#32453
Ti-Tai Wang (titaiwangms) wants to merge 2 commits into
microsoft:mainfrom
titaiwangms:titaiwangms/gqa-workspace-xqa-flash

Conversation

@titaiwangms

Copy link
Copy Markdown
Contributor

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 AI and others added 2 commits September 4, 2026 21:10
Introduce graph-free checked workspace recipes for GroupQueryAttention preparation buffers, including windowed cache staging, compaction, sequence metadata, and QKV preprocessing layouts.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: b24d04bb-12ab-41a0-8941-76eb9769005d
Add checked graph-free XQA and Flash workspace recipes with Flash split parity and documented non-monotonicity.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: b24d04bb-12ab-41a0-8941-76eb9769005d
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants