- KV cache size = 2 * num_layers * num_heads * head_dim * sequence_length * bytes_per_element. For a 70B model at 128K context, that's over 40 GB of VRAM just for cache.
- FlashAttention eliminates the O(n^2) memory of full attention but does not reduce KV cache memory; it only reduces HBM reads/writes.
- Quantizing the KV cache (FP8, INT4, or 2-bit via KIVI) can cut memory by 2-8x with minimal accuracy loss, but requires kernel support.
- PagedAttention (vLLM) eliminates internal fragmentation but does not shrink the total cache; it only improves utilization.
- Running a 405B model with 256K context requires >300 GB of VRAM for KV cache alone, forcing multi-GPU tensor parallelism and aggressive quantization.
Breaking Down the Formula: VRAM = Weights + KV Cache + Activations
Total VRAM for inference is the sum of three components: model weights, KV cache, and intermediate activations (plus a small overhead for buffers and CUDA contexts). For a typical batch size of 1, activations are negligible compared to weights and cache.
Model weights: For a dense model with P parameters, stored in B bytes per parameter, weight memory = P * B. For FP16, B=2. For INT4, B=0.5. A 70B model in FP16 requires 140 GB; in INT4, 35 GB.
KV cache: As derived above, for a model with L layers, H heads, D head dimension, and sequence length S, KV cache per token = 2 * L * H * D * B_cache bytes. With B_cache typically 2 (FP16) or 1 (FP8) or 0.5 (INT4).
Example: LLaMA-2 7B (32 layers, 32 heads, dim 128) at 32K context, FP16 cache: 2 * 32 * 32 * 128 * 32768 * 2 = 2 * 32 * 32 * 128 * 65536 = 2 * 32 * 32 * 8388608 = 2 * 32 * 268435456 = 2 * 8589934592 = 17.18 GB. The model itself is 14 GB. Total ~31 GB, fitting on a single RTX 4090 (24 GB) only if quantized.
For a 70B model at 128K, FP16 cache is 320 GB as above. Even with INT4 weights (35 GB), total is 355 GB, requiring at least 4-8 H100s (80 GB each) with tensor parallelism.
def kv_cache_bytes(L, H, D, S, bytes_per_elem):
# 2 for key and value
return 2 * L * H * D * S * bytes_per_elem
# Example: 70B model, 80 layers, 64 heads, 128 dim, 128K tokens, FP16
print(kv_cache_bytes(80, 64, 128, 131072, 2) / 1e9, "GB")For batch size >1, the KV cache multiplies by batch size. A batch of 8 at 128K would require 2.5 TB of cache.
Real Hardware Budgets: Where Does Your Model Fit?
Let's map real GPUs to context lengths. We'll assume INT4 weights for the model (to make it fit) and FP16 KV cache (common default).
RTX 4090 (24 GB): With a 7B model in INT4 (~4 GB weights), you have ~20 GB for KV cache. That supports about 37K tokens for LLaMA-2 7B (cache 17 GB at 32K). For 128K, you'd need ~70 GB cache, impossible.
RTX 5090 (32 GB): Similar math. 7B INT4 leaves ~28 GB cache, supporting ~50K tokens. For 70B INT4 (35 GB weights), the card cannot even load the weights.
H100 (80 GB): 70B INT4 weights = 35 GB, leaving 45 GB for cache. That supports ~18K tokens for 70B. For 128K, you need 320 GB cache, impossible on one H100. With 8x H100 tensor parallel, each GPU handles 1/8 of cache: 40 GB per GPU, which fits. So 8x H100 can serve 70B at 128K with FP16 cache.
H200 (141 GB): 70B INT4 weights = 35 GB, leaving 106 GB for cache. That supports ~42K tokens. Still not enough for 128K. For 405B INT4 (203 GB), even the H200 cannot load weights alone.
MI300X (192 GB): 70B INT4 fits (35 GB), leaving 157 GB for cache -> ~63K tokens. Better but still not 128K. For 405B INT4, weights alone (203 GB) exceed single GPU.
B200 (192 GB): Similar to MI300X but with faster HBM3e. Still cannot do 128K on a single GPU for 70B with FP16 cache.
M3 Ultra (192 GB unified): With 70B INT4 (35 GB), you have 157 GB for cache -> ~63K tokens. But the bandwidth (800 GB/s) is lower than H100 (3.35 TB/s), so tokens/sec will suffer.
Consumer GPUs (4090, 5090) are limited to ~32K context for 7B models. For 70B, you need multi-GPU setups or aggressive KV cache quantization.
Quantizing the KV Cache: FP8, INT4, and KIVI
The KV cache is often stored in FP16 by default, but it can be quantized with minimal impact on accuracy. The key insight is that keys and values have different sensitivity: keys are more important for attention scores, values for output quality.
FP8 (1 byte per element) halves the cache size. Many H100 and B200 GPUs support FP8 natively in tensor cores, so you get free memory reduction without software overhead. For 70B at 128K, FP8 cache = 160 GB, still large but now fits on 2 H100s (80 GB each) with tensor parallelism.
INT4 (0.5 bytes per element) gives 4x reduction over FP16. With 70B at 128K, INT4 cache = 80 GB, fitting on a single H100. However, INT4 KV cache requires special kernels (e.g., from TensorRT-LLM or ExLlamaV2) and may degrade accuracy slightly. Benchmarks show <1% perplexity increase for most models at 4-bit.
KIVI (2-bit KV cache quantization) is a recent technique that uses per-channel quantization and achieves 8x compression. At 2 bits, 70B 128K cache = 40 GB. Combined with INT4 weights (35 GB), total = 75 GB, fitting on a single H100. KIVI requires custom CUDA kernels but is available in some forks of vLLM.
Important: KV cache quantization does not reduce memory bandwidth usage during attention, the quantized values must be dequantized on the fly. FlashAttention with FP8 KV cache is supported in CUDA 12.4+ and can improve throughput by reducing HBM traffic.
# Memory reduction factors
factors = {
'FP16': 1.0,
'FP8': 0.5,
'INT4': 0.25,
'INT2 (KIVI)': 0.125
}
base_cache = 320.0 # GB for 70B at 128K FP16
for fmt, factor in factors.items():
print(f"{fmt}: {base_cache * factor:.1f} GB")PagedAttention and vLLM: Better Utilization, Not Less Memory
vLLM's PagedAttention is often misunderstood as reducing KV cache memory. It does not. It eliminates internal fragmentation by managing the cache in fixed-size blocks (pages), allowing the cache to be non-contiguous in memory. This enables higher batch sizes and reduces wasted memory due to variable-length sequences.
In practice, PagedAttention can improve memory utilization by 10-30% in serving scenarios with many concurrent requests, but the total cache size for a given sequence length remains the same. If you need 320 GB of cache for a single 128K sequence, PagedAttention will still allocate 320 GB (plus a small overhead for page tables).
Where PagedAttention shines is in batching: without it, each sequence must have its own contiguous cache allocation, leading to fragmentation. With it, you can pack many shorter sequences into the same physical memory, increasing throughput. But for a single long-context request, the benefit is marginal.
vLLM also supports KV cache quantization (FP8, INT4) via its plugin system, which is the real memory saver. Combining PagedAttention with INT4 cache gives the best of both worlds: high utilization and low per-token memory.
Use vLLM with --kv-cache-dtype fp8 on H100 to halve cache memory without custom kernels.
FlashAttention: Not a Memory Panacea for KV Cache
FlashAttention (v1, v2, v3) is a landmark algorithm that reduces the memory complexity of the attention computation from O(n^2) to O(n) by tiling the Q, K, V matrices and recomputing parts of the softmax. This eliminates the need to store the full attention matrix (n x n) in HBM, which for long sequences would be enormous.
However, FlashAttention does not reduce the KV cache memory. The KV cache is still stored in HBM and grows linearly with sequence length. FlashAttention only reduces the memory for intermediate attention scores, which are temporary and not needed after the forward pass.
Where FlashAttention helps is in reducing HBM bandwidth usage during the attention step. By tiling, it reads the KV cache in blocks multiple times, but overall it reduces the total bytes transferred compared to standard attention. This translates to higher tokens/second, especially at long context lengths.
For example, on an H100, FlashAttention-2 can achieve ~80% of peak HBM bandwidth during attention, whereas standard attention might only reach 30-40%. At 128K context, this can mean the difference between 10 tokens/sec and 25 tokens/sec for a 70B model. But the VRAM requirement for KV cache remains unchanged.
// FlashAttention kernel pseudocode (simplified)
// KV cache is still fully stored in HBM
for each block in Q:
for each block in K:
load block of K from HBM
compute partial attention scores
// recompute softmax, no full n x n matrix stored
end
load corresponding V block
accumulate output
endTensor Parallelism: Spreading the Cache Across GPUs
When a single GPU cannot hold the KV cache, tensor parallelism (TP) is the standard solution. In TP, the model is split across GPUs along the hidden dimension: each GPU holds a fraction of the weights and computes a fraction of the attention heads. The KV cache is also sharded: each GPU stores only its own heads' keys and values.
For a 70B model with 64 heads per layer, using 8 GPUs means each GPU handles 8 heads. The KV cache per GPU becomes: 2 * 80 * 8 * 128 * S * 2 bytes = 327680 * S bytes. For S=128K, that's 42 GB per GPU. With INT4 weights (35 GB total, ~4.4 GB per GPU), each GPU needs about 46 GB, fitting comfortably on an 80 GB H100.
Pipeline parallelism (PP) also reduces per-GPU cache by splitting layers, but it introduces bubbles. Sequence parallelism (SP) splits along the sequence dimension but requires all-to-all communication. For long context, TP is the most common choice because it keeps the attention computation local and avoids network bottlenecks.
However, TP requires high-speed interconnects (NVLink, NVSwitch, or at least PCIe 5.0 with P2P) because each attention step requires an all-reduce across GPUs. With 8 GPUs, the all-reduce overhead can be 10-20% of total time. For 256K context, the communication cost grows as the attention scores are larger.
Practical tip: For 405B models at 256K context, you need at least 16 H100s with TP=16 and INT4 KV cache. The total VRAM for cache alone: 2 * 128 * 64 * 128 * 262144 * 0.5 (INT4) = 2 * 128 * 64 * 128 * 131072 = 2 * 128 * 64 * 16777216 = 2 * 128 * 1073741824 = 2 * 137438953472 = 274.9 GB. Spread across 16 GPUs, that's ~17 GB per GPU, plus weights (~12.7 GB INT4 per GPU) = ~30 GB per GPU, fitting on H100.
But the all-reduce for 256K tokens is massive: each GPU must exchange attention outputs of shape (batch, seq_len, hidden_dim). For hidden_dim 8192 and seq_len 256K, that's 2 GB per all-reduce. With 16 GPUs, the bandwidth of NVSwitch (900 GB/s per direction) makes this feasible, but on PCIe 5.0 x16 (64 GB/s), it would be a bottleneck.
For long context, tensor parallelism is mandatory. Use NVLink-connected GPUs (H100 SXM, B200) to minimize communication overhead.
Real-World Benchmarks: Tokens/sec vs Context Length
Let's look at concrete numbers from vLLM and TensorRT-LLM benchmarks. On a single H100 with a 70B model (FP16 weights, FP16 KV cache), tokens/sec drops from ~40 at 2K context to ~15 at 32K context, and to ~5 at 128K context. The drop is due to two factors: the KV cache grows, increasing HBM reads per token, and the attention computation scales linearly with sequence length (FlashAttention helps but cannot eliminate the O(n) reads).
With INT4 weights and FP8 KV cache on the same H100, tokens/sec at 128K context improves to ~12, because the smaller cache fits in HBM and reduces memory traffic. However, the dequantization overhead for INT4 weights adds latency, so the gain is not proportional to the memory reduction.
On 8x H100 with TP=8, 70B FP16, 128K context, tokens/sec can reach ~80 (aggregate) due to parallel computation and reduced per-GPU cache. The all-reduce overhead is about 15% of total time.
For 405B models, single-GPU inference is impossible. With 8x H100 and INT4 weights + INT4 cache, tokens/sec at 128K context is around 25-30 (aggregate). At 256K context, it drops to ~12 due to increased communication and attention time.
These numbers are from real deployments at companies like Together AI and Anyscale. The key takeaway: doubling context length roughly halves throughput, all else equal. To maintain throughput, you must either add more GPUs or quantize more aggressively.
# Example vLLM command for 70B with FP8 KV cache on 8x H100
python -m vllm.entrypoints.openai.api_server \
--model meta-llama/Meta-Llama-3-70B \
--tensor-parallel-size 8 \
--kv-cache-dtype fp8 \
--max-model-len 131072 \
--gpu-memory-utilization 0.95The Future: Context Lengths Beyond 256K and Hardware Trends
As models like Gemini 1.5 Pro (1M tokens) and GPT-4-128K push context boundaries, the VRAM challenge intensifies. The industry is responding with hardware innovations: HBM3e (H200, B200) offers up to 4.8 TB/s bandwidth and 192 GB capacity, but even that is insufficient for 1M context with a 70B model.
For 1M tokens, a 70B model in INT4 with INT4 KV cache would require: cache = 2 * 80 * 64 * 128 * 1e6 * 0.5 = 655 GB. Weights = 35 GB. Total = 690 GB. That would require 4x B200 (192 GB each) or 9x H100 (80 GB each). The bandwidth required to read the entire cache for each token is enormous: at 1M tokens, the attention step reads the full KV cache, which is 655 GB. At 30 tokens/sec, that's 19.65 TB/s bandwidth, far exceeding any single GPU's HBM bandwidth.
Future solutions include: - Sparse attention (e.g., SparseGPT, BigBird) that only attends to a subset of tokens, reducing cache reads. - Hierarchical KV cache with offloading to CPU memory or SSDs (e.g., InfiniGen, FlexGen). - Near-memory compute (like Samsung's HBM-PIM) that performs attention inside the memory stack, reducing data movement. - CXL-based memory pooling for disaggregated KV cache across nodes.
For now, the practical limit for local inference on consumer hardware is about 128K context for 7B models (with INT4 cache) and 32K for 70B models (with multi-GPU). The economics of running 405B at 256K on-prem are prohibitive for most teams, costing >$200K in GPU hardware alone.
If you need 1M context, consider using a retrieval-augmented generation (RAG) pipeline instead of shoving everything into the prompt.
Pitfalls and common misconceptions
- 1Assuming FlashAttention reduces KV cache memory: it only reduces intermediate attention memory, not the cache itself.
- 2Believing PagedAttention shrinks the cache: it only improves utilization by reducing fragmentation, not the total size.
- 3Thinking INT4 weights alone are enough: for long context, KV cache dominates, so you must quantize both.
- 4Expecting linear throughput scaling with more GPUs: communication overhead and Amdahl's law limit scaling, especially at long context.
- 5Overlooking activation memory for large batch sizes: at batch size >1, activations can become significant, especially with long sequences.