KV Cache Deep Dive: Correctness, Memory and Serving Budgets

Why a cache exists, and when reuse is correct

A causal decoder's past hidden states cannot depend on later tokens. With fixed model/adapters, positions, attention semantics and evaluation behavior, old keys and values at each layer can therefore be reused. Compute Q, K and V for each new token, append new K/V and let its query attend to the allowed prefix. Saying “only compute Q” omits the new token's K/V. A multi-layer model needs a distinct cache per attention layer.

The cache is memoization mathematically and a capacity constraint operationally. Recomputing the full prefix can produce equivalent logits but costs extra work; evicting old tokens without recomputing them changes the information available and is not the same operation. Changing a prefix token, model weights, adapter or relevant position mapping invalidates reuse. Training dropout and bidirectional attention do not inherit this simple equivalence argument.

For full-prefix naive attention, a step at length t forms a t×t score matrix; one-token cached attention forms a 1×t row. Summing over a growing sequence gives cubic versus quadratic attention-score work in total. Weight projections/MLPs have their own costs; these big-O terms do not predict a constant speedup or a flat cached-step timing curve. Cached attention still reads a growing prefix. Preallocation removes repeated allocation/copy overhead, not this attention work.

The restaurant-ticket/notebook analogy is therefore “write reusable facts once, consult them again”, not “never reread the history”. Earlier anecdotes about one particular A10/A100 deployment, the invention of caching by GPT-2, and fixed dollar-per-hour thresholds were unsupported. Primary cache documentation explains reuse; the PagedAttention paper addresses allocation and sharing, not every historical origin claim.

Derive the exact byte formula

At one ordinary attention layer, both K and V have logical shape (B,Hkv,T,dh). Count elements first, then bytes:

The factor two is K and V; L is the number of relevant attention layers; Hkv is KV heads, not query heads; s is bytes per stored element. For unequal request lengths replace B*T with the sum of lengths. For heterogeneous layers sum their individual contributions. Quantized layouts add scale/zero-point metadata and alignment; using s=0.5 for INT4 describes payload only, not an arbitrary implementation's allocated footprint.

For a small hand-counted example, L=2,Hkv=2,dh=4,T=3,B=1,s=2 gives 192 bytes, not 96. Each layer has 24 K and 24 V elements, so two layers have 96 elements at two bytes each. This hand count is a useful negative control against a memorized formula.

Use GB = 10^9 bytes, GiB = 2^30 bytes, KiB = 1024 and MiB = 2^20. Seven billion two-byte parameters are 14 GB, about 13.039 GiB, not 14 GiB. Parameters, reserved allocator memory and actual free HBM are separate observations.

MHA, MQA, GQA and configuration reading

MHA uses one KV head per query head. MQA shares one K/V head across query heads. GQA groups query heads around several KV heads; with 32 queries and eight KV heads, ordinary cache payload is four times smaller than MHA with the same layer count/head width/precision. Quality is an empirical training/uptraining result, not guaranteed by assuming redundant heads.

Read architecture/config fields, including explicit head_dim where present. hidden_size // num_attention_heads is not universally the attention head width. A missing num_key_value_heads may have a documented architecture default; do not silently guess for unfamiliar architectures. Local/sliding-window layers, MLA, recurrent layers and cross-attention require their own accounting.

For ordinary full-context BF16 caches the following are configuration-derived payload calculations, not server measurements:

A 4096-token Llama-2-7B request is 2 GiB of K/V; the Llama-3-8B shape is 512 MiB. FP8 payload halves the listed BF16 bytes before scales and implementation details. For GPT-2 XL,307200bytes is300KiB; bytes and KiB are not interchangeable.

For Gemma-2-27B, the primary report specifies 46 layers, 32 query heads, 16 KV heads, head width 128, alternating 4096-token local and 8192-token global spans. At contexts up to4096, the full participating cache payload grows at376832bytes/token = 368KiB, not184KiB. Beyond4096, local layers need not keep growing: for23local and23global layers at8192 tokens, the bounded payload is 2*23*16*128*2*(4096+8192) =2315255808bytes =2.15625GiB, before implementation overhead. The explicit head width is128, not4608/32=144.

The original Jamba report uses an attention-to-Mamba ratio1:7, not25%attention layers. Its published comparison gives4GB of KV at256Kcontext versus32GB for a stated Transformer comparator, with recurrent state a separate term. Do not transfer those numbers to Jamba1.5: read the specific variant configuration and account for both attention cache and recurrent state before budgeting a deployment.

Gemma configurations especially illustrate why this is not a universal scoreboard: head width can differ from hidden_size/heads, and alternating local/global layers retain different lengths. MoE routing changes weight/compute accounting but does not by itself change the number of attention-cache tensors. Hybrid models such as Jamba retain attention caches and recurrent state; “no KV cache anywhere” applies only to an attention-free design, not every hybrid.

Roofline reasoning: prefill and decode are tendencies

For a dense weight matrix with N participating weights and B one-token requests, a simplified step performs about 2NB FLOPs and reads Ns_w weight bytes. Weight-only arithmetic intensity is 2B/s_w: B FLOPs/byte for two-byte weights, not universally two. It also reads K/V and performs attention work, writes state, launches kernels and may communicate across GPUs.

A first-order roofline lower bound is

Peak specifications are optimistic ceilings, not observed effective rates. For an explicitly hypothetical device with 312 TFLOP/s and 2 TB/s, crossover intensity is 156 FLOPs/byte. Ignoring cache and all overhead, BF16 weight reuse reaches that at B≈156, not a guarantee that a real batch of 156 fits or becomes compute-bound. Prefill processes many prompt rows together and often has higher weight reuse; short prompts, long-context attention, quantized kernels and communication can change the bottleneck. Decode at high concurrency is not necessarily bandwidth-bound.

For an 8B illustrative dense model with two-byte weights and the 128-KiB/token cache shape, a 4096-token request reads roughly 16 GB weights plus 0.537 GB cache per step under this simplified model. A 2-TB/s bandwidth ceiling gives at least about 8.27 ms just to move those bytes, before other costs. At B=32 the weight traffic is amortized but cache traffic is about 17.18 GB and attention work also grows. This is why “16× batch gives exactly16× throughput for free” is false. Measure aggregate throughput, TTFT and per-stream ITL together.

Protected exact-model cache experiment

Save the transformer-block lesson's exact TinyDecoder, norm and softmax implementation, including its NumPy import, as tiny_transformer.py. Copy the definitions without the training loop; the cache test creates a fresh deterministic model and does not require a trained checkpoint. These protected NumPy payloads remain byte-identical. The one-block learned-position model is deliberately small: it compares every logit vector against full-prefix computation and rejects context overflow. It uses repeated concatenation for clarity, not an efficient allocator.

import numpy as np
from tiny_transformer import TinyDecoder, norm, softmax
 
model = TinyDecoder()
p = model.p
cache_k, cache_v = [], []
 
def decode_step(token):
    position = len(cache_k)
    if position >= model.context:
        raise ValueError('context capacity reached; no silent position reset')
    x = (p['E'][token] + p['P'][position])[None, None, :]
    n, _ = norm(x, p['g1'], p['b1'])
    q, k, v = [n @ p[name] for name in ('Q','K','V')]
    cache_k.append(k); cache_v.append(v)
    keys = np.concatenate(cache_k, axis=1)
    values = np.concatenate(cache_v, axis=1)
    weights = softmax(q @ keys.transpose(0,2,1) / np.sqrt(model.width))
    h = x + (weights @ values) @ p['O']
    m, _ = norm(h, p['g2'], p['b2'])
    z = h + np.maximum(m @ p['F'], 0) @ p['G']
    return (z @ p['U'])[0,0]
 
ids = [0,1,2,0,1,2,0,1]
for i, token in enumerate(ids):
    cached = decode_step(token)
    full = model.forward(np.array([ids[:i+1]]))[0][0,-1]
    np.testing.assert_allclose(cached, full, atol=1e-10, rtol=1e-10)
try:
    decode_step(2)
except ValueError:
    pass
else:
    raise AssertionError('context overflow must be explicit')
print('eight cached/full-prefix comparisons passed')

Changing the learned position to zero at each step is a useful negative control. Repeated tokens must still use different position indices. Matching sampled argmax tokens is weaker than matching logits: distributions can disagree while selecting the same token. For stochastic sampling, compare distributions or use matched random draws.

Preallocated grouped-query cache with chunked prefill

This second complete CPU program isolates attention with fixed projection matrices; it is not a full language model and does not claim language quality or GPU speed. It handles unequal query/KV head counts, a new chunk's absolute positions, capacity checks, exact batch shape and mutation only after validation. repeat_interleave materializes expanded tensors here—it is not a free view or the intended optimized GQA kernel.

import math
import torch
torch.manual_seed(66)
torch.set_num_threads(1)
B,H,HK,DH,D,T=2,4,2,3,12,7
x=torch.randn(B,T,D,dtype=torch.float64)
Wq=torch.randn(D,H*DH,dtype=torch.float64)/math.sqrt(D)
Wk=torch.randn(D,HK*DH,dtype=torch.float64)/math.sqrt(D)
Wv=torch.randn(D,HK*DH,dtype=torch.float64)/math.sqrt(D)
Wo=torch.randn(H*DH,D,dtype=torch.float64)/math.sqrt(D)
def project(z,w,h):
    return (z@w).reshape(B,z.shape[1],h,DH).transpose(1,2)
def attend(q,k,v,start):
    k=k.repeat_interleave(H//HK,1); v=v.repeat_interleave(H//HK,1)
    qi=torch.arange(start,start+q.shape[2])[:,None]
    ki=torch.arange(k.shape[2])[None,:]
    score=q@k.transpose(-1,-2)/math.sqrt(DH)
    prob=score.masked_fill(ki>qi,-torch.inf).softmax(-1)
    return (prob@v).transpose(1,2).reshape(B,q.shape[2],H*DH)@Wo
class Cache:
    def __init__(self,capacity):
        self.k=torch.empty(B,HK,capacity,DH,dtype=torch.float64)
        self.v=torch.empty_like(self.k); self.used=0
    def step(self,z):
        if z.ndim!=3 or z.shape[0]!=B or z.shape[2]!=D or z.shape[1]<1:
            raise ValueError('shape')
        if z.dtype!=self.k.dtype or z.device!=self.k.device: raise ValueError('dtype/device')
        end=self.used+z.shape[1]
        if end>self.k.shape[2]: raise ValueError('capacity')
        start=self.used
        q,k,v=project(z,Wq,H),project(z,Wk,HK),project(z,Wv,HK)
        self.k[:,:,start:end]=k; self.v[:,:,start:end]=v
        self.used=end
        return attend(q,self.k[:,:,:end],self.v[:,:,:end],start)
full=attend(project(x,Wq,H),project(x,Wk,HK),project(x,Wv,HK),0)
for widths in ([1]*T,[3,2,2],[7]):
    c=Cache(T); pieces=[]; start=0
    for width in widths:
        pieces.append(c.step(x[:,start:start+width])); start+=width
    assert torch.allclose(torch.cat(pieces,1),full,atol=1e-12,rtol=1e-12)
    before=c.used
    try: c.step(x[:,:1])
    except ValueError: pass
    else: raise AssertionError('overflow accepted')
    assert c.used==before
# Wrong rectangular mask: new query treated as absolute position zero.
wrong=attend(project(x[:,-1:],Wq,H),project(x,Wk,HK),project(x,Wv,HK),0)
assert not torch.allclose(wrong,full[:,-1:])
assert 2*2*1*2*3*4*2==192
print('GQA full/token/chunk parity; wrong-mask and capacity controls passed')

Read the rectangular mask by hand. If three tokens are already cached and the next chunk has two tokens, its absolute query positions are 3 and 4; key positions are 0 through 4. Query 3 may see keys 0–3 but must mask key 4. Query 4 may see all five. Thus the allowed rows are [1,1,1,1,0] and [1,1,1,1,1], not the top-left two rows of a 5×5 triangle. In the program, start=self.used supplies that offset before used is advanced. For a single new token, all existing keys plus itself are valid; for a chunk, later members of the chunk are still future tokens. This explains why a non-square causal flag with the wrong alignment can pass an uncached one-token smoke test and fail chunked prefill.

The fixture fixes H=4,HK=2, so divisibility is built into the example rather than checked dynamically. A reusable constructor must reject incompatible head counts. It is CPU, dense, equal-length, finite-input attention: no ragged request mask, sliding-window policy, RoPE implementation or general nonfinite-input guarantee is supplied. These are scope limits, not reasons to alter the protected exact programs.

For RoPE, rotate a new Q/K at its true absolute/cache position before storing K; rotate old K exactly once. The scalar dot-product identity depends on relative offsets. A production cache also needs per-request lengths, attention windows, ownership and safe slot reuse. This rectangular-mask test catches one part of that contract; the protected learned-position model checks position use independently.

From payload bytes to admission capacity

Available KV memory is measured usable memory minus weights, peak activations/workspaces, graph buffers, runtime reserves and safety margin. There is no universal two-GiB overhead. For block size b, reserve b*ceil(T_i/b) token slots per live request, plus retained reusable prefixes and lookahead as appropriate. Generation grows the cache, so admission must reserve capacity, enforce a bound or support preemption. Average context does not protect a tail request.

Paging and growth, worked without a server. With four-token blocks, request A at length five can use physical blocks [7,2]: logical positions 0–3 live in block 7 and position 4 in block 2. Request B at length three can use [9]. The block table—not physical adjacency—preserves token order for attention. Twelve slots are reserved for eight live tokens, so four slots are tail waste. If A grows to nine tokens while B stays at three, A needs a third block: total reservation rises to sixteen slots for twelve live tokens. A twelve-slot pool that admitted the first state cannot promise that growth without a free block, preemption or a stricter output bound.

Identical immutable prefix blocks may be shared only under the correct model/adapter/tenant identity. If two continuations share a partially filled block, writing into it in place corrupts the sibling: a sharing allocator must copy that block before either continuation mutates it. Reference counts prevent reclaiming a block still used by another request. This explains copy-on-write conceptually; the cited vLLM V1 automatic prefix cache caches full blocks, so do not infer partial-block reuse from this example. Paging itself preserves attention information; eviction deliberately discards information. Neither a block table nor a smaller storage dtype certifies a latency or quality gain.

Worked scenario, intentionally using binary units consistently: assume a 24-GiB usable pool, exactly 7B BF16 parameters and a measured 2-GiB non-KV reserve. Weight payload is 13.0385 GiB and remaining KV is about 8.9615 GiB. At 2 GiB per 4096-token request, the payload upper bound is four requests. An actual “24 GB” device may expose a different byte count; read the device and allocator instead of treating this scenario as a certified A10 capacity.

For 405B two-byte weights: 810 GB is about 754.37 GiB; ideal TP=8 divides this into about 94.30 GiB per rank, already beyond a hypothetical 80-GiB rank before cache/workspace. FP8 weight payload would be about 47.15 GiB/rank before metadata and nonquantized tensors. The 504-KiB/token BF16 cache is ideally63 KiB/rank when eight KV heads shard evenly over eight ranks; FP8 payload gives31.5 KiB/rank. Some architectures/TP degrees replicate heads or state, so dividing everything by TP is not universally correct. The arithmetic establishes a necessary fit check, not 32-user throughput or a complete hardware prescription.

For the Qwen2.5-32B exercise, use its configuration (L=64,Hkv=8,dh=128) and a nominal32B two-byte weight estimate. Cache/token is262144 bytes;20 requests at8192 tokens consume40 GiB globally. Weight payload is about59.60 GiB, total99.60 GiB before overhead. Under ideal two-way sharding and two80-GiB usable devices, the arithmetic leaves about60.40 GiB globally for everything else. That is a conditional capacity margin, not proof that the actual checkpoint count, kernel layouts, output growth or latency SLO fits.

B, L, T, Hkv, dh, item_bytes = 1, 32, 4096, 8, 128, 2
size = 2*B*L*T*Hkv*dh*item_bytes
assert size == 512*1024**2
arrival = 0.
token_times = [.12, .16, .21, .25]  # illustrative seconds, not a benchmark
first_token = token_times[0]-arrival
mean_interval = (token_times[-1]-token_times[0])/(len(token_times)-1)
assert abs(first_token-.12) < 1e-12
assert abs(mean_interval-.0433333333333333) < 1e-12
print('KV MiB', size/1024**2, 'TTFT', first_token, 'mean ITL', mean_interval)

The protected timestamp example is illustrative, not a timing benchmark: TTFT is0.12 seconds and mean ITL is about0.04333 seconds. Distinguish first network byte, first generated token, token gaps and end-to-end time; proxy buffering can change client observations without changing GPU math.

MLA, quantization, eviction and recurrent alternatives

MLA in DeepSeek-V3 stores a compressed KV latent and a decoupled rotary-key part. For the report's512+64 elements across61 layers at two bytes, payload is70272 bytes/token =68.625 KiB. Comparing its per-layer576 elements to2048 for the eight-head GQA example gives about3.556×, but whole-model ratios also depend on layer counts. MLA can absorb projection matrices for inference; “one extra decompression matmul always makes prefill slower and decode faster” is not a universal implementation model.

FP8/INT8/INT4 caches trade payload for scales, quantization/dequantization, supported kernels and possible accuracy changes. Exponent/mantissa format and calibrated scales matter. Half storage does not guarantee double throughput or a fixed0.05-point MMLU loss. FP8 cache support is not categorically H100-only; check the serving release's device/backend matrix. FP16 overflow is a range issue, not something that inevitably appears at32K because cached vectors accumulate with context. Inspect finite activations, dtype and kernels rather than clipping blindly.

StreamingLLM/attention sinks retain selected initial tokens plus a recent window. They support long streaming operation with bounded attention state under their method; they do not provide exact recall of discarded middle tokens or arbitrary million-token reasoning. H2O chooses heavy-hitter/recent entries based on attention importance. Its empirical results are configuration-specific, not proof that80% of every cache is useless. Eviction changes attention and needs retrieval/long-range tests. Do not invent an attn_sinks=True vLLM flag.

Sparse attention such as Longformer/BigBird changes the attention pattern. Mamba/RWKV use recurrent/state-space mechanisms; Jamba/Griffin are hybrid designs with their own state and attention windows. These are distinct architecture families, not all descendants of the same sparse-attention kernel. Fixed state can bound recurrent-state storage at inference, but training cost, model size and hybrid attention still matter. No universal two-point quality penalty follows.

Prefix caching reuses work across identical token prefixes under the same model, adapter, positions and relevant conditioning. SGLang's RadixAttention organizes prefix reuse in a radix tree; vLLM uses release-specific block hashing. Reuse can lower prefill work and memory duplication, not make uncached suffix generation or networking free. Include model/tokenizer/template/adapter identity and tenant isolation in cache keys. Full-block reuse and partial-block copy-on-write are separate implementation choices. Hash collisions, private-prefix side channels and stale adapters are correctness/security concerns.

Operational recipes and failure diagnosis

For a later authorized GPU experiment, start with an already available model and pinned serving environment. Record configuration and effective byte counts; run full-prefix and cached logits on short deterministic inputs first; then measure prefill and decode separately with device synchronization or device events, warmup, repeated trials and a bounded context/batch sweep. Include input/output lengths, dtype, attention backend, allocation strategy and communication. Plot step latency versus context; do not expect cached attention to be perfectly flat. Change head width only with compatible projection shapes—otherwise the experiment changes the model or is invalid.

OOM after load can come from growing cache, workspaces, graph captures, another process or fragmentation. Check allocation/preemption logs and actual reservations rather than assuming fragmentation. A17-token request with16-token blocks reserves32 slots and wastes15, but smaller blocks increase metadata/kernel costs and are not supported by every backend. Increasing max_num_batched_tokens is not a general instruction to admit fewer requests. SDPA kernel dispatch depends on dtype/device/layout/mask/version; a cache slice alone neither guarantees nor forbids FlashAttention. Use backend diagnostics, not a promised fivefold speedup.

Bound prompt bytes/tokens, output tokens, active requests, waiting queue and deadlines. Reclaim ownership on completion/cancellation, prevent cross-request attention, and fail on capacity rather than resetting positions silently. No GPU timing, full-model serving or endpoint is executed here.

Worked recall and decision rules

1. Cache formula? Count K and V at each relevant layer with KV heads, retained tokens and storage bytes; then add metadata/other memory separately.
2. Why decode often bandwidth-limited? Few new rows reuse large weights poorly, and cache reads grow with context. For BF16 weight-only batch1 the simplified intensity is1 FLOP/byte, not an invariant2.
3. What does paging solve? Noncontiguous fixed-size allocation reduces reservation/external-fragmentation waste, retaining a last-block tail and finite capacity.
4. Why GQA quality can hold? It is a trained architectural tradeoff demonstrated on specified models, not guaranteed head redundancy.32/8 gives4× payload savings under equal other dimensions.
5. How many4K requests fit the24-GiB scenario? Four by the stated payload budget; recompute with actual free bytes, reserve and output bound before admission.
6. When choose each intervention? Fix avoidable allocation/reuse waste first when applicable; quantize only with verified kernels and quality tests; shorten/evict only when loss of old information is acceptable; choose GQA/MLA/hybrids with suitable trained weights; add/shard hardware only after accounting for communication and per-rank memory. These are tradeoffs, not an unconditional ranking.
7. What does the exact test prove? Protected decoder logit parity and overflow; separate GQA attention parity for token/chunk paths; it does not reproduce vLLM or establish GPU latency.

Spaced practice: tomorrow draw the two-row mask for a cached prefix of three and a two-token chunk, then recount the 192-byte toy by enumerating K and V. In a week, recompute the paging example with A growing to thirteen tokens: A needs four blocks and B one, reserving twenty slots. Check an actual model configuration before reusing any formula, and report throughput, TTFT and ITL as separate observations.

Primary sources

Transformers4.57.1 cache explanation, GQA, Llama3 report, DeepSeek-V3, PagedAttention, StreamingLLM, H2O, Mamba, and SGLang. Configuration identities and calculations are distinguished from measured deployment results in the accompanying evidence ledger.

Course navigation

Course map · Lesson 16 of 18

Previous: Adaptation with LoRA: low-rank updates and masked supervision · Next: RAG as a testable evidence pipeline, not a truth switch

Optional reading

Use these links for additional background. Version-specific behavior is documented in the primary references above.

• Optional resource 1
• Optional resource 2
• Optional resource 3
• Optional resource 4