teeny_quant/quant/
mod.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! Weight quantization primitives: INT8/INT4 affine quantization (`affine`, `pack4`) and
18//! `F8_E4M3`/`F8_E5M2` (`fp8`), each parameterized by a [`Granularity`].
19
20pub mod affine;
21pub mod fp8;
22pub mod granularity;
23pub(crate) mod groups;
24pub mod pack4;
25
26pub use affine::{AffineParams, QuantizedAffine, dequantize_affine, quantize_affine};
27
28/// Computes the flat-index -> group-id table for a tensor of `shape` under `granularity`, and
29/// the resulting group count. Quantizing functions in this module compute and use this
30/// internally; it's exposed so callers reconstructing a quantized tensor (e.g.
31/// [`crate::validate`], or an external consumer of `teeny-quant`'s output) can recompute the
32/// same element -> group mapping without duplicating the grouping logic.
33pub fn compute_groups(shape: &[usize], granularity: Granularity) -> (Vec<u32>, usize) {
34    groups::assign_groups(shape, granularity)
35}
36pub use fp8::{Fp8Variant, QuantizedFp8, dequantize_fp8, f8_to_f32, f32_to_f8, quantize_fp8};
37pub use granularity::Granularity;
38
39/// Which quantization scheme to apply to a tensor.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum Scheme {
42    /// 8-bit affine integer quantization.
43    Int8 {
44        /// Symmetric (`zero_point = 0`) vs asymmetric.
45        symmetric: bool,
46    },
47    /// 4-bit affine integer quantization, nibble-packed into `U8` (see [`pack4`]).
48    Int4 {
49        /// Symmetric (`zero_point = 0`) vs asymmetric.
50        symmetric: bool,
51    },
52    /// 8-bit floating point.
53    Fp8 {
54        /// Which OCP FP8 encoding.
55        variant: Fp8Variant,
56    },
57}
58
59impl Scheme {
60    /// The compressed-tensors `weights.type` field for this scheme.
61    pub fn type_name(self) -> &'static str {
62        match self {
63            Scheme::Int8 { .. } | Scheme::Int4 { .. } => "int",
64            Scheme::Fp8 { .. } => "float",
65        }
66    }
67
68    /// The compressed-tensors `weights.num_bits` field for this scheme.
69    pub fn num_bits(self) -> u8 {
70        match self {
71            Scheme::Int8 { .. } => 8,
72            Scheme::Int4 { .. } => 4,
73            Scheme::Fp8 { .. } => 8,
74        }
75    }
76
77    /// Whether this scheme is symmetric (`zero_point = 0` / no zero-point tensor written).
78    pub fn is_symmetric(self) -> bool {
79        match self {
80            Scheme::Int8 { symmetric } | Scheme::Int4 { symmetric } => symmetric,
81            // FP8 quantization here is amax-scaled with no zero-point, i.e. always symmetric.
82            Scheme::Fp8 { .. } => true,
83        }
84    }
85
86    /// A short, stable name used in CLI output and file names (e.g. `int8`, `int4`, `fp8_e4m3`).
87    pub fn short_name(self) -> String {
88        match self {
89            Scheme::Int8 { .. } => "int8".to_string(),
90            Scheme::Int4 { .. } => "int4".to_string(),
91            Scheme::Fp8 {
92                variant: Fp8Variant::E4M3,
93            } => "fp8_e4m3".to_string(),
94            Scheme::Fp8 {
95                variant: Fp8Variant::E5M2,
96            } => "fp8_e5m2".to_string(),
97        }
98    }
99}
teeny_quant/quant/mod.rs

teeny_quant/quant/
mod.rs
