[Ti-Tai Wang (titaiwangms)]

Summary

• Size sliding-window GroupQueryAttention Flash split workspace using the resident or staged KV extent instead of the unbounded absolute sequence length.
• Size and launch the unfused fallback with the same effective KV length.
• Preserve existing non-windowed behavior and all dispatch decisions.
• Add host coverage for the effective-length calculation and CUDA regression coverage for T >> C on unfused and multi-token staged Flash routes.

Stack

Depends on #32454. The new change for this PR is commit c9a579a6e8; earlier commits belong to the existing GQA workspace-recipe stack.

This is a runtime allocation-reduction prerequisite for the subsequent GQA Level-1/Level-2 workspace-estimation PR. It does not add planner opt-in, planned-pointer access, root slicing, or estimator wiring.

Validation

• CUDA provider test target built successfully.
• CUDA internal module: 239 passed, 2 unrelated LeanAttention tests skipped.
• Targeted GQA runtime tests: 5 passed, including windowed unfused T=64, C=8 and staged Flash T=64, C=8, S=3.
• Changed-file lintrunner passed.
• git diff --check passed.

[Copilot]

🟡 Changes recommended

The Flash regression geometry never changes its split workspace, so it cannot detect restoration of unbounded sizing.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Bounds windowed GQA Flash and unfused workspace usage to the resident or staged KV extent.

Changes:

• Adds effective KV-length calculation and applies it to runtime sizing and launches.
• Completes graph-free GQA workspace recipes.
• Adds host/CUDA tests and design documentation.

File summaries

File	Description
group_query_attention_workspace_test.cc	Tests preparation and effective KV lengths.
group_query_attention_workspace_mea_unfused_test.cc	Tests backend and complete-route recipes.
group_query_attention_workspace_header_test.cc	Extends standalone-header checks.
group_query_attention_op_test.cc	Adds windowed CUDA regressions.
group_query_attention.cc	Bounds Flash and unfused allocations.
group_query_attention_workspace.h	Defines workspace APIs and recipe types.
group_query_attention_workspace.cc	Models preparation and past preservation.
group_query_attention_workspace_routes.cc	Composes selected-route recipes.
group_query_attention_workspace_mea_unfused.cc	Implements MEA and unfused recipes.
group_query_attention_impl.cu	Launches unfused attention with effective length.
attention_workspace_estimation.md	Documents GQA workspace composition.

Review details

• Files reviewed: 11/11 changed files
• Comments generated: 1
• Review effort level: Balanced

---

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

[Copilot]

🟡 Changes recommended

The new single-token windowed regression calculates an expected cache layout inconsistent with the runtime’s drifting eviction algorithm.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 12/12 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟢 Approval recommended

The bounded runtime dimensions, recipe calculations, validation, and regression coverage are consistent with the physical cache layout.

Review details

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

[Chi Lo (chilo-ms)]

Found two actionable issues:

1. [P2] Complete recipes do not apply the runtime’s windowed KV bound —  onnxruntime/contrib_ops/cuda/bert/group_query_attention_workspace_routes.cc:131-150 . Flash receives  route.flash.total_sequence_length  unchanged, and unfused receives  route.unfused.total_sequence_length , while runtime first applies  min(total_sequence_length, resident/staged capacity)  at  group_query_attention.cc:509-512 . For  T=257, C=8, S=3 , runtime sizes against 11 rows, but a naturally populated complete route sizes against 257, defeating exact recipe parity and potentially rejecting valid constrained-memory plans. Normalize both backend lengths inside  GetGQACompleteWorkspaceRecipe  using  preparation.recipe.effective_kv_cache_capacity , and add windowed Flash/unfused parity tests.
2. [P2] The unfused regression does not verify the allocation reduction —  onnxruntime/test/contrib_ops/group_query_attention_op_test.cc:3416 . Restoring the old  S_kv = total_sequence_length  allocation still produces identical outputs, so this test would pass despite regressing the PR’s main unfused behavior. Populate the existing  AttentionKernelDebugInfo::effective_kv_length_bound  field and assert  EffectiveKvLengthBound=8 , or verify the request through an instrumented/constrained allocator.

[Ti-Tai Wang (titaiwangms)]

Addressed both P2 findings in fbc8dadc70 (rebased onto current main):

1. GetGQACompleteWorkspaceRecipe now applies the same windowed effective-KV-length policy as runtime before composing Flash and unfused recipes, using preparation's checked resident/staged capacity. The tests cover all policy branches for both backends: T > C+S, single-token C, T < C+S, and non-windowed raw T.
2. GQA now emits EffectiveKvLengthBound from the actual runtime allocation extent. The CUDA regressions assert 8 for the single-token unfused case and 11 for the multi-token Flash case, in addition to the existing Flash split assertion.

Validation:

• lintrunner and git diff --check: passed
• targeted CUDA GQA runtime regressions: 5/5 passed
• CUDA internal workspace suite: the new complete-route parity test passed; 247/250 passed overall, with 2 existing LeanAttention skips and the unrelated AllocatorTest.CUDAAllocatorFallbackTest failing because it attempted a 42.5 GB allocation.

[Copilot]

🟢 Approval recommended

The bounded sizing matches runtime cache semantics and is covered by targeted host and CUDA regression tests.

Review details

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