rlx_runtime/
kv_cache.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Per-layer K/V cache for autoregressive decode (Whisper, Qwen, Gemma, …).
17
18use crate::compile_cache::pad_rows;
19
20/// Layer-wise past K/V tensors in row-major `[past_len * kv_dim]` layout per layer.
21#[derive(Debug, Clone)]
22pub struct LayerKvCache {
23    pub past_len: usize,
24    pub layers_k: Vec<Vec<f32>>,
25    pub layers_v: Vec<Vec<f32>>,
26    /// Absolute token index of row `0` in each layer's K/V buffer after
27    /// sliding-window trim (Gemma 3/4 ISWA). Zero when the buffer holds
28    /// the full prefix from position `0`.
29    pub layers_kv_base: Vec<usize>,
30}
31
32impl LayerKvCache {
33    pub fn from_layer_outputs(
34        num_layers: usize,
35        batch: usize,
36        past_seq: usize,
37        kv_dim: usize,
38        outputs: &[Vec<f32>],
39    ) -> Result<Self, String> {
40        let dims: Vec<usize> = vec![kv_dim; num_layers];
41        Self::from_layer_outputs_per_layer(num_layers, batch, past_seq, &dims, outputs)
42    }
43
44    /// Like [`Self::from_layer_outputs`] but accepts a per-layer
45    /// `kv_dim` vector. Gemma 4 12B's full-attention layers have
46    /// `kv_dim = 1 * 512 = 512` while sliding layers have `8 * 256 =
47    /// 2048`; this constructor handles that heterogeneity.
48    pub fn from_layer_outputs_per_layer(
49        num_layers: usize,
50        batch: usize,
51        past_seq: usize,
52        kv_dims: &[usize],
53        outputs: &[Vec<f32>],
54    ) -> Result<Self, String> {
55        if outputs.len() != 2 * num_layers {
56            return Err(format!(
57                "from_layer_outputs_per_layer: expected {} K/V tensors, got {}",
58                2 * num_layers,
59                outputs.len()
60            ));
61        }
62        if kv_dims.len() != num_layers {
63            return Err(format!(
64                "from_layer_outputs_per_layer: expected {} kv_dims, got {}",
65                num_layers,
66                kv_dims.len()
67            ));
68        }
69        let mut layers_k = Vec::with_capacity(num_layers);
70        let mut layers_v = Vec::with_capacity(num_layers);
71        for layer in 0..num_layers {
72            let kv_dim = kv_dims[layer];
73            let expected = batch * past_seq * kv_dim;
74            let k = &outputs[2 * layer];
75            let v = &outputs[2 * layer + 1];
76            if k.len() != expected || v.len() != expected {
77                return Err(format!(
78                    "layer {layer}: k.len={} v.len={} expected {expected} (kv_dim={kv_dim})",
79                    k.len(),
80                    v.len()
81                ));
82            }
83            layers_k.push(k.clone());
84            layers_v.push(v.clone());
85        }
86        Ok(Self {
87            past_len: past_seq,
88            layers_k,
89            layers_v,
90            layers_kv_base: vec![0; num_layers],
91        })
92    }
93
94    /// Pad each layer's K/V to `upper` rows along the sequence axis (`kv_dim` inner).
95    pub fn pad_layers_to_upper(&self, upper: u64, kv_dim: usize) -> (Vec<Vec<f32>>, Vec<Vec<f32>>) {
96        let dims: Vec<usize> = vec![kv_dim; self.layers_k.len()];
97        self.pad_layers_to_upper_per_layer(upper, &dims)
98    }
99
100    /// Like [`Self::pad_layers_to_upper`] but pads each layer to its
101    /// own `kv_dim`. The number of dims must equal the number of
102    /// cached layers.
103    pub fn pad_layers_to_upper_per_layer(
104        &self,
105        upper: u64,
106        kv_dims: &[usize],
107    ) -> (Vec<Vec<f32>>, Vec<Vec<f32>>) {
108        assert_eq!(
109            kv_dims.len(),
110            self.layers_k.len(),
111            "pad_layers_to_upper_per_layer: kv_dims len {} != layers {}",
112            kv_dims.len(),
113            self.layers_k.len(),
114        );
115        let padded_k = self
116            .layers_k
117            .iter()
118            .zip(kv_dims.iter())
119            .map(|(k, &d)| pad_rows(k, d, upper))
120            .collect();
121        let padded_v = self
122            .layers_v
123            .iter()
124            .zip(kv_dims.iter())
125            .map(|(v, &d)| pad_rows(v, d, upper))
126            .collect();
127        (padded_k, padded_v)
128    }
129
130    /// Update cache from decode outputs: `[logits, k0, v0, k1, v1, …]` (bucket-padded).
131    pub fn advance_from_decode_outputs(
132        &mut self,
133        outputs: Vec<Vec<f32>>,
134        batch: usize,
135        kv_dim: usize,
136    ) -> Result<(), String> {
137        let dims: Vec<usize> = vec![kv_dim; self.layers_k.len()];
138        self.advance_from_decode_outputs_per_layer(outputs, batch, &dims)
139    }
140
141    /// Trim each layer's K/V history to at most `window` rows on
142    /// the sequence axis, keeping the most recent rows. Used by
143    /// Gemma 3/4 sliding-attention layers — long contexts can keep
144    /// only the last `window` (e.g. 1024) tokens per sliding layer
145    /// without affecting attention semantics (those layers mask out
146    /// older positions anyway).
147    ///
148    /// `kv_dims_keep` selects which layers to trim and at what dim:
149    /// `kv_dims_keep[i] = Some((dim, window))` trims layer `i`,
150    /// `None` leaves the layer untouched. Pass-through for layers
151    /// whose attention is full-causal.
152    ///
153    /// Note: `past_len` is unchanged — the per-layer K/V buffers
154    /// just hold fewer real rows now; the decode flow's per-layer
155    /// `past_k_{i}` input shape will see the trimmed length. Caller
156    /// is responsible for ensuring the graph's declared `past_seq`
157    /// matches the trimmed length OR the trimmed layer is bound
158    /// dynamically.
159    pub fn trim_sliding_window_per_layer(
160        &mut self,
161        kv_dims_keep: &[Option<(usize, usize)>],
162    ) -> Result<(), String> {
163        if kv_dims_keep.len() != self.layers_k.len() {
164            return Err(format!(
165                "trim_sliding_window_per_layer: kv_dims_keep len {} != layers {}",
166                kv_dims_keep.len(),
167                self.layers_k.len(),
168            ));
169        }
170        for (i, spec) in kv_dims_keep.iter().enumerate() {
171            let Some((kv_dim, window)) = spec else {
172                continue;
173            };
174            let kv_dim = *kv_dim;
175            let window = *window;
176            if window == 0 || kv_dim == 0 {
177                continue;
178            }
179            let rows = self.layers_k[i].len() / kv_dim;
180            if rows <= window {
181                continue;
182            }
183            let drop_rows = rows - window;
184            let drop_bytes = drop_rows * kv_dim;
185            self.layers_k[i].drain(..drop_bytes);
186            self.layers_v[i].drain(..drop_bytes);
187            if self.layers_kv_base.len() <= i {
188                self.layers_kv_base.resize(self.layers_k.len(), 0);
189            }
190            self.layers_kv_base[i] = self.layers_kv_base[i].saturating_add(drop_rows);
191        }
192        Ok(())
193    }
194
195    /// Per-layer variant of [`Self::advance_from_decode_outputs`].
196    pub fn advance_from_decode_outputs_per_layer(
197        &mut self,
198        outputs: Vec<Vec<f32>>,
199        _batch: usize,
200        kv_dims: &[usize],
201    ) -> Result<(), String> {
202        let n = self.layers_k.len();
203        if outputs.len() != 1 + 2 * n {
204            return Err(format!(
205                "advance_from_decode_outputs_per_layer: expected {} outputs, got {}",
206                1 + 2 * n,
207                outputs.len()
208            ));
209        }
210        if kv_dims.len() != n {
211            return Err(format!(
212                "advance_from_decode_outputs_per_layer: kv_dims len {} != layers {n}",
213                kv_dims.len()
214            ));
215        }
216        let new_len = self.past_len + 1;
217        let mut iter = outputs.into_iter();
218        let _logits = iter.next().ok_or("missing logits")?;
219        for i in 0..n {
220            let k = iter.next().ok_or("missing k")?;
221            let v = iter.next().ok_or("missing v")?;
222            let real_len = new_len * kv_dims[i];
223            self.layers_k[i] = k[..real_len.min(k.len())].to_vec();
224            self.layers_v[i] = v[..real_len.min(v.len())].to_vec();
225        }
226        self.past_len = new_len;
227        Ok(())
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn sliding_window_trim_keeps_last_w_rows() {
237        // 3 layers, each storing 6 rows of kv_dim=4 = 24 floats.
238        let kv_dim = 4;
239        let rows = 6;
240        let mut cache = LayerKvCache {
241            past_len: rows,
242            layers_k: vec![(0..(rows * kv_dim)).map(|x| x as f32).collect(); 3],
243            layers_v: vec![(0..(rows * kv_dim)).map(|x| x as f32).collect(); 3],
244            layers_kv_base: vec![0; 3],
245        };
246        // Trim layer 0 to last 2 rows; layer 1 untouched; layer 2 to last 4.
247        let spec = [Some((kv_dim, 2)), None, Some((kv_dim, 4))];
248        cache.trim_sliding_window_per_layer(&spec).unwrap();
249        assert_eq!(cache.layers_k[0].len(), 2 * kv_dim);
250        assert_eq!(cache.layers_kv_base[0], 4);
251        assert_eq!(cache.layers_kv_base[1], 0);
252        assert_eq!(cache.layers_kv_base[2], 2);
253        // Layer 0 should now hold the LAST 2 rows: rows 4 and 5.
254        assert_eq!(
255            cache.layers_k[0],
256            vec![16., 17., 18., 19., 20., 21., 22., 23.]
257        );
258        assert_eq!(
259            cache.layers_k[1].len(),
260            6 * kv_dim,
261            "untouched layer keeps full history"
262        );
263        assert_eq!(cache.layers_k[2].len(), 4 * kv_dim);
264    }
265
266    #[test]
267    fn sliding_window_trim_no_op_when_under_window() {
268        let kv_dim = 4;
269        let rows = 3;
270        let mut cache = LayerKvCache {
271            past_len: rows,
272            layers_k: vec![vec![1.0f32; rows * kv_dim]],
273            layers_v: vec![vec![2.0f32; rows * kv_dim]],
274            layers_kv_base: vec![0],
275        };
276        cache
277            .trim_sliding_window_per_layer(&[Some((kv_dim, 10))])
278            .unwrap();
279        assert_eq!(cache.layers_k[0].len(), rows * kv_dim);
280    }
281}
rlx_runtime/kv_cache.rs

rlx_runtime/
kv_cache.rs
