1use crate::kv_cache::LayerKvCache;
15use crate::linear_core::{GdnCfg, GdnWeights, VmfPhaseCfg, VmfPhaseWeights};
16use crate::pipeline::{
17 AttnKind, DenseFfn, FfnKind, LayerWeights, MoeFfn, MtpModule, Pipeline, PipelineWeights,
18};
19use crate::qtensor::QTensor;
20use crate::sampler::SamplerConfig;
21use crate::tokenizer::Tokenizer;
22use cortiq_core::quant::dequant_tensor;
23use cortiq_core::{CmfError, CmfModel, LayerType, ModelArch};
24use std::sync::Arc;
25
26pub enum Overlay<'a> {
29 None,
30 One(&'a str),
31 Blend(&'a [(String, f32)]),
33}
34
35impl Overlay<'_> {
36 fn blend_touches(&self, model: &CmfModel, name: &str) -> bool {
37 match self {
38 Overlay::Blend(list) => list
39 .iter()
40 .any(|(sid, _)| model.tensor(&format!("skill.{sid}.{name}")).is_some()),
41 _ => false,
42 }
43 }
44}
45
46fn dequant_by_name(model: &CmfModel, name: &str) -> Result<Vec<f32>, String> {
47 let entry = model
48 .tensor(name)
49 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
50 let mut out = vec![0.0f32; entry.n_elems()];
51 dequant_tensor(entry, model.entry_bytes(entry), &mut out)?;
52 Ok(out)
53}
54
55fn blend_f32(model: &CmfModel, name: &str, list: &[(String, f32)]) -> Result<Vec<f32>, String> {
58 let mut acc: Option<Vec<f32>> = None;
59 for (sid, w) in list {
60 let sname = format!("skill.{sid}.{name}");
61 let src = if model.tensor(&sname).is_some() { &sname } else { name };
62 let t = dequant_by_name(model, src)?;
63 match &mut acc {
64 None => {
65 let mut t = t;
66 for v in t.iter_mut() {
67 *v *= w;
68 }
69 acc = Some(t);
70 }
71 Some(a) => {
72 for (av, tv) in a.iter_mut().zip(&t) {
73 *av += w * tv;
74 }
75 }
76 }
77 }
78 acc.ok_or_else(|| "empty blend".into())
79}
80
81fn load_f32(model: &CmfModel, name: &str, ov: &Overlay) -> Result<Vec<f32>, String> {
83 if ov.blend_touches(model, name) {
84 if let Overlay::Blend(list) = ov {
85 return blend_f32(model, name, list);
86 }
87 }
88 let skill = match ov {
89 Overlay::One(s) => Some(*s),
90 _ => None,
91 };
92 let entry = model
93 .resolve_tensor(name, skill)
94 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
95 let bytes = model.entry_bytes(entry);
96 let mut out = vec![0.0f32; entry.n_elems()];
97 dequant_tensor(entry, bytes, &mut out)?;
98 Ok(out)
99}
100
101pub(crate) fn build_layer_ffn(
107 model: &Arc<CmfModel>,
108 arch: &ModelArch,
109 li: usize,
110 force_f32: bool,
111 ov: &Overlay,
112) -> Result<FfnKind, CmfError> {
113 let prefix = format!("model.layers.{li}.");
114 let load_dense = |p: &str| -> Result<DenseFfn, CmfError> {
115 Ok(DenseFfn {
116 gate_proj: load_matrix(model, &format!("{p}gate_proj.weight"), force_f32, ov)?,
117 up_proj: load_matrix(model, &format!("{p}up_proj.weight"), force_f32, ov)?,
118 down_proj: load_matrix(model, &format!("{p}down_proj.weight"), force_f32, ov)?,
119 })
120 };
121 let router_name = format!("{prefix}mlp.gate.weight");
122 if model.tensor(&router_name).is_none() {
123 return Ok(FfnKind::Dense(load_dense(&format!("{prefix}mlp."))?));
124 }
125 let cfg = arch.moe.as_ref().ok_or_else(|| {
126 CmfError::Parse(format!("{router_name} present but header has no arch.moe block"))
127 })?;
128 let experts = (0..cfg.num_experts)
129 .map(|e| load_dense(&format!("{prefix}mlp.experts.{e}.")))
130 .collect::<Result<Vec<_>, _>>()?;
131 let shared = if model
132 .tensor(&format!("{prefix}mlp.shared_expert.gate_proj.weight"))
133 .is_some()
134 {
135 Some((
136 load_dense(&format!("{prefix}mlp.shared_expert."))?,
137 load_matrix(model, &format!("{prefix}mlp.shared_expert_gate.weight"), force_f32, ov)?,
138 ))
139 } else {
140 None
141 };
142 Ok(FfnKind::Moe(MoeFfn {
143 router: load_matrix(model, &router_name, force_f32, ov)?,
144 experts,
145 top_k: cfg.top_k,
146 norm_topk_prob: cfg.norm_topk_prob,
147 shared,
148 stats: std::cell::RefCell::new(Vec::new()),
149 }))
150}
151
152fn load_matrix(
153 model: &Arc<CmfModel>,
154 name: &str,
155 force_f32: bool,
156 ov: &Overlay,
157) -> Result<QTensor, CmfError> {
158 if ov.blend_touches(model, name) {
162 if let Overlay::Blend(list) = ov {
163 let entry = model
164 .tensor(name)
165 .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
166 let data = blend_f32(model, name, list)
167 .map_err(|e| CmfError::Parse(format!("blend: {e}")))?;
168 return Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]));
169 }
170 }
171 let skill = match ov {
172 Overlay::One(s) => Some(*s),
173 _ => None,
174 };
175 let name: &str = &match skill {
178 Some(sid) if model.tensor(&format!("skill.{sid}.{name}")).is_some() => {
179 format!("skill.{sid}.{name}")
180 }
181 _ => name.to_string(),
182 };
183 let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
184 if force_f32 {
185 let entry = model
186 .tensor(name)
187 .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
188 if entry.shape.len() != 2 {
189 return Err(err(format!("'{name}' is not 2-D")));
190 }
191 let data = load_f32(model, name, &Overlay::None).map_err(err)?;
192 Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]))
193 } else {
194 QTensor::from_model(model, name).map_err(err)
195 }
196}
197
198impl Pipeline {
199 pub fn from_model(model: &Arc<CmfModel>, sampler_config: SamplerConfig) -> Result<Self, CmfError> {
201 Self::from_model_with_skill(model, sampler_config, None)
202 }
203
204 pub fn from_model_with_skill(
210 model: &Arc<CmfModel>,
211 sampler_config: SamplerConfig,
212 skill: Option<&str>,
213 ) -> Result<Self, CmfError> {
214 match skill {
215 Some(s) => Self::from_model_with_overlay(model, sampler_config, &Overlay::One(s)),
216 None => Self::from_model_with_overlay(model, sampler_config, &Overlay::None),
217 }
218 }
219
220 pub fn from_model_with_blend(
223 model: &Arc<CmfModel>,
224 sampler_config: SamplerConfig,
225 blend: &[(String, f32)],
226 ) -> Result<Self, CmfError> {
227 Self::from_model_with_overlay(model, sampler_config, &Overlay::Blend(blend))
228 }
229
230 fn from_model_with_overlay(
231 model: &Arc<CmfModel>,
232 sampler_config: SamplerConfig,
233 ov: &Overlay,
234 ) -> Result<Self, CmfError> {
235 let skill = match ov {
236 Overlay::One(s) => Some(*s),
237 _ => None,
238 };
239 if let Some(sid) = skill {
240 let known = model.header.skills.iter().any(|s| s.id == sid)
241 || model.skill_tensors(sid).next().is_some();
242 if !known {
243 return Err(CmfError::Parse(format!(
244 "skill '{sid}' not in this container (header.skills: {:?})",
245 model.header.skills.iter().map(|s| &s.id).collect::<Vec<_>>()
246 )));
247 }
248 tracing::info!(
249 "skill '{sid}': {} replacement tensors overlaid",
250 model.skill_tensors(sid).count()
251 );
252 }
253 let arch = model.arch().clone();
254 let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
255
256 let masks_present = !model.masks.masks.is_empty();
262 let force_f32 = masks_present; let mut tokenizer = if let Some(vocab_bytes) = &model.vocab {
266 Tokenizer::from_bytes(vocab_bytes)
267 .map_err(|e| CmfError::Parse(format!("embedded tokenizer: {e}")))?
268 } else {
269 let sidecar = model.path.with_file_name("tokenizer.json");
270 if sidecar.exists() {
271 Tokenizer::from_file(&sidecar)
272 .map_err(|e| CmfError::Parse(format!("sidecar tokenizer: {e}")))?
273 } else {
274 tracing::warn!("no tokenizer in file or sidecar — using byte-level fallback");
275 Tokenizer::byte_level()
276 }
277 };
278 if let Some(tc) = &model.header.tokenizer_config {
280 tokenizer.chat_template = tc.chat_template.clone();
281 tokenizer.extra_eos.extend(tc.eos_token_ids.iter().copied());
282 if tokenizer.bos_token_id.is_none() {
283 tokenizer.bos_token_id = tc.bos_token_id;
284 }
285 tracing::info!(
286 "chat bundle: template {} chars, {} stop ids",
287 tc.chat_template.as_deref().map(str::len).unwrap_or(0),
288 tc.eos_token_ids.len()
289 );
290 }
291
292 let embed_tokens = load_matrix(model, "model.embed_tokens.weight", false, ov)?;
294 let final_norm = load_f32(model, "model.norm.weight", ov).map_err(err)?;
295 let lm_head = if model.tensor("lm_head.weight").is_some() {
296 load_matrix(model, "lm_head.weight", false, ov)?
297 } else if arch.tie_word_embeddings {
298 load_matrix(model, "model.embed_tokens.weight", false, ov)?
300 } else {
301 return Err(CmfError::MissingTensor(
302 "lm_head.weight (and tie_word_embeddings is false)".into(),
303 ));
304 };
305
306 let has_linear = arch
308 .layer_types
309 .iter()
310 .any(|t| matches!(t, LayerType::LinearAttention));
311 let mut vmf_cfg = None;
312 let mut gdn_cfg = None;
313 if has_linear {
314 let lc = arch.linear_core.as_ref().ok_or_else(|| {
315 CmfError::Parse(
316 "model has LinearAttention layers but no arch.linear_core — \
317 reconvert with the current converter"
318 .into(),
319 )
320 })?;
321 let need = |v: Option<usize>, name: &str| {
322 v.ok_or_else(|| CmfError::Parse(format!("linear core needs arch.{name}")))
323 };
324 match lc.kind.as_str() {
325 "vmf_phase" => {
326 vmf_cfg = Some(VmfPhaseCfg {
327 num_heads: lc.num_heads,
328 nphase: need(lc.nphase, "linear_core.nphase")?,
329 value_head_dim: lc.value_head_dim,
330 hidden_size: arch.hidden_size,
331 phase_mass: std::env::var("CMF_PHASE_MASS")
334 .ok()
335 .and_then(|v| v.parse().ok())
336 .unwrap_or(0.0),
337 });
338 }
339 "gated_delta_net" => {
340 gdn_cfg = Some(GdnCfg {
341 num_v_heads: lc.num_heads,
342 num_k_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
343 key_head_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
344 value_head_dim: lc.value_head_dim,
345 conv_kernel: need(
346 arch.linear_conv_kernel_dim,
347 "linear_conv_kernel_dim",
348 )?,
349 hidden_size: arch.hidden_size,
350 rms_eps: arch.rms_norm_eps as f64,
351 });
352 }
353 other => {
354 return Err(CmfError::Parse(format!(
355 "unknown linear core '{other}' (this runtime executes: \
356 gated_delta_net, vmf_phase)"
357 )));
358 }
359 }
360 }
361
362 let load_full_attn = |prefix: &str| -> Result<AttnKind, CmfError> {
364 let t = |suffix: &str| load_matrix(model, &format!("{prefix}{suffix}"), force_f32, ov);
365 let n = |suffix: &str| -> Option<Vec<f32>> {
366 model
367 .tensor(&format!("{prefix}{suffix}"))
368 .and_then(|_| load_f32(model, &format!("{prefix}{suffix}"), ov).ok())
369 };
370 let wq = t("self_attn.q_proj.weight")?;
371 let output_gate = wq.rows() == 2 * arch.num_attention_heads * arch.head_dim;
373 let bias = match (
375 n("self_attn.q_proj.bias"),
376 n("self_attn.k_proj.bias"),
377 n("self_attn.v_proj.bias"),
378 ) {
379 (Some(a), Some(b), Some(c)) => Some((a, b, c)),
380 _ => None,
381 };
382 Ok(AttnKind::Full {
383 wq,
384 wk: t("self_attn.k_proj.weight")?,
385 wv: t("self_attn.v_proj.weight")?,
386 wo: t("self_attn.o_proj.weight")?,
387 q_norm: n("self_attn.q_norm.weight"),
388 k_norm: n("self_attn.k_norm.weight"),
389 output_gate,
390 bias,
391 })
392 };
393
394 let load_linear_attn = |prefix: &str| -> Result<AttnKind, CmfError> {
395 if gdn_cfg.is_some() {
396 let t = |suffix: &str| {
398 load_matrix(model, &format!("{prefix}linear_attn.{suffix}"), force_f32, ov)
399 };
400 let f = |suffix: &str| {
401 load_f32(model, &format!("{prefix}linear_attn.{suffix}"), ov).map_err(err)
402 };
403 return Ok(AttnKind::LinearGdn(GdnWeights {
404 in_proj_qkv: t("in_proj_qkv.weight")?,
405 in_proj_z: t("in_proj_z.weight")?,
406 in_proj_a: t("in_proj_a.weight")?,
407 in_proj_b: t("in_proj_b.weight")?,
408 conv1d: f("conv1d.weight")?,
409 a_log: f("A_log")?,
410 dt_bias: f("dt_bias")?,
411 norm: f("norm.weight")?,
412 out_proj: t("out_proj.weight")?,
413 }));
414 }
415 let t = |suffix: &str| load_matrix(model, &format!("{prefix}vmf_attn.{suffix}"), force_f32, ov);
416 let a_log = load_f32(model, &format!("{prefix}vmf_attn.A_log"), ov).map_err(err)?;
417 Ok(AttnKind::Linear(VmfPhaseWeights {
418 thq: t("thq.weight")?,
419 thk: t("thk.weight")?,
420 v_proj: t("v_proj.weight")?,
421 out_proj: t("out_proj.weight")?,
422 decay: a_log
423 .iter()
424 .map(|&a| (-(a as f64).exp()).exp())
425 .collect(),
426 }))
427 };
428
429 let mut layers = Vec::with_capacity(arch.num_layers);
430 for li in 0..arch.num_layers {
431 let prefix = format!("model.layers.{li}.");
432 let attn = match arch.layer_types.get(li) {
433 Some(LayerType::LinearAttention) => load_linear_attn(&prefix)?,
434 _ => load_full_attn(&prefix)?,
435 };
436 layers.push(LayerWeights {
437 input_norm: load_f32(model, &format!("{prefix}input_layernorm.weight"), ov).map_err(err)?,
438 post_norm: load_f32(model, &format!("{prefix}post_attention_layernorm.weight"), ov)
439 .map_err(err)?,
440 ffn: build_layer_ffn(model, &arch, li, false, ov)?,
442 attn,
443 });
444 }
445
446 let mtp = if let Some(cfg) = &arch.mtp {
448 if cfg.num_layers != 1 {
449 return Err(CmfError::Parse(format!(
450 "MTP with {} blocks not supported yet (only 1)",
451 cfg.num_layers
452 )));
453 }
454 let p = "model.mtp.";
455 let attn = load_full_attn("model.mtp.layers.0.")?;
456 Some(MtpModule {
457 enorm: load_f32(model, &format!("{p}enorm.weight"), ov).map_err(err)?,
458 hnorm: load_f32(model, &format!("{p}hnorm.weight"), ov).map_err(err)?,
459 eh_proj: load_matrix(model, &format!("{p}eh_proj.weight"), false, ov)?,
460 layer: LayerWeights {
461 input_norm: load_f32(model, &format!("{p}layers.0.input_layernorm.weight"), ov)
462 .map_err(err)?,
463 post_norm: load_f32(
464 model,
465 &format!("{p}layers.0.post_attention_layernorm.weight"),
466 ov,
467 )
468 .map_err(err)?,
469 ffn: FfnKind::Dense(DenseFfn {
470 gate_proj: load_matrix(model, &format!("{p}layers.0.mlp.gate_proj.weight"), false, ov)?,
471 up_proj: load_matrix(model, &format!("{p}layers.0.mlp.up_proj.weight"), false, ov)?,
472 down_proj: load_matrix(model, &format!("{p}layers.0.mlp.down_proj.weight"), false, ov)?,
473 }),
474 attn,
475 },
476 final_norm: load_f32(model, &format!("{p}norm.weight"), ov).map_err(err)?,
477 kv: LayerKvCache::new(arch.num_kv_heads, arch.head_dim),
478 })
479 } else {
480 None
481 };
482
483 tracing::info!(
484 "Pipeline loaded: {} | {}L ({} linear) | {:.2}B params | storage: {} | MTP: {}",
485 arch.arch_name,
486 arch.num_layers,
487 arch.layer_types
488 .iter()
489 .filter(|t| matches!(t, LayerType::LinearAttention))
490 .count(),
491 model.total_param_count() as f64 / 1e9,
492 if force_f32 { "f32 (masked)" } else { "quantized mmap" },
493 if mtp.is_some() { "yes" } else { "no" }
494 );
495
496 let cap = std::env::var("CMF_MAX_SEQ")
499 .ok()
500 .and_then(|v| v.parse::<usize>().ok())
501 .unwrap_or(8192);
502 let max_seq_len = arch.max_position_embeddings.min(cap);
503
504 let mut pipeline = Pipeline::new(
505 tokenizer,
506 PipelineWeights {
507 embed_tokens,
508 layers,
509 lm_head,
510 final_norm,
511 },
512 arch.hidden_size,
513 arch.intermediate_size,
514 arch.num_attention_heads,
515 arch.num_kv_heads,
516 arch.head_dim,
517 arch.num_layers,
518 arch.vocab_size,
519 arch.rms_norm_eps,
520 arch.rope_theta as f32,
521 arch.norm_style,
522 max_seq_len,
523 sampler_config,
524 );
525 let rotary = ((arch.head_dim as f32 * arch.partial_rotary_factor) as usize).max(2);
526 pipeline.set_rotary(rotary, arch.rope_theta as f32);
527 pipeline.vmf_cfg = vmf_cfg;
528 pipeline.gdn_cfg = gdn_cfg;
529 pipeline.mtp = mtp;
530 pipeline.install_dynamic_routing(model, false);
531 match ov {
535 Overlay::One(sid) => {
536 pipeline.dyn_active = model.header.skills.iter().position(|s| &s.id == sid);
537 }
538 Overlay::Blend(_) => pipeline.dyn_blend_loaded = true,
539 Overlay::None => {}
540 }
541 if let Some(c) = &model.header.calibration {
544 pipeline.set_calib_temp(c.temperature);
545 }
546 Ok(pipeline)
547 }
548
549 pub(crate) fn install_dynamic_routing(
554 &mut self,
555 model: &Arc<CmfModel>,
556 force_f32: bool,
557 ) {
558 self.model = Some(model.clone());
559 self.dyn_force_f32 = force_f32;
560 let mut per_skill = Vec::with_capacity(model.header.skills.len());
561 for sk in &model.header.skills {
562 let mut ffn_layers = std::collections::BTreeSet::new();
563 let mut non_ffn = false;
564 let prefix = format!("skill.{}.", sk.id);
565 for t in model.skill_tensors(&sk.id) {
566 let rel = &t.name[prefix.len()..]; let toks: Vec<&str> = rel.split('.').collect();
568 if toks.len() >= 5
569 && toks[0] == "model"
570 && toks[1] == "layers"
571 && toks[3] == "mlp"
572 {
573 if let Ok(li) = toks[2].parse::<usize>() {
574 ffn_layers.insert(li);
575 continue;
576 }
577 }
578 non_ffn = true; }
580 if non_ffn {
581 tracing::warn!(
582 "skill '{}' replaces non-FFN tensors — excluded from dynamic \
583 routing (static overlay still works)",
584 sk.id
585 );
586 per_skill.push(None);
587 } else {
588 per_skill.push(Some(ffn_layers.into_iter().collect::<Vec<_>>()));
589 }
590 }
591 self.dyn_skill_layers = per_skill;
592 }
593
594 pub fn set_active_skill(&mut self, idx: Option<usize>) -> Result<(), CmfError> {
601 if self.dyn_active == idx {
602 return Ok(());
603 }
604 let model = self
605 .model
606 .clone()
607 .ok_or_else(|| CmfError::Parse("dynamic routing needs a model-backed pipeline".into()))?;
608 let mut union: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
609 if let Some(old) = self.dyn_active {
610 if let Some(Some(ls)) = self.dyn_skill_layers.get(old) {
611 union.extend(ls.iter().copied());
612 }
613 }
614 let new_id: Option<String> = match idx {
615 Some(n) => match self.dyn_skill_layers.get(n) {
616 Some(Some(ls)) => {
617 union.extend(ls.iter().copied());
618 Some(model.header.skills[n].id.clone())
619 }
620 _ => {
621 return Err(CmfError::Parse(format!(
622 "skill index {n} not dynamic-eligible"
623 )))
624 }
625 },
626 None => None,
627 };
628 let ov = match &new_id {
629 Some(s) => Overlay::One(s),
630 None => Overlay::None,
631 };
632 let arch = model.arch();
633 for li in union {
634 self.weights.layers[li].ffn =
635 build_layer_ffn(&model, arch, li, self.dyn_force_f32, &ov)?;
636 }
637 self.dyn_active = idx;
638 Ok(())
639 }
640}