[CPU] Add windowed kv cache for GQA

[Tianlei Wu (tianleiwu)]

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

# 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

• Tests added/updated
• No breaking changes
• Documentation updated (docs/contrib_ops/cuda/gqa.md)

[Copilot]

Pull request overview

This PR adds CPU Execution Provider support for sliding_window_cache=1 on com.microsoft::GroupQueryAttention, enabling cache-relative (windowed) KV caches with a drifting-layout compaction strategy to reduce per-token decode latency while preserving numerical parity with a full-length cache.

Changes:

• Implement windowed KV cache planning/compaction/staging in the CPU GQA kernel Compute() to reuse existing attention paths with cache-relative seqlens_k.
• Add CPU parity coverage by reusing the existing CUDA windowed-cache test suite on CPUExecutionProvider (fp32 cache) and add stress scenarios for frequent compactions/staging.
• Update the GQA documentation to reflect the corrected capacity requirement (C >= W) and describe CPU behavior/limitations.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc	Implements windowed KV cache planning and memory movement (compaction + staging) for CPU GQA.
onnxruntime/contrib_ops/cpu/bert/group_query_attention.h	Adds a private sliding_window_cache_ attribute flag on the CPU kernel.
onnxruntime/test/python/transformers/test_gqa.py	Adds CPU windowed-cache parity tests by subclassing the existing suite; extends scenarios for small-slack compaction and staged chunked prefill.
docs/contrib_ops/cuda/gqa.md	Updates windowed cache capacity guidance and documents CPU support + limitations.

Comments suppressed due to low confidence (1)

onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc:94

• PlanWindowedKvCache computes past_sequence_length + sequence_length and seqlens_k[b] + 1 - sequence_length using int, which can overflow when seqlens_k[b] is near INT32_MAX. Even if such inputs are invalid, overflowing before validation is undefined behavior. Do the arithmetic in int64_t/SafeInt before passing values to WindowedCacheEnd.

  for (int b = 0; b < batch_size; ++b) {
    const int past_sequence_length = seqlens_k[b] + 1 - sequence_length;  // P
    const int end_before = WindowedCacheEnd(past_sequence_length, capacity, gap);
    const int end_after = WindowedCacheEnd(past_sequence_length + sequence_length, capacity, gap);
    const int kept = end_after - sequence_length;  // resident rows that survive this step

[kunal-vaishnavi]

This looks like a valid issue.

[Tianlei Wu (tianleiwu)]

Good catch — fixed in d3c1a4d on the parent branch (#29904).

WindowedCacheEnd now takes int64_t parameters and does the intermediate math in 64-bit, and PlanWindowedKvCache computes past_sequence_length as std::max<int64_t>(int64_t{seqlens_k[b]} + 1 - sequence_length, 0). seqlens_k is caller-supplied, so seqlens_k[b] + 1 and overflow + gap - 1 could signed-overflow near INT32_MAX — UB, even though later validation would reject such inputs. The result is bounded by capacity, so narrowing back to int is safe. The clamp at zero also prevents a bogus (too small) seqlens_k from producing a negative row count downstream.

The equivalent CUDA helper in group_query_attention_impl.cu got the same treatment in the same commit.

[kunal-vaishnavi]

From Copilot:

In PlanWindowedKvCache, we’re doing seqlens_k[b] + 1 - sequence_length and past_sequence_length + sequence_length in int. Could we switch this arithmetic to int64_t or SafeInt before converting back? That would avoid signed-overflow UB on invalid/extreme inputs before validation kicks in.

[Tianlei Wu (tianleiwu)]

Fixed in d3c1a4d on the parent branch (#29904): past_sequence_length is now computed as std::max<int64_t>(int64_t{seqlens_k[b]} + 1 - sequence_length, 0), so neither the + 1 nor the subtraction can overflow, and a bogus seqlens_k cannot yield a negative row count. WindowedCacheEnd and the matching CUDA __device__ helper were widened to int64_t in the same commit.

[kunal-vaishnavi]

Is output_qk implemented inside GQA? From what I recall, we added it to MHA because GQA's optimized kernels (e.g. flash attention, memory efficient attention) would never accumulate the QK matrix.

[Tianlei Wu (tianleiwu)]

Yes, qk_output is implemented in GQA. The schema registers the qk_output attribute and runs shape inference for output index 3 (GroupQueryAttentionTypeAndShapeInference(ctx, 3, qk_output_index) in core/graph/contrib_ops/bert_defs.cc), and the CPU kernel writes that output through the unfused attention path in gqa_attention_base.h.

The reason it feels restricted is that the fast paths opt out of it: both the flash-attention and memory-efficient-attention branches are only selected when output_qk == nullptr. So requesting QK output forces the unfused kernel. That is also why the windowed cache rejects the combination for now — the unfused path would need separate handling for the drifting cache layout.

[kunal-vaishnavi]

I added some small notes before I noticed this was merging into another PR. I merged the CPU PR into the CUDA PR so that all changes can be made in one branch.