KV caching: correct positions, masks, memory and reuse

The story: reuse follows causality

An autoregressive decoder predicts the next token from a fixed prefix. When a new token arrives, it cannot change earlier hidden states in a causal model: those states were functions only of earlier tokens. At every layer we can therefore retain the earlier keys and values and compute the new position alone. The old queries have already done their job; future positions read old keys and values, not old queries.

Think of a meeting notebook. Appending one speaker's note avoids rewriting the notebook, but answering a new question can still require reading all earlier notes. A KV cache saves repeated projections and feed-forward work; it does not eliminate the attention read over history. The original claim of 500 total attention operations for 500 generated positions confused newly computed token states with query–key interactions.

Preparation: understand causal self-attention, residual blocks and next-token logits from DL036–041. Here the complete random-weight CPU decoder makes the cache contract testable without a trained checkpoint. The outcomes are cached/full-prefix equivalence, correct chunk masks and positions, memory arithmetic, a defensible benchmark, and an informed choice among paging, sharing, compression and eviction.

1 · Work saved, work that remains

For fixed layer count and width, recomputing a prefix of length t costs an attention term proportional to t squared plus a projection/FFN term proportional to t. Across N growing prefixes, these become cubic and quadratic totals. A cached decoder computes one new state per step but scores it against t keys: attention is linear per step and quadratic in total; projection/FFN work is constant per step and linear in total. Constants include model width, heads, batch and vocabulary projection. A fixed prompt of length P adds prefill work and a decode attention sum proportional to NP plus N squared.

For N=500, the square sum is 41,791,750; the linear sum is 125,250. Those are score-grid counts in a simplified accounting model, not FLOPs, measured latency or an 85,000-fold serving improvement. For N=2048, use the exact formula below and include whether the implementation computes a full square or only valid causal pairs.

A sliding fixed-size context changes this asymptotic model. Cropping a learned-position decoder and restarting positions at zero also changes the function: dropping the first cache entry is not equivalent to recomputing the cropped prefix. Earlier hidden states still contain information from evicted tokens. Define an actual sliding-attention architecture or recompute at the boundary; do not assert equivalence to full context.

The immutable-prefix argument assumes identical weights, tokenization, attention policy, positional convention and evaluation behavior. Dropout, parameter updates or inconsistent masks invalidate the simple equality. Decoder cross-attention can also cache projections of a fixed encoder output; it is false that every encoder-related workload has no useful caching. A bidirectional encoder's self-attention states generally change when its input is extended.

2 · Complete decoder with per-layer caches

Each cache has shape (batch, heads, cached_length, head_dim). A new chunk of T queries starts at P, the cached length. Its allowed key condition is key_position <= P + query_index. A square lower triangle applied to a rectangular score array is wrong: it hides the prefix from the first new query. For a single new token, all available keys are visible, provided there is no padding or other restriction.

The following fences form one CPU program. This reference deliberately uses concatenation for inspectability, not a production allocator. It checks chunked prefill as well as one-token decoding, retains every layer's cache, rejects inconsistent lengths and uses explicit learned absolute positions. It has no hidden dependency on the earlier nanoGPT lesson.

import math, statistics, time
import torch
from torch import nn
import torch.nn.functional as F
 
torch.set_num_threads(1)
torch.manual_seed(42)
 
class CachedAttention(nn.Module):
    def __init__(self, width=24, heads=3):
        super().__init__()
        if width % heads:
            raise ValueError('width must divide into heads')
        self.width, self.heads, self.dim = width, heads, width // heads
        self.qkv = nn.Linear(width, 3 * width)
        self.out = nn.Linear(width, width)
 
    def forward(self, x, cache=None):
        b, t, c = x.shape
        q, k, v = [z.reshape(b, t, self.heads, self.dim).transpose(1, 2)
                   for z in self.qkv(x).chunk(3, dim=-1)]
        past = 0
        if cache is not None:
            pk, pv = cache
            expected = (b, self.heads, pk.shape[2], self.dim)
            if pk.shape != expected or pv.shape != expected:
                raise ValueError('cache shape mismatch')
            if any(z.device != x.device or z.dtype != x.dtype for z in (pk, pv)):
                raise ValueError('cache device/dtype mismatch')
            past = pk.shape[2]
            k, v = torch.cat([pk, k], 2), torch.cat([pv, v], 2)
        query_pos = past + torch.arange(t, device=x.device)
        key_pos = torch.arange(past + t, device=x.device)
        allow = key_pos[None, :] <= query_pos[:, None]
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=allow,
                                           dropout_p=0.0, is_causal=False)
        return self.out(y.transpose(1, 2).reshape(b, t, c)), (k, v)
 
class CachedBlock(nn.Module):
    def __init__(self, width, heads):
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(width), nn.LayerNorm(width)
        self.attn = CachedAttention(width, heads)
        self.ffn = nn.Sequential(nn.Linear(width, 4*width), nn.GELU(),
                                 nn.Linear(4*width, width))
    def forward(self, x, cache):
        y, new = self.attn(self.n1(x), cache)
        x = x + y
        return x + self.ffn(self.n2(x)), new
 
class CachedDecoder(nn.Module):
    def __init__(self, vocab=19, width=24, heads=3, layers=2, capacity=160):
        super().__init__()
        self.capacity = capacity
        self.token, self.position = nn.Embedding(vocab, width), nn.Embedding(capacity, width)
        self.blocks = nn.ModuleList([CachedBlock(width, heads) for _ in range(layers)])
        self.norm, self.head = nn.LayerNorm(width), nn.Linear(width, vocab)
    def forward(self, ids, caches=None, wrong_positions=False):
        if ids.ndim != 2 or ids.shape[1] == 0:
            raise ValueError('nonempty [batch,time] IDs required')
        past = 0
        if caches is not None:
            if len(caches) != len(self.blocks):
                raise ValueError('one cache per layer required')
            lengths = [c[0].shape[2] for c in caches]
            if len(set(lengths)) != 1:
                raise ValueError('layer cache lengths disagree')
            past = lengths[0]
        if past + ids.shape[1] > self.capacity:
            raise ValueError('position capacity exceeded; no silent crop')
        start = 0 if wrong_positions else past
        pos = torch.arange(start, start + ids.shape[1], device=ids.device)
        x = self.token(ids) + self.position(pos)[None, :, :]
        new = []
        for i, block in enumerate(self.blocks):
            x, c = block(x, None if caches is None else caches[i])
            new.append(c)
        return self.head(self.norm(x)), new
 
@torch.inference_mode()
def generate(model, prompt, count, cached=True, temperature=0.0, top_k=None):
    if count < 0 or temperature < 0:
        raise ValueError('count and temperature must be nonnegative')
    if model.training:
        raise ValueError('call eval before inference')
    if count == 0:
        return prompt.clone()
    if prompt.shape[1] + count - 1 > model.capacity:
        raise ValueError('generation exceeds position capacity')
    ids, caches = prompt.clone(), None
    logits, caches = model(ids)
    for step in range(count):
        scores = logits[:, -1, :]
        if temperature == 0:
            nxt = scores.argmax(-1, keepdim=True)
        else:
            scores = scores / temperature
            if top_k is not None:
                if not 1 <= top_k <= scores.shape[-1]:
                    raise ValueError('invalid top_k')
                threshold = scores.topk(top_k).values[:, -1:]
                scores = scores.masked_fill(scores < threshold, -torch.inf)
            nxt = torch.multinomial(scores.softmax(-1), 1)
        ids = torch.cat([ids, nxt], 1)
        if step + 1 < count:  # do not compute an unused final forward
            logits, caches = model(nxt, caches) if cached else model(ids)
    return ids
 
model = CachedDecoder().double().eval()
ids = torch.randint(0, 19, (2, 17))
with torch.inference_mode():
    full, _ = model(ids)
    chunks, cache = [], None
    for a, b in [(0, 5), (5, 8), (8, 9), (9, 17)]:
        logits, cache = model(ids[:, a:b], cache)
        chunks.append(logits)
    torch.testing.assert_close(torch.cat(chunks, 1), full, rtol=1e-10, atol=1e-10)
    prefix, cache = model(ids[:, :5])
    wrong, _ = model(ids[:, 5:8], cache, wrong_positions=True)
    assert not torch.allclose(wrong, full[:, 5:8])
    assert all(c[0].shape == (2, 3, 5, 8) for c in cache)
assert torch.equal(generate(model, ids[:, :5], 12, True),
                   generate(model, ids[:, :5], 12, False))
assert torch.equal(generate(model, ids[:, :5], 0), ids[:, :5])
print('cached/full logits, chunked prefill, greedy generation and wrong-position control pass')

To integrate into a trained decoder, replace the attention call and propagate (logits, caches) through every block, as this program does. Prefill once on the full prompt, then pass only the latest sampled token. Retain temperature/top-k semantics unchanged for a fair comparison. Greedy ID parity is a useful check; close logits alone do not guarantee identical sampled text under floating-point changes.

For production, allocate fixed buffers or pages and write into slices rather than repeatedly copying the full cache with cat. Maintain per-request valid lengths, padding masks, beam reorder indices and capacity checks. A cache belongs to a request/model/configuration, not a global variable reused after an unrelated prompt.

3 · Benchmark without inventing a speedup

Measure time-to-first-token separately from steady decode. Keep prompt, weights, precision, output length, thread count and generation policy fixed. Check outputs before timing, warm up, run repeats and report median and spread. Longer cached decode still reads more KV bytes; at small batch it can be bandwidth-limited, while batching raises reuse of weights and arithmetic intensity. Neither prefill-always-compute-bound nor decode-always-bandwidth-bound is a theorem. Measure batch and context independently.

prompt = ids[:1, :5]
for n in (8, 24, 48):
    result = {}
    for cached in (False, True):
        generate(model, prompt, n, cached)
        times = []
        for _ in range(3):
            t0 = time.perf_counter()
            generate(model, prompt, n, cached)
            times.append(time.perf_counter() - t0)
        result[str(cached)] = statistics.median(times)
    print('CPU tokens', n, 'seconds', result,
          'observed ratio', result['False'] / result['True'])

These are actual local timings of random-weight toy inference, not A100 or useful language-generation results. Plot log(time) against log(N) only after accounting for prompt work, width terms, kernel launch overhead and memory traffic: cubic versus quadratic total attention work is not the original claimed quadratic-versus-linear total runtime curve.

GPU procedure, unexecuted: use an already available compatible model/device, move model and prompt together, and recheck logits at a dtype-appropriate tolerance. Synchronize CUDA before and after every timed region, warm each path, reset peak-memory statistics, record baseline allocated memory, run repeated fixed-length greedy generations, and report peak allocated minus baseline as well as total allocated/reserved memory. Sweep batch and context without silently shortening failed runs; record OOM and unsupported cases. Include device, driver, torch version and precision. No hardware rental or GPU result is implied.

4 · Memory arithmetic and architecture-specific layouts

For ordinary MHA/GQA with equal key and value width, raw KV bytes are:

B is batch, L layers, T cached tokens, H-KV the number of KV heads, d-h head width and s bytes per element. Use KV heads, not query heads, for GQA. For unequal key/value widths replace 2*d_h by their sum. Allocation padding, quantization scales, workspace and allocator overhead are additional.

These rows compare raw storage at a common hypothetical length, not supported context limits. In particular, the GPT-2 XL configuration has n_positions=1024; its 4096-token row is a storage extrapolation, not a runnable stock-model configuration. At its 1024-position limit, the same FP16 cache arithmetic gives 314,572,800 bytes (300 MiB).

The last row is not 2*512: V3 stores a shared 512-element compressed KV latent plus a 64-element decoupled rotary key, at two bytes each. A comparison storing expanded per-head keys of width 192 and values of width 128 for 128 heads needs 81,920 bytes/token/layer: the compressed layout saves 98.59375% relative to that specifically defined baseline. Comparing only the nonrotary 128-wide K/V baseline gives 65,536 bytes and 98.2421875%. Neither comparison is a universal “93% MLA saving”; specify exactly what is stored and which model is compared.

A nominal 7B FP16 weight array is about 14 billion bytes, not 27GB. Capacity is (usable_memory - weights - workspace - reservations) / per_request_cache; a 40GB device does not automatically have 13GB free. A 70B FP16 dense model's weights alone exceed one 80GB GPU, regardless of its GQA cache. Quantization or sharding changes that deployment question. The cache can dominate at large batch/context, but there is no universal 2k-token crossover.

def kv_bytes(batch, layers, length, kv_heads, dim, bytes_per_element=2):
    return 2 * batch * layers * length * kv_heads * dim * bytes_per_element
assert kv_bytes(1, 32, 4096, 8, 128) == 536870912
assert kv_bytes(1, 80, 4096, 8, 128) == 1342177280
assert 61 * 4096 * (512 + 64) * 2 == 287834112
for n in (500, 2048):
    print('score counts', n, n*(n+1)//2, n*(n+1)*(2*n+1)//6)

5 · Failure fingerprints and serving tradeoffs

Wrong positions: cached and uncached logits diverge immediately after the prefix, even though shapes are valid. Test a multi-token chunk with nonzero prefix, not just the first token. Wrong concatenation dimension: head width grows instead of time, eventually breaking matmul. Prompt reprocessing: outputs may remain right while cost stays high; count forwarded positions. Wrong mask: is_causal=True uses an upper-left causal alignment for nonsquare SDPA in torch2.9, not this chunk's offset triangle. Training mode or stale model: reuse no longer represents the current function. Padding: valid cache length and absolute position need not be inferred from tensor allocation length.

The useful mental model is a growing, repeatedly read ledger. Optimize row size, sharing, allocation and eviction separately:

• PagedAttention: logical token blocks map to physical blocks, enabling incremental allocation and sharing with copy-on-write. It reduces reserved slack and external fragmentation; partially filled blocks and metadata still cost memory. The 2023 vLLM study's throughput gains are workload comparisons, not guaranteed 96% utilization or four times the concurrency everywhere.
• Prefix caching / RadixAttention: index reusable states by an exact token prefix and relevant model/configuration identity. A changed system prompt invalidates later states. Sharing requires compatible positions, adapters and isolation policy; sensitive tenant data must not be exposed through reuse or timing. SGLang's radix tree is a prefix-management mechanism, not a universal 5–10x win.
• MQA/GQA: share one or several KV heads among query heads. The raw reduction is query-head count divided by KV-head count relative to the same MHA layout. Changing a trained checkpoint's projections requires a supported conversion/uptraining strategy and evaluation, not a serving toggle.
• MLA: cache a compressed joint latent and a separate rotary key; absorb nonrotary up-projections into query/output work. This changes the model parameterization. The extra rotary component is why the cache is larger than the latent alone.
• Quantized KV / KIVI: fewer bits reduce payload bytes, but scales, groups, residual full-precision tokens and conversion work remain. KIVI distinguishes per-channel key and per-token value quantization. Two-bit versus sixteen-bit payload is an eightfold ideal reduction, not fourfold; the paper's total memory result includes other storage. Evaluate long-context retrieval and generation, not only short perplexity.
• Speculative decoding: propose several tokens, verify under the target and handle accept/reject with a correct algorithm. Exact variants preserve the target distribution; speed depends on acceptance, proposal cost and verification batching. Roll back rejected KV positions. Multi-token prediction heads can supply proposals but do not guarantee gains.
• Sliding window plus attention sinks: retain initial sink tokens and recent context to bound stored state. StreamingLLM studies stability under that policy; discarded mid-history is not still available. Do not equate this with Mistral's trained sliding-window architecture or advertise unlimited recall.

Start with exact paging/sharing where applicable, then evaluate numerical or architectural changes. Eviction changes available information; quantization changes numerics; GQA/MLA changes the architecture. These are not interchangeable “free” optimizations.

6 · Recall and worked answers

1. Why cubic versus quadratic? Sum full-prefix score-grid sizes t squared versus cached rows t. Add the separate projection/FFN terms and prompt prefill; the CPU benchmark need not have those slopes.
2. What is Llama 3 70B's 4k cache? 2*80*4096*8*128*2 = 1,342,177,280 bytes, 1.25GiB per sequence. Batch multiplies it; query heads do not replace KV heads.
3. Two common bugs? Position reset breaks logit parity after prefill; recomputing the prompt preserves redundant work. The wrong-position negative control above distinguishes correctness from performance.
4. How much fragmentation does paging save? Measure allocated live blocks, partly filled final blocks and previously reserved slack for the actual request distribution. There is no fixed percentage to memorize.
5. When omit the cache? A single full-sequence classification/scoring pass with no continuation has no cross-step reuse. Training normally uses a full causal pass and backprop rather than this inference cache. Memory/latency tradeoffs may favor recomputation, but long prefill alone does not prove it cheaper; benchmark the intended workload.
6. What is cached in V3 MLA? The joint 512-dimensional latent and 64-dimensional rotary key. At FP16 that is1152 bytes/layer/token. Explain the expanded baseline before quoting a percentage. The noncommutation of position-dependent rotation with a shared low-rank up-projection motivates decoupled RoPE, developed next.

Exit exercise: add a cache-capacity error test, compare arbitrary chunk partitions to full logits, then change only batch or context in the benchmark. On tomorrow's recall, derive the mask and memory formula without the table. On day seven, recompute the table rather than memorizing its old incorrect numbers.

Sources and optional navigation

Primary references: Transformers4.57.1 cache contract, torch2.9 SDPA, PagedAttention, SGLang, KIVI, StreamingLLM, Llama3, DeepSeek-V3 §2.1.1 and §4.2. Accessed/snapshot identities and exact claim passages are retained in the batch evidence; CPU arithmetic is separate from paper results.

Optional inherited resources, not evidence for speed guarantees: KV cache video, inference techniques video, Pope et al., KV explanation, Character.AI engineering. Course map · Previous DL041 · Next DL043.