Do the KV Cache Math Before You Pick a Context Length

AI Studio — Tue, 10 Nov 2026 19:17:55 +0000

Your model weights fit comfortably into GPU memory, the inference server boots cleanly, and your initial smoke test returns tokens without a hiccup. You feel good about your setup. Then real traffic arrives with longer prompts and multiple users hitting the endpoint at once. Suddenly the process crashes with an out-of-memory error, or the serving engine quietly throttles your effective batch size down to a crawl. The model weights did not change size after startup. The culprit eating up the rest of your LLM GPU memory is the KV cache.

During autoregressive generation, a decoder-only transformer predicts one token at a time. Thanks to causal masking, the attention keys and values for earlier tokens only depend on the tokens that came before them, which means past states remain fixed as new tokens get appended. Rather than running the entire prompt and generated sequence through every layer on every single decoding step, the runtime stores one Key tensor and one Value tensor per token inside each decoder layer. Reusing those stored tensors turns a painful quadratic recomputation problem into a much faster linear lookup per step, at the cost of linear memory growth across sequence length and active requests.

Deriving the Memory Formula from Tensor Shapes

Before spinning up an instance or hardcoding server flags, you can calculate the exact footprint of those cached states from the model architecture. Inside a single transformer layer, the key cache and the value cache each have the logical tensor shape [batch, kv_heads, seq_len, head_dim]. Multiplying those dimensions across all layers and accounting for both keys and values gives the baseline formula for KV cache memory.

bytes ≈ 2 × layers × kv_heads × head_dim × seq_len × bytes_per_element × batch

Each parameter maps directly to a field in the model configuration or your runtime flags:

2 accounts for storing two separate tensors per layer, one for Keys and one for Values.
layers is the number of transformer decoder blocks in the model (num_hidden_layers in most config files).
kv_heads is the number of attention heads dedicated to keys and values per layer (num_key_value_heads), which may be smaller than the query head count.
head_dim is the vector dimension of a single attention head, usually computed as the model's hidden dimension divided by the number of query heads.
seq_len is the total number of tokens currently held in context for a request, combining prompt tokens and generated output tokens.
bytes_per_element is the storage width of each number in the cache, such as 2 bytes for FP16 or BF16 and 1 byte for 8-bit integer or floating-point formats.
batch is the number of concurrent sequences holding active context in memory at the same time.

Treat the output of this equation as a strict lower bound. Real serving runtimes deal with memory allocator padding, block size rounding, page tables, and temporary workspace buffers for attention kernels. If your back-of-the-envelope calculation leaves zero breathing room above the model weights, the server will fail under load before reaching your theoretical token limit.

A Tiny Python Sizing Helper

You do not need a heavy framework to run these numbers before editing a deployment manifest. This short script uses standard Python to return the raw byte floor and format it in gibibytes (GiB, where 1 GiB is 1024^3 bytes).

def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, batch=1, bytes_per_elem=2):
    return 2 * layers * kv_heads * head_dim * seq_len * bytes_per_elem * batch

def print_kv_gib(layers, kv_heads, head_dim, seq_len, batch=1, bytes_per_elem=2):
    total_bytes = kv_cache_bytes(layers, kv_heads, head_dim, seq_len, batch, bytes_per_elem)
    gib = total_bytes / (1024 ** 3)
    print(f"KV cache floor: {gib:.2f} GiB ({total_bytes:,} bytes)")

if __name__ == "__main__":
    # Illustrative GQA config with 8 concurrent sequences at 32,768 tokens
    print_kv_gib(layers=32, kv_heads=8, head_dim=128, seq_len=32768, batch=8, bytes_per_elem=2)

Worked Examples on an Illustrative Architecture

To see how quickly KV cache size scales in practice, consider a purely hypothetical model configuration. This is an illustrative config for clean math rather than any specific real-world release. Assume a transformer with 32 layers, 32 query heads, a head dimension of 128, and a standard FP16 or BF16 cache taking 2 bytes per element.

If this hypothetical model uses standard multi-head attention (MHA), every query head gets its own dedicated key head and value head, making kv_heads equal to 32. Plugging a single token (seq_len = 1, batch = 1) into the formula yields 2 * 32 * 32 * 128 * 1 * 2 * 1, which equals 524,288 bytes, or exactly 0.5 MiB per token. Half a mebibyte sounds harmless until you multiply it by modern context lengths. A single sequence running at 8,192 tokens consumes 4 GiB of KV cache memory. Push that single sequence out to 32,768 tokens, and the cache alone demands 16 GiB before a second user even connects.

Now take that exact same 32-layer, 128-head-dim architecture and switch the attention mechanism to grouped-query attention (GQA) with 8 KV heads (kv_heads = 8). Instead of 32 pairs of key and value projections per layer, groups of four query heads share a single key and value head. That architectural change drops the per-token cost by a factor of four down to 131,072 bytes, or 128 KiB per token. With 8 KV heads in FP16 or BF16, a single sequence at 32,768 tokens drops from 16 GiB down to 4 GiB. If you serve eight concurrent sequences that all hit the full 32,768-token mark simultaneously, the combined KV cache footprint lands at 32 GiB.

We can push the math one step further by changing the storage data type. If your runtime supports an 8-bit KV cache (bytes_per_element = 1) on that same 8-head GQA configuration, every figure gets cut in half. Those eight concurrent sequences at 32,768 tokens now require 16 GiB instead of 32 GiB.

Attention & Cache Setup	KV Heads	Bytes per Element	Per-Token Footprint	1 Sequence at 32,768 Tokens	8 Sequences at 32,768 Tokens
Multi-Head Attention (FP16/BF16)	32	2	0.5 MiB (524,288 B)	16 GiB	N/A in example
Grouped-Query Attention (FP16/BF16)	8	2	128 KiB (131,072 B)	4 GiB	32 GiB
Grouped-Query Attention (8-bit)	8	1	64 KiB (65,536 B)	2 GiB	16 GiB

Which Knobs You Actually Control

When you are trying to fit an open-weight model onto fixed hardware without crippling concurrency, you have five practical levers to pull.

Context Length Cap

Just because a model checkpoint advertises a massive maximum context window does not mean your serving deployment should allow it. Capping the maximum sequence length in your engine configuration puts a hard ceiling on seq_len. If your application only processes short chat turns or small code snippets, setting a conservative context limit prevents a single rogue request from claiming gigabytes of cache space.

Max Concurrent Sequences

Your maximum batch size directly multiplies the per-sequence cache requirement. Lowering the concurrency cap bounds peak memory usage, though it forces excess incoming requests to wait in a queue. Sizing this parameter requires balancing your worst-case prompt lengths against the memory left over after loading weights and runtime buffers.

KV Cache Precision

Many modern inference engines let you quantize the KV cache down to 8-bit formats, and some experimental or specialized runtimes even offer 4-bit cache options. Halving or quartering bytes_per_element provides immediate relief for long-context workloads. However, lower-precision caches introduce quantization noise into the attention mechanism. Quality tradeoffs are always possible, particularly on tasks requiring precise retrieval across long contexts or subtle mathematical reasoning. Never treat KV cache quantization as lossless; run evaluations on your own workload before committing to it in production.

Choosing Architectures with GQA or MQA

When you are still in the model selection phase, check num_key_value_heads in the model's config.json before downloading the weights. Multi-Query Attention (MQA), introduced in the paper "Fast Transformer Decoding: One Write-Head is All You Need" by Noam Shazeer (arXiv 1911.02150), takes head sharing to the extreme by using a single key and value head per layer across all query heads. Grouped-Query Attention, detailed in "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" by Ainslie et al. (arXiv 2305.13245), sits between full MHA and MQA by assigning one KV head to each subgroup of query heads. Picking a checkpoint trained with GQA or MQA shrinks kv_heads at the architectural level, giving you a smaller per-token multiplier before you touch a single runtime setting.

Paged and Block KV Allocation in Serving Engines

Naive inference implementations pre-allocate a contiguous chunk of GPU memory sized to the maximum possible sequence length for every active request. Because most requests finish well below the maximum length cap, contiguous allocation wastes huge amounts of memory on empty reserved slots and internal fragmentation. Modern serving engines solve this with block-level allocation. PagedAttention, the core mechanism behind vLLM introduced in "Efficient Memory Management for Large Language Model Serving with PagedAttention" by Kwon et al. (arXiv 2309.06180), partitions the KV cache into fixed-size blocks that get allocated on demand as new tokens arrive, much like virtual memory pages in an operating system. This approach cuts most of that fragmentation and allows multiple requests to share physical memory blocks when they start with the exact same prompt prefix. In that paper's experiments, Kwon et al. reported 2 to 4 times higher throughput than the baseline systems they compared against at similar latency, representing their benchmarked outcome rather than a universal guarantee for every deployment.

How the Same Rule Shapes Prompt Caching on Hosted APIs

Even if you decide against self-hosting and call managed LLM APIs instead, the mechanics of the KV cache still dictate how you should structure your prompts. You never manage GPU tensors directly when calling an external endpoint, yet commercial prompt caching features rely on the exact same underlying causal dependency.

In a causal transformer, the Key and Value vectors computed for token position t depend on the exact sequence of tokens from position 1 through t. If a serving cluster caches the KV blocks for a long system prompt, a subsequent request can only reuse those precomputed blocks if its starting tokens match the cached sequence token for token. Injecting a dynamic value near the top of a prompt, such as a current timestamp, a unique request UUID, or a per-user profile block at the beginning of the system instructions, alters the early token stream and invalidates every cached KV state that follows it. To get high cache hit rates whether you run your own paged engine or pay per token on an API, put static system instructions, tool schemas, and unchanging reference documents at the very front of the prompt, and append dynamic per-request input at the very end.

Pre-Deployment Sizing Checklist

Inspect config.json for your target checkpoint and pull out num_hidden_layers, num_attention_heads, num_key_value_heads, and hidden_size to confirm head_dim and whether the architecture uses MHA, GQA, or MQA.
Run the kv_cache_bytes calculation using your target data type, desired context length cap, and target maximum concurrent batch size to find the bare minimum cache floor.
Subtract model weight memory and a safety buffer for activation spikes, CUDA contexts, and block allocator overhead from your available LLM GPU memory before deciding how much space the KV cache pool can actually claim.
Set explicit flags in your serving engine for maximum model length and maximum concurrent sequences so traffic spikes queue gracefully instead of triggering out-of-memory crashes.
Benchmark task accuracy on representative long-context prompts if you enable an 8-bit or 4-bit quantized KV cache to verify that precision reduction does not degrade output quality.
Order your prompt templates with static prefixes first and variable user or session data last so block-sharing engines can deduplicate prefix memory across concurrent requests.

DEV Community: AI Studio