Skip to content
vivek pal

wiki / tiny llms

memory math for small models

How to estimate weights, KV cache, and working memory before you download a model.

updated · 4 min read

three buckets

A model running locally needs memory for three things: the weights, the KV cache, and working buffers for the forward pass. Weights are fixed once you pick a model and a precision. The KV cache grows with every token in the context. Working buffers depend on how many tokens go through the model at once.

The two models used below are Llama 3.2 1B (1.23B parameters) and Qwen3-4B (4.0B parameters). Every number comes from their config.json and model cards.

weights

weight bytes = parameters × bits per parameter / 8

Quantized formats store each weight as a small integer, plus a scale and a bias for each group of consecutive weights. MLX’s default affine mode uses groups of 64, with the scale and bias kept in the input dtype. With bf16 that is 32 extra bits per 64 weights, or 0.5 bits per weight, so “4-bit” really costs 4.5 bits. Groups of 32 cost 5 bits per weight and groups of 128 cost 4.25. Smaller groups track the original weights more closely and cost more memory.

precisionbits per weightLlama 3.2 1BQwen3-4B
bf16162.46 GB8.00 GB
8-bit, group 648.51.31 GB4.25 GB
4-bit, group 644.50.69 GB2.25 GB

Real files land a little off these numbers. MLX’s quantize only accepts arrays with two or more dimensions, so 1-D tensors such as norm weights stay in bf16. They are a tiny share of the total. Both models tie their input and output embeddings, so the embedding matrix is counted once.

kv cache

Each layer keeps a key vector and a value vector for every token in the context, for every KV head.

kv bytes per token = 2 × layers × kv_heads × head_dim × bytes per value

The 2 is one key plus one value. At fp16 or bf16, bytes per value is 2.

Grouped-query attention (GQA) matters here. Standard multi-head attention gives every query head its own key and value head. GQA shares each KV head across a group of query heads. Both example models have 32 query heads and 8 KV heads, which makes the cache 4x smaller than full multi-head attention would.

Llama 3.2 1B has 16 layers, 8 KV heads, and a head dim of 64.

2 × 16 × 8 × 64 × 2 = 32,768 bytes per token
 4,096 tokens × 32,768 = 0.13 GB
32,768 tokens × 32,768 = 1.07 GB

With 32 KV heads instead of 8, the 32k cache would be 4.3 GB.

Qwen3-4B has 36 layers, 8 KV heads, and a head dim of 128.

2 × 36 × 8 × 128 × 2 = 147,456 bytes per token
 4,096 tokens × 147,456 = 0.60 GB
32,768 tokens × 147,456 = 4.83 GB

At 32k tokens, Qwen3-4B’s cache is about twice the size of its 4-bit weights. On a 16 GB laptop, long context runs you out of memory before parameter count does.

mlx_lm.generate has two levers for this. --kv-bits quantizes the cache (off by default, groups of 64 via --kv-group-size), and --max-kv-size switches to a fixed-size rotating cache. At 8 bits the 32k Qwen3-4B cache drops to about 2.57 GB (8.5 / 16 of 4.83 GB). At 4 bits it drops to about 1.36 GB. Check output quality on your own task before keeping either.

working memory

The forward pass needs temporary buffers sized by the number of tokens processed at once. mlx-lm reads long prompts in chunks set by --prefill-step-size, 2048 tokens by default, so peak memory stays bounded while the prompt is read.

For Llama 3.2 1B (hidden size 2048, MLP size 8192), one 2048-token chunk at bf16 needs:

hidden states:     2048 × 2048 × 2 =  8.4 MB
MLP intermediate:  2048 × 8192 × 2 = 33.6 MB

A naive attention implementation would also materialize a chunk × context score matrix per head, which grows fast with context. That is one reason runtimes use fused attention kernels. Logits for one token are vocab × 4 bytes at fp32: 128,256 × 4 is about 0.5 MB. Leave headroom on top of all this for the runtime and the rest of the OS.

worked totals

Weights at 4-bit, group 64, plus the bf16 KV cache, before working buffers:

model4k context32k context
Llama 3.2 1B0.69 + 0.13 = 0.82 GB0.69 + 1.07 = 1.76 GB
Qwen3-4B2.25 + 0.60 = 2.85 GB2.25 + 4.83 = 7.08 GB

doing it for any model

Open the model’s config.json and read num_hidden_layers, num_key_value_heads, and head_dim. If head_dim is missing, it is usually hidden_size / num_attention_heads. Then:

def model_memory_gb(params, bits, layers, kv_heads, head_dim, context, kv_bytes=2):
    weights = params * bits / 8
    kv = 2 * layers * kv_heads * head_dim * kv_bytes * context
    return round(weights / 1e9, 2), round(kv / 1e9, 2)

print(model_memory_gb(4.0e9, 4.5, 36, 8, 128, 32_768))  # (2.25, 4.83)

sources

← all notes