KV Cache Optimization

The KV cache is the memory a Transformer keeps of every key and value vector it has already computed, so that generating the next token does not require recomputing attention over the entire sequence from scratch. KV cache optimization is the umbrella term for the techniques that keep that memory from becoming the bottleneck: quantizing it to fewer bits, sharing keys and values across attention heads, evicting or compressing entries that stop mattering, and managing the underlying memory so it never fragments. For any deployment serving long context windows or many concurrent users, the KV cache, not the model’s weights, is usually what runs out of GPU memory first.

Why the Cache Exists at All

Transformer self-attention is causal: token N attends to every token from 1 to N. Without a cache, generating token N+1 would mean recomputing the key and value projections for all N prior tokens again, an O(N²) total cost across a full generation. The KV cache trades memory for that recomputation: once a token’s key and value vectors are computed, they are stored, and every later step only computes the new token’s own K and V, then attends against everything cached so far. This is what makes autoregressive decoding an O(N) incremental process instead of a quadratic one, and it is inseparable from why decoding is memory-bandwidth bound: every step reads the entire cache back out of GPU memory, and that read grows with every token generated.

%%{init: {'theme': 'base'}}%%
graph TD
    classDef default fill:#ffffff,stroke:#4338CA,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef data    fill:#EEF0F7,stroke:#0D9488,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef process fill:#F7F8FC,stroke:#6366F1,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef output  fill:#4338CA,stroke:#4338CA,stroke-width:2px,color:#ffffff,rx:8px,ry:8px;

    PROMPT([Prompt tokens 1..N]):::data --> PREFILL[Prefill: one parallel pass, compute K,V for all N tokens]:::process
    PREFILL --> CACHE[(KV cache: K,V for tokens 1..N)]:::data
    CACHE --> DECODE[Decode step: compute K,V for token N+1 only]:::process
    DECODE -->|"attend against full cache"| OUT[Emit token N+1]:::output
    OUT -.->|append K,V to cache| CACHE

The Memory Math

The cache’s size is deterministic and easy to compute, which is exactly why it is such a predictable planning problem:

KV cache bytes = 2 x layers x kv_heads x head_dim x seq_len x batch_size x bytes_per_element

The leading 2 accounts for storing both keys and values. Plug in numbers for a Llama-3.3-70B-class model (80 layers, 8 KV heads under grouped-query attention, head dimension 128) at fp16 (2 bytes per element), a single sequence at 128K context:

2 x 80 x 8 x 128 x 131072 x 1 x 2 bytes ≈ 34.3 GB

That is 34 GB for one sequence’s cache, on top of roughly 140 GB just to hold the model’s own weights in fp16. Double the batch size and the cache doubles again. This is the arithmetic behind the industry’s shift, covered below, toward architectures that shrink kv_heads by design and serving stacks that shrink bytes_per_element at runtime.

Interactive: Cache Size Against Context Length, Batch, and Precision

Slide kv_heads from 8 up to 64 and watch every bar scale linearly, because head count is a pure multiplier on cache size: this is the exact lever grouped-query attention pulls at the architecture level, before any runtime optimization is even applied. Slide precision from fp16 down to int4 and the bars fall by the same 4x, for free at inference time, no retraining required, though real deployments trade some of that gain for a measured accuracy cost. And notice how quickly plain fp16 at full multi-head attention crosses the 80GB line: this is precisely why context windows past roughly 32K tokens were rare in production until GQA, quantization, and paged memory management became standard rather than optional.

The Optimization Toolkit

# Scenario: applying KIVI-style KV cache quantization to fit a longer
# context window on the same GPU, using a library that implements the
# asymmetric per-channel (key) / per-token (value) scheme from the paper.
from transformers import AutoModelForCausalLM
from kivi import KiviConfig  # illustrative: wraps the model's cache

kivi_config = KiviConfig(
    k_bits=2,          # keys quantized per-channel
    v_bits=2,          # values quantized per-token
    group_size=32,     # calibration group size, smaller improves accuracy
    residual_length=32 # keep the most recent tokens in full precision
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B", cache_config=kivi_config
)
# The residual window matters: quantizing every token uniformly loses more
# quality than keeping the most recent, most-attended tokens exact and only
# compressing older, already-settled context.

# Scenario: two requests that share a long system prompt and reference
# document. Prefix caching avoids re-running prefill for the shared part.
shared_prefix = open("product_policy.md").read()  # 6,000 tokens, identical
# reused across requests

request_a = shared_prefix + "\n\nCustomer asks: can I return item X?"
request_b = shared_prefix + "\n\nCustomer asks: is item X covered under warranty?"
# A serving stack with prefix caching (vLLM's automatic prefix caching,
# Anthropic's prompt caching, Gemini's context caching) computes the KV
# cache for shared_prefix once and reuses it for both requests, paying
# full prefill cost only for the short, request-specific suffix.

What’s New (2025-2026)

• MLA became a defining DeepSeek-era idea, not a one-off. Multi-head Latent Attention, introduced in DeepSeek-V2 and carried through DeepSeek-V3 and R1, pushed KV cache compression further than GQA alone by caching a compressed latent representation instead of full keys and values. Research through 2025-2026 on retrofitting MLA into other Transformer-based models, rather than only training it in from scratch, points to cache-efficient attention becoming a design goal teams actively optimize for, not an afterthought bolted on at serving time.
• Low-bit KV cache quantization became a standard serving option, not a research curiosity. FP8 KV cache is now a supported flag in vLLM, SGLang, and TensorRT-LLM, and 2-4 bit schemes following KIVI’s asymmetric per-channel/per-token approach have moved from papers into production-adjacent tooling as teams push context windows past 100K tokens on fixed hardware.
• Prefix and prompt caching became a standard API feature. Anthropic’s prompt caching, Gemini’s context caching, and equivalent features from other providers turned prefix reuse, previously a serving-stack implementation detail, into a billed, documented API capability, because the KV cache savings translate directly into cost savings passed to the customer.
• Disaggregated prefill/decode architectures went from research to production infrastructure. Systems following Mooncake’s KVCache-centric design, separating compute-bound prefill from bandwidth-bound decode onto different hardware pools and transferring the cache between them, moved from academic proposals to how several large-scale inference providers actually run their fleets, particularly for long-context workloads.
• “Inference-aware” pretraining made cache efficiency a day-one architecture decision. Rather than retrofitting GQA or MLA onto an existing design, more 2025-2026 model releases chose their attention variant explicitly to bound KV cache growth, treating the memory cost of serving the model as a pretraining-time constraint alongside accuracy and compute budget.

Practical Guidance

The throughline across every technique here is the same: the KV cache is not incidental bookkeeping, it is usually the actual constraint on how long a context window or how many concurrent users a given GPU can serve. Treating cache efficiency as an architecture decision (GQA, MLA) and a serving decision (quantization, paging, eviction, disaggregation) at the same time, rather than reaching for one lever alone, is what separates a deployment that scales cleanly to long context from one that hits a memory wall well before the model’s advertised context limit.

from vllm import LLM, SamplingParams

# Scenario: serving a 70B chat model where users routinely paste long
# documents, so the KV cache, not the model weights, is what runs out
# of GPU memory first under concurrent load.
llm = LLM(
    model="meta-llama/Llama-3.3-70B-Instruct",
    # PagedAttention (on by default in vLLM): the KV cache is stored
    # in fixed-size, non-contiguous blocks instead of one long
    # per-sequence buffer, eliminating fragmentation from padding
    # sequences to the same length.
    gpu_memory_utilization=0.90,
    block_size=16,
    # FP8 KV cache: halves the cache's memory footprint relative to
    # fp16 with a measured accuracy hit typically under 1%.
    kv_cache_dtype="fp8",
    max_model_len=131072,
)

params = SamplingParams(temperature=0.7, max_tokens=1024)
out = llm.generate("Summarize the attached 40-page contract...", params)
print(out[0].outputs[0].text)