02.07 · Walkthrough
Where the Memory Went
Build a memory budget for a transformer inference run: weights, KV cache, activations. Identify which component dominates at which batch size and context length.
Transformer inference memory is mostly weights at small batch and short context, then KV cache as live tokens accumulate, with activations and runtime workspace adding phase-specific pressure. A useful budget separates resident weights, per-token KV, prefill activations, and fragmentation so admission control is based on memory, not just compute.
What this lesson answers
- how to budget transformer inference memory
- when does KV cache dominate inference memory
- why long context causes inference OOM
Notes
A transformer inference memory budget is the sum of resident weights, resident KV cache, transient activations/workspace, and allocator fragmentation: . For dense decoder-only models, weight memory is , so a 70B model in FP16 needs GB before tensor-parallel sharding.
References
Common questions
- What takes memory during transformer inference?
- The main buckets are model weights, KV cache, activations, workspace, and allocator overhead. Weights are resident for the model shard. KV cache grows with active sequences and stored tokens. Activations are usually smaller during decode, but can matter during prefill, especially with long prompts or less memory-efficient attention kernels.
- Why can a model fit but still run out of memory?
- Fitting the weights only proves the base model can be resident. Serving also needs KV cache for every live token, plus buffers for attention, logits, communication, graph capture, and allocator slack. A run that looks safe at load time can fail once batch size or context length grows.
- How does paged KV cache help serving?
- Paged KV cache allocates storage in blocks for tokens that actually exist, instead of reserving a full maximum-length region for every request. This reduces wasted memory when live sequence lengths vary. It helps most under mixed-length traffic, and least when all requests are already near the configured maximum.