Account for weights, workspace, and the KV state actually resident on each rank

For a conventional dense decoder with equal key and value widths, each retained token contributes a key and a value at every attention layer. Multiply the number of layers by local KV heads, head width, element bytes, and the number of retained token positions. Grouped-query attention has fewer KV heads than query heads; substituting query-head count overestimates that particular payload. Architectures with different key/value widths, sliding windows, compressed latent states, or unequal layers require their own sum rather than this simplified product. Specify the architecture and cache dtype independently of the weight dtype.

A parameter estimate similarly multiplies stored elements by bytes only when the format is truly uniform and unpacked. Quantized weights can add scales, zero points, padding, and special full-precision tensors. In ordinary inference there is no reason to automatically add training optimizer states, but an accidental gradient graph or retained activation can still consume memory. Include logits, attention workspace, communication buffers, CUDA graph pools, and runtime allocations where applicable. Prefill and decode can have different peaks, so observing a quiet decode iteration does not bound the memory required to admit a new long prompt.

When KV heads divide evenly across tensor-parallel ranks, each rank can retain its head shard. Once the rank count exceeds the number of KV heads, an implementation may replicate heads rather than store fractional heads. The tagged vLLM 0.9.1 Llama attention code explicitly partitions in the former case and replicates in the latter, with divisibility checks. For example, eight KV heads across sixteen ranks leave one resident KV head per rank in that implementation, not half a head; aggregate storage is then twice the unreplicated payload.

Pipeline placement, replicated embeddings, communication workspaces, and uneven layer assignment can further change per-rank peaks. Compute the budget for the busiest rank, not just total cluster memory divided by worker count. A shared prompt also has a logical token count larger than its unique physical KV storage when correctly reused, while incompatible prefixes cannot be merged merely because token lengths match. Record the exact rank topology, model revision, dtype, and implementation version with a memory estimate so later capacity measurements can be compared against the same assumptions.

Compare predicted payload growth with measured allocated memory, pool occupancy, and device-level usage while controlling batch shape and prefix sharing. Differences can reveal block rounding, replicated heads, cached workspaces, or allocations outside the framework profiler. Leave explicit headroom for transient peaks and concurrent work. A measured average is not an admission invariant, and a bandwidth-limited decode kernel can still fail because capacity is exhausted. Keep capacity in bytes separate from throughput in bytes per second or floating-point operations per second.

The roofline lab isolates that final distinction. Reset, explicitly set BOTH intensity I=8 and bandwidth B=80, predict a ceiling of 640 GFLOP/s, then Play or Step. Reset and re-set BOTH I=8 and B=160; predict the fixed 1000-GFLOP/s compute cap. Edits retain t; playback traces the roof without changing the selected kernel. B uses decimal GB/s. This synthetic ceiling has no model tensors, KV heads, devices, or allocator, and it cannot determine whether either configuration has enough memory to hold one request.