What is a KV cache in transformer inference?

The KV cache stores the key and value projections for every token generated so far, letting the model attend to its own history without recomputing attention over the full context at each step. This turns the quadratic cost of self-attention into linear incremental work during autoregressive generation. Without it, every new token would require a full forward pass over all previous tokens, making long-context inference impractical.

How the KV cache works

In a standard transformer decoder, each layer computes attention by projecting the input into queries, keys, and values. During training or prompt processing, the model sees the entire sequence at once and computes attention in parallel. During autoregressive generation, tokens arrive one at a time. The query for the new token must attend to keys and values from all prior tokens — including those generated in previous steps.

The KV cache solves this by materializing and storing the key and value tensors for every layer after each forward pass. When the next token arrives, the model:

Computes the query projection for the new token only
Retrieves the cached keys and values for all previous positions
Computes attention between the new query and the full cached key/value history
Appends the new token’s keys and values to the cache

This avoids recomputing key and value projections for the entire history at every step. The memory cost grows linearly with sequence length: two tensors per layer per token, each of shape [batch_size, num_heads, head_dim].

# Simplified KV cache update per layer
# cache_k: [batch, num_kv_heads, seq_len, head_dim]
# cache_v: [batch, num_kv_heads, seq_len, head_dim]

def update_kv_cache(cache_k, cache_v, new_k, new_v):
    # new_k, new_v: [batch, num_kv_heads, 1, head_dim]
    cache_k = torch.cat([cache_k, new_k], dim=2)
    cache_v = torch.cat([cache_v, new_v], dim=2)
    return cache_k, cache_v

Why it matters for inference latency

The KV cache shifts the bottleneck from compute to memory bandwidth. During prefill (prompt processing), the model is compute-bound: matrix multiplications dominate. During decode (token generation), the model becomes memory-bound — each step reads the entire cache from VRAM to compute attention for a single new token.

This has concrete implications:

Prefill scales with prompt length squared (attention over all pairs)
Decode scales with prompt length linearly (reading the cache)
Cache size determines maximum context — a 70B model with 32K context needs ~1.5 GB of KV cache per request at FP16

The memory footprint formula for a single request:

cache_bytes = 2 * num_layers * num_kv_heads * head_dim * seq_len * bytes_per_element

For LLaMA-70B (80 layers, 8 KV heads, 128 head_dim) at 32K context in FP16:

2 * 80 * 8 * 128 * 32768 * 2 ≈ 1.07 GB per request

At batch size 32, that’s 34 GB just for the cache — before model weights or activations. This is why KV cache management dominates serving infrastructure design.

Concrete example: generation loop with cache

@torch.no_grad()
def generate(model, input_ids, max_new_tokens, kv_cache=None):
    """
    input_ids: [batch, seq_len]
    kv_cache: list of (k_cache, v_cache) per layer, or None for fresh start
    """
    batch_size, seq_len = input_ids.shape
    past_len = 0 if kv_cache is None else kv_cache[0][0].shape[2]
    
    # Prefill: process full prompt, populate initial cache
    if kv_cache is None:
        logits, kv_cache = model(input_ids, use_cache=True)
        next_token = logits[:, -1:].argmax(dim=-1)
        generated = [next_token]
    else:
        # Resume from existing cache (e.g., continued conversation)
        logits, kv_cache = model(input_ids[:, past_len:], kv_cache=kv_cache, use_cache=True)
        next_token = logits[:, -1:].argmax(dim=-1)
        generated = [next_token]
    
    # Decode loop: one token at a time
    for _ in range(max_new_tokens - 1):
        logits, kv_cache = model(next_token, kv_cache=kv_cache, use_cache=True)
        next_token = logits[:, -1:].argmax(dim=-1)
        generated.append(next_token)
    
    return torch.cat(generated, dim=1), kv_cache

Notice the model signature: it accepts an optional kv_cache and returns an updated one. The first call processes the full prompt (prefill). Subsequent calls pass only the single new token (decode). The cache grows by one position each iteration.

Multi-query and grouped-query attention change the math

Standard multi-head attention (MHA) uses separate key/value heads for each query head. Multi-query attention (MQA) shares a single key/value head across all query heads. Grouped-query attention (GQA) sits between them — multiple query heads share each key/value head.

This directly reduces KV cache size:

Attention type	KV heads	Cache reduction vs MHA
MHA	num_heads	1x (baseline)
GQA (8:1)	num_heads / 8	8x smaller
MQA	1	num_heads x smaller

LLaMA-2-70B uses GQA with 8 query heads per KV head. LLaMA-3-70B uses 8:1 GQA. This is not an implementation detail — it’s a fundamental architectural choice that determines whether your 32K context fits in a single GPU.

# GQA cache shape: [batch, num_kv_heads, seq_len, head_dim]
# For LLaMA-70B: num_kv_heads = 8 (vs 64 query heads)
# MHA equivalent would need 64 KV heads → 8x more cache memory

Cache eviction and sliding window attention

When context exceeds the model’s trained window (or your GPU memory), you must evict. Naive eviction — dropping the oldest tokens — breaks attention for tasks needing early context (system prompts, few-shot examples, document starts).

Common strategies:

Sliding window attention (Mistral, Longformer): each token attends only to the previous W tokens. The cache becomes a fixed-size ring buffer. No eviction decisions needed — old tokens simply fall off. Trade-off: the model cannot attend to arbitrarily distant tokens.

Attention sinks (StreamingLLM): keep the first few tokens (the “sink”) plus a sliding window. The initial tokens act as global anchors. Works surprisingly well for streaming chat.

Importance-based eviction: score tokens by attention weight magnitude or gradient-based saliency, evict lowest-scoring. Expensive to compute, rarely used in production.

Chunked prefill with cache reuse: split long prompts into chunks, compute KV cache for each chunk, then discard intermediate activations. The final cache represents the full prompt. n4n.ai uses this pattern to handle prompts that exceed single-batch memory limits without OOM.

Quantization: FP8, INT8, and INT4 KV cache

Model weights get quantized to INT4/INT8 routinely. KV cache quantization is trickier — it’s activation data, not static weights, and errors accumulate across decoding steps.

FP8 (E4M3/E5M2): Near-lossless for most models. 2x memory savings vs FP16. Requires Hopper (H100) or Blackwell for native support; emulated on Ampere with modest overhead.

INT8 asymmetric: 2x savings. Needs per-token or per-channel scaling factors stored alongside the cache. Adds ~0.5% memory overhead for scales. Quality degradation appears around 16K+ context for some models.

INT4: 4x savings. Generally too lossy for KV cache without calibration or mixed-precision schemes (keep recent tokens in FP8/INT8, older in INT4). Research area — not production-ready for general use.

# Pseudocode for INT8 KV cache with per-token scales
def quantize_kv_cache(k_cache, v_cache):
    # k_cache: [batch, heads, seq_len, head_dim] in FP16
    k_scale = k_cache.abs().amax(dim=-1, keepdim=True) / 127  # [batch, heads, seq_len, 1]
    v_scale = v_cache.abs().amax(dim=-1, keepdim=True) / 127
    
    k_int8 = (k_cache / k_scale).round().clamp(-128, 127).to(torch.int8)
    v_int8 = (v_cache / v_scale).round().clamp(-128, 127).to(torch.int8)
    
    return k_int8, v_int8, k_scale, v_scale  # scales stored in FP16

During attention, dequantize on the fly: k_fp16 = k_int8 * k_scale. Modern kernels fuse this with the attention computation.

Common misconceptions

“KV cache is just an optimization”

It’s not optional. Autoregressive generation without KV cache is O(n²) per token — each step recomputes attention over all prior tokens. A 100-token generation would cost ~5000 attention computations instead of 100. No production system runs without it.

“Cache size equals context length”

Context length is the maximum sequence length the model was trained to handle (or can extrapolate to). Cache size is the actual memory allocated for a specific request. You can have a 128K-context model but only allocate 4K cache for a short request. The cache grows dynamically during generation until it hits the context limit or memory budget.

“All layers have the same cache shape”

Not with GQA/MQA. Early layers may use different head configurations than later layers (rare but possible). More commonly, the number of KV heads differs from query heads. The cache shape per layer is [batch, num_kv_heads, seq_len, head_dim] — num_kv_heads varies by architecture.

“FP16 cache is fine for everything”

At 32K+ context, FP16 cache dominates VRAM. A single H100 (80 GB) running LLaMA-3-70B can serve roughly:

32K context, FP16 cache: ~4 concurrent requests
32K context, FP8 cache: ~8 concurrent requests
128K context, FP8 cache: ~2 concurrent requests

Quantization isn’t a nice-to-have — it’s the difference between viable and unviable throughput.

“Prefix caching and KV cache are the same thing”

Prefix caching reuses KV cache across requests that share a common prefix (system prompt, few-shot examples, document). The KV cache is the data structure; prefix caching is a serving-layer optimization that deduplicates it. They compose: prefix caching stores KV cache entries in a shared pool keyed by prefix hash.

# Prefix cache lookup (simplified)
prefix_hash = hash(system_prompt + few_shot_examples)
if prefix_hash in prefix_cache:
    kv_cache = prefix_cache[prefix_hash].clone()  # copy-on-write for safety
    input_ids = user_message_tokens  # only process the unique suffix
else:
    kv_cache = None
    input_ids = full_prompt_tokens

What this means for system design

If you’re building an inference stack, the KV cache dictates:

Memory planning: Cache is the largest per-request allocation. Size it for your max context and batch size, then add 20% headroom for fragmentation.
Batching strategy: Continuous batching (interleaving prefill and decode in the same batch) requires cache allocation/deallocation per request. Fragmentation becomes real. Use a memory pool with slab allocation.
Scheduling: Long-context requests hold cache for their entire lifetime. They block memory that could serve many short requests. Consider separate queues or priority bands.
Observability: Track cache hit rate (for prefix caching), cache memory usage per request, and decode throughput vs. cache size. The decode latency curve should be flat — if it grows with context, something’s wrong (likely cache not being reused, or memory pressure causing swapping).
Model selection: GQA/MQA architectures are not just “more efficient” — they enable longer context at the same hardware cost. For serving, this is often the deciding factor between model families.

The KV cache is the central data structure of transformer inference. Every optimization — quantization, prefix caching, sliding window, paged attention — operates on it. Understanding its shape, growth, and access patterns is table stakes for shipping LLMs in production.