fuel_quantized/
cpu.rs

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

fuel_quantized/
cpu.rs
