fuel_quantized/
k_quants.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// Derived from candle (https://github.com/huggingface/candle, MIT OR Apache-2.0): `candle-core/src/quantized/k_quants.rs`
3// at commit 549eacb011f25b2c6fc765d6c3dbd26e42f72ee6, modified here. candle's code is itself a port of ggml / llama.cpp (MIT), which its own comments cite by commit.
4// Contributors, how this was verified and the licence texts: THIRD-PARTY-NOTICES.md (section "candle").
5use crate::utils::{
6    get_scale_min_k4, group_for_dequantization, group_for_quantization, make_q3_quants,
7    make_qkx1_quants, make_qkx3_quants, make_qp_quants, make_qx_quants, nearest_int,
8};
9use byteorder::{ByteOrder, LittleEndian};
10use fuel_ir::Result;
11use fuel_ir::quantized::GgmlDType;
12use half::{bf16, f16, slice::HalfFloatSliceExt};
13use rayon::prelude::*;
14
15// Default to QK_K 256 rather than 64.
16pub const QK_K: usize = 256;
17pub const K_SCALE_SIZE: usize = 12;
18
19pub const QK4_0: usize = 32;
20pub const QK4_1: usize = 32;
21pub const QK5_0: usize = 32;
22pub const QK5_1: usize = 32;
23pub const QK8_0: usize = 32;
24pub const QK8_1: usize = 32;
25
26/// A GGML block type (or a dense scalar standing in for one).
27///
28/// `bytemuck::Pod` is a supertrait on purpose: every implementor is
29/// `#[repr(C)]` plain data with no padding and no invalid bit pattern, so a
30/// byte slice of the right length and alignment IS a valid `[Self]`, and
31/// all-zero bytes ARE a valid `Self`. Making that a bound the compiler checks
32/// (the derive refuses a type with padding) replaces the hand-written
33/// `unsafe` casts that each asserted it in a comment.
34pub trait GgmlType: Sized + Clone + Send + Sync + bytemuck::Pod {
35    const DTYPE: GgmlDType;
36    const BLCK_SIZE: usize;
37
38    /// The all-zero value. Safe by construction: `Self: Pod` implies `Zeroable`.
39    fn zeros() -> Self {
40        <Self as bytemuck::Zeroable>::zeroed()
41    }
42    fn to_float(xs: &[Self], ys: &mut [f32]);
43    fn from_float(xs: &[f32], ys: &mut [Self]);
44    fn from_float_imatrix(
45        _xs: &[f32],
46        _ys: &mut [Self],
47        _imatrix_weights: &[f32],
48        _n_per_row: usize,
49    ) {
50        panic!(
51            "`from_float_imatrix` is unimplemented for {:?}",
52            Self::DTYPE
53        );
54    }
55}
56
57#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
58#[repr(C)]
59pub struct BlockQ4_0 {
60    pub(crate) d: f16,
61    pub(crate) qs: [u8; QK4_0 / 2],
62}
63const _: () = assert!(std::mem::size_of::<BlockQ4_0>() == 18);
64
65#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
66#[repr(C)]
67pub struct BlockQ4_1 {
68    pub(crate) d: f16,
69    pub(crate) m: f16,
70    pub(crate) qs: [u8; QK4_1 / 2],
71}
72const _: () = assert!(std::mem::size_of::<BlockQ4_1>() == 20);
73
74#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
75#[repr(C)]
76pub struct BlockQ5_0 {
77    pub(crate) d: f16,
78    pub(crate) qh: [u8; 4],
79    pub(crate) qs: [u8; QK5_0 / 2],
80}
81const _: () = assert!(std::mem::size_of::<BlockQ5_0>() == 22);
82
83#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
84#[repr(C)]
85pub struct BlockQ5_1 {
86    pub(crate) d: f16,
87    pub(crate) m: f16,
88    pub(crate) qh: [u8; 4],
89    pub(crate) qs: [u8; QK5_1 / 2],
90}
91const _: () = assert!(std::mem::size_of::<BlockQ5_1>() == 24);
92
93#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
94#[repr(C)]
95pub struct BlockQ8_0 {
96    pub(crate) d: f16,
97    pub(crate) qs: [i8; QK8_0],
98}
99const _: () = assert!(std::mem::size_of::<BlockQ8_0>() == 34);
100
101#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
102#[repr(C)]
103pub struct BlockQ8_1 {
104    pub(crate) d: f16,
105    pub(crate) s: f16,
106    pub(crate) qs: [i8; QK8_1],
107}
108const _: () = assert!(std::mem::size_of::<BlockQ8_1>() == 36);
109
110#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
111#[repr(C)]
112pub struct BlockQ2K {
113    pub(crate) scales: [u8; QK_K / 16],
114    pub(crate) qs: [u8; QK_K / 4],
115    pub(crate) d: f16,
116    pub(crate) dmin: f16,
117}
118const _: () = assert!(QK_K / 16 + QK_K / 4 + 2 * 2 == std::mem::size_of::<BlockQ2K>());
119
120#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
121#[repr(C)]
122pub struct BlockQ3K {
123    pub(crate) hmask: [u8; QK_K / 8],
124    pub(crate) qs: [u8; QK_K / 4],
125    pub(crate) scales: [u8; 12],
126    pub(crate) d: f16,
127}
128const _: () = assert!(QK_K / 8 + QK_K / 4 + 12 + 2 == std::mem::size_of::<BlockQ3K>());
129
130#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
131// https://github.com/ggerganov/llama.cpp/blob/468ea24fb4633a0d681f7ac84089566c1c6190cb/k_quants.h#L82
132#[repr(C)]
133pub struct BlockQ4K {
134    pub(crate) d: f16,
135    pub(crate) dmin: f16,
136    pub(crate) scales: [u8; K_SCALE_SIZE],
137    pub(crate) qs: [u8; QK_K / 2],
138}
139const _: () = assert!(QK_K / 2 + K_SCALE_SIZE + 2 * 2 == std::mem::size_of::<BlockQ4K>());
140
141#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
142#[repr(C)]
143pub struct BlockQ5K {
144    pub(crate) d: f16,
145    pub(crate) dmin: f16,
146    pub(crate) scales: [u8; K_SCALE_SIZE],
147    pub(crate) qh: [u8; QK_K / 8],
148    pub(crate) qs: [u8; QK_K / 2],
149}
150const _: () =
151    assert!(QK_K / 8 + QK_K / 2 + 2 * 2 + K_SCALE_SIZE == std::mem::size_of::<BlockQ5K>());
152
153#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
154#[repr(C)]
155pub struct BlockQ6K {
156    pub(crate) ql: [u8; QK_K / 2],
157    pub(crate) qh: [u8; QK_K / 4],
158    pub(crate) scales: [i8; QK_K / 16],
159    pub(crate) d: f16,
160}
161const _: () = assert!(3 * QK_K / 4 + QK_K / 16 + 2 == std::mem::size_of::<BlockQ6K>());
162
163#[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
164#[repr(C)]
165pub struct BlockQ8K {
166    pub(crate) d: f32,
167    pub(crate) qs: [i8; QK_K],
168    pub(crate) bsums: [i16; QK_K / 16],
169}
170const _: () = assert!(4 + QK_K + QK_K / 16 * 2 == std::mem::size_of::<BlockQ8K>());
171
172/// View a block slice as its raw bytes: GGUF's on-disk layout, byte for byte.
173///
174/// Infallible: `u8` has alignment 1 and every `T: GgmlType` is `Pod`, so this
175/// is a plain reinterpretation with no length or alignment condition.
176pub fn blocks_as_bytes<T: GgmlType>(blocks: &[T]) -> &[u8] {
177    bytemuck::cast_slice(blocks)
178}
179
180/// Reinterpret GGUF bytes as a dense `[T]` array.
181///
182/// Borrowed when `bytes` happens to be aligned for `T`, **copied** into an
183/// aligned `Vec<T>` otherwise: a `&[u8]` guarantees only 1-byte alignment, and
184/// a sub-slice taken at a tensor offset inside an mmap need not be aligned for
185/// the `f16`/`f32` fields of a block, so assuming it is would be unsound.
186///
187/// A trailing partial block is dropped, matching the pre-existing
188/// per-model `cpu_dequant_*_bytes` callers; callers that must reject a ragged
189/// length check `bytes.len() % size_of::<T>()` themselves.
190pub fn blocks_from_bytes_cow<T: GgmlType>(bytes: &[u8]) -> std::borrow::Cow<'_, [T]> {
191    let size = std::mem::size_of::<T>();
192    let whole = &bytes[..(bytes.len() / size) * size];
193    match bytemuck::try_cast_slice::<u8, T>(whole) {
194        Ok(blocks) => std::borrow::Cow::Borrowed(blocks),
195        Err(_) => std::borrow::Cow::Owned(
196            whole
197                .chunks_exact(size)
198                .map(bytemuck::pod_read_unaligned::<T>)
199                .collect(),
200        ),
201    }
202}
203
204impl GgmlType for BlockQ4_0 {
205    const DTYPE: GgmlDType = GgmlDType::Q4_0;
206    const BLCK_SIZE: usize = QK4_0;
207
208    // https://github.com/ggerganov/llama.cpp/blob/468ea24fb4633a0d681f7ac84089566c1c6190cb/ggml.c#L1525
209    fn to_float(xs: &[Self], ys: &mut [f32]) {
210        let k = ys.len();
211        let qk = Self::BLCK_SIZE;
212        debug_assert!(
213            k.is_multiple_of(qk),
214            "dequantize_row_q4_0: {k} is not divisible by {qk}"
215        );
216
217        let nb = k / qk;
218        for i in 0..nb {
219            let d = xs[i].d.to_f32();
220
221            for j in 0..(qk / 2) {
222                let x0 = (xs[i].qs[j] & 0x0F) as i16 - 8;
223                let x1 = (xs[i].qs[j] >> 4) as i16 - 8;
224
225                ys[i * qk + j] = (x0 as f32) * d;
226                ys[i * qk + j + qk / 2] = (x1 as f32) * d;
227            }
228        }
229    }
230
231    fn from_float(xs: &[f32], ys: &mut [Self]) {
232        // quantize_row_q4_0
233        let qk = Self::BLCK_SIZE;
234        let k = xs.len();
235        debug_assert!(k.is_multiple_of(qk), "{k} is not divisible by {qk}");
236        debug_assert_eq!(
237            ys.len(),
238            k / qk,
239            "size mismatch {} {} {}",
240            xs.len(),
241            ys.len(),
242            qk,
243        );
244        for (i, ys) in ys.iter_mut().enumerate() {
245            let mut amax = 0f32;
246            let mut max = 0f32;
247
248            let xs = &xs[i * qk..(i + 1) * qk];
249            for &x in xs.iter() {
250                if amax < x.abs() {
251                    amax = x.abs();
252                    max = x;
253                }
254            }
255            let d = max / -8.0;
256            let id = if d != 0f32 { 1. / d } else { 0. };
257            ys.d = f16::from_f32(d);
258
259            for (j, q) in ys.qs.iter_mut().enumerate() {
260                let x0 = xs[j] * id;
261                let x1 = xs[qk / 2 + j] * id;
262                let xi0 = u8::min(15, (x0 + 8.5) as u8);
263                let xi1 = u8::min(15, (x1 + 8.5) as u8);
264                *q = xi0 | (xi1 << 4)
265            }
266        }
267    }
268
269    // https://github.com/ggerganov/llama.cpp/blob/b5ffb2849d23afe73647f68eec7b68187af09be6/ggml.c#L2361C10-L2361C122
270}
271
272impl GgmlType for BlockQ4_1 {
273    const DTYPE: GgmlDType = GgmlDType::Q4_1;
274    const BLCK_SIZE: usize = QK4_1;
275
276    fn from_float(xs: &[f32], ys: &mut [Self]) {
277        // quantize_row_q4_1
278        let qk = Self::BLCK_SIZE;
279
280        debug_assert_eq!(
281            ys.len() * qk,
282            xs.len(),
283            "size mismatch {} {} {}",
284            xs.len(),
285            ys.len(),
286            qk,
287        );
288        for (i, ys) in ys.iter_mut().enumerate() {
289            let xs = &xs[i * qk..(i + 1) * qk];
290
291            let mut min = f32::INFINITY;
292            let mut max = f32::NEG_INFINITY;
293            for &x in xs.iter() {
294                min = f32::min(x, min);
295                max = f32::max(x, max);
296            }
297            let d = (max - min) / ((1 << 4) - 1) as f32;
298            let id = if d != 0f32 { 1. / d } else { 0. };
299            ys.d = f16::from_f32(d);
300            ys.m = f16::from_f32(min);
301
302            for (j, q) in ys.qs.iter_mut().take(qk / 2).enumerate() {
303                let x0 = (xs[j] - min) * id;
304                let x1 = (xs[qk / 2 + j] - min) * id;
305
306                let xi0 = u8::min(15, (x0 + 0.5) as u8);
307                let xi1 = u8::min(15, (x1 + 0.5) as u8);
308
309                *q = xi0 | (xi1 << 4);
310            }
311        }
312    }
313
314    // https://github.com/ggerganov/llama.cpp/blob/468ea24fb4633a0d681f7ac84089566c1c6190cb/ggml.c#L1545
315    fn to_float(xs: &[Self], ys: &mut [f32]) {
316        let k = ys.len();
317        debug_assert!(
318            k.is_multiple_of(QK4_1),
319            "dequantize_row_q4_1: {k} is not divisible by {QK4_1}"
320        );
321
322        let nb = k / QK4_1;
323        for i in 0..nb {
324            let d = xs[i].d.to_f32();
325            let m = xs[i].m.to_f32();
326
327            for j in 0..(QK4_1 / 2) {
328                let x0 = xs[i].qs[j] & 0x0F;
329                let x1 = xs[i].qs[j] >> 4;
330
331                ys[i * QK4_1 + j] = (x0 as f32) * d + m;
332                ys[i * QK4_1 + j + QK4_1 / 2] = (x1 as f32) * d + m;
333            }
334        }
335    }
336}
337
338impl GgmlType for BlockQ5_0 {
339    const DTYPE: GgmlDType = GgmlDType::Q5_0;
340    const BLCK_SIZE: usize = QK5_0;
341
342    fn from_float(xs: &[f32], ys: &mut [Self]) {
343        // quantize_row_q5_0
344        debug_assert_eq!(
345            ys.len() * Self::BLCK_SIZE,
346            xs.len(),
347            "size mismatch {} {} {}",
348            xs.len(),
349            ys.len(),
350            Self::BLCK_SIZE,
351        );
352        for (i, ys) in ys.iter_mut().enumerate() {
353            let xs = &xs[i * Self::BLCK_SIZE..(i + 1) * Self::BLCK_SIZE];
354
355            let mut amax = 0f32;
356            let mut max = 0f32;
357            for &x in xs.iter() {
358                if amax < x.abs() {
359                    amax = x.abs();
360                    max = x;
361                }
362            }
363            let d = max / -16.;
364            let id = if d != 0f32 { 1. / d } else { 0. };
365            ys.d = f16::from_f32(d);
366            let mut qh = 0u32;
367            for j in 0..Self::BLCK_SIZE / 2 {
368                let x0 = xs[j] * id;
369                let x1 = xs[j + Self::BLCK_SIZE / 2] * id;
370                let xi0 = ((x0 + 16.5) as i8).min(31) as u8;
371                let xi1 = ((x1 + 16.5) as i8).min(31) as u8;
372                ys.qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4);
373                qh |= ((xi0 as u32 & 0x10) >> 4) << j;
374                qh |= ((xi1 as u32 & 0x10) >> 4) << (j + Self::BLCK_SIZE / 2);
375            }
376            LittleEndian::write_u32(&mut ys.qh, qh)
377        }
378    }
379
380    // https://github.com/ggerganov/llama.cpp/blob/468ea24fb4633a0d681f7ac84089566c1c6190cb/ggml.c#L1566
381    fn to_float(xs: &[Self], ys: &mut [f32]) {
382        let k = ys.len();
383        debug_assert!(
384            k.is_multiple_of(QK5_0),
385            "dequantize_row_q5_0: {k} is not divisible by {QK5_0}"
386        );
387        let nb = k / QK5_0;
388        for i in 0..nb {
389            let d = xs[i].d.to_f32();
390            let qh: u32 = LittleEndian::read_u32(&xs[i].qh);
391
392            for j in 0..(QK5_0 / 2) {
393                let xh_0 = (((qh >> j) << 4) & 0x10) as u8;
394                let xh_1 = ((qh >> (j + 12)) & 0x10) as u8;
395
396                let x0 = ((xs[i].qs[j] & 0x0F) | xh_0) as i32 - 16;
397                let x1 = ((xs[i].qs[j] >> 4) | xh_1) as i32 - 16;
398
399                ys[i * QK5_0 + j] = (x0 as f32) * d;
400                ys[i * QK5_0 + j + QK5_0 / 2] = (x1 as f32) * d;
401            }
402        }
403    }
404}
405
406impl GgmlType for BlockQ5_1 {
407    const DTYPE: GgmlDType = GgmlDType::Q5_1;
408    const BLCK_SIZE: usize = QK5_1;
409
410    fn from_float(xs: &[f32], ys: &mut [Self]) {
411        // quantize_row_q5_1
412        let qk = Self::BLCK_SIZE;
413        debug_assert_eq!(
414            ys.len() * qk,
415            xs.len(),
416            "size mismatch {} {} {}",
417            xs.len(),
418            ys.len(),
419            qk,
420        );
421        for (i, ys) in ys.iter_mut().enumerate() {
422            let xs = &xs[i * qk..(i + 1) * qk];
423
424            let mut min = f32::INFINITY;
425            let mut max = f32::NEG_INFINITY;
426            for &x in xs.iter() {
427                min = f32::min(x, min);
428                max = f32::max(x, max);
429            }
430            let d = (max - min) / ((1 << 5) - 1) as f32;
431            let id = if d != 0f32 { 1. / d } else { 0. };
432            ys.d = f16::from_f32(d);
433            ys.m = f16::from_f32(min);
434
435            let mut qh = 0u32;
436            for (j, q) in ys.qs.iter_mut().take(qk / 2).enumerate() {
437                let x0 = (xs[j] - min) * id;
438                let x1 = (xs[qk / 2 + j] - min) * id;
439
440                let xi0 = (x0 + 0.5) as u8;
441                let xi1 = (x1 + 0.5) as u8;
442
443                *q = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4);
444                // get the 5-th bit and store it in qh at the right position
445                qh |= ((xi0 as u32 & 0x10) >> 4) << j;
446                qh |= ((xi1 as u32 & 0x10) >> 4) << (j + qk / 2);
447            }
448            LittleEndian::write_u32(&mut ys.qh, qh);
449        }
450    }
451
452    // https://github.com/ggerganov/llama.cpp/blob/468ea24fb4633a0d681f7ac84089566c1c6190cb/ggml.c#L1592
453    fn to_float(xs: &[Self], ys: &mut [f32]) {
454        let k = ys.len();
455        debug_assert!(
456            k.is_multiple_of(QK5_1),
457            "dequantize_row_q5_1: {k} is not divisible by {QK5_1}"
458        );
459
460        let nb = k / QK5_1;
461        for i in 0..nb {
462            let d = xs[i].d.to_f32();
463            let m = xs[i].m.to_f32();
464            let qh: u32 = LittleEndian::read_u32(&xs[i].qh);
465
466            for j in 0..(QK5_1 / 2) {
467                let xh_0 = (((qh >> j) << 4) & 0x10) as u8;
468                let xh_1 = ((qh >> (j + 12)) & 0x10) as u8;
469
470                let x0 = (xs[i].qs[j] & 0x0F) | xh_0;
471                let x1 = (xs[i].qs[j] >> 4) | xh_1;
472
473                ys[i * QK5_1 + j] = (x0 as f32) * d + m;
474                ys[i * QK5_1 + j + QK5_1 / 2] = (x1 as f32) * d + m;
475            }
476        }
477    }
478}
479
480impl GgmlType for BlockQ8_0 {
481    const DTYPE: GgmlDType = GgmlDType::Q8_0;
482    const BLCK_SIZE: usize = QK8_0;
483
484    // https://github.com/ggerganov/llama.cpp/blob/468ea24fb4633a0d681f7ac84089566c1c6190cb/ggml.c#L1619
485    fn to_float(xs: &[Self], ys: &mut [f32]) {
486        let k = ys.len();
487        debug_assert!(
488            k.is_multiple_of(QK8_0),
489            "dequantize_row_q8_0: {k} is not divisible by {QK8_0}"
490        );
491
492        let nb = k / QK8_0;
493
494        for i in 0..nb {
495            let d = xs[i].d.to_f32();
496
497            for j in 0..QK8_0 {
498                ys[i * QK8_0 + j] = xs[i].qs[j] as f32 * d;
499            }
500        }
501    }
502
503    fn from_float(xs: &[f32], ys: &mut [Self]) {
504        // quantize_row_q8_0
505        let k = xs.len();
506        debug_assert!(
507            k.is_multiple_of(Self::BLCK_SIZE),
508            "{k} is not divisible by {}",
509            Self::BLCK_SIZE
510        );
511        debug_assert_eq!(
512            ys.len(),
513            k / Self::BLCK_SIZE,
514            "size mismatch {} {} {}",
515            xs.len(),
516            ys.len(),
517            Self::BLCK_SIZE
518        );
519        for (i, ys) in ys.iter_mut().enumerate() {
520            let mut amax = 0f32;
521            let xs = &xs[i * Self::BLCK_SIZE..(i + 1) * Self::BLCK_SIZE];
522            for &x in xs.iter() {
523                amax = amax.max(x.abs())
524            }
525            let d = amax / ((1 << 7) - 1) as f32;
526            let id = if d != 0f32 { 1. / d } else { 0. };
527            ys.d = f16::from_f32(d);
528            for (y, &x) in ys.qs.iter_mut().zip(xs.iter()) {
529                *y = f32::round(x * id) as i8
530            }
531        }
532    }
533}
534
535impl GgmlType for BlockQ8_1 {
536    const DTYPE: GgmlDType = GgmlDType::Q8_1;
537    const BLCK_SIZE: usize = QK8_1;
538
539    fn from_float(xs: &[f32], ys: &mut [Self]) {
540        // quantize_row_q8_1
541        debug_assert_eq!(
542            ys.len() * Self::BLCK_SIZE,
543            xs.len(),
544            "size mismatch {} {} {}",
545            xs.len(),
546            ys.len(),
547            Self::BLCK_SIZE
548        );
549        for (i, ys) in ys.iter_mut().enumerate() {
550            let mut amax = 0f32;
551            let xs = &xs[i * Self::BLCK_SIZE..(i + 1) * Self::BLCK_SIZE];
552            for &x in xs.iter() {
553                amax = amax.max(x.abs())
554            }
555            let d = amax / ((1 << 7) - 1) as f32;
556            let id = if d != 0f32 { 1. / d } else { 0. };
557            ys.d = f16::from_f32(d);
558            let mut sum = 0i32;
559            for j in 0..Self::BLCK_SIZE / 2 {
560                let v0 = xs[j] * id;
561                let v1 = xs[j + Self::BLCK_SIZE / 2] * id;
562                ys.qs[j] = f32::round(v0) as i8;
563                ys.qs[j + Self::BLCK_SIZE / 2] = f32::round(v1) as i8;
564                sum += ys.qs[j] as i32 + ys.qs[j + Self::BLCK_SIZE / 2] as i32;
565            }
566            ys.s = f16::from_f32(sum as f32) * ys.d;
567        }
568    }
569
570    fn to_float(_xs: &[Self], _ys: &mut [f32]) {
571        unimplemented!("no support for vec-dot on Q8_1") // GAP(GAP-125)
572    }
573}
574
575impl GgmlType for BlockQ2K {
576    const DTYPE: GgmlDType = GgmlDType::Q2K;
577    const BLCK_SIZE: usize = QK_K;
578
579    // https://github.com/ggerganov/llama.cpp/blob/8183159cf3def112f6d1fe94815fce70e1bffa12/k_quants.c#L279
580    fn from_float(xs: &[f32], ys: &mut [Self]) {
581        const Q4SCALE: f32 = 15.0;
582
583        for (block, x) in group_for_quantization(xs, ys) {
584            //calculate scales and mins
585            let mut mins: [f32; QK_K / 16] = [0.0; QK_K / 16];
586            let mut scales: [f32; QK_K / 16] = [0.0; QK_K / 16];
587
588            for (j, x_scale_slice) in x.chunks(16).enumerate() {
589                (scales[j], mins[j]) = make_qkx1_quants(3, 5, x_scale_slice);
590            }
591            // get max scale and max min and ensure they are >= 0.0
592            let max_scale = scales.iter().fold(0.0, |max, &val| val.max(max));
593            let max_min = mins.iter().fold(0.0, |max, &val| val.max(max));
594
595            if max_scale > 0.0 {
596                let iscale = Q4SCALE / max_scale;
597                for (j, scale) in scales.iter().enumerate().take(QK_K / 16) {
598                    block.scales[j] = nearest_int(iscale * scale) as u8;
599                }
600                block.d = f16::from_f32(max_scale / Q4SCALE);
601            } else {
602                for j in 0..QK_K / 16 {
603                    block.scales[j] = 0;
604                }
605                block.d = f16::from_f32(0.0);
606            }
607
608            if max_min > 0.0 {
609                let iscale = Q4SCALE / max_min;
610                for (j, scale) in block.scales.iter_mut().enumerate() {
611                    let l = nearest_int(iscale * mins[j]) as u8;
612                    *scale |= l << 4;
613                }
614                block.dmin = f16::from_f32(max_min / Q4SCALE);
615            } else {
616                block.dmin = f16::from_f32(0.0);
617            }
618
619            let mut big_l: [u8; QK_K] = [0; QK_K];
620
621            for j in 0..QK_K / 16 {
622                let d = block.d.to_f32() * (block.scales[j] & 0xF) as f32;
623                if d == 0.0 {
624                    continue;
625                }
626                let dm = block.dmin.to_f32() * (block.scales[j] >> 4) as f32;
627                for ii in 0..16 {
628                    let ll = nearest_int((x[16 * j + ii] + dm) / d).clamp(0, 3);
629                    big_l[16 * j + ii] = ll as u8;
630                }
631            }
632
633            for j in (0..QK_K).step_by(128) {
634                for ll in 0..32 {
635                    block.qs[j / 4 + ll] = big_l[j + ll]
636                        | (big_l[j + ll + 32] << 2)
637                        | (big_l[j + ll + 64] << 4)
638                        | (big_l[j + ll + 96] << 6);
639                }
640            }
641        }
642    }
643
644    fn from_float_imatrix(xs: &[f32], ys: &mut [Self], imatrix_weights: &[f32], n_per_row: usize) {
645        for (sblk_idx, (block, x)) in group_for_quantization(xs, ys).into_iter().enumerate() {
646            let mut mins: [f32; QK_K / 16] = [0.0; QK_K / 16];
647            let mut scales: [f32; QK_K / 16] = [0.0; QK_K / 16];
648            let mut weights: [f32; 16] = [0.0; 16];
649            let mut sw: [f32; QK_K / 16] = [0.0; QK_K / 16];
650            let mut ls: [u8; QK_K / 16] = [0; QK_K / 16];
651            let mut lm: [u8; QK_K / 16] = [0; QK_K / 16];
652
653            let sum_x2 = x.iter().map(|x| x * x).sum::<f32>();
654            let sigma2 = sum_x2 / QK_K as f32;
655            for (j, x_scale_slice) in x.as_chunks::<16>().0.iter().enumerate() {
656                for (l, (w_elem, x_elem)) in weights.iter_mut().zip(x_scale_slice).enumerate() {
657                    let imatrix_row = sblk_idx % (n_per_row / QK_K);
658                    let imatrix_w = imatrix_weights[imatrix_row * QK_K + 16 * j + l];
659                    *w_elem = imatrix_w * (sigma2 + x_elem * x_elem).sqrt();
660                }
661                let sumw = weights.iter().sum::<f32>();
662                sw[j] = sumw;
663                (scales[j], mins[j]) =
664                    make_qkx3_quants(3, x_scale_slice, Some(&weights), -0.9, 0.05, 36, false);
665            }
666
667            let d_block = make_qp_quants(QK_K / 16, 15, &scales, &mut ls, &sw);
668            let m_block = make_qp_quants(QK_K / 16, 15, &mins, &mut lm, &sw);
669
670            block.d = f16::from_f32(d_block);
671            block.dmin = f16::from_f32(m_block);
672
673            for j in 0..QK_K / 16 {
674                block.scales[j] = ls[j] | (lm[j] << 4);
675            }
676
677            let mut big_l: [u8; QK_K] = [0; QK_K];
678
679            for j in 0..QK_K / 16 {
680                let d = block.d.to_f32() * (block.scales[j] & 0xF) as f32;
681                if d == 0.0 {
682                    continue;
683                }
684                let dm = block.dmin.to_f32() * (block.scales[j] >> 4) as f32;
685                for ii in 0..16 {
686                    let ll = nearest_int((x[16 * j + ii] + dm) / d).clamp(0, 3);
687                    big_l[16 * j + ii] = ll as u8;
688                }
689            }
690
691            for j in (0..QK_K).step_by(128) {
692                for ll in 0..32 {
693                    block.qs[j / 4 + ll] = big_l[j + ll]
694                        | (big_l[j + ll + 32] << 2)
695                        | (big_l[j + ll + 64] << 4)
696                        | (big_l[j + ll + 96] << 6);
697                }
698            }
699        }
700    }
701    // https://github.com/ggerganov/llama.cpp/blob/8183159cf3def112f6d1fe94815fce70e1bffa12/k_quants.c#L354
702    fn to_float(xs: &[Self], ys: &mut [f32]) {
703        for (block, y) in group_for_dequantization(xs, ys) {
704            let d = block.d.to_f32();
705            let min = block.dmin.to_f32();
706
707            let mut is = 0;
708
709            for (y_block, qs) in y
710                .as_chunks_mut::<128>()
711                .0
712                .iter_mut()
713                .zip(block.qs.as_chunks::<32>().0.iter())
714            {
715                // Step by 32 over q.
716                let mut shift = 0;
717                let mut y_block_index = 0;
718                for _j in 0..4 {
719                    let sc = block.scales[is];
720                    is += 1;
721                    let dl = d * (sc & 0xF) as f32;
722                    let ml = min * (sc >> 4) as f32;
723                    for q in &qs[..16] {
724                        let y = dl * ((q >> shift) & 3) as f32 - ml;
725                        y_block[y_block_index] = y;
726                        y_block_index += 1;
727                    }
728
729                    let sc = block.scales[is];
730                    is += 1;
731                    let dl = d * (sc & 0xF) as f32;
732                    let ml = min * (sc >> 4) as f32;
733                    for q in &qs[16..] {
734                        let y = dl * ((q >> shift) & 3) as f32 - ml;
735                        y_block[y_block_index] = y;
736                        y_block_index += 1;
737                    }
738
739                    shift += 2;
740                }
741            }
742        }
743    }
744}
745
746impl GgmlType for BlockQ3K {
747    const DTYPE: GgmlDType = GgmlDType::Q3K;
748    const BLCK_SIZE: usize = QK_K;
749
750    fn from_float(xs: &[f32], ys: &mut [Self]) {
751        for (block, x) in group_for_quantization(xs, ys) {
752            let mut scales: [f32; QK_K / 16] = [0.0; QK_K / 16];
753            for (j, x_scale_slice) in x.as_chunks::<16>().0.iter().enumerate() {
754                scales[j] = make_q3_quants(x_scale_slice, 4, true);
755            }
756
757            // Get max scale by absolute value.
758            let mut max_scale: f32 = 0.0;
759            for &scale in scales.iter() {
760                if scale.abs() > max_scale.abs() {
761                    max_scale = scale;
762                }
763            }
764
765            block.scales.fill(0);
766
767            if max_scale != 0.0 {
768                let iscale = -32.0 / max_scale;
769                for (j, scale) in scales.iter().enumerate() {
770                    let l_val = nearest_int(iscale * scale);
771                    let l_val = l_val.clamp(-32, 31) + 32;
772                    if j < 8 {
773                        block.scales[j] = (l_val & 0xF) as u8;
774                    } else {
775                        block.scales[j - 8] |= ((l_val & 0xF) << 4) as u8;
776                    }
777                    let l_val = l_val >> 4;
778                    block.scales[j % 4 + 8] |= (l_val << (2 * (j / 4))) as u8;
779                }
780                block.d = f16::from_f32(1.0 / iscale);
781            } else {
782                block.d = f16::from_f32(0.0);
783            }
784
785            let mut l: [i8; QK_K] = [0; QK_K];
786
787            for j in 0..QK_K / 16 {
788                let sc = if j < 8 {
789                    block.scales[j] & 0xF
790                } else {
791                    block.scales[j - 8] >> 4
792                };
793                let sc = (sc | (((block.scales[8 + j % 4] >> (2 * (j / 4))) & 3) << 4)) as i8 - 32;
794                let d = block.d.to_f32() * sc as f32;
795                if d != 0.0 {
796                    for ii in 0..16 {
797                        let l_val = nearest_int(x[16 * j + ii] / d);
798                        l[16 * j + ii] = (l_val.clamp(-4, 3) + 4) as i8;
799                    }
800                }
801            }
802
803            block.hmask.fill(0);
804            let mut m = 0;
805            let mut hm = 1;
806
807            for ll in l.iter_mut() {
808                if *ll > 3 {
809                    block.hmask[m] |= hm;
810                    *ll -= 4;
811                }
812                m += 1;
813                if m == QK_K / 8 {
814                    m = 0;
815                    hm <<= 1;
816                }
817            }
818
819            for j in (0..QK_K).step_by(128) {
820                for l_val in 0..32 {
821                    block.qs[j / 4 + l_val] = (l[j + l_val]
822                        | (l[j + l_val + 32] << 2)
823                        | (l[j + l_val + 64] << 4)
824                        | (l[j + l_val + 96] << 6))
825                        as u8;
826                }
827            }
828        }
829    }
830
831    fn from_float_imatrix(xs: &[f32], ys: &mut [Self], imatrix_weights: &[f32], n_per_row: usize) {
832        for (sblk_idx, (block, x)) in group_for_quantization(xs, ys).into_iter().enumerate() {
833            let mut scales: [f32; QK_K / 16] = [0.0; QK_K / 16];
834            let mut weights: [f32; 16] = [0.0; 16];
835            let mut sw: [f32; QK_K / 16] = [0.0; QK_K / 16];
836            let mut ls: [i8; QK_K / 16] = [0; QK_K / 16];
837            let mut l: [i8; QK_K] = [0; QK_K];
838
839            let sum_x2 = x.iter().map(|x| x * x).sum::<f32>();
840            let sigma2 = 2. * sum_x2 / QK_K as f32;
841
842            for (j, x_scale_slice) in x.as_chunks::<16>().0.iter().enumerate() {
843                for (l_idx, (w_elem, x_elem)) in weights.iter_mut().zip(x_scale_slice).enumerate() {
844                    let imatrix_row = sblk_idx % (n_per_row / QK_K);
845                    let imatrix_w = imatrix_weights[imatrix_row * QK_K + 16 * j + l_idx];
846                    *w_elem = imatrix_w * (sigma2 + x_elem * x_elem).sqrt();
847                }
848                let sumw = weights.iter().sum::<f32>();
849                sw[j] = sumw;
850                scales[j] = unsafe {
851                    make_qx_quants(
852                        16,
853                        4,
854                        x_scale_slice.as_ptr(),
855                        l.as_mut_ptr().add(16 * j),
856                        1,
857                        weights.as_ptr(),
858                    )
859                };
860            }
861
862            block.scales.fill(0);
863            let d_block = unsafe {
864                make_qx_quants(
865                    QK_K / 16,
866                    32,
867                    scales.as_ptr(),
868                    ls.as_mut_ptr(),
869                    1,
870                    sw.as_ptr(),
871                )
872            };
873            block.d = f16::from_f32(d_block);
874            for (j, l_val) in ls.iter().enumerate().take(QK_K / 16) {
875                if j < 8 {
876                    block.scales[j] = (l_val & 0xF) as u8;
877                } else {
878                    block.scales[j - 8] |= ((l_val & 0xF) << 4) as u8;
879                }
880                let l_val = l_val >> 4;
881                block.scales[j % 4 + 8] |= (l_val << (2 * (j / 4))) as u8;
882            }
883
884            for j in 0..QK_K / 16 {
885                let sc = if j < 8 {
886                    block.scales[j] & 0xF
887                } else {
888                    block.scales[j - 8] >> 4
889                };
890                let sc = (sc | (((block.scales[8 + j % 4] >> (2 * (j / 4))) & 3) << 4)) as i8 - 32;
891                let d = block.d.to_f32() * sc as f32;
892                if d != 0.0 {
893                    for ii in 0..16 {
894                        let l_val = nearest_int(x[16 * j + ii] / d);
895                        l[16 * j + ii] = (l_val.clamp(-4, 3) + 4) as i8;
896                    }
897                }
898            }
899
900            block.hmask.fill(0);
901            let mut m = 0;
902            let mut hm = 1;
903
904            for ll in l.iter_mut() {
905                if *ll > 3 {
906                    block.hmask[m] |= hm;
907                    *ll -= 4;
908                }
909                m += 1;
910                if m == QK_K / 8 {
911                    m = 0;
912                    hm <<= 1;
913                }
914            }
915
916            for j in (0..QK_K).step_by(128) {
917                for l_val in 0..32 {
918                    block.qs[j / 4 + l_val] = (l[j + l_val]
919                        | (l[j + l_val + 32] << 2)
920                        | (l[j + l_val + 64] << 4)
921                        | (l[j + l_val + 96] << 6))
922                        as u8;
923                }
924            }
925        }
926    }
927
928    // https://github.com/ggerganov/llama.cpp/blob/8183159cf3def112f6d1fe94815fce70e1bffa12/k_quants.c#L533
929    fn to_float(xs: &[Self], ys: &mut [f32]) {
930        const KMASK1: u32 = 0x03030303;
931        const KMASK2: u32 = 0x0f0f0f0f;
932
933        for (block, y) in group_for_dequantization(xs, ys) {
934            //Reconstruct the scales
935            let mut aux = [0; 4];
936            LittleEndian::read_u32_into(&block.scales, &mut aux[0..3]);
937
938            let tmp = aux[2];
939            aux[2] = ((aux[0] >> 4) & KMASK2) | (((tmp >> 4) & KMASK1) << 4);
940            aux[3] = ((aux[1] >> 4) & KMASK2) | (((tmp >> 6) & KMASK1) << 4);
941            aux[0] = (aux[0] & KMASK2) | (((tmp) & KMASK1) << 4);
942            aux[1] = (aux[1] & KMASK2) | (((tmp >> 2) & KMASK1) << 4);
943
944            //Transfer the scales into an i8 array
945            let scales: &mut [i8] =
946                unsafe { std::slice::from_raw_parts_mut(aux.as_mut_ptr() as *mut i8, 16) };
947
948            let d_all = block.d.to_f32();
949            let mut m = 1;
950            let mut is = 0;
951
952            // Dequantize both 128 long blocks
953            // 32 qs values per 128 long block
954            // Each 16 elements get a scale
955            for (y, qs) in y
956                .as_chunks_mut::<128>()
957                .0
958                .iter_mut()
959                .zip(block.qs.as_chunks::<32>().0.iter())
960            {
961                let mut shift = 0;
962                for shift_scoped_y in y.as_chunks_mut::<32>().0.iter_mut() {
963                    for (scale_index, scale_scoped_y) in shift_scoped_y
964                        .as_chunks_mut::<16>()
965                        .0
966                        .iter_mut()
967                        .enumerate()
968                    {
969                        let dl = d_all * (scales[is] as f32 - 32.0);
970                        for (i, inner_y) in scale_scoped_y.iter_mut().enumerate() {
971                            let new_y = dl
972                                * (((qs[i + 16 * scale_index] >> shift) & 3) as i8
973                                    - if (block.hmask[i + 16 * scale_index] & m) == 0 {
974                                        4
975                                    } else {
976                                        0
977                                    }) as f32;
978                            *inner_y = new_y;
979                        }
980                        // 16 block finished => advance scale index
981                        is += 1;
982                    }
983                    // 32 block finished => increase shift and m
984                    shift += 2;
985                    m <<= 1;
986                }
987            }
988        }
989    }
990}
991
992impl GgmlType for BlockQ4K {
993    const DTYPE: GgmlDType = GgmlDType::Q4K;
994    const BLCK_SIZE: usize = QK_K;
995
996    fn from_float(xs: &[f32], ys: &mut [Self]) {
997        for (block, x) in group_for_quantization(xs, ys) {
998            let mut mins: [f32; QK_K / 32] = [0.0; QK_K / 32];
999            let mut scales: [f32; QK_K / 32] = [0.0; QK_K / 32];
1000
1001            for (j, x_scale_slice) in x.as_chunks::<32>().0.iter().enumerate() {
1002                (scales[j], mins[j]) = make_qkx1_quants(15, 5, x_scale_slice);
1003            }
1004
1005            // get max scale and max min and ensure they are >= 0.0
1006            let max_scale = scales.iter().fold(0.0, |max, &val| val.max(max));
1007            let max_min = mins.iter().fold(0.0, |max, &val| val.max(max));
1008
1009            let inv_scale = if max_scale > 0.0 {
1010                63.0 / max_scale
1011            } else {
1012                0.0
1013            };
1014            let inv_min = if max_min > 0.0 { 63.0 / max_min } else { 0.0 };
1015
1016            for j in 0..QK_K / 32 {
1017                let ls = nearest_int(inv_scale * scales[j]).min(63) as u8;
1018                let lm = nearest_int(inv_min * mins[j]).min(63) as u8;
1019                if j < 4 {
1020                    block.scales[j] = ls;
1021                    block.scales[j + 4] = lm;
1022                } else {
1023                    block.scales[j + 4] = (ls & 0xF) | ((lm & 0xF) << 4);
1024                    block.scales[j - 4] |= (ls >> 4) << 6;
1025                    block.scales[j] |= (lm >> 4) << 6;
1026                }
1027            }
1028
1029            block.d = f16::from_f32(max_scale / 63.0);
1030            block.dmin = f16::from_f32(max_min / 63.0);
1031
1032            let mut l: [u8; QK_K] = [0; QK_K];
1033
1034            for j in 0..QK_K / 32 {
1035                let (sc, m) = get_scale_min_k4(j, &block.scales);
1036                let d = block.d.to_f32() * sc as f32;
1037                if d != 0.0 {
1038                    let dm = block.dmin.to_f32() * m as f32;
1039                    for ii in 0..32 {
1040                        let l_val = nearest_int((x[32 * j + ii] + dm) / d);
1041                        l[32 * j + ii] = l_val.clamp(0, 15) as u8;
1042                    }
1043                }
1044            }
1045
1046            let q = &mut block.qs;
1047            for j in (0..QK_K).step_by(64) {
1048                for l_val in 0..32 {
1049                    let offset_index = (j / 64) * 32 + l_val;
1050                    q[offset_index] = l[j + l_val] | (l[j + l_val + 32] << 4);
1051                }
1052            }
1053        }
1054    }
1055
1056    fn from_float_imatrix(xs: &[f32], ys: &mut [Self], imatrix_weights: &[f32], n_per_row: usize) {
1057        for (sblk_idx, (block, x)) in group_for_quantization(xs, ys).into_iter().enumerate() {
1058            let mut mins: [f32; QK_K / 32] = [0.0; QK_K / 32];
1059            let mut scales: [f32; QK_K / 32] = [0.0; QK_K / 32];
1060            let mut weights: [f32; 32] = [0.0; 32];
1061            let mut sw: [f32; QK_K / 32] = [0.0; QK_K / 32];
1062            let mut ls: [u8; QK_K / 32] = [0; QK_K / 32];
1063            let mut lm: [u8; QK_K / 32] = [0; QK_K / 32];
1064
1065            let sum_x2 = x.iter().map(|x| x * x).sum::<f32>();
1066            let sigma2 = 2. * sum_x2 / QK_K as f32;
1067
1068            for (j, x_scale_slice) in x.as_chunks::<32>().0.iter().enumerate() {
1069                for (l, (w_elem, x_elem)) in weights.iter_mut().zip(x_scale_slice).enumerate() {
1070                    let imatrix_row = sblk_idx % (n_per_row / QK_K);
1071                    let imatrix_w = imatrix_weights[imatrix_row * QK_K + 32 * j + l];
1072                    *w_elem = imatrix_w * (sigma2 + x_elem * x_elem).sqrt();
1073                }
1074                let sumw = weights.iter().sum::<f32>();
1075                sw[j] = sumw;
1076                (scales[j], mins[j]) =
1077                    make_qkx3_quants(15, x_scale_slice, Some(&weights), -0.9, 0.05, 36, false);
1078            }
1079
1080            let d_block = make_qp_quants(QK_K / 32, 63, &scales, &mut ls, &sw);
1081            let m_block = make_qp_quants(QK_K / 32, 63, &mins, &mut lm, &sw);
1082            for j in 0..QK_K / 32 {
1083                let ls_val = ls[j];
1084                let lm_val = lm[j];
1085                if j < 4 {
1086                    block.scales[j] = ls_val;
1087                    block.scales[j + 4] = lm_val;
1088                } else {
1089                    block.scales[j + 4] = (ls_val & 0xF) | ((lm_val & 0xF) << 4);
1090                    block.scales[j - 4] |= (ls_val >> 4) << 6;
1091                    block.scales[j] |= (lm_val >> 4) << 6;
1092                }
1093            }
1094
1095            block.d = f16::from_f32(d_block);
1096            block.dmin = f16::from_f32(m_block);
1097
1098            let mut l: [u8; QK_K] = [0; QK_K];
1099            for j in 0..QK_K / 32 {
1100                let (sc, m) = get_scale_min_k4(j, &block.scales);
1101                let d = block.d.to_f32() * sc as f32;
1102                if d != 0.0 {
1103                    let dm = block.dmin.to_f32() * m as f32;
1104                    for ii in 0..32 {
1105                        let l_val = nearest_int((x[32 * j + ii] + dm) / d);
1106                        l[32 * j + ii] = l_val.clamp(0, 15) as u8;
1107                    }
1108                }
1109            }
1110
1111            let q = &mut block.qs;
1112            for j in (0..QK_K).step_by(64) {
1113                for l_val in 0..32 {
1114                    let offset_index = (j / 64) * 32 + l_val;
1115                    q[offset_index] = l[j + l_val] | (l[j + l_val + 32] << 4);
1116                }
1117            }
1118        }
1119    }
1120    // https://github.com/ggerganov/llama.cpp/blob/8183159cf3def112f6d1fe94815fce70e1bffa12/k_quants.c#L735
1121    fn to_float(xs: &[Self], ys: &mut [f32]) {
1122        for (block, y) in group_for_dequantization(xs, ys) {
1123            let d = block.d.to_f32();
1124            let min = block.dmin.to_f32();
1125            let q = &block.qs;
1126            let mut is = 0;
1127            let mut ys_index = 0;
1128
1129            for j in (0..QK_K).step_by(64) {
1130                let q = &q[j / 2..j / 2 + 32];
1131                let (sc, m) = get_scale_min_k4(is, &block.scales);
1132                let d1 = d * sc as f32;
1133                let m1 = min * m as f32;
1134                let (sc, m) = get_scale_min_k4(is + 1, &block.scales);
1135                let d2 = d * sc as f32;
1136                let m2 = min * m as f32;
1137                for q in q {
1138                    y[ys_index] = d1 * (q & 0xF) as f32 - m1;
1139                    ys_index += 1;
1140                }
1141                for q in q {
1142                    y[ys_index] = d2 * (q >> 4) as f32 - m2;
1143                    ys_index += 1;
1144                }
1145                is += 2;
1146            }
1147        }
1148    }
1149}
1150
1151// https://github.com/ggerganov/llama.cpp/blob/8183159cf3def112f6d1fe94815fce70e1bffa12/k_quants.c#L928
1152impl GgmlType for BlockQ5K {
1153    const DTYPE: GgmlDType = GgmlDType::Q5K;
1154    const BLCK_SIZE: usize = QK_K;
1155
1156    // https://github.com/ggerganov/llama.cpp/blob/8183159cf3def112f6d1fe94815fce70e1bffa12/k_quants.c#L793
1157    fn from_float(xs: &[f32], ys: &mut [Self]) {
1158        for (block, x) in group_for_quantization(xs, ys) {
1159            let mut mins: [f32; QK_K / 32] = [0.0; QK_K / 32];
1160            let mut scales: [f32; QK_K / 32] = [0.0; QK_K / 32];
1161
1162            for (j, x_scale_slice) in x.as_chunks::<32>().0.iter().enumerate() {
1163                (scales[j], mins[j]) = make_qkx1_quants(31, 5, x_scale_slice);
1164            }
1165
1166            // get max scale and max min and ensure they are >= 0.0
1167            let max_scale = scales.iter().fold(0.0, |max, &val| val.max(max));
1168            let max_min = mins.iter().fold(0.0, |max, &val| val.max(max));
1169
1170            let inv_scale = if max_scale > 0.0 {
1171                63.0 / max_scale
1172            } else {
1173                0.0
1174            };
1175            let inv_min = if max_min > 0.0 { 63.0 / max_min } else { 0.0 };
1176            for j in 0..QK_K / 32 {
1177                let ls = nearest_int(inv_scale * scales[j]).min(63) as u8;
1178                let lm = nearest_int(inv_min * mins[j]).min(63) as u8;
1179                if j < 4 {
1180                    block.scales[j] = ls;
1181                    block.scales[j + 4] = lm;
1182                } else {
1183                    block.scales[j + 4] = (ls & 0xF) | ((lm & 0xF) << 4);
1184                    block.scales[j - 4] |= (ls >> 4) << 6;
1185                    block.scales[j] |= (lm >> 4) << 6;
1186                }
1187            }
1188            block.d = f16::from_f32(max_scale / 63.0);
1189            block.dmin = f16::from_f32(max_min / 63.0);
1190
1191            let mut l: [u8; QK_K] = [0; QK_K];
1192            for j in 0..QK_K / 32 {
1193                let (sc, m) = get_scale_min_k4(j, &block.scales);
1194                let d = block.d.to_f32() * sc as f32;
1195                if d == 0.0 {
1196                    continue;
1197                }
1198                let dm = block.dmin.to_f32() * m as f32;
1199                for ii in 0..32 {
1200                    let ll = nearest_int((x[32 * j + ii] + dm) / d);
1201                    l[32 * j + ii] = ll.clamp(0, 31) as u8;
1202                }
1203            }
1204
1205            let qh = &mut block.qh;
1206            let ql = &mut block.qs;
1207            qh.fill(0);
1208
1209            let mut m1 = 1;
1210            let mut m2 = 2;
1211            for n in (0..QK_K).step_by(64) {
1212                let offset = (n / 64) * 32;
1213                for j in 0..32 {
1214                    let mut l1 = l[n + j];
1215                    if l1 > 15 {
1216                        l1 -= 16;
1217                        qh[j] |= m1;
1218                    }
1219                    let mut l2 = l[n + j + 32];
1220                    if l2 > 15 {
1221                        l2 -= 16;
1222                        qh[j] |= m2;
1223                    }
1224                    ql[offset + j] = l1 | (l2 << 4);
1225                }
1226                m1 <<= 2;
1227                m2 <<= 2;
1228            }
1229        }
1230    }
1231
1232    fn from_float_imatrix(xs: &[f32], ys: &mut [Self], imatrix_weights: &[f32], n_per_row: usize) {
1233        for (sblk_idx, (block, x)) in group_for_quantization(xs, ys).into_iter().enumerate() {
1234            let mut mins: [f32; QK_K / 32] = [0.0; QK_K / 32];
1235            let mut scales: [f32; QK_K / 32] = [0.0; QK_K / 32];
1236            let mut weights: [f32; 32] = [0.0; 32];
1237            let mut sw: [f32; QK_K / 32] = [0.0; QK_K / 32];
1238            let mut ls: [u8; QK_K / 32] = [0; QK_K / 32];
1239            let mut lm: [u8; QK_K / 32] = [0; QK_K / 32];
1240
1241            let sum_x2 = x.iter().map(|x| x * x).sum::<f32>();
1242            let sigma2 = 2. * sum_x2 / QK_K as f32;
1243
1244            for (j, x_scale_slice) in x.as_chunks::<32>().0.iter().enumerate() {
1245                for (l, (w_elem, x_elem)) in weights.iter_mut().zip(x_scale_slice).enumerate() {
1246                    let imatrix_row = sblk_idx % (n_per_row / QK_K);
1247                    let imatrix_w = imatrix_weights[imatrix_row * QK_K + 32 * j + l];
1248                    *w_elem = imatrix_w * (sigma2 + x_elem * x_elem).sqrt();
1249                }
1250                let sumw = weights.iter().sum::<f32>();
1251                sw[j] = sumw;
1252                (scales[j], mins[j]) =
1253                    make_qkx3_quants(31, x_scale_slice, Some(&weights), -0.9, 0.05, 36, false);
1254            }
1255
1256            let d_block = make_qp_quants(QK_K / 32, 63, &scales, &mut ls, &sw);
1257            let m_block = make_qp_quants(QK_K / 32, 63, &mins, &mut lm, &sw);
1258            for j in 0..QK_K / 32 {
1259                let ls_val = ls[j].min(63);
1260                let lm_val = lm[j].min(63);
1261                if j < 4 {
1262                    block.scales[j] = ls_val;
1263                    block.scales[j + 4] = lm_val;
1264                } else {
1265                    block.scales[j + 4] = (ls_val & 0xF) | ((lm_val & 0xF) << 4);
1266                    block.scales[j - 4] |= (ls_val >> 4) << 6;
1267                    block.scales[j] |= (lm_val >> 4) << 6;
1268                }
1269            }
1270
1271            block.d = f16::from_f32(d_block);
1272            block.dmin = f16::from_f32(m_block);
1273
1274            let mut l: [u8; QK_K] = [0; QK_K];
1275            for j in 0..QK_K / 32 {
1276                let (sc, m) = get_scale_min_k4(j, &block.scales);
1277                let d = block.d.to_f32() * sc as f32;
1278                if d != 0.0 {
1279                    let dm = block.dmin.to_f32() * m as f32;
1280                    for ii in 0..32 {
1281                        let l_val = nearest_int((x[32 * j + ii] + dm) / d);
1282                        l[32 * j + ii] = l_val.clamp(0, 31) as u8;
1283                    }
1284                }
1285            }
1286
1287            let qh = &mut block.qh;
1288            let ql = &mut block.qs;
1289            qh.fill(0);
1290
1291            let mut m1 = 1;
1292            let mut m2 = 2;
1293            for n in (0..QK_K).step_by(64) {
1294                let offset = (n / 64) * 32;
1295                for j in 0..32 {
1296                    let mut l1 = l[n + j];
1297                    if l1 > 15 {
1298                        l1 -= 16;
1299                        qh[j] |= m1;
1300                    }
1301                    let mut l2 = l[n + j + 32];
1302                    if l2 > 15 {
1303                        l2 -= 16;
1304                        qh[j] |= m2;
1305                    }
1306                    ql[offset + j] = l1 | (l2 << 4);
1307                }
1308                m1 <<= 2;
1309                m2 <<= 2;
1310            }
1311        }
1312    }
1313
1314    // https://github.com/ggerganov/llama.cpp/blob/8183159cf3def112f6d1fe94815fce70e1bffa12/k_quants.c#L928
1315    fn to_float(xs: &[Self], ys: &mut [f32]) {
1316        for (block, y) in group_for_dequantization(xs, ys) {
1317            let d = block.d.to_f32();
1318            let min = block.dmin.to_f32();
1319            let ql = &block.qs;
1320            let qh = &block.qh;
1321            let mut is = 0;
1322            let mut u1 = 1;
1323            let mut u2 = 2;
1324            let mut ys_index = 0;
1325
1326            for j in (0..QK_K).step_by(64) {
1327                let ql = &ql[j / 2..j / 2 + 32];
1328                let (sc, m) = get_scale_min_k4(is, &block.scales);
1329                let d1 = d * sc as f32;
1330                let m1 = min * m as f32;
1331                let (sc, m) = get_scale_min_k4(is + 1, &block.scales);
1332                let d2 = d * sc as f32;
1333                let m2 = min * m as f32;
1334                for (ql, qh) in ql.iter().zip(qh) {
1335                    let to_add = if qh & u1 != 0 { 16f32 } else { 0f32 };
1336                    y[ys_index] = d1 * ((ql & 0xF) as f32 + to_add) - m1;
1337                    ys_index += 1;
1338                }
1339                for (ql, qh) in ql.iter().zip(qh) {
1340                    let to_add = if qh & u2 != 0 { 16f32 } else { 0f32 };
1341                    y[ys_index] = d2 * ((ql >> 4) as f32 + to_add) - m2;
1342                    ys_index += 1;
1343                }
1344                is += 2;
1345                u1 <<= 2;
1346                u2 <<= 2;
1347            }
1348        }
1349    }
1350}
1351
1352impl GgmlType for BlockQ6K {
1353    const DTYPE: GgmlDType = GgmlDType::Q6K;
1354    const BLCK_SIZE: usize = QK_K;
1355
1356    fn from_float(xs: &[f32], ys: &mut [Self]) {
1357        debug_assert_eq!(
1358            xs.len(),
1359            ys.len() * Self::BLCK_SIZE,
1360            "quantize_row_q6k: size mismatch {} {} {}",
1361            xs.len(),
1362            ys.len(),
1363            Self::BLCK_SIZE
1364        );
1365        let mut l = [0i8; QK_K];
1366        let mut scales = [0f32; QK_K / 16];
1367        let mut x = xs.as_ptr();
1368        let l = l.as_mut_ptr();
1369        unsafe {
1370            for y in ys.iter_mut() {
1371                let mut max_scale = 0f32;
1372                let mut max_abs_scale = 0f32;
1373                for (ib, scale_) in scales.iter_mut().enumerate() {
1374                    let scale =
1375                        make_qx_quants(16, 32, x.add(16 * ib), l.add(16 * ib), 1, std::ptr::null());
1376                    *scale_ = scale;
1377                    let abs_scale = scale.abs();
1378                    if abs_scale > max_abs_scale {
1379                        max_abs_scale = abs_scale;
1380                        max_scale = scale
1381                    }
1382                }
1383
1384                let iscale = -128f32 / max_scale;
1385                y.d = f16::from_f32(1.0 / iscale);
1386
1387                for (y_scale, scale) in y.scales.iter_mut().zip(scales.iter()) {
1388                    *y_scale = nearest_int(iscale * scale).min(127) as i8
1389                }
1390
1391                for (j, &y_scale) in y.scales.iter().enumerate() {
1392                    let d = y.d.to_f32() * y_scale as f32;
1393                    if d == 0. {
1394                        continue;
1395                    }
1396                    for ii in 0..16 {
1397                        let ll = nearest_int(*x.add(16 * j + ii) / d).clamp(-32, 31);
1398                        *l.add(16 * j + ii) = (ll + 32) as i8
1399                    }
1400                }
1401
1402                let mut ql = y.ql.as_mut_ptr();
1403                let mut qh = y.qh.as_mut_ptr();
1404
1405                for j in (0..QK_K).step_by(128) {
1406                    for l_idx in 0..32 {
1407                        let q1 = *l.add(j + l_idx) & 0xF;
1408                        let q2 = *l.add(j + l_idx + 32) & 0xF;
1409                        let q3 = *l.add(j + l_idx + 64) & 0xF;
1410                        let q4 = *l.add(j + l_idx + 96) & 0xF;
1411                        *ql.add(l_idx) = (q1 | (q3 << 4)) as u8;
1412                        *ql.add(l_idx + 32) = (q2 | (q4 << 4)) as u8;
1413                        *qh.add(l_idx) = ((*l.add(j + l_idx) >> 4)
1414                            | ((*l.add(j + l_idx + 32) >> 4) << 2)
1415                            | ((*l.add(j + l_idx + 64) >> 4) << 4)
1416                            | ((*l.add(j + l_idx + 96) >> 4) << 6))
1417                            as u8;
1418                    }
1419                    ql = ql.add(64);
1420                    qh = qh.add(32);
1421                }
1422
1423                x = x.add(QK_K)
1424            }
1425        }
1426    }
1427
1428    fn from_float_imatrix(xs: &[f32], ys: &mut [Self], imatrix_weights: &[f32], n_per_row: usize) {
1429        debug_assert_eq!(
1430            xs.len(),
1431            ys.len() * Self::BLCK_SIZE,
1432            "quantize_row_q6k imatrix: size mismatch {} {} {}",
1433            xs.len(),
1434            ys.len(),
1435            Self::BLCK_SIZE
1436        );
1437        let mut l = [0i8; QK_K];
1438        let mut scales = [0f32; QK_K / 16];
1439        let mut x = xs.as_ptr();
1440        let imatrix_weights = imatrix_weights.as_ptr();
1441        let l = l.as_mut_ptr();
1442        unsafe {
1443            for (sblk_idx, y) in ys.iter_mut().enumerate() {
1444                let mut max_scale = 0f32;
1445                let mut max_abs_scale = 0f32;
1446                for (ib, scale_) in scales.iter_mut().enumerate() {
1447                    let imatrix_row = sblk_idx % (n_per_row / QK_K);
1448                    let scale = make_qx_quants(
1449                        16,
1450                        32,
1451                        x.add(16 * ib),
1452                        l.add(16 * ib),
1453                        1,
1454                        imatrix_weights.add(QK_K * imatrix_row + 16 * ib),
1455                    );
1456                    *scale_ = scale;
1457                    let abs_scale = scale.abs();
1458                    if abs_scale > max_abs_scale {
1459                        max_abs_scale = abs_scale;
1460                        max_scale = scale
1461                    }
1462                }
1463
1464                let iscale = -128f32 / max_scale;
1465                y.d = f16::from_f32(1.0 / iscale);
1466
1467                for (y_scale, scale) in y.scales.iter_mut().zip(scales.iter()) {
1468                    *y_scale = nearest_int(iscale * scale).min(127) as i8
1469                }
1470
1471                for (j, &y_scale) in y.scales.iter().enumerate() {
1472                    let d = y.d.to_f32() * y_scale as f32;
1473                    if d == 0. {
1474                        continue;
1475                    }
1476                    for ii in 0..16 {
1477                        let ll = nearest_int(*x.add(16 * j + ii) / d).clamp(-32, 31);
1478                        *l.add(16 * j + ii) = (ll + 32) as i8
1479                    }
1480                }
1481
1482                let mut ql = y.ql.as_mut_ptr();
1483                let mut qh = y.qh.as_mut_ptr();
1484
1485                for j in (0..QK_K).step_by(128) {
1486                    for l_idx in 0..32 {
1487                        let q1 = *l.add(j + l_idx) & 0xF;
1488                        let q2 = *l.add(j + l_idx + 32) & 0xF;
1489                        let q3 = *l.add(j + l_idx + 64) & 0xF;
1490                        let q4 = *l.add(j + l_idx + 96) & 0xF;
1491                        *ql.add(l_idx) = (q1 | (q3 << 4)) as u8;
1492                        *ql.add(l_idx + 32) = (q2 | (q4 << 4)) as u8;
1493                        *qh.add(l_idx) = ((*l.add(j + l_idx) >> 4)
1494                            | ((*l.add(j + l_idx + 32) >> 4) << 2)
1495                            | ((*l.add(j + l_idx + 64) >> 4) << 4)
1496                            | ((*l.add(j + l_idx + 96) >> 4) << 6))
1497                            as u8;
1498                    }
1499                    ql = ql.add(64);
1500                    qh = qh.add(32);
1501                }
1502
1503                x = x.add(QK_K)
1504            }
1505        }
1506    }
1507
1508    // https://github.com/ggerganov/llama.cpp/blob/8183159cf3def112f6d1fe94815fce70e1bffa12/k_quants.c#L1067
1509    fn to_float(xs: &[Self], ys: &mut [f32]) {
1510        let k = ys.len();
1511        debug_assert!(
1512            k.is_multiple_of(QK_K),
1513            "dequantize_row_q6k: {k} is not divisible by {QK_K}"
1514        );
1515
1516        for (idx_x, x) in xs.iter().enumerate() {
1517            let d = x.d.to_f32();
1518            let ql = &x.ql;
1519            let qh = &x.qh;
1520            let sc = &x.scales;
1521            for n in (0..QK_K).step_by(128) {
1522                let idx = n / 128;
1523                let ys = &mut ys[idx_x * QK_K + n..];
1524                let sc = &sc[8 * idx..];
1525                let ql = &ql[64 * idx..];
1526                let qh = &qh[32 * idx..];
1527                for l in 0..32 {
1528                    let is = l / 16;
1529                    let q1 = ((ql[l] & 0xF) | ((qh[l] & 3) << 4)) as i8 - 32;
1530                    let q2 = ((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) as i8 - 32;
1531                    let q3 = ((ql[l] >> 4) | (((qh[l] >> 4) & 3) << 4)) as i8 - 32;
1532                    let q4 = ((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) as i8 - 32;
1533                    ys[l] = d * sc[is] as f32 * q1 as f32;
1534                    ys[l + 32] = d * sc[is + 2] as f32 * q2 as f32;
1535                    ys[l + 64] = d * sc[is + 4] as f32 * q3 as f32;
1536                    ys[l + 96] = d * sc[is + 6] as f32 * q4 as f32;
1537                }
1538            }
1539        }
1540    }
1541}
1542
1543impl GgmlType for BlockQ8K {
1544    const DTYPE: GgmlDType = GgmlDType::Q8K;
1545    const BLCK_SIZE: usize = QK_K;
1546
1547    fn from_float(xs: &[f32], ys: &mut [Self]) {
1548        let k = xs.len();
1549        debug_assert!(
1550            k.is_multiple_of(QK_K),
1551            "quantize_row_q8k: {k} is not divisible by {QK_K}"
1552        );
1553        for (i, y) in ys.iter_mut().enumerate() {
1554            let mut max = 0f32;
1555            let mut amax = 0f32;
1556            let xs = &xs[i * QK_K..(i + 1) * QK_K];
1557            for &x in xs.iter() {
1558                if amax < x.abs() {
1559                    amax = x.abs();
1560                    max = x;
1561                }
1562            }
1563            if amax == 0f32 {
1564                y.d = 0f32;
1565                y.qs.fill(0)
1566            } else {
1567                let iscale = -128f32 / max;
1568                for (j, q) in y.qs.iter_mut().enumerate() {
1569                    // ggml uses nearest_int with bit magic here, maybe we want the same
1570                    // but we would have to test and benchmark it.
1571                    let v = (iscale * xs[j]).round();
1572                    *q = v.min(127.) as i8
1573                }
1574                for j in 0..QK_K / 16 {
1575                    let mut sum = 0i32;
1576                    for ii in 0..16 {
1577                        sum += y.qs[j * 16 + ii] as i32
1578                    }
1579                    y.bsums[j] = sum as i16
1580                }
1581                y.d = 1.0 / iscale
1582            }
1583        }
1584    }
1585
1586    fn to_float(xs: &[Self], ys: &mut [f32]) {
1587        let k = ys.len();
1588        debug_assert!(
1589            k.is_multiple_of(QK_K),
1590            "dequantize_row_q8k: {k} is not divisible by {QK_K}"
1591        );
1592        for (i, x) in xs.iter().enumerate() {
1593            for (j, &q) in x.qs.iter().enumerate() {
1594                ys[i * QK_K + j] = x.d * q as f32
1595            }
1596        }
1597    }
1598}
1599
1600/// Plain f32 dot product, eight independent lanes so the compiler can vectorise it.
1601fn dot_f32(a: &[f32], b: &[f32]) -> f32 {
1602    let mut acc = [0f32; 8];
1603    let ((ca, ra), (cb, rb)) = (a.as_chunks::<8>(), b.as_chunks::<8>());
1604    let tail: f32 = ra.iter().zip(rb).map(|(x, y)| x * y).sum();
1605    for (xa, xb) in ca.iter().zip(cb) {
1606        for l in 0..8 {
1607            acc[l] += xa[l] * xb[l];
1608        }
1609    }
1610    acc.iter().sum::<f32>() + tail
1611}
1612
1613/// `dst[m, n] = lhs[m, k] x rhs_t[n, k]^T` where `rhs_t` holds quantized weight rows.
1614///
1615/// The one matmul route for quantized weights: every weight row is converted to f32 (`to_float`; the
1616/// block scale is applied to each value here, BEFORE any multiply) and dotted with the f32
1617/// activation row. There is no activation quantisation, no integer multiply, and no scaling of an
1618/// integer block sum: the integer dot kernels (Q8_0 x Q8_0, Q4_0 x Q8_0, the K-quant two-level
1619/// paths and their AVX2 and NEON versions) were removed from every published crate on 2026-11-01
1620/// (ROADMAP board 307).
1621pub fn matmul<T: GgmlType>(
1622    (m, k, n): (usize, usize, usize),
1623    lhs: &[f32],
1624    rhs_t: &[T],
1625    dst: &mut [f32],
1626) -> Result<()> {
1627    if m * k != lhs.len() || m * n != dst.len() {
1628        fuel_ir::bail!(
1629            "quantized matmul: bad lengths lhs {} dst {} for ({m},{k},{n})",
1630            lhs.len(),
1631            dst.len()
1632        );
1633    }
1634    let k_in_blocks = k.div_ceil(T::BLCK_SIZE);
1635    // Result laid out column-major (one chunk of `m` values per weight row) so each parallel task
1636    // owns a disjoint chunk; transposed into `dst` afterwards (a no-op copy for m == 1).
1637    let mut tmp = vec![0f32; m * n];
1638    tmp.par_chunks_mut(m)
1639        .enumerate()
1640        .with_min_len(64)
1641        .for_each_init(
1642            || vec![0f32; k_in_blocks * T::BLCK_SIZE],
1643            |buf, (col, out)| {
1644                T::to_float(&rhs_t[col * k_in_blocks..(col + 1) * k_in_blocks], buf);
1645                for (row, o) in out.iter_mut().enumerate() {
1646                    *o = dot_f32(&buf[..k], &lhs[row * k..(row + 1) * k]);
1647                }
1648            },
1649        );
1650    for col in 0..n {
1651        for row in 0..m {
1652            dst[row * n + col] = tmp[col * m + row];
1653        }
1654    }
1655    Ok(())
1656}
1657
1658/// [`matmul`] with an f16 activation and an f16 result (computed in f32).
1659pub fn matmul_f16<T: GgmlType>(
1660    mkn: (usize, usize, usize),
1661    lhs: &[f16],
1662    rhs_t: &[T],
1663    dst: &mut [f16],
1664) -> Result<()> {
1665    let (m, k, n) = mkn;
1666    if m * k != lhs.len() {
1667        fuel_ir::bail!("unexpected lhs length {} {mkn:?}", lhs.len());
1668    }
1669    let lhs32: Vec<f32> = lhs.iter().map(|x| x.to_f32()).collect();
1670    let mut dst32 = vec![0f32; m * n];
1671    matmul(mkn, &lhs32, rhs_t, &mut dst32)?;
1672    for (d, s) in dst.iter_mut().zip(dst32) {
1673        *d = f16::from_f32(s);
1674    }
1675    Ok(())
1676}
1677
1678impl GgmlType for f32 {
1679    const DTYPE: GgmlDType = GgmlDType::F32;
1680    const BLCK_SIZE: usize = 1;
1681
1682    fn from_float(xs: &[f32], ys: &mut [Self]) {
1683        debug_assert_eq!(
1684            xs.len(),
1685            ys.len(),
1686            "size mismatch xs {} != ys {}",
1687            xs.len(),
1688            ys.len()
1689        );
1690        ys.copy_from_slice(xs);
1691    }
1692
1693    fn to_float(xs: &[Self], ys: &mut [f32]) {
1694        debug_assert_eq!(
1695            xs.len(),
1696            ys.len(),
1697            "size mismatch xs {} != ys {}",
1698            xs.len(),
1699            ys.len()
1700        );
1701        ys.copy_from_slice(xs);
1702    }
1703}
1704
1705impl GgmlType for f16 {
1706    const DTYPE: GgmlDType = GgmlDType::F16;
1707    const BLCK_SIZE: usize = 1;
1708
1709    fn from_float(xs: &[f32], ys: &mut [Self]) {
1710        debug_assert_eq!(
1711            xs.len(),
1712            ys.len(),
1713            "size mismatch xs {} != ys {}",
1714            xs.len(),
1715            ys.len()
1716        );
1717        ys.convert_from_f32_slice(xs);
1718    }
1719
1720    fn to_float(xs: &[Self], ys: &mut [f32]) {
1721        debug_assert_eq!(
1722            xs.len(),
1723            ys.len(),
1724            "size mismatch xs {} != ys {}",
1725            xs.len(),
1726            ys.len()
1727        );
1728        xs.convert_to_f32_slice(ys);
1729    }
1730}
1731
1732impl GgmlType for bf16 {
1733    const DTYPE: GgmlDType = GgmlDType::BF16;
1734    const BLCK_SIZE: usize = 1;
1735
1736    fn from_float(xs: &[f32], ys: &mut [Self]) {
1737        debug_assert_eq!(
1738            xs.len(),
1739            ys.len(),
1740            "size mismatch xs {} != ys {}",
1741            xs.len(),
1742            ys.len()
1743        );
1744        ys.convert_from_f32_slice(xs);
1745    }
1746
1747    fn to_float(xs: &[Self], ys: &mut [f32]) {
1748        debug_assert_eq!(
1749            xs.len(),
1750            ys.len(),
1751            "size mismatch xs {} != ys {}",
1752            xs.len(),
1753            ys.len()
1754        );
1755        xs.convert_to_f32_slice(ys);
1756    }
1757}
1758
1759#[cfg(test)]
1760mod matmul_tests {
1761    use super::*;
1762
1763    /// A small deterministic matrix: values in about [-1, 1], no two rows alike.
1764    fn fill(len: usize, seed: u32) -> Vec<f32> {
1765        (0..len as u32)
1766            .map(|i| {
1767                let x = (i
1768                    .wrapping_mul(2654435761)
1769                    .wrapping_add(seed.wrapping_mul(40503)))
1770                    >> 8;
1771                (x % 2001) as f32 / 1000.0 - 1.0
1772            })
1773            .collect()
1774    }
1775
1776    /// The matmul is the f32 matmul of the DEQUANTISED weights with the f32 activations (so it
1777    /// agrees with a plain f32 reference to rounding), for several rows of activations at once.
1778    #[test]
1779    fn matmul_equals_the_f32_matmul_of_the_dequantised_weights() {
1780        let (m, k, n) = (3usize, 64usize, 8usize);
1781        let w = fill(n * k, 7);
1782        let x = fill(m * k, 11);
1783        let mut blocks = vec![BlockQ4_0::zeros(); n * k / QK4_0];
1784        BlockQ4_0::from_float(&w, &mut blocks);
1785        let mut deq = vec![0f32; n * k];
1786        BlockQ4_0::to_float(&blocks, &mut deq);
1787        let mut want = vec![0f32; m * n];
1788        for r in 0..m {
1789            for c in 0..n {
1790                want[r * n + c] = (0..k).map(|i| x[r * k + i] * deq[c * k + i]).sum();
1791            }
1792        }
1793        let mut got = vec![0f32; m * n];
1794        matmul((m, k, n), &x, &blocks, &mut got).expect("quantized matmul");
1795        assert!(
1796            want.iter().any(|v| v.abs() > 0.1),
1797            "control: the reference is not all zeros"
1798        );
1799        for (g, w) in got.iter().zip(&want) {
1800            assert!(
1801                (g - w).abs() <= 1e-3 * w.abs().max(1.0),
1802                "matmul {g} vs reference {w}"
1803            );
1804        }
1805    }
1806
1807    /// The f16 entry point is the f32 one with the activation rounded to f16 and the result rounded
1808    /// back, so it agrees with the f32 matmul to f16 precision.
1809    #[test]
1810    fn matmul_f16_agrees_with_matmul_to_f16_precision() {
1811        let (m, k, n) = (2usize, 64usize, 8usize);
1812        let w = fill(n * k, 3);
1813        let x = fill(m * k, 5);
1814        let mut blocks = vec![BlockQ4_0::zeros(); n * k / QK4_0];
1815        BlockQ4_0::from_float(&w, &mut blocks);
1816        let x16: Vec<f16> = x.iter().map(|v| f16::from_f32(*v)).collect();
1817        let x_back: Vec<f32> = x16.iter().map(|v| v.to_f32()).collect();
1818        let mut want = vec![0f32; m * n];
1819        matmul((m, k, n), &x_back, &blocks, &mut want).expect("matmul");
1820        let mut got = vec![f16::ZERO; m * n];
1821        matmul_f16((m, k, n), &x16, &blocks, &mut got).expect("matmul_f16");
1822        assert!(want.iter().any(|v| v.abs() > 0.1), "control: not all zeros");
1823        for (g, w) in got.iter().zip(&want) {
1824            assert!(
1825                (g.to_f32() - w).abs() <= 2e-3 * w.abs().max(1.0),
1826                "matmul_f16 {g} vs {w}"
1827            );
1828        }
1829    }
1830}
fuel_quantized/k_quants.rs

fuel_quantized/
k_quants.rs
