teeny_quant/quant/
fp8.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//! `F8_E4M3`/`F8_E5M2` weight quantization.
18//!
19//! Both are natively supported `safetensors` dtypes (see the pinned `safetensors` 0.7's
20//! `Dtype::F8_E4M3`/`Dtype::F8_E5M2`), so unlike INT4 (see [`crate::quant::pack4`]) no packing
21//! convention is needed -- this module only has to do the bit-level `f32 <-> f8` conversion and
22//! per-group `amax`-based scale computation, since the crate depends on neither the `half` crate
23//! (f16/bf16 only, no f8) nor any other f8 implementation.
24//!
25//! Encoding follows the OCP FP8 spec: `E4M3` is the "FN" variant (no infinities; the single
26//! exponent=`1111`/mantissa=`111` bit pattern is reserved for NaN, freeing up the rest of that
27//! exponent for finite values up to `448`). `E5M2` is IEEE-754-like (has infinities). Rounding is
28//! round-to-nearest-even; out-of-range magnitudes saturate (to `448`/`57344` for `E4M3`/`E5M2`
29//! respectively, or to infinity for `E5M2`, which has one). **Subnormal outputs are flushed to
30//! zero** rather than rounded into the target format's subnormal range -- an accepted
31//! simplification for weight quantization, where values within about one ULP of the smallest
32//! normal (`2^-9` for `E4M3`, `2^-16` for `E5M2`) are negligible relative to the scale factor
33//! applied before conversion. Decoding handles subnormals and NaN/Inf fully, since it also needs
34//! to correctly read back bytes this module didn't itself produce.
35
36use crate::quant::granularity::Granularity;
37use crate::quant::groups::assign_groups;
38
39/// Which OCP FP8 encoding to target.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum Fp8Variant {
42    /// 4 exponent bits, 3 mantissa bits, no infinities (`safetensors` `F8_E4M3`).
43    E4M3,
44    /// 5 exponent bits, 2 mantissa bits, IEEE-like (`safetensors` `F8_E5M2`).
45    E5M2,
46}
47
48impl Fp8Variant {
49    fn mantissa_bits(self) -> u32 {
50        match self {
51            Fp8Variant::E4M3 => 3,
52            Fp8Variant::E5M2 => 2,
53        }
54    }
55
56    fn exp_bits(self) -> u32 {
57        match self {
58            Fp8Variant::E4M3 => 4,
59            Fp8Variant::E5M2 => 5,
60        }
61    }
62
63    fn bias(self) -> i32 {
64        match self {
65            Fp8Variant::E4M3 => 7,
66            Fp8Variant::E5M2 => 15,
67        }
68    }
69
70    fn has_infinity(self) -> bool {
71        matches!(self, Fp8Variant::E5M2)
72    }
73
74    /// Largest unbiased exponent used by a *finite* value in this format.
75    fn max_finite_unbiased_exp(self) -> i32 {
76        match self {
77            Fp8Variant::E4M3 => 8,  // biased 15, mantissa < 0b111 (448 = 1.75 * 2^8)
78            Fp8Variant::E5M2 => 15, // biased 30, mantissa <= 0b11 (57344 = 1.75 * 2^15)
79        }
80    }
81
82    fn max_finite_value(self) -> f32 {
83        match self {
84            Fp8Variant::E4M3 => 448.0,
85            Fp8Variant::E5M2 => 57344.0,
86        }
87    }
88
89    fn exp_mask(self) -> u8 {
90        ((1u32 << self.exp_bits()) - 1) as u8
91    }
92}
93
94/// Encodes `x` as an `f8` byte in `variant`. See the module docs for rounding/saturation
95/// behavior.
96pub fn f32_to_f8(x: f32, variant: Fp8Variant) -> u8 {
97    let bits = x.to_bits();
98    let sign: u8 = ((bits >> 31) & 1) as u8;
99
100    if x == 0.0 {
101        return sign << 7;
102    }
103    if x.is_nan() {
104        let mantissa_bits = variant.mantissa_bits();
105        let all_ones_mantissa = ((1u32 << mantissa_bits) - 1) as u8;
106        return (sign << 7) | (variant.exp_mask() << mantissa_bits) | all_ones_mantissa;
107    }
108
109    let m_bits = variant.mantissa_bits();
110    let bias = variant.bias();
111    let max_unbiased = variant.max_finite_unbiased_exp();
112    let min_unbiased = 1 - bias; // smallest *normal* exponent in the target format
113
114    let saturate = || -> u8 {
115        if variant.has_infinity() {
116            (sign << 7) | (variant.exp_mask() << m_bits)
117        } else {
118            // E4M3FN: no infinity -- saturate to the largest finite value (exp=1111, mantissa=110).
119            (sign << 7) | (variant.exp_mask() << m_bits) | (((1u32 << m_bits) - 2) as u8)
120        }
121    };
122
123    if x.is_infinite() {
124        return saturate();
125    }
126
127    let exp_field_f32 = (bits >> 23) & 0xFF;
128    let unbiased_exp = exp_field_f32 as i32 - 127;
129    let mantissa_f32 = bits & 0x007F_FFFF;
130
131    if unbiased_exp > max_unbiased {
132        return saturate();
133    }
134    if unbiased_exp < min_unbiased {
135        // Too small even for the target format's smallest normal -- flush to zero (see module
136        // docs: subnormal outputs aren't supported).
137        return sign << 7;
138    }
139
140    // 24-bit significand: implicit leading 1 + the 23 mantissa bits.
141    let significand = (1u32 << 23) | mantissa_f32;
142    let shift = 23 - m_bits;
143
144    // Round-to-nearest-even at `shift`.
145    let half = 1u32 << (shift - 1);
146    let mask = (1u32 << shift) - 1;
147    let remainder = significand & mask;
148    let mut rounded = significand >> shift;
149    if remainder > half || (remainder == half && (rounded & 1) == 1) {
150        rounded += 1;
151    }
152
153    let implicit_bit = 1u32 << m_bits;
154    let mut exp_out = unbiased_exp;
155    if rounded == (implicit_bit << 1) {
156        // Rounding carried all the way through the mantissa: bump the exponent instead.
157        rounded = implicit_bit;
158        exp_out += 1;
159    }
160    let mantissa_field = rounded - implicit_bit;
161
162    if exp_out > max_unbiased {
163        return saturate();
164    }
165
166    let biased_exp = (exp_out + bias) as u8;
167    (sign << 7) | (biased_exp << m_bits) | (mantissa_field as u8)
168}
169
170/// Decodes an `f8` byte (`variant`) back to `f32`. Unlike [`f32_to_f8`], this fully handles
171/// subnormals, since it may be asked to decode bytes this module didn't itself produce.
172pub fn f8_to_f32(byte: u8, variant: Fp8Variant) -> f32 {
173    let m_bits = variant.mantissa_bits();
174    let bias = variant.bias();
175    let exp_mask = variant.exp_mask();
176
177    let sign = (byte >> 7) & 1;
178    let sign_f: f32 = if sign == 1 { -1.0 } else { 1.0 };
179    let biased_exp = (byte >> m_bits) & exp_mask;
180    let mantissa = (byte & ((1u32 << m_bits) - 1) as u8) as u32;
181
182    if variant == Fp8Variant::E4M3 && biased_exp == exp_mask && mantissa == (1 << m_bits) - 1 {
183        return f32::NAN;
184    }
185    if variant == Fp8Variant::E5M2 && biased_exp == exp_mask {
186        return if mantissa == 0 {
187            sign_f * f32::INFINITY
188        } else {
189            f32::NAN
190        };
191    }
192
193    if biased_exp == 0 {
194        if mantissa == 0 {
195            return sign_f * 0.0;
196        }
197        let val = (mantissa as f32) / ((1u32 << m_bits) as f32) * 2f32.powi(1 - bias);
198        return sign_f * val;
199    }
200
201    let unbiased = biased_exp as i32 - bias;
202    let val = (1.0 + (mantissa as f32) / ((1u32 << m_bits) as f32)) * 2f32.powi(unbiased);
203    sign_f * val
204}
205
206/// An `f32` tensor quantized to `f8` bytes, one [`f32`] scale per group (the tensor is scaled by
207/// `1 / scale` before conversion, matching the affine schemes' convention of storing a
208/// multiplicative dequantization scale).
209#[derive(Debug, Clone)]
210pub struct QuantizedFp8 {
211    /// One `f8`-encoded byte per input element, same order as the input.
212    pub qvalues: Vec<u8>,
213    /// One scale per group.
214    pub scales: Vec<f32>,
215    /// `qvalues[i]`'s group is `groups[i]`.
216    pub groups: Vec<u32>,
217    /// Which variant `qvalues` was encoded with.
218    pub variant: Fp8Variant,
219}
220
221/// Quantizes `data` (row-major, shape `shape`) to `variant`, scaling each group so its `amax`
222/// maps to the format's largest finite value.
223pub fn quantize_fp8(
224    data: &[f32],
225    shape: &[usize],
226    granularity: Granularity,
227    variant: Fp8Variant,
228) -> QuantizedFp8 {
229    let (groups, ngroups) = assign_groups(shape, granularity);
230
231    let mut amax = vec![0f32; ngroups];
232    for (i, &x) in data.iter().enumerate() {
233        let g = groups[i] as usize;
234        amax[g] = amax[g].max(x.abs());
235    }
236
237    let target_max = variant.max_finite_value();
238    let scales: Vec<f32> = amax
239        .iter()
240        .map(|&a| if a == 0.0 { 1.0 } else { a / target_max })
241        .collect();
242
243    let qvalues: Vec<u8> = data
244        .iter()
245        .enumerate()
246        .map(|(i, &x)| {
247            let scale = scales[groups[i] as usize];
248            f32_to_f8(x / scale, variant)
249        })
250        .collect();
251
252    QuantizedFp8 {
253        qvalues,
254        scales,
255        groups,
256        variant,
257    }
258}
259
260/// Reconstructs `f32` values from a [`QuantizedFp8`].
261pub fn dequantize_fp8(q: &QuantizedFp8) -> Vec<f32> {
262    q.qvalues
263        .iter()
264        .enumerate()
265        .map(|(i, &v)| f8_to_f32(v, q.variant) * q.scales[q.groups[i] as usize])
266        .collect()
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    fn round_trip(x: f32, variant: Fp8Variant, tol: f32) {
274        let byte = f32_to_f8(x, variant);
275        let back = f8_to_f32(byte, variant);
276        assert!(
277            (x - back).abs() <= tol,
278            "{x:?} -> {byte:#04x} -> {back:?} (variant {variant:?})"
279        );
280    }
281
282    #[test]
283    fn zero_and_signs_round_trip_exactly() {
284        for variant in [Fp8Variant::E4M3, Fp8Variant::E5M2] {
285            assert_eq!(f8_to_f32(f32_to_f8(0.0, variant), variant), 0.0);
286            round_trip(1.0, variant, 0.0);
287            round_trip(-1.0, variant, 0.0);
288            round_trip(2.0, variant, 0.0);
289            round_trip(-2.0, variant, 0.0);
290        }
291    }
292
293    #[test]
294    fn e4m3_max_finite_round_trips_exactly() {
295        round_trip(448.0, Fp8Variant::E4M3, 0.0);
296        round_trip(-448.0, Fp8Variant::E4M3, 0.0);
297    }
298
299    #[test]
300    fn e4m3_overflow_saturates_to_max_finite() {
301        let byte = f32_to_f8(1.0e6, Fp8Variant::E4M3);
302        assert_eq!(f8_to_f32(byte, Fp8Variant::E4M3), 448.0);
303        let byte = f32_to_f8(f32::INFINITY, Fp8Variant::E4M3);
304        assert_eq!(f8_to_f32(byte, Fp8Variant::E4M3), 448.0);
305    }
306
307    #[test]
308    fn e5m2_max_finite_round_trips_exactly() {
309        round_trip(57344.0, Fp8Variant::E5M2, 0.0);
310    }
311
312    #[test]
313    fn e5m2_overflow_saturates_to_infinity() {
314        let byte = f32_to_f8(1.0e6, Fp8Variant::E5M2);
315        assert!(f8_to_f32(byte, Fp8Variant::E5M2).is_infinite());
316    }
317
318    #[test]
319    fn small_values_flush_to_zero() {
320        // Well below E4M3's smallest normal (2^-6).
321        let byte = f32_to_f8(1.0e-10, Fp8Variant::E4M3);
322        assert_eq!(f8_to_f32(byte, Fp8Variant::E4M3), 0.0);
323    }
324
325    #[test]
326    fn mid_range_round_trips_within_quantization_error() {
327        for variant in [Fp8Variant::E4M3, Fp8Variant::E5M2] {
328            for x in [0.1f32, 0.3, 1.5, 3.7, -12.25, 100.0] {
329                let byte = f32_to_f8(x, variant);
330                let back = f8_to_f32(byte, variant);
331                let tol = x.abs() * 0.2 + 0.01; // f8 has very few mantissa bits
332                assert!(
333                    (x - back).abs() < tol,
334                    "{x} -> {back} (variant {variant:?})"
335                );
336            }
337        }
338    }
339
340    #[test]
341    fn per_channel_quantization_uses_independent_scales() {
342        let data = vec![0.1f32, -0.1, 200.0, -200.0];
343        let shape = [2usize, 2];
344        let q = quantize_fp8(
345            &data,
346            &shape,
347            Granularity::PerChannel { axis: 0 },
348            Fp8Variant::E4M3,
349        );
350        assert_eq!(q.scales.len(), 2);
351        let deq = dequantize_fp8(&q);
352        for (orig, back) in data.iter().zip(deq.iter()) {
353            let tol = orig.abs() * 0.3 + 0.01;
354            assert!((orig - back).abs() < tol, "{orig} vs {back}");
355        }
356    }
357}
teeny_quant/quant/fp8.rs

teeny_quant/quant/
fp8.rs
