rlx_cpu/
dequant_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//! Cache dequantized GGUF weight bytes for static params.
16//!
17//! Qwen3.5 decode with `--packed` was re-dequantizing every K-quant
18//! weight on every matmul (hundreds of times per token). Keys are
19//! `(k, n, scheme, bytes_hash)` — stable for identical GGUF bytes regardless
20//! of arena offset (multiple compiled graphs reuse offsets).
21
22use rlx_ir::quant::QuantScheme;
23use std::collections::HashMap;
24use std::sync::{Arc, OnceLock, RwLock};
25
26#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
27struct DequantKey {
28    k: u32,
29    n: u32,
30    scheme: u8,
31    /// Content hash — `w_off` alone collides across compiled graphs.
32    bytes_hash: u64,
33}
34
35fn weight_bytes_hash(w_bytes: &[u8]) -> u64 {
36    use std::hash::{Hash, Hasher};
37    let mut hasher = std::collections::hash_map::DefaultHasher::new();
38    w_bytes.hash(&mut hasher);
39    hasher.finish()
40}
41
42fn scheme_tag(scheme: QuantScheme) -> u8 {
43    match scheme {
44        QuantScheme::GgufQ4K => 1,
45        QuantScheme::GgufQ5K => 2,
46        QuantScheme::GgufQ6K => 3,
47        QuantScheme::GgufQ8K => 4,
48        QuantScheme::GgufQ4_0 => 5,
49        QuantScheme::GgufQ8_0 => 6,
50        QuantScheme::GgufQ2K => 7,
51        QuantScheme::GgufQ3K => 8,
52        QuantScheme::GgufIQ4NL => 9,
53        QuantScheme::GgufIQ4XS => 10,
54        QuantScheme::GgufIQ2XXS => 11,
55        QuantScheme::GgufIQ2XS => 12,
56        QuantScheme::GgufIQ2S => 13,
57        QuantScheme::GgufIQ3XXS => 14,
58        QuantScheme::GgufIQ3S => 15,
59        QuantScheme::GgufIQ1S => 16,
60        QuantScheme::GgufIQ1M => 17,
61        QuantScheme::GgufTQ1_0 => 18,
62        QuantScheme::GgufTQ2_0 => 19,
63        QuantScheme::GgufMXFP4 => 20,
64        QuantScheme::GgufNVFP4 => 21,
65        QuantScheme::GgufQ4_1 => 22,
66        QuantScheme::GgufQ5_0 => 23,
67        QuantScheme::GgufQ5_1 => 24,
68        QuantScheme::GgufQ1_0 => 25,
69        QuantScheme::GgufQ2_0 => 26,
70        _ => 255,
71    }
72}
73
74fn dequant_gguf(w_bytes: &[u8], k: usize, n: usize, scheme: QuantScheme) -> Vec<f32> {
75    let n_elems = k * n;
76    match scheme {
77        QuantScheme::GgufQ4K => rlx_gguf::dequant_q4_k(w_bytes, n_elems),
78        QuantScheme::GgufQ5K => rlx_gguf::dequant_q5_k(w_bytes, n_elems),
79        QuantScheme::GgufQ6K => rlx_gguf::dequant_q6_k(w_bytes, n_elems),
80        QuantScheme::GgufQ8K => rlx_gguf::dequant_q8_k(w_bytes, n_elems),
81        QuantScheme::GgufQ2K => rlx_gguf::dequant_q2_k(w_bytes, n_elems),
82        QuantScheme::GgufQ3K => rlx_gguf::dequant_q3_k(w_bytes, n_elems),
83        QuantScheme::GgufQ4_0 => rlx_gguf::dequant_q4_0(w_bytes, n_elems),
84        QuantScheme::GgufQ4_1 => rlx_gguf::dequant_q4_1(w_bytes, n_elems),
85        QuantScheme::GgufQ5_0 => rlx_gguf::dequant_q5_0(w_bytes, n_elems),
86        QuantScheme::GgufQ5_1 => rlx_gguf::dequant_q5_1(w_bytes, n_elems),
87        QuantScheme::GgufQ8_0 => rlx_gguf::dequant_q8_0(w_bytes, n_elems),
88        QuantScheme::GgufIQ4NL => rlx_gguf::iq_dequant::dequant_iq4_nl(w_bytes, n_elems),
89        QuantScheme::GgufIQ4XS => rlx_gguf::iq_dequant::dequant_iq4_xs(w_bytes, n_elems),
90        QuantScheme::GgufIQ2XXS => rlx_gguf::iq_dequant::dequant_iq2_xxs(w_bytes, n_elems),
91        QuantScheme::GgufIQ2XS => rlx_gguf::iq_dequant::dequant_iq2_xs(w_bytes, n_elems),
92        QuantScheme::GgufIQ2S => rlx_gguf::iq_dequant::dequant_iq2_s(w_bytes, n_elems),
93        QuantScheme::GgufIQ3XXS => rlx_gguf::iq_dequant::dequant_iq3_xxs(w_bytes, n_elems),
94        QuantScheme::GgufIQ3S => rlx_gguf::iq_dequant::dequant_iq3_s(w_bytes, n_elems),
95        QuantScheme::GgufIQ1S => rlx_gguf::iq_dequant::dequant_iq1_s(w_bytes, n_elems),
96        QuantScheme::GgufIQ1M => rlx_gguf::iq_dequant::dequant_iq1_m(w_bytes, n_elems),
97        QuantScheme::GgufTQ1_0 => rlx_gguf::tq_dequant::dequant_tq1_0(w_bytes, n_elems),
98        QuantScheme::GgufTQ2_0 => rlx_gguf::tq_dequant::dequant_tq2_0(w_bytes, n_elems),
99        QuantScheme::GgufMXFP4 => rlx_gguf::mx_dequant::dequant_mxfp4(w_bytes, n_elems),
100        QuantScheme::GgufNVFP4 => rlx_gguf::mx_dequant::dequant_nvfp4(w_bytes, n_elems),
101        QuantScheme::GgufQ1_0 => rlx_gguf::q1_dequant::dequant_q1_0(w_bytes, n_elems),
102        QuantScheme::GgufQ2_0 => rlx_gguf::q2_dequant::dequant_q2_0(w_bytes, n_elems),
103        other => panic!("dequant_cache: unsupported GGUF scheme {other:?}"),
104    }
105    .expect("GGUF dequant failed")
106}
107
108static CACHE: OnceLock<RwLock<HashMap<DequantKey, Arc<[f32]>>>> = OnceLock::new();
109
110fn cache_enabled() -> bool {
111    !matches!(
112        rlx_ir::env::var("RLX_DEQUANT_CACHE").as_deref(),
113        Some("0") | Some("false") | Some("off")
114    )
115}
116
117/// Return dense `[k×n]` weights (GGUF row-major `[n,k]` layout) for `w_bytes`.
118pub fn gguf_weight_f32(
119    _w_off: usize,
120    w_bytes: &[u8],
121    k: usize,
122    n: usize,
123    scheme: QuantScheme,
124) -> Arc<[f32]> {
125    if !cache_enabled() {
126        return Arc::from(dequant_gguf(w_bytes, k, n, scheme).into_boxed_slice());
127    }
128    let key = DequantKey {
129        k: k as u32,
130        n: n as u32,
131        scheme: scheme_tag(scheme),
132        bytes_hash: weight_bytes_hash(w_bytes),
133    };
134    let cache = CACHE.get_or_init(|| RwLock::new(HashMap::new()));
135    if let Some(hit) = cache.read().expect("dequant cache poisoned").get(&key) {
136        return Arc::clone(hit);
137    }
138    let dense = dequant_gguf(w_bytes, k, n, scheme);
139    let arc: Arc<[f32]> = Arc::from(dense.into_boxed_slice());
140    cache
141        .write()
142        .expect("dequant cache poisoned")
143        .insert(key, Arc::clone(&arc));
144    arc
145}
146
147/// Drop cached dequantized weights (e.g. between model loads in tests).
148pub fn clear_dequant_cache() {
149    if let Some(c) = CACHE.get() {
150        c.write().expect("dequant cache poisoned").clear();
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn gguf_dequant_cache_hits_on_second_lookup() {
160        clear_dequant_cache();
161        const QK_K: usize = 256;
162        let mut packed = Vec::new();
163        packed.extend_from_slice(&half::f16::from_f32(1.0).to_le_bytes());
164        packed.extend_from_slice(&half::f16::from_f32(1.0).to_le_bytes());
165        let mut scales = [0u8; 12];
166        for s in &mut scales[0..4] {
167            *s = 0x01;
168        }
169        packed.extend_from_slice(&scales);
170        packed.extend(std::iter::repeat_n(0x77u8, QK_K / 2));
171        let k = 256;
172        let n = 1;
173        let w_off = 4096;
174        let hash = weight_bytes_hash(&packed);
175        let a = gguf_weight_f32(w_off, &packed, k, n, QuantScheme::GgufQ4K);
176        let b = gguf_weight_f32(w_off + 999, &packed, k, n, QuantScheme::GgufQ4K);
177        assert!(
178            Arc::ptr_eq(&a, &b),
179            "same bytes at different offsets should hit"
180        );
181        let mut other = packed.clone();
182        other[0] ^= 0x01;
183        let c = gguf_weight_f32(w_off, &other, k, n, QuantScheme::GgufQ4K);
184        assert!(!Arc::ptr_eq(&a, &c), "different bytes should miss: {hash}");
185    }
186
187    #[test]
188    fn gguf_dequant_cache_q4_1_q5_hits() {
189        clear_dequant_cache();
190        let w: Vec<f32> = (0..512).map(|i| (i as f32 * 0.01).sin()).collect();
191        for (scheme, ggml) in [
192            (QuantScheme::GgufQ4_1, rlx_gguf::GgmlType::Q4_1),
193            (QuantScheme::GgufQ5_0, rlx_gguf::GgmlType::Q5_0),
194            (QuantScheme::GgufQ5_1, rlx_gguf::GgmlType::Q5_1),
195        ] {
196            let packed = rlx_gguf::quantize(&w, ggml).unwrap();
197            let a = gguf_weight_f32(0, &packed, 32, 16, scheme);
198            let b = gguf_weight_f32(8192, &packed, 32, 16, scheme);
199            assert!(Arc::ptr_eq(&a, &b), "cache hit expected for {scheme:?}");
200        }
201    }
202}
rlx_cpu/dequant_cache.rs

rlx_cpu/
dequant_cache.rs
