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

fuel_quantized/
cpu.rs
