onnxruntime
bc8e7ed7 - Add GQA workspace estimation (#32617)

Commit
2 days ago
Add GQA workspace estimation (#32617) ## Description Adds conservative Level-1 and Level-2 workspace estimation for CUDA contrib `GroupQueryAttention`, tracking #29775. This is PR4, the final GQA workspace-estimation PR in the current stack. It depends on #32602, which bounds windowed runtime workspace by the resident/staged KV extent. The estimator: - inventories all reachable GQA CUDA backend routes over dynamic shape bounds; - composes checked preparation and backend workspace recipes; - uses `C` for single-token windowed decode and checked `C + S` for multi-token staging; - bounds Flash split workspace independently of the non-monotonic split heuristic; - reports non-windowed and reachable cuDNN cases as unavailable when no sound graph-free bound exists; - reports Level-1 `runtime_workspace_bytes`; and - declares one operator-owned slot-0 Level-2 root with 256-byte alignment. This PR does not change runtime allocation topology, opt into planner-owned allocation, or slice a planned root. Those remain follow-up integration work. ## Validation - Fresh Release CUDA build on A100 (SM80) - CUDA internal module: 264 passed, 2 unrelated LeanAttention skips - GQA workspace estimate/bounds tests: 18/18 passed - GQA workspace recipe/arithmetic tests: 19/19 passed - Lintrunner and `git diff --check` passed - Deep route-dominance review swept 2.7M complete routes and 25.6M Flash boundary points with no underestimation ## Stack - Depends on #32602 - Tracks #29775 --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading