burn_std/tensor/
quantization.rs

1//! Quantization data representation.
2
3// Re-exported types
4pub use cubecl_common::quant::scheme::{
5    BlockScale, BlockSize, QuantMode, QuantScheme, QuantStore, QuantValue, ScaleDtype,
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 `ScaleDtype` 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 cubecl_common::e4m3;
18use cubecl_common::quant::scheme::{f32_to_ue8m0, ue8m0_to_f32};
19use num_traits::PrimInt;
20use serde::{Deserialize, Serialize};
21
22use crate::{DType, Metadata, Shape, bytes::Bytes};
23
24/// Configuration for a device quantization behavior.
25///
26/// This configuration determines how tensors are quantized and how quantization rules
27/// propagate through operations on a given device. It is applied once during device
28/// initialization. See also the [device settings](crate::DeviceSettings).
29#[derive(new, Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
30pub struct QuantConfig {
31    /// Defines how a tensor is quantized.
32    pub scheme: QuantScheme,
33    /// How quantization is propagated during computation.
34    pub propagation: QuantPropagation,
35    // NOTE: accumulation is currently unused, only scheme and propagation have an impact
36    // /// The precision used for the accumulation in various kernels.
37    // pub acc: QuantAcc,
38}
39
40#[derive(
41    Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default,
42)]
43/// The precision of accumulating elements.
44pub enum QuantAcc {
45    /// Full precision.
46    #[default]
47    F32,
48    /// Half precision.
49    F16,
50    /// bfloat16 precision.
51    BF16,
52}
53
54/// Calibration method used to compute the quantization range mapping.
55pub enum Calibration {
56    /// Computes quantization range mapping based on the min and max values.
57    MinMax,
58    /// Absolute-mean calibration for BitNet b1.58-style `{-1, 0, +1}` weight quantization.
59    ///
60    /// The range is `[-γ, +γ]` where γ = `mean(|W|)` per tensor or per block (BitNet b1.58
61    /// §3.1). Use with `QuantValue::Q2S` and `QuantStore::PackedU32` for 2-bit packed storage.
62    AbsMean,
63}
64
65/// Specify if the output of an operation is quantized using the scheme of the input
66/// or returned unquantized.
67#[derive(
68    Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default,
69)]
70pub enum QuantPropagation {
71    /// The output is quantized using the scheme of the input.
72    Propagate,
73    /// The output is not quantized.
74    #[default]
75    Inhibit,
76}
77
78/// The quantization tensor data parameters.
79#[derive(Clone, Debug)]
80pub struct QParams<S> {
81    /// The scaling factor.
82    pub scales: S,
83    /// The per-tensor scale [`scales`](Self::scales) are relative to, for a two-level scheme.
84    pub global: Option<S>,
85}
86
87/// Scales recovered from a quantized byte buffer.
88#[derive(Clone, Debug, PartialEq)]
89pub struct DecodedScales {
90    /// One scale per block, or a single entry for a per-tensor level.
91    pub block: Vec<f32>,
92    /// The per-tensor scale, for a level that carries one.
93    pub global: Option<f32>,
94}
95
96/// A quantization parameter tensor descriptor.
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct QParamTensor {
99    /// Start of the tensor in the buffer
100    pub offset_start: usize,
101    /// Offset of tensor end from the end of the buffer
102    pub offset_end: usize,
103    /// Metadata of the tensor
104    pub metadata: Metadata,
105    /// Data type of the tensor
106    pub dtype: DType,
107}
108
109/// Whether a backend can quantize against this scheme's scales.
110///
111/// A backend answers `supports_dtype` with this, so an unsupported scheme is declined where it is
112/// chosen rather than at the first quantize. Each condition is asserted again where it is relied
113/// on, so bypassing this reports the specific rule rather than this general one.
114pub fn quantizable(scheme: &QuantScheme) -> bool {
115    match (scheme.block_scale(), global_scale_dtype(scheme)) {
116        (Some(block), Some(global)) => {
117            // The per-tensor scale is the largest block scale over the block dtype's maximum, so a
118            // block dtype reaching f32's range drives it subnormal. Any precision the per-tensor
119            // scale itself loses becomes a mismatch applied to every block.
120            block.dtype.max_representable() <= crate::f16::MAX.to_f32() && global == ScaleDtype::F32
121        }
122        _ => true,
123    }
124}
125
126/// The dtype of the per-tensor scale block scales are normalized against, for a two-level scheme.
127///
128/// [`QuantScheme::tensor_scale`] answers for a per-tensor scheme too, where the scale is the whole
129/// reconstruction rather than a factor over block scales.
130pub fn global_scale_dtype(scheme: &QuantScheme) -> Option<ScaleDtype> {
131    scheme.block_scale().and(scheme.tensor_scale())
132}
133
134/// Calculate the shape of the block scale grid for a given tensor and scheme.
135///
136/// A two-level scheme's per-tensor scale is not part of this grid.
137pub fn params_shape(data_shape: &Shape, scheme: &QuantScheme) -> Shape {
138    match scheme.block_size() {
139        None => Shape::new([1]),
140        Some(block_size) => Shape::from(block_size.num_blocks(data_shape.as_slice())),
141    }
142}
143
144/// Which block each element of a tensor falls in, for a block scheme. A block is a rectangle, so
145/// its members are a run of the flat storage only when it spans the trailing dimension; anything
146/// walking values against block scales asks here rather than chunking.
147#[derive(Debug, Clone)]
148pub struct BlockLayout {
149    shape: Shape,
150    block: Vec<u8>,
151    blocks: Shape,
152}
153
154impl BlockLayout {
155    /// How `block` tiles a tensor of `shape`.
156    pub fn new(shape: &Shape, block: &BlockSize) -> Self {
157        Self {
158            shape: shape.clone(),
159            block: block.to_dim_vec(shape.num_dims()),
160            blocks: Shape::from(block.num_blocks(shape.as_slice())),
161        }
162    }
163
164    /// How many blocks the tensor holds, which is how many block scales it has.
165    pub fn num_blocks(&self) -> usize {
166        self.blocks.num_elements()
167    }
168
169    /// Whether every dimension is a whole number of blocks.
170    pub fn divides(&self) -> bool {
171        self.shape
172            .iter()
173            .zip(&self.block)
174            .all(|(&dim, &extent)| dim.is_multiple_of(extent as usize))
175    }
176
177    /// The row-major index of the block holding row-major element `index`.
178    pub fn block_of(&self, mut index: usize) -> usize {
179        let mut block = 0;
180        let mut stride = 1;
181        for dim in (0..self.shape.num_dims()).rev() {
182            let coordinate = index % self.shape[dim];
183            index /= self.shape[dim];
184            block += coordinate / self.block[dim] as usize * stride;
185            stride *= self.blocks[dim];
186        }
187        block
188    }
189}
190
191/// Quantized data bytes representation.
192///
193/// # Notes
194/// 1) The quantized values are packed into 32-bit unsigned integers. For example, int8
195///    quantized values pack 4 grouped values into a single `u32`. When unpacking these values,
196///    we make sure to retrieve only the meaningful values (and ignore the alignment padding).
197/// 2) Quantization parameters are appended to the tensor data.
198///    As such, the last bytes always correspond to the scale parameter.
199///    If the quantization scheme includes an offset (zero-point) parameter, it is next to last.
200pub struct QuantizedBytes {
201    /// The quantized values and quantization parameters represented as bytes.
202    pub bytes: Bytes,
203    /// The quantization scheme.
204    pub scheme: QuantScheme,
205    /// The shape of the quantized tensor. The block count, and so the scale count, follows from
206    /// it per axis: a block that does not span the trailing dimension is not a run of elements.
207    pub shape: Shape,
208}
209
210impl QuantizedBytes {
211    /// Creates a new quantized bytes representation.
212    ///
213    /// `global` is the per-tensor scale, required by a two-level scheme and rejected by a
214    /// one-level one.
215    pub fn new<E: bytemuck::CheckedBitPattern + bytemuck::NoUninit>(
216        value: Vec<E>,
217        shape: impl Into<Shape>,
218        scheme: QuantScheme,
219        scales: &[f32],
220        global: Option<f32>,
221    ) -> Self {
222        let shape = shape.into();
223        assert_eq!(
224            value.len(),
225            shape.num_elements(),
226            "{} quantized values do not fill a tensor of shape {shape:?}",
227            value.len()
228        );
229        // Only used for 8-bit quantization data comparison in tests
230        if TypeId::of::<E>() != TypeId::of::<i8>() {
231            panic!("Invalid quantized type");
232        }
233
234        // Re-interpret `Vec<E>` as `Vec<i8>` with `Vec::from_raw_parts`
235        let i8s: Vec<i8> = bytemuck::allocation::cast_vec(value);
236        let mut bytes = Bytes::from_elems(i8s);
237
238        let scales = match scheme.block_size() {
239            None => &scales[..1],
240            Some(_) => scales,
241        };
242        let scale_bytes = encode_scales(scales, scheme.scale_dtype());
243        bytes.extend_from_byte_slice_aligned(scale_bytes.as_slice(), QPARAM_ALIGN);
244
245        // Last, so a reader can peel it off the end before the block scales it normalizes.
246        match (global_scale_dtype(&scheme), global) {
247            (Some(dtype), Some(global)) => {
248                // Encoding the per-tensor scale narrower would round it, and the block scales were
249                // normalized against the unrounded one.
250                assert_eq!(
251                    dtype,
252                    ScaleDtype::F32,
253                    "a two-level scheme stores its per-tensor scale as f32, got {scheme:?}"
254                );
255                let global_bytes = encode_scales(&[global], dtype);
256                bytes.extend_from_byte_slice_aligned(global_bytes.as_slice(), QPARAM_ALIGN);
257            }
258            (Some(_), None) => panic!("{scheme:?} requires a per-tensor scale"),
259            (None, Some(_)) => panic!("{scheme:?} does not take a per-tensor scale"),
260            (None, None) => {}
261        }
262
263        Self {
264            bytes,
265            scheme,
266            shape,
267        }
268    }
269
270    /// The number of quantized elements.
271    pub fn num_elements(&self) -> usize {
272        self.shape.num_elements()
273    }
274
275    /// Returns the int8 quantized values with the quantization parameters.
276    pub fn into_vec_i8(self) -> (Vec<i8>, DecodedScales) {
277        let scheme = self.scheme;
278        let (values, (qparams, num_params)) = self.split_values_off();
279
280        // Laid out as `[block scale, ...]` optionally followed by the per-tensor scale.
281        let global_bytes = global_scale_size(&scheme);
282        let block_end = qparams
283            .len()
284            .checked_sub(global_bytes)
285            .expect("quantized parameter buffer is shorter than the scheme's global scale");
286        let block_start = block_end
287            .checked_sub(scale_size(scheme.scale_dtype()) * num_params)
288            .expect("quantized parameter buffer is shorter than the scheme's block scales");
289
290        let block = decode_scales(&qparams[block_start..block_end], scheme.scale_dtype());
291        let global =
292            global_scale_dtype(&scheme).map(|dtype| decode_scales(&qparams[block_end..], dtype)[0]);
293
294        (values, DecodedScales { block, global })
295    }
296
297    fn split_i8_values(self, scale_bytes: usize) -> (Vec<i8>, Vec<u8>) {
298        let mut values = read_bytes_to_i8(self.bytes);
299
300        let values_end = values
301            .len()
302            .checked_sub(scale_bytes)
303            .expect("quantized tensor data is shorter than its scheme's parameters");
304        let qparams = values.split_off(values_end);
305
306        (values, bytemuck::cast_vec(qparams))
307    }
308
309    /// Splits the quantized values of the tensor from the quantization parameters.
310    ///
311    /// Returns the values in i8 and a newly allocated vector containing the
312    /// quantization parameter bytes.
313    fn split_values_off(self) -> (Vec<i8>, (Vec<u8>, usize)) {
314        let num_params = params_shape(&self.shape, &self.scheme).num_elements();
315        let scale_bytes =
316            scale_size(self.scheme.scale_dtype()) * num_params + global_scale_size(&self.scheme);
317
318        if let QuantStore::PackedU32(packed_dim) = self.scheme.store {
319            assert_eq!(
320                packed_dim, 0,
321                "Packing must be on innermost dimension for splitting off values"
322            );
323        }
324
325        let (values, qparams) = match self.scheme.store {
326            QuantStore::Native => self.split_i8_values(scale_bytes),
327            QuantStore::PackedU32(_) => match self.scheme.value {
328                QuantValue::Q8F | QuantValue::Q8S => self.split_i8_values(scale_bytes),
329                QuantValue::Q4F | QuantValue::Q4S | QuantValue::Q2F | QuantValue::Q2S => {
330                    let split_at =
331                        self.bytes.len().checked_sub(scale_bytes).expect(
332                            "quantized tensor data is shorter than its scheme's parameters",
333                        );
334                    let qparams = self.bytes[split_at..].to_vec();
335                    let values = bytemuck::cast_slice::<_, u32>(&self.bytes[..split_at]);
336                    // Sub-byte values are unpacked as i8s for value equality tests
337                    let values = unpack_q_to_i8s(values, self.num_elements(), &self.scheme.value);
338                    (values, qparams)
339                }
340                QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1 => {
341                    unimplemented!("Not yet supported")
342                }
343            },
344            QuantStore::PackedNative(_) => unimplemented!("Not yet supported"),
345        };
346
347        (values, (qparams, num_params))
348    }
349}
350
351/// Round a scale up to the smallest value representable by the scale dtype that is no smaller.
352///
353/// Backends that keep scales in `f32` must apply this when quantizing, so that the scale they
354/// divide by is the one that will actually be stored. Otherwise a tensor dequantizes differently
355/// after a save/load round trip.
356///
357/// Up rather than to nearest, because a scale is derived from the largest magnitude it has to
358/// cover. Rounding down puts that value past the end of the quantized range, where it clips, which
359/// measured several times worse than the coarser step rounding up costs.
360pub fn scale_to_dtype(scale: f32, dtype: ScaleDtype) -> f32 {
361    dtype.round_up(scale)
362}
363
364/// Bytes taken by the per-tensor scale, zero for a scheme that does not carry one over blocks.
365pub fn global_scale_size(scheme: &QuantScheme) -> usize {
366    global_scale_dtype(scheme).map_or(0, scale_size)
367}
368
369/// Number of storage elements a tensor of `shape` occupies under `scheme`.
370///
371/// A packed store divides only the packed dimension, rounding that extent up on its own and
372/// leaving the others intact, so a non-divisible extent pads once per line rather than once
373/// over the flattened tensor. This mirrors the storage shape the allocation actually uses (see
374/// `CubeTensor::quantized_storage` in burn-cubecl); flattening first would under-count, e.g. a
375/// `[3, 3]` Q4 `PackedU32` tensor occupies `3 * ceil(3 / 8) = 3` words, not `ceil(9 / 8) = 2`.
376fn storage_elements(scheme: &QuantScheme, shape: &Shape) -> usize {
377    let num_quants = scheme.num_quants();
378
379    match scheme.store {
380        QuantStore::PackedU32(packed_dim) | QuantStore::PackedNative(packed_dim)
381            if num_quants > 1 && !shape.is_empty() =>
382        {
383            let packed_dim = shape.num_dims() - packed_dim - 1;
384            let mut storage = shape.clone();
385            storage[packed_dim] = storage[packed_dim].div_ceil(num_quants);
386            storage.num_elements()
387        }
388        _ => shape.num_elements().div_ceil(num_quants),
389    }
390}
391
392/// Total bytes a tensor of `shape` occupies under `scheme`, laid out as [`QuantizedBytes::new`]
393/// writes it: values, then block scales, then (for a two-level scheme) the per-tensor scale.
394pub fn quantized_data_len(scheme: &QuantScheme, shape: &Shape) -> usize {
395    let value_bytes = storage_elements(scheme, shape) * scheme.size_bits_stored().div_ceil(8);
396
397    let num_params = params_shape(shape, scheme).num_elements();
398    let scale_bytes = num_params * scale_size(scheme.scale_dtype());
399
400    value_bytes + scale_bytes + global_scale_size(scheme)
401}
402
403/// Bytes per stored scale entry for the given scale dtype.
404pub fn scale_size(dtype: ScaleDtype) -> usize {
405    match dtype {
406        ScaleDtype::F32 => 4,
407        ScaleDtype::F16 | ScaleDtype::BF16 => 2,
408        ScaleDtype::UE8M0 | ScaleDtype::UE4M3 => 1,
409    }
410}
411
412/// Decode stored scale entries into f32.
413fn decode_scales(bytes: &[u8], dtype: ScaleDtype) -> Vec<f32> {
414    match dtype {
415        ScaleDtype::F32 => bytes
416            .as_chunks::<4>()
417            .0
418            .iter()
419            .map(|c| f32::from_ne_bytes([c[0], c[1], c[2], c[3]]))
420            .collect(),
421        ScaleDtype::F16 => bytes
422            .as_chunks::<2>()
423            .0
424            .iter()
425            .map(|c| crate::f16::from_ne_bytes([c[0], c[1]]).to_f32())
426            .collect(),
427        ScaleDtype::BF16 => bytes
428            .as_chunks::<2>()
429            .0
430            .iter()
431            .map(|c| crate::bf16::from_ne_bytes([c[0], c[1]]).to_f32())
432            .collect(),
433        ScaleDtype::UE4M3 => bytes.iter().map(|b| e4m3::from_bits(*b).to_f32()).collect(),
434        ScaleDtype::UE8M0 => bytes.iter().map(|b| ue8m0_to_f32(*b)).collect(),
435    }
436}
437
438/// Encode f32 scales at the scale dtype for serialization.
439fn encode_scales(scales: &[f32], dtype: ScaleDtype) -> Vec<u8> {
440    match dtype {
441        ScaleDtype::F32 => scales.iter().flat_map(|s| s.to_ne_bytes()).collect(),
442        ScaleDtype::F16 => scales
443            .iter()
444            .flat_map(|s| crate::f16::from_f32(*s).to_ne_bytes())
445            .collect(),
446        ScaleDtype::BF16 => scales
447            .iter()
448            .flat_map(|s| crate::bf16::from_f32(*s).to_ne_bytes())
449            .collect(),
450        ScaleDtype::UE4M3 => scales
451            .iter()
452            .map(|s| e4m3::from_f32(*s).to_bits())
453            .collect(),
454        // `f32_to_ue8m0` rounds up, which is the rule for a scale and matches `scale_to_dtype`;
455        // a scale reaching here has already been rounded onto the grid anyway.
456        ScaleDtype::UE8M0 => scales.iter().map(|s| f32_to_ue8m0(*s)).collect(),
457    }
458}
459
460fn read_bytes_to_i8(bytes: Bytes) -> Vec<i8> {
461    match bytes.try_into_vec::<i8>() {
462        Ok(val) => val,
463        // Safety,
464        //
465        // `Vec<u8>` can be Re-interpreted as `Vec<i8>` since they share the same alignment.
466        Err(bytes) => unsafe { core::mem::transmute::<Vec<u8>, Vec<i8>>(bytes.to_vec()) },
467    }
468}
469
470/// Pack signed 8-bit integer values into a sequence of unsigned 32-bit integers.
471pub fn pack_i8s_to_u32s(values: Vec<i8>) -> Vec<u32> {
472    // Shift and combine groups of four 8-bit values into a u32.
473    // Same as doing this:
474    //     let result = (d_u8 & 0xFF) << 24 | (c_u8 & 0xFF) << 16 | (b_u8 & 0xFF) << 8 | (a_u8 & 0xFF);
475    #[cfg(target_endian = "big")]
476    {
477        values
478            .chunks(4)
479            .map(|x| {
480                x.iter()
481                    .enumerate()
482                    .fold(0u32, |acc, (i, x)| acc | (*x as u32 & 0xFF) << (i * 8))
483            })
484            .collect()
485    }
486
487    // The order of bytes in little endian matches the above description, we just need to
488    // handle padding when the number of values is not a factor of 4
489    #[cfg(target_endian = "little")]
490    {
491        let mut values = values;
492        let remainder = values.len() % 4;
493        if remainder != 0 {
494            // Pad with zeros
495            values.extend(core::iter::repeat_n(0, 4 - remainder));
496        }
497
498        let len = values.len() / 4;
499        let capacity = values.capacity() / 4;
500
501        // Pre-forget the old vec and re-interpret as u32
502        let mut values = core::mem::ManuallyDrop::new(values);
503        let ptr = values.as_mut_ptr() as *mut u32;
504
505        unsafe { Vec::from_raw_parts(ptr, len, capacity) }
506    }
507}
508
509/// Unpack integer values into a sequence of signed 8-bit integers.
510pub(crate) fn unpack_q_to_i8s<Q: PrimInt>(
511    values: &[Q],
512    numel: usize,
513    value: &QuantValue,
514) -> Vec<i8> {
515    let size_store = size_of::<Q>() * 8;
516    let size_quant = value.size_bits();
517    let num_quants = size_store / size_quant;
518    let mask = Q::from((1 << size_quant) - 1).unwrap();
519    let sign_shift = 8 - size_quant; // sign extension for sub-byte values
520    values
521        .iter()
522        .enumerate()
523        .flat_map(|(i, &packed)| {
524            // A single u32 could contain less than four 8-bit values...
525            let n = core::cmp::min(num_quants, numel - i * num_quants);
526            // Extract each 8-bit segment from u32 and cast back to i8
527            // Same as doing this (when 4 values are fully packed):
528            //     let a = (packed & 0xFF) as i8;
529            //     let b = ((packed >> 8) & 0xFF) as i8;
530            //     let c = ((packed >> 16) & 0xFF) as i8;
531            //     let d = ((packed >> 24) & 0xFF) as i8;
532            (0..n).map(move |i| {
533                let raw = (packed >> (i * size_quant) & mask).to_u8().unwrap();
534                ((raw << sign_shift) as i8) >> sign_shift
535            })
536        })
537        .collect()
538}
539
540#[cfg(test)]
541mod tests {
542
543    use super::*;
544    use alloc::vec;
545
546    #[test]
547    fn should_pack_i8s_to_u32() {
548        let packed = pack_i8s_to_u32s(vec![-128, 2, -3, 127]);
549
550        assert_eq!(packed, vec![2147287680]);
551    }
552
553    #[test]
554    fn should_pack_i8s_to_u32_padded() {
555        let packed = pack_i8s_to_u32s(vec![-128, 2, -3, 127, 55]);
556        let packed_padded = pack_i8s_to_u32s(vec![-128, 2, -3, 127, 55, 0, 0, 0]);
557
558        assert_eq!(packed, vec![2147287680, 55]);
559        assert_eq!(packed, packed_padded);
560    }
561
562    #[test]
563    fn should_unpack_u32s_to_i8s() {
564        let unpacked = unpack_q_to_i8s(&[2147287680u32], 4, &QuantValue::Q8S);
565
566        assert_eq!(unpacked, vec![-128, 2, -3, 127]);
567    }
568
569    #[test]
570    fn should_unpack_u32s_to_i8s_padded() {
571        let unpacked = unpack_q_to_i8s(&[55u32], 1, &QuantValue::Q8S);
572
573        assert_eq!(unpacked, vec![55]);
574    }
575
576    #[test]
577    fn should_unpack_u32s_to_i8s_arange() {
578        let unpacked = unpack_q_to_i8s(
579            &[
580                0u32, 286331136, 286331153, 572657937, 572662306, 857874978, 858993459, 858993459,
581                1145324612, 1145324612, 1431655748, 1431655765, 1717982549, 1717986918, 2003199590,
582                2004318071,
583            ],
584            128,
585            &QuantValue::Q4S,
586        );
587
588        assert_eq!(
589            unpacked,
590            vec![
591                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,
592                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,
593                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,
594                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,
595                6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7
596            ]
597        );
598    }
599
600    #[test]
601    fn should_pack_unpack_quantization_parameters_per_tensor_symmetric() {
602        // Quantized [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
603        let scale = 0.03937008;
604        let values = vec![0i8, 25, 51, 76, 102, 127];
605
606        let q_bytes = QuantizedBytes::new(
607            values.clone(),
608            [2, 3],
609            QuantScheme::default()
610                .with_value(QuantValue::Q8S)
611                .with_store(QuantStore::Native),
612            &[scale],
613            None,
614        );
615
616        let (q_values, qparams) = q_bytes.into_vec_i8();
617
618        assert_eq!(qparams.block, vec![scale]);
619
620        assert_eq!(q_values, values);
621    }
622
623    /// Backends divide by what `scale_to_dtype` returns and hand that same value to
624    /// `encode_scales`. If encoding moved it, a tensor would dequantize differently after a
625    /// save/load round trip, so the codec has to leave an already-rounded scale alone.
626    #[test]
627    fn scale_to_dtype_survives_the_codec() {
628        // Includes values that saturate (500), land in e4m3's subnormals (1e-3), and underflow
629        // it entirely (7.7e-4).
630        let scales = [0.5f32, 0.3, 1.0 / 3.0, 500.0, 1e-3, 7.7e-4];
631
632        for dtype in [
633            ScaleDtype::F32,
634            ScaleDtype::F16,
635            ScaleDtype::BF16,
636            ScaleDtype::UE4M3,
637            ScaleDtype::UE8M0,
638        ] {
639            let rounded: Vec<f32> = scales.iter().map(|s| scale_to_dtype(*s, dtype)).collect();
640            let via_codec = decode_scales(&encode_scales(&rounded, dtype), dtype);
641
642            assert_eq!(
643                rounded, via_codec,
644                "the codec moves a scale {dtype:?} can already represent"
645            );
646            // 500 is past what e4m3 can hold, so it saturates rather than rounding up.
647            for (scale, rounded) in scales.iter().zip(&rounded).filter(|(s, _)| **s < 500.0) {
648                assert!(
649                    rounded >= scale,
650                    "{scale} rounded down to {rounded} for {dtype:?}"
651                );
652            }
653        }
654    }
655
656    /// The two scale regions have different widths, and the length assertion pins the layout as
657    /// dense: nothing is padded between them.
658    #[test]
659    fn should_pack_unpack_two_level_scales() {
660        // Exactly representable, so this pins the layout rather than the formats' rounding.
661        let block_scales = [0.5f32, 0.125];
662        let global = 3.0f32;
663        let values = vec![0i8, 25, 51, 76, 102, 127, -128, -1];
664
665        let scheme = QuantScheme::default()
666            .with_value(QuantValue::Q8S)
667            .with_store(QuantStore::Native)
668            .per_block([4], ScaleDtype::UE4M3)
669            .per_tensor(ScaleDtype::F32);
670
671        let q_bytes = QuantizedBytes::new(values.clone(), [8], scheme, &block_scales, Some(global));
672
673        // 8 values, one byte per UE4M3 block scale, a 4 byte f32 per-tensor scale.
674        assert_eq!(q_bytes.bytes.len(), 8 + 2 + 4);
675
676        let (q_values, scales) = q_bytes.into_vec_i8();
677
678        assert_eq!(q_values, values);
679        assert_eq!(scales.block, block_scales);
680        assert_eq!(scales.global, Some(global));
681    }
682
683    /// MXFP4's scale level: one `ue8m0` byte per block, no per-tensor scale. Serialization has to
684    /// carry it like any other, or a model in that format cannot be written or read back.
685    ///
686    /// The scales here are powers of two, which is every value `ue8m0` has — so the round trip is
687    /// exact and a mismatch is a real defect rather than the format's own rounding.
688    #[test]
689    fn should_pack_unpack_ue8m0_block_scales() {
690        let block_scales = [0.25f32, 8.0];
691        let values = vec![0i8, 25, 51, 76, 102, 127, -128, -1];
692
693        let scheme = QuantScheme::default()
694            .with_value(QuantValue::Q8S)
695            .with_store(QuantStore::Native)
696            .per_block([4], ScaleDtype::UE8M0);
697
698        let q_bytes = QuantizedBytes::new(values.clone(), [8], scheme, &block_scales, None);
699
700        // 8 values and one byte per block scale; no per-tensor scale rides along.
701        assert_eq!(q_bytes.bytes.len(), 8 + 2);
702
703        let (q_values, scales) = q_bytes.into_vec_i8();
704
705        assert_eq!(q_values, values);
706        assert_eq!(scales.block, block_scales);
707        assert_eq!(scales.global, None);
708    }
709
710    #[test]
711    #[should_panic(expected = "requires a per-tensor scale")]
712    fn two_level_scheme_without_a_global_scale_is_rejected() {
713        let scheme = QuantScheme::default()
714            .with_value(QuantValue::Q8S)
715            .with_store(QuantStore::Native)
716            .per_block([4], ScaleDtype::F32)
717            .per_tensor(ScaleDtype::F32);
718
719        QuantizedBytes::new(vec![0i8; 8], [8], scheme, &[0.5, 0.125], None);
720    }
721
722    #[test]
723    #[should_panic(expected = "stores its per-tensor scale as f32")]
724    fn a_narrower_per_tensor_scale_is_rejected() {
725        let scheme = QuantScheme::default()
726            .with_value(QuantValue::Q8S)
727            .with_store(QuantStore::Native)
728            .per_block([4], ScaleDtype::UE4M3)
729            .per_tensor(ScaleDtype::F16);
730
731        QuantizedBytes::new(vec![0i8; 8], [8], scheme, &[0.5, 0.125], Some(3.0));
732    }
733
734    /// What a backend answers `supports_dtype` with, so a scheme it declines here is one no path
735    /// reaches: each of these panics further in, where the rule is relied on.
736    #[test]
737    fn quantizable_declines_what_no_backend_can_store() {
738        assert!(quantizable(&QuantScheme::default()));
739        assert!(quantizable(
740            &QuantScheme::default().per_block([4], ScaleDtype::F16)
741        ));
742        assert!(quantizable(
743            &QuantScheme::default()
744                .per_block([4], ScaleDtype::UE4M3)
745                .per_tensor(ScaleDtype::F32)
746        ));
747
748        // `ue8m0` stores a bare exponent, so its round-up is to the next power of two.
749        assert!(quantizable(
750            &QuantScheme::default().per_block([4], ScaleDtype::UE8M0)
751        ));
752        assert!(quantizable(
753            &QuantScheme::default().per_tensor(ScaleDtype::UE8M0)
754        ));
755        // MXFP4, whose whole definition is that pairing.
756        assert!(quantizable(
757            &QuantScheme::default()
758                .with_value(QuantValue::E2M1)
759                .per_block([32], ScaleDtype::UE8M0)
760        ));
761
762        // Block scales reaching f32's range leave the per-tensor scale subnormal, and a narrower
763        // per-tensor scale rounds away precision every block was normalized against.
764        assert!(!quantizable(
765            &QuantScheme::default()
766                .per_block([4], ScaleDtype::F32)
767                .per_tensor(ScaleDtype::F32)
768        ));
769        assert!(!quantizable(
770            &QuantScheme::default()
771                .per_block([4], ScaleDtype::UE4M3)
772                .per_tensor(ScaleDtype::BF16)
773        ));
774    }
775
776    /// `scale_size` is what the readers use to locate the scales in the buffer, so an encoding
777    /// wider or narrower than it claims silently misreads every scale.
778    #[test]
779    fn encoded_scale_width_matches_scale_size() {
780        // Powers of two, so `ue8m0` carries them exactly like the wider dtypes do.
781        let scales = [0.5f32, 0.25, 0.125];
782
783        for dtype in [
784            ScaleDtype::F32,
785            ScaleDtype::F16,
786            ScaleDtype::BF16,
787            ScaleDtype::UE4M3,
788            ScaleDtype::UE8M0,
789        ] {
790            assert_eq!(
791                encode_scales(&scales, dtype).len(),
792                scale_size(dtype) * scales.len(),
793                "encoded width disagrees with scale_size for {dtype:?}"
794            );
795        }
796    }
797
798    #[test]
799    fn should_pack_unpack_ue4m3_block_scales() {
800        // Exactly representable in e4m3, so the round trip is lossless and the test pins the
801        // layout rather than the format's rounding.
802        let scales = [0.5f32, 0.125];
803        let values = vec![0i8, 25, 51, 76, 102, 127, -128, -1];
804
805        let q_bytes = QuantizedBytes::new(
806            values.clone(),
807            [8],
808            QuantScheme::default()
809                .with_value(QuantValue::Q8S)
810                .with_store(QuantStore::Native)
811                .per_block([4], ScaleDtype::UE4M3),
812            &scales,
813            None,
814        );
815
816        let (q_values, qparams) = q_bytes.into_vec_i8();
817
818        assert_eq!(qparams.block, scales);
819        assert_eq!(q_values, values);
820    }
821}
burn_std/tensor/quantization.rs

burn_std/tensor/
quantization.rs
