Inference memory model
Estimate capacity before running a server
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Bytes per datatype
- Transformer parameters
- KV-cache concept
The 7-billion-parameter service starts on one 80 GiB GPU. A process-level memory query reports 15 GiB allocated immediately after BF16 model weights load, 18 GiB after runtime initialization and graph capture, and a steadily growing allocation as prompts are admitted. A 32,000-token request can consume gigabytes of additional attention state. When free blocks reach zero, the scheduler stops admitting requests even though the GPU's arithmetic units are not fully busy.
Record memory after each controlled event: process start, weight load, graph capture, one 1000-token prefill, and each additional 1000 generated or retained tokens. Compare the observed deltas with the model architecture. The experiment separates fixed model storage, per-request metadata, per-token attention state, temporary workspaces, allocator reservation, and safety reserve before the lesson assigns names or formulas to those categories.
Memory capacity can limit inference throughput. Parameters use a large fixed allocation. Each active sequence adds KV-cache state based on layer count, KV heads, head dimension, token count, and bytes per value. Runtime workspaces, graph pools, communication buffers, allocator fragmentation, and safety margin use additional capacity. Include all categories when you calculate admission limits.
Separate resident, per-request, per-token, and transient bytes. Weight quantization reduces fixed resident memory. Grouped-query attention reduces per-token KV memory. Paged allocation can reduce unused reservations. An attention workspace can increase peak memory without changing long-lived capacity. The scheduler needs the incremental memory cost of each admitted sequence as well as a reserve for transient peaks.
The purpose of the calculation is not to predict one static allocation exactly. The calculation identifies which term grows when the workload changes. Longer contexts grow KV state. More concurrent requests add block tails and scheduler metadata. A larger graph-capture set can reserve more fixed memory. Speculative verification and long prefill can increase transient workspace.
Calculate KV bytes per retained token
For a decoder layer, attention retains one key vector and one value vector for each token and KV head. The approximate bytes per token are 2 × layer count × KV-head count × head dimension × bytes per cache element. Multiply by the number of physical retained tokens. Query-head count does not appear directly because grouped-query and multi-query attention allow several query heads to share one KV head.
Apply sharding only when the runtime actually partitions the corresponding KV heads or dimensions across ranks. Some metadata, buffers, or layers can remain replicated. For pipeline parallelism, each rank normally stores cache for its assigned layers. For tensor parallelism, the exact cache distribution depends on the model and attention implementation. Confirm the per-rank allocation rather than dividing by world size automatically.
Paged caches allocate fixed token blocks. A sequence whose length is not a multiple of block size leaves unused positions in its final block. Shared prefixes can make several logical sequences refer to the same physical blocks. Branches or beam candidates can share blocks until copy-on-write is required. Admission must calculate physical ownership, not only the sum of logical lengths.
Account for fixed and per-request memory
Weight memory includes packed values, scale and zero-point metadata, duplicated embeddings or output heads, adapter weights, and any alternative copy retained for unsupported operations. Quantization labels such as 4-bit do not imply exactly half a byte per parameter after grouping, alignment, and metadata. Use the actual artifact and runtime representation.
Per-request state outside KV can include token arrays, sampling state, logits buffers, block tables, sequence metadata, and network buffers. The per-request allocations are usually smaller than weights or long-context KV, but they matter at high sequence counts and on the CPU scheduler. Record both device and host memory.
Graph capture, compilation, libraries, communication, and caching allocators reserve memory that may not appear as a live model tensor. Measure memory after initialization, after graph capture, during long prefill, during steady decode, and during the largest supported mixed batch. The difference between allocated, reserved, and free memory depends on the runtime API.
Use peak memory to set admission headroom
The steady-state KV calculation does not include every peak. Prefill can require attention, logits, or multimodal workspaces. Speculative decoding evaluates several positions and can create larger temporary batches. Cache transfer, recomputation, or compaction can need a destination buffer before old storage is released. Distributed collectives also reserve communication buffers.
Keep a free-block and byte reserve large enough for these peaks and for in-flight requests to finish. If the scheduler assigns every block to new admissions, an active sequence may fail when it generates its next token. A capacity plan should distinguish blocks available for new requests from blocks reserved for guaranteed progress of admitted requests.
Compare the formula with server-reported cache capacity and with a measured high-water mark under the target workload. Differences are useful: they identify metadata, replication, allocator rounding, or a mistaken sharding assumption. Preserve the calculation so a change in cache dtype, model architecture, or parallel plan can be evaluated before deployment.
bytes = 2 × layers × KV heads × head dim × bytes/value × tokensAdd page tails, metadata, replication, workspaces, and safety headroom; subtract only real sharding or sharing.
Worked example: Estimating KV capacity
A model has 32 layers, 8 KV heads, head dimension 128, and BF16 cache values. The server wants 128,000 retained tokens per request at concurrency eight.
- Per token, KV uses 2 × 32 × 8 × 128 × 2 bytes = 131,072 bytes, or 128 KiB before sharding.
- One 128,000-token request therefore needs about 15.26 GiB. Eight such fully populated requests would exceed 122 GiB before weights and overhead.
- Apply actual tensor-parallel cache partitioning if each rank stores only a subset of KV heads, then add block-tail and allocator overhead. Prefix sharing may reduce physical tokens when prompts repeat.
- Compare with measured allocation under realistic length distribution. Most requests may be shorter, so admission should reserve against bounded risk rather than a fictional all-maximum average.
Conclusion: The formula makes long-context concurrency concrete and explains why fewer KV heads, lower cache precision, paging, and admission policy are system-level levers.
Common errors
- Multiplying parameter count by dtype and calling it server memory. KV and runtime allocations can dominate long-context operation.
- Assuming every logical token owns unique KV. Prefix sharing and beam or branch semantics change physical ownership.
- Using average length without tail protection. A few long requests can consume the free-block reserve and stall all admissions.
Reference summary
KV bytes/token ≈ 2 × layers × KV_heads × head_dim × bytes/valueServing memory contains fixed weights, per-token KV state, per-request metadata, runtime reserves, graph-capture pools, communication buffers, temporary workspaces, and allocator overhead. Calculate the terms separately so a change in context length, concurrency, precision, or graph configuration has a predictable effect.
Admission needs both current free capacity and future progress capacity. Preserve enough KV blocks and temporary bytes for admitted sequences to generate their next tokens and for peak operations to complete. A configuration that allocates every block to new requests can fail after those requests begin decode.
- Weight memory may include packed/quantized values, scales, and duplicate formats.
- KV capacity is governed by cache dtype, layers, KV heads, head dimension, and tokens.
- Tensor parallelism shards some weights and caches but adds communication and duplicated metadata/workspace.
- Prefix sharing and copy-on-write reduce duplicate KV when prompts overlap.
- Eviction and recomputation policies convert memory shortage into latency or compute.
Knowledge check
Why does grouped-query attention reduce KV-cache bytes without reducing the number of query heads?
Show the answer guide
- Every query head still produces its own query vector and output contribution, so grouped-query attention does not reduce the declared number of query heads.
- Several query heads are assigned to one shared key head and one shared value head. The model stores and reuses one K vector and one V vector for the group instead of one pair per query head.
- KV-cache bytes per retained token are proportional to the number of KV heads, not the number of query heads: 2 times layers times KV heads times head dimension times bytes per element.
- If 32 query heads use 8 KV heads, four query heads share each KV pair and the K/V head contribution to cache size is one quarter of full multi-head attention, while all 32 query projections remain.
Verify KV-cache growth for MHA and GQA
Instantiate or analytically model two otherwise equal 32-layer attention configurations with 32 query heads, head dimension 128, and BF16 cache: one with 32 KV heads and one with 8 KV heads. Allocate caches for batch sizes 1 and 8 at retained lengths 128, 2048, and 8192, and measure device-memory deltas if a runtime is available.
Evidence to keep: Submit the calculation code, a table of predicted and measured bytes for every batch-length pair, the explicit cache tensor shapes, and a plot of bytes versus retained tokens. Explain any difference from the four-to-one prediction using allocator blocks, metadata, sharding, or measurement reserve.
Memory capacity bounds active work. The serving roofline explains how phase, batch, and reuse determine whether that work is limited by weights, KV traffic, or compute.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.
- MLPerf Inference Rules↗
- Orca: Iteration-Level Scheduling↗
- PagedAttention / vLLM↗
- NVIDIA Nsight Systems User Guide↗
Official reference for capturing and reading CPU, CUDA API, GPU workload, and communication timelines.