Skip to main content

cherenkov/model/
preparation.rs

1use super::*;
2use serde::{Deserialize, Serialize};
3
4/// Representation requirements for loading with the current Qwen4-exp engine.
5#[derive(Debug, Clone, Serialize, Deserialize)]
6#[serde(tag = "status", rename_all = "snake_case")]
7pub enum Preparation {
8    /// The described representation can be consumed without conversion.
9    Direct,
10    /// Supported after the listed preparation steps.
11    Required { steps: Vec<PreparationStep> },
12    /// The current engine cannot prepare this representation.
13    Unsupported { reasons: Vec<CompatibilityIssue> },
14}
15
16/// A transformation needed to produce the engine's prepared representation.
17#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum PreparationStep {
20    /// Normalize native Qwen4-exp tensor names and layout.
21    NormalizeQwen4Exp,
22    /// Convert supported dense weights to affine 4-bit storage.
23    QuantizeAffineQ4,
24    /// Group expert weights into records for streaming and residency.
25    PackExpertRecords,
26    /// Interleave n-gram codes, scales, and biases by row.
27    InterleaveNgramRows,
28}
29
30/// A representation constraint that prevents preparation by the current engine.
31#[derive(Debug, Clone, Serialize, Deserialize)]
32#[serde(tag = "kind", rename_all = "snake_case")]
33pub enum CompatibilityIssue {
34    Architecture,
35    Conventions,
36    GgufConversion,
37    Encoding { tensor: String },
38    GroupSize { tensor: String, group_size: u64 },
39    Shape { tensor: String },
40}
41
42impl ModelDescription {
43    /// Requirements for the current Cherenkov Qwen4-exp packed format. This
44    /// checks representation compatibility, not hardware or model completeness.
45    pub fn preparation(&self) -> Preparation {
46        if matches!(self.format.container, ContainerFormat::Gguf { .. }) {
47            return Preparation::Unsupported {
48                reasons: vec![CompatibilityIssue::GgufConversion],
49            };
50        }
51
52        if !matches!(self.architecture, Architecture::Qwen4Exp) {
53            return Preparation::Unsupported {
54                reasons: vec![CompatibilityIssue::Architecture],
55            };
56        }
57
58        let native = matches!(
59            self.format.conventions,
60            CheckpointConventions::Qwen4ExpHuggingFace
61        );
62        let known = native
63            || matches!(
64                self.format.conventions,
65                CheckpointConventions::Qwen4ExpMlx | CheckpointConventions::Qwen4ExpCherenkov
66            );
67
68        if !known {
69            return Preparation::Unsupported {
70                reasons: vec![CompatibilityIssue::Conventions],
71            };
72        }
73
74        let reasons: Vec<_> = self
75            .tensors
76            .iter()
77            .filter_map(|t| compatibility(t, native))
78            .collect();
79
80        if !reasons.is_empty() {
81            return Preparation::Unsupported { reasons };
82        }
83
84        if matches!(
85            self.format.container,
86            ContainerFormat::CherenkovPacked { version: 1 }
87        ) {
88            return Preparation::Direct;
89        }
90
91        let mut steps = Vec::new();
92
93        if native {
94            steps.extend([
95                PreparationStep::NormalizeQwen4Exp,
96                PreparationStep::QuantizeAffineQ4,
97            ]);
98        }
99
100        steps.extend([
101            PreparationStep::PackExpertRecords,
102            PreparationStep::InterleaveNgramRows,
103        ]);
104
105        Preparation::Required { steps }
106    }
107}
108
109fn compatibility(tensor: &Tensor, native: bool) -> Option<CompatibilityIssue> {
110    let quantized_role = matches!(
111        tensor.role,
112        TensorRole::Projection
113            | TensorRole::Embedding
114            | TensorRole::Expert
115            | TensorRole::NgramEmbedding
116    );
117
118    if !quantized_role {
119        return None;
120    }
121
122    let width = tensor.shape().and_then(|s| s.last()).copied().unwrap_or(0);
123    let alignment = if tensor.role == TensorRole::NgramEmbedding {
124        8
125    } else {
126        64
127    };
128
129    if width == 0 || !width.is_multiple_of(alignment) {
130        return Some(CompatibilityIssue::Shape {
131            tensor: tensor.name.clone(),
132        });
133    }
134
135    match &tensor.encoding {
136        TensorEncoding::Dense { tensor } if native && tensor.dtype == Dtype::Bf16 => None,
137        TensorEncoding::Affine {
138            bits: 4,
139            group_size,
140            scales,
141            offset: AffineOffset::Bias { tensor: biases },
142            ..
143        } if scales.dtype == Dtype::Bf16 && biases.dtype == Dtype::Bf16 => {
144            let valid = if tensor.role == TensorRole::NgramEmbedding {
145                matches!(group_size, 8 | 16 | 32 | 64)
146            } else {
147                *group_size == 64
148            };
149
150            (!valid).then(|| CompatibilityIssue::GroupSize {
151                tensor: tensor.name.clone(),
152                group_size: *group_size,
153            })
154        }
155        _ => Some(CompatibilityIssue::Encoding {
156            tensor: tensor.name.clone(),
157        }),
158    }
159}