Memory Management: Fitting a Model and Deciding Concurrency

A model fitting on a GPU is only the first memory check. Serving also needs temporary buffers and a KV cache for every active request. A configuration that loads successfully can still run out of memory when longer prompts arrive or a larger batch starts decoding.

The useful calculation is a budget: subtract everything the runtime already needs, then translate the remaining bytes into stored token positions. This gives a capacity bound to test, rather than a promised concurrency level.

What occupies memory

Weights depend on parameter count and representation. A rounded 70-billion-parameter model has about 140 decimal GB of raw 16-bit weights, 70 GB at eight bits, or 35 GB at four bits. Quantization scales, zero points, mixed-precision tensors, alignment, and runtime copies add overhead.

Runtime memory includes activations, attention workspaces, communication buffers, graph captures, and allocator reservations. Its size can change with prefill chunk size, batch shape, and backend. KV cache is not the only allocation that grows with workload.

KV cache stores keys and values for processed token positions. Paged allocation reduces unused reservations, but occupied entries still cost bytes. Finally, leave headroom for allocations that the simplified calculation misses.

All GB figures here use $10^9$ bytes. MiB and GiB instead use powers of two; mixing the units can turn apparent spare capacity into an out-of-memory error.

Bytes per stored token

For ordinary full-context attention with the same KV shape at every layer,

where $L$ is layer count, $H_{\text{KV}}$ is the number of KV heads, $d_h$ is head dimension, and $s$ is bytes per stored value. The leading two counts keys and values.

For an 80-layer architecture with eight KV heads, dimension 128, and 16-bit storage:

These are the Llama 3.1 70B KV dimensions. With full cache storage, 4,096 tokens cost about 1.34 GB for one sequence before block padding and metadata. Quantized KV, sliding windows, hybrid layers, latent attention, and distributed placement need their own accounting.

For requests with lengths $T_1,\ldots,T_C$, the unshared payload is $b_{\text{token}}\sum_i T_i$. The common shortcut $C\times T\times b_{\text{token}}$ assumes every request has the same stored length. Shared prefixes reduce duplicated entries; paging rounds each private allocation to block boundaries.

From bytes to a concurrency bound

If 14 GB remains for KV cache after weights, runtime memory, and your chosen reserve, the ideal capacity is

stored token positions. Twenty requests at 2,000 positions use 40,000 of them, leaving only about 6.4% of that KV budget free. Eighteen such requests use 36,000 positions and leave about 15.7%. These percentages concern the KV allocation, not total GPU memory.

A scheduler also needs room for requests to grow. Admitting everything that fits now may force preemption later. Conversely, reserving every request's worst-case completion wastes capacity. Real admission policies balance that risk against the workload's length distribution.

The calculator is not a compatibility checker. A model's supported context length and a backend's quantization formats are separate constraints. Tensor parallelism may replicate KV heads when the parallel degree exceeds the number of heads; total cluster memory is not always perfectly divisible among GPUs.

Long context changes the budget

At 131,072 positions, the 70B example needs about 42.95 GB of FP16 KV payload for one sequence. That is separate from its weights. A smaller model may have a smaller cache, but parameter count alone does not determine it.

Using the paper’s 405B dimensions—126 layers, eight KV heads, and dimension 128—the same context requires about 67.65 GB. Use the checkpoint’s actual KV configuration when sizing a deployment.

What offloading costs

Moving tensors to host RAM or storage expands the available capacity, but introduces a different bottleneck. If a design transfers $D$ bytes for each step over a link sustaining $B$ bytes per second, transfer alone takes at least $D/B$ seconds. Transfers can overlap computation, but cannot escape the bandwidth requirement.

CPU execution is not the same as copying all weights to the GPU every token. A runtime can execute some layers on the CPU and transfer their activations. Likewise, memory-mapping a model file does not mean every step reads it from an SSD: the operating system may keep pages in RAM. Identify which tensors move, how often, and where computation runs before estimating speed.

Fast GPU interconnects help distributed execution, but their headline bandwidth is not directly comparable with HBM bandwidth without checking direction, topology, and the collective being used.

A capacity calculation to adapt

from math import floor
 
layers, kv_heads, head_dim, bytes_per_value = 80, 8, 128, 2
bytes_per_token = 2 * layers * kv_heads * head_dim * bytes_per_value
kv_budget_bytes = 14 * 10**9  # after other allocations and reserve
slots = floor(kv_budget_bytes / bytes_per_token)
assert bytes_per_token == 327_680
assert slots == 42_724
 
stored_lengths = [2000] * 18
payload = sum(stored_lengths) * bytes_per_token
print(f"KV payload: {payload / 10**9:.3f} GB")
print(f"Remaining KV budget: {(kv_budget_bytes - payload) / 10**9:.3f} GB")

To reduce memory, choose a compatible lower-precision representation, shorten retained context, reuse valid prefixes, or distribute the model. GQA is an architectural choice that usually requires a model trained or adapted for it, rather than an engine switch. The next post examines quantization, including its overhead and quality tradeoffs.

Additional reading

• The Llama 3 Herd of Models, Dubey et al. Architecture and context-length details.
• Llama 3.1 70B model card. Checkpoint and deployment context.
• PagedAttention, Kwon et al. Cache allocation and sharing.
• vLLM optimization guidance. Runtime memory and scheduling controls.