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/vto 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.