Skip to content

Reduce local-attention KV storage with per-layer rings - #1017

Closed
Mikyx-1 wants to merge 1 commit into
google:mainfrom
Mikyx-1:kv-cache-local-windows
Closed

Reduce local-attention KV storage with per-layer rings#1017
Mikyx-1 wants to merge 1 commit into
google:mainfrom
Mikyx-1:kv-cache-local-windows

Conversation

@Mikyx-1

@Mikyx-1 Mikyx-1 commented Sep 6, 2026

Copy link
Copy Markdown

Closes #1016

Local-attention layers currently allocate KV storage for the entire context. This change gives each layer its own FP32 buffer: local layers use min(seq_len, window - 1 + max(1, prefill_batch_size)) rows, and global layers keep full-context storage.

The extra batch rows prevent QKV prefill writes from overwriting history needed by earlier queries. Runtime growth preserves cached history, copies remain independent, and attention reads use each layer's ring capacity while score-buffer indexing and arithmetic remain unchanged.

Validation on Gemma 3 270M and 1B: all 10 A/B comparisons produced byte-identical decoded output and identical token IDs (1,280 generated tokens per build). Runs used a shared baseline-calibrated matrix plan to isolate the cache change from timing-dependent autotuning.

Model KV capacity, MiB Peak process RAM, MiB Prefill, tokens/s Decode, tokens/s
270M 288.5 → 72.7 644.9 → 565.6 1013.3 → 1082.3 40.48 → 42.30
1B 416.5 → 100.0 1574.3 → 1458.5 199.2 → 231.7 13.79 → 15.00

Medians from three alternating pairs with a 3,367-token prompt, 128 generated tokens, 8,192-token capacity, 256-token prefill batches, and six pinned workers on an i5-12400F. CPU load varied; speed improvements are measured observations, not a universal performance guarantee. RAM savings were consistent.

  • CMake cache/attention regression tests: 8/8 passed, covering wraparound, runtime growth, prefix attention, global layers, context limits, copies, and exact FP32 attention comparisons on AVX2 and EMU128.
  • AddressSanitizer with leak detection: 8/8 passed. Changed KV/attention translation units and tests were instrumented; supporting libraries were reused.
  • CLI, single benchmark, API server, and batch benchmark built successfully.

C++ API change: direct users of the former public KVCache::kv_cache matrix must use LayerCache(layer) or Row(layer, position). Repository callers are updated.

@Mikyx-1

Mikyx-1 commented Sep 6, 2026

Copy link
Copy Markdown
Author

Superseded by #1020, which ports the local-cache changes to dev and includes new dev-baseline validation.

@Mikyx-1 Mikyx-1 closed this Sep 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Reduce KV memory for local-attention layers

1 participant