KV Cache Memory Calculations: Context Length vs. Batch Size

Most OOMs that kill serving servers come from exceeding the KV cache budget. Weights are pinned in memory right after loading; everything that changes after that is the KV cache. The common failure pattern is burning GPU time on experiments instead of calculating in advance why memory explodes the moment you increase batch size or context length. The formula itself is straightforward.

The Formula and Real Numbers

Six variables determine KV cache size.

KV_cache_bytes = 2 × L × H_kv × d_head × S × B × P

• 2 — separate storage for Key and Value
• L — number of Transformer layers
• H_kv — number of KV heads (far fewer than Q heads in GQA models)
• d_head — head dimension
• S — number of tokens processed so far (seq_len)
• B — number of concurrent requests (batch size)
• P — bytes per element (FP16 → 2, FP8 → 1, INT4 → 0.5)

Let's plug in Llama-3 8B: 32 layers, 8 KV heads (GQA), head dimension 128.

# Llama-3 8B, FP16, seq_len=4096, batch=32
2 × 32 × 8 × 128 × 4096 × 32 × 2 = 17,179,869,184 bytes ≈ 16 GB

On an A100 80GB, Llama-3 8B weights take roughly 16 GB in FP16. The KV cache has to fit inside the remaining 64 GB, so batch=32, seq=4096 still leaves plenty of headroom. But the story changes as context grows longer.

Llama-3 70B has 80 layers, 64 Q heads, and only 8 KV heads (GQA). At batch=4, seq=4096:

# Actual GQA (kv_heads=8)
2 × 80 × 8 × 128 × 4096 × 4 × 2 = 5,368,709,120 bytes ≈ 5 GB

# If it had been designed with MHA (kv_heads=64)
2 × 80 × 64 × 128 × 4096 × 4 × 2 = 42,949,672,960 bytes ≈ 40 GB

That's an 8× difference under the same conditions, entirely due to GQA. With straight MHA, loading the 70B weights (~140 GB) onto two A100s (160 GB) and then finding an additional 40 GB for the KV budget would be nearly impossible.

Context Length vs. Batch Size: A Zero-Sum Relationship

In the formula, S and B are multiplied together. Once the GPU memory budget is fixed, the ceiling on S × B is fixed too.

Llama-3 8B, FP16, 16 GB KV budget:

Scaling seq_len 8× from 4k to 32k drops batch size from 32 to 4 — an 8× throughput loss. This is exactly what people mean when they say "long-context serving is expensive": context and batch size cannibalize each other within the same KV budget.

This relationship directly drives deployment decisions. If a product requirement comes in to raise average input length to 16k while maintaining the same throughput target, the options are: add more GPUs, reduce H_kv (GQA → MQA or MLA), or reduce P (quantization). This formula is where that calculation starts.

MHA → GQA → MQA → MLA: Ways to Reduce H_kv

The structural approaches to cutting KV memory either reduce H_kv directly or change the cache representation to an equivalent effect.

GQA was proposed in Ainslie et al. 2023. Q heads are grouped, and each group shares a single KV head. Llama-3 70B's Q=64, KV=8 means 8 groups each sharing 8 Q heads. The savings ratio relative to MHA is H_q / H_kv = 64/8 = 8×.

MQA takes it to the extreme: reduce KV heads to 1. Memory savings are maximal, but quality degradation is noticeable. In practice, GQA has become the sweet spot for large models, and MQA is rarely used.

MLA (Multi-head Latent Attention) was introduced in the DeepSeek-V2 technical report. The approach differs from GQA. Instead of storing K and V per head, it caches a low-dimensional latent vector shared across all heads and recovers each head's K and V via an up-projection at inference time. Mathematically this is equivalent to reducing both H_kv and d_head simultaneously, but because it uses a shared low-dimensional space, quality is better preserved. DeepSeek-V2 used this architecture to cut the KV cache by 93.3% compared to the earlier DeepSeek 67B.

In DeepSeek's ablations, GQA underperformed MHA on quality, while MLA matched MHA. All three approaches reduce the H_kv term in the formula (or an equivalent term), and the savings ratio can be computed exactly from the formula — that's what makes it practically useful.

KV Quantization: Reducing P, but Know Where the Errors Show Up

Reducing P is quantization. The savings are straightforward.

Enabling FP8 KV cache in vLLM is a single flag:

vllm serve <model> --kv-cache-dtype fp8

Internally, KV tensors are stored in fp8_e4m3 format, and the QK·ScoreV matrix multiplications also run in FP8. This requires CUDA 11.8 or later; on H100 and B200, hardware FP8 arithmetic provides direct acceleration.

According to numbers published on the vLLM official blog (2026.04), FP8 increases output throughput by 14.9% over BF16 on Llama-3.1-8B and reduces the inter-token latency slope to 54% of BF16's. On a 1M-token long-context benchmark, FP8 matched baseline AUC, with up to 1–2 point losses reported on reasoning tasks.

Dropping to INT4 yields larger savings but a different error profile. As KIVI (ICML 2024) showed, at 2-bit quantization errors concentrate around "attention sink" tokens — the small number of positions (typically the first token in a sequence) that most other tokens attend to heavily. Quantization error at these positions propagates through the entire attention output. KIVI reports 2.6× memory reduction with roughly 3% accuracy loss on GSM8K.

At short contexts (<2k), sink tokens account for a small fraction, so INT4 is relatively safe. Beyond 32k, accumulated error weighting becomes significant and perplexity degradation becomes pronounced. This is why FP8 has settled in as the practical compromise.

Three OOM Scenarios

Scenario 1 — Fixed batch, growing context: A single request's input during the prefill phase is longer than expected. If other requests in the batch have already populated their KV entries and one request tries to process a 128k document, the KV budget spikes instantly. vLLM handles this with request preemption rather than a hard CUDA OOM, but running torch directly will just crash.

Checking memory state:

import torch
print(torch.cuda.memory_summary(abbreviated=True))

If reserved_bytes.all.peak is significantly larger than allocated_bytes.all.current, the KV cache spiked and was released at some point in the past.

Scenario 2 — Fixed context, growing batch: The system runs fine at batch=16 initially, but as traffic surges, concurrent requests push the batch to 32 or 48. Weights stay constant while KV grows linearly with batch size until it exceeds the remaining memory.

nvidia-smi --query-gpu=memory.used,memory.free --format=csv -l 1

Sampling at 1-second intervals shows memory.used stepping up as the batch accumulates.

Scenario 3 — Prefill spike: During decode, sequences grow one token at a time, so KV grows gradually. Prefill is different. When multiple long prompts enter prefill simultaneously, KV for all of them is allocated at once, causing a sharp instantaneous memory spike. If the system is stable most of the time but produces intermittent OOMs at specific moments, suspect this pattern. In vLLM, the --max-num-batched-tokens parameter limits the number of tokens processed per step and suppresses this spike.

Quick Calculation Sheet

def kv_cache_gb(num_layers, num_kv_heads, head_dim, seq_len, batch_size, dtype_bytes=2):
    total = 2 * num_layers * num_kv_heads * head_dim * seq_len * batch_size * dtype_bytes
    return total / (1024 ** 3)

# Llama-3 8B, FP16, seq_len=4096, batch=32
kv_cache_gb(32, 8, 128, 4096, 32, 2)   # → 16.0 GB

# seq_len=32768, batch=4 — same 16 GB budget
kv_cache_gb(32, 8, 128, 32768, 4, 2)   # → 16.0 GB

# Switch to FP8 (P=1)
kv_cache_gb(32, 8, 128, 32768, 4, 1)   # → 8.0 GB → can double batch to 8

PagedAttention does not change the total governed by this formula. It dynamically allocates KV pages to reduce internal fragmentation, but the total budget is still determined by this formula. That's distinct from GQA or FP8, which directly reduce terms in the formula.

How much more batch size you can actually get when moving from GQA (kv_heads=8) to MLA at equivalent quality is something you have to measure empirically. The natural next question is how DeepSeek-V2's 93.3% reduction figure translates into concrete serving batch configurations in practice.