fuel_quantized/
cpu.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/mod.rs`
3// at commit 549eacb011f25b2c6fc765d6c3dbd26e42f72ee6, modified here. Only IN PART: the `QuantizedType` trait and the per-format scalar ops behind it were extracted from there; the rest of this file is ours.
4// Contributors, how this was verified and the licence texts: THIRD-PARTY-NOTICES.md (section "candle").
5//! Format-agnostic CPU helpers for quantized storage.
6//!
7//! Holds the dyn-boxable `QuantizedType` trait (per-block-format scalar
8//! ops) plus `cpu_zeros` / `cpu_from_data` constructors used by the
9//! `fuel-cpu-backend`-side `CpuQStorage` adapter and by the file-format
10//! readers in `fuel-core/src/quantized/`.
11
12use crate::iq_grid_types::{
13    BlockIq1M, BlockIq1S, BlockIq2S, BlockIq2Xs, BlockIq2Xxs, BlockIq3S, BlockIq3Xxs,
14};
15use crate::iq4::{BlockIq4Nl, BlockIq4Xs};
16use crate::k_quants::{
17    self, BlockQ2K, BlockQ3K, BlockQ4_0, BlockQ4_1, BlockQ4K, BlockQ5_0, BlockQ5_1, BlockQ5K,
18    BlockQ6K, BlockQ8_0, BlockQ8_1, BlockQ8K, GgmlType,
19};
20use crate::tq_mx_types::{BlockMxfp4, BlockNvfp4, BlockTq1_0, BlockTq2_0};
21use fuel_ir::quantized::GgmlDType;
22use fuel_ir::{HostBuffer, Result};
23use half::{bf16, f16};
24use std::borrow::Cow;
25
26pub trait QuantizedType: Send + Sync {
27    fn dtype(&self) -> GgmlDType;
28    fn matmul_t(&self, mkn: (usize, usize, usize), lhs: &[f32], dst: &mut [f32]) -> Result<()>;
29    fn matmul_t_f16(&self, mkn: (usize, usize, usize), lhs: &[f16], dst: &mut [f16]) -> Result<()>;
30    fn dequantize(&self, elem_count: usize) -> Result<HostBuffer>;
31    fn storage_size_in_bytes(&self) -> usize;
32    fn as_ptr(&self) -> *const u8;
33    fn block_size(&self) -> usize;
34    #[allow(clippy::wrong_self_convention)]
35    fn from_float(&mut self, xs: &[f32]);
36    #[allow(clippy::wrong_self_convention)]
37    fn from_float_imatrix(&mut self, xs: &[f32], imatrix_weights: &[f32], n_per_row: usize);
38    fn size(&self) -> usize;
39}
40
41impl<T: GgmlType + Send + Sync> QuantizedType for Vec<T> {
42    fn matmul_t(&self, mkn: (usize, usize, usize), lhs: &[f32], dst: &mut [f32]) -> Result<()> {
43        k_quants::matmul(mkn, lhs, self.as_slice(), dst)
44    }
45    fn matmul_t_f16(&self, mkn: (usize, usize, usize), lhs: &[f16], dst: &mut [f16]) -> Result<()> {
46        k_quants::matmul_f16(mkn, lhs, self.as_slice(), dst)
47    }
48    fn size(&self) -> usize {
49        self.len() * core::mem::size_of::<T>()
50    }
51    fn from_float(&mut self, xs: &[f32]) {
52        T::from_float(xs, self)
53    }
54    fn from_float_imatrix(&mut self, xs: &[f32], imatrix_weights: &[f32], n_per_row: usize) {
55        T::from_float_imatrix(xs, self, imatrix_weights, n_per_row)
56    }
57    fn dtype(&self) -> GgmlDType {
58        T::DTYPE
59    }
60    fn block_size(&self) -> usize {
61        T::BLCK_SIZE
62    }
63    fn dequantize(&self, elem_count: usize) -> Result<HostBuffer> {
64        // GAP-333: without this, too many elements panicked on a block index
65        // and a partial block left a silently-zero tail (`to_float` only
66        // `debug_assert`s the block multiple).
67        T::DTYPE.check_dequant_count("dequantize", elem_count, self.storage_size_in_bytes())?;
68        let mut ys = vec![0.0f32; elem_count];
69        T::to_float(self.as_slice(), &mut ys);
70        Ok(HostBuffer::F32(ys))
71    }
72    fn storage_size_in_bytes(&self) -> usize {
73        self.len() * std::mem::size_of::<T>()
74    }
75    fn as_ptr(&self) -> *const u8 {
76        self.as_ptr() as *const u8
77    }
78}
79
80/// GAP-336: copies `data` into owned blocks, one `read_unaligned` per block.
81///
82/// This replaces `as_t_slice`, which took the `Cow` by value and returned a
83/// slice into it. For `Cow::Owned` that slice outlived the buffer it pointed
84/// into (a use-after-free, even with valid input), and it also panicked on a
85/// byte buffer that was not aligned for `T`. Nothing is borrowed here, so the
86/// result depends on neither the input's lifetime nor its alignment.
87fn blocks_from_bytes<T: GgmlType>(data: &[u8]) -> Result<Vec<T>> {
88    let size = std::mem::size_of::<T>();
89    if !data.len().is_multiple_of(size) {
90        return Err(fuel_ir::Error::Msg(format!(
91            "cpu_from_data: {} bytes is not a whole number of {:?} blocks ({size} bytes each)",
92            data.len(),
93            T::DTYPE,
94        ))
95        .bt());
96    }
97    Ok(data
98        .chunks_exact(size)
99        // Each chunk is exactly `size_of::<T>()` bytes and `T: Pod` (a
100        // supertrait of `GgmlType`), so any byte pattern is a valid `T`.
101        .map(bytemuck::pod_read_unaligned::<T>)
102        .collect())
103}
104
105pub fn cpu_zeros(dtype: GgmlDType, elem_count: usize) -> Box<dyn QuantizedType> {
106    match dtype {
107        GgmlDType::F32 => Box::new(vec![f32::zeros(); elem_count]),
108        GgmlDType::F16 => Box::new(vec![f16::zeros(); elem_count]),
109        GgmlDType::BF16 => Box::new(vec![bf16::zeros(); elem_count]),
110        GgmlDType::Q4_0 => Box::new(vec![BlockQ4_0::zeros(); elem_count / BlockQ4_0::BLCK_SIZE]),
111        GgmlDType::Q4_1 => Box::new(vec![BlockQ4_1::zeros(); elem_count / BlockQ4_1::BLCK_SIZE]),
112        GgmlDType::Q5_0 => Box::new(vec![BlockQ5_0::zeros(); elem_count / BlockQ5_0::BLCK_SIZE]),
113        GgmlDType::Q5_1 => Box::new(vec![BlockQ5_1::zeros(); elem_count / BlockQ5_1::BLCK_SIZE]),
114        GgmlDType::Q8_0 => Box::new(vec![BlockQ8_0::zeros(); elem_count / BlockQ8_0::BLCK_SIZE]),
115        GgmlDType::Q8_1 => Box::new(vec![BlockQ8_1::zeros(); elem_count / BlockQ8_1::BLCK_SIZE]),
116        GgmlDType::Q2K => Box::new(vec![BlockQ2K::zeros(); elem_count / BlockQ2K::BLCK_SIZE]),
117        GgmlDType::Q3K => Box::new(vec![BlockQ3K::zeros(); elem_count / BlockQ3K::BLCK_SIZE]),
118        GgmlDType::Q4K => Box::new(vec![BlockQ4K::zeros(); elem_count / BlockQ4K::BLCK_SIZE]),
119        GgmlDType::Q5K => Box::new(vec![BlockQ5K::zeros(); elem_count / BlockQ5K::BLCK_SIZE]),
120        GgmlDType::Q6K => Box::new(vec![BlockQ6K::zeros(); elem_count / BlockQ6K::BLCK_SIZE]),
121        GgmlDType::Q8K => Box::new(vec![BlockQ8K::zeros(); elem_count / BlockQ8K::BLCK_SIZE]),
122        GgmlDType::Iq4Nl => Box::new(vec![
123            BlockIq4Nl::zeros();
124            elem_count / BlockIq4Nl::BLCK_SIZE
125        ]),
126        GgmlDType::Iq4Xs => Box::new(vec![
127            BlockIq4Xs::zeros();
128            elem_count / BlockIq4Xs::BLCK_SIZE
129        ]),
130        GgmlDType::Iq2Xxs => Box::new(vec![
131            BlockIq2Xxs::zeros();
132            elem_count / BlockIq2Xxs::BLCK_SIZE
133        ]),
134        GgmlDType::Iq2Xs => Box::new(vec![
135            BlockIq2Xs::zeros();
136            elem_count / BlockIq2Xs::BLCK_SIZE
137        ]),
138        GgmlDType::Iq3Xxs => Box::new(vec![
139            BlockIq3Xxs::zeros();
140            elem_count / BlockIq3Xxs::BLCK_SIZE
141        ]),
142        GgmlDType::Iq1S => Box::new(vec![BlockIq1S::zeros(); elem_count / BlockIq1S::BLCK_SIZE]),
143        GgmlDType::Iq3S => Box::new(vec![BlockIq3S::zeros(); elem_count / BlockIq3S::BLCK_SIZE]),
144        GgmlDType::Iq2S => Box::new(vec![BlockIq2S::zeros(); elem_count / BlockIq2S::BLCK_SIZE]),
145        GgmlDType::Iq1M => Box::new(vec![BlockIq1M::zeros(); elem_count / BlockIq1M::BLCK_SIZE]),
146        GgmlDType::Tq1_0 => Box::new(vec![
147            BlockTq1_0::zeros();
148            elem_count / BlockTq1_0::BLCK_SIZE
149        ]),
150        GgmlDType::Tq2_0 => Box::new(vec![
151            BlockTq2_0::zeros();
152            elem_count / BlockTq2_0::BLCK_SIZE
153        ]),
154        GgmlDType::Mxfp4 => Box::new(vec![
155            BlockMxfp4::zeros();
156            elem_count / BlockMxfp4::BLCK_SIZE
157        ]),
158        GgmlDType::Nvfp4 => Box::new(vec![
159            BlockNvfp4::zeros();
160            elem_count / BlockNvfp4::BLCK_SIZE
161        ]),
162    }
163}
164
165/// Builds CPU quantized storage from raw block bytes. The bytes are copied,
166/// so `data` may be owned or borrowed and need not be aligned for `dtype`.
167/// Returns `Err` if `data` is not a whole number of blocks.
168pub fn cpu_from_data(dtype: GgmlDType, data: Cow<'_, [u8]>) -> Result<Box<dyn QuantizedType>> {
169    let data: &[u8] = &data;
170    Ok(match dtype {
171        GgmlDType::F32 => Box::new(blocks_from_bytes::<f32>(data)?),
172        GgmlDType::F16 => Box::new(blocks_from_bytes::<f16>(data)?),
173        GgmlDType::BF16 => Box::new(blocks_from_bytes::<bf16>(data)?),
174        GgmlDType::Q4_0 => Box::new(blocks_from_bytes::<BlockQ4_0>(data)?),
175        GgmlDType::Q4_1 => Box::new(blocks_from_bytes::<BlockQ4_1>(data)?),
176        GgmlDType::Q5_0 => Box::new(blocks_from_bytes::<BlockQ5_0>(data)?),
177        GgmlDType::Q5_1 => Box::new(blocks_from_bytes::<BlockQ5_1>(data)?),
178        GgmlDType::Q8_0 => Box::new(blocks_from_bytes::<BlockQ8_0>(data)?),
179        GgmlDType::Q8_1 => Box::new(blocks_from_bytes::<BlockQ8_1>(data)?),
180        GgmlDType::Q2K => Box::new(blocks_from_bytes::<BlockQ2K>(data)?),
181        GgmlDType::Q3K => Box::new(blocks_from_bytes::<BlockQ3K>(data)?),
182        GgmlDType::Q4K => Box::new(blocks_from_bytes::<BlockQ4K>(data)?),
183        GgmlDType::Q5K => Box::new(blocks_from_bytes::<BlockQ5K>(data)?),
184        GgmlDType::Q6K => Box::new(blocks_from_bytes::<BlockQ6K>(data)?),
185        GgmlDType::Q8K => Box::new(blocks_from_bytes::<BlockQ8K>(data)?),
186        GgmlDType::Iq4Nl => Box::new(blocks_from_bytes::<BlockIq4Nl>(data)?),
187        GgmlDType::Iq4Xs => Box::new(blocks_from_bytes::<BlockIq4Xs>(data)?),
188        GgmlDType::Iq2Xxs => Box::new(blocks_from_bytes::<BlockIq2Xxs>(data)?),
189        GgmlDType::Iq2Xs => Box::new(blocks_from_bytes::<BlockIq2Xs>(data)?),
190        GgmlDType::Iq3Xxs => Box::new(blocks_from_bytes::<BlockIq3Xxs>(data)?),
191        GgmlDType::Iq1S => Box::new(blocks_from_bytes::<BlockIq1S>(data)?),
192        GgmlDType::Iq3S => Box::new(blocks_from_bytes::<BlockIq3S>(data)?),
193        GgmlDType::Iq2S => Box::new(blocks_from_bytes::<BlockIq2S>(data)?),
194        GgmlDType::Iq1M => Box::new(blocks_from_bytes::<BlockIq1M>(data)?),
195        GgmlDType::Tq1_0 => Box::new(blocks_from_bytes::<BlockTq1_0>(data)?),
196        GgmlDType::Tq2_0 => Box::new(blocks_from_bytes::<BlockTq2_0>(data)?),
197        GgmlDType::Mxfp4 => Box::new(blocks_from_bytes::<BlockMxfp4>(data)?),
198        GgmlDType::Nvfp4 => Box::new(blocks_from_bytes::<BlockNvfp4>(data)?),
199    })
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    fn f32_bytes(xs: &[f32]) -> Vec<u8> {
207        xs.iter().flat_map(|x| x.to_le_bytes()).collect()
208    }
209
210    fn dequant(q: &dyn QuantizedType, n: usize) -> Vec<f32> {
211        q.dequantize(n).unwrap().as_slice::<f32>().unwrap().to_vec()
212    }
213
214    /// `Cow::Owned` is the input `as_t_slice` turned into a dangling slice.
215    /// This pins the values; the use-after-free itself is not observable from
216    /// a plain test (it needs Miri), so the evidence for it is the lifetime
217    /// argument in `blocks_from_bytes`'s doc.
218    #[test]
219    fn owned_f32_bytes_round_trip() {
220        let xs = [1.5_f32, -2.0, 3.25];
221        let q = cpu_from_data(GgmlDType::F32, Cow::Owned(f32_bytes(&xs))).unwrap();
222        assert_eq!(dequant(q.as_ref(), 3), xs);
223    }
224
225    #[test]
226    fn owned_q4_0_blocks_round_trip() {
227        // Two hand-built blocks, checked against the format's own rule
228        // (`value = (nibble - 8) * d`, low nibbles first), so the test trusts
229        // neither the quantizer nor a raw-pointer read.
230        let mut bytes = Vec::new();
231        for (d, q) in [(1.0_f32, 0x3A_u8), (0.5, 0x81)] {
232            bytes.extend_from_slice(&f16::from_f32(d).to_le_bytes());
233            bytes.extend(std::iter::repeat_n(q, 16));
234        }
235        let back = cpu_from_data(GgmlDType::Q4_0, Cow::Owned(bytes)).unwrap();
236        // Block 0: low 0xA -> 2, high 0x3 -> -5, scale 1. Block 1: low 0x1 -> -7,
237        // high 0x8 -> 0, scale 0.5.
238        let want: Vec<f32> = [2.0, -5.0, -3.5, 0.0]
239            .iter()
240            .flat_map(|&v| std::iter::repeat_n(v, 16))
241            .collect();
242        assert_eq!(dequant(back.as_ref(), 64), want);
243    }
244
245    /// GAP-333: `dequantize` checks the count against the stored blocks.
246    #[test]
247    fn dequantize_declines_a_count_the_blocks_cannot_supply() {
248        let q = cpu_zeros(GgmlDType::Q4_0, 64); // 2 blocks, 36 bytes
249        assert_eq!(dequant(q.as_ref(), 64), vec![0.0; 64]);
250        let too_many = q.dequantize(96).expect_err("3 blocks from 2").to_string();
251        assert!(
252            too_many.contains("96 elements need 54 bytes of Q4_0 blocks, but the buffer holds 36"),
253            "{too_many}"
254        );
255        let partial = q.dequantize(40).expect_err("a partial block").to_string();
256        assert!(
257            partial.contains("element count 40 is not a multiple of Q4_0's block size 32"),
258            "{partial}"
259        );
260    }
261
262    #[test]
263    fn a_partial_trailing_block_is_declined() {
264        let err = cpu_from_data(GgmlDType::Q4_0, Cow::Owned(vec![0_u8; 18 + 5]))
265            .err()
266            .expect("a partial block must be declined")
267            .to_string();
268        assert!(
269            err.contains("23 bytes is not a whole number of Q4_0 blocks (18 bytes each)"),
270            "wrong decline: {err}"
271        );
272    }
273
274    /// `as_t_slice` panicked on this input ("Data pointer must be aligned").
275    #[test]
276    fn misaligned_bytes_are_accepted() {
277        let xs = [0.5_f32, 7.0, -1.25, 2.0];
278        let body = f32_bytes(&xs);
279        let mut buf = vec![0_u8; 3 + body.len()];
280        let off = (0..3)
281            .find(|o| !(buf.as_ptr() as usize + o).is_multiple_of(std::mem::align_of::<f32>()))
282            .expect("one of three offsets is misaligned for f32");
283        buf[off..off + body.len()].copy_from_slice(&body);
284        let bytes = &buf[off..off + body.len()];
285        assert!(!(bytes.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()));
286        let q = cpu_from_data(GgmlDType::F32, Cow::Borrowed(bytes)).unwrap();
287        assert_eq!(dequant(q.as_ref(), 4), xs);
288    }
289}
fuel_quantized/cpu.rs

fuel_quantized/
cpu.rs
