Deep Learning · Transformers

KV Cache Arithmetic: Bytes per Token, Layers, and the Batch Size Ceiling

How the decode-time key and value cache is sized, why it scales with every layer rather than only the last, and how its growth sets the maximum batch on a fixed device.

Autoregressive decoding would recompute, at step $t$, the keys and values for tokens $1 \dots t$ that were already computed at steps $1 \dots t-1$. The key and value projections alone cost $2 \cdot t \cdot d_{\text{model}}^2$ multiply-accumulates per layer, and that entire amount is redundant work when the results are thrown away and recomputed. Storing the results instead turns the recomputation into a memory read, and the array that holds them is the KV cache. Its size is a fixed function of the architecture, the dtype, and the sequence length, and that function is usually what limits how many requests a device can serve at once.

The size formula and where each factor comes from

For one sequence, one layer, the cache holds two tensors of shape $n_{\text{kv}} \times d_{\text{model}}$ if the key and value projections are full-width, or $n_{\text{kv}} \times h \cdot d_k$ in the same thing written per head. The number of layers multiplies it, because every layer has its own key and value projections and therefore its own cache. The total is

$$\text{bytes} = 2 \cdot L \cdot n_{\text{kv}} \cdot d_{\text{model}} \cdot s$$

where $L$ is the layer count, $n_{\text{kv}}$ the number of cached positions, $d_{\text{model}}$ the width of the key and value projections combined, and $s$ the bytes per element: 2 for 16-bit, 1 for 8-bit.

The factor of 2 is keys plus values. The factor $L$ is the part most often underestimated: a 32-layer model caches 32 separate key tensors and 32 separate value tensors, and no layer can reuse another’s, because the projections differ. It is not a cache of the residual stream; it is a cache of every layer’s intermediate attention state.

A concrete instance makes the scale clear. Take $L = 32$, $d_{\text{model}} = 4096$, 16-bit storage. Per token that is $2 \times 32 \times 4096 \times 2 = 524{,}288$ bytes, about 512 KiB. At 4,096 cached positions the sequence occupies 2 GiB. At 32 concurrent sequences of that length the cache alone needs 64 GiB, before any weights or activations. The weights of the same model, at roughly 7 billion parameters in 16-bit, take about 13 GiB. So the cache overtakes the weights at a context length of roughly 27,000 tokens for a single sequence, and much sooner once batching begins.

This is the reason the decode phase is described as memory-bound rather than compute-bound. Per generated token, the arithmetic is a handful of matrix-vector products whose total work is on the order of $2 \times$ the parameter count, tens of billions of operations for a large model, while the data that must be read is the entire weight set plus the entire cache. On a device with a fixed memory bandwidth, the time per token is set by bytes moved divided by bandwidth, and the cache’s contribution grows linearly with context while the weights’ contribution is constant. Long-context decoding therefore gets slower per token as the conversation continues, even though the per-token compute is unchanged.

Why the cache exists at all, and what it trades

The alternative to caching is recomputation. Without a cache, step $t$ recomputes keys and values for all $t$ positions, at a cost that grows linearly with $t$ and sums to $O(n^2)$ over a full generation. With a cache, the same step reads $t$ positions from memory, at a cost that also grows linearly in bytes but at a much higher effective rate, because a memory read of a contiguous array is far cheaper than the projection that produced it. The projection is a matrix multiply over $d_{\text{model}} \times d_{\text{model}}$ weights per position; the read is $2 d_{\text{model}}$ values per position.

The cache also changes what the attention operation looks like. At decode time there is exactly one query row per sequence, so the score matrix is $1 \times n_{\text{kv}}$, and the operation is a dot product against a long array followed by a weighted sum of another long array. That shape has almost no arithmetic intensity: each cached element is used once per query. A tiled implementation designed for the $n \times n$ training case has nothing to tile over in the query dimension, so the decode kernel is a bandwidth-bound streaming kernel, and the tiling that helps training does not help here. The kernel design consequences are developed in IO-aware attention tiling.

Prefill, by contrast, processes the whole prompt at once and looks like training: many query rows, an $n \times n$ score matrix, high arithmetic intensity. The two phases have different bottlenecks, and a serving system’s throughput is usually set by whichever phase dominates for its workload. A long-prompt, short-generation workload is compute-bound in prefill; a short-prompt, long-generation workload is bandwidth-bound in decode.

Sharing is the main lever

The cache size formula has four free factors, and each has been attacked.

Reducing $s$ by storing the cache in 8-bit or lower halves or quarters the bytes directly. The cost is quantization error in keys and values, which propagates into every subsequent attention read. Keys are more sensitive than values, because an error in a key changes the score and therefore the softmax weights, while an error in a value changes the averaged output more benignly. Schemes that keep keys at higher precision than values exist for that reason.

Reducing the number of cached heads is the second lever, and the largest. If $g$ query heads share one key head and one value head, the cache shrinks by a factor of $h/g$. This is the mechanism of grouped-query attention and its single-group extreme, multi-query attention, and it is examined in grouped-query attention and KV cache sharing. The saving is exact and requires no approximation: the shared key is a real key, computed by a real projection, and every query head uses it. What is lost is the per-head key subspace, which the score computation previously had.

Reducing $L$ is not available without changing the model, but it explains why layer count dominates the constant. A model with twice the layers at the same width doubles the cache, and depth is often the cheaper axis to scale for quality per parameter, a tension that shows up directly in serving cost.

Reducing $n_{\text{kv}}$ by evicting old positions is the fourth lever, and the only one that changes the model’s function rather than its representation. Eviction policies differ in what they keep: a sliding window keeps the most recent $W$ positions, which bounds the cache at $W$ per sequence regardless of generation length; attention sinks keep a small number of initial positions plus a window, on the observation that models place substantial attention mass on the first few tokens; and scoring-based policies keep positions whose accumulated attention weight is high. Each is a real change to what the model can see, and each has a failure mode where the evicted information was load-bearing. The sliding-window case is in sliding-window attention and effective context reach.

def kv_cache_bytes(layers, d_model, seq_len, batch=1, bytes_per=2, kv_heads=None, heads=None):
    # kv_heads < heads models grouped-query sharing: only kv_heads are cached
    factor = 1.0 if kv_heads is None else kv_heads / heads
    per_token = 2 * layers * d_model * factor * bytes_per
    return int(per_token * seq_len * batch)

# 32 layers, width 4096, 16-bit, full multi-head, 4096 positions
print(kv_cache_bytes(32, 4096, 4096) / 2**30, "GiB")   # ~2 GiB

Fragmentation and the allocation strategy

An array sized for the maximum context is wasteful, and an array grown by reallocation copies. Serving systems therefore allocate the cache in fixed-size blocks and map a sequence’s logical positions onto physical blocks, the way a virtual memory system maps pages. A block holds a fixed number of positions, commonly a small power of two, and sequences of different lengths consume different numbers of blocks. The waste per sequence is at most one partially filled block, so the internal fragmentation is bounded by the block size rather than by the maximum context.

The trade-off is indirection. A contiguous cache lets the attention kernel stream one array; a block table requires the kernel to gather each block, which costs an extra lookup per block and breaks the assumption of a single linear run. At a block size of 16 positions and a context of 4,096, there are 256 blocks per sequence per layer, so the indirection is amortized over 16 positions’ worth of data per block and is small relative to the bytes moved. Smaller blocks reduce waste and increase indirection; larger blocks do the reverse. This is the same space-versus-lookup trade that page tables make, with the same resolution: a block size in the low tens of positions.

A second-order effect is that the cache’s memory is allocated per request and freed when the request ends, so a server’s peak concurrency is set by the worst-case total rather than the average. Systems that reserve the maximum context per request therefore serve far fewer concurrent sequences than the memory would allow, and systems that grow the cache on demand pay a reallocation or a block-table update per growth step. Both are visible in the relationship between batch size and latency: batching more sequences amortizes the weight read across them, which raises throughput, until the cache fills the device and further sequences must wait. The point where throughput stops improving is the point at which the device’s memory is fully occupied by caches, and it is reachable at batch sizes far below what the compute would allow.

The arithmetic above also explains why the cache is not usually the right target for a first optimization. Shrinking it changes how many sequences fit; making the kernel faster changes how long each step takes. On a bandwidth-bound decode, those are the same lever viewed from two directions, and the one with the larger factor is the one to pull. Halving the cache doubles the batch that fits and therefore roughly doubles throughput, provided the added batch does not itself saturate bandwidth, which it eventually does, because the weight read is shared but the cache read is not.

References

  1. Fast Transformer Decoding: One Write-Head is All You Need — Noam Shazeer, 2019
  2. Efficiently Scaling Transformer Inference — Pope et al., MLSys 2023
  3. Attention Is All You Need — Vaswani et al., NeurIPS 2017