KV Cache Memory › Batch Size Effects

Michael BrenndoerferJanuary 7, 20266 min read
Language AI HandbookAll chapters
Part 42 of 67Inference Optimization
Chapter 314KV Cache Memory › Batch Size Effects
Reading section 6 of 14 · Batch Size Effects · 6 minSection 6 of 14 · 6 min

Book contents

Section 6 of 14 · 6 min read

KV Cache Memory › Batch Size EffectsOne section · Sections in this chapter6 / 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

Batch size scales linearly with cache memory, following the same multiplicative relationship we observed with sequence length. Each sequence in the batch maintains its own independent cache. Tokens in one sequence do not attend to tokens in another sequence, so cached values are not shared across batch elements. Think of batch processing as running BB conversations simultaneously in BB isolated rooms. Each room has its own filing cabinet, and they are completely separate. You cannot combine the cabinets to save space.

This independence changes the interpretation for memory planning. If you want to process 8 concurrent requests, you need 8 complete copies of the KV cache, one for each sequence. The batch dimension in our cache tensors is not a clever abstraction that enables sharing; it simply organizes multiple independent caches for efficient parallel processing. The GPU executes them in parallel across its compute units, which amortizes the latency of loading model weights, but every cache must be allocated separately in memory.

The batch size tradeoff is at the heart of inference economics. Serving individual requests one by one wastes GPU resources because the arithmetic units sit idle while waiting for data transfers to complete. Batching multiple requests allows the GPU's parallel processing capabilities to run at higher utilization. But this throughput gain comes at the direct cost of proportionally more cache memory. Doubling batch size doubles throughput potential but also doubles cache requirements. The optimal batch size is always the largest that fits in available memory after accounting for model weights, activation memory, and framework overhead.

In[12]:
Code
# Calculate cache for various batch sizes
batch_sizes = [1, 2, 4, 8, 16, 32, 64]
cache_by_batch = []

for batch in batch_sizes:
    cache = calculate_kv_cache_memory(
        num_layers=llama_7b["num_layers"],
        num_kv_heads=llama_7b["num_kv_heads"],
        head_dim=llama_7b["head_dim"],
        batch_size=batch,
        seq_len=4096,
        dtype_bytes=2,
    )
    cache_by_batch.append(cache["GB"])
Out[13]:
Visualization
Bar chart showing KV cache memory increasing with batch size from 1 to 64.
Linear scaling of KV cache memory for LLaMA 7B at 4,096-token context in FP16 across batch sizes from 1 to 64. Memory increases from 2.0 GB with one sequence to 128.0 GB with 64 concurrent sequences. Each sequence maintains its own independent cache. The dashed line shows the 13.0 GB model-weight reference; cache memory first exceeds it at batch size 8.

This visualization clearly demonstrates the linear relationship between batch size and KV cache memory. The batch size tradeoff is fundamental to inference economics, as it directly determines the throughput your system can achieve in production deployment. Doubling batch size doubles memory requirements, but it also doubles the number of requests you can process in parallel. This creates a direct tension between throughput and memory availability:

  • Single-sequence inference minimizes cache memory but leaves GPU compute underutilized
  • Large batches maximize throughput (tokens per second) but require proportionally more memory
  • The optimal batch size is typically the largest that fits in available memory after accounting for model weights and activation memory

GPU compute is expensive, and leaving it idle while processing a single sequence is wasteful. Batching amortizes the fixed costs of loading model weights and allows the GPU's parallel processing capabilities to be fully utilized. But this amortization is only possible if you have sufficient memory to hold multiple caches simultaneously.

In[14]:
Code
# What batch size fits in 80GB (A100) for LLaMA 7B?
total_gpu_memory = 80  # GB
model_weights = weight_gb
activation_overhead = 2  # GB (rough estimate for intermediate computations)
available_for_cache = total_gpu_memory - model_weights - activation_overhead

# Cache per sequence at 4096 tokens
cache_per_seq = calculate_kv_cache_memory(
    num_layers=llama_7b["num_layers"],
    num_kv_heads=llama_7b["num_kv_heads"],
    head_dim=llama_7b["head_dim"],
    batch_size=1,
    seq_len=4096,
    dtype_bytes=2,
)["GB"]

max_batch = int(available_for_cache / cache_per_seq)
Out[15]:
Console
GPU Memory Budget: 80 GB
  Model weights:     13.0 GB
  Activation buffer: 2.0 GB
  Available for KV:  65.0 GB
  Cache per sequence: 2.00 GB
  Maximum batch size: 32

This calculation reveals excellent throughput potential for LLaMA 7B on high-end hardware. With an A100 80GB GPU running LLaMA 7B at 4096-token context, the available memory after weights and overhead can support a batch size that significantly exceeds what most consumer hardware could achieve. The 80 GB memory capacity of the A100 is precisely calibrated for workloads like this, which is why it remains the workhorse of LLM serving infrastructure.

Dynamic Batch Size Management

Production serving systems cannot always predict the sequence length of incoming requests in advance. A request that starts with a short prompt may generate a long response, gradually consuming more cache memory than expected. This variability complicates static batch size planning. Modern inference servers like vLLM and TensorRT-LLM address this with continuous batching and dynamic cache management, which we explore in the Paged Attention chapter. For now, the key insight is that naive static batch size planning is too conservative: it must reserve worst-case memory for every sequence, wasting capacity for short requests.

The safest approach for deployment planning is to calculate the maximum cache memory using the longest sequence you expect to encounter, then size your batch accordingly. If your 99th percentile request requires 8K tokens, plan your batch size for that length, even if most requests are much shorter. This conservative approach guarantees you will not run out of memory, but it leaves efficiency on the table. Dynamic allocation systems can recover that efficiency by reallocating unused cache blocks from short sequences to long ones.