Attention and the KV Cache: The Arithmetic of a Forward Pass

A large language model feels like it is thinking hard, so it is natural to assume that generating each token is an expensive computation. For the arithmetic that actually matters during serving, the opposite is true: decoding one token is cheap in floating-point operations and expensive in memory movement. The object that makes this true — and that quietly decides how many users you can serve at once — is the key-value cache. This article builds that claim from the ground up: we start from the attention equation itself, count the operations a forward pass performs, show why caching is not an optimization but a necessity, and then compute the cache's size in bytes precisely enough to plan a deployment.

The claim: decoding is memory-bound, not compute-bound

Inference happens in two distinct regimes. The first is prefill: the model reads the entire prompt at once, computing the internal state for every prompt token in parallel. The second is decode: the model emits the answer one token at a time, and each new token must attend to everything that came before it. Prefill is a wide, parallel matrix multiply and is genuinely compute-bound — the hardware's floating-point units are the bottleneck. Decode is the opposite. Each decode step processes a single new token, so there is very little arithmetic to do, but to do it the hardware must stream the model's weights and the accumulated attention state out of memory. The bottleneck is bandwidth, not compute.

This distinction is the thesis of the article, and every number that follows is a consequence of it. If decoding were compute-bound, we would optimize for faster arithmetic. Because it is memory-bound, we optimize for fewer bytes moved — which is why numeric precision, attention grouping, and how the cache is laid out in memory matter more to throughput than the raw speed of the multiply-accumulate units. Define the two regimes now because we will return to them repeatedly: prefill fills the cache, decode reads and grows it.

Scaled dot-product attention, defined

Attention is a weighted average. For a set of query vectors packed as rows of a matrix Q, keys K, and values V, the output is a softmax-weighted combination of the value rows, where the weights come from how well each query matches each key. Vaswani et al. (2017) define it as scaled dot-product attention: the query-key dot products are divided by the square root of the key dimension before the softmax, which keeps the gradients well-conditioned as the dimension grows.

Two facts about this equation drive everything downstream. First, the output for a query depends on the keys and values of every position it is allowed to attend to. Second — and this is the hinge of the whole article — the keys and values for a given token do not depend on any token that comes after it. A token's K and V are computed once from that token's own hidden state, and they never change. That property is what makes caching correct, not merely convenient.

Multi-head shapes and the cost of one token

Real models run attention in parallel heads. The model width d_model is split into h heads of size d_k = d_model / h; each head computes its own attention and the results are concatenated and projected back to width d_model. The projections that produce Q, K, and V are the parameters that matter for our accounting: for each token, the model multiplies that token's hidden state by the weight matrices W_q, W_k, and W_v.

For counting operations, a useful synthesis holds: a single forward pass over one token performs on the order of two floating-point operations per parameter — one multiply and one add per weight touched. So the compute per generated token scales with the parameter count P, while, as we are about to see, the memory the step must touch scales with something else entirely. Holding these two quantities side by side is what reveals the bottleneck.

Why we cache: recomputation is quadratic

Consider decoding without any cache. To produce token n, the model needs the keys and values of tokens 1 through n. With no stored state, it must recompute all of them from scratch at every step. Producing an m-token answer then recomputes the K and V projections roughly 1 + 2 + ... + m times — quadratic work in the sequence length, almost all of it redundant, because those earlier keys and values are identical every time.

The KV cache removes the redundancy. Because a token's K and V never change once computed, we store them the first time and reuse them forever. Prefill computes and stores K and V for every prompt token in one parallel pass. Each decode step then computes K and V for only the single new token, appends them to the cache, and attends against the stored rows. The per-step attention cost drops from recomputing an n-by-d block to reading it: the quadratic collapses to linear growth in stored state.

The pseudocode below shows the loop for a single layer; a real model repeats the same append-and-attend pattern once per layer. Note the one rule that makes it correct and fast: append the new row, never recompute the old ones.

cache_K, cache_V = [], []            # grow by one row per decode step

# Prefill: fill the cache for all prompt tokens in parallel
for h in hidden_states(prompt_tokens):
    cache_K.append(W_k @ h)
    cache_V.append(W_v @ h)

# Decode: one new token at a time
while not done:
    h_t = hidden_state(last_token)
    q   = W_q @ h_t
    cache_K.append(W_k @ h_t)          # append new row ...
    cache_V.append(W_v @ h_t)          # ... never recompute old rows
    scores = softmax((q @ stack(cache_K).T) / sqrt(d_k))
    attn   = scores @ stack(cache_V)
    last_token = sample(project(attn))

How big is the cache? The exact bytes

Now the payload. The cache stores two vectors — a key and a value — per token, per layer. Each vector has width d_model (summed across heads), and each element occupies p bytes, where p depends on numeric precision: 2 bytes for 16-bit, 1 byte for 8-bit, and so on. Multiplying these together gives the cache size for a sequence of n tokens across L layers.

The formula is small but its consequences are not. It is linear in context length n, so doubling the conversation doubles the cache; linear in layers and width, so bigger models pay more per token; and linear in precision p, which is the lever quantization pulls. Work a concrete example: a 7-billion-parameter-class model with L = 32 layers and d_model = 4096, in 16-bit, stores 2 x 32 x 4096 x 2 bytes = 512 KB for every single token. At an 8,000-token context that is about 4 GB of cache — for one sequence. The model weights are fixed, but this grows with every user and every token they generate.

This is why the KV cache, not the weights, usually caps how many sequences you can batch. It also explains two popular levers directly from the formula. Grouped-query attention shares one set of K,V across several query heads, cutting the effective width in the cache by the grouping factor (an 8-to-1 grouping turns our 512 KB/token into 64 KB/token). Quantizing the cache from 16-bit to 8-bit halves p, and therefore halves the bytes, at some cost in precision.

Where the cache breaks: fragmentation

The formula tells you the cache's ideal size, but a running server rarely achieves it, because of how the memory is allocated. A conversation does not announce its final length in advance, so a naive server reserves a contiguous block sized for the maximum context it will allow. Most conversations end early, so most of that reserved block sits empty and unusable by anyone else. Kwon et al. (2023) measured this internal and external fragmentation and found that a large fraction of KV memory was wasted this way — capacity you paid for but could not batch into.

PagedAttention borrows the operating-system idea of virtual memory. Instead of one contiguous reservation, the cache is split into fixed-size blocks that are allocated on demand as a sequence grows and freed when it ends, with a small table mapping a sequence's logical positions to physical blocks. The sequence sees a contiguous cache; the hardware sees packed blocks with almost no waste. The result is that far more sequences fit in the same memory, which for a memory-bound workload translates almost directly into higher throughput.

The lesson generalizes beyond one system: when a resource is your binding constraint, how you allocate it matters as much as how much of it you have.

The roofline: when memory wins

We can make 'memory-bound' precise with arithmetic intensity: the ratio of floating-point operations performed to bytes moved from memory. Every processor has a balance point — the intensity at which its compute throughput and its memory bandwidth are matched. Below that point a kernel is memory-bound and adding compute does nothing; above it the kernel is compute-bound.

Decode sits firmly below the balance point. A single token does on the order of 2P operations but must read the weights (P elements) and the entire KV cache from memory, so its intensity is low and roughly constant regardless of how fast the arithmetic units are. This is the formal statement of the article's thesis, and it tells you where to spend effort: to speed up decode you reduce bytes moved — smaller precision (p), grouped K,V (effective d_model), and batching that reuses the weights across many sequences per fetch — not faster floating-point. It also tells you when the analysis flips: prefill, and very large batches, push intensity up and can become compute-bound, at which point the opposite optimizations apply.

That boundary — memory-bound decode versus compute-bound prefill — is the single most useful mental model for reasoning about inference performance, and it falls straight out of the attention equation and the byte count we derived.

What the arithmetic tells a systems designer

Put the pieces together. Attention makes a token's key and value fixed for life, which licenses the cache; the cache turns quadratic recomputation into linear growth; the cache's size in bytes is a small linear formula in layers, width, context, and precision; fragmentation inflates that ideal size until you allocate it in pages; and arithmetic intensity confirms that at decode time you are moving bytes, not doing math.

The practical consequences are concrete. Your maximum batch size at long context is set by KV memory, not weights, so estimate it from the byte formula before you provision. The highest-leverage knobs are the ones that shrink p (cache quantization), shrink effective width (grouped-query attention), and shrink waste (paged allocation) — in that rough order of how much engineering they cost. And when you profile, expect decode to be bandwidth-limited; a change that adds arithmetic but saves memory traffic is usually a win. The rest of this series takes the same measure-it-precisely posture to the other levers of inference: the accept-reject math of speculative decoding, the error introduced by quantization, and the recall-versus-latency trade-off of the vector indexes that feed retrieval.

1. Size the KV cache with 2 L n d_model p before choosing batch size — weights are not the constraint at long context.
2. Reach for the memory levers first: cache quantization (smaller p), grouped-query attention (smaller effective width), paged allocation (less waste).
3. Read profiles through the roofline: decode is memory-bound, so trade compute for bytes, not the reverse.

Practitioner Toolkit

Copy-paste, strictly defensive artifacts you can use today. Nothing here attacks a real system.

def kv_bytes(layers, d_model, tokens, bytes_per_elem=2, kv_group=1):
    # 2 = one key + one value vector per token per layer
    # kv_group = query-heads per KV-head (1 = multi-head, 8 = GQA 8:1)
    return 2 * layers * d_model * tokens * bytes_per_elem // kv_group

# ~7B-class model, 8K context, fp16, multi-head
print(kv_bytes(32, 4096, 8192))          # -> 4,294,967,296  (~4.0 GB)
print(kv_bytes(32, 4096, 8192, kv_group=8))  # -> ~0.5 GB with GQA 8:1

Glossary

Prefill

The phase that computes and caches the attention state for every prompt token in one parallel pass.

Decode

The autoregressive phase that generates one token at a time, appending to and reading from the KV cache.

KV cache

The stored key and value vectors for every past token, kept so attention need not recompute them each step.

d_model

The model's hidden width; the total dimension of the key and value vectors summed across attention heads.

Scaled dot-product attention

Attention whose query-key scores are divided by the square root of the key dimension before the softmax.

Grouped-query attention

Sharing one set of keys and values across several query heads, shrinking the KV cache by the grouping factor.

PagedAttention

Allocating the KV cache in fixed-size blocks on demand, like virtual memory, to eliminate fragmentation waste.

Arithmetic intensity

The ratio of floating-point operations to bytes moved; low intensity means a kernel is memory-bound.

References

1. Vaswani et al. (2017), Attention Is All You Need
2. Kwon et al. (2023), Efficient Memory Management for LLM Serving with PagedAttention (vLLM)
3. Leviathan et al. (2023), Fast Inference from Transformers via Speculative Decoding
4. Frantar et al. (2022), GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers
5. Lin et al. (2023), AWQ: Activation-aware Weight Quantization for LLM Compression