30. KV-Cache with Sliding-Window Eviction

Problem

Sliding-window attention means a token more than window_size positions in the past can never be attended to again — so a serving system built around it has no reason to keep those old key/value tensors in memory at all. A sliding-window KV-cache exploits this: instead of growing without bound like an ordinary autoregressive KV-cache, it evicts the oldest entries once it holds window_size tokens, keeping GPU memory for the cache flat regardless of how long the generated sequence gets. This is precisely the memory-efficiency payoff that makes Mistral's sliding-window attention practical to serve at long context lengths, not just cheaper to compute.

Implement a cache object with a single update(k_new, v_new) method: append this step's new key/value (one token's worth) to the cache, evict from the front if the cache now exceeds window_size tokens, and return the current full cached (k, v) tensors.

Source: src/30_kv_cache_sliding_window_eviction.py

class SlidingWindowKVCache:
    def __init__(self, window_size: int): ...
    def update(self, k_new: torch.Tensor, v_new: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: ...

>>> cache = SlidingWindowKVCache(window_size=3)
>>> k1, v1 = torch.randn(1, 2, 1, 4), torch.randn(1, 2, 1, 4)
>>> ck, cv = cache.update(k1, v1)
>>> ck.shape
torch.Size([1, 2, 1, 4])

Step-by-Step Approach

1. Store window_size and initialize self.k/self.v to None in __init__ — there is nothing to concatenate onto before the first call.
2. On the first update, since the cache is empty, just assign self.k, self.v = k_new, v_new directly rather than concatenating onto an empty tensor.
3. On subsequent calls, append along the sequence axis: self.k = torch.cat([self.k, k_new], dim=2) — the new token's k/v has shape (batch, heads, 1, head_dim).
4. After appending, check if self.k.shape[2] > window_size; if so, evict from the front by slicing to the last window_size entries: self.k[:, :, -window_size:, :].
5. Return the (possibly just-evicted) self.k, self.v so the caller always attends against exactly what the cache currently holds.
6. Sanity-check that the cache length never exceeds window_size no matter how many steps are run — that flat memory ceiling is the entire point of the eviction policy.

The key insight is that eviction is just a slice, not a different data structure: because sliding-window attention only ever needs the most recent window_size tokens, a simple "keep the last W" tensor slice after each append is sufficient -- no ring buffer or linked-list bookkeeping is required for correctness, only for avoiding the O(window_size) copy cost of re-slicing every step in a real high-throughput system.

Reference solution

import torch


class SlidingWindowKVCache:
    def __init__(self, window_size: int):
        self.window_size = window_size
        self.k = None
        self.v = None

    def update(self, k_new: torch.Tensor, v_new: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if self.k is None:
            # first call: nothing to concatenate onto yet
            self.k, self.v = k_new, v_new
        else:
            # append along the sequence axis (dim=2): (batch, heads, t, hd) -> (batch, heads, t+1, hd)
            self.k = torch.cat([self.k, k_new], dim=2)
            self.v = torch.cat([self.v, v_new], dim=2)

        # evict from the front once the cache exceeds the window -- a plain
        # slice is O(window_size) to copy but keeps the cache tensor
        # contiguous, which matters for the matmuls that read it downstream
        if self.k.shape[2] > self.window_size:
            self.k = self.k[:, :, -self.window_size:, :]
            self.v = self.v[:, :, -self.window_size:, :]

        return self.k, self.v

Key Functions & Tricks

• torch.cat([self.k, k_new], dim=2) — grows the cache along the sequence axis, the standard KV-cache append operation.
• self.k[:, :, -window_size:, :] — negative-index slicing along the sequence axis is the entire eviction policy; no explicit index bookkeeping needed.
• self.k is None guard — avoids a special-cased "empty tensor" shape for the very first call, which is easy to get wrong if you instead try to pre-allocate an empty (batch, heads, 0, head_dim) tensor.
• Instance state (self.k, self.v) — the cache is inherently stateful across calls, unlike the pure functions in most other problems in this set; this mirrors how a real inference server's per-request cache object behaves across a multi-step decode loop.

How to Recognize This Pattern