metrale_core/
dtype.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Element types for weights and scales, and the quantization layout of a weight tensor.
4//!
5//! Owner: core.
6//! Invariants: none beyond the types.
7
8use serde::{Deserialize, Serialize};
9
10/// 2026-09-25: Element type of a weight or scale tensor; the width is
11/// [`DType::element_size_bits`].
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
13pub enum DType {
14    /// 2026-09-25: 4-bit float, 1 sign + 2 exponent + 1 mantissa bits; the
15    /// weight type of [`QuantConfig::nvfp4`].
16    E2M1,
17
18    /// 2026-09-25: 8-bit float, 1 sign + 4 exponent + 3 mantissa bits; the
19    /// block-scale type of [`QuantConfig::nvfp4`] and the weight type of
20    /// [`QuantConfig::fp8`].
21    FP8E4M3,
22
23    /// 2026-09-25: 8-bit float, 1 sign + 5 exponent + 2 mantissa bits.
24    FP8E5M2,
25
26    BF16,
27
28    FP16,
29
30    FP32,
31}
32
33impl DType {
34    /// 2026-09-25: Width of one element in bits (4 for `E2M1`).
35    pub const fn element_size_bits(&self) -> usize {
36        match self {
37            DType::E2M1 => 4,
38            DType::FP8E4M3 | DType::FP8E5M2 => 8,
39            DType::BF16 | DType::FP16 => 16,
40            DType::FP32 => 32,
41        }
42    }
43
44    pub const fn elements_per_u32(&self) -> usize {
45        32 / self.element_size_bits()
46    }
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct QuantConfig {
51    pub weight_type: DType,
52
53    pub scale_type: DType,
54
55    /// 2026-09-25: Weights per scale factor; 0 means one scale for the whole
56    /// tensor ([`QuantConfig::fp8`]).
57    pub group_size: usize,
58
59    pub global_scale_type: DType,
60}
61
62impl QuantConfig {
63    /// 2026-09-25: NVFP4: E2M1 weights, one FP8E4M3 scale per 16 weights, and
64    /// an FP32 tensor scale.
65    pub fn nvfp4() -> Self {
66        Self {
67            weight_type: DType::E2M1,
68            scale_type: DType::FP8E4M3,
69            group_size: 16,
70            global_scale_type: DType::FP32,
71        }
72    }
73
74    /// 2026-09-25: FP8E4M3 weights with FP32 per-tensor scales.
75    pub fn fp8() -> Self {
76        Self {
77            weight_type: DType::FP8E4M3,
78            scale_type: DType::FP32,
79            group_size: 0,
80            global_scale_type: DType::FP32,
81        }
82    }
83}
metrale_core/dtype.rs

metrale_core/
dtype.rs
