cherenkov/model/
preparation.rs1use super::*;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Serialize, Deserialize)]
6#[serde(tag = "status", rename_all = "snake_case")]
7pub enum Preparation {
8 Direct,
10 Required { steps: Vec<PreparationStep> },
12 Unsupported { reasons: Vec<CompatibilityIssue> },
14}
15
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum PreparationStep {
20 NormalizeQwen4Exp,
22 QuantizeAffineQ4,
24 PackExpertRecords,
26 InterleaveNgramRows,
28}
29
30#[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 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}