versorc_math/quantization/
turboquant.rs

1// ============================================================================
2// VERSORC-MATH: TURBOQUANT GEOMETRIC QUANTIZER (3-BIT LLOYD-MAX + 1-BIT QJL)
3// License: Apache-2.0
4//
5// Implementation of TurboQuant (Amir Zandieh et al. / Google Research)
6// coupled with Clifford Versor Isometry:
7// Stage 1: (b - 1) bit Lloyd-Max optimal scalar codebook on isotropized coordinates.
8// Stage 2: 1-bit Quantized Johnson-Lindenstrauss (QJL) sketch on the residual vector:
9//          r = R(w) - Q_base(R(w)).
10//
11// Guarantees unbiased inner-product estimation and eliminates quantization bias in attention:
12// <w, x> = <Q_base(w_rot), x_rot> + s_res * sum_i (b_i * x_rot_i).
13// ============================================================================
14
15use super::clifford_3bit::{
16    bf16_to_f32, f32_to_bf16, LLOYD_MAX_3BIT_BOUNDS, LLOYD_MAX_3BIT_CENTROIDS,
17};
18use crate::versor::DyadicWalshVersor;
19
20/// Packed 32-weight TurboQuant block.
21/// - 32 x 3-bit Lloyd-Max base codes: 96 bits (12 bytes)
22/// - 32 x 1-bit QJL residual signs: 32 bits (4 bytes)
23/// - 1 BF16 base scale: 16 bits (2 bytes)
24/// - 1 BF16 residual scale: 16 bits (2 bytes)
25/// Total: 20 bytes (160 bits) for 32 weights.
26#[repr(C)]
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub struct TurboQuantBlock32 {
29    /// 3-bit base codes packed into 12 bytes (96 bits)
30    pub base_codes: [u8; 12],
31    /// 1-bit QJL residual signs packed into 4 bytes (32 bits)
32    pub residual_signs: u32,
33    /// Base standard deviation scale in BF16
34    pub base_scale: u16,
35    /// Residual scale (sqrt(2/pi) * std(r)) in BF16
36    pub residual_scale: u16,
37}
38
39impl Default for TurboQuantBlock32 {
40    fn default() -> Self {
41        Self {
42            base_codes: [0u8; 12],
43            residual_signs: 0,
44            base_scale: 0,
45            residual_scale: 0,
46        }
47    }
48}
49
50impl TurboQuantBlock32 {
51    /// Quantizes 32 float weights using 3-bit Lloyd-Max base + 1-bit QJL residual.
52    pub fn pack(weights: &[f32; 32]) -> Self {
53        // 1. Calculate base scale (standard deviation)
54        let sum_sq: f32 = weights.iter().map(|&w| w * w).sum();
55        let std_dev = (sum_sq / 32.0).sqrt().max(1e-8);
56        let inv_scale = 1.0 / std_dev;
57
58        let mut base_codes = [0u8; 12];
59        let mut residual = [0.0f32; 32];
60
61        // Pack 32 x 3-bit codes into 12 bytes (each byte holds 2 codes + 2/3 of a code)
62        // We pack every group of 8 weights into 3 bytes (8 x 3 = 24 bits = 3 bytes)
63        for g in 0..4 {
64            let mut g_bits: u32 = 0;
65            for i in 0..8 {
66                let idx = g * 8 + i;
67                let w = weights[idx];
68                let norm = w * inv_scale;
69                let sign_bit = if norm < 0.0 { 1u32 } else { 0u32 };
70                let abs_val = norm.abs();
71
72                let mag_code = if abs_val < LLOYD_MAX_3BIT_BOUNDS[0] {
73                    0u32
74                } else if abs_val < LLOYD_MAX_3BIT_BOUNDS[1] {
75                    1u32
76                } else if abs_val < LLOYD_MAX_3BIT_BOUNDS[2] {
77                    2u32
78                } else {
79                    3u32
80                };
81
82                let code_3bit = (sign_bit << 2) | mag_code;
83                g_bits |= code_3bit << (i * 3);
84
85                // Compute exact base reconstruction and residual
86                let recon_sign = if sign_bit == 1 { -1.0f32 } else { 1.0f32 };
87                let recon_base = recon_sign * LLOYD_MAX_3BIT_CENTROIDS[mag_code as usize] * std_dev;
88                residual[idx] = w - recon_base;
89            }
90
91            base_codes[g * 3] = (g_bits & 0xFF) as u8;
92            base_codes[g * 3 + 1] = ((g_bits >> 8) & 0xFF) as u8;
93            base_codes[g * 3 + 2] = ((g_bits >> 16) & 0xFF) as u8;
94        }
95
96        // 2. QJL 1-bit residual signs
97        let mut residual_signs: u32 = 0;
98        for (i, &r) in residual.iter().enumerate() {
99            if r >= 0.0 {
100                residual_signs |= 1u32 << i;
101            }
102        }
103
104        // 3. Optimal Joint 2x2 Gram Projection:
105        // Solves min ||w - (alpha_1 * u_norm + alpha_2 * v_sign)||^2
106        let mut u_norm = [0.0f32; 32];
107        let mut v_sign = [0.0f32; 32];
108        for g in 0..4 {
109            let b0 = base_codes[g * 3] as u32;
110            let b1 = base_codes[g * 3 + 1] as u32;
111            let b2 = base_codes[g * 3 + 2] as u32;
112            let g_bits = b0 | (b1 << 8) | (b2 << 16);
113            for i in 0..8 {
114                let idx = g * 8 + i;
115                let code = ((g_bits >> (i * 3)) & 0x7) as usize;
116                let sign = if (code & 0x4) != 0 { -1.0f32 } else { 1.0f32 };
117                let mag_idx = code & 0x3;
118                u_norm[idx] = sign * LLOYD_MAX_3BIT_CENTROIDS[mag_idx];
119
120                let res_bit = (residual_signs >> idx) & 1;
121                v_sign[idx] = if res_bit == 1 { 1.0f32 } else { -1.0f32 };
122            }
123        }
124
125        let g11: f32 = u_norm.iter().map(|&u| u * u).sum();
126        let g12: f32 = u_norm.iter().zip(v_sign.iter()).map(|(&u, &v)| u * v).sum();
127        let g22: f32 = 32.0f32; // ||v_sign||^2 = sum_1^32 (+-1)^2 = 32
128
129        let rhs1: f32 = weights.iter().zip(u_norm.iter()).map(|(&w, &u)| w * u).sum();
130        let rhs2: f32 = weights.iter().zip(v_sign.iter()).map(|(&w, &v)| w * v).sum();
131
132        let det = g11 * g22 - g12 * g12;
133        let (alpha1, alpha2) = if det.abs() > 1e-12 {
134            let a1 = (rhs1 * g22 - rhs2 * g12) / det;
135            let a2 = (rhs2 * g11 - rhs1 * g12) / det;
136            (a1, a2)
137        } else {
138            (std_dev, std_dev * 0.2)
139        };
140
141        Self {
142            base_codes,
143            residual_signs,
144            base_scale: f32_to_bf16(alpha1),
145            residual_scale: f32_to_bf16(alpha2),
146        }
147    }
148
149    /// Unpacks all 32 weights into output slice.
150    pub fn unpack(&self, out: &mut [f32; 32]) {
151        let b_scale = bf16_to_f32(self.base_scale);
152        let r_scale = bf16_to_f32(self.residual_scale);
153
154        for g in 0..4 {
155            let b0 = self.base_codes[g * 3] as u32;
156            let b1 = self.base_codes[g * 3 + 1] as u32;
157            let b2 = self.base_codes[g * 3 + 2] as u32;
158            let g_bits = b0 | (b1 << 8) | (b2 << 16);
159
160            for i in 0..8 {
161                let idx = g * 8 + i;
162                let code = ((g_bits >> (i * 3)) & 0x7) as usize;
163                let sign = if (code & 0x4) != 0 { -1.0f32 } else { 1.0f32 };
164                let mag_idx = code & 0x3;
165                let base_val = sign * LLOYD_MAX_3BIT_CENTROIDS[mag_idx] * b_scale;
166
167                let res_bit = (self.residual_signs >> idx) & 1;
168                let res_sign = if res_bit == 1 { 1.0f32 } else { -1.0f32 };
169                let res_val = res_sign * r_scale;
170
171                out[idx] = base_val + res_val;
172            }
173        }
174    }
175
176    /// Direct dot product of packed TurboQuant block with activation slice.
177    #[inline(always)]
178    pub fn dot_with_slice(&self, x: &[f32; 32]) -> f32 {
179        let b_scale = bf16_to_f32(self.base_scale);
180        let r_scale = bf16_to_f32(self.residual_scale);
181        let mut total = 0.0f32;
182
183        for g in 0..4 {
184            let b0 = self.base_codes[g * 3] as u32;
185            let b1 = self.base_codes[g * 3 + 1] as u32;
186            let b2 = self.base_codes[g * 3 + 2] as u32;
187            let g_bits = b0 | (b1 << 8) | (b2 << 16);
188
189            for i in 0..8 {
190                let idx = g * 8 + i;
191                let code = ((g_bits >> (i * 3)) & 0x7) as usize;
192                let sign = if (code & 0x4) != 0 { -1.0f32 } else { 1.0f32 };
193                let mag_idx = code & 0x3;
194                let base_val = sign * LLOYD_MAX_3BIT_CENTROIDS[mag_idx] * b_scale;
195
196                let res_bit = (self.residual_signs >> idx) & 1;
197                let res_sign = if res_bit == 1 { 1.0f32 } else { -1.0f32 };
198                let res_val = res_sign * r_scale;
199
200                total += (base_val + res_val) * x[idx];
201            }
202        }
203
204        total
205    }
206}
207
208/// 64-byte Cache-Line Aligned TurboQuant Block holding 64 weights.
209/// Uses 2 x TurboQuantBlock32 (40 bytes total), padded and aligned to 64 bytes.
210#[repr(C, align(64))]
211#[derive(Debug, Clone, Copy, PartialEq, Eq)]
212pub struct TurboQuantCacheLine64 {
213    pub blocks: [TurboQuantBlock32; 2],
214    pub _padding: [u8; 24], // 40 + 24 = 64 bytes
215}
216
217impl Default for TurboQuantCacheLine64 {
218    fn default() -> Self {
219        Self {
220            blocks: [TurboQuantBlock32::default(); 2],
221            _padding: [0u8; 24],
222        }
223    }
224}
225
226impl TurboQuantCacheLine64 {
227    pub const CAPACITY: usize = 64;
228
229    /// Packs 64 float weights into a 64-byte cache line block.
230    pub fn pack_64(weights: &[f32; 64]) -> Self {
231        let b0: &[f32; 32] = weights[0..32].try_into().unwrap();
232        let b1: &[f32; 32] = weights[32..64].try_into().unwrap();
233        Self {
234            blocks: [TurboQuantBlock32::pack(b0), TurboQuantBlock32::pack(b1)],
235            _padding: [0u8; 24],
236        }
237    }
238
239    /// Evaluates direct dot product of this 64-weight block with activation slice.
240    #[inline(always)]
241    pub fn dot_64(&self, x: &[f32; 64]) -> f32 {
242        let x0: &[f32; 32] = x[0..32].try_into().unwrap();
243        let x1: &[f32; 32] = x[32..64].try_into().unwrap();
244        self.blocks[0].dot_with_slice(x0) + self.blocks[1].dot_with_slice(x1)
245    }
246}
247
248/// High-level TurboQuant Tensor with Clifford Dyadic Versor Isometry.
249#[derive(Debug, Clone)]
250pub struct TurboQuantTensor {
251    pub rows: usize,
252    pub cols: usize,
253    pub cache_lines: Vec<TurboQuantCacheLine64>,
254}
255
256impl TurboQuantTensor {
257    /// Compiles an unquantized FP32 weight matrix into the AOT rotated TurboQuant representation.
258    pub fn compile_aot(weights: &[f32], rows: usize, cols: usize, versor: &DyadicWalshVersor) -> Self {
259        assert_eq!(weights.len(), rows * cols);
260        assert_eq!(cols % 64, 0, "Columns must be a multiple of 64 for cache-line alignment");
261        assert_eq!(versor.dimension(), cols);
262
263        let lines_per_row = cols / 64;
264        let mut cache_lines = Vec::with_capacity(rows * lines_per_row);
265        let mut row_buffer = vec![0.0f32; cols];
266
267        for r in 0..rows {
268            let row_slice = &weights[r * cols..(r + 1) * cols];
269            row_buffer.copy_from_slice(row_slice);
270
271            // Isotropic rotation: R(w)
272            versor.apply_in_place(&mut row_buffer);
273
274            for l in 0..lines_per_row {
275                let chunk: &[f32; 64] = row_buffer[l * 64..(l + 1) * 64].try_into().unwrap();
276                cache_lines.push(TurboQuantCacheLine64::pack_64(chunk));
277            }
278        }
279
280        Self {
281            rows,
282            cols,
283            cache_lines,
284        }
285    }
286
287    /// Evaluates Matrix-Vector Product y = W * x using Dual Isometry with TurboQuant unbiased projection.
288    pub fn matvec(&self, x: &[f32], out: &mut [f32], versor: &DyadicWalshVersor) {
289        assert_eq!(x.len(), self.cols);
290        assert_eq!(out.len(), self.rows);
291
292        // 1. Single activation rotation: x_rot = R(x)
293        let mut x_rot = x.to_vec();
294        versor.apply_in_place(&mut x_rot);
295
296        // 2. Direct streaming dot products
297        let lines_per_row = self.cols / 64;
298        for r in 0..self.rows {
299            let mut row_dot = 0.0f32;
300            let line_offset = r * lines_per_row;
301            for l in 0..lines_per_row {
302                let cache_line = &self.cache_lines[line_offset + l];
303                let chunk_x: &[f32; 64] = x_rot[l * 64..(l + 1) * 64].try_into().unwrap();
304                row_dot += cache_line.dot_64(chunk_x);
305            }
306            out[r] = row_dot;
307        }
308    }
309}
versorc_math/quantization/turboquant.rs

versorc_math/quantization/
turboquant.rs
