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

eredu_checkpoint/
lib.rs
