1use crate::kv_cache::LayerKvCache;
15use crate::linear_core::{
16 GdnCfg, GdnWeights, ShortConvCfg, ShortConvWeights, VmfPhaseCfg, VmfPhaseWeights,
17};
18use crate::pipeline::{
19 AttnKind, DenseFfn, FfnKind, LayerWeights, MoeFfn, MtpModule, Pipeline, PipelineWeights,
20};
21use crate::qtensor::QTensor;
22use crate::sampler::SamplerConfig;
23use crate::tokenizer::Tokenizer;
24use cortiq_core::quant::dequant_tensor;
25use cortiq_core::{CmfError, CmfModel, LayerType, ModelArch};
26use std::sync::Arc;
27
28pub enum Overlay<'a> {
31 None,
32 One(&'a str),
33 Blend(&'a [(String, f32)]),
35}
36
37impl Overlay<'_> {
38 fn blend_touches(&self, model: &CmfModel, name: &str) -> bool {
39 match self {
40 Overlay::Blend(list) => list
41 .iter()
42 .any(|(sid, _)| model.tensor(&format!("skill.{sid}.{name}")).is_some()),
43 _ => false,
44 }
45 }
46}
47
48fn dequant_by_name(model: &CmfModel, name: &str) -> Result<Vec<f32>, String> {
49 let entry = model
50 .tensor(name)
51 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
52 let mut out = vec![0.0f32; entry.n_elems()];
53 dequant_tensor(entry, model.entry_bytes(entry), &mut out)?;
54 Ok(out)
55}
56
57fn blend_f32(model: &CmfModel, name: &str, list: &[(String, f32)]) -> Result<Vec<f32>, String> {
60 let mut acc: Option<Vec<f32>> = None;
61 for (sid, w) in list {
62 let sname = format!("skill.{sid}.{name}");
63 let src = if model.tensor(&sname).is_some() {
64 &sname
65 } else {
66 name
67 };
68 let t = dequant_by_name(model, src)?;
69 match &mut acc {
70 None => {
71 let mut t = t;
72 for v in t.iter_mut() {
73 *v *= w;
74 }
75 acc = Some(t);
76 }
77 Some(a) => {
78 for (av, tv) in a.iter_mut().zip(&t) {
79 *av += w * tv;
80 }
81 }
82 }
83 }
84 acc.ok_or_else(|| "empty blend".into())
85}
86
87pub(crate) fn load_f32(model: &CmfModel, name: &str, ov: &Overlay) -> Result<Vec<f32>, String> {
89 if ov.blend_touches(model, name) {
90 if let Overlay::Blend(list) = ov {
91 return blend_f32(model, name, list);
92 }
93 }
94 let skill = match ov {
95 Overlay::One(s) => Some(*s),
96 _ => None,
97 };
98 let entry = model
99 .resolve_tensor(name, skill)
100 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
101 let bytes = model.entry_bytes(entry);
102 let mut out = vec![0.0f32; entry.n_elems()];
103 dequant_tensor(entry, bytes, &mut out)?;
104 Ok(out)
105}
106
107pub(crate) fn build_layer_ffn(
113 model: &Arc<CmfModel>,
114 arch: &ModelArch,
115 li: usize,
116 force_f32: bool,
117 ov: &Overlay,
118) -> Result<FfnKind, CmfError> {
119 build_ffn_at(model, arch, &format!("model.layers.{li}."), force_f32, ov)
120}
121
122pub(crate) fn build_ffn_at(
127 model: &Arc<CmfModel>,
128 arch: &ModelArch,
129 prefix: &str,
130 force_f32: bool,
131 ov: &Overlay,
132) -> Result<FfnKind, CmfError> {
133 let prefix = prefix.to_string();
134 let load_dense = |p: &str| -> Result<DenseFfn, CmfError> {
135 let gate_proj = load_matrix(model, &format!("{p}gate_proj.weight"), force_f32, ov)?;
136 let up_proj = load_matrix(model, &format!("{p}up_proj.weight"), force_f32, ov)?;
137 let down_proj = load_matrix(model, &format!("{p}down_proj.weight"), force_f32, ov)?;
138 let inter = gate_proj.rows();
142 if up_proj.rows() != inter || down_proj.cols() != inter {
143 return Err(CmfError::Parse(format!(
144 "{p}: FFN dims disagree (gate.rows={inter}, up.rows={}, \
145 down.cols={}); all three must equal inter'",
146 up_proj.rows(),
147 down_proj.cols()
148 )));
149 }
150 if down_proj.rows() != arch.hidden_size {
151 return Err(CmfError::Parse(format!(
152 "{p}: down_proj.rows={} != hidden_size={}",
153 down_proj.rows(),
154 arch.hidden_size
155 )));
156 }
157 Ok(DenseFfn {
158 gate_proj,
159 up_proj,
160 down_proj,
161 act: crate::pipeline::Act::from_arch_full(arch),
162 })
163 };
164 let router_name = format!("{prefix}mlp.gate.weight");
165 let resonance_moe = model.tensor(&router_name).is_none()
169 && arch.moe.as_ref().is_some_and(|m| m.router_resonance)
170 && model
171 .tensor(&format!("{prefix}mlp.experts.0.gate_proj.weight"))
172 .is_some();
173 if model.tensor(&router_name).is_none() && !resonance_moe {
174 return Ok(FfnKind::Dense(load_dense(&format!("{prefix}mlp."))?));
175 }
176 let cfg = arch.moe.as_ref().ok_or_else(|| {
177 CmfError::Parse(format!(
178 "{router_name} present but header has no arch.moe block"
179 ))
180 })?;
181 let mut experts = Vec::new();
186 for e in 0..cfg.num_experts {
187 if model
188 .tensor(&format!("{prefix}mlp.experts.{e}.gate_proj.weight"))
189 .is_none()
190 {
191 break;
192 }
193 experts.push(load_dense(&format!("{prefix}mlp.experts.{e}."))?);
194 }
195 if experts.is_empty() {
196 return Err(CmfError::Parse(format!(
197 "{prefix}: router present but no expert tensors"
198 )));
199 }
200 let shared = if model
201 .tensor(&format!("{prefix}mlp.shared_expert.gate_proj.weight"))
202 .is_some()
203 {
204 let gate_name = format!("{prefix}mlp.shared_expert_gate.weight");
205 Some((
206 load_dense(&format!("{prefix}mlp.shared_expert."))?,
207 if model.tensor(&gate_name).is_some() {
208 Some(load_matrix(model, &gate_name, force_f32, ov)?)
209 } else {
210 None
211 },
212 ))
213 } else {
214 None
215 };
216 let bias_name = format!("{prefix}mlp.expert_bias");
219 let expert_bias = if model.tensor(&bias_name).is_some() {
220 Some(load_f32(model, &bias_name, ov).map_err(CmfError::Parse)?)
221 } else {
222 None
223 };
224 let top_k = std::env::var("CMF_MOE_TOPK")
230 .ok()
231 .and_then(|v| v.parse::<usize>().ok())
232 .filter(|&k| k >= 1 && k <= cfg.top_k)
233 .inspect(|k| tracing::info!("MoE top_k override: {} (header {})", k, cfg.top_k))
234 .unwrap_or(cfg.top_k);
235 let route_tau = std::env::var("CMF_MOE_TAU")
237 .ok()
238 .and_then(|v| v.parse::<f32>().ok())
239 .filter(|&t| t > 0.0 && t < 1.0)
240 .inspect(|t| tracing::info!("MoE adaptive routing: tau {t}"));
241 let mask = moe_task_mask(&prefix, experts.len());
242 let router = if resonance_moe {
243 QTensor::from_f32(vec![0.0; experts.len() * arch.hidden_size], experts.len(), arch.hidden_size)
245 } else {
246 load_matrix(model, &router_name, force_f32, ov)?
247 };
248 if router.rows() != experts.len() {
249 return Err(CmfError::Parse(format!(
250 "{router_name}: {} rows != {} experts",
251 router.rows(),
252 experts.len()
253 )));
254 }
255 let top_k = top_k.min(experts.len());
256 let pes_name = format!("{prefix}mlp.per_expert_scale");
260 let per_expert_scale = if model.tensor(&pes_name).is_some() {
261 Some(load_f32(model, &pes_name, ov).map_err(CmfError::Parse)?)
262 } else {
263 None
264 };
265 let router_input_norm = per_expert_scale.is_some();
266 let per_expert = model.tensor(&format!("{prefix}mlp.experts.0.desc.mu")).is_some();
271 let resonance = if per_expert {
272 let ne_d = experts.len();
273 let hidden = arch.hidden_size;
274 let mut mu = Vec::with_capacity(ne_d * hidden);
275 let mut u = Vec::new();
276 let mut bias = Vec::with_capacity(ne_d);
277 let mut k = 0usize;
278 for e in 0..ne_d {
279 let m = load_f32(model, &format!("{prefix}mlp.experts.{e}.desc.mu"), ov).map_err(CmfError::Parse)?;
280 if m.len() != hidden {
281 return Err(CmfError::Parse(format!("{prefix}mlp.experts.{e}.desc.mu: {} != {hidden}", m.len())));
282 }
283 mu.extend_from_slice(&m);
284 let un = format!("{prefix}mlp.experts.{e}.desc.u");
285 if model.tensor(&un).is_some() {
286 let ue = load_f32(model, &un, ov).map_err(CmfError::Parse)?;
287 let ke = ue.len() / hidden.max(1);
288 if e == 0 {
289 k = ke;
290 }
291 if ke != k {
292 return Err(CmfError::Parse(format!("{un}: rank {ke} != {k}")));
293 }
294 u.extend_from_slice(&ue);
295 }
296 let bn = format!("{prefix}mlp.experts.{e}.desc.bias");
297 bias.push(if model.tensor(&bn).is_some() {
298 load_f32(model, &bn, ov).map_err(CmfError::Parse)?.first().copied().unwrap_or(0.0)
299 } else {
300 0.0
301 });
302 }
303 Some(crate::pipeline::Resonance { mu, u, k, bias })
304 } else if model.tensor(&format!("{prefix}mlp.desc.mu")).is_some() {
305 let mu = load_f32(model, &format!("{prefix}mlp.desc.mu"), ov).map_err(CmfError::Parse)?;
306 let ne_d = experts.len();
307 let hidden = arch.hidden_size;
308 if mu.len() != ne_d * hidden {
309 return Err(CmfError::Parse(format!("{prefix}mlp.desc.mu: {} != {ne_d}×{hidden}", mu.len())));
310 }
311 let u_name = format!("{prefix}mlp.desc.u");
312 let (u, k) = if model.tensor(&u_name).is_some() {
313 let u = load_f32(model, &u_name, ov).map_err(CmfError::Parse)?;
314 let k = u.len() / (ne_d * hidden).max(1);
315 (u, k)
316 } else {
317 (Vec::new(), 0)
318 };
319 let b_name = format!("{prefix}mlp.desc.bias");
320 let bias = if model.tensor(&b_name).is_some() {
321 load_f32(model, &b_name, ov).map_err(CmfError::Parse)?
322 } else {
323 vec![0.0; ne_d]
324 };
325 Some(crate::pipeline::Resonance { mu, u, k, bias })
326 } else {
327 None
328 };
329 let moe = MoeFfn {
330 router,
331 experts,
332 top_k,
333 route_tau,
334 norm_topk_prob: cfg.norm_topk_prob,
335 router_sigmoid: cfg.router_sigmoid,
336 expert_bias,
337 routed_scaling: cfg.routed_scaling_factor.unwrap_or(1.0),
338 shared,
339 stats: std::cell::RefCell::new(Vec::new()),
340 act_sq: std::cell::RefCell::new(Vec::new()),
341 act_rows: std::cell::RefCell::new(Vec::new()),
342 mask,
343 per_expert_scale,
344 router_input_norm,
345 resonance,
346 };
347 if model
350 .tensor(&format!("{prefix}mlp.gate_proj.weight"))
351 .is_some()
352 {
353 let norm = |suffix: &str| -> Result<Vec<f32>, CmfError> {
354 load_f32(model, &format!("{prefix}{suffix}.weight"), ov).map_err(CmfError::Parse)
355 };
356 return Ok(FfnKind::DenseMoe(Box::new(crate::pipeline::DenseMoeFfn {
357 dense: load_dense(&format!("{prefix}mlp."))?,
358 moe,
359 post_norm_1: norm("post_feedforward_layernorm_1")?,
360 pre_norm_2: norm("pre_feedforward_layernorm_2")?,
361 post_norm_2: norm("post_feedforward_layernorm_2")?,
362 })));
363 }
364 Ok(FfnKind::Moe(moe))
365}
366
367pub(crate) fn moe_task_mask(prefix: &str, ne: usize) -> Option<Vec<bool>> {
375 use std::sync::OnceLock;
376 static CFG: OnceLock<Option<(std::collections::HashMap<usize, Vec<u64>>, f64)>> =
377 OnceLock::new();
378 let cfg = CFG.get_or_init(|| {
379 let path = std::env::var("CMF_MOE_MASK").ok()?;
380 let cover = std::env::var("CMF_MOE_MASK_COVER")
381 .ok()
382 .and_then(|v| v.parse::<f64>().ok())
383 .filter(|&c| c > 0.0 && c <= 1.0)
384 .unwrap_or(0.9);
385 let text = std::fs::read_to_string(&path)
386 .map_err(|e| tracing::warn!("CMF_MOE_MASK: cannot read {path}: {e}"))
387 .ok()?;
388 let map: std::collections::HashMap<String, Vec<u64>> = serde_json::from_str(&text)
389 .map_err(|e| tracing::warn!("CMF_MOE_MASK: bad JSON in {path}: {e}"))
390 .ok()?;
391 tracing::info!("MoE task mask: {path}, cover {cover}");
392 Some((
393 map.into_iter()
394 .filter_map(|(k, v)| Some((k.parse::<usize>().ok()?, v)))
395 .collect(),
396 cover,
397 ))
398 });
399 let (stats, cover) = cfg.as_ref()?;
400 let li: usize = prefix
402 .split("layers.")
403 .nth(1)?
404 .split('.')
405 .next()?
406 .parse()
407 .ok()?;
408 let counts = stats.get(&li)?;
409 if counts.len() != ne {
410 tracing::warn!(
411 "CMF_MOE_MASK: layer {li} has {} counts, model has {ne} experts — skipped",
412 counts.len()
413 );
414 return None;
415 }
416 let total: u64 = counts.iter().sum();
417 if total == 0 {
418 return None;
419 }
420 let mut order: Vec<usize> = (0..ne).collect();
421 order.sort_unstable_by_key(|&e| std::cmp::Reverse(counts[e]));
422 let mut mask = vec![false; ne];
423 let mut acc = 0u64;
424 let mut kept = 0usize;
425 for &e in &order {
426 mask[e] = true;
427 acc += counts[e];
428 kept += 1;
429 if (acc as f64) >= cover * (total as f64) {
430 break;
431 }
432 }
433 tracing::info!(
434 "MoE task mask L{li}: {kept}/{ne} experts for {:.0}% mass",
435 cover * 100.0
436 );
437 Some(mask)
438}
439
440fn load_matrix(
441 model: &Arc<CmfModel>,
442 name: &str,
443 force_f32: bool,
444 ov: &Overlay,
445) -> Result<QTensor, CmfError> {
446 if ov.blend_touches(model, name) {
450 if let Overlay::Blend(list) = ov {
451 let entry = model
452 .tensor(name)
453 .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
454 let data =
455 blend_f32(model, name, list).map_err(|e| CmfError::Parse(format!("blend: {e}")))?;
456 return Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]));
457 }
458 }
459 let skill = match ov {
460 Overlay::One(s) => Some(*s),
461 _ => None,
462 };
463 let name: &str = &match skill {
466 Some(sid) if model.tensor(&format!("skill.{sid}.{name}")).is_some() => {
467 format!("skill.{sid}.{name}")
468 }
469 _ => name.to_string(),
470 };
471 let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
472 if force_f32 {
473 let entry = model
474 .tensor(name)
475 .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
476 if entry.shape.len() != 2 {
477 return Err(err(format!("'{name}' is not 2-D")));
478 }
479 let data = load_f32(model, name, &Overlay::None).map_err(err)?;
480 Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]))
481 } else {
482 QTensor::from_model(model, name).map_err(err)
483 }
484}
485
486impl Pipeline {
487 pub fn from_model(
489 model: &Arc<CmfModel>,
490 sampler_config: SamplerConfig,
491 ) -> Result<Self, CmfError> {
492 Self::from_model_with_skill(model, sampler_config, None)
493 }
494
495 pub fn from_model_with_skill(
501 model: &Arc<CmfModel>,
502 sampler_config: SamplerConfig,
503 skill: Option<&str>,
504 ) -> Result<Self, CmfError> {
505 match skill {
506 Some(s) => Self::from_model_with_overlay(model, sampler_config, &Overlay::One(s)),
507 None => Self::from_model_with_overlay(model, sampler_config, &Overlay::None),
508 }
509 }
510
511 pub fn from_model_with_blend(
514 model: &Arc<CmfModel>,
515 sampler_config: SamplerConfig,
516 blend: &[(String, f32)],
517 ) -> Result<Self, CmfError> {
518 Self::from_model_with_overlay(model, sampler_config, &Overlay::Blend(blend))
519 }
520
521 fn skill_file_guard(model: &CmfModel) -> Result<(), CmfError> {
522 if model.required_features & cortiq_core::format::features::SKILL_FILE != 0 {
525 return Err(CmfError::Parse(
526 "this file is a standalone SKILL, not a runnable model — attach it: \
527 cortiq skill apply <base.cmf> <this file> -o specialist.cmf"
528 .into(),
529 ));
530 }
531 Ok(())
532 }
533
534 fn from_model_with_overlay(
535 model: &Arc<CmfModel>,
536 sampler_config: SamplerConfig,
537 ov: &Overlay,
538 ) -> Result<Self, CmfError> {
539 if let Some(dir) = model.path.parent() {
544 crate::gpu::set_cache_dir(dir.to_path_buf());
545 }
546 crate::gpu::graph_unsupported_reset();
549 Self::skill_file_guard(model)?;
550 let skill = match ov {
551 Overlay::One(s) => Some(*s),
552 _ => None,
553 };
554 if let Some(sid) = skill {
555 let known = model.header.skills.iter().any(|s| s.id == sid)
556 || model.skill_tensors(sid).next().is_some();
557 if !known {
558 return Err(CmfError::Parse(format!(
559 "skill '{sid}' not in this container (header.skills: {:?})",
560 model
561 .header
562 .skills
563 .iter()
564 .map(|s| &s.id)
565 .collect::<Vec<_>>()
566 )));
567 }
568 tracing::info!(
569 "skill '{sid}': {} replacement tensors overlaid",
570 model.skill_tensors(sid).count()
571 );
572 }
573 let arch = model.arch().clone();
574 let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
575 if let Some(heads) = &arch.attention_heads_per_layer {
576 if heads.len() != arch.num_layers {
577 return Err(CmfError::Parse(format!(
578 "arch.attention_heads_per_layer has {} entries, expected {}",
579 heads.len(),
580 arch.num_layers
581 )));
582 }
583 if let Some((li, &nh)) = heads
584 .iter()
585 .enumerate()
586 .find(|(_, nh)| **nh == 0 || **nh % arch.num_kv_heads != 0)
587 {
588 return Err(CmfError::Parse(format!(
589 "layer {li} has {nh} Q heads, which must be nonzero and divisible by {} KV heads",
590 arch.num_kv_heads
591 )));
592 }
593 }
594 if arch
595 .layer_types
596 .iter()
597 .any(|t| matches!(t, LayerType::SlidingAttention))
598 && arch.sliding_window.is_none()
599 {
600 return Err(CmfError::Parse(
601 "model has SlidingAttention layers but no arch.sliding_window".into(),
602 ));
603 }
604
605 let heads_masked = model.masks.masks.iter().any(|m| {
615 m.head_masks.iter().any(|row| {
616 let mut bits = 0usize;
617 for &b in row.iter() {
618 bits += b.count_ones() as usize;
619 }
620 !row.is_empty() && bits < arch.num_attention_heads
621 })
622 });
623 let force_f32 = heads_masked; let mut tokenizer = if let Some(vocab_bytes) = &model.vocab {
627 Tokenizer::from_bytes(vocab_bytes)
628 .map_err(|e| CmfError::Parse(format!("embedded tokenizer: {e}")))?
629 } else {
630 let sidecar = model.path.with_file_name("tokenizer.json");
631 if sidecar.exists() {
632 Tokenizer::from_file(&sidecar)
633 .map_err(|e| CmfError::Parse(format!("sidecar tokenizer: {e}")))?
634 } else {
635 tracing::warn!("no tokenizer in file or sidecar — using byte-level fallback");
636 Tokenizer::byte_level()
637 }
638 };
639 if let Some(tc) = &model.header.tokenizer_config {
641 tokenizer.chat_template = tc.chat_template.clone();
642 tokenizer.extra_eos.extend(tc.eos_token_ids.iter().copied());
643 if tokenizer.bos_token_id.is_none() {
644 tokenizer.bos_token_id = tc.bos_token_id;
645 }
646 tracing::info!(
647 "chat bundle: template {} chars, {} stop ids",
648 tc.chat_template.as_deref().map(str::len).unwrap_or(0),
649 tc.eos_token_ids.len()
650 );
651 }
652 if arch.arch_name.to_lowercase().contains("gemma") && tokenizer.bos_token_id.is_some() {
656 tokenizer.add_bos = true;
657 }
658
659 let embed_tokens = load_matrix(model, "model.embed_tokens.weight", false, ov)?;
661 let final_norm = load_f32(model, "model.norm.weight", ov).map_err(err)?;
662 let lm_head = if model.tensor("lm_head.weight").is_some() {
663 load_matrix(model, "lm_head.weight", false, ov)?
664 } else if arch.tie_word_embeddings {
665 load_matrix(model, "model.embed_tokens.weight", false, ov)?
667 } else {
668 return Err(CmfError::MissingTensor(
669 "lm_head.weight (and tie_word_embeddings is false)".into(),
670 ));
671 };
672
673 let has_linear = arch
675 .layer_types
676 .iter()
677 .any(|t| matches!(t, LayerType::LinearAttention));
678 let mut vmf_cfg = None;
679 let mut gdn_cfg = None;
680 if has_linear {
681 let lc = arch.linear_core.as_ref().ok_or_else(|| {
682 CmfError::Parse(
683 "model has LinearAttention layers but no arch.linear_core — \
684 reconvert with the current converter"
685 .into(),
686 )
687 })?;
688 let need = |v: Option<usize>, name: &str| {
689 v.ok_or_else(|| CmfError::Parse(format!("linear core needs arch.{name}")))
690 };
691 match lc.kind.as_str() {
692 "vmf_phase" => {
693 vmf_cfg = Some(VmfPhaseCfg {
694 num_heads: lc.num_heads,
695 nphase: need(lc.nphase, "linear_core.nphase")?,
696 value_head_dim: lc.value_head_dim,
697 hidden_size: arch.hidden_size,
698 phase_mass: std::env::var("CMF_PHASE_MASS")
701 .ok()
702 .and_then(|v| v.parse().ok())
703 .unwrap_or(0.0),
704 });
705 }
706 "gated_delta_net" => {
707 gdn_cfg = Some(GdnCfg {
708 num_v_heads: lc.num_heads,
709 num_k_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
710 key_head_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
711 value_head_dim: lc.value_head_dim,
712 conv_kernel: need(arch.linear_conv_kernel_dim, "linear_conv_kernel_dim")?,
713 hidden_size: arch.hidden_size,
714 rms_eps: arch.rms_norm_eps,
715 });
716 }
717 other => {
718 return Err(CmfError::Parse(format!(
719 "unknown linear core '{other}' (this runtime executes: \
720 gated_delta_net, vmf_phase)"
721 )));
722 }
723 }
724 }
725
726 let has_kda = arch.layer_types.iter().any(|t| matches!(t, LayerType::Kda));
728 let kda_cfg = if has_kda {
729 let need = |v: Option<usize>, name: &str| {
730 v.ok_or_else(|| CmfError::Parse(format!("KDA core needs arch.{name}")))
731 };
732 Some(crate::linear_core::KdaCfg {
733 num_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
734 head_k_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
735 head_v_dim: need(arch.linear_value_head_dim, "linear_value_head_dim")?,
736 conv_kernel: need(arch.linear_conv_kernel_dim, "linear_conv_kernel_dim")?,
737 hidden_size: arch.hidden_size,
738 rms_eps: arch.rms_norm_eps,
739 })
740 } else {
741 None
742 };
743
744 let has_short_conv = arch
746 .layer_types
747 .iter()
748 .any(|t| matches!(t, LayerType::ShortConv));
749 let short_conv_cfg = if has_short_conv {
750 Some(ShortConvCfg {
751 hidden_size: arch.hidden_size,
752 kernel: arch.linear_conv_kernel_dim.ok_or_else(|| {
753 CmfError::Parse(
754 "model has ShortConv layers but no arch.linear_conv_kernel_dim — \
755 reconvert with the current converter"
756 .into(),
757 )
758 })?,
759 })
760 } else {
761 None
762 };
763
764 let load_full_attn = |prefix: &str, layer: Option<usize>| -> Result<AttnKind, CmfError> {
766 let t = |suffix: &str| load_matrix(model, &format!("{prefix}{suffix}"), force_f32, ov);
767 let n = |suffix: &str| -> Option<Vec<f32>> {
768 model
769 .tensor(&format!("{prefix}{suffix}"))
770 .and_then(|_| load_f32(model, &format!("{prefix}{suffix}"), ov).ok())
771 };
772 if let Some(mla) = arch.mla.as_ref() {
774 let (q_proj, q_a, q_a_norm) = if mla.q_lora_rank.is_some() {
776 (
777 t("self_attn.q_b_proj.weight")?,
778 Some(t("self_attn.q_a_proj.weight")?),
779 Some(n("self_attn.q_a_layernorm.weight").ok_or_else(|| {
780 CmfError::Parse(format!("{prefix}: MLA needs q_a_layernorm"))
781 })?),
782 )
783 } else {
784 (t("self_attn.q_proj.weight")?, None, None)
785 };
786 let hd = mla.qk_rope_head_dim + mla.qk_nope_head_dim;
787 let nh = q_proj.rows() / hd;
788 let mut scale = 1.0 / (hd as f32).sqrt();
791 if let Some(y) = arch.yarn.as_ref() {
792 if let Some(m) = y.mscale_all_dim.filter(|&m| m > 0.0) {
793 let ms = 0.1 * m * y.factor.ln() + 1.0;
794 scale *= ms * ms;
795 }
796 }
797 return Ok(AttnKind::Mla(Box::new(crate::pipeline::MlaWeights {
798 q_proj,
799 q_a,
800 q_a_norm,
801 kv_a: t("self_attn.kv_a_proj_with_mqa.weight")?,
802 kv_a_norm: n("self_attn.kv_a_layernorm.weight").ok_or_else(|| {
803 CmfError::Parse(format!("{prefix}: MLA needs kv_a_layernorm"))
804 })?,
805 kv_b: t("self_attn.kv_b_proj.weight")?,
806 o_proj: t("self_attn.o_proj.weight")?,
807 nh,
808 qk_rope: mla.qk_rope_head_dim,
809 qk_nope: mla.qk_nope_head_dim,
810 v_dim: mla.v_head_dim,
811 lora: mla.kv_lora_rank,
812 scale,
813 nope: mla.nope,
814 })));
815 }
816 let wq = t("self_attn.q_proj.weight")?;
817 let nh = layer
818 .and_then(|li| {
819 arch.attention_heads_per_layer
820 .as_ref()
821 .and_then(|v| v.get(li).copied())
822 })
823 .unwrap_or(arch.num_attention_heads);
824 let output_gate = arch.global_head_dim.is_none() && wq.rows() == 2 * nh * arch.head_dim;
828 let is_global_layer = arch.global_head_dim.is_some()
831 && layer.is_some_and(|li| {
832 arch.sliding_window_pattern
833 .is_some_and(|p| p > 0 && (li + 1) % p == 0)
834 });
835 let expect = if is_global_layer {
836 nh * arch.global_head_dim.unwrap_or(arch.head_dim)
837 } else {
838 nh * arch.head_dim
839 };
840 if !output_gate && wq.rows() != expect {
841 return Err(CmfError::Parse(format!(
842 "{prefix}self_attn.q_proj.weight rows={} != heads({nh}) * head_dim({})",
843 wq.rows(),
844 expect / nh.max(1)
845 )));
846 }
847 let gate_name = format!("{prefix}self_attn.g_proj.weight");
848 let softplus_gate = if model.tensor(&gate_name).is_some() {
849 let gate = load_matrix(model, &gate_name, force_f32, ov)?;
850 if gate.cols() != arch.hidden_size {
851 return Err(CmfError::Parse(format!(
852 "{gate_name} cols={} != hidden_size ({})",
853 gate.cols(),
854 arch.hidden_size
855 )));
856 }
857 let per_head = if gate.rows() == nh {
858 true
859 } else if gate.rows() == nh * arch.head_dim {
860 false
861 } else {
862 return Err(CmfError::Parse(format!(
863 "{gate_name} rows={} must equal heads ({nh}) or heads*head_dim ({})",
864 gate.rows(),
865 nh * arch.head_dim
866 )));
867 };
868 Some((gate, per_head))
869 } else {
870 None
871 };
872 let bias = match (
874 n("self_attn.q_proj.bias"),
875 n("self_attn.k_proj.bias"),
876 n("self_attn.v_proj.bias"),
877 ) {
878 (Some(a), Some(b), Some(c)) => Some((a, b, c)),
879 _ => None,
880 };
881 Ok(AttnKind::Full {
882 wq,
883 wk: t("self_attn.k_proj.weight")?,
884 wv: t("self_attn.v_proj.weight")?,
885 wo: t("self_attn.o_proj.weight")?,
886 q_norm: n("self_attn.q_norm.weight"),
887 k_norm: n("self_attn.k_norm.weight"),
888 output_gate,
889 softplus_gate,
890 bias,
891 })
892 };
893
894 let load_linear_attn = |prefix: &str| -> Result<AttnKind, CmfError> {
895 if gdn_cfg.is_some() {
896 let t = |suffix: &str| {
898 load_matrix(
899 model,
900 &format!("{prefix}linear_attn.{suffix}"),
901 force_f32,
902 ov,
903 )
904 };
905 let f = |suffix: &str| {
906 load_f32(model, &format!("{prefix}linear_attn.{suffix}"), ov).map_err(err)
907 };
908 return Ok(AttnKind::LinearGdn(GdnWeights {
909 in_proj_qkv: t("in_proj_qkv.weight")?,
910 in_proj_z: t("in_proj_z.weight")?,
911 in_proj_a: t("in_proj_a.weight")?,
912 in_proj_b: t("in_proj_b.weight")?,
913 conv1d: f("conv1d.weight")?,
914 a_log: f("A_log")?,
915 dt_bias: f("dt_bias")?,
916 norm: f("norm.weight")?,
917 out_proj: t("out_proj.weight")?,
918 }));
919 }
920 let t = |suffix: &str| {
921 load_matrix(model, &format!("{prefix}vmf_attn.{suffix}"), force_f32, ov)
922 };
923 let a_log = load_f32(model, &format!("{prefix}vmf_attn.A_log"), ov).map_err(err)?;
924 let k_gate = if model
928 .tensor(&format!("{prefix}vmf_attn.k_gate.weight"))
929 .is_some()
930 {
931 Some((
932 t("k_gate.weight")?,
933 load_f32(model, &format!("{prefix}vmf_attn.k_gate.bias"), ov).map_err(err)?,
934 ))
935 } else {
936 None
937 };
938 Ok(AttnKind::Linear(VmfPhaseWeights {
939 thq: t("thq.weight")?,
940 thk: t("thk.weight")?,
941 v_proj: t("v_proj.weight")?,
942 out_proj: t("out_proj.weight")?,
943 decay: a_log.iter().map(|&a| (-(a as f64).exp()).exp()).collect(),
944 k_gate,
945 }))
946 };
947
948 let load_short_conv = |prefix: &str| -> Result<AttnKind, CmfError> {
952 let t = |suffix: &str| {
953 load_matrix(
954 model,
955 &format!("{prefix}short_conv.{suffix}"),
956 force_f32,
957 ov,
958 )
959 };
960 Ok(AttnKind::ShortConv(ShortConvWeights {
961 in_proj: t("in_proj.weight")?,
962 conv: load_f32(model, &format!("{prefix}short_conv.conv.weight"), ov)
963 .map_err(err)?,
964 out_proj: t("out_proj.weight")?,
965 }))
966 };
967
968 let load_kda = |prefix: &str| -> Result<AttnKind, CmfError> {
972 let t = |suffix: &str| {
973 load_matrix(model, &format!("{prefix}kda_attn.{suffix}"), force_f32, ov)
974 };
975 let f = |suffix: &str| {
976 load_f32(model, &format!("{prefix}kda_attn.{suffix}"), ov).map_err(err)
977 };
978 let gate = if model
979 .tensor(&format!("{prefix}kda_attn.g_proj.weight"))
980 .is_some()
981 {
982 crate::linear_core::KdaOutGate::Full(t("g_proj.weight")?)
983 } else {
984 crate::linear_core::KdaOutGate::LowRank(
985 t("g_a_proj.weight")?,
986 t("g_b_proj.weight")?,
987 )
988 };
989 Ok(AttnKind::Kda(Box::new(crate::linear_core::KdaWeights {
990 q_proj: t("q_proj.weight")?,
991 k_proj: t("k_proj.weight")?,
992 v_proj: t("v_proj.weight")?,
993 conv_q: f("q_conv1d.weight")?,
994 conv_k: f("k_conv1d.weight")?,
995 conv_v: f("v_conv1d.weight")?,
996 f_a: t("f_a_proj.weight")?,
997 f_b: t("f_b_proj.weight")?,
998 dt_bias: f("dt_bias")?,
999 a_log: f("A_log")?,
1000 b_proj: t("b_proj.weight")?,
1001 gate,
1002 o_norm: f("o_norm.weight")?,
1003 o_proj: t("o_proj.weight")?,
1004 gate_lower_bound: arch.kda_gate_lower_bound.map(|v| v as f32),
1005 })))
1006 };
1007
1008 fn anyhow_like(ok: bool) -> Result<(), ()> {
1009 if ok { Ok(()) } else { Err(()) }
1010 }
1011 let mut layers = Vec::with_capacity(arch.num_layers);
1012 let is_g3n = arch.g3n.is_some();
1013 let owns_its_layers = is_g3n || arch.arch_name == "deepseek_v4";
1018 for li in 0..(if owns_its_layers { 0 } else { arch.num_layers }) {
1019 let prefix = format!("model.layers.{li}.");
1020 let attn = match arch.layer_types.get(li) {
1021 Some(LayerType::LinearAttention) => load_linear_attn(&prefix)?,
1022 Some(LayerType::Kda) => load_kda(&prefix)?,
1023 Some(LayerType::ShortConv) => load_short_conv(&prefix)?,
1024 _ => load_full_attn(&prefix, Some(li))?,
1025 };
1026 let pre_ffn = format!("{prefix}pre_feedforward_layernorm.weight");
1030 let sandwich = model.tensor(&pre_ffn).is_some();
1031 layers.push(LayerWeights {
1032 input_norm: load_f32(model, &format!("{prefix}input_layernorm.weight"), ov)
1033 .map_err(err)?,
1034 post_norm: if sandwich {
1035 load_f32(model, &pre_ffn, ov).map_err(err)?
1036 } else {
1037 load_f32(
1038 model,
1039 &format!("{prefix}post_attention_layernorm.weight"),
1040 ov,
1041 )
1042 .map_err(err)?
1043 },
1044 attn_out_norm: if sandwich {
1045 Some(
1046 load_f32(
1047 model,
1048 &format!("{prefix}post_attention_layernorm.weight"),
1049 ov,
1050 )
1051 .map_err(err)?,
1052 )
1053 } else {
1054 None
1055 },
1056 ffn_out_norm: if sandwich {
1057 Some(
1058 load_f32(
1059 model,
1060 &format!("{prefix}post_feedforward_layernorm.weight"),
1061 ov,
1062 )
1063 .map_err(err)?,
1064 )
1065 } else {
1066 None
1067 },
1068 layer_scale: model
1070 .tensor(&format!("{prefix}layer_scalar"))
1071 .and_then(|_| {
1072 load_f32(model, &format!("{prefix}layer_scalar"), ov)
1073 .ok()
1074 .and_then(|v| v.first().copied())
1075 }),
1076 ffn: build_layer_ffn(model, &arch, li, false, ov)?,
1078 attn,
1079 });
1080 }
1081
1082 let mtp_present = model
1091 .tensor("model.mtp.layers.0.self_attn.q_proj.weight")
1092 .is_some()
1093 || model.tensor("model.mtp.eh_proj.weight").is_some();
1094 let dsv4_mtp = model.tensor("model.mtp.0.main_proj.weight").is_some();
1099 if arch.mtp.is_some() && !mtp_present && !dsv4_mtp {
1100 tracing::info!(
1101 "header declares an MTP head but the file carries none — \
1102 loading without it"
1103 );
1104 }
1105 let mtp = if let Some(cfg) = arch.mtp.as_ref().filter(|_| mtp_present) {
1106 if cfg.num_layers != 1 {
1107 return Err(CmfError::Parse(format!(
1108 "MTP with {} blocks not supported yet (only 1)",
1109 cfg.num_layers
1110 )));
1111 }
1112 let p = "model.mtp.";
1113 let attn = load_full_attn("model.mtp.layers.0.", None)?;
1114 Some(MtpModule {
1115 enorm: load_f32(model, &format!("{p}enorm.weight"), ov).map_err(err)?,
1116 hnorm: load_f32(model, &format!("{p}hnorm.weight"), ov).map_err(err)?,
1117 eh_proj: load_matrix(model, &format!("{p}eh_proj.weight"), false, ov)?,
1118 layer: LayerWeights {
1119 attn_out_norm: None,
1120 ffn_out_norm: None,
1121 layer_scale: None,
1122 input_norm: load_f32(model, &format!("{p}layers.0.input_layernorm.weight"), ov)
1123 .map_err(err)?,
1124 post_norm: load_f32(
1125 model,
1126 &format!("{p}layers.0.post_attention_layernorm.weight"),
1127 ov,
1128 )
1129 .map_err(err)?,
1130 ffn: build_ffn_at(model, &arch, &format!("{p}layers.0."), false, ov)?,
1134 attn,
1135 },
1136 final_norm: load_f32(model, &format!("{p}norm.weight"), ov).map_err(err)?,
1137 kv: LayerKvCache::new(arch.num_kv_heads, arch.head_dim),
1138 })
1139 } else {
1140 None
1141 };
1142
1143 tracing::info!(
1144 "Pipeline loaded: {} | {}L ({} linear) | {:.2}B params | storage: {} | MTP: {}",
1145 arch.arch_name,
1146 arch.num_layers,
1147 arch.layer_types
1148 .iter()
1149 .filter(|t| matches!(t, LayerType::LinearAttention))
1150 .count(),
1151 model.total_param_count() as f64 / 1e9,
1152 if force_f32 {
1153 "f32 (masked)"
1154 } else {
1155 "quantized mmap"
1156 },
1157 if mtp.is_some() { "yes" } else { "no" }
1158 );
1159
1160 let cap = std::env::var("CMF_MAX_SEQ")
1171 .ok()
1172 .and_then(|v| v.parse::<usize>().ok())
1173 .unwrap_or(32_768);
1174 let max_seq_len = arch.max_position_embeddings.min(cap);
1175
1176 let total_layers = arch.num_layers * arch.num_loops;
1178
1179 let mut pipeline = Pipeline::new(
1180 tokenizer,
1181 PipelineWeights {
1182 embed_tokens,
1183 layers,
1184 lm_head,
1185 final_norm,
1186 },
1187 arch.hidden_size,
1188 arch.intermediate_size,
1189 arch.num_attention_heads,
1190 arch.num_kv_heads,
1191 arch.head_dim,
1192 total_layers,
1193 arch.num_layers, arch.loop_final_norm,
1195 arch.vocab_size,
1196 arch.rms_norm_eps,
1197 arch.rope_theta as f32,
1198 arch.norm_style,
1199 max_seq_len,
1200 sampler_config,
1201 );
1202 let rotary = ((arch.head_dim as f32 * arch.partial_rotary_factor) as usize).max(2);
1203 pipeline.set_rotary(rotary, arch.rope_theta as f32);
1204 pipeline.attention_heads_per_layer = arch.attention_heads_per_layer.clone();
1205 if let Some(yarn) = &arch.yarn {
1206 pipeline.inv_freq = std::sync::Arc::new(crate::attention::yarn_inv_freq(
1207 rotary,
1208 arch.rope_theta as f32,
1209 yarn.factor,
1210 yarn.original_max_position_embeddings,
1211 yarn.beta_fast,
1212 yarn.beta_slow,
1213 ));
1214 pipeline.rope_scale = yarn.attention_factor;
1215 }
1216 pipeline.embed_multiplier = arch.embed_multiplier;
1220 pipeline.logit_multiplier = arch.logit_multiplier;
1221 if let Some(qpas) = arch.query_pre_attn_scalar {
1222 pipeline.attn_scale = 1.0 / (qpas as f32).sqrt();
1223 }
1224 if let (Some(w), Some(p)) = (arch.sliding_window, arch.sliding_window_pattern) {
1225 pipeline.swa = Some((w, p));
1226 if let Some(base) = arch.rope_local_base_freq {
1227 pipeline.inv_freq_local = Some(std::sync::Arc::new(
1228 crate::attention::rope_inv_freq(rotary, base as f32),
1229 ));
1230 }
1231 }
1232 let explicit_sliding: Vec<bool> = arch
1233 .layer_types
1234 .iter()
1235 .map(|t| matches!(t, cortiq_core::LayerType::SlidingAttention))
1236 .collect();
1237 if explicit_sliding.iter().any(|&v| v) {
1238 pipeline.sliding_layers = Some(explicit_sliding);
1239 if let Some(w) = arch.sliding_window {
1240 pipeline.swa = Some((w, usize::MAX));
1241 }
1242 let local_rotary = ((arch.head_dim as f32
1243 * arch
1244 .local_partial_rotary_factor
1245 .unwrap_or(arch.partial_rotary_factor))
1246 as usize)
1247 .max(2);
1248 pipeline.rotary_dim_local = Some(local_rotary);
1249 if let Some(base) = arch.rope_local_base_freq {
1250 pipeline.inv_freq_local = Some(std::sync::Arc::new(
1251 crate::attention::rope_inv_freq(local_rotary, base as f32),
1252 ));
1253 }
1254 }
1255 if let (Some(ghd), Some(gkv)) = (arch.global_head_dim, arch.num_global_kv_heads) {
1259 pipeline.global_attn = Some((ghd, gkv));
1260 let prf = arch.global_partial_rotary_factor.unwrap_or(1.0);
1261 let half = ghd / 2;
1262 let ra = (((prf * ghd as f32) as usize) / 2).min(half);
1263 let mut f = vec![0.0f32; half];
1264 for (i, slot) in f.iter_mut().enumerate().take(ra) {
1265 *slot = 1.0 / (arch.rope_theta as f32).powf(2.0 * i as f32 / ghd as f32);
1266 }
1267 pipeline.inv_freq_global = Some(std::sync::Arc::new(f));
1268 let global_at = |li: usize| -> bool {
1273 match &pipeline.sliding_layers {
1274 Some(map) => !map.get(li).copied().unwrap_or(false),
1275 None => pipeline
1276 .swa
1277 .map(|(_, p)| p > 0 && p != usize::MAX && (li + 1) % p == 0)
1278 .unwrap_or(false),
1279 }
1280 };
1281 for li in 0..arch.num_layers {
1282 if global_at(li) {
1283 pipeline.kv_cache.layers[li] = crate::kv_cache::LayerKvCache::new(gkv, ghd);
1284 }
1285 }
1286 }
1287 if let Some(mla) = arch.mla.as_ref() {
1290 let hd = mla.qk_rope_head_dim + mla.qk_nope_head_dim;
1291 pipeline.head_dim = hd;
1292 pipeline.num_kv_heads = arch.num_attention_heads;
1293 pipeline.rotary_dim = mla.qk_rope_head_dim;
1294 let half = mla.qk_rope_head_dim / 2;
1295 let mut f = vec![0.0f32; half];
1296 for (i, slot) in f.iter_mut().enumerate() {
1297 *slot = 1.0
1298 / (arch.rope_theta as f32).powf(2.0 * i as f32 / mla.qk_rope_head_dim as f32);
1299 }
1300 pipeline.inv_freq = std::sync::Arc::new(f);
1301 for li in 0..arch.num_layers {
1302 pipeline.kv_cache.layers[li] =
1303 crate::kv_cache::LayerKvCache::new(arch.num_attention_heads, hd);
1304 }
1305 }
1306 if let Some(fac) = &arch.rope_freq_factors {
1310 let mut f = pipeline.inv_freq.as_ref().clone();
1311 for (i, v) in f.iter_mut().enumerate() {
1312 if let Some(&d) = fac.get(i) {
1313 *v /= d as f32;
1314 }
1315 }
1316 pipeline.inv_freq = std::sync::Arc::new(f);
1317 }
1318 pipeline.attn_v_norm = arch.attn_v_norm;
1319 pipeline.final_softcap = arch.final_logit_softcapping.map(|c| c as f32);
1320 if let Some(ncl) = arch.head_clusters {
1322 let cm = load_f32(model, "lm_head.clusters.weight", ov).map_err(err)?;
1323 if cm.len() != ncl * arch.hidden_size {
1324 return Err(CmfError::Parse(format!(
1325 "lm_head.clusters.weight: {} != {ncl}×{}",
1326 cm.len(),
1327 arch.hidden_size
1328 )));
1329 }
1330 pipeline.head_clusters = Some(std::sync::Arc::new(cm));
1331 }
1332 pipeline.attn_softcap = arch.attn_logit_softcapping.unwrap_or(0.0) as f32;
1333 pipeline.vmf_cfg = vmf_cfg;
1334 pipeline.gdn_cfg = gdn_cfg;
1335 pipeline.kda_cfg = kda_cfg;
1336 if let Some(gc) = arch.g3n.as_ref() {
1337 use crate::g3n::{G3nAltUp, G3nGlobals, G3nLaurel, G3nLayer};
1338 anyhow_like(gc.altup_num_inputs == crate::g3n::ALTUP_N).map_err(|_| {
1339 CmfError::Parse(format!(
1340 "g3n: altup_num_inputs {} != supported {}",
1341 gc.altup_num_inputs,
1342 crate::g3n::ALTUP_N
1343 ))
1344 })?;
1345 let t = |name: &str| load_matrix(model, name, force_f32, ov);
1346 let f = |name: &str| load_f32(model, name, ov).map_err(err);
1347 let mut altup_proj = Vec::new();
1348 let mut altup_unembed = Vec::new();
1349 for i in 0..crate::g3n::ALTUP_N - 1 {
1350 altup_proj.push(t(&format!("model.altup_projections.{i}.weight"))?);
1351 altup_unembed.push(t(&format!("model.altup_unembed_projections.{i}.weight"))?);
1352 }
1353 let first_shared = arch.num_layers.saturating_sub(gc.num_kv_shared_layers);
1354 let sliding_of = |li: usize| {
1355 matches!(
1356 arch.layer_types.get(li),
1357 Some(cortiq_core::LayerType::SlidingAttention)
1358 )
1359 };
1360 let mut g3n_layers = Vec::with_capacity(arch.num_layers);
1361 for li in 0..arch.num_layers {
1362 let pfx = format!("model.layers.{li}.");
1363 let shared = li >= first_shared && first_shared > 0;
1364 let share_src = if shared {
1365 let want = sliding_of(li);
1366 (0..first_shared).rev().find(|&j| sliding_of(j) == want)
1367 } else {
1368 None
1369 };
1370 g3n_layers.push(G3nLayer {
1371 altup: G3nAltUp {
1372 router_norm: f(&format!("{pfx}altup.router_norm.weight"))?,
1373 modality_router: t(&format!("{pfx}altup.modality_router.weight"))?,
1374 prediction_coefs: t(&format!("{pfx}altup.prediction_coefs.weight"))?,
1375 correction_coefs: t(&format!("{pfx}altup.correction_coefs.weight"))?,
1376 correct_output_scale: f(&format!("{pfx}altup.correct_output_scale"))?,
1377 },
1378 laurel: G3nLaurel {
1379 left: t(&format!("{pfx}laurel.linear_left.weight"))?,
1380 right: t(&format!("{pfx}laurel.linear_right.weight"))?,
1381 post_norm: f(&format!("{pfx}laurel.post_laurel_norm.weight"))?,
1382 },
1383 input_norm: f(&format!("{pfx}input_layernorm.weight"))?,
1384 post_attn_norm: f(&format!("{pfx}post_attention_layernorm.weight"))?,
1385 pre_ffw_norm: f(&format!("{pfx}pre_feedforward_layernorm.weight"))?,
1386 post_ffw_norm: f(&format!("{pfx}post_feedforward_layernorm.weight"))?,
1387 wq: t(&format!("{pfx}self_attn.q_proj.weight"))?,
1388 wk: if shared {
1389 None
1390 } else {
1391 Some(t(&format!("{pfx}self_attn.k_proj.weight"))?)
1392 },
1393 wv: if shared {
1394 None
1395 } else {
1396 Some(t(&format!("{pfx}self_attn.v_proj.weight"))?)
1397 },
1398 wo: t(&format!("{pfx}self_attn.o_proj.weight"))?,
1399 q_norm: f(&format!("{pfx}self_attn.q_norm.weight"))?,
1400 k_norm: if shared {
1401 None
1402 } else {
1403 Some(f(&format!("{pfx}self_attn.k_norm.weight"))?)
1404 },
1405 kv_share_src: share_src,
1406 sliding: sliding_of(li),
1407 gate: t(&format!("{pfx}mlp.gate_proj.weight"))?,
1408 up: t(&format!("{pfx}mlp.up_proj.weight"))?,
1409 down: t(&format!("{pfx}mlp.down_proj.weight"))?,
1410 sparsity: gc.activation_sparsity.get(li).copied().unwrap_or(0.0),
1411 ple_gate: t(&format!("{pfx}per_layer_input_gate.weight"))?,
1412 ple_proj: t(&format!("{pfx}per_layer_projection.weight"))?,
1413 post_ple_norm: f(&format!("{pfx}post_per_layer_input_norm.weight"))?,
1414 });
1415 }
1416 let hd = arch.head_dim;
1417 let globals = G3nGlobals {
1418 altup_proj,
1419 altup_unembed,
1420 ple_embed: t("model.embed_tokens_per_layer.weight")?,
1421 ple_model_proj: t("model.per_layer_model_projection.weight")?,
1422 ple_norm: f("model.per_layer_projection_norm.weight")?,
1423 ple_vocab: gc.ple_vocab,
1424 ple_dim: gc.ple_dim,
1425 num_layers: arch.num_layers,
1426 hidden: arch.hidden_size,
1427 rms_eps: arch.rms_norm_eps,
1428 inv_freq_local: crate::attention::rope_inv_freq(
1429 hd,
1430 arch.rope_local_base_freq.unwrap_or(10_000.0) as f32,
1431 ),
1432 inv_freq_global: crate::attention::rope_inv_freq(hd, arch.rope_theta as f32),
1433 window: arch.sliding_window.unwrap_or(512),
1434 };
1435 pipeline.g3n = Some(Box::new((globals, g3n_layers)));
1436 }
1437 if arch.arch_name == "deepseek_v4" {
1442 let moe = arch
1443 .moe
1444 .as_ref()
1445 .ok_or_else(|| CmfError::Parse("deepseek_v4: no moe config".into()))?;
1446 let cfg = crate::dsv4::Dsv4Cfg {
1447 dim: arch.hidden_size,
1448 n_heads: arch.num_attention_heads,
1449 head_dim: arch.head_dim,
1450 rope_head_dim: if arch.partial_rotary_factor < 1.0 {
1457 (((arch.head_dim as f32 * arch.partial_rotary_factor) as usize) & !1)
1458 .clamp(2, arch.head_dim)
1459 } else {
1460 64.min(arch.head_dim)
1461 },
1462 q_lora_rank: 0,
1466 o_lora_rank: 0,
1467 o_groups: 8,
1474 hc_mult: 4,
1475 hc_sinkhorn_iters: 20,
1476 hc_eps: 1e-6,
1477 norm_eps: arch.rms_norm_eps as f32,
1478 n_routed_experts: moe.num_experts,
1479 top_k: moe.top_k,
1480 moe_inter: moe.moe_intermediate_size,
1481 route_scale: moe.routed_scaling_factor.unwrap_or(1.0),
1482 swiglu_limit: 10.0,
1488 window: arch.sliding_window.unwrap_or(128),
1489 index_topk: 512,
1490 vocab: arch.vocab_size,
1491 };
1492 let (g, dl) = crate::dsv4::load(model, &cfg, arch.num_layers)
1493 .map_err(|e| CmfError::Parse(format!("deepseek_v4: {e}")))?;
1494 let mut cfg = cfg;
1499 if let Some(l0) = dl.first() {
1500 cfg.q_lora_rank = l0.wq_a.rows();
1501 let attn_width = arch.num_attention_heads * arch.head_dim;
1502 if l0.wo_a.cols() > 0 && attn_width % l0.wo_a.cols() == 0 {
1503 cfg.o_groups = (attn_width / l0.wo_a.cols()).max(1);
1504 }
1505 cfg.o_lora_rank = l0.wo_b.cols() / cfg.o_groups.max(1);
1506 cfg.hc_mult = (l0.hc_attn_fn.len() / l0.hc_attn_base.len().max(1)) / cfg.dim.max(1);
1507 if cfg.hc_mult == 0 {
1508 cfg.hc_mult = 4;
1509 }
1510 }
1511 let (yf, yo, ybf, ybs) = match &arch.yarn {
1524 Some(y) => (
1525 y.factor,
1526 y.original_max_position_embeddings,
1527 y.beta_fast,
1528 y.beta_slow,
1529 ),
1530 None => {
1531 tracing::warn!(
1532 "deepseek_v4: the header carries no YaRN profile — \
1533 falling back to the release's (factor 16, original \
1534 65536, beta 32/1). Re-converting with a build that \
1535 reads rope_scaling.type would make this exact."
1536 );
1537 (16.0, 65536, 32.0, 1.0)
1538 }
1539 };
1540 pipeline.inv_freq = std::sync::Arc::new(crate::attention::yarn_inv_freq(
1541 cfg.rope_head_dim,
1542 arch.rope_theta as f32,
1543 yf,
1544 yo,
1545 ybf,
1546 ybs,
1547 ));
1548 if let Ok(stats) = std::env::var("CMF_MOE_PIN") {
1553 let cover = std::env::var("CMF_MOE_PIN_COVER")
1554 .ok()
1555 .and_then(|v| v.parse::<f64>().ok())
1556 .filter(|&c| c > 0.0 && c <= 1.0)
1557 .unwrap_or(0.95);
1558 let hot = crate::pin::hot_experts(&stats, cover);
1559 let mut names: Vec<String> = Vec::new();
1560 for e in &model.tensors {
1561 let is_expert = e.name.contains(".mlp.experts.");
1562 if !is_expert {
1563 names.push(e.name.clone()); }
1565 }
1566 let mut kept_experts = 0usize;
1567 if let Some(hot) = &hot {
1568 for (li, experts) in hot {
1569 for e in experts {
1570 for w in ["gate_proj", "up_proj", "down_proj"] {
1571 names.push(format!("model.layers.{li}.mlp.experts.{e}.{w}.weight"));
1572 }
1573 kept_experts += 1;
1574 }
1575 }
1576 }
1577 let r = crate::pin::pin_tensors(model, &names);
1578 tracing::info!(
1579 "закреплено {:.1} ГБ ({} тензоров, горячих экспертов {kept_experts}, покрытие {cover}); лимит {}",
1580 r.bytes as f64 / 1e9,
1581 r.tensors,
1582 r.limit
1583 .map(|l| format!("{:.1} ГБ", l as f64 / 1e9))
1584 .unwrap_or_else(|| "неизвестен".into())
1585 );
1586 if r.skipped > 0 {
1587 tracing::warn!("не закреплено тензоров: {}", r.skipped);
1588 }
1589 }
1590 let st = crate::dsv4::Dsv4State::new(arch.num_layers);
1591 let depth = std::env::var("CMF_DSV4_MTP_DEPTH")
1595 .ok()
1596 .and_then(|v| v.parse::<usize>().ok())
1597 .unwrap_or(3);
1598 pipeline.dsv4_mtp = crate::dsv4::load_mtp(model, &cfg, depth);
1599 crate::dsv4::dspark_reserve_note(&pipeline.dsv4_mtp, &cfg, &dl);
1601 pipeline.dsv4 = Some(Box::new((g, dl, cfg, st)));
1602 }
1603 pipeline.short_conv_cfg = short_conv_cfg;
1604 pipeline.mtp = mtp;
1605 pipeline.install_dynamic_routing(model, false);
1606 match ov {
1610 Overlay::One(sid) => {
1611 pipeline.dyn_active = model.header.skills.iter().position(|s| &s.id == sid);
1612 }
1613 Overlay::Blend(_) => pipeline.dyn_blend_loaded = true,
1614 Overlay::None => {}
1615 }
1616 if let Some(c) = &model.header.calibration {
1619 pipeline.set_calib_temp(c.temperature);
1620 }
1621 let o1 = match crate::nystrom::o1_from_env() {
1627 crate::nystrom::O1Env::Off => None,
1628 crate::nystrom::O1Env::On(cfg) => Some(cfg),
1629 crate::nystrom::O1Env::Unset => model
1630 .header
1631 .provenance
1632 .as_ref()
1633 .and_then(|p| p.get("o1_attn"))
1634 .and_then(crate::nystrom::O1Cfg::from_json),
1635 };
1636 if o1.is_some() {
1637 if pipeline.attn_softcap > 0.0 {
1638 return Err(CmfError::Parse(
1639 "--o1 with attention-logit soft-capping (Gemma-2) is not supported: \
1640 the streaming operator has no capped-score form"
1641 .into(),
1642 ));
1643 }
1644 pipeline.set_o1(o1);
1645 }
1646 Ok(pipeline)
1647 }
1648
1649 pub(crate) fn install_dynamic_routing(&mut self, model: &Arc<CmfModel>, force_f32: bool) {
1654 self.model = Some(model.clone());
1655 self.dyn_force_f32 = force_f32;
1656 let mut per_skill = Vec::with_capacity(model.header.skills.len());
1657 for sk in &model.header.skills {
1658 let mut ffn_layers = std::collections::BTreeSet::new();
1659 let mut non_ffn = false;
1660 let prefix = format!("skill.{}.", sk.id);
1661 for t in model.skill_tensors(&sk.id) {
1662 let rel = &t.name[prefix.len()..]; let toks: Vec<&str> = rel.split('.').collect();
1664 if toks.len() >= 5 && toks[0] == "model" && toks[1] == "layers" && toks[3] == "mlp"
1665 {
1666 if let Ok(li) = toks[2].parse::<usize>() {
1667 ffn_layers.insert(li);
1668 continue;
1669 }
1670 }
1671 non_ffn = true; }
1673 if non_ffn {
1674 tracing::warn!(
1675 "skill '{}' replaces non-FFN tensors — excluded from dynamic \
1676 routing (static overlay still works)",
1677 sk.id
1678 );
1679 per_skill.push(None);
1680 } else {
1681 per_skill.push(Some(ffn_layers.into_iter().collect::<Vec<_>>()));
1682 }
1683 }
1684 self.dyn_skill_layers = per_skill;
1685 }
1686
1687 pub fn set_active_skill(&mut self, idx: Option<usize>) -> Result<(), CmfError> {
1694 self.kv_cache.clear();
1696 self.kv_history.clear();
1697 if self.dyn_active == idx {
1698 return Ok(());
1699 }
1700 let model = self.model.clone().ok_or_else(|| {
1701 CmfError::Parse("dynamic routing needs a model-backed pipeline".into())
1702 })?;
1703 let mut union: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
1704 if let Some(old) = self.dyn_active {
1705 if let Some(Some(ls)) = self.dyn_skill_layers.get(old) {
1706 union.extend(ls.iter().copied());
1707 }
1708 }
1709 let new_id: Option<String> = match idx {
1710 Some(n) => match self.dyn_skill_layers.get(n) {
1711 Some(Some(ls)) => {
1712 union.extend(ls.iter().copied());
1713 Some(model.header.skills[n].id.clone())
1714 }
1715 _ => {
1716 return Err(CmfError::Parse(format!(
1717 "skill index {n} not dynamic-eligible"
1718 )));
1719 }
1720 },
1721 None => None,
1722 };
1723 let ov = match &new_id {
1724 Some(s) => Overlay::One(s),
1725 None => Overlay::None,
1726 };
1727 let arch = model.arch();
1728 for li in union {
1729 self.weights.layers[li].ffn =
1730 build_layer_ffn(&model, arch, li, self.dyn_force_f32, &ov)?;
1731 }
1732 self.dyn_active = idx;
1733 Ok(())
1734 }
1735}