KVCacheEstimator

A tiny Swift library that sizes the key-value cache of a transformer model from its config.json, for anyone planning on-device or server memory for LLM inference.

kv_bytes = 2 (K and V) x bytes_per_element x num_hidden_layers x num_key_value_heads x head_dim
           x context_length x batch

import KVCacheEstimator

let shape = try ModelShape(configJSON: Data(contentsOf: configURL))
shape.bytesPerToken(.fp16)                                   // Llama 3.1 8B: 131,072 bytes (128 KiB)
shape.bytes(contextLength: 8192, batch: 4).gibibytes         // 4.0
shape.naiveMultiHeadBytes(contextLength: 8192, batch: 4)?.gibibytes  // 16.0, the common 4x mistake

Two terms are the ones people get wrong:

• Use num_key_value_heads, not num_attention_heads. Grouped-query attention keeps fewer K/V heads than query heads (Ainslie et al., arXiv 2305.13245). naiveMultiHeadBytes shows how far off the query-head shortcut is.
• Read head_dim when the config publishes it. hidden_size / num_attention_heads is only the fallback; some real configs publish a different head_dim. The parser refuses a config it cannot size instead of guessing.

Precisions: .fp16 (bf16/fp16, 2 bytes), .fp8 (1 byte), .int4 (0.5 byte).

An interactive version with presets for 19 open models is the ml0x KV cache calculator.

Install

.package(url: "https://github.com/theluckystrike/KVCacheEstimator.git", from: "0.1.0")

No dependencies beyond Foundation. MIT licensed.