onnxruntime
8ed3bea2 - Add GQA XQA and Flash workspace recipes (#32453)

Commit
7 days ago
Add GQA XQA and Flash workspace recipes (#32453) ## 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. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b24d04bb-12ab-41a0-8941-76eb9769005d
Author
Parents
Loading