kv-cache-size 0.1.1

Exact KV-cache arithmetic for transformer inference: bytes per token, total cache bytes, and the context length that fits a memory budget.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

# kv-cache-size

Exact KV-cache arithmetic for transformer inference. No dependencies, no allocation,
no model downloads, no network.

```text
kv_bytes = 2 x bytes_per_element x num_hidden_layers x num_key_value_heads
           x head_dim x context_length x batch_size
```

The leading `2` is one K tensor plus one V tensor.

## The two terms people get wrong

1. **`num_key_value_heads`, not `num_attention_heads`.** Grouped-query attention keeps
   fewer key/value heads than query heads (Ainslie et al., *GQA*, arXiv:2305.13245), so
   using the query-head count overstates the cache by the GQA group size — 4x on
   Llama 3.1 8B, 8x on Qwen2.5-7B.
2. **`head_dim` as published, not derived.** `hidden_size / num_attention_heads` is a
   habit that breaks on configs where the published `head_dim` disagrees with it.
   `head_dim_from_hidden` exists for configs that genuinely omit the field, and is
   deliberately separate from the main path so the fallback is visible at the call site.

## Example

```rust
use kv_cache_size::{KvCacheConfig, KvPrecision};

// Llama 3.1 8B: 32 layers, 8 key-value heads, head_dim 128.
let cfg = KvCacheConfig::new(32, 8, 128).unwrap();

assert_eq!(cfg.bytes_per_token(KvPrecision::Bf16), 131_072.0); // 128 KiB per token
assert_eq!(cfg.total_bytes(KvPrecision::Bf16, 8192, 1).unwrap(), 1_073_741_824.0);

// How much context fits in 16 GiB of spare VRAM at fp8, batch 4?
let budget = 16.0 * 1024.0 * 1024.0 * 1024.0;
assert_eq!(cfg.max_context(KvPrecision::Fp8, budget, 4).unwrap(), 65_536);
```

## What this crate does not do

It sizes the KV cache only. Weights, activations, CUDA context and allocator
fragmentation are not included, so the number is a floor for planning, not a
capacity guarantee. Interactive version and per-model configs:
<https://ml0x.com/calculators/kv-cache-size-calculator.html>.

Licensed under MIT OR Apache-2.0.