Throughput Optimization › Batch Size Tuning

Choosing the right batch size is an empirical exercise, but several principles guide the search. You cannot simply pick the largest batch that fits in memory: you need to understand the relationship between batch size, GPU compute utilization, and the memory-bandwidth bottleneck that governs autoregressive decoding.

The first constraint is GPU memory. For each request in a batch, you must hold:

• The model weights (fixed cost, shared across all requests)
• The activation tensors during the forward pass (proportional to batch size and sequence length)
• The KV cache for all previous tokens (proportional to batch size and sequence length)

As batch size grows, the KV cache and activations grow proportionally. At some point, you exceed available GPU memory and must either use a smaller batch or switch to a smaller model.

The approximate memory budget for a batch is:

$$M_{\text{total}} = M_{\text{weights}} + B \times (M_{\text{activations}} + M_{\text{kv\_cache}})$$

where:

• $M_{\text{weights}}$: fixed memory for model parameters (roughly $2 \times \text{params}$ bytes in FP16)
• $B$: batch size
• $M_{\text{activations}}$: activation memory per request (proportional to sequence length)
• $M_{\text{kv\_cache}}$: KV cache memory per request (proportional to context length and layer count)

For a 7B model on an A100 80GB GPU, model weights consume about 14 GB, leaving roughly 60 GB for batch data. If each request uses 600 MB at a 2048-token context, maximum batch size is about 100 requests. In practice, you leave a safety margin of 10-15% to avoid out-of-memory errors from measurement imprecision and batch composition variance.

This calculation clarifies why KV cache quantization (computing and storing KV entries in INT8 or FP8 rather than FP16) has such a large impact. Halving KV cache memory roughly doubles the batch size you can support at a given context length, which doubles throughput for memory-bound workloads.

GPUs have two bottlenecks: compute (how fast arithmetic operations run) and memory bandwidth (how fast data moves from HBM memory to the compute units). Large language model inference is almost entirely memory-bandwidth-bound during token generation, because the prefill of a single token reads the entire model weight matrix once but performs only a small number of arithmetic operations per byte of weight data.

To understand why, consider a single linear layer with weight matrix $W \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}$ during decode. For a batch of $B$ tokens, the operation is $y = xW$ where $x \in \mathbb{R}^{B \times d_{\text{in}}}$. The number of floating-point operations is $2 \times B \times d_{\text{in}} \times d_{\text{out}}$. The number of bytes read from HBM is approximately $d_{\text{in}} \times d_{\text{out}} \times 2$ bytes (for FP16 weights, loaded once). The ratio of operations to memory access is:

$$\text{arithmetic\_intensity} = \frac{2 \times B \times d_{\text{in}} \times d_{\text{out}}}{d_{\text{in}} \times d_{\text{out}} \times 2} = B$$

The arithmetic intensity equals the batch size. At $B = 1$, you perform 1 floating-point operation per byte read. At $B = 100$, you perform 100 operations per byte read. The A100 GPU has a peak arithmetic intensity (the "roofline") of roughly 156 FLOP/byte at its memory bandwidth and FP16 peak throughput. You need a batch size of about 156 just to reach the boundary between memory-bound and compute-bound operation.

Increasing batch size helps by amortizing the weight-read cost over more requests. With batch size $B$, each weight-read generates $B$ tokens instead of one. The effective memory efficiency is:

$$\text{arithmetic\_intensity} = \frac{\text{FLOPs per batch}}{\text{bytes accessed per batch}} \approx \frac{2 \times \text{params} \times B}{\text{param\_bytes} + B \times \text{activation\_bytes\_per\_token}}$$

At small $B$, the denominator is dominated by $\text{param\_bytes}$ (reading the weights once), so arithmetic intensity is low and the GPU is memory-bound. As $B$ grows, the numerator grows faster, and eventually the GPU becomes compute-bound. This crossover point is the sweet spot for batch size: you want to be just above it, where adding more to the batch no longer helps because you have already saturated the compute units.

The crossover point shows the minimum batch size needed to saturate GPU compute. Below this, you are underutilizing the GPU because the hardware spends most of its time waiting for memory reads. Above this, you are in the compute-bound regime and additional batch size yields diminishing throughput gains.

This analysis has practical implications beyond just picking a batch size. It tells you that:

• Small batch sizes (1-8) are particularly wasteful because the weight memory reads dominate. Moving from batch size 1 to 8 improves throughput almost eightfold with minimal added latency.
• The return on increasing batch size flattens significantly beyond the crossover point. Going from batch size 200 to 400 improves throughput much less than going from 1 to 2.
• Quantization (reducing bytes per weight) shifts the roofline intensity downward, meaning you reach the compute-bound regime at a smaller batch size. INT8 quantization roughly halves the bytes per weight, halving the crossover batch size and making it easier to saturate GPU compute.

The analysis above applies specifically to the decode phase. During prefill, the arithmetic intensity is much higher because you process a full prompt of $T$ tokens simultaneously. For a 256-token prompt, the effective arithmetic intensity during prefill is approximately $256 \times$ higher than for a single decode step at the same batch size.

This difference explains why prefill and decode have different optimal batch sizes. For prefill, even batch size 1 may be compute-bound on long prompts. For decode, you often need batch sizes in the hundreds to reach compute saturation. This asymmetry motivates prefill-decode disaggregation, which we cover later in the chapter.

In practice, you find the optimal batch size by profiling, not by calculation alone. The memory estimates above involve many approximations, and the actual throughput depends on framework overhead, kernel selection, and driver behavior. A practical search procedure:

1. Start with a conservative batch size (e.g., 8) that definitely fits in memory.
2. Measure throughput (tokens per second) and GPU utilization with nvidia-smi or a profiler.
3. Double the batch size and repeat.
4. Stop when either GPU memory is exhausted or throughput improvement falls below 10%.
5. Back off slightly from the maximum to leave headroom for request length variance.

This procedure takes 30-60 minutes for a typical model and hardware combination, and the results are invaluable for production capacity planning. You will often discover that your initial batch size was 5-10x below optimal.