[Tianlei Wu (tianleiwu)]

Description

Specifies the sliding_window_cache KV-cache contract for com.microsoft::GroupQueryAttention so that callers can index the cache themselves, and so that multi-token (speculative decoding, chunked prefill) steps are well defined. No kernel behavior changes; this documents and tests the contract the CPU and CUDA kernels already implement.

The schema previously said the operator "keeps only the min(total_sequence_length, cache_capacity) most recent tokens, contiguously, using cache-relative indexing and evicting from the front as needed." That describes minimum-overflow eviction, but both kernels evict in blocks: after the buffer fills, the resident count sawtooths between local_window_size and the capacity instead of staying at min(T, C). A consumer that trusts the old wording reads the buffer at the wrong offsets whenever the capacity exceeds the window.

The contract is now stated explicitly. With capacity C (dimension 2 of past_key), window W = local_window_size, total length T, and G = C - W + 1:

L(T) = T                            if T <= C
L(T) = T - G * ceil((T - C) / G)    otherwise

After a step, rows [0, L) of present_key/present_value hold the L most recent positions in order, so row i holds absolute position T - L + i; rows [L, C) are unspecified. Hence min(T, W) <= L(T) <= min(T, C).

Because L depends only on T, the layout is independent of how tokens were split into steps: a speculative verification step of S tokens leaves exactly the layout the same tokens produce one at a time. Any S >= 1 is accepted, including S > C, via internal staging. Rejecting k draft tokens requires no cache edit and is exact when L(T - k) == L(T) - k, which callers can evaluate from the published formula; otherwise the positions the shorter layout needs have already been evicted. Sizing C >= W + S - 1 — the rule the ONNX Runtime GenAI model builder already applies when it sizes a windowed paged ring from the prefill chunk size — removes the staging copy and widens the range of exact rollbacks.

The attention_bias last dimension is clarified as well: it stays total_sequence_length in absolute key coordinates under a windowed cache. It is not reduced to the capacity or to local_window_size; the kernel selects the columns of the resident positions.

Summary of Changes

File	Change
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Adds a Windowed KV Cache section to the operator description covering capacity, resident range, eviction, multi-token steps and rollback. Rewrites the sliding_window_cache attribute text, and clarifies past_key, past_value and attention_bias.
docs/ContribOperators.md	Regenerated for the schema strings above.
docs/contrib_ops/cpu/gqa.md	New normative Windowed (Sliding-Window) KV Cache section: notation, resident-range formula with a worked table, speculative-decoding guidance, the rollback rule, CPU implementation notes (drifting layout, block compaction, staging, absolute attention_bias offsets) and constraints.
docs/contrib_ops/cuda/gqa.md	Points at the normative section, corrects the "evicts from the front" wording, and drops a stale claim that CPU rejects attention_bias with sliding_window_cache=1.
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc	Three CPU tests locking the contract.

New tests, each driving a windowed cache across steps with C = 6, W = 4 (so G = 3, crossing two eviction boundaries) and values that encode their absolute position:

• WindowedCacheResidentRangeMatchesSpec_CPU — checks present_value[0, L) against an independent implementation of the published formula after every step, plus the windowed output value. Fails if the kernel kept min(T, C).
• WindowedCacheMultiTokenStepMatchesSingleTokenSteps_CPU — a 4-token step after an 8-token prefill (which itself exercises the staging path, S > C) produces byte-identical present buffers and outputs to four 1-token steps.
• WindowedCacheRejectedDraftTokensNeedNoCacheEdit_CPU — verifies 4 drafts, accepts 2, resumes from the untouched buffer, and matches the path where only the accepted tokens were ever appended.

Motivation and Context

The eviction policy is observable through present_key/present_value and therefore part of the operator contract, but it was underspecified in a way that affected both EP implementations and external cache manipulation. ONNX Runtime GenAI's CheckWindowedKvCacheRewind documents the resident range as [T - min(T, C), T), i.e. the old wording rather than the implemented policy; it is safe only because it refuses every rewind once the cache fills.

Publishing the layout as a closed-form function of T, C and W gives callers what they need without adding an output that reports the resident count: speculative decoding can compute where each absolute position lives, decide whether a partial-accept rollback is exact, and size the capacity accordingly. It also fixes the reference for EP authors, who must materialize the contiguous layout before returning even if they use a ring internally, since no output carries a ring offset.

Open questions intentionally left as-is: multi-token steps with past context still require batch_size == 1 (an operator-wide restriction, now documented in the windowed section), and rollbacks that cross an eviction boundary remain unsupported because the evicted keys are genuinely gone.

Testing

• Built onnxruntime_provider_test from a fresh CPU-only Release configuration.
• GroupQueryAttentionTest.WindowedCache*: 5 tests passed (3 new, 2 pre-existing).
• GroupQueryAttentionTest.*: 59 passed, 0 failed (CUDA/WebGPU cases skipped in a CPU build).
• Verified docs/ContribOperators.md matches the schema strings byte-for-byte under gen_contrib_doc.py semantics, including the two-space prefix the generator emits on blank documentation lines, so the Doc Gen check stays green.
• Ran lintrunner on all changed files; no lint issues.

Not covered locally: the CUDA kernel implements the same contract, but the new tests are CPU-only and CUDA parity for windowed caches remains covered by TestGQAWindowedKvCache in onnxruntime/test/python/transformers/test_gqa.py.

[Copilot]

🟡 Changes recommended

The documented cross-provider contract conflicts with CUDA behavior, and the tests do not validate key-cache layout.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Documents the GroupQueryAttention windowed KV-cache contract and adds CPU contract tests.

Changes:

• Defines cache layout, eviction, multi-token, rollback, and attention-bias semantics.
• Updates CPU, CUDA, and generated operator documentation.
• Adds CPU tests for layout invariance and rollback.

File summaries

File	Review
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc	Moderate: Keys remain zero, so tests do not validate key compaction or staging. Encode positions and validate present_key.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Moderate: Chunk invariance conflicts with CUDA behavior. Nits: Define lengths per batch element; qualify batch-size restrictions, CUDA staging, and unsupported CUDA windowed attention bias.
docs/ContribOperators.md	Regenerated schema documentation must reflect the corrected contract and restrictions.
docs/contrib_ops/cuda/gqa.md	Moderate: The documented block-eviction layout is not implemented by CUDA. Update the implementation and add direct layout tests, or scope the contract appropriately.
docs/contrib_ops/cpu/gqa.md	Moderate: The CPU formula is incorrectly presented as a shared CUDA contract. Nit: Define and apply T independently per batch element for ragged batches.

Review details

Suppressed comments (3)

docs/contrib_ops/cpu/gqa.md:117

• T must be per batch element. The scalar total_sequence_length is documented by the schema as the batch maximum, while the CPU planner computes each slice from seqlens_k[b] + 1; ragged batches can therefore have different L values and row origins. Define T_b = seqlens_k[b] + 1 here and apply the formulas independently per batch element.

| `T` | `total_sequence_length`, i.e. the absolute number of tokens processed so far including this step. |

onnxruntime/core/graph/contrib_ops/bert_defs.cc:1362

• This describes windowed attention_bias as supported by the operator, but the CUDA EP still calls CheckCustomAttentionInputs without enabling support_windowed_attention_bias (cuda/bert/group_query_attention.cc:334-337), and test_gqa.py:2913-2940 explicitly verifies that this combination is rejected. Models following this schema will therefore fail on CUDA. Qualify the behavior as CPU-only and retain the CUDA restriction in its docs, or implement and test the absolute-column offset on CUDA.

               "The last dimension is indexed by absolute key position and stays total_sequence_length when "
               "sliding_window_cache is 1: it is not reduced to the cache capacity or to local_window_size. The "
               "operator reads the columns of the positions that are resident in the cache and ignores the rest.",

onnxruntime/core/graph/contrib_ops/bert_defs.cc:1224

• “Any S >= 1 is accepted” omits an active operator restriction: CheckInputs rejects sequence_length > 1 with past context when batch_size != 1 (group_query_attention_helper.h:401-417). The detailed CPU notes mention this, but the public schema and attribute text promise unconditional support. Qualify multi-token support with the batch-size restriction in both schema locations.

Because `L` depends only on `T`, the resulting layout is independent of how the tokens were split into steps: a multi-token step of `S` tokens (speculative decoding, chunked prefill) leaves exactly the layout that the same tokens would produce one at a time. Any `S >= 1` is accepted, including `S > C`; a step that would evict positions it still has to read is staged internally, so the capacity does not have to cover the step.

• Files reviewed: 5/5 changed files
• Comments generated: 6
• Review effort level: Balanced

---

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

[kunal-vaishnavi]

For all of these contrib op docs, should we consider consolidating per EP and have one doc per contrib op with sections per EP? I feel it would be easier to maintain and avoid drfit across EP docs. For example, we don't define op schemas per EP. Instead, we have one op schema that is used by all EPs.

[github-actions]

You can commit the suggested changes from lintrunner.