GQA support sliding_window_cache by tianleiwu · Pull Request #29904 · microsoft/onnxruntime

tianleiwu · 2026-07-27T16:14:26Z

Description

Adds a new sliding_window_cache attribute to com.microsoft::GroupQueryAttention so that layers using local (sliding-window) attention can be backed by a window-sized KV cache instead of a full-length one. Today local_window_size narrows the attention mask only — the cache is still allocated for the whole sequence, so a model like gpt-oss pays for max_length positions on every layer even though half of its layers can never look back more than 128 tokens.

The attribute is opt-in and defaults to 0, so every existing model keeps the exact current code path. When set to 1, the operator treats the bound past/present buffer as a fixed-capacity ring of the most recent tokens and does all cache addressing in cache-relative coordinates, while RoPE keeps using absolute positions. Because both the causal mask and the local-window mask depend only on the query/key distance, translating the cache origin does not change the math.

Design

For a windowed layer with capacity C, the invariant maintained is:

the cache holds the L = min(T, C) most recent tokens, contiguously, at cache indices [0, L), where T is the absolute total sequence length.

Appending S new tokens onto past length P:

Lp = min(P, C)                # tokens currently resident
E  = max(0, Lp + S - C)       # tokens evicted by this step
shift:  cache[E .. Lp)  ->  cache[0 .. Lp-E)
write:  new tokens      ->  cache[Lp-E .. Lp-E+S)

Two details are worth calling out:

Everything is derived on device from seqlens_k. is_first_prompt cannot be used here: it is computed from sequence_length == total_sequence_length (input 6), and onnxruntime-genai binds input 6 to the allocated buffer width, not the running total. seqlens_k (input 5) is the true running total, so GetSequenceLengths emits both absolute lengths (for RoPE) and cache-relative lengths (for cache addressing) plus a per-batch eviction count, all on device. This also keeps the launches CUDA-graph capturable.
Multi-token steps run against a staging cache. A naive in-place scheme would require C >= local_window_size + sequence_length - 1, coupling the cache size to the prefill chunk size. Instead any step with S > 1 is seeded into a temporary C + S staging cache and the surviving window is copied back afterwards, which reduces the requirement to just C >= local_window_size. Decode (S = 1) runs in place with a compaction shift of at most one row.

Summary of Changes

Op schema and docs

File	Change
`onnxruntime/core/graph/contrib_ops/bert_defs.cc`	New `sliding_window_cache` INT attribute (default `0`); shape inference keeps `present_*` at the past buffer's own sequence dim instead of growing it to `total_sequence_length`.
`docs/ContribOperators.md`, `docs/contrib_ops/cuda/gqa.md`	Regenerated op docs plus a new "Sliding Window (Windowed) KV Cache" section covering capacity, cache-relative indexing, RoPE bounds, staging, and constraints.

Parameters and validation

File	Change
`onnxruntime/contrib_ops/cpu/bert/attention_parameters.h`	Adds `is_windowed_kv_cache`, `kv_cache_capacity`, `kv_cache_real_capacity`, `rotary_max_position`.
`onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h`	`CheckInputs` takes the new attribute; in windowed mode `present_sequence_length = past_sequence_length` (the bound buffer is the capacity); validates `C >= local_window_size`, `past_present_share_buffer`, and `local_window_size > 0`; adds `rotary_max_position` so RoPE stays bounded when rotary caches are absent. Rejects `attention_bias` combined with a windowed cache.
`onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc`, `js/...`, `webgpu/...`	CPU / JS / WebGPU GQA return `NOT_IMPLEMENTED` for `sliding_window_cache=1` rather than silently computing the wrong result.

CUDA kernel

File	Change
`group_query_attention_impl.cu` / `.h`	`GetSequenceLengths` additionally emits `cache_past_seq_lens`, `cache_total_seq_lens` and `evict_counts`; new fused `LaunchCompactKvCache` (in-place eviction shift) and `LaunchCopyKvCacheWindow` (staging seed / copy-back) row-movement kernels — K and V are moved in a single launch, one row-tile per block.
`group_query_attention_qkv.cuh`	KV append uses the cache-relative offset; RoPE keeps the absolute position.
`group_query_attention.cc` / `.h`, `attention_data.h`	Allocates the staging cache and compaction scratch, wires the windowed buffers through, and disables the Flash-Attention fast-decode path for windowed caches (its RoPE/cache indexing assumptions do not hold).

Tests

File	Change
`onnxruntime/test/python/transformers/test_gqa.py`	New `TestGQAWindowedKvCache` suite (11 tests) plus `test_windowed_cache_attention_bias_is_rejected`, and a `has_fp8_xqa()` SM>=8.9 guard so the fp8 XQA sliding-window parity tests skip on older GPUs instead of failing.

Testing

pytest onnxruntime/test/python/transformers/test_gqa.py -k TestGQAWindowedKvCache
pytest onnxruntime/test/python/transformers/test_gqa.py

TestGQAWindowedKvCache covers: prompt shorter/longer than capacity followed by decode, chunked prefill, eviction boundaries, RoPE, gpt-oss shape with head_sink, int8 quantized cache, batch-1 packed head layout, capacity exactly equal to the window size, and the two negative cases (capacity too small, local_window_size missing).

End-to-end validation on gpt-oss-20b (24 layers, 12 of them sliding with window 128), same weights with and without the attribute:

	GPQA-diamond (198 q, `reasoning=medium`)	KV cache @ `max_length=131072`, int8 KV
full-length cache	0.6010 (119/198)	5138 MiB
windowed cache	0.5960 (118/198)	3090 MiB (-39.9%)

Paired McNemar over the 198 questions gives p = 1.0000 (both correct 102, only-baseline 17, only-windowed 16), i.e. the accuracy difference is not significant. Results are bit-identical to a full-length cache while nothing has been evicted; once the cache origin shifts, flash accumulates the same keys in a different block order, and on an MoE model that rounding difference can flip an expert route — hence the aggregate/paired comparison rather than greedy token parity.

Decode cost of the extra row movement is ~1-2% (measured on A100 with CUDA graphs enabled, prompt 512-65536); the copies are launch-bound rather than bandwidth-bound.

Backward compatibility: sliding_window_cache defaults to 0, which is byte-for-byte the existing code path. No existing model or exported graph changes behavior.

Motivation and Context

Reduce KV cache memory for LLM models that use sliding-window attention (gpt-oss and similar). For gpt-oss, half the layers are sliding-window, so this saves roughly 50% of the KV cache for long contexts such as 128k — measured at ~40% of the total KV allocation on the int8-KV variant above.

The corresponding onnxruntime-genai change (model-builder emits sliding_window for the CUDA EP, and the KV-cache allocator sizes windowed layers to the window) is a separate PR.

Copilot

Pull request overview

This PR adds a windowed (“sliding-window”, cache-relative) KV cache mode to the CUDA GroupQueryAttention kernel, including cache compaction/staging logic, RoPE position handling decoupled from cache capacity, and Python parity tests to validate windowed vs full-cache behavior.

Changes:

Added sliding_window_cache attribute to the contrib op schema and updated shape inference to keep present KV shapes capacity-based in windowed mode.
Implemented CUDA-side windowed KV cache eviction/compaction, cache-relative sequence-length handling, and disabled FlashAttention fast-decode for windowed caches.
Extended Python GQA tests with windowed-cache parity coverage and added RoPE cache length control for windowed setups.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_gqa.py	Adds `sliding_window_cache`/`rope_cache_length` test config and a new windowed KV cache parity test suite.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Adds `sliding_window_cache` schema attribute and adjusts GQA shape inference for windowed cache capacity semantics.
onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h	Explicitly rejects `sliding_window_cache=1` on WebGPU.
onnxruntime/contrib_ops/js/bert/group_query_attention.h	Explicitly rejects `sliding_window_cache=1` on JS.
onnxruntime/contrib_ops/cuda/bert/group_query_attention.h	Stores the new `sliding_window_cache_` attribute on the CUDA kernel.
onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc	Adds windowed-cache validation, staging for multi-token steps, compaction scratch, and disables Flash fast-decode for windowed caches.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh	Splits RoPE max-position from cache capacity and uses cache-relative append offsets.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h	Declares new windowed-cache compaction/copy helpers and extends sequence-length computation API.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu	Implements cache-relative length computation, cache compaction shift, staging-window copy, and uses cache-relative lengths across backends.
onnxruntime/contrib_ops/cuda/bert/attention_data.h	Adds cache-relative seq-length buffers + compaction scratch pointer to GQA runtime data.
onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc	Rejects `sliding_window_cache=1` on CPU to avoid silent incorrect behavior.
onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h	Extends `CheckInputs` to validate and propagate windowed-cache/RoPE max-position parameters.
onnxruntime/contrib_ops/cpu/bert/attention_parameters.h	Extends `GroupQueryAttentionParameters` with windowed-cache and RoPE max-position fields.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

…_window_cache

…_cache support - Add sliding_window_cache attribute to the operator schema table - Document the new sliding window (windowed) KV cache mode and its key behaviors - Include constraints and interactions with other features like Flash Attention

- Use CudaKernel::GetComputeStream(context) for the windowed KV staging scratch buffers; OpKernelContext::GetComputeStream() does not exist in the plugin EP adapter, which broke the Linux/Windows CUDA Plugin EP builds. - Restore the MatMulBlockQuantizedFp4Weight sections that were accidentally dropped from ContribOperators.md / OperatorKernels.md, fixing the Windows GPU Kernel Documentation Validation check.

Add CPU support for `sliding_window_cache=1` on `com.microsoft::GroupQueryAttention`, matching the existing CUDA/XQA contract: the cache buffer holds only the most recent tokens, indexed in cache-relative coordinates; results are numerically identical to a full-length cache. The second commit introduces a drifting-layout optimisation that amortises the per-step compaction cost, reducing per-token decode latency by 1.35–1.81× depending on buffer size. ## Summary of Changes ### CPU Kernel (`group_query_attention.cc / .h`) | File | Change | |---|---| | `contrib_ops/cpu/bert/group_query_attention.cc` | Three new anonymous-namespace helpers (`WindowedCacheEnd`, `PlanWindowedKvCache`, `MoveKvCacheRows`) plus the windowed block inside `Compute()`. | | `contrib_ops/cpu/bert/group_query_attention.h` | Add private `sliding_window_cache_` field. | **Key design decisions:** - **Single intercept point.** The CPU kernel has four attention paths (naive, quantized, quantized-flash, MLAS-flash). Instead of modifying each, the entire windowed-cache logic lives in `Compute()` and passes a compacted buffer as both `past_key/value` and `present_key/value`. The CPU kernel infers `past_present_share_buffer` from pointer equality, so `ConcatStateChunkGQA` only appends new tokens. - **Drifting layout (commit 2).** Resident entries live at `[0, end)` with the append point tracked via: ``` gap = C - W + 1 end(T) = T while T ≤ C = T − gap × ⌈(T − C) / gap⌉ afterwards (so W ≤ end ≤ C) ``` This is still a pure function of the absolute past length `T`, so the op is stateless. A whole `gap`-sized block is reclaimed at once, so compaction runs once every `gap` decode steps instead of every step. Safety: entries in `[0, end − W)` are older than the window and are already zeroed by the local-window mask — the same property that makes any `C > W` cache correct. - **Skip compaction on drifting steps.** When `seed_offset == 0` (the common decode step) the `MoveKvCacheRows` calls are bypassed entirely; a no-op thread-pool dispatch is not free (~65 µs/token). - **Staging for hazardous steps.** A multi-token step that would compact below the history its own first query reads runs against a temporary staging cache of `max_b(end_before) + S` rows; the surviving tail is written back afterwards. - **Thread-pool parallel `MoveKvCacheRows`.** Compaction is parallelised over `(batch, kv_head)` via `ThreadPool::TryParallelFor` with per-slice `memmove`. At `C = W = 128` (0.50 MiB moved per token) the serial version cost ~0.32 ms/token, making windowed decode 3× slower than the full cache; parallelised over 8 threads it drops to noise. ### Tests (`test_gqa.py`) | Test class / method | What it covers | |---|---| | `TestGQAWindowedKvCacheCpu` | Subclasses the existing CUDA suite; runs the same 11 test scenarios on `CPUExecutionProvider` with `float32` cache. | | `test_small_slack_many_compactions` | `C = W + 8`, 400 decode steps → ~44 compaction events; exercises the drifting layout exhaustively. | | `test_chunked_prefill_small_slack` | 32-token chunks at `C = W + 8`; staging from a drifted append point (the gap is < chunk size). | ### Documentation (`docs/contrib_ops/cuda/gqa.md`) - Corrected the "Cache capacity" bullet: `C ≥ W` is the minimum; modest surplus is recommended for CPU performance. - Added "CPU support" bullet describing the drifting layout and staging contract. - Corrected the PR-level doc note that falsely claimed amortisation was impossible under statelessness. ## Performance Results All measurements on AMD EPYC 7V12, 96 cores, 8 intra-op threads, fp32, 64 Q-heads / 8 KV-heads × 64 (`head_size`), `local_window_size = 128`. ### Before this PR — full-cache is unchanged Windowed decode latency (decode p10 ms) **before** the drifting layout was always-compact-on-every-step. The baseline (no sliding window) did not regress: | Prompt | Full cache (before) | Full cache (after) | |---|---|---| | 512 | 0.129–0.163 | 0.122–0.162 | | 2048 | 0.395–0.427 | 0.356–0.395 | | 8192 | — | within noise | ### Windowed vs full cache — decode throughput (decode p10 ms) `C = W + 256` (slack 256, the old default): | Prompt | Full cache | Windowed | Speedup | |---|---|---|---| | 512 | 0.169 | 0.145 | 1.17× | | 2048 | 0.352 | 0.167 | 2.11× | | 8192 | 1.301 | 0.155 | 8.4× | Windowed decode is **flat in context length** (attention scans only `W` effective keys). The CPU full-cache path does not exploit `local_window_size` and scans the entire cache. KV memory: full cache at 8192 tokens = 33.2 MiB; windowed (`C = 384`) = 1.50 MiB (fixed). ### Drifting layout — capacity sweep (prefill 2048 + 800 decode steps, decode p50 ms, 3 runs) A/B baseline pins `gap = 1` in the planner, which is equivalent to always-compact-on-every-step. | C | slack | gap | before | after | speedup | |---|---|---|---|---|---| | 128 | 0 | 1 | 0.0789 | *(same code, noise)* | — | | 132 | 4 | 5 | 0.0909 | 0.0654 | 1.39× | | **144** | **16** | **17** | 0.0964 | **0.0649** | **1.43×** | | 160 | 32 | 33 | 0.0957 | 0.0675 | 1.42× | | 256 | 128 | 129 | 0.1097 | 0.0773 | 1.42× | | 384 | 256 | 257 | 0.1573 | 0.0972 | 1.62× | | 512 | 384 | 385 | 0.2010 | 0.1109 | 1.81× | **The new decode-latency optimum is `C = W + 16`** (decode p50 0.065 ms, 1.35× better than `C = W` before this change). The curve is now a U-shape with a flat floor over slack 4–128; without drift it was monotonically worse with capacity because every decode step shifted `C − 1` rows. The analytic optimum matches: minimising `attention_slope × slack/2 + shift_cost/gap` over `gap` gives `slack* = sqrt(2 × shift_cost / attention_slope) ≈ 17`. ## Testing ```bash # Full windowed suite (13 CPU + 13 CUDA tests): pytest onnxruntime/test/python/transformers/test_gqa.py -k "TestGQAWindowedKvCache" -q # CPU GQA regression tests: pytest onnxruntime/test/python/transformers/test_gqa_cpu.py -q ``` Expected: 26 windowed tests pass, 4 CPU regression tests pass. ## Limitations (CPU only) The following are rejected with a clear error at session creation or compute time when combined with `sliding_window_cache=1` on CPU: - `attention_bias` input - `qk_output` attribute - Shared QKV layout (`kv_sequence_length == 0`) These limitations do not apply to the CUDA path. ## Checklist - [x] Tests added/updated - [x] No breaking changes - [x] Documentation updated (`docs/contrib_ops/cuda/gqa.md`)

github-actions

You can commit the suggested changes from lintrunner.

Comment-only follow-up; no functional change. - UnpackRoPEAppend: expand the rope_max_pos / cache_capacity docs to state explicitly that rope_max_pos is an ABSOLUTE position limit (the cos/sin cache length) and must not be conflated with cache_capacity. They only coincide for a full-length cache; using the capacity as the RoPE bound on a windowed layer would silently skip RoPE for every absolute position >= capacity -- an accuracy bug rather than a crash. - Correct the per-channel scale shape in the q_fold_scale and ScaleHeadsByChannelScaleKernel comments to (kv_num_heads, 1, head_size), matching the operator spec in docs/contrib_ops/cuda/gqa.md, and note that the indexing flattens the singleton middle dimension. - test_gqa.py: drop the now-stale enumeration of quantized XQA selection constraints (the same-tensor k/v scale requirement was lifted and the head_size/group_size sets are configuration dependent).

- Widen the windowed-cache end-pointer math to int64_t in both the CPU (WindowedCacheEnd, PlanWindowedKvCache) and CUDA kernels. seqlens_k is caller-supplied, so `seqlens_k[b] + 1` and `overflow + gap - 1` could signed-overflow near INT32_MAX before validation rejects the input. Results stay bounded by capacity, so narrowing back to int is safe. - Clamp the derived past_sequence_length at zero so a bogus seqlens_k cannot produce a negative row count downstream. - Hoist the CUDA windowed_end lambda out of GetSequenceLengths into a file-scope __device__ WindowedCacheEnd() helper that mirrors the CPU implementation, and document that the two must stay in sync. - Fix lintrunner comment alignment in GetSequenceLengths.

github-actions

You can commit the suggested changes from lintrunner.

GQA support sliding_window_cache

1e9c47b

tianleiwu marked this pull request as draft July 27, 2026 16:14

tianleiwu requested a review from Copilot July 27, 2026 16:14

tianleiwu changed the title ~~GQA support sliding_window_cache~~ [CUDA] GQA support sliding_window_cache Jul 27, 2026

Copilot started reviewing on behalf of tianleiwu July 27, 2026 16:15 View session

Copilot AI reviewed Jul 27, 2026

View reviewed changes

Comment thread onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h

Comment thread onnxruntime/core/graph/contrib_ops/bert_defs.cc

tianleiwu and others added 6 commits July 27, 2026 16:27

skip fp8 tests on device < 89

4ce9527

Add checks that sliding_window_cache=1 requires local_window_size > 0

c0360ba

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

update op spec

f047948

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Merge remote-tracking branch 'origin/main' into tlwu/20260727/sliding…

c28472c

…_window_cache

tianleiwu marked this pull request as ready for review July 29, 2026 00:28

tianleiwu added 2 commits July 29, 2026 01:57

reject attemtion_bias and windowed kv cache

c40bf67

Optimize Windowed KV cache compaction kernel

ec2bad9

tianleiwu mentioned this pull request Jul 29, 2026

Support windowed KV cache for sliding-window layers on CUDA microsoft/onnxruntime-genai#2334

Open

3 tasks

github-actions Bot reviewed Jul 29, 2026

View reviewed changes

Comment thread onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu Outdated

kunal-vaishnavi changed the title ~~[CUDA] GQA support sliding_window_cache~~ GQA support sliding_window_cache Jul 29, 2026

kunal-vaishnavi reviewed Jul 30, 2026

View reviewed changes

Comment thread onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc

kunal-vaishnavi reviewed Jul 30, 2026

View reviewed changes

Comment thread onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc

kunal-vaishnavi reviewed Jul 30, 2026

View reviewed changes

Comment thread onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu Outdated

tianleiwu added 2 commits July 30, 2026 07:28

github-actions Bot reviewed Jul 30, 2026

View reviewed changes

Comment thread onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh Outdated

lintrunner

8f3daa0

tianleiwu mentioned this pull request Jul 30, 2026

[CUDA] Add state_window to LinearAttention and CausalConvWithState to support MTP #31157

Open

kunal-vaishnavi approved these changes Jul 30, 2026

View reviewed changes

kunal-vaishnavi enabled auto-merge (squash) July 30, 2026 14:56

kunal-vaishnavi merged commit 7b89931 into main Jul 30, 2026
87 checks passed

kunal-vaishnavi deleted the tlwu/20260727/sliding_window_cache branch July 30, 2026 14:56

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

GQA support sliding_window_cache - #29904

GQA support sliding_window_cache#29904
kunal-vaishnavi merged 13 commits into
mainfrom
tlwu/20260727/sliding_window_cache

tianleiwu commented Jul 27, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

github-actions Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

tianleiwu commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Design

Summary of Changes

Testing

Motivation and Context

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tianleiwu commented Jul 27, 2026 •

edited

Loading