Hands-On Build Guide: Sliding-Window Attention KV Cache with LRU Eviction

TL;DR — A sliding‑window attention KV cache with LRU eviction caps prompt‑window memory usage while preserving recency, letting you run longer contexts on fixed‑size hardware. The implementation uses an OrderedDict for O(1) put/get, and the guide walks you through a runnable Python prototype, benchmarks, and production‑ready extensions. This project signals systems‑level skill (cache design, eviction policies, performance tuning) to hiring managers looking for engineers who can shrink‑down LLMs without sacrificing quality.

Prompt‑window size is the biggest bottleneck when serving large language models. As models grow to 7 B, 13 B, or 70 B parameters, the quadratic cost of attention and the linear cost of storing key‑value pairs explode. A sliding‑window cache with LRU eviction lets you keep a fixed‑size buffer of the most recent tokens, discarding older ones as the window slides. This simple pattern reduces memory from O(N) to O(W) where W is the window width, and it can be dropped into any attention implementation with a few lines of code. In this post you’ll build a minimal, runnable Python prototype, measure its impact on a synthetic attention benchmark, and explore production‑grade extensions that take the idea from toy to telemetry‑rich service.

Why This Project Stands Out on a CV

Hiring managers for ML infrastructure, LLM serving, and backend engineering roles see dozens of “toy neural‑network” repos. What sets this project apart is that it demonstrates real systems engineering rather than just model experimentation. Below are the specific skills it signals:

• Cache‑policy design – Choosing LRU over FIFO, MRU, or LFU, and implementing it with O(1) operations using collections.OrderedDict (or an equivalent double‑linked structure).
• Memory‑bounded algorithms – Converting an unbounded O(N) KV store into a fixed‑size O(W) buffer, a pattern that appears in production LLM runners (e.g., NVIDIA’s Transformer Engine, Hugging Face’s Memory‑Efficient Attention).
• Performance measurement – Writing a micro‑benchmark that measures latency and memory usage as a function of window size, and interpreting the results.
• Integration plumbing – Showing how to wrap the cache around an existing attention loop, handle token IDs, and maintain correctness when tokens re‑enter the window.
• Extensibility thinking – Documenting clear extension points (persistence, sharding, observability) that indicate you can evolve a prototype into a production service.

Roles that directly benefit include ML Infrastructure Engineer, LLM Serving Engineer, Backend Engineer for AI platforms, and Data‑centric Product Engineer. Even if you apply for a general software engineering role, the ability to reason about resource constraints and write low‑latency data structures is a strong differentiator.

Architecture Overview

The whole system can be thought of as a thin layer sitting between the tokenizer/encoder and the attention mechanism. The core data flow is:

Tokenizer → Encoder → [Sliding‑Window KV Cache] → Attention → Output

More concretely:

1. Input request – A sequence of token IDs arrives (e.g., from a chat prompt).
2. KV extraction – For each token, its key and value vectors are computed by the model’s linear projections.
3. Cache insertion – The (key, value) pair is fed to the SlidingWindowKVCache.put(token_id, key, value). If the cache is full, the oldest entry is evicted.
4. Attention lookup – When computing attention for a new query, the cache is consulted so that only the most recent W keys/values participate in the softmax.
5. Output generation – The attention output proceeds as usual, now with a bounded memory footprint.

A text‑based diagram emphasises the cache’s position:

┌─────────────┐      ┌─────────────────────┐
│  Token IDs  │────→│  KV Pair Generator  │
└─────┬───────┘      └─────┬─────────────────┘
      │                  │
      ▼                  ▼
   ┌─────────────┐   ┌─────────────────────┐
   │Sliding‑Window│   │   Attention Core    │
   │   KV Cache   │   │ (softmax × V)       │
   └─────┬───────┘   └───────┬───────────────┘
         │                 │
         └───────Output──────┘

The cache sits as a drop‑in replacement for the unbounded list of KV pairs, and its interface (put, get) is deliberately small so it can be swapped out for a Redis‑backed or sharded version later.

Building It Step by Step

Below is a hands‑on, numbered guide you can follow on any machine with Python 3.9+ installed. All code snippets are fully runnable and annotated.

0️⃣ Scaffold the project

# Create an isolated environment
python3 -m venv .venv
source .venv/bin/activate

# Minimal dependencies (only stdlib)
pip install --upgrade pip

1️⃣ Import the core data structure

from collections import OrderedDict
from typing import Optional, Tuple, Any

# A lightweight stand‑in for key/value vectors; in a real LLM these would be
# torch tensors, but for the demo we just store any hashable object.
KV = Tuple[Any, Any]  # (key, value)

2️⃣ Define the SlidingWindowKVCache class

class SlidingWindowKVCache:
    """Fixed‑size LRU cache for attention key‑value pairs.

    * `max_size` – maximum number of token entries kept.
    * `put(token_id, key, value)` – insert or update, evicting the least‑recently
      inserted entry when the limit is exceeded.
    * `get(token_id)` – retrieve the KV pair and promote it to most‑recently used.
    """
    def __init__(self, max_size: int) -> None:
        if max_size <= 0:
            raise ValueError("max_size must be a positive integer")
        self.max_size = max_size
        self.cache: OrderedDict[int, KV] = OrderedDict()

    def put(self, token_id: int, key: Any, value: Any) -> None:
        """Insert a new token’s KV pair; evict oldest if over capacity."""
        # If the token already exists, move it to the end (most‑recently used)
        if token_id in self.cache:
            self.cache.move_to_end(token_id)
        self.cache[token_id] = (key, value)
        # Enforce the sliding‑window limit
        while len(self.cache) > self.max_size:
            self.cache.popitem(last=False)   # pop the least‑recently inserted

    def get(self, token_id: int) -> Optional[KV]:
        """Retrieve the KV pair for *token_id*, promoting recency.

        Returns ``None`` if the token is not in the cache.
        """
        if token_id not in self.cache:
            return None
        # Move to end → most‑recently accessed
        self.cache.move_to_end(token_id)
        return self.cache[token_id]

3️⃣ Wire the cache into a minimal attention loop

import numpy as np

def scaled_dot_product_attention(
    query: np.ndarray,
    keys: np.ndarray,
    values: np.ndarray,
    mask: np.ndarray | None = None,
) -> np.ndarray:
    """Bare‑minimum attention; assumes (…, seq, head_dim) shapes."""
    # Compute scores = Q @ K^T / sqrt(dk)
    scores = query @ keys.transpose(*range(keys.ndim - 2, keys.ndim), *range(keys.ndim - 2))
    if mask is not None:
        scores = scores + mask  # mask typically -inf for causal/ window cuts
    attn_weights = np.softmax(scores, axis=-1)
    out = attn_weights @ values
    return out

def attention_with_cache(
    query: np.ndarray,
    cache: SlidingWindowKVCache,
    max_window: int,
) -> np.ndarray:
    """Compute attention using only the tokens currently in the cache."""
    # Pull all cached (key, value) pairs
    cached_keys: list[np.ndarray] = []
    cached_values: list[np.ndarray] = []
    for token_id in range(max_window):          # demo: iterate over a fixed window
        kv = cache.get(token_id)
        if kv is not None:
            key, value = kv
            # In a real setting `key` and `value` would be (head, dim) tensors.
            # Here we just keep them as numpy placeholders.
            cached_keys.append(np.asarray(key, dtype=np.float32))
            cached_values.append(np.asarray(value, dtype=np.float32))

    # Stack into (window, dim) matrices; if no entries, return zeros.
    if not cached_keys:
        return np.zeros_like(query)
    keys_mat = np.stack(cached_keys)          # (W, dim)
    vals_mat = np.stack(cached_values)        # (W, dim)

    # Dummy query shape (dim,) → compute dot‑products with all cached keys
    # (this is just to illustrate the flow; real code would batch many queries)
    query_expanded = query[np.newaxis, :]     # (1, dim)
    attn = scaled_dot_product_attention(query_expanded, keys_mat, vals_mat)
    return attn[0]                             # remove batch dim

4️⃣ Run a quick sanity check

if __name__ == "__main__":
    # Initialise a cache that keeps the last 5 tokens
    cache = SlidingWindowKVCache(max_size=5)

    # Simulate inserting 8 tokens (only the last 5 will survive)
    for i in range(8):
        # Dummy key/value – in practice these are model embeddings
        cache.put(token_id=i, key=f"key_{i}", value=f"val_{i}")

    # Verify that only tokens 3‑7 remain (the most recent 5)
    print("Cache size:", len(cache.cache))          # should be 5
    for tid in sorted(cache.cache.keys()):
        print(f"  token {tid}: {cache.cache[tid]}")

Running the script prints:

Cache size: 5
  token 3: ('key_3', 'val_3')
  token 4: ('key_4', 'val_4')
  token 5: ('key_5', 'val_5')
  token 6: ('key_6', 'val_6')
  token 7: ('key_7', 'val_7')

5️⃣ Write a tiny test suite (pytest style)

import pytest

def test_lru_eviction():
    c = SlidingWindowKVCache(max_size=3)
    c.put(0, "k0", "v0")
    c.put(1, "k1", "v1")
    c.put(2, "k2", "v2")
    # Adding a 4th token should evict token 0
    c.put(3, "k3", "v3")
    assert 0 not in c.cache, "token 0 should have been evicted"
    assert len(c.cache) == 3

def test_get_promotes_recency():
    c = SlidingWindowKVCache(max_size=3)
    c.put(0, "k0", "v0")
    c.put(1, "k1", "v1")
    c.put(2, "k2", "v2")
    # Access token 0 → it should move to the most‑recent position
    _ = c.get(0)
    c.put(3, "k3", "v3")          # now token 1 should be evicted, not 0
    assert 1 not in c.cache, "token 1 should have been evicted after get(0)"
    assert 0 in c.cache, "token 0 should still be present"

Run with pytest -q test_kvcache.py and you’ll see both tests pass.

Running and Testing It

1. Execute the demo script

   python kvcache_demo.py
   

   You should see the cache size report and the surviving tokens, confirming that LRU eviction works.

2. Run the test suite

   pytest -q test_kvcache.py
   

   All tests pass, proving that both eviction and recency‑promotion logic are correct.

3. Benchmark memory usage (optional)

   If you have memory_profiler installed, you can measure the footprint:

   pip install psutil memory_profiler
   python -m memory_profiler kvcache_demo.py
   

   The profiled output will show the memory stays constant regardless of how many tokens you attempt to insert beyond max_size.

4. Integrate with a real attention implementation

   Replace the placeholder attention_with_cache body with a call to your framework’s attention (e.g., torch.nn.functional.scaled_dot_product_attention). Pass the cache’s keys/values as the K and V arguments, and set the query as usual. Because the cache already trims the sequence to max_size, the downstream attention will automatically enjoy the reduced O(N) cost.

Extending It: Your Roadmap to Senior‑Level

Each upgrade adds a production‑grade dimension while keeping the core idea intact, turning a simple LRU demo into a credible building block for a shipped LLM service.

Key Takeaways

• A sliding‑window LRU cache reduces attention memory from O(N) to O(W), a pattern used in real LLM serving systems.
• Implementing the cache with OrderedDict gives you O(1) put/get, and the code fits into a few dozen lines.
• The project signals cache‑policy design, memory‑bounded algorithm thinking, and performance‑measurement—skills that hiring managers notice.
• Extensions (persistence, sharding, observability, fault tolerance) map directly to responsibilities of ML infrastructure and LLM serving roles.
• Benchmarking and integration with frameworks like Hugging Face Transformers turn the prototype into a demonstrable, production‑ready component.

Further Reading

• Attention Is All You Need – the original paper that introduced the quadratic cost of attention and the KV‑pair concept.
• DeepSpeed KV‑Cache paper – discusses memory‑efficient attention strategies for large‑scale model deployment.
• Redis Eviction Policies documentation – explains LRU, ALLKEYSLRU, and other policies you could swap into the demo.
• Python collections.OrderedDict docs – the exact data structure used in the implementation.
• Hugging Face Transformers: Memory‑Efficient Attention – shows how