The KV cache

An application prompt is mostly boilerplate, and the model has no memory of having read it before. Instructions, tool schemas, worked examples, a set of retrieved documents and a conversation's accumulated history are identical on every call, they routinely run to tens of thousands of tokens, and without reuse the GPU processes all of them from the beginning every time a user types a sentence. Prefix caching removes that cost for the portion of the prompt that did not change and leaves everything else exactly as it was.

This chapter is about what gets kept and why keeping it is exact rather than an approximation, which is the mechanical core of everything else in this topic. Prefill — the phase that processes the whole prompt in parallel before any output token is produced — is where the work happens, and the question here is narrower than how fast prefill runs: of the work prefill does, how much is a function of the prompt alone, such that two requests sharing an opening are literally performing the same arithmetic twice? You will meet the answer in two packagings. Hosted APIs expose it as prompt caching, where you mark how far the stable region of the prompt extends and the response reports how many tokens were read from cache and how many were written into it. Self-hosted engines such as vLLM, SGLang and TensorRT-LLM implement it as automatic prefix caching, which is commonly on by default and requires no change to the request at all, so a team can depend on it for months before anyone discovers that it is there. Either way it is a serving-layer mechanism rather than something the model knows about, and it either fires or does not fire depending on properties of your prompt that the rest of this topic is about.

What a decoder holds on to while it generates#

A language model reads a prompt one position at a time, and at each position it has to decide what else in the sequence is worth looking at. The word "it" in the ninetieth token of a paragraph means nothing on its own, and whatever it refers to is somewhere behind it. Attention is the mechanism that does that finding. It works by having every position advertise what it contains and separately ask for what it needs, with both the advertisement and the request expressed as lists of numbers derived from the token by multiplying it against a matrix of trained weights. How much position j matters to position i is then decided by how well i's request lines up with j's advertisement.

Three such lists exist per position, and their names are borrowed from databases rather than from linguistics. The query is the request: what position i is looking for. The key is the advertisement: what position j offers to anyone looking. The value is the payload: what position j actually contributes once it has been selected. A model is not one of these mechanisms but dozens of them stacked in layers, each layer refining the representation the layer below produced, and inside each layer the work is split across several heads that attend to different things in parallel. The running representation every layer reads and writes — one vector per position, carrying everything the model has worked out about that position so far — is the residual stream.

Inside every self-attention layer, the residual stream at each position is projected through three weight matrices into a query, a key, and a value. The query for position i is used exactly once — it is dotted against the keys of positions 0 through i, the resulting weights are applied to those positions' values, and then that query is of no further use to anybody. The keys and values are different. When the model later computes position i+1, it needs the keys and values of every position at or before i+1, including all the ones it already computed. When it computes position i+2, it needs them again. Over a 500-token generation, the key and value vectors of the first prompt token are consulted 500 times.

The word quadratic in the next paragraph is worth unpacking, because the size of that number is what forces the design. Producing output token t without a cache means running the forward pass over all t positions, since every one of them needs a key and a value present, and producing token t+1 means running it over t+1 positions. Adding those costs up across a whole generation gives a total that grows with the square of the length rather than with the length itself, so doubling the number of output tokens roughly quadruples the arithmetic. Keeping the keys and values instead makes each step cost one position's worth of work, so the total grows in proportion to the length, and the gap between those two curves is the entire reason the data structure exists.

Recomputing them 500 times would mean re-running the full forward pass over the whole sequence for every single output token, which is quadratic in the sequence length and enormously wasteful. So implementations keep them. Every layer holds two tensors that grow by one row per token generated, and that pair of tensors, across all layers, is the KV cache in its exact form. For a batch of one at layer l, each has the shape [n_kv_heads, seq_len, head_dim].

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map="cuda"
)

ids = tok("The refund policy for enterprise accounts", return_tensors="pt").to("cuda")
out = model(**ids, use_cache=True)

cache = out.past_key_values
print(len(cache))                 # 32   -> one entry per decoder layer
print(cache[0][0].shape)          # torch.Size([1, 8, 7, 128])
#                                     batch, kv_heads, seq_len, head_dim
print(cache[0][0].dtype)          # torch.bfloat16

Seven tokens in, thirty-two layers deep, eight key/value heads per layer, one hundred twenty-eight dimensions per head. That is the entire durable residue of a prefill: not the logits — the vocabulary-sized vector of scores from which the next token is chosen — not the attention weights, not the intermediate MLP activations, all of which are computed and discarded. Just K and V, per layer, per head, per position.

The dependency that makes a prefix reusable#

Now the property the whole topic rests on, walked one link at a time. Every link below is separately checkable, and skipping any of them leaves a rule a reader can recite and cannot apply to a case they have not been shown.

Start at the top of the stack. The key vector at position i in layer l is produced by multiplying the residual stream at position i, as it enters layer l, by that layer's key matrix. Two things go into that multiplication and nothing else: one vector, and one fixed set of weights. So whatever the key turns out to be, it is decided entirely by what the residual stream at position i held on the way in.

Step down one layer and ask where that residual stream came from. It is what layer l−1 wrote at position i, which is the sum of what layer l−1 read at position i and the attention output it computed there. That attention output is a weighted mixture of the values at the positions i was allowed to look at, and the set of positions i was allowed to look at is 0 through i inclusive. Nothing from i+1 onwards entered the mixture.

That restriction is the causal mask, and it is a deliberate constraint written into the attention computation rather than a property of the text. Before the attention weights are normalised, the score between position i and any position after it is set to negative infinity, which drives its weight to zero. It is there because the model is trained to predict each token from the ones before it, and a model that could see the next token during training would learn to copy it instead of predicting it. The mask is what makes the training task non-trivial, and the property this chapter needs is its side effect.

Now repeat those two steps all the way down. Every layer's output at position i depends on that layer's input at positions 0 through i. Every layer's input at those positions is the layer below's output at those same positions. And the bottom of the stack is the embedding lookup, which turns token ID i into a vector by indexing a table, with no reference to anything else in the sequence at all. There is no rung on that ladder where information from a later position can enter, because the mask closes the only door it could have come through.

The conclusion therefore has no escape hatches: K and V at position i are a pure function of token IDs 0 through i and the model weights. Nothing that arrives after position i can reach back and alter them. Neither can the sampling temperature, the requested output length, the other sequences sharing the batch, or the user who sent the request.

In plain terms, a token's key and value are its contribution to everything behind it, computed from everything in front of it. The past does not change, so neither does the contribution. That is why "recompute" and "reuse" are not two approximations of one another but two names for the same number.

The chain also says exactly what would break it, which is more useful than the rule on its own. Any mechanism that lets a later position affect an earlier one's key or value cuts the argument at its first link. A bidirectional encoder, where every position attends to the whole sequence, has no reusable prefix at all, because appending a token changes every key already computed. The same goes for any scheme that normalises across the full sequence length before the mask is applied. Reuse of the kind in this chapter is a property of causal decoders specifically, and "does anything downstream reach back" is the question to put to an unfamiliar architecture before assuming the property survives.

That is why serving a repeated prefix is cheap, and it is worth being precise about the claim. It is not that reuse is a good approximation of recomputation, nor that the difference is small enough not to matter. Given identical token IDs and identical weights, the cached tensors are the tensors a recompute would produce. The arithmetic was already done and its result is sitting in memory. Feeding the model a cached prefix plus one new token costs a single forward step over one position, rather than a forward pass over thousands of them.

Two qualifications belong here so they do not surprise you later. Bit-for-bit equality across a different batch shape is not guaranteed, because GPU matmul kernels reduce in different orders at different tile sizes and floating-point addition is not associative; the differences land around 1e-3 in bf16, which is noise relative to the logit gaps that decide sampling. And the reuse holds only within one model and one precision — a cache written by an fp8 build of the same checkpoint is not usable by the bf16 build, because those are different functions of the same tokens.

What a token of cache costs in memory, and what it saves in arithmetic#

Both sides of the trade are calculable from the architecture. The memory per token is 2 × layers × kv_heads × head_dim × bytes_per_value — the leading 2 being one tensor for K and one for V.

Each of the other four factors is there for a reason worth naming, because a reader who knows the reasons can predict the figure for a model that is not in the table. Layers appears because every layer computes its own keys and values and every layer needs them again on the next step, so the store is not one tensor but one pair per layer, and a deeper model pays proportionally more. It is KV heads rather than the total number of attention heads because grouped-query attention lets several query heads share one key/value head, and only the shared ones are ever stored. Head dim is how many numbers make up a single head's key vector, so it converts a count of heads into a count of values. And bytes per value is the precision the tensors are held in, which is why the same architecture served in bf16 and in fp8 gives two answers to the same question. Multiply those four for the number of values a token contributes, double it because there are two tensors, and the result is in bytes.

Run it for the second row of the table. Llama 3 8B has 32 layers, 8 key/value heads and 128 dimensions per head, and bf16 holds each value in 2 bytes, which comes to 128 KiB for every token in the sequence. That figure is per token rather than per request, so the cost of holding a 6,100-token prefix is 128 KiB multiplied by 6,100 tokens, which is the 762 MiB in the last column. Nothing in that calculation refers to the prompt's content or the user, which is why cache capacity is arithmetic rather than observation.

The gap between the first row and the second is grouped-query attention, which shares one key/value head across several query heads and was adopted in large part because the full multi-head arrangement made the cache ruinous at long context. A 32k-token conversation on Llama 2 7B would need 16 GiB of cache for a single sequence, which is most of a 24 GB card before a single weight is loaded.

Turning the other side of the trade into a number needs one rule of thumb about what a forward pass costs. Nearly all of a transformer's arithmetic is matrix multiplication against its weights, and in a matrix multiplication each weight is used exactly once per token: multiplied by an input, then added into a running total. That is two floating-point operations per weight per token, which is where the leading 2 in the estimate below comes from. Multiplying it by the parameter count and then by the token count accounts for the whole prefill. The estimate ignores the attention score computation itself, which is why it is a rule of thumb rather than an identity, but at these prompt lengths the weight multiplications dominate and the figure is close enough to reason with.

The arithmetic saved is the other half. Prefill costs roughly 2 × n_params × n_tokens FLOPs, so a 6,100-token prefix on an 8B model is about 9.8 × 10¹³ floating-point operations. An A100 sustaining perhaps 150 TFLOP/s on that shape gets through it in around 0.65 seconds, which is the bulk of a cold time to first token — TTFT, the delay between the request arriving and the first character of the response leaving — once you add tokenizer, scheduling, and network overhead. Skipping it leaves a 40-token suffix, which is 6.4 × 10¹¹ FLOPs and finishes in single-digit milliseconds. That is the entire trick: 762 MiB of resident memory buys the elimination of two thirds of a second of GPU time, every time that prefix is served again.

Both of the times in that paragraph are a single division, and doing it in the open is worth more than trusting the result. A TFLOP/s is a trillion floating-point operations a second, so dividing 9.8 × 10¹³ operations by 150 × 10¹² operations a second leaves seconds, and what comes out is the 0.65 quoted. The 40-token suffix is the same division with a much smaller numerator, since 40 tokens is a small fraction of 6,100 and the operation count falls in the same proportion, which is how 6.4 × 10¹¹ finishes in single-digit milliseconds. Neither figure is a benchmark. Both are an operation count divided by a throughput, and the shape of the answer — two thirds of a second against a handful of milliseconds — survives being wrong about that throughput by a factor of two.

What appending one token actually costs#

Watching the shapes change makes the asymmetry concrete. Feed the model a prefix, keep the cache, then feed it a single additional token and let it reuse what is there.

prefix = tok("You are a support agent. " * 400, return_tensors="pt").to("cuda")
out = model(**prefix, use_cache=True)
print(out.past_key_values[0][0].shape)   # [1, 8, 2000, 128]  — 2,000 tokens cached

nxt = tok(" Refund", return_tensors="pt").to("cuda")
step = model(
    input_ids=nxt.input_ids,                       # exactly one token
    past_key_values=out.past_key_values,
    use_cache=True,
)
print(step.logits.shape)                 # [1, 1, 128256]
print(step.past_key_values[0][0].shape)  # [1, 8, 2001, 128]  — grew by one row

Inside each attention layer of that second call, the query matrix has a single row and the key matrix has 2,001, so the score computation is one vector against two thousand rather than two thousand against two thousand. The MLP blocks, which dominate the parameter count, run over a single position rather than two thousand of them. That is why the second call finishes in about seven milliseconds while the first took roughly two hundred: the work is linear in new tokens, not in sequence length.

The MLP half of that explanation is where most of the saving actually lives, and it is worth separating from the attention half. A decoder block contains the attention projections and a feed-forward network, both of them matrix multiplications against fixed weights, and the feed-forward network holds the larger share of the parameters. Its cost is therefore proportional to how many positions are pushed through it, with no dependence on sequence length whatsoever: feed it two thousand positions and it does two thousand positions of work. Attention is the part whose cost does depend on how much history is present, and even there the cached step compares a single query row against the resident keys rather than every row against every row.

The query is the thing that is never cached, and the asymmetry is not an implementation choice. A query is consumed the instant it is used — it asks "what should position i look at" and, once the answer is computed, position i never asks again. Keys and values are answers rather than questions: they describe what a position offers to any later position that asks, so every future token consults them. Cache the answers, discard the questions. A serving stack that stored Q as well would double its memory footprint to hold tensors that nothing will ever read.

One consequence of this shape is worth carrying into the rest of the topic. Because the cached step touches every layer's full key and value tensors, the cost of a decode step grows with how much cache is resident, slowly but measurably. Reuse makes the first token dramatically cheaper and makes each subsequent token slightly more expensive, and both effects come from the same tensors being present.

The limits of a cache hit#

The word "cache" is what makes the next four paragraphs necessary. In almost every other system an engineer meets, a cache stores a finished answer under the request that produced it, and a hit means that answer comes back without the work being done. None of that describes this store, and reading across from the familiar meaning produces predictions that are confident and false.

It is not a response cache. Nothing about the model's output is stored, and a second request still samples freshly from the logits — a cache hit on a 6,100-token prefix followed by two different questions produces two entirely different answers, and the same question asked twice at temperature 0.8 can still diverge. What is reused is a stage of the computation, not a conclusion.

It is not compression either. The cache does not make the prefix smaller in any sense; it makes it larger. Six thousand tokens that occupied about 24 KB as integer IDs become 762 MiB as key and value tensors, a thirty-thousand-fold expansion in exchange for not having to derive them again. Every design decision later in this topic — eviction, block sizes, why a batch runs out of room — follows from the fact that this is a memory-for-compute trade with a brutal exchange rate.

It is not model-agnostic, either, and this catches people who think of it as a data cache. The tensors are a function of the weights, so they are valid for exactly one checkpoint at one precision. Swap in a fine-tuned variant, promote a new revision of the same model, or serve the same weights quantised to fp8, and every entry becomes garbage — not slightly stale, but arithmetically unrelated to what the new model would have produced. Nothing in the tensors themselves records which model made them, which is why the identity of the weights has to be carried in whatever key indexes them, a point that returns with some force in chapter 8, on entries that are keyed wrongly.

Finally, it does nothing for the per-token generation rate. Decode already keeps its own KV cache within a single request; that is how generation works at all. The reuse discussed here is about carrying that cache across requests so a repeated prefix does not need a second prefill. Once generation starts, each new token still costs one serial forward step, and the arrival rate of tokens after the first is exactly what it was before. Chapter 6 takes that apart with measurements of the two phases, because it is the single most over-claimed benefit in the area.

The KV cache is therefore two ideas stacked on one property. The property is that K and V at a position depend on nothing but the tokens before it and the weights, which makes them safe to store and reuse without qualification. The first idea is to keep them within a generation, which every implementation does because the alternative is quadratic. The second is to keep them past the end of a request so that the next prompt beginning with the same tokens skips its prefill, which is the one this topic is about — and which costs 128 KiB per token on a typical 8B model to save roughly 16 GFLOPs per token, a trade that is excellent when the prefix is re-served and merely expensive when it is not.

None of this is worth anything unless the reuse is genuinely exact. A serving stack that returns nearly the right keys and values does not produce an error; it produces slightly different logits, which produce a different token, which produces a different continuation, and the whole thing looks like a plausible answer to a question nobody asked. Before trusting a cache, you want a test that pins the appended token's logits against a full recompute — and when that test fails, the cache is usually innocent.