metrale_cache/
kv_cache.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The paged KV cache: the `KvCacheDtype` storage formats, the
4//! per-layer block geometry in `KvCacheConfig`, and the `PagedKvCache` block
5//! pool with its per-block reference counts.
6//!
7//! Owner: cache.
8//! Invariants:
9//! - Until `release`, every attention layer `i < num_layers` has its own K pool
10//!   and V pool of `num_blocks` blocks, sized by that layer's K-side and V-side
11//!   block bytes.
12//! - A block index is pushed back onto the free list only when a decrement
13//!   (`dec_ref`, `return_evicted_block`) takes its count from 1 to 0. A
14//!   decrement at 0 never pushes: it is logged, and `dec_ref` also panics in
15//!   a debug build.
16
17use anyhow::{Result, bail};
18use metrale_gpu_runtime::gpu::DevicePtr;
19
20pub(crate) const NVFP4_GROUP_SIZE: usize = 16;
21
22/// 2026-09-25: KV cache storage format. A symmetric variant uses one format
23/// for K and V; a `*K*V` variant names the K format, then the V format
24/// (`kv_pair`).
25#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub enum KvCacheDtype {
27    /// 2026-09-25: BF16, 2 bytes per element.
28    Bf16,
29    /// 2026-09-25: FP8 E4M3, 1 byte per element, with one scale per tensor
30    /// (the `k_scale`/`v_scale` arguments of `reshape_and_cache_flash_fp8`).
31    Fp8,
32    /// 2026-09-25: E2M1 values, two per byte, plus one FP8 E4M3 scale byte per
33    /// 16-element group: 4.5 bits per element.
34    Nvfp4,
35    /// 2026-09-25: WHT-rotated 16-level Lloyd-Max codebook (`TURBO4_LUT`) in
36    /// the `Nvfp4` byte layout.
37    Turbo4,
38    /// 2026-09-25: WHT-rotated 8-level Lloyd-Max codebook (`TURBO3_LUT`), 8
39    /// values in 3 bytes plus one FP8 scale byte per group: 3.5 bits per
40    /// element, 22% smaller than `Turbo4`.
41    Turbo3,
42    /// 2026-09-25: WHT-rotated 4-level codebook, 4 values per byte plus one FP8
43    /// scale byte per group: 2.5 bits per element, 6.4x smaller than `Bf16`.
44    /// The server's `auto_high_precision_layers` gives this dtype a larger
45    /// high-precision boundary than the other symmetric turbo dtypes.
46    Turbo2,
47    /// 2026-09-25: WHT-rotated FP8 E4M3 data plus one BF16 scale per group: 9
48    /// bits per element, 12.5% more than `Fp8`.
49    Turbo8,
50    /// 2026-09-25: K `Turbo4`, V `Turbo3`: 8 bits per K+V element pair, against
51    /// 9 for `Turbo4`.
52    Turbo4KTurbo3V,
53    /// 2026-09-25: K `Turbo4`, V `Turbo8`: 13.5 bits per K+V element pair,
54    /// against 18 for `Turbo8`.
55    Turbo4KTurbo8V,
56    /// 2026-09-25: K `Turbo3`, V `Turbo8`.
57    Turbo3KTurbo8V,
58    /// 2026-09-25: K `Bf16`, V `Turbo4`.
59    Bf16KTurbo4V,
60    /// 2026-09-25: K `Bf16`, V `Turbo3`.
61    Bf16KTurbo3V,
62    /// 2026-09-25: K `Fp8`, V `Turbo4`.
63    Fp8KTurbo4V,
64    /// 2026-09-25: K `Fp8`, V `Turbo3`.
65    Fp8KTurbo3V,
66    /// 2026-09-25: K `Bf16`, V `Turbo2`.
67    Bf16KTurbo2V,
68    /// 2026-09-25: K `Fp8`, V `Turbo2`.
69    Fp8KTurbo2V,
70}
71
72impl std::fmt::Display for KvCacheDtype {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        // 2026-09-25: The spelling lives only in `name` (catalog.rs), whose
75        // match has no wildcard arm.
76        f.write_str(self.name())
77    }
78}
79
80impl KvCacheDtype {
81    /// 2026-09-25: The (K, V) formats. A symmetric variant returns
82    /// `(self, self)`.
83    pub fn kv_pair(self) -> (KvCacheDtype, KvCacheDtype) {
84        match self {
85            KvCacheDtype::Turbo4KTurbo3V => (KvCacheDtype::Turbo4, KvCacheDtype::Turbo3),
86            KvCacheDtype::Turbo4KTurbo8V => (KvCacheDtype::Turbo4, KvCacheDtype::Turbo8),
87            KvCacheDtype::Turbo3KTurbo8V => (KvCacheDtype::Turbo3, KvCacheDtype::Turbo8),
88            KvCacheDtype::Bf16KTurbo4V => (KvCacheDtype::Bf16, KvCacheDtype::Turbo4),
89            KvCacheDtype::Bf16KTurbo3V => (KvCacheDtype::Bf16, KvCacheDtype::Turbo3),
90            KvCacheDtype::Fp8KTurbo4V => (KvCacheDtype::Fp8, KvCacheDtype::Turbo4),
91            KvCacheDtype::Fp8KTurbo3V => (KvCacheDtype::Fp8, KvCacheDtype::Turbo3),
92            KvCacheDtype::Bf16KTurbo2V => (KvCacheDtype::Bf16, KvCacheDtype::Turbo2),
93            KvCacheDtype::Fp8KTurbo2V => (KvCacheDtype::Fp8, KvCacheDtype::Turbo2),
94            other => (other, other),
95        }
96    }
97
98    /// 2026-09-25: True for the symmetric turbo formats, whose cache contents
99    /// are in the WHT-rotated basis: `write_kv_cache.rs` runs
100    /// `wht_bf16_inplace` on K and V before the turbo write. The decode
101    /// attention calls it on each side of `kv_pair()`, applying the WHT to Q
102    /// when K is rotated and the inverse WHT to the output when V is rotated.
103    /// A `*K*V` variant itself returns false, so call it on a side.
104    pub fn is_wht_rotated(self) -> bool {
105        matches!(
106            self,
107            KvCacheDtype::Turbo2
108                | KvCacheDtype::Turbo3
109                | KvCacheDtype::Turbo4
110                | KvCacheDtype::Turbo8
111        )
112    }
113
114    /// 2026-09-25: True if K and V use different formats.
115    pub fn is_asymmetric(self) -> bool {
116        matches!(
117            self,
118            KvCacheDtype::Turbo4KTurbo3V
119                | KvCacheDtype::Turbo4KTurbo8V
120                | KvCacheDtype::Turbo3KTurbo8V
121                | KvCacheDtype::Bf16KTurbo4V
122                | KvCacheDtype::Bf16KTurbo3V
123                | KvCacheDtype::Fp8KTurbo4V
124                | KvCacheDtype::Fp8KTurbo3V
125                | KvCacheDtype::Bf16KTurbo2V
126                | KvCacheDtype::Fp8KTurbo2V
127        )
128    }
129}
130
131impl std::str::FromStr for KvCacheDtype {
132    type Err = anyhow::Error;
133    fn from_str(s: &str) -> Result<Self> {
134        match s {
135            "bf16" => Ok(KvCacheDtype::Bf16),
136            "fp8" => Ok(KvCacheDtype::Fp8),
137            "nvfp4" => Ok(KvCacheDtype::Nvfp4),
138            "turbo4" => Ok(KvCacheDtype::Turbo4),
139            "turbo3" => Ok(KvCacheDtype::Turbo3),
140            "turbo2" => Ok(KvCacheDtype::Turbo2),
141            "turbo8" => Ok(KvCacheDtype::Turbo8),
142            "turbo4k_turbo3v" | "turbo4k3v" => Ok(KvCacheDtype::Turbo4KTurbo3V),
143            "turbo4k_turbo8v" | "turbo4k8v" => Ok(KvCacheDtype::Turbo4KTurbo8V),
144            "turbo3k_turbo8v" | "turbo3k8v" => Ok(KvCacheDtype::Turbo3KTurbo8V),
145            "bf16k_turbo4v" | "bf16k4v" => Ok(KvCacheDtype::Bf16KTurbo4V),
146            "bf16k_turbo3v" | "bf16k3v" => Ok(KvCacheDtype::Bf16KTurbo3V),
147            "fp8k_turbo4v" | "fp8k4v" => Ok(KvCacheDtype::Fp8KTurbo4V),
148            "fp8k_turbo3v" | "fp8k3v" => Ok(KvCacheDtype::Fp8KTurbo3V),
149            "bf16k_turbo2v" | "bf16k2v" => Ok(KvCacheDtype::Bf16KTurbo2V),
150            "fp8k_turbo2v" | "fp8k2v" => Ok(KvCacheDtype::Fp8KTurbo2V),
151            other => bail!(
152                "Unsupported --kv-cache-dtype '{other}'. Symmetric: 'bf16', 'fp8', 'nvfp4', 'turbo4', 'turbo3', 'turbo8'. \
153                Asymmetric (TQ+): turbo*_turbo*v, bf16k_turbo[34]v (safer asym: K baseline, V compressed), fp8k_turbo[34]v."
154            ),
155        }
156    }
157}
158
159/// 2026-09-25: Geometry and formats of the paged KV cache.
160pub struct KvCacheConfig {
161    /// 2026-09-25: Tokens per block.
162    pub block_size: usize,
163    /// 2026-09-25: KV heads per layer, where `layer_dims` does not override it.
164    pub num_kv_heads: usize,
165    /// 2026-09-25: Head dimension, where `layer_dims` does not override it.
166    pub head_dim: usize,
167    /// 2026-09-25: Attention layers that get a K pool and a V pool.
168    pub num_layers: usize,
169    /// 2026-09-25: Format of every layer that `layer_dtypes` does not cover.
170    pub dtype: KvCacheDtype,
171    /// 2026-09-25: Per-layer format: layer `i` uses `layer_dtypes[i]` when
172    /// `i < layer_dtypes.len()`, else `dtype`.
173    pub layer_dtypes: Vec<KvCacheDtype>,
174    /// 2026-09-25: Per-layer `(num_kv_heads, head_dim)`, by the same rule as
175    /// `layer_dtypes`, for models whose attention layers differ in shape (the
176    /// Gemma-4 loader's `kv_layer_dims`). Pool sizes and block strides use it.
177    pub layer_dims: Vec<(usize, usize)>,
178    /// 2026-09-25: Per-sequence cap on HBM-resident blocks under
179    /// `--high-speed-swap`. With `Some(N)` the model engine slides a
180    /// sequence's window at `N` blocks (`ensure_blocks_through_decode`);
181    /// `None` means no cap.
182    pub cache_blocks_per_seq: Option<u32>,
183}
184
185impl KvCacheConfig {
186    /// 2026-09-25: The format of attention layer `layer_idx`.
187    pub fn dtype_for_layer(&self, layer_idx: usize) -> KvCacheDtype {
188        if layer_idx < self.layer_dtypes.len() {
189            self.layer_dtypes[layer_idx]
190        } else {
191            self.dtype
192        }
193    }
194
195    /// 2026-09-25: Bytes of one side of one block for `dtype` and
196    /// `(nkv, hd)`. A `*K*V` variant gives its K-side size.
197    fn block_bytes_dims(&self, dtype: KvCacheDtype, nkv: usize, hd: usize) -> usize {
198        let elems = self.block_size * nkv * hd;
199        match dtype {
200            KvCacheDtype::Bf16
201            | KvCacheDtype::Bf16KTurbo4V
202            | KvCacheDtype::Bf16KTurbo3V
203            | KvCacheDtype::Bf16KTurbo2V => elems * 2,
204            KvCacheDtype::Fp8
205            | KvCacheDtype::Fp8KTurbo4V
206            | KvCacheDtype::Fp8KTurbo3V
207            | KvCacheDtype::Fp8KTurbo2V => elems,
208            KvCacheDtype::Nvfp4
209            | KvCacheDtype::Turbo4
210            | KvCacheDtype::Turbo4KTurbo3V
211            | KvCacheDtype::Turbo4KTurbo8V => {
212                // 2026-09-25: Two 4-bit values per byte, then one FP8 scale
213                // byte per group. NVFP4 and Turbo4 differ only in codebook.
214                let data = elems / 2;
215                let num_groups = elems / NVFP4_GROUP_SIZE;
216                data + num_groups
217            }
218            KvCacheDtype::Turbo3 | KvCacheDtype::Turbo3KTurbo8V => {
219                // 2026-09-25: 8 values in 3 bytes, then one FP8 scale byte per
220                // group.
221                let data = elems * 3 / 8;
222                let num_groups = elems / NVFP4_GROUP_SIZE;
223                data + num_groups
224            }
225            KvCacheDtype::Turbo2 => {
226                // 2026-09-25: 4 values per byte, then one FP8 scale byte per
227                // group.
228                let data = elems / 4;
229                let num_groups = elems / NVFP4_GROUP_SIZE;
230                data + num_groups
231            }
232            KvCacheDtype::Turbo8 => {
233                // 2026-09-25: One FP8 byte per element, then one 2-byte BF16
234                // scale per group, where the layouts above use a 1-byte FP8
235                // scale.
236                let num_groups = elems / NVFP4_GROUP_SIZE;
237                elems + num_groups * 2
238            }
239        }
240    }
241
242    /// 2026-09-25: V-side bytes per block: `block_bytes_dims` of
243    /// `dtype.kv_pair().1`.
244    #[allow(dead_code)]
245    pub fn v_block_bytes_dims(&self, dtype: KvCacheDtype, nkv: usize, hd: usize) -> usize {
246        let (_, v) = dtype.kv_pair();
247        self.block_bytes_dims(v, nkv, hd)
248    }
249
250    /// 2026-09-25: K-side bytes per block: `block_bytes_dims` of
251    /// `dtype.kv_pair().0` (for `Bf16KTurbo3V`, the BF16 size).
252    #[allow(dead_code)]
253    pub fn k_block_bytes_dims(&self, dtype: KvCacheDtype, nkv: usize, hd: usize) -> usize {
254        let (k, _) = dtype.kv_pair();
255        self.block_bytes_dims(k, nkv, hd)
256    }
257
258    /// 2026-09-25: K-side bytes per block of attention layer `layer_idx`, from
259    /// its dims and the K side of its format. `PagedKvCache::new` sizes the K
260    /// pool with it.
261    pub fn k_block_bytes_for_layer(&self, layer_idx: usize) -> usize {
262        let (nkv, hd) = self.dims_for_layer(layer_idx);
263        let (k, _) = self.dtype_for_layer(layer_idx).kv_pair();
264        self.block_bytes_dims(k, nkv, hd)
265    }
266
267    /// 2026-09-25: V-side bytes per block of attention layer `layer_idx`.
268    /// Equals `k_block_bytes_for_layer` for a symmetric format.
269    pub fn v_block_bytes_for_layer(&self, layer_idx: usize) -> usize {
270        let (nkv, hd) = self.dims_for_layer(layer_idx);
271        let (_, v) = self.dtype_for_layer(layer_idx).kv_pair();
272        self.block_bytes_dims(v, nkv, hd)
273    }
274
275    /// 2026-09-25: `block_bytes_dims` with the global `num_kv_heads` and
276    /// `head_dim`, ignoring `layer_dims`.
277    fn block_bytes_for_dtype(&self, dtype: KvCacheDtype) -> usize {
278        self.block_bytes_dims(dtype, self.num_kv_heads, self.head_dim)
279    }
280
281    /// 2026-09-25: Bytes of one side of one block with the uniform `dtype` and
282    /// the global dims. For a `*K*V` dtype this is the K side.
283    pub fn block_bytes(&self) -> usize {
284        self.block_bytes_for_dtype(self.dtype)
285    }
286
287    /// 2026-09-25: `(num_kv_heads, head_dim)` of attention layer `layer_idx`:
288    /// `layer_dims[layer_idx]` when present, else the global values.
289    pub fn dims_for_layer(&self, layer_idx: usize) -> (usize, usize) {
290        if layer_idx < self.layer_dims.len() {
291            self.layer_dims[layer_idx]
292        } else {
293            (self.num_kv_heads, self.head_dim)
294        }
295    }
296
297    /// 2026-09-25: Bytes of one side of one block of attention layer
298    /// `layer_idx`, from its own dims and format. For a `*K*V` format this is
299    /// the K side.
300    pub fn block_bytes_for_layer(&self, layer_idx: usize) -> usize {
301        let (nkv, hd) = self.dims_for_layer(layer_idx);
302        self.block_bytes_dims(self.dtype_for_layer(layer_idx), nkv, hd)
303    }
304
305    /// 2026-09-25: Twice `block_bytes()`. That is K + V for a symmetric
306    /// uniform dtype, but twice the K side for a `*K*V` dtype.
307    pub fn block_bytes_kv(&self) -> usize {
308        self.block_bytes() * 2
309    }
310
311    /// 2026-09-25: K plus V bytes of one block index summed over all
312    /// layers, per layer's own dims and K and V formats.
313    pub fn block_bytes_kv_all_layers(&self) -> usize {
314        (0..self.num_layers)
315            .map(|i| self.k_block_bytes_for_layer(i) + self.v_block_bytes_for_layer(i))
316            .sum()
317    }
318
319    /// 2026-09-25: Elements per block and side, from the global dims:
320    /// `block_size * num_kv_heads * head_dim`.
321    pub fn cache_stride_elements(&self) -> usize {
322        self.block_size * self.num_kv_heads * self.head_dim
323    }
324
325    /// 2026-09-25: NVFP4 data bytes per block (two E2M1 values per byte),
326    /// from the global dims.
327    pub fn nvfp4_data_bytes(&self) -> usize {
328        self.block_size * self.num_kv_heads * self.head_dim / 2
329    }
330
331    /// 2026-09-25: NVFP4 scale bytes per block (one FP8 byte per group), from
332    /// the global dims.
333    pub fn nvfp4_scale_bytes(&self) -> usize {
334        self.block_size * self.num_kv_heads * self.head_dim / NVFP4_GROUP_SIZE
335    }
336
337    /// 2026-09-25: Turbo4 data bytes per block, equal to the NVFP4 value.
338    pub fn turbo4_data_bytes(&self) -> usize {
339        self.nvfp4_data_bytes()
340    }
341
342    /// 2026-09-25: Turbo4 scale bytes per block, equal to the NVFP4 value.
343    pub fn turbo4_scale_bytes(&self) -> usize {
344        self.nvfp4_scale_bytes()
345    }
346
347    /// 2026-09-25: Turbo3 data bytes per block (8 values in 3 bytes).
348    pub fn turbo3_data_bytes(&self) -> usize {
349        let elems = self.block_size * self.num_kv_heads * self.head_dim;
350        elems * 3 / 8
351    }
352
353    /// 2026-09-25: Turbo3 scale bytes per block, equal to the NVFP4 value.
354    pub fn turbo3_scale_bytes(&self) -> usize {
355        self.nvfp4_scale_bytes()
356    }
357
358    /// 2026-09-25: Turbo2 data bytes per block (4 values per byte).
359    pub fn turbo2_data_bytes(&self) -> usize {
360        let elems = self.block_size * self.num_kv_heads * self.head_dim;
361        elems / 4
362    }
363
364    /// 2026-09-25: Turbo2 scale bytes per block, equal to the NVFP4 value.
365    pub fn turbo2_scale_bytes(&self) -> usize {
366        self.nvfp4_scale_bytes()
367    }
368
369    /// 2026-09-25: Turbo8 data bytes per block (one FP8 byte per element).
370    pub fn turbo8_data_bytes(&self) -> usize {
371        self.block_size * self.num_kv_heads * self.head_dim
372    }
373
374    /// 2026-09-25: Turbo8 scale bytes per block: one 2-byte BF16 scale per
375    /// group, twice the NVFP4 value.
376    pub fn turbo8_scale_bytes(&self) -> usize {
377        self.nvfp4_scale_bytes() * 2
378    }
379}
380
381/// 2026-09-25: One attention layer's K and V pools.
382struct LayerPool {
383    k_pool: DevicePtr,
384    v_pool: DevicePtr,
385    /// 2026-09-25: Bytes between K blocks; differs from V for a `*K*V` format.
386    k_block_stride: usize,
387    /// 2026-09-25: Bytes between V blocks.
388    v_block_stride: usize,
389    /// 2026-09-25: This layer's format.
390    dtype: KvCacheDtype,
391}
392
393/// 2026-09-25: Paged KV cache across all attention layers: the pools and
394/// the block free list with per-block reference counts.
395pub struct PagedKvCache {
396    layers: Vec<LayerPool>,
397    num_blocks: usize,
398    free_blocks: Vec<u32>,
399    /// 2026-09-25: Per-block reference count, so a block can be shared (the
400    /// prefix cache holds one). Set to 1 on alloc; the block returns to the
401    /// free list when a decrement reaches 0.
402    block_ref_counts: Vec<u32>,
403    config: KvCacheConfig,
404    /// 2026-09-25: Per-block refcount event history (`METRALE_KV_TRACE=1`;
405    /// empty otherwise).
406    trace: block_trace::BlockTrace,
407}
408
409mod block_trace;
410mod catalog;
411mod paged_impl;
412/// 2026-09-25: Release both pools of every layer, and clear the free list and
413/// ref counts with them, so a released cache cannot hand out a block of freed
414/// memory: `alloc_block` then fails. Every pool is freed even after an error;
415/// the first error is returned.
416impl metrale_core::scope::ModelResource<dyn metrale_gpu_runtime::gpu::GpuBackend> for PagedKvCache {
417    fn label(&self) -> &'static str {
418        "kv cache"
419    }
420
421    fn release(&mut self, gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend) -> anyhow::Result<()> {
422        let mut first_error = None;
423        for layer in self.layers.drain(..) {
424            for ptr in [layer.k_pool, layer.v_pool] {
425                if let Err(e) = gpu.free(ptr)
426                    && first_error.is_none()
427                {
428                    first_error = Some(e);
429                }
430            }
431        }
432        self.free_blocks.clear();
433        self.block_ref_counts.clear();
434        self.num_blocks = 0;
435        match first_error {
436            Some(e) => Err(e),
437            None => Ok(()),
438        }
439    }
440}
441
442#[cfg(test)]
443mod tests;
444
445#[cfg(test)]
446mod tests_tq_plus;
metrale_cache/kv_cache.rs

metrale_cache/
kv_cache.rs
