4. Inference Optimization

KV Cache

seniorai-mlcaching

Without a KV cache, generating token N of a response means recomputing attention over all N tokens from scratch -- generating a 1000-token response would do roughly 1000x that much redundant work. A KV cache fixes this: each token's key/value vectors are computed once, cached, and reused -- producing token N only computes ONE new query against the (already-cached) keys/values of every earlier token.

dot/softmax are given (same as the Attention exercise). Implement KVCache.attend(q, k, v): process one new token whose own key/value vectors are k/v -- append them to the cache -- then return this token's attention output, querying with q against ALL cached keys/values so far (including the one you just appended).

Example

cache = KVCache()
cache.attend([1.0], [1.0], [42.0])   # first token: attends only to itself
# -> [42.0]   (only one key/value exists so far, weight = 1.0)

Constraints

  • This is the property that actually matters: calling attend() once per token, in order, must produce EXACTLY the same output as running full (non-incremental) causal attention over the whole sequence-so-far at every step. A KV cache is a pure performance optimization -- it must never change the result, only how much work it takes to get there.
  • Append k/v to the cache BEFORE computing scores -- the current token attends to itself too.

The missing link: the cache can't grow forever

A generation running long enough eventually runs into a hard memory budget -- the "append once, reuse forever" idea above is correct but incomplete, because "forever" isn't actually an option. Real systems (Mistral's sliding-window attention, llama.cpp's context shifting) cap the cache at a fixed max_size: once full, appending a new entry evicts the OLDEST one first. This trades away long-range history for a bounded, predictable memory footprint -- the model can only ever attend to the most recent max_size tokens, not the whole conversation.

Implement this in __init__(self, max_size=None): max_size=None means unbounded (today's behavior, unchanged). When max_size is set, attend must evict the oldest cached entry BEFORE scoring, whenever the cache would otherwise exceed max_size -- so an evicted token can never appear in that step's (or any later step's) attention output.

Real-world grounding

  • llama.cpp's KV cache implementation, including context-shifting/eviction logic: src/llama-kv-cache.cpp
  • vLLM's KV cache manager (paired with PagedAttention, which stores the cache in non-contiguous memory pages): vllm/v1/core/kv_cache_manager.py
  • Mistral 7B's sliding-window attention (the paper that popularized a fixed attention window as a deliberate design choice, not just a memory hack): search "sliding window attention" in mistralai/mistral-src

Real KV caches add even more on top of this: paging non-contiguous memory, sharing cache pages across requests with a common prefix, and quantizing cached values to save memory -- but "evict oldest when full" is the core idea underneath the eviction policy specifically.

hints
kv_cache.pysign in to run
Loading...
test result

This is a preview. Sign in to run this exercise against the real tests.