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

fuel_quantized/
cpu.rs
