burn_std/tensor/
quantization.rs

1//! Quantization data representation.
2
3// Re-exported types
4pub use cubecl_common::quant::scheme::{
5    BlockSize, QuantLevel, QuantMode, QuantParam, QuantScheme, QuantStore, QuantValue,
6};
7
8/// Alignment (in bytes) for quantization parameters in serialized tensor data.
9///
10/// NOTE: This is currently f32-based since scales were originally always f32.
11/// With `QuantParam` now supporting different precisions (F16, BF16, etc.),
12/// this alignment may need to be revisited in the future.
13pub const QPARAM_ALIGN: usize = core::mem::align_of::<f32>();
14
15use alloc::vec::Vec;
16use core::any::TypeId;
17use num_traits::PrimInt;
18use serde::{Deserialize, Serialize};
19
20use crate::{DType, Metadata, Shape, bytes::Bytes};
21
22/// Configuration for a device quantization behavior.
23///
24/// This configuration determines how tensors are quantized and how quantization rules
25/// propagate through operations on a given device. It is applied once during device
26/// initialization. See also the [device settings](crate::DeviceSettings).
27#[derive(new, Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
28pub struct QuantConfig {
29    /// Defines how a tensor is quantized.
30    pub scheme: QuantScheme,
31    /// How quantization is propagated during computation.
32    pub propagation: QuantPropagation,
33    // NOTE: accumulation is currently unused, only scheme and propagation have an impact
34    // /// The precision used for the accumulation in various kernels.
35    // pub acc: QuantAcc,
36}
37
38#[derive(
39    Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default,
40)]
41/// The precision of accumulating elements.
42pub enum QuantAcc {
43    /// Full precision.
44    #[default]
45    F32,
46    /// Half precision.
47    F16,
48    /// bfloat16 precision.
49    BF16,
50}
51
52/// Calibration method used to compute the quantization range mapping.
53pub enum Calibration {
54    /// Computes quantization range mapping based on the min and max values.
55    MinMax,
56    /// Absolute-mean calibration for BitNet b1.58-style `{-1, 0, +1}` weight quantization.
57    ///
58    /// The range is `[-γ, +γ]` where γ = `mean(|W|)` per tensor or per block (BitNet b1.58
59    /// §3.1). Use with `QuantValue::Q2S` and `QuantStore::PackedU32` for 2-bit packed storage.
60    AbsMean,
61}
62
63/// Specify if the output of an operation is quantized using the scheme of the input
64/// or returned unquantized.
65#[derive(
66    Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default,
67)]
68pub enum QuantPropagation {
69    /// The output is quantized using the scheme of the input.
70    Propagate,
71    /// The output is not quantized.
72    #[default]
73    Inhibit,
74}
75
76/// The quantization tensor data parameters.
77#[derive(Clone, Debug)]
78pub struct QParams<S> {
79    /// The scaling factor.
80    pub scales: S,
81}
82
83/// A quantization parameter tensor descriptor.
84#[derive(Debug, Clone, PartialEq, Eq)]
85pub struct QParamTensor {
86    /// Start of the tensor in the buffer
87    pub offset_start: usize,
88    /// Offset of tensor end from the end of the buffer
89    pub offset_end: usize,
90    /// Metadata of the tensor
91    pub metadata: Metadata,
92    /// Data type of the tensor
93    pub dtype: DType,
94}
95
96/// Calculate the shape of the quantization parameters for a given tensor and level
97pub fn params_shape(data_shape: &Shape, level: QuantLevel) -> Shape {
98    match level {
99        QuantLevel::Tensor => Shape::new([1]),
100        QuantLevel::Block(block_size) => {
101            let mut params_shape = data_shape.clone();
102            let block_size = block_size.to_dim_vec(data_shape.num_dims());
103
104            for (shape, block_size) in params_shape.iter_mut().zip(block_size) {
105                *shape = (*shape).div_ceil(block_size as usize);
106            }
107
108            params_shape
109        }
110    }
111}
112
113/// Quantized data bytes representation.
114///
115/// # Notes
116/// 1) The quantized values are packed into 32-bit unsigned integers. For example, int8
117///    quantized values pack 4 grouped values into a single `u32`. When unpacking these values,
118///    we make sure to retrieve only the meaningful values (and ignore the alignment padding).
119/// 2) Quantization parameters are appended to the tensor data.
120///    As such, the last bytes always correspond to the scale parameter.
121///    If the quantization scheme includes an offset (zero-point) parameter, it is next to last.
122pub struct QuantizedBytes {
123    /// The quantized values and quantization parameters represented as bytes.
124    pub bytes: Bytes,
125    /// The quantization scheme.
126    pub scheme: QuantScheme,
127    /// The number of quantized elements.
128    pub num_elements: usize,
129}
130
131impl QuantizedBytes {
132    /// Creates a new quantized bytes representation.
133    pub fn new<E: bytemuck::CheckedBitPattern + bytemuck::NoUninit>(
134        value: Vec<E>,
135        scheme: QuantScheme,
136        scales: &[f32],
137    ) -> Self {
138        let num_elements = value.len();
139        // Only used for 8-bit quantization data comparison in tests
140        if TypeId::of::<E>() != TypeId::of::<i8>() {
141            panic!("Invalid quantized type");
142        }
143
144        // Re-interpret `Vec<E>` as `Vec<i8>` with `Vec::from_raw_parts`
145        let i8s: Vec<i8> = bytemuck::allocation::cast_vec(value);
146        let mut bytes = Bytes::from_elems(i8s);
147
148        let scales = match scheme.level {
149            QuantLevel::Tensor => &scales[..1],
150            QuantLevel::Block(_block_size) => scales,
151        };
152        let scale_bytes = encode_scales(scales, scheme.param);
153        bytes.extend_from_byte_slice_aligned(scale_bytes.as_slice(), QPARAM_ALIGN);
154
155        Self {
156            bytes,
157            scheme,
158            num_elements,
159        }
160    }
161
162    /// Returns the int8 quantized values with the quantization parameters.
163    pub fn into_vec_i8(self) -> (Vec<i8>, QParams<Vec<f32>>) {
164        let param = self.scheme.param;
165        let (values, (qparams, num_params)) = self.split_values_off();
166
167        // Quantization parameters are added at the end of the tensor data.
168        // As such, the last bytes always correspond to the scale parameter(s),
169        // stored at the scheme's param dtype. For example, per-block
170        // quantization can have multiple parameters for a single tensor:
171        // [scale, scale, scale, ...]
172        let scales_size = scale_size(param) * num_params;
173        let scales = decode_scales(&qparams[qparams.len() - scales_size..], param);
174
175        (values, QParams { scales })
176    }
177
178    fn split_i8_values(self, scale_bytes: usize) -> (Vec<i8>, Vec<u8>) {
179        let mut values = read_bytes_to_i8(self.bytes);
180
181        let values_end = values.len() - scale_bytes;
182        let qparams = values.split_off(values_end);
183
184        (values, bytemuck::cast_vec(qparams))
185    }
186
187    /// Splits the quantized values of the tensor from the quantization parameters.
188    ///
189    /// Returns the values in i8 and a newly allocated vector containing the
190    /// quantization parameter bytes.
191    fn split_values_off(self) -> (Vec<i8>, (Vec<u8>, usize)) {
192        let num_params = match self.scheme.level {
193            QuantLevel::Tensor => 1,
194            QuantLevel::Block(block_size) => self.num_elements / block_size.num_elements(),
195        };
196        let scale_bytes = scale_size(self.scheme.param) * num_params;
197
198        if let QuantStore::PackedU32(packed_dim) = self.scheme.store {
199            assert_eq!(
200                packed_dim, 0,
201                "Packing must be on innermost dimension for splitting off values"
202            );
203        }
204
205        let (values, qparams) = match self.scheme.store {
206            QuantStore::Native => self.split_i8_values(scale_bytes),
207            QuantStore::PackedU32(_) => match self.scheme.value {
208                QuantValue::Q8F | QuantValue::Q8S => self.split_i8_values(scale_bytes),
209                QuantValue::Q4F | QuantValue::Q4S | QuantValue::Q2F | QuantValue::Q2S => {
210                    let split_at = self.bytes.len() - scale_bytes;
211                    let qparams = self.bytes[split_at..].to_vec();
212                    let values = bytemuck::cast_slice::<_, u32>(&self.bytes[..split_at]);
213                    // Sub-byte values are unpacked as i8s for value equality tests
214                    let values = unpack_q_to_i8s(values, self.num_elements, &self.scheme.value);
215                    (values, qparams)
216                }
217                QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1 => {
218                    unimplemented!("Not yet supported")
219                }
220            },
221            QuantStore::PackedNative(_) => unimplemented!("Not yet supported"),
222        };
223
224        (values, (qparams, num_params))
225    }
226}
227
228/// Bytes per stored scale entry for the given param dtype.
229fn scale_size(param: QuantParam) -> usize {
230    match param {
231        QuantParam::F32 => 4,
232        QuantParam::F16 | QuantParam::BF16 => 2,
233        QuantParam::UE8M0 | QuantParam::UE4M3 => 1,
234    }
235}
236
237/// Decode stored scale entries into f32.
238fn decode_scales(bytes: &[u8], param: QuantParam) -> Vec<f32> {
239    match param {
240        QuantParam::F32 => bytes
241            .chunks_exact(4)
242            .map(|c| f32::from_ne_bytes([c[0], c[1], c[2], c[3]]))
243            .collect(),
244        QuantParam::F16 => bytes
245            .chunks_exact(2)
246            .map(|c| crate::f16::from_ne_bytes([c[0], c[1]]).to_f32())
247            .collect(),
248        QuantParam::BF16 => bytes
249            .chunks_exact(2)
250            .map(|c| crate::bf16::from_ne_bytes([c[0], c[1]]).to_f32())
251            .collect(),
252        QuantParam::UE8M0 | QuantParam::UE4M3 => unimplemented!("Not yet supported"),
253    }
254}
255
256/// Encode f32 scales at the param dtype for serialization.
257fn encode_scales(scales: &[f32], param: QuantParam) -> Vec<u8> {
258    match param {
259        QuantParam::F32 => scales.iter().flat_map(|s| s.to_ne_bytes()).collect(),
260        QuantParam::F16 => scales
261            .iter()
262            .flat_map(|s| crate::f16::from_f32(*s).to_ne_bytes())
263            .collect(),
264        QuantParam::BF16 => scales
265            .iter()
266            .flat_map(|s| crate::bf16::from_f32(*s).to_ne_bytes())
267            .collect(),
268        QuantParam::UE8M0 | QuantParam::UE4M3 => unimplemented!("Not yet supported"),
269    }
270}
271
272fn read_bytes_to_i8(bytes: Bytes) -> Vec<i8> {
273    match bytes.try_into_vec::<i8>() {
274        Ok(val) => val,
275        // Safety,
276        //
277        // `Vec<u8>` can be Re-interpreted as `Vec<i8>` since they share the same alignment.
278        Err(bytes) => unsafe { core::mem::transmute::<Vec<u8>, Vec<i8>>(bytes.to_vec()) },
279    }
280}
281
282/// Pack signed 8-bit integer values into a sequence of unsigned 32-bit integers.
283pub fn pack_i8s_to_u32s(values: Vec<i8>) -> Vec<u32> {
284    // Shift and combine groups of four 8-bit values into a u32.
285    // Same as doing this:
286    //     let result = (d_u8 & 0xFF) << 24 | (c_u8 & 0xFF) << 16 | (b_u8 & 0xFF) << 8 | (a_u8 & 0xFF);
287    #[cfg(target_endian = "big")]
288    {
289        values
290            .chunks(4)
291            .map(|x| {
292                x.iter()
293                    .enumerate()
294                    .fold(0u32, |acc, (i, x)| acc | (*x as u32 & 0xFF) << (i * 8))
295            })
296            .collect()
297    }
298
299    // The order of bytes in little endian matches the above description, we just need to
300    // handle padding when the number of values is not a factor of 4
301    #[cfg(target_endian = "little")]
302    {
303        let mut values = values;
304        let remainder = values.len() % 4;
305        if remainder != 0 {
306            // Pad with zeros
307            values.extend(core::iter::repeat_n(0, 4 - remainder));
308        }
309
310        let len = values.len() / 4;
311        let capacity = values.capacity() / 4;
312
313        // Pre-forget the old vec and re-interpret as u32
314        let mut values = core::mem::ManuallyDrop::new(values);
315        let ptr = values.as_mut_ptr() as *mut u32;
316
317        unsafe { Vec::from_raw_parts(ptr, len, capacity) }
318    }
319}
320
321/// Unpack integer values into a sequence of signed 8-bit integers.
322pub(crate) fn unpack_q_to_i8s<Q: PrimInt>(
323    values: &[Q],
324    numel: usize,
325    value: &QuantValue,
326) -> Vec<i8> {
327    let size_store = size_of::<Q>() * 8;
328    let size_quant = value.size_bits();
329    let num_quants = size_store / size_quant;
330    let mask = Q::from((1 << size_quant) - 1).unwrap();
331    let sign_shift = 8 - size_quant; // sign extension for sub-byte values
332    values
333        .iter()
334        .enumerate()
335        .flat_map(|(i, &packed)| {
336            // A single u32 could contain less than four 8-bit values...
337            let n = core::cmp::min(num_quants, numel - i * num_quants);
338            // Extract each 8-bit segment from u32 and cast back to i8
339            // Same as doing this (when 4 values are fully packed):
340            //     let a = (packed & 0xFF) as i8;
341            //     let b = ((packed >> 8) & 0xFF) as i8;
342            //     let c = ((packed >> 16) & 0xFF) as i8;
343            //     let d = ((packed >> 24) & 0xFF) as i8;
344            (0..n).map(move |i| {
345                let raw = (packed >> (i * size_quant) & mask).to_u8().unwrap();
346                ((raw << sign_shift) as i8) >> sign_shift
347            })
348        })
349        .collect()
350}
351
352#[cfg(test)]
353mod tests {
354
355    use super::*;
356    use alloc::vec;
357
358    #[test]
359    fn should_pack_i8s_to_u32() {
360        let packed = pack_i8s_to_u32s(vec![-128, 2, -3, 127]);
361
362        assert_eq!(packed, vec![2147287680]);
363    }
364
365    #[test]
366    fn should_pack_i8s_to_u32_padded() {
367        let packed = pack_i8s_to_u32s(vec![-128, 2, -3, 127, 55]);
368        let packed_padded = pack_i8s_to_u32s(vec![-128, 2, -3, 127, 55, 0, 0, 0]);
369
370        assert_eq!(packed, vec![2147287680, 55]);
371        assert_eq!(packed, packed_padded);
372    }
373
374    #[test]
375    fn should_unpack_u32s_to_i8s() {
376        let unpacked = unpack_q_to_i8s(&[2147287680u32], 4, &QuantValue::Q8S);
377
378        assert_eq!(unpacked, vec![-128, 2, -3, 127]);
379    }
380
381    #[test]
382    fn should_unpack_u32s_to_i8s_padded() {
383        let unpacked = unpack_q_to_i8s(&[55u32], 1, &QuantValue::Q8S);
384
385        assert_eq!(unpacked, vec![55]);
386    }
387
388    #[test]
389    fn should_unpack_u32s_to_i8s_arange() {
390        let unpacked = unpack_q_to_i8s(
391            &[
392                0u32, 286331136, 286331153, 572657937, 572662306, 857874978, 858993459, 858993459,
393                1145324612, 1145324612, 1431655748, 1431655765, 1717982549, 1717986918, 2003199590,
394                2004318071,
395            ],
396            128,
397            &QuantValue::Q4S,
398        );
399
400        assert_eq!(
401            unpacked,
402            vec![
403                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
404                2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
405                3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5,
406                5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6,
407                6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7
408            ]
409        );
410    }
411
412    #[test]
413    fn should_pack_unpack_quantization_parameters_per_tensor_symmetric() {
414        // Quantized [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
415        let scale = 0.03937008;
416        let values = vec![0i8, 25, 51, 76, 102, 127];
417
418        let q_bytes = QuantizedBytes::new(
419            values.clone(),
420            QuantScheme::default()
421                .with_value(QuantValue::Q8S)
422                .with_store(QuantStore::Native),
423            &[scale],
424        );
425
426        let (q_values, qparams) = q_bytes.into_vec_i8();
427
428        assert_eq!(qparams.scales, vec![scale]);
429
430        assert_eq!(q_values, values);
431    }
432}
burn_std/tensor/quantization.rs

burn_std/tensor/
quantization.rs
