ferrox_core/
turboquant.rs

1//! Walsh–Hadamard transform helpers for TurboQuant-style KV compression.
2//!
3//! Host FWHT + group quant; Metal `FERROX_CTK=turbo4` stores the 4-bit
4//! groups (WHT optional on host upload). turbo8 aliases ggml Q8_0 on Metal.
5//! Algorithm follows the public TurboQuant line (randomized Hadamard +
6//! per-group scale).
7//!
8//! In-place FWHT on `x` of length `n = 2^k`. Output is unnormalized
9//! (each butterfly is `a+b`, `a-b`); callers that need orthonormal
10//! Hadamard divide by `sqrt(n)`.
11
12/// In-place Fast Walsh–Hadamard Transform. `x.len()` must be a power of two.
13pub fn fwht_inplace(x: &mut [f32]) {
14    let n = x.len();
15    assert!(
16        n.is_power_of_two() && n > 0,
17        "fwht length must be 2^k, got {n}"
18    );
19    let mut h = 1usize;
20    while h < n {
21        let step = h * 2;
22        for i in (0..n).step_by(step) {
23            for j in i..i + h {
24                let a = x[j];
25                let b = x[j + h];
26                x[j] = a + b;
27                x[j + h] = a - b;
28            }
29        }
30        h = step;
31    }
32}
33
34/// Orthonormal FWHT: `fwht_inplace` then scale by `1/sqrt(n)`.
35pub fn fwht_orthonormal_inplace(x: &mut [f32]) {
36    let n = x.len();
37    fwht_inplace(x);
38    let inv = 1.0 / (n as f32).sqrt();
39    for v in x.iter_mut() {
40        *v *= inv;
41    }
42}
43
44/// Per-group absmax quantize after optional WHT (turbo4-style sketch).
45///
46/// Packs each group of `group` floats into `group/2` bytes (two 4-bit
47/// codes per byte) plus one f32 scale. Returns `(packed, scales)`.
48pub fn quantize_turbo4_groups(x: &[f32], group: usize) -> (Vec<u8>, Vec<f32>) {
49    assert!(group >= 2 && group.is_power_of_two());
50    assert_eq!(x.len() % group, 0);
51    let n_groups = x.len() / group;
52    let mut packed = Vec::with_capacity(n_groups * (group / 2));
53    let mut scales = Vec::with_capacity(n_groups);
54    for g in 0..n_groups {
55        let chunk = &x[g * group..(g + 1) * group];
56        let mut amax = 0.0f32;
57        for &v in chunk {
58            amax = amax.max(v.abs());
59        }
60        let scale = if amax > 0.0 { amax / 7.0 } else { 1.0 };
61        scales.push(scale);
62        let inv = 1.0 / scale;
63        for pair in chunk.as_chunks::<2>().0 {
64            let q0 = (pair[0] * inv).round().clamp(-8.0, 7.0) as i8;
65            let q1 = (pair[1] * inv).round().clamp(-8.0, 7.0) as i8;
66            let n0 = (q0 as u8) & 0x0f;
67            let n1 = (q1 as u8) & 0x0f;
68            packed.push(n0 | (n1 << 4));
69        }
70    }
71    (packed, scales)
72}
73
74/// Inverse of [`quantize_turbo4_groups`] (ignores WHT; dequant only).
75pub fn dequantize_turbo4_groups(packed: &[u8], scales: &[f32], group: usize) -> Vec<f32> {
76    assert!(group >= 2 && group.is_power_of_two());
77    let n_groups = scales.len();
78    assert_eq!(packed.len(), n_groups * (group / 2));
79    let mut out = Vec::with_capacity(n_groups * group);
80    for (g, &scale) in scales.iter().enumerate() {
81        let bytes = &packed[g * (group / 2)..(g + 1) * (group / 2)];
82        for &b in bytes {
83            let q0 = ((b & 0x0f) as i8) << 4 >> 4; // sign-extend 4-bit
84            let q1 = ((b >> 4) as i8) << 4 >> 4;
85            out.push(q0 as f32 * scale);
86            out.push(q1 as f32 * scale);
87        }
88    }
89    out
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn fwht_involutory_up_to_scale() {
98        let mut x = vec![1.0, 2.0, 3.0, 4.0, -1.0, 0.5, 0.25, -0.5];
99        let orig = x.clone();
100        fwht_inplace(&mut x);
101        fwht_inplace(&mut x);
102        // Unnormalized FWHT twice → n * identity.
103        let n = orig.len() as f32;
104        for (a, b) in orig.iter().zip(x.iter()) {
105            assert!((a * n - b).abs() < 1e-4, "{a} vs {b}");
106        }
107    }
108
109    #[test]
110    fn orthonormal_fwht_is_involution() {
111        let mut x = vec![0.1, -0.2, 0.3, -0.4];
112        let orig = x.clone();
113        fwht_orthonormal_inplace(&mut x);
114        fwht_orthonormal_inplace(&mut x);
115        for (a, b) in orig.iter().zip(x.iter()) {
116            assert!((a - b).abs() < 1e-5, "{a} vs {b}");
117        }
118    }
119
120    #[test]
121    fn turbo4_roundtrip_reasonable() {
122        let mut x: Vec<f32> = (0..16).map(|i| (i as f32 * 0.3).sin()).collect();
123        fwht_orthonormal_inplace(&mut x);
124        let (packed, scales) = quantize_turbo4_groups(&x, 8);
125        let mut y = dequantize_turbo4_groups(&packed, &scales, 8);
126        fwht_orthonormal_inplace(&mut y);
127        // After inverse WHT, compare to pre-WHT... we mutated x in place.
128        // Rebuild original for error check.
129        let orig: Vec<f32> = (0..16).map(|i| (i as f32 * 0.3).sin()).collect();
130        let mut err = 0.0f32;
131        for (a, b) in orig.iter().zip(y.iter()) {
132            err += (a - b).abs();
133        }
134        err /= orig.len() as f32;
135        assert!(err < 0.15, "mean abs err {err}");
136    }
137}
ferrox_core/turboquant.rs

ferrox_core/
turboquant.rs
