Why is there no cache for the queries in a transformer?

Because nothing reuses a past query. A token's query is consumed once to produce that token's own output, and causal masking ensures later tokens never attend back through an earlier query.

Implementing a KV Cache From Scratch: Pure PyTorch

Q: Does a KV cache change the model's output?

No. A KV cache is an exact mathematical identity, not an approximation. The cached keys and values are precisely the ones the uncached path would recompute, guaranteeing the exact same token generation.

Q: If the KV cache makes decoding cheap, why is long context still slow?

Because the attention operation at each step is still a query against the entire cache, meaning its cost rises linearly with sequence length. Additionally, the memory footprint of the cache grows linearly, making inference heavily memory-bandwidth-bound.

The waste in naive generation

Autoregressive generation produces one token at a time, and each token is predicted from the whole sequence so far. The obvious way to write the loop is to append the new token and feed the entire growing sequence back into the model. That works, but it does an enormous amount of repeated work. Inside each attention layer, the model projects every token in the sequence to a key and a value and computes attention over all of them, and it does this from scratch on every step. Yet causal attention means a token only ever attends to itself and the tokens before it, so the key and value of any earlier token are exactly the same on step fifty as they were on step ten. A naive approach would recompute them anyway, wasting a large amount of compute.

Caching

On each step we project only the new token to its key and value, append them to a stored cache, and let the new token's query attend over the whole cache. The earlier keys and values are read, never recomputed. This is not an approximation: the numbers that come out are the same, because we are reusing exactly the values the naive loop would have recomputed.

Cached attention in practice

Practically, the change occurs in the attention head. The cached version differs from a naive one in two ways: it accepts the previous keys and values and concatenates the new ones onto them, and it returns the extended pair so the caller can store it. The causal mask needs a little change, because the incoming chunk may be a long prompt during the first pass or a single token afterward. A query at absolute position p may attend to every key up to and including p, which a single triangular mask offset by the cache length expresses for both cases at once.

import torch, torch.nn as nn
from torch.nn import functional as F

class Head(nn.Module):
    def __init__(self, n_embd, head_size):
        super().__init__()
        self.key   = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)

    def forward(self, x, past=None):
        q, k, v = self.query(x), self.key(x), self.value(x)     # (B, T_new, head_size)
        if past is not None:
            pk, pv = past
            k = torch.cat([pk, k], dim=1)                       # read + extend the cache
            v = torch.cat([pv, v], dim=1)
        T_new, T_tot = q.size(1), k.size(1)
        wei = (q @ k.transpose(-2, -1)) * k.size(-1) ** -0.5    # (B, T_new, T_tot)
        # query row i is at absolute position (T_tot - T_new) + i, and may see keys 0..that
        mask = torch.tril(torch.ones(T_new, T_tot, device=x.device), diagonal=T_tot - T_new)
        wei = wei.masked_fill(mask == 0, float("-inf"))
        out = F.softmax(wei, dim=-1) @ v
        return out, (k, v)

The rest of the model just threads the cache through. Each multi-head block keeps a list of its heads' key-value pairs, each transformer block passes the cache to its attention, and the top-level model holds one entry per layer. The single subtlety is the position embedding: with a cache, the new token is not at position zero of what the model sees, it is at the absolute position given by how much is already cached, so we offset the position indices accordingly.

class MultiHead(nn.Module):
    def __init__(self, n_head, n_embd):
        super().__init__()
        self.heads = nn.ModuleList([Head(n_embd, n_embd // n_head) for _ in range(n_head)])
        self.proj  = nn.Linear(n_embd, n_embd)

    def forward(self, x, past=None):
        outs, new = [], []
        for i, h in enumerate(self.heads):
            o, kv = h(x, None if past is None else past[i])
            outs.append(o); new.append(kv)
        return self.proj(torch.cat(outs, dim=-1)), new          # new = list of (k, v) per head

class Block(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        self.sa  = MultiHead(n_head, n_embd)
        self.ff  = nn.Sequential(nn.Linear(n_embd, 4 * n_embd), nn.ReLU(),
                                 nn.Linear(4 * n_embd, n_embd))
        self.ln1, self.ln2 = nn.LayerNorm(n_embd), nn.LayerNorm(n_embd)

    def forward(self, x, past=None):
        a, new = self.sa(self.ln1(x), past)
        x = x + a
        x = x + self.ff(self.ln2(x))
        return x, new

class TinyTransformer(nn.Module):
    def __init__(self, vocab_size, n_embd, n_head, n_layer, block_size):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.blocks  = nn.ModuleList([Block(n_embd, n_head) for _ in range(n_layer)])
        self.ln_f    = nn.LayerNorm(n_embd)
        self.head    = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, past=None):
        T = idx.size(1)
        past_len = 0 if past is None else past[0][0][0].size(1)  # cached sequence length so far
        pos = torch.arange(past_len, past_len + T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        new_past = []
        for i, block in enumerate(self.blocks):
            x, np_ = block(x, None if past is None else past[i])
            new_past.append(np_)
        logits = self.head(self.ln_f(x))
        return logits, new_past

Toy model for generation

The cache is an exact identity for any weights, so training is not strictly necessary to demonstrate it, but a trained model generates real sentences and makes the equality check more convincing.

import random
random.seed(0); torch.manual_seed(0)

DET, ADJ = ["the", "a"], ["red", "blue", "tall", "small", "quick", "slow", "bright", "calm"]
NOUN = ["fox", "cat", "dog", "bird", "tree", "rock", "lake", "hill"]
VERB = ["sees", "watches", "chases", "holds", "finds", "likes"]
def sentence():
    return (f"{random.choice(DET)} {random.choice(ADJ)} {random.choice(NOUN)} "
            f"{random.choice(VERB)} {random.choice(DET)} {random.choice(ADJ)} {random.choice(NOUN)}. ")

corpus = "".join(sentence() for _ in range(3000))
chars = sorted(set(corpus)); vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}; itos = {i: c for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in corpus])

block_size = 256
device = "cuda" if torch.cuda.is_available() else "cpu"

def get_batch(bs=64):
    ix = torch.randint(len(data) - block_size, (bs,))
    x = torch.stack([data[i:i + block_size]         for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    return x.to(device), y.to(device)

model = TinyTransformer(vocab_size, 96, 4, 3, block_size).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
for step in range(2000):
    xb, yb = get_batch()
    logits, _ = model(xb)
    loss = F.cross_entropy(logits.view(-1, vocab_size), yb.view(-1))
    opt.zero_grad(set_to_none=True); loss.backward(); opt.step()

Prefill and decode

With the cache in place, generation splits into two phases. The first, prefill, runs the prompt through the model once to populate the cache, a single parallel pass over all the prompt tokens. The second, decode, generates one token at a time, each step feeding only the newest token and extending the cache. The naive function below keeps no cache and re-runs the whole sequence; the cached one shows the prefill-then-decode structure. Both decode greedily so the comparison is exact.

@torch.no_grad()
def generate_naive(model, idx, n_new):
    model.eval()
    for _ in range(n_new):
        logits, _ = model(idx)                          # re-run the whole sequence, no cache
        nxt = logits[:, -1].argmax(-1, keepdim=True)
        idx = torch.cat([idx, nxt], dim=1)
    return idx

@torch.no_grad()
def generate_cached(model, idx, n_new):
    model.eval()
    logits, past = model(idx)                           # PREFILL: process the prompt once, fill the cache
    nxt = logits[:, -1].argmax(-1, keepdim=True)
    out = [idx, nxt]
    for _ in range(n_new - 1):
        logits, past = model(nxt, past)                 # DECODE: feed one token, extend the cache
        nxt = logits[:, -1].argmax(-1, keepdim=True)
        out.append(nxt)
    return torch.cat(out, dim=1)

Equality check

Caching is supposed to change the speed, not the output, which is always a good sanity check. Generated greedily from the same prompt, the two functions must produce the identical token sequence.

prompt = torch.tensor([[stoi[c] for c in "the quick fox "]], device=device)
a = generate_naive(model, prompt, 40)
b = generate_cached(model, prompt, 40)

print("identical tokens:", torch.equal(a, b))
print("naive :", "".join(itos[i] for i in a[0].tolist()))
print("cached:", "".join(itos[i] for i in b[0].tolist()))

The output validates outputs are the same:

identical tokens: True
naive : the quick fox chases a slow tree. a bright tree holds 
cached: the quick fox chases a slow tree. a bright tree holds

Computational benefit

The saving comes from not redoing the earlier tokens. Without a cache, generating the token at position t reprojects all t keys and values and computes a full t-by-t attention, so the work on each step grows with the length already produced. With a cache, that step projects one new key and value and computes a single query against t cached keys. Summed over a generation of length n, the difference is an order in n.

Work to generate n tokens

K/V projections attention scores

naive sum of t = O(n²) sum of t² = O(n³)

cached n = O(n) sum of t = O(n²)

Practically, we can measure the speed up as follows:

import time

@torch.no_grad()
def timed(fn, model, prompt, n_new, reps=3):
    fn(model, prompt, 8)                                 # warmup
    t0 = time.time()
    for _ in range(reps): fn(model, prompt, n_new)
    return (time.time() - t0) / reps

prompt = torch.tensor([[stoi[c] for c in "the quick fox "]], device=device)
for n in [32, 64, 128]:
    tn = timed(generate_naive,  model, prompt, n)
    tc = timed(generate_cached, model, prompt, n)
    print(f"generate {n:3d} tokens   naive={tn*1e3:7.1f} ms   cached={tc*1e3:7.1f} ms   speedup={tn/tc:.1f}x")

On a small model and short sequences the absolute numbers are tiny and the speedup is modest, since the toy spends much of its time in fixed Python overhead; the point is the trend, that the ratio grows with length. Results for reference:

generate  32 tokens   naive=   58.1 ms   cached=   13.5 ms   speedup=4.3x
generate  64 tokens   naive=  111.9 ms   cached=   25.7 ms   speedup=4.4x
generate 128 tokens   naive=  276.0 ms   cached=   49.2 ms   speedup=5.6x

Memory overhead

We now hold a key and a value for every token, in every head, in every layer, and that store grows linearly with the sequence length. The size is two, for keys and values, times the number of layers, times the model dimension, times the sequence length, times the batch, times the bytes per number.

@torch.no_grad()
def cache_kib(model, seq_len):
    _, past = model(torch.zeros(1, seq_len, dtype=torch.long, device=device))
    total = sum(k.numel() * k.element_size() + v.numel() * v.element_size()
                for layer in past for (k, v) in layer)
    return total / 1024

for n in [32, 64, 128]:
    print(f"seq_len={n:4d}   measured KV cache = {cache_kib(model, n):6.1f} KiB")

def cache_estimate_mb(n_layer, d_model, seq_len, batch=1, bytes_per=2):
    return 2 * n_layer * d_model * seq_len * batch * bytes_per / 1024 / 1024

# extrapolate the same formula to a production-sized model (fp16, full multi-head attention)
for n in [2048, 8192, 32768]:
    print(f"7B-ish (32 layers, d=4096) at {n:5d} tokens ~ {cache_estimate_mb(32, 4096, n):7.1f} MB")

A seven-billion-parameter model in half precision with full multi-head attention spends on the order of half a megabyte of cache per token, so a few thousand tokens of context is already gigabytes, and that is per sequence in the batch. At long context and high batch, the KV cache, not the weights, is the memory that runs out first, which is why so much inference engineering is about managing it. Here is the output for our small toy model:

seq_len=  32   measured KV cache =   72.0 KiB
seq_len=  64   measured KV cache =  144.0 KiB
seq_len= 128   measured KV cache =  288.0 KiB
7B-ish (32 layers, d=4096) at  2048 tokens ~  1024.0 MB
7B-ish (32 layers, d=4096) at  8192 tokens ~  4096.0 MB
7B-ish (32 layers, d=4096) at 32768 tokens ~ 16384.0 MB

Where this departs from production

The mechanism is exactly the production one, but the surrounding engineering is where the real work lies. Our toy example grows the cache by concatenating a tensor each step, which reallocates; real implementations preallocate a fixed buffer and write each new key and value into the next slot. The prefill and decode split we sketched is a load-bearing distinction at scale: prefill is compute-bound and parallel over the prompt, while decode is memory-bandwidth-bound and strictly sequential, one token at a time, which is why decoding is the slow part and why the cache, which makes decode cheap per step, matters.

Because the cache is usually the binding memory constraint, much of modern serving is built around shrinking or managing it. Grouped-query and multi-query attention share keys and values across query heads, cutting the cache by a large factor, which is why most recent large models use them. Paged attention, the idea behind vLLM, stores the cache in fixed-size pages so that many sequences of different lengths can share memory without fragmentation, raising the batch size a server can hold. And the cache can be compressed directly, by quantizing the stored keys and values to eight or fewer bits, or bounded by keeping only a sliding window of recent tokens. All of these are refinements of the one thing you built above: store the keys and values instead of recomputing them.

Common questions

Does a KV cache change the model's output?

No. It is an exact identity, not an approximation: the cached keys and values are precisely the ones the uncached path would recompute, so the logits and the sampled tokens are the same. If you ever see a cached and uncached run disagree, that is a bug, usually in the position offset or the causal mask.

Why is there no cache for the queries?

Because nothing reuses a past query. A token's query is consumed once, to produce that token's own output, and the causal mask means later tokens never attend back through an earlier query. Keys and values are different: every future token's query will attend over them, so they are worth keeping. The new token still computes its own query fresh each step; there is simply nothing to cache.

If the cache makes decoding cheap, why is long context still slow and expensive?

Because two things still grow with length. The attention at each step is a query against the whole cache, so its cost rises linearly with how much you have generated, and the cache memory itself grows linearly too. The cache removes the redundant recomputation of earlier tokens, turning a cubic total into a quadratic one, but it does not make attention free, and it shifts the pressure onto memory, which is why long-context inference is dominated by the size and bandwidth of the KV cache.

Implementing a KV Cache From Scratch

Quick Answer / TL;DR