KV Cache Memory › Combined Effects: The Memory Budget
Language AI HandbookAll chapters KV Cache Memory › Combined Effects: The Memory BudgetOne section · Sections in this chapter7 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
In practice, we must account for all memory consumers simultaneously. The GPU has a fixed amount of memory, and every component of inference must fit within this budget. Omitting any component from memory accounting causes out-of-memory errors and crashes inference. Accurate memory estimation is essential for reliable deployment, and there is no substitute for working through the full accounting before you commit to a hardware configuration.
The total GPU memory budget must accommodate all major contributors working at the same time. Model weights are fixed once loaded and remain constant regardless of input or batch size, representing the irreducible baseline cost of using a particular model. The KV cache is the dynamic component that varies most dramatically with your batch size and sequence length choices. Activation memory holds intermediate results like attention scores and feed-forward layer outputs, growing with batch size but not with sequence length in the same way as the cache. Framework overhead, including CUDA context, PyTorch allocator buffers, and communication buffers in multi-GPU settings, represents a fixed cost that exists regardless of what computation you perform.
The interplay between these components creates a constrained optimization problem. Model weights cannot be reduced without using a smaller model or applying quantization. Framework overhead is similarly fixed. This leaves the KV cache and activation memory as the variables that respond to your choices about batch size and sequence length. When a deployment configuration fails to fit in GPU memory, the first lever you should reach for is cache optimization: reducing the number of KV heads via GQA, quantizing cache values to INT8 or INT4, or implementing page-level eviction when cache exceeds available memory.
Let us build a memory estimator that accounts for all these factors.
def estimate_inference_memory(
model_params: int,
num_layers: int,
num_kv_heads: int,
head_dim: int,
batch_size: int,
seq_len: int,
dtype_bytes: int = 2,
activation_factor: float = 0.1, # Fraction of weights
) -> dict:
"""
Estimate total GPU memory for inference.
activation_factor: Rough multiplier for activation memory (varies by implementation)
"""
# Model weights
weight_bytes = model_params * dtype_bytes
weight_gb = weight_bytes / (1024**3)
# KV cache
cache = calculate_kv_cache_memory(
num_layers, num_kv_heads, head_dim, batch_size, seq_len, dtype_bytes
)
# Activation memory (rough estimate)
activation_gb = weight_gb * activation_factor
# Framework overhead (rough estimate)
overhead_gb = 0.5 # CUDA context, allocator buffers
total_gb = weight_gb + cache["GB"] + activation_gb + overhead_gb
return {
"weights_gb": weight_gb,
"kv_cache_gb": cache["GB"],
"activation_gb": activation_gb,
"overhead_gb": overhead_gb,
"total_gb": total_gb,
"kv_cache_pct": (cache["GB"] / total_gb) * 100,
}# Memory breakdown for different configurations
configs = [
{"name": "Short context (2K)", "batch": 1, "seq": 2048},
{"name": "Medium context (8K)", "batch": 1, "seq": 8192},
{"name": "Long context (32K)", "batch": 1, "seq": 32768},
{"name": "Batched (4K x 8)", "batch": 8, "seq": 4096},
]
memory_results = []
for config in configs:
mem = estimate_inference_memory(
model_params=llama_7b["total_params"],
num_layers=llama_7b["num_layers"],
num_kv_heads=llama_7b["num_kv_heads"],
head_dim=llama_7b["head_dim"],
batch_size=config["batch"],
seq_len=config["seq"],
)
memory_results.append((config["name"], mem))LLaMA 7B Memory Breakdown (FP16) ====================================================================== Short context (2K): Weights: 13.0 GB KV Cache: 1.00 GB (6% of total) Activation: 1.3 GB Overhead: 0.5 GB TOTAL: 15.8 GB Medium context (8K): Weights: 13.0 GB KV Cache: 4.00 GB (21% of total) Activation: 1.3 GB Overhead: 0.5 GB TOTAL: 18.8 GB Long context (32K): Weights: 13.0 GB KV Cache: 16.00 GB (52% of total) Activation: 1.3 GB Overhead: 0.5 GB TOTAL: 30.8 GB Batched (4K x 8): Weights: 13.0 GB KV Cache: 16.00 GB (52% of total) Activation: 1.3 GB Overhead: 0.5 GB TOTAL: 30.8 GB
This breakdown demonstrates how dramatically memory requirements shift across different usage patterns. For short 2K contexts, model weights account for roughly 82% of total memory, with KV cache representing about 6%. At 32K tokens, the cache grows to approximately 52% of total memory and overtakes the weights. Batching 8 sequences at 4K context produces the same 16 GB cache and the same breakdown as one 32K sequence because batch size and sequence length enter the formula as a product. This transition from weight-dominated to cache-dominated memory profiles changes deployment strategies and hardware selection.
The percentage breakdown reveals how memory allocation shifts with configuration. At short contexts, weights dominate at about 85% of total memory. At long contexts or large batches, the KV cache becomes the primary memory consumer. This shift demonstrates why cache optimization becomes critical for long-context applications. An engineering team that sizes hardware for a short-context model and then later extends the context window may discover that the hardware that seemed adequate is suddenly insufficient, even though the model itself has not changed.