How much memory does a model need?

A running language model is two things sitting in memory: a few billion fixed numbers that are the model itself, and a pile of notes it keeps about every conversation it is in the middle of, which grows with every word. This is how to size both on paper, before you download a single file.

That arithmetic decides which models your Spark can run at all, how long each conversation can get, and how many people can use it at once. Get it wrong and the server refuses to start. On the Spark it can be worse: the GPU and the rest of the computer share one pool of memory, so overshooting doesn’t just fail your script, it can make Linux start killing programs. And “how much memory does a 70B model need?” is a classic interview question with a crisp answer that most people fumble.

You have already met pieces of this. Why ~17 tokens per second? showed how many bytes each past word of a conversation costs. What does the OS do with my memory? showed how the Spark’s shared memory behaves when it fills up. Here you put them into one budget and check it against your own server from the baseline. Next, The KV cache, by hand has you build those conversation notes yourself, and batching turns “how many people fit” into “how fast each of them gets served”.

Today

• Say, to within a gigabyte and without looking anything up, how much memory a model needs to serve a given number of people at a given conversation length, split into four parts (the model itself, each conversation’s running notes, the scratch space used during one step, and overhead), and say which part is biggest.
• Count a model’s learned numbers (its parameters) by hand from the small settings file that ships with it, config.json, matrix by matrix, and land on Llama 3.1 8B’s 8,030,261,248 exactly.
• Turn a storage format, from plain 16-bit numbers down to the squeezed 4-bit ones (BF16, FP8, Q8_0, AWQ, Q4_K_M), into bytes per number, including the small scale factors each squeezed format quietly stores alongside.
• Read the memory numbers PyTorch and the Spark report, and say why PyTorch’s two counters (memory_allocated, memory_reserved) and the operating system all give different answers.
• Predict how many conversations, at what length, your Spark can hold before it runs out, then find the real edge and explain the gap.

Before this, you want Why ~17 tokens per second? (the per-word memory formula, and the trick of letting attention heads share their stored keys and values), Baseline: my Spark vs Fireworks (your server and benchmark), and What does the OS do with my memory? (the Spark’s shared memory, Linux’s file cache, and the out-of-memory killer). Budget one evening of arithmetic and one session on the Spark.

The Problem

Here is the question from the diagnostic: how much memory does an 8B model need? The honest answer you gave was a shrug, and the common wrong answer is “8 GB, one byte per parameter”. Neither will get you far.

Now make it real. Your team wants to use the Spark as a shared assistant: Llama 3.1 8B, long chats, 32K tokens of context each. You start vLLM, it thinks for a minute, and prints a line like maximum concurrency for 32,768 tokens per request: 24x. Is that good? Is it right? Where did 24 come from? Then you start it again with a 70B model in 4 bits at 128K context and it prints something like 2x. Then a colleague runs a notebook on the same box and something gets killed.

Every one of those numbers comes from a few lines of arithmetic. By the end of this lesson you will do that arithmetic faster than vLLM boots.

The Idea

Four buckets

Think of a restaurant kitchen during service. There is equipment bolted to the floor: ovens, fridges, the walk-in. That is the same size whether you serve one table or forty. Then each table has an order ticket and a tray of their own food in progress, and that space grows with the number of tables and how long each one stays. While a dish is being plated there is a burst of counter space in use, gone a minute later. And there is the space you can never use: walkways, the fire-code clearance, the manager’s desk.

GPU memory during inference has exactly those four buckets.

1. Weights. The model itself. Bolted to the floor: loaded once, shared by every user. Size = parameters × bytes per parameter.
2. KV cache. Each user’s private tray. Grows with every token of every conversation. Size = bytes per token × context length × users.
3. Activations. The temporary numbers inside one forward pass: the vector flowing through a layer, the wide MLP intermediate, the scores at the end. Freed as soon as the layer is done. Small in decode, a short spike during prefill.
4. Overhead. Memory that is in use but holds none of the above: the CUDA runtime’s own setup, library scratch space, memory the allocator holds but isn’t using, and on the Spark, the whole operating system.

The whole lesson is one equation and the details that make it honest:

users that fit  =  (memory − OS − weights − overhead − activation peak)  ÷  (KV bytes per token × context)

Weights and KV cache are pure arithmetic, so we start there.

Bytes, GB and GiB

One unit trap first, because it will otherwise eat 7% of every answer. A GBGBA gigabyte: 1,000,000,000 bytes (10⁹). Used by disk makers, Hugging Face file sizes, and network speeds. is a billion bytes. A GiBGiBA gibibyte: 2³⁰ = 1,073,741,824 bytes, about 7% more than a GB. Memory chips come in powers of two, so RAM and GPU memory are really sized in GiB, and PyTorch and nvidia-smi (which prints MiB) report in binary units. is 2³⁰ bytes. Llama 3.1 8B in bf16 is 16.06 GB and 14.96 GiB; both are right. Memory is built in powers of two, so when you compare against what a machine holds, work in GiB. The KV cache formula is naturally binary (128 KiB per token, so 8,192 tokens is exactly 1 GiB), which is one more reason to stay there.

Counting parameters from config.json

Every model on Hugging Face ships a config.json. It is a dozen numbers that completely determine the shape of every weight matrix. Here are the ones that matter for Llama 3.1 8B:

{
  "vocab_size": 128256,
  "hidden_size": 4096,
  "intermediate_size": 14336,
  "num_hidden_layers": 32,
  "num_attention_heads": 32,
  "num_key_value_heads": 8,
  "head_dim": 128,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16"
}

You have written this model, so you know every piece. The job now is just to say how big each one is. Call hidden_size h, vocab_size V, intermediate_size I, heads n, KV heads k, head_dim d.

The embedding table. One row of h numbers per vocabulary entry: V × h = 128,256 × 4,096 = 525,336,576. That is half a billion parameters before you reach the first layer, which surprises everyone the first time.

Attention, per layer. Four projection matrices. The query projection maps h inputs to n × d outputs. The key and value projections map h inputs to only k × d outputs, because with GQAGQAGrouped-query attention: the model has fewer key/value heads than query heads, and each K/V head is shared by a group of query heads. Llama 3.1 8B has 32 query heads and 8 K/V heads, so groups of 4. there are only 8 key/value heads. The output projection maps n × d back to h.

q_proj   4096 × (32 × 128) = 4096 × 4096 = 16,777,216
k_proj   4096 × ( 8 × 128) = 4096 × 1024 =  4,194,304
v_proj   4096 × ( 8 × 128) = 4096 × 1024 =  4,194,304
o_proj   (32 × 128) × 4096 = 4096 × 4096 = 16,777,216
                                   total  = 41,943,040

That asymmetry, K and V a quarter the size of Q, is Shazeer’s idea from the box above, sitting in your parameter count.

The MLP, per layer. Llama uses SwiGLUSwiGLUA feed-forward block with three matrices instead of two: a “gate” and an “up” projection run side by side, the gate goes through the SiLU function and multiplies the up output elementwise, then a “down” projection maps back. More expressive per parameter than the old two-matrix MLP., so there are three matrices, not two: gate (h × I), up (h × I), down (I × h).

3 × 4096 × 14336 = 176,160,768

The MLP is four times the attention block. Most of an LLM’s parameters live here.

Norms. Each layer has two RMSNormRMSNormA normalization that divides each vector by its root-mean-square and multiplies by a learned scale. The scale is one number per hidden dimension, so it has h parameters and no bias. layers (one before attention, one before the MLP), each with h = 4,096 learned scales: 8,192 per layer. One more after the last layer: 4,096. Rounding error, but count them; it is how you get an exact match.

The output head. The lm_head maps h back to V scores, another 525,336,576. Unless tie_word_embeddings is true, in which case the model reuses the embedding table and this costs zero. Llama 3.1 8B: false, so it is a separate matrix. Llama 3.2 1B: true, so it is shared. The flag is worth 525 million parameters (1 GB in bf16); always read it.

Add it up:

per layer      = 41,943,040 + 176,160,768 + 8,192         =   218,112,000
32 layers      = 32 × 218,112,000                          = 6,979,584,000
embed + head   = 2 × 525,336,576                           = 1,050,673,152
final norm                                                 =         4,096
total                                                      = 8,030,261,248

8.03 billion. I checked this against transformers itself: building the model from this config on PyTorch’s meta device (which creates tensor shapes without allocating any memory) and summing numel() over its parameters gives 8,030,261,248, to the parameter. The same method gives 1,235,814,400 for Llama 3.2 1B (tied head) and 70,553,706,496 for Llama 3.1 70B. You will write the counter yourself in Build It.

Two useful facts fall out of the breakdown. The MLP is 70% of the model, attention 17%, embedding and head 13%. And the shortcut for a quick estimate: per layer, attention is about 2h² + 2hkd and the MLP is 3hI, so for this family the layers come to roughly L × (2.5h² + 3hI). You can do that in your head in an interview: 32 × (2.5 × 16.8M + 176M) ≈ 7.0B, plus 1.05B of vocabulary matrices.

Bytes per parameter

Weights in memory = parameters × bytes per parameter. The bytes depend on the format.

• FP32FP3232-bit floating point, 4 bytes per number. The default for training math in the old days; almost nobody serves in it now.: 4 bytes. 8B model: 32.1 GB.
• BF16 or FP16: 2 bytes. 16.1 GB. This is what most checkpoints ship as (torch_dtype in the config). Why bf16 and not fp16 is the number formats lesson.
• FP8FP8An 8-bit floating point format, 1 byte per number, with a scale factor per tensor or per channel so the small range still covers the weights.: 1 byte for the layers. Popular FP8 checkpoints usually leave the embedding and lm_head in bf16, so an 8B FP8 file is about 9.1 GB, not 8.0.
• Integer formats with block scales. You cannot just round a weight to one of 16 values (4 bits) without a sense of scale. So these formats chop each matrix into small blocks and store one or two scale numbers per block alongside the tiny integers. The scales are real bytes, and they are why “4-bit” is never 4.0 bits.

The cleanest example is llama.cpp’s GGUFGGUFThe single-file model format used by llama.cpp. It holds the weights (in many possible quantized types), the tokenizer, and the config in one file. type Q8_0: blocks of 32 weights, each stored as a signed 8-bit integer, plus one 16-bit scale per block. 32 bytes + 2 bytes = 34 bytes per 32 weights = 8.5 bits per weight.

The other two 4-bit formats you will meet:

• AWQAWQActivation-aware Weight Quantization (Lin et al., 2023). Stores layer weights as 4-bit integers in groups of 128, each group with a 16-bit scale and a 4-bit zero point, choosing the scales so the weights that matter most for real activations lose least. checkpoints quantize the transformer layers in groups of 128 weights. Each group adds a 16-bit scale and a 4-bit zero point: 20 extra bits per 128 weights, so about 4.16 bits per layer weight. But the embedding and lm_head usually stay in 16-bit. For Llama 3.1 8B that is 1.05 billion parameters at 2 bytes (2.1 GB) plus 6.98 billion at 4.16 bits (3.6 GB): about 5.7 GB, which is what published AWQ files of this model weigh. The “4-bit model” averages 5.7 bits per weight.
• Q4_K_M in GGUF uses “K-quant” super-blocks of 256 weights, split into 8 sub-blocks of 32 with their own 6-bit scales and minimums. The basic 4-bit K type works out to 4.5 bits per weight. The “_M” mix keeps some of the most sensitive tensors (parts of the value and down projections, and the output head) in a 6-bit type. The average across the file lands around 4.8 to 4.9 bits, so an 8B file is about 4.9 GB.

The general rule: file size on disk is a very good estimate of weight memory. Almost every byte of a safetensors or GGUF file is weight. When you do not know a format’s internals, ls -l the file and divide by the parameter count to get its true bits per weight.

The KV cache: per token, per request, per batch

You met the formula in the bandwidth lesson. Here it is again, because now it is the main event:

KV bytes per token  =  2 (K and V) × layers × kv_heads × head_dim × bytes per value

Llama 3.1 8B, bf16:   2 × 32 × 8 × 128 × 2  =  131,072 B  =  128 KiB

Read each factor. The 2 is because every layer stores a key vector and a value vector for every past token. Layers multiply because each layer has its own attention. And only the KV heads count, not the query heads, because queries are recomputed from the current token every step and never stored. That last point is why GQA’s 8-instead-of-32 heads is a 4x discount on memory.

From the per-token number, everything else is multiplication:

Notice the 70B’s 32-user row: the KV cache alone would fill an H100. Notice the 8B at 128K: one user’s cache weighs more than the whole model. And notice the MoE: 30 billion parameters but a smaller cache per token than the 8B, because it has only 4 KV heads. Cache size depends on the attention shape, not on the parameter count.

Two practical modifiers. First, store the cache in FP8 and every number in that table halves; most serving engines offer it as one flag, with some accuracy cost you would measure. Second, “context” means the whole conversation so far, prompt plus everything generated. A request that allows 8K tokens grows toward that over its life; a serving engine has to decide whether to reserve the worst case up front or allocate as it goes. That decision is what paged attention is about.

Activations: small, temporary, spiky

During a forward pass, each layer produces temporary tensors: the normed input, q, k and v, the attention output, and the MLP’s wide intermediates. Those are activationsactivationsThe intermediate values computed during a forward pass. In training you keep them for the backward pass; in inference you can throw each one away as soon as the next layer has used it.. At inference, nobody needs them after the next layer runs, so PyTorch frees them right away. Only one layer’s worth is alive at a time.

In decode, a step processes one new token per user. For Llama 3.1 8B the biggest temporaries are the gate and up projections, 14,336 numbers each. Per token, everything alive in one layer is roughly (4h + 3I) × 2 bytes ≈ 120 KB. Times 64 users is under 8 MB. Noise.

In prefill, the step processes the whole prompt at once, so every temporary is N tokens tall. At N = 8,192 the MLP alone briefly holds about 8,192 × 3 × 14,336 × 2 bytes ≈ 0.7 GB, and the full set is about 1 GB. It lasts milliseconds, but it is the peak that has to fit, not the average. Two things make the peak much worse if you are careless:

• Computing logits for every position. The output head turns each token’s vector into 128,256 scores. If you run the model’s forward on a 32K prompt and ask for logits at every position, that is 32,768 × 128,256 × 2 bytes ≈ 8.4 GB of scores, of which you need exactly one row (the last). Per token that is about 250 KB, twice the KV cache. Hugging Face’s generate avoids this by asking for only the last position (it passes logits_to_keep=1; I checked this in the source of a current transformers 5 release); a bare model(input_ids) call does not.
• Naive attention. Textbook attention builds an N × N score matrix per head. At 16K tokens with 32 heads in bf16 that is 32 × 16,384² × 2 bytes = 16 GiB, per layer, before softmax makes an fp32 copy. Fused kernels (PyTorch’s SDPA, FlashAttention) never build that matrix, so their memory grows with N, not N². That is the FlashAttention lesson.

Real engines cap the activation peak on purpose. vLLM, for example, splits long prompts into chunks and never runs more than a set number of tokens through the model in one forward pass (its max_num_batched_tokens setting), so the spike has a known size no matter how long the prompt.

Overhead: the memory nobody asked for

This is the bucket that makes careful predictions wrong by a gigabyte or three, so it is worth knowing its parts.

The CUDA context. The first time a process touches the GPU, the CUDA driver sets up its state: loaded GPU code for every kernel library you link (PyTorch ships many), internal tables, scratch areas. That memory belongs to the driver, not to PyTorch, so PyTorch’s own counters never show it. On discrete GPUs it is typically a few hundred MB to around 1 GB per process, depending on versions and on whether the driver loads kernel code lazily. Measure yours; it is per process, so two Python processes using the GPU pay it twice.

The caching allocator. Asking the CUDA driver for memory (cudaMalloc) and giving it back is slow and forces the GPU to synchronize. So PyTorch keeps a caching allocatorcaching allocatorPyTorch’s GPU memory manager. When a tensor is freed, it keeps the block for reuse instead of returning it to the driver, so the next allocation of a similar size is instant.: freed blocks stay in PyTorch’s pool for the next tensor. That gives you two numbers:

• torch.cuda.memory_allocated(): bytes currently held by live tensors.
• torch.cuda.memory_reserved(): bytes PyTorch has taken from the driver, including cached free blocks. Always ≥ allocated.

The gap between them is memory that is free to PyTorch but not to anyone else. If the pool’s free blocks are the wrong sizes for a new request (a 300 MB tensor needs 300 contiguous MB within a block, and the pool has ten 100 MB holes), PyTorch has to ask the driver for more. That waste is fragmentationfragmentationFree memory broken into pieces too small or oddly placed to satisfy a request. The total free is enough; no single piece is., and it is why “out of memory” can happen when the numbers say there should be room. A tensor that grows by concatenation is the worst case: Hugging Face’s default DynamicCache appends each new token’s keys and values with torch.cat, which allocates a slightly bigger tensor and frees the old one, every step, in every layer.

I watched the caching behaviour on my Mac’s GPU backend, which has the same design: allocating a 256 MiB tensor made the driver hold 1,024 MiB (it reserves in big chunks); after del, live tensors were 0 MiB while the driver still held 1,024 MiB; only torch.mps.empty_cache() gave it back. On CUDA, torch.cuda.empty_cache() does the same job. You will see the CUDA version of this on the Spark.

Everything else. cuBLAS scratch space, NCCL buffers if you use several GPUs, and in serving engines, memory captured by CUDA graphs (pre-recorded sequences of kernel launches that cut Python overhead, each with its own fixed buffers). Budget about 1 to 3 GiB in total for a single-GPU server and then measure.

The Spark: one pool, shared with everything

On an H100 or a 4090, the GPU’s memory is its own. The operating system, your browser, your Python process’s CPU-side objects all live in separate system RAM. The 80 GB is yours.

The Spark has no separate GPU memory. As you saw in the virtual memory lesson, its 128 GB is one pool shared by the CPU and GPU. So the budget has a new line at the top: the OS, the desktop if one is running, every CPU process, and the page cache all come out of the same 128 GB. Some of that is reclaimable (the kernel will drop cached file pages under pressure); some is not. And the failure mode is different: on a discrete card, running out gives you a polite Python exception. On the Spark, the pressure lands on the whole machine, and you may meet the kernel’s OOM killer instead, or a frozen desktop while the system thrashes.

Two consequences for measuring. Tools built for discrete GPUs may not know what to report: on unified-memory systems nvidia-smi’s memory columns can read “Not Supported” or show something other than what you expect, so check what yours prints and do not trust it blindly. And the honest “how much can I still use?” number comes from Linux: free -h’s available column, or MemAvailable in /proc/meminfo. PyTorch’s torch.cuda.mem_get_info() asks the CUDA driver the same question; compare the two on your machine, because on unified memory the driver’s idea of “free” may not count file cache the kernel would happily drop.

Serving engines take it all up front

A serving engine cannot allocate KV cache memory one request at a time with torch.cat and hope; it would fragment, and it would discover it was out of memory in the middle of a user’s answer. So engines like vLLM grab the cache’s memory once, at startup, and hand out fixed-size pieces of it.

vLLM’s version, roughly (the details move between versions; check the docs for yours):

1. You set gpu_memory_utilizationgpu_memory_utilizationA vLLM setting (default 0.9): the fraction of the GPU’s total memory this vLLM instance is allowed to use for everything, weights and cache included. (default 0.9). That is the instance’s budget: 0.9 × total memory.
2. It loads the weights.
3. It runs a fake forward pass at the largest size it will ever run (max_num_batched_tokens tokens) and records the peak activation memory, plus whatever non-PyTorch memory has appeared.
4. Whatever is left of the budget becomes the KV cache pool, divided into blocks of a few tokens each. It logs the pool’s size in tokens, and the “maximum concurrency” it implies: pool tokens ÷ max_model_len.

This explains the opening scene. vLLM is not guessing when it says “maximum concurrency 24x”; it is dividing the pool it carved by the context you allowed. It also explains the colleague’s notebook. On the Spark, gpu_memory_utilization is not “how much of the GPU”, it is “how much of the whole computer”. NVIDIA’s DGX Spark playbooks show the settings they test with; look there before choosing yours, and leave real headroom. Newer vLLM versions also let you give the KV cache size directly in bytes instead of a fraction; check vllm serve --help for what your version supports.

Putting it together

Now play with all of it. Each bar is drawn to its own machine’s capacity. The Spark’s first segment is the operating system and everything else that lives in its shared pool; it is set to 10 GiB as an assumption until you measure it. Overhead is 1.5 GiB of runtime plus a computed activation peak for an 8K-token prefill chunk.

Things to try:

• Defaults (8B, BF16, 8K, 8 users). The Spark and the H100 have room to spare: about 100 and 62 users at this context. The 4090 misses by about 1.4 GiB; it tops out at 6. A 24 GB card is a weights-plus-a-few-chats machine.
• Drag context to 128K. The 4090 cannot hold even one user. The H100 holds 3. One user’s cache is now bigger than the model.
• Now flip the KV cache to FP8. Every “max users” number roughly doubles. This is the cheapest capacity win there is, which is why every engine offers it.
• Llama 3.1 70B, BF16. 131 GiB of weights. Nothing here holds it, including the Spark. Switch to Q4_K_M: now the Spark and the H100 can hold it (the 4090 still cannot), and at 32K context the Spark fits about 7 users, the H100 about 3.
• Qwen3 30B-A3B, BF16. 57 GiB of weights, which the H100 can hold with room for about 27 users at 8K, but the 4090 can’t load at all. As the bandwidth lesson showed, it decodes like a 3B model. Big-but-sparse on a big-but-slow memory is the Spark’s niche.

One thing this figure deliberately ignores: speed. A user who fits is not a user who is served well. On the Spark, 100 users at 8K means every decode step reads 15 GiB of weights plus 100 GiB of caches through 273 GB/s, about 0.45 seconds per step, or about 2 tokens per second each. The envelope lesson did that half. Capacity tells you who can be in the room; bandwidth tells you how fast the room moves.

Build It

Everything goes in inference/memory/ next to last lesson’s server and harness. The first three tasks run anywhere (do them on your Mac); the last three need the Spark.

# fetch_configs.py
import json, pathlib
from huggingface_hub import hf_hub_download

REPOS = [
    "meta-llama/Llama-3.2-1B",
    "meta-llama/Llama-3.1-8B",
    "meta-llama/Llama-3.1-70B",
    "Qwen/Qwen2.5-7B",
    "Qwen/Qwen3-8B",
    "Qwen/Qwen3-30B-A3B",
]
out = pathlib.Path("configs")
out.mkdir(exist_ok=True)
for repo in REPOS:
    name = repo.split("/")[1]
    cfg = json.load(open(hf_hub_download(repo, "config.json")))
    (out / f"{name}.json").write_text(json.dumps(cfg, indent=2))
    try:  # sharded safetensors checkpoints have an index whose metadata gives total bytes
        idx = json.load(open(hf_hub_download(repo, "model.safetensors.index.json")))
        cfg_total = idx.get("metadata", {}).get("total_size")
        print(f"{name:20s} total_size={cfg_total}")
    except Exception as e:
        print(f"{name:20s} no index ({type(e).__name__})")

def count_params(cfg: dict) -> dict[str, int]:
    """Return {'embed': ..., 'attn': ..., 'mlp': ..., 'norms': ..., 'lm_head': ..., 'total': ...}
    summed over all layers. Dense Llama/Qwen-style decoders only (MoE is a stretch goal)."""
    ...

EXPECTED = {
    "Llama-3.2-1B": 1_235_814_400,
    "Llama-3.1-8B": 8_030_261_248,
    "Llama-3.1-70B": 70_553_706_496,
    "Qwen2.5-7B": 7_615_616_512,
    "Qwen3-8B": 8_190_735_360,
}

GiB = 2**30

FORMATS = {  # effective bits per weight: (layers, embedding + lm_head)
    "bf16": (16, 16), "fp8": (8, 16), "q8_0": (8.5, 8.5), "awq": (4.16, 16), "q4_k_m": (4.9, 4.9),
}

def weight_bytes(cfg: dict, fmt: str) -> float: ...
def kv_bytes_per_token(cfg: dict, kv_bytes: int = 2) -> int: ...
def plan(cfg, fmt, ctx, users, capacity_gib, reserve_gib=0.0, overhead_gib=1.5, kv_bytes=2) -> dict:
    """Return each bucket in GiB, the total, whether it fits, and max_users at this ctx."""
    ...

# probe.py: one line of every memory number we care about.
import os, time, torch

GiB = 2**30

def meminfo_available() -> float:
    with open("/proc/meminfo") as f:
        for line in f:
            if line.startswith("MemAvailable:"):
                return int(line.split()[1]) * 1024 / GiB  # value is in KiB
    return float("nan")

def rss() -> float:
    with open(f"/proc/{os.getpid()}/status") as f:
        for line in f:
            if line.startswith("VmRSS:"):
                return int(line.split()[1]) * 1024 / GiB
    return float("nan")

def snap(label: str) -> dict:
    if torch.cuda.is_initialized():
        torch.cuda.synchronize()
        free, total = torch.cuda.mem_get_info()
        alloc, resv, peak = (torch.cuda.memory_allocated(), torch.cuda.memory_reserved(),
                             torch.cuda.max_memory_allocated())
    else:
        free = total = alloc = resv = peak = float("nan")
    row = dict(label=label, t=time.time(), alloc=alloc / GiB, reserved=resv / GiB, peak=peak / GiB,
               cuda_free=free / GiB, cuda_total=total / GiB, os_available=meminfo_available(), rss=rss())
    print(" ".join(f"{k}={v:.2f}" if isinstance(v, float) else f"{k}={v}" for k, v in row.items()))
    return row

def fits(model, batch: int, ctx: int) -> bool:
    """Prefill `batch` random prompts of length ctx, decode 8 tokens, keep the cache.
    Return False on torch.cuda.OutOfMemoryError (clean up: del, gc.collect(), empty_cache())."""
    ...

def find_edge(model, ctx: int, lo: int, hi: int) -> int:
    """Binary search for the largest batch that fits at ctx."""
    ...

Measure

Fill these in. Predictions go in before you run anything.

Parameter counts (your counter vs the checkpoint’s total_size ÷ bytes per param):

The probe on the Spark, Llama 3.1 8B bf16:

The edge:

vLLM: predicted KV pool ___ GiB / ___ tokens; logged ___ tokens; max concurrency predicted ___, logged ___.

What the numbers should look like, and why. These are expectations to check, not results.

• Weights should be the most boring number on the page: within a few MB of 8,030,261,248 × 2 bytes. The few MB are small buffers that are not parameters (rotary embedding frequencies, for example). If yours is off by a whole GiB or more, a tensor got upcast or duplicated.
• The CUDA context shows up as memory gone from the OS with nothing allocated in PyTorch. Somewhere from a few hundred MB to a GB or so is normal on discrete GPUs; I don’t know the Spark’s number, which is why it is the first row you measure. Run two Python processes that each touch CUDA and watch it double.
• KV growth should be very close to L × 128 KiB per sequence, because DynamicCache holds exactly the keys and values and nothing else. The peak during each prefill will sit well above the steady value: that is the activation spike, plus the brief double copy when the cache is concatenated. Expect reserved to climb to the largest peak and stay there.
• After del and empty_cache(), alloc returns to the weights and reserved drops close to it. If os_available does not come back fully, part of the gap is the context, which only goes away when the process exits.
• The edge should land a little below your prediction, not above it. The common reasons: fragmentation (reserved well above alloc at failure), the prefill spike at large B × L (push prompts through in chunks to remove it), and the OS reserve being larger than you assumed (a desktop session, a browser, a forgotten notebook). If you measure more than predicted, check whether something in your stack is quantizing the cache.
• vLLM’s pool should match your prediction within a few percent if you used the same overhead and activation numbers vLLM profiled. The log tells you how much it set aside for activations; plug that in and the match should tighten.

Break It

Ask for every logit. Replace generate with a plain model(input_ids) on a 32K prompt. Predict first: 32,768 × 128,256 × 2 bytes. Watch peak. The fix is one argument: request logits only for the last position (logits_to_keep=1 in current transformers; check the model’s forward signature in your version). Lesson: a single tensor nobody needed can be bigger than the whole KV cache.

Use textbook attention. Load the model with attn_implementation="eager" and prefill 8K, then 16K tokens. Predict the score matrix: heads × N² × 2 bytes per layer, plus an fp32 softmax copy. At 16K that is tens of GiB for a single layer’s temporaries; it will fail where SDPA sailed through. Lesson: the activation peak depends on the algorithm, not just the model, which is the whole reason FlashAttention exists.

Run two at once. Start two vLLM instances on the Spark, each with the default gpu_memory_utilization. The second one will fail at startup, or, worse, the first one’s budget plus the OS will leave nothing for the second and the machine will start reclaiming and swapping. Then try a notebook that loads a model while one vLLM instance is running. Lesson: on unified memory, a serving engine’s budget is a claim on the whole machine. Write the reserve you need into your launch script.

Fragment the pool. Generate 4,000 tokens one at a time at batch 8 with DynamicCache, snapping every 250 tokens, and plot reserved − alloc. Then repeat with a pre-sized cache (transformers’ StaticCache, which allocates the full length once; check the docs for how to pass it to generate) and compare. Try also with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, which lets the allocator grow a segment in place instead of making new ones. Lesson: allocating once and slicing beats growing, which is exactly what serving engines do.

Explain It Back

Answer each out loud before opening it.

Go Deeper

• “Transformer Inference Arithmetic”, kipply (Carol Chen), 2022. The calculations in this lesson (parameters, KV cache, FLOPs, batch size) worked through for real models. Read the memory sections alongside your calculator.
• “Efficient Memory Management for Large Language Model Serving with PagedAttention”, Woosuk Kwon et al., SOSP 2023. The vLLM paper. Its opening measurements of how much KV cache memory earlier servers wasted to fragmentation and over-reservation are this lesson’s overhead section, at scale.
• “Fast Transformer Decoding: One Write-Head is All You Need”, Noam Shazeer, 2019. Multi-query attention, and the clearest early statement that KV memory is the decoding problem.
• “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, Joshua Ainslie et al., EMNLP 2023. Why Llama 3’s attention looks the way it does in your parameter count.
• “AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration”, Ji Lin et al., MLSys 2024. Where the group-of-128 scales come from, and why some weights deserve more care.
• PyTorch docs, “CUDA semantics” (memory management section) and “Understanding CUDA Memory Usage”. The caching allocator, memory_reserved, PYTORCH_CUDA_ALLOC_CONF, and the memory snapshot tool that shows every allocation on a timeline. Use the snapshot when t5’s numbers don’t add up.
• vLLM docs, engine arguments and the memory/optimization pages. The current meaning of gpu_memory_utilization, max_num_batched_tokens, kv_cache_dtype and friends. These move between versions; this is where to check.
• The GGUF and quantization type descriptions in the ggml / llama.cpp repositories. The exact block layouts of Q8_0 and the K-quants, if you want to compute bits per weight from first principles instead of from file sizes.

Connections

Built on: Why ~17 tokens per second? (the KV formula and GQA), Back-of-the-envelope (users per GPU, and why bandwidth is the Spark’s first wall), and virtual memory (unified memory, the page cache, and the OOM killer you will meet if you skip the container limit).

Next: The KV cache, by hand, where you build the cache this lesson only measured and watch your own GPT-2 speed up. Then batching, where “max users” turns into throughput, and paged attention, where the pool vLLM carved gets split into blocks and shared. Quantization is the other side of the weights bucket: you now know what a 4-bit file costs in bytes; that lesson measures what it costs in quality. And when you are ready for the bigger version of this lesson, Why training needs so much memory adds gradients, optimizer states and the activations that inference gets to throw away.