The KV Cache

TL;DR: The KV cache keeps attention keys and values from tokens already processed, so later decode steps reuse them. It saves repeated computation but consumes memory as context grows. For Llama 3.1 70B with BF16 cache, each retained token costs 327,680 bytes across the model, before allocation overhead.

Follow one token

A typical situation: your model loads successfully, but requests start waiting when several users paste long documents. The model weights still fit. Each conversation has created its own growing attention state beside them in GPU memory, also called HBM.

A token is a unit of model input, often a word fragment. At each attention layer, a processed token produces a query, a key and a value. The query scores which positions to attend to; keys are the vectors it scores against; values carry the information combined using those scores.

For ordinary causal self-attention, a token cannot see future tokens. Its saved keys and values therefore stay valid as the same sequence grows. Reuse requires the same model and preceding context, not merely the same word at the same position in a different conversation.

Read the diagram downward. D was predicted on the previous step; this step consumes it and predicts E. Prefill vs Decode explains how processing the prompt produces the first output token. The emitted token gets its own cache entry when the model consumes it on a later step. Queries need not persist for future decode steps. The Transformers caching explanation shows this per-layer update.

Caching removes repeated processing of old positions. It does not make a whole growing response linear-time: with full attention, each new query still reads an increasing history. Sliding-window, compressed-latent and hybrid models need different state accounting.

Calculate the memory, with units

For uniform full-attention MHA or GQA layers, count the stored vectors:

bytes per retained token = 2 × layers × KV heads × head dimension × bytes/element
Llama 3.1 70B, BF16 KV = 2 × 80 × 8 × 128 × 2
                      = 327,680 bytes = 320 KiB
8,192 retained tokens = 327,680 × 8,192 = 2,684,354,560 bytes
                      = 2.5 GiB ≈ 2.68 GB per sequence
sanity: four such sequences need 10 GiB of KV payload alone.

Three terms in the formula need a picture. Attention runs as several independent heads per layer, each looking at the sequence in its own way. The head dimension is the length of one head's key or value vector, here 128 numbers. Multi-head attention (MHA) gives every query head its own key and value head. Grouped-query attention (GQA) lets several query heads share one KV head, which is why this model stores 8 KV heads while computing 64 query heads. Had it used MHA, the same formula would give 2 × 80 × 64 × 128 × 2 = 2,621,440 bytes, 2.5 MiB per token, eight times more.

The leading 2 counts keys and values. BF16 contributes a separate 2 bytes per element. The Llama 3.1 70B architecture in NVIDIA’s published configuration specifies 80 layers, 8 KV heads and 64 query heads; 8,192 hidden dimensions divided by 64 gives head dimension 128. Our calculation assumes BF16 cache independently of the weight format. Use KV heads, not query heads.

These are payload-only upper bounds with no prefix sharing. The budget is what remains after weights and runtime allocations. Reserve room for output growth and block rounding. In a sharded deployment, check each GPU's placement and remaining capacity; a sum across GPUs can hide one full rank. KV Cache Sizing carries the planning procedure.

Why a cache can still be slow

The 8,192-position example holds 2.5 GiB. At 32 active sequences, the logical K/V payload attended to is 80 GiB per decode step. That is not a measured HBM transfer count: kernels, reuse, sharding and dtype determine physical traffic. It explains why memory capacity and memory bandwidth are separate constraints. A request may fit and still miss its latency target.

Start with normal KV reuse for autoregressive inference. When capacity binds, bound admitted work and output growth before increasing concurrency. Cache quantization or offload can help, but validate quality and transfer latency. The reversal condition for disabling caching is a workload with no useful reuse, such as a single forward pass, or a correctness/debugging experiment that requires recomputation.

What changes when a request ends

PagedAttention manages physical blocks beneath the logical sequence. Completion releases the request's references; a prefix cache may retain reusable blocks. Reserved GPU memory can therefore remain high even when no request is running. Prefix Caching and KV Reuse explains that separate lifetime.

For a deployment using vLLM, compare its reported KV token capacity with the workload's retained positions, then inspect waiting requests and preemption trends. A rising preemption count signals insufficient cache allocation for scheduled work, but does not by itself prove that compute has spare capacity. Weight quantization and KV quantization are separate choices; changing one dtype need not change the other.

What interviewers are listening for

They want the vector count and the scope of the formula before the multiplication. If asked whether a mixture-of-experts model has a smaller cache because fewer experts run, explain that expert activity describes the feed-forward layers. Attention configuration determines attention state.

Check yourself: with the hypothetical 20 GiB budget, can eight sequences grow from 8,192 to 9,216 retained positions? Each adds 1,024 × 320 KiB = 320 MiB; eight add 2.5 GiB. The payload becomes 22.5 GiB, so those admissions left no growth room.

Key takeaways

• Llama 3.1 70B BF16 KV costs 320 KiB per retained position, before overhead.
• An 8,192-position sequence holds 2.5 GiB; count prompt state and output growth.
• Saved K/V avoids repeated computation; full-attention history still gets read.
• Use per-rank allocation and latency evidence before increasing concurrency.