eredu_checkpoint/
lib.rs

1//! Backend-neutral checkpoint contracts.
2//!
3//! This crate describes stored tensor encodings and load-time intent. It does
4//! not allocate backend tensors or execute accelerator operations.
5
6#![warn(missing_docs)]
7
8use eredu_gguf::{Endian, GgmlType};
9use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
10
11/// Composite model-artifact and component schemas.
12pub mod composite;
13/// Backend-neutral logical GGUF storage and portable encoded leases.
14pub mod expert;
15pub mod gguf_store;
16pub mod recipe;
17/// Canonical SafeTensors index parsing and shard-path admission.
18pub mod safetensors;
19pub mod schema;
20/// Backend-neutral checkpoint stores, selections, and encoded leases.
21pub mod store;
22/// Header-only validation of declarative SafeTensors and GGUF plans.
23pub mod validation;
24
25pub use recipe::{AtomicMatrixRecipeFamily, MatrixRecipeMember, RecipeAlias};
26
27/// Backend-neutral description of a checkpoint's stored scalar encoding.
28#[derive(Debug, Clone, Eq, PartialEq)]
29pub enum StoredDtype {
30    /// Boolean values.
31    Bool,
32    /// Unsigned 8-bit integers.
33    U8,
34    /// Signed 8-bit integers.
35    I8,
36    /// Signed 16-bit integers.
37    I16,
38    /// Unsigned 16-bit integers.
39    U16,
40    /// IEEE half-precision floating point.
41    F16,
42    /// Brain floating point.
43    BF16,
44    /// Signed 32-bit integers.
45    I32,
46    /// Unsigned 32-bit integers.
47    U32,
48    /// IEEE single-precision floating point.
49    F32,
50    /// IEEE double-precision floating point.
51    F64,
52    /// Signed 64-bit integers.
53    I64,
54    /// Unsigned 64-bit integers.
55    U64,
56    /// Complex values with two 32-bit floating-point components.
57    C64,
58    /// Encoded FP8 E4M3 bytes.
59    F8E4M3,
60    /// Packed FP4 E2M1 values.
61    F4,
62    /// Unsigned E8M0 scale bytes used by MX formats.
63    F8E8M0,
64    /// Encoded FP8 E5M2 bytes.
65    F8E5M2,
66    /// Another storage encoding not represented by a named variant.
67    Other(String),
68}
69
70/// Physical tensor encoding recorded by an admitted artifact catalog.
71///
72/// This remains distinct from [`LinearFormat`], which describes the format
73/// selected for an executable neural operator.
74#[derive(Debug, Clone, Eq, PartialEq)]
75#[non_exhaustive]
76pub enum SourceTensorEncoding {
77    /// Scalar storage in a SafeTensors payload.
78    Safetensors(StoredDtype),
79    /// One physical GGML block encoding in a GGUF shard.
80    Gguf {
81        /// Exact GGML tensor encoding.
82        ggml_type: GgmlType,
83        /// Byte order declared by the containing shard.
84        endian: Endian,
85    },
86}
87
88/// Invalid backend-neutral checkpoint metadata.
89#[derive(Debug, Clone, thiserror::Error, Eq, PartialEq)]
90#[error("{0}")]
91pub struct Error(String);
92
93impl Error {
94    /// Creates a checkpoint metadata error.
95    pub fn invalid(message: impl Into<String>) -> Self {
96        Self(message.into())
97    }
98}
99
100/// Per-group affine integer quantization stored alongside a checkpoint.
101#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
102pub struct AffineQuantization {
103    /// Number of adjacent input values sharing one scale and bias.
104    pub group_size: i32,
105    /// Packed bit width for each weight value.
106    pub bits: i32,
107    /// Quantization mode.
108    #[serde(default = "default_affine_mode")]
109    pub mode: AffineQuantizationMode,
110}
111
112impl Default for AffineQuantization {
113    fn default() -> Self {
114        Self {
115            group_size: 64,
116            bits: 4,
117            mode: AffineQuantizationMode::Affine,
118        }
119    }
120}
121
122impl AffineQuantization {
123    /// Creates and validates an affine encoding.
124    pub fn new(group_size: i32, bits: i32) -> Result<Self, Error> {
125        let value = Self {
126            group_size,
127            bits,
128            mode: AffineQuantizationMode::Affine,
129        };
130        value.validate()?;
131        Ok(value)
132    }
133
134    /// Validates the portable affine storage geometry.
135    pub fn validate(self) -> Result<(), Error> {
136        if self.mode != AffineQuantizationMode::Affine {
137            return Err(Error::invalid(
138                "only affine integer quantization is supported",
139            ));
140        }
141        if self.group_size != 16 && (self.group_size <= 0 || self.group_size % 32 != 0) {
142            return Err(Error::invalid(format!(
143                "group_size must be 16 or a positive multiple of 32, got {}",
144                self.group_size
145            )));
146        }
147        if !matches!(self.bits, 2 | 3 | 4 | 5 | 6 | 8) {
148            return Err(Error::invalid(format!(
149                "bits must be one of 2, 3, 4, 5, 6, or 8, got {}",
150                self.bits
151            )));
152        }
153        Ok(())
154    }
155}
156
157const fn default_affine_mode() -> AffineQuantizationMode {
158    AffineQuantizationMode::Affine
159}
160
161/// Portable affine quantization mode.
162#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
163#[serde(rename_all = "lowercase")]
164pub enum AffineQuantizationMode {
165    /// Per-group scale-and-bias affine quantization.
166    Affine,
167}
168
169/// Packed physical encoding of a model weight. Dense storage is represented
170/// by `None` at the use site.
171#[derive(Debug, Clone, Copy, PartialEq, Eq)]
172pub enum WeightQuantization {
173    /// Per-group affine integer storage.
174    Affine(AffineQuantization),
175    /// Microscaling FP4 with E2M1 values and E8M0 scales.
176    MxFp4,
177    /// Checkpoint-native GGML blocks.
178    GgufIQuant {
179        /// Native GGML tensor encoding.
180        ggml_type: GgmlType,
181        /// Byte order declared by the GGUF container.
182        endian: Endian,
183    },
184}
185
186/// Physical encoding of the scale companion for an E4M3 block-FP8 matrix.
187#[derive(Debug, Clone, Copy, PartialEq, Eq)]
188pub enum BlockFp8ScaleEncoding {
189    /// Floating-point inverse scales (F16, BF16, or F32 in an artifact).
190    FloatingPoint,
191    /// Unsigned exponent-only E8M0 inverse scales.
192    Ue8m0,
193}
194
195/// Geometry and companion encoding for an E4M3 block-FP8 matrix.
196#[derive(Debug, Clone, Copy, PartialEq, Eq)]
197pub struct BlockFp8Format {
198    /// Number of output rows represented by one scale.
199    pub block_rows: i32,
200    /// Number of input columns represented by one scale.
201    pub block_columns: i32,
202    /// Physical encoding of each inverse scale.
203    pub scale_encoding: BlockFp8ScaleEncoding,
204}
205
206impl BlockFp8Format {
207    /// Creates a validated block-FP8 format.
208    pub fn new(
209        block_rows: i32,
210        block_columns: i32,
211        scale_encoding: BlockFp8ScaleEncoding,
212    ) -> Result<Self, Error> {
213        let format = Self {
214            block_rows,
215            block_columns,
216            scale_encoding,
217        };
218        format.validate()?;
219        Ok(format)
220    }
221
222    /// Validates positive two-dimensional block geometry.
223    pub fn validate(self) -> Result<(), Error> {
224        if self.block_rows <= 0 || self.block_columns <= 0 {
225            return Err(Error::invalid(format!(
226                "block-FP8 geometry must be positive, got [{}, {}]",
227                self.block_rows, self.block_columns
228            )));
229        }
230        Ok(())
231    }
232}
233
234/// Complete physical encoding selected for one linear matrix.
235///
236/// Unlike [`WeightQuantization`], this type includes dense and block-FP8
237/// storage, so a neural-layer specification never needs an architecture-owned
238/// format enum or an out-of-band quantization flag.
239#[derive(Debug, Clone, Copy, PartialEq, Eq)]
240pub enum LinearFormat {
241    /// Ordinary floating-point matrix storage.
242    Dense,
243    /// Per-group affine integer storage.
244    Affine(AffineQuantization),
245    /// Microscaling FP4 with E2M1 values and E8M0 scales.
246    MxFp4,
247    /// Checkpoint-native GGML blocks.
248    GgufIQuant {
249        /// Native GGML tensor encoding.
250        ggml_type: GgmlType,
251        /// Byte order declared by the GGUF container.
252        endian: Endian,
253    },
254    /// E4M3 values with one inverse scale per two-dimensional block.
255    E4M3BlockFp8(BlockFp8Format),
256}
257
258impl LinearFormat {
259    /// Validates the selected physical encoding and its geometry.
260    pub fn validate(self) -> Result<(), Error> {
261        match self {
262            Self::Dense => Ok(()),
263            Self::Affine(config) => config.validate(),
264            Self::MxFp4 => WeightQuantization::MxFp4.validate(),
265            Self::GgufIQuant { ggml_type, endian } => {
266                WeightQuantization::GgufIQuant { ggml_type, endian }.validate()
267            }
268            Self::E4M3BlockFp8(format) => format.validate(),
269        }
270    }
271
272    /// Returns the packed-quantization descriptor when this format is
273    /// represented by the standard affine/GGUF materializer.
274    pub const fn weight_quantization(self) -> Option<WeightQuantization> {
275        match self {
276            Self::Dense | Self::E4M3BlockFp8(_) => None,
277            Self::Affine(config) => Some(WeightQuantization::Affine(config)),
278            Self::MxFp4 => Some(WeightQuantization::MxFp4),
279            Self::GgufIQuant { ggml_type, endian } => {
280                Some(WeightQuantization::GgufIQuant { ggml_type, endian })
281            }
282        }
283    }
284}
285
286impl From<WeightQuantization> for LinearFormat {
287    fn from(value: WeightQuantization) -> Self {
288        match value {
289            WeightQuantization::Affine(config) => Self::Affine(config),
290            WeightQuantization::MxFp4 => Self::MxFp4,
291            WeightQuantization::GgufIQuant { ggml_type, endian } => {
292                Self::GgufIQuant { ggml_type, endian }
293            }
294        }
295    }
296}
297
298impl From<Option<WeightQuantization>> for LinearFormat {
299    fn from(value: Option<WeightQuantization>) -> Self {
300        value.map_or(Self::Dense, Into::into)
301    }
302}
303
304impl WeightQuantization {
305    /// MXFP4 group size fixed by the format.
306    pub const MXFP4_GROUP_SIZE: i32 = 32;
307    /// MXFP4 packed value width fixed by the format.
308    pub const MXFP4_BITS: i32 = 4;
309
310    /// Returns the grouping used by packed execution.
311    pub fn group_size(self) -> i32 {
312        match self {
313            Self::Affine(config) => config.group_size,
314            Self::MxFp4 => Self::MXFP4_GROUP_SIZE,
315            Self::GgufIQuant { ggml_type, .. } => {
316                ggml_type.block_and_bytes().expect("validated GGML type").0 as i32
317            }
318        }
319    }
320
321    /// Returns the packed storage width.
322    pub fn bits(self) -> i32 {
323        match self {
324            Self::Affine(config) => config.bits,
325            Self::MxFp4 => Self::MXFP4_BITS,
326            Self::GgufIQuant { ggml_type, .. } => {
327                ggml_type.block_and_bytes().expect("validated GGML type").1 as i32
328            }
329        }
330    }
331
332    /// Returns whether the encoding stores affine bias companions.
333    pub const fn has_biases(self) -> bool {
334        matches!(self, Self::Affine(_))
335    }
336
337    /// Returns checkpoint-native GGML metadata, when present.
338    pub const fn gguf_iquant(self) -> Option<(GgmlType, Endian)> {
339        match self {
340            Self::GgufIQuant { ggml_type, endian } => Some((ggml_type, endian)),
341            _ => None,
342        }
343    }
344
345    /// Validates portable storage geometry.
346    pub fn validate(self) -> Result<(), Error> {
347        match self {
348            Self::Affine(config) => config.validate(),
349            Self::MxFp4 => Ok(()),
350            Self::GgufIQuant { ggml_type, .. } => ggml_type
351                .block_and_bytes()
352                .map(|_| ())
353                .map_err(|error| Error::invalid(error.to_string())),
354        }
355    }
356}
357
358impl From<AffineQuantization> for WeightQuantization {
359    fn from(value: AffineQuantization) -> Self {
360        Self::Affine(value)
361    }
362}
363
364#[derive(Serialize, Deserialize)]
365struct WeightQuantizationMetadata {
366    group_size: i32,
367    bits: i32,
368    #[serde(default = "default_quantization_mode")]
369    mode: String,
370}
371
372fn default_quantization_mode() -> String {
373    "affine".into()
374}
375
376impl Serialize for WeightQuantization {
377    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
378    where
379        S: Serializer,
380    {
381        let mode = match self {
382            Self::Affine(_) => "affine",
383            Self::MxFp4 => "mxfp4",
384            Self::GgufIQuant { .. } => {
385                return Err(serde::ser::Error::custom(
386                    "checkpoint-native GGML block metadata is not serializable",
387                ))
388            }
389        };
390        WeightQuantizationMetadata {
391            group_size: self.group_size(),
392            bits: self.bits(),
393            mode: mode.into(),
394        }
395        .serialize(serializer)
396    }
397}
398
399impl<'de> Deserialize<'de> for WeightQuantization {
400    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
401    where
402        D: Deserializer<'de>,
403    {
404        let metadata = WeightQuantizationMetadata::deserialize(deserializer)?;
405        match metadata.mode.as_str() {
406            "affine" => AffineQuantization::new(metadata.group_size, metadata.bits)
407                .map(Self::Affine)
408                .map_err(de::Error::custom),
409            "mxfp4"
410                if metadata.group_size == Self::MXFP4_GROUP_SIZE
411                    && metadata.bits == Self::MXFP4_BITS =>
412            {
413                Ok(Self::MxFp4)
414            }
415            "mxfp4" => Err(de::Error::custom(format!(
416                "MXFP4 requires group_size=32 and bits=4, got group_size={} bits={}",
417                metadata.group_size, metadata.bits
418            ))),
419            mode => Err(de::Error::custom(format!(
420                "unsupported quantization mode {mode:?}"
421            ))),
422        }
423    }
424}
eredu_checkpoint/lib.rs

eredu_checkpoint/
lib.rs
