KV Cache Bytes per Token
easyEvery running request holds a KV cache, the scratch memory a model uses to continue an answer from a prompt it has already read: after processing the prompt, the model has written a keyed summary of it, and each new token extends that summary instead of re-reading everything. That scratch space is the real resource being scheduled, so interviewers start here: how big is one token of it? The formula is 2 x bytes x layers x kv_heads x head_dim. The first 2 is because the cache stores both a key and a value per layer per token; bytes are 2 for bf16, the half-precision format models run in; kv_heads is the number of key-value attention heads, which modern models shrink below the query heads to save exactly this memory; head_dim is the width of each head. For an 8B-class model with 32 layers, 8 kv heads and 128 head_dim, one token of cache is 131,072 bytes. For a 70B-class model with 80 layers, also 8 kv heads and 128 head_dim, it is 327,680 bytes, and every seat in the batch pays that times its full sequence length (prompt plus answer so far) out of what remains after the weights. Implement kv_bytes_per_token(layers, kv_heads, head_dim): return the bytes one token of KV cache occupies for a bf16 model with that shape.
Run your code to see output