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 if model.tensor(&router_name).is_none() {
166 return Ok(FfnKind::Dense(load_dense(&format!("{prefix}mlp."))?));
167 }
168 let cfg = arch.moe.as_ref().ok_or_else(|| {
169 CmfError::Parse(format!(
170 "{router_name} present but header has no arch.moe block"
171 ))
172 })?;
173 let mut experts = Vec::new();
178 for e in 0..cfg.num_experts {
179 if model
180 .tensor(&format!("{prefix}mlp.experts.{e}.gate_proj.weight"))
181 .is_none()
182 {
183 break;
184 }
185 experts.push(load_dense(&format!("{prefix}mlp.experts.{e}."))?);
186 }
187 if experts.is_empty() {
188 return Err(CmfError::Parse(format!(
189 "{prefix}: router present but no expert tensors"
190 )));
191 }
192 let shared = if model
193 .tensor(&format!("{prefix}mlp.shared_expert.gate_proj.weight"))
194 .is_some()
195 {
196 let gate_name = format!("{prefix}mlp.shared_expert_gate.weight");
197 Some((
198 load_dense(&format!("{prefix}mlp.shared_expert."))?,
199 if model.tensor(&gate_name).is_some() {
200 Some(load_matrix(model, &gate_name, force_f32, ov)?)
201 } else {
202 None
203 },
204 ))
205 } else {
206 None
207 };
208 let bias_name = format!("{prefix}mlp.expert_bias");
211 let expert_bias = if model.tensor(&bias_name).is_some() {
212 Some(load_f32(model, &bias_name, ov).map_err(CmfError::Parse)?)
213 } else {
214 None
215 };
216 let top_k = std::env::var("CMF_MOE_TOPK")
222 .ok()
223 .and_then(|v| v.parse::<usize>().ok())
224 .filter(|&k| k >= 1 && k <= cfg.top_k)
225 .inspect(|k| tracing::info!("MoE top_k override: {} (header {})", k, cfg.top_k))
226 .unwrap_or(cfg.top_k);
227 let route_tau = std::env::var("CMF_MOE_TAU")
229 .ok()
230 .and_then(|v| v.parse::<f32>().ok())
231 .filter(|&t| t > 0.0 && t < 1.0)
232 .inspect(|t| tracing::info!("MoE adaptive routing: tau {t}"));
233 let mask = moe_task_mask(&prefix, experts.len());
234 let router = load_matrix(model, &router_name, force_f32, ov)?;
235 if router.rows() != experts.len() {
236 return Err(CmfError::Parse(format!(
237 "{router_name}: {} rows != {} experts",
238 router.rows(),
239 experts.len()
240 )));
241 }
242 let top_k = top_k.min(experts.len());
243 let pes_name = format!("{prefix}mlp.per_expert_scale");
247 let per_expert_scale = if model.tensor(&pes_name).is_some() {
248 Some(load_f32(model, &pes_name, ov).map_err(CmfError::Parse)?)
249 } else {
250 None
251 };
252 let router_input_norm = per_expert_scale.is_some();
253 let moe = MoeFfn {
254 router,
255 experts,
256 top_k,
257 route_tau,
258 norm_topk_prob: cfg.norm_topk_prob,
259 router_sigmoid: cfg.router_sigmoid,
260 expert_bias,
261 routed_scaling: cfg.routed_scaling_factor.unwrap_or(1.0),
262 shared,
263 stats: std::cell::RefCell::new(Vec::new()),
264 act_sq: std::cell::RefCell::new(Vec::new()),
265 act_rows: std::cell::RefCell::new(Vec::new()),
266 mask,
267 per_expert_scale,
268 router_input_norm,
269 };
270 if model
273 .tensor(&format!("{prefix}mlp.gate_proj.weight"))
274 .is_some()
275 {
276 let norm = |suffix: &str| -> Result<Vec<f32>, CmfError> {
277 load_f32(model, &format!("{prefix}{suffix}.weight"), ov).map_err(CmfError::Parse)
278 };
279 return Ok(FfnKind::DenseMoe(Box::new(crate::pipeline::DenseMoeFfn {
280 dense: load_dense(&format!("{prefix}mlp."))?,
281 moe,
282 post_norm_1: norm("post_feedforward_layernorm_1")?,
283 pre_norm_2: norm("pre_feedforward_layernorm_2")?,
284 post_norm_2: norm("post_feedforward_layernorm_2")?,
285 })));
286 }
287 Ok(FfnKind::Moe(moe))
288}
289
290pub(crate) fn moe_task_mask(prefix: &str, ne: usize) -> Option<Vec<bool>> {
298 use std::sync::OnceLock;
299 static CFG: OnceLock<Option<(std::collections::HashMap<usize, Vec<u64>>, f64)>> =
300 OnceLock::new();
301 let cfg = CFG.get_or_init(|| {
302 let path = std::env::var("CMF_MOE_MASK").ok()?;
303 let cover = std::env::var("CMF_MOE_MASK_COVER")
304 .ok()
305 .and_then(|v| v.parse::<f64>().ok())
306 .filter(|&c| c > 0.0 && c <= 1.0)
307 .unwrap_or(0.9);
308 let text = std::fs::read_to_string(&path)
309 .map_err(|e| tracing::warn!("CMF_MOE_MASK: cannot read {path}: {e}"))
310 .ok()?;
311 let map: std::collections::HashMap<String, Vec<u64>> = serde_json::from_str(&text)
312 .map_err(|e| tracing::warn!("CMF_MOE_MASK: bad JSON in {path}: {e}"))
313 .ok()?;
314 tracing::info!("MoE task mask: {path}, cover {cover}");
315 Some((
316 map.into_iter()
317 .filter_map(|(k, v)| Some((k.parse::<usize>().ok()?, v)))
318 .collect(),
319 cover,
320 ))
321 });
322 let (stats, cover) = cfg.as_ref()?;
323 let li: usize = prefix
325 .split("layers.")
326 .nth(1)?
327 .split('.')
328 .next()?
329 .parse()
330 .ok()?;
331 let counts = stats.get(&li)?;
332 if counts.len() != ne {
333 tracing::warn!(
334 "CMF_MOE_MASK: layer {li} has {} counts, model has {ne} experts — skipped",
335 counts.len()
336 );
337 return None;
338 }
339 let total: u64 = counts.iter().sum();
340 if total == 0 {
341 return None;
342 }
343 let mut order: Vec<usize> = (0..ne).collect();
344 order.sort_unstable_by_key(|&e| std::cmp::Reverse(counts[e]));
345 let mut mask = vec![false; ne];
346 let mut acc = 0u64;
347 let mut kept = 0usize;
348 for &e in &order {
349 mask[e] = true;
350 acc += counts[e];
351 kept += 1;
352 if (acc as f64) >= cover * (total as f64) {
353 break;
354 }
355 }
356 tracing::info!(
357 "MoE task mask L{li}: {kept}/{ne} experts for {:.0}% mass",
358 cover * 100.0
359 );
360 Some(mask)
361}
362
363fn load_matrix(
364 model: &Arc<CmfModel>,
365 name: &str,
366 force_f32: bool,
367 ov: &Overlay,
368) -> Result<QTensor, CmfError> {
369 if ov.blend_touches(model, name) {
373 if let Overlay::Blend(list) = ov {
374 let entry = model
375 .tensor(name)
376 .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
377 let data =
378 blend_f32(model, name, list).map_err(|e| CmfError::Parse(format!("blend: {e}")))?;
379 return Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]));
380 }
381 }
382 let skill = match ov {
383 Overlay::One(s) => Some(*s),
384 _ => None,
385 };
386 let name: &str = &match skill {
389 Some(sid) if model.tensor(&format!("skill.{sid}.{name}")).is_some() => {
390 format!("skill.{sid}.{name}")
391 }
392 _ => name.to_string(),
393 };
394 let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
395 if force_f32 {
396 let entry = model
397 .tensor(name)
398 .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
399 if entry.shape.len() != 2 {
400 return Err(err(format!("'{name}' is not 2-D")));
401 }
402 let data = load_f32(model, name, &Overlay::None).map_err(err)?;
403 Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]))
404 } else {
405 QTensor::from_model(model, name).map_err(err)
406 }
407}
408
409impl Pipeline {
410 pub fn from_model(
412 model: &Arc<CmfModel>,
413 sampler_config: SamplerConfig,
414 ) -> Result<Self, CmfError> {
415 Self::from_model_with_skill(model, sampler_config, None)
416 }
417
418 pub fn from_model_with_skill(
424 model: &Arc<CmfModel>,
425 sampler_config: SamplerConfig,
426 skill: Option<&str>,
427 ) -> Result<Self, CmfError> {
428 match skill {
429 Some(s) => Self::from_model_with_overlay(model, sampler_config, &Overlay::One(s)),
430 None => Self::from_model_with_overlay(model, sampler_config, &Overlay::None),
431 }
432 }
433
434 pub fn from_model_with_blend(
437 model: &Arc<CmfModel>,
438 sampler_config: SamplerConfig,
439 blend: &[(String, f32)],
440 ) -> Result<Self, CmfError> {
441 Self::from_model_with_overlay(model, sampler_config, &Overlay::Blend(blend))
442 }
443
444 fn skill_file_guard(model: &CmfModel) -> Result<(), CmfError> {
445 if model.required_features & cortiq_core::format::features::SKILL_FILE != 0 {
448 return Err(CmfError::Parse(
449 "this file is a standalone SKILL, not a runnable model — attach it: \
450 cortiq skill apply <base.cmf> <this file> -o specialist.cmf"
451 .into(),
452 ));
453 }
454 Ok(())
455 }
456
457 fn from_model_with_overlay(
458 model: &Arc<CmfModel>,
459 sampler_config: SamplerConfig,
460 ov: &Overlay,
461 ) -> Result<Self, CmfError> {
462 Self::skill_file_guard(model)?;
463 let skill = match ov {
464 Overlay::One(s) => Some(*s),
465 _ => None,
466 };
467 if let Some(sid) = skill {
468 let known = model.header.skills.iter().any(|s| s.id == sid)
469 || model.skill_tensors(sid).next().is_some();
470 if !known {
471 return Err(CmfError::Parse(format!(
472 "skill '{sid}' not in this container (header.skills: {:?})",
473 model
474 .header
475 .skills
476 .iter()
477 .map(|s| &s.id)
478 .collect::<Vec<_>>()
479 )));
480 }
481 tracing::info!(
482 "skill '{sid}': {} replacement tensors overlaid",
483 model.skill_tensors(sid).count()
484 );
485 }
486 let arch = model.arch().clone();
487 let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
488 if let Some(heads) = &arch.attention_heads_per_layer {
489 if heads.len() != arch.num_layers {
490 return Err(CmfError::Parse(format!(
491 "arch.attention_heads_per_layer has {} entries, expected {}",
492 heads.len(),
493 arch.num_layers
494 )));
495 }
496 if let Some((li, &nh)) = heads
497 .iter()
498 .enumerate()
499 .find(|(_, nh)| **nh == 0 || **nh % arch.num_kv_heads != 0)
500 {
501 return Err(CmfError::Parse(format!(
502 "layer {li} has {nh} Q heads, which must be nonzero and divisible by {} KV heads",
503 arch.num_kv_heads
504 )));
505 }
506 }
507 if arch
508 .layer_types
509 .iter()
510 .any(|t| matches!(t, LayerType::SlidingAttention))
511 && arch.sliding_window.is_none()
512 {
513 return Err(CmfError::Parse(
514 "model has SlidingAttention layers but no arch.sliding_window".into(),
515 ));
516 }
517
518 let heads_masked = model.masks.masks.iter().any(|m| {
528 m.head_masks.iter().any(|row| {
529 let mut bits = 0usize;
530 for &b in row.iter() {
531 bits += b.count_ones() as usize;
532 }
533 !row.is_empty() && bits < arch.num_attention_heads
534 })
535 });
536 let force_f32 = heads_masked; let mut tokenizer = if let Some(vocab_bytes) = &model.vocab {
540 Tokenizer::from_bytes(vocab_bytes)
541 .map_err(|e| CmfError::Parse(format!("embedded tokenizer: {e}")))?
542 } else {
543 let sidecar = model.path.with_file_name("tokenizer.json");
544 if sidecar.exists() {
545 Tokenizer::from_file(&sidecar)
546 .map_err(|e| CmfError::Parse(format!("sidecar tokenizer: {e}")))?
547 } else {
548 tracing::warn!("no tokenizer in file or sidecar — using byte-level fallback");
549 Tokenizer::byte_level()
550 }
551 };
552 if let Some(tc) = &model.header.tokenizer_config {
554 tokenizer.chat_template = tc.chat_template.clone();
555 tokenizer.extra_eos.extend(tc.eos_token_ids.iter().copied());
556 if tokenizer.bos_token_id.is_none() {
557 tokenizer.bos_token_id = tc.bos_token_id;
558 }
559 tracing::info!(
560 "chat bundle: template {} chars, {} stop ids",
561 tc.chat_template.as_deref().map(str::len).unwrap_or(0),
562 tc.eos_token_ids.len()
563 );
564 }
565 if arch.arch_name.to_lowercase().contains("gemma") && tokenizer.bos_token_id.is_some() {
569 tokenizer.add_bos = true;
570 }
571
572 let embed_tokens = load_matrix(model, "model.embed_tokens.weight", false, ov)?;
574 let final_norm = load_f32(model, "model.norm.weight", ov).map_err(err)?;
575 let lm_head = if model.tensor("lm_head.weight").is_some() {
576 load_matrix(model, "lm_head.weight", false, ov)?
577 } else if arch.tie_word_embeddings {
578 load_matrix(model, "model.embed_tokens.weight", false, ov)?
580 } else {
581 return Err(CmfError::MissingTensor(
582 "lm_head.weight (and tie_word_embeddings is false)".into(),
583 ));
584 };
585
586 let has_linear = arch
588 .layer_types
589 .iter()
590 .any(|t| matches!(t, LayerType::LinearAttention));
591 let mut vmf_cfg = None;
592 let mut gdn_cfg = None;
593 if has_linear {
594 let lc = arch.linear_core.as_ref().ok_or_else(|| {
595 CmfError::Parse(
596 "model has LinearAttention layers but no arch.linear_core — \
597 reconvert with the current converter"
598 .into(),
599 )
600 })?;
601 let need = |v: Option<usize>, name: &str| {
602 v.ok_or_else(|| CmfError::Parse(format!("linear core needs arch.{name}")))
603 };
604 match lc.kind.as_str() {
605 "vmf_phase" => {
606 vmf_cfg = Some(VmfPhaseCfg {
607 num_heads: lc.num_heads,
608 nphase: need(lc.nphase, "linear_core.nphase")?,
609 value_head_dim: lc.value_head_dim,
610 hidden_size: arch.hidden_size,
611 phase_mass: std::env::var("CMF_PHASE_MASS")
614 .ok()
615 .and_then(|v| v.parse().ok())
616 .unwrap_or(0.0),
617 });
618 }
619 "gated_delta_net" => {
620 gdn_cfg = Some(GdnCfg {
621 num_v_heads: lc.num_heads,
622 num_k_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
623 key_head_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
624 value_head_dim: lc.value_head_dim,
625 conv_kernel: need(arch.linear_conv_kernel_dim, "linear_conv_kernel_dim")?,
626 hidden_size: arch.hidden_size,
627 rms_eps: arch.rms_norm_eps,
628 });
629 }
630 other => {
631 return Err(CmfError::Parse(format!(
632 "unknown linear core '{other}' (this runtime executes: \
633 gated_delta_net, vmf_phase)"
634 )));
635 }
636 }
637 }
638
639 let has_kda = arch.layer_types.iter().any(|t| matches!(t, LayerType::Kda));
641 let kda_cfg = if has_kda {
642 let need = |v: Option<usize>, name: &str| {
643 v.ok_or_else(|| CmfError::Parse(format!("KDA core needs arch.{name}")))
644 };
645 Some(crate::linear_core::KdaCfg {
646 num_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
647 head_k_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
648 head_v_dim: need(arch.linear_value_head_dim, "linear_value_head_dim")?,
649 conv_kernel: need(arch.linear_conv_kernel_dim, "linear_conv_kernel_dim")?,
650 hidden_size: arch.hidden_size,
651 rms_eps: arch.rms_norm_eps,
652 })
653 } else {
654 None
655 };
656
657 let has_short_conv = arch
659 .layer_types
660 .iter()
661 .any(|t| matches!(t, LayerType::ShortConv));
662 let short_conv_cfg = if has_short_conv {
663 Some(ShortConvCfg {
664 hidden_size: arch.hidden_size,
665 kernel: arch.linear_conv_kernel_dim.ok_or_else(|| {
666 CmfError::Parse(
667 "model has ShortConv layers but no arch.linear_conv_kernel_dim — \
668 reconvert with the current converter"
669 .into(),
670 )
671 })?,
672 })
673 } else {
674 None
675 };
676
677 let load_full_attn = |prefix: &str, layer: Option<usize>| -> Result<AttnKind, CmfError> {
679 let t = |suffix: &str| load_matrix(model, &format!("{prefix}{suffix}"), force_f32, ov);
680 let n = |suffix: &str| -> Option<Vec<f32>> {
681 model
682 .tensor(&format!("{prefix}{suffix}"))
683 .and_then(|_| load_f32(model, &format!("{prefix}{suffix}"), ov).ok())
684 };
685 if let Some(mla) = arch.mla.as_ref() {
687 let (q_proj, q_a, q_a_norm) = if mla.q_lora_rank.is_some() {
689 (
690 t("self_attn.q_b_proj.weight")?,
691 Some(t("self_attn.q_a_proj.weight")?),
692 Some(n("self_attn.q_a_layernorm.weight").ok_or_else(|| {
693 CmfError::Parse(format!("{prefix}: MLA needs q_a_layernorm"))
694 })?),
695 )
696 } else {
697 (t("self_attn.q_proj.weight")?, None, None)
698 };
699 let hd = mla.qk_rope_head_dim + mla.qk_nope_head_dim;
700 let nh = q_proj.rows() / hd;
701 let mut scale = 1.0 / (hd as f32).sqrt();
704 if let Some(y) = arch.yarn.as_ref() {
705 if let Some(m) = y.mscale_all_dim.filter(|&m| m > 0.0) {
706 let ms = 0.1 * m * y.factor.ln() + 1.0;
707 scale *= ms * ms;
708 }
709 }
710 return Ok(AttnKind::Mla(Box::new(crate::pipeline::MlaWeights {
711 q_proj,
712 q_a,
713 q_a_norm,
714 kv_a: t("self_attn.kv_a_proj_with_mqa.weight")?,
715 kv_a_norm: n("self_attn.kv_a_layernorm.weight").ok_or_else(|| {
716 CmfError::Parse(format!("{prefix}: MLA needs kv_a_layernorm"))
717 })?,
718 kv_b: t("self_attn.kv_b_proj.weight")?,
719 o_proj: t("self_attn.o_proj.weight")?,
720 nh,
721 qk_rope: mla.qk_rope_head_dim,
722 qk_nope: mla.qk_nope_head_dim,
723 v_dim: mla.v_head_dim,
724 lora: mla.kv_lora_rank,
725 scale,
726 nope: mla.nope,
727 })));
728 }
729 let wq = t("self_attn.q_proj.weight")?;
730 let nh = layer
731 .and_then(|li| {
732 arch.attention_heads_per_layer
733 .as_ref()
734 .and_then(|v| v.get(li).copied())
735 })
736 .unwrap_or(arch.num_attention_heads);
737 let output_gate = arch.global_head_dim.is_none() && wq.rows() == 2 * nh * arch.head_dim;
741 let is_global_layer = arch.global_head_dim.is_some()
744 && layer.is_some_and(|li| {
745 arch.sliding_window_pattern
746 .is_some_and(|p| p > 0 && (li + 1) % p == 0)
747 });
748 let expect = if is_global_layer {
749 nh * arch.global_head_dim.unwrap_or(arch.head_dim)
750 } else {
751 nh * arch.head_dim
752 };
753 if !output_gate && wq.rows() != expect {
754 return Err(CmfError::Parse(format!(
755 "{prefix}self_attn.q_proj.weight rows={} != heads({nh}) * head_dim({})",
756 wq.rows(),
757 expect / nh.max(1)
758 )));
759 }
760 let gate_name = format!("{prefix}self_attn.g_proj.weight");
761 let softplus_gate = if model.tensor(&gate_name).is_some() {
762 let gate = load_matrix(model, &gate_name, force_f32, ov)?;
763 if gate.cols() != arch.hidden_size {
764 return Err(CmfError::Parse(format!(
765 "{gate_name} cols={} != hidden_size ({})",
766 gate.cols(),
767 arch.hidden_size
768 )));
769 }
770 let per_head = if gate.rows() == nh {
771 true
772 } else if gate.rows() == nh * arch.head_dim {
773 false
774 } else {
775 return Err(CmfError::Parse(format!(
776 "{gate_name} rows={} must equal heads ({nh}) or heads*head_dim ({})",
777 gate.rows(),
778 nh * arch.head_dim
779 )));
780 };
781 Some((gate, per_head))
782 } else {
783 None
784 };
785 let bias = match (
787 n("self_attn.q_proj.bias"),
788 n("self_attn.k_proj.bias"),
789 n("self_attn.v_proj.bias"),
790 ) {
791 (Some(a), Some(b), Some(c)) => Some((a, b, c)),
792 _ => None,
793 };
794 Ok(AttnKind::Full {
795 wq,
796 wk: t("self_attn.k_proj.weight")?,
797 wv: t("self_attn.v_proj.weight")?,
798 wo: t("self_attn.o_proj.weight")?,
799 q_norm: n("self_attn.q_norm.weight"),
800 k_norm: n("self_attn.k_norm.weight"),
801 output_gate,
802 softplus_gate,
803 bias,
804 })
805 };
806
807 let load_linear_attn = |prefix: &str| -> Result<AttnKind, CmfError> {
808 if gdn_cfg.is_some() {
809 let t = |suffix: &str| {
811 load_matrix(
812 model,
813 &format!("{prefix}linear_attn.{suffix}"),
814 force_f32,
815 ov,
816 )
817 };
818 let f = |suffix: &str| {
819 load_f32(model, &format!("{prefix}linear_attn.{suffix}"), ov).map_err(err)
820 };
821 return Ok(AttnKind::LinearGdn(GdnWeights {
822 in_proj_qkv: t("in_proj_qkv.weight")?,
823 in_proj_z: t("in_proj_z.weight")?,
824 in_proj_a: t("in_proj_a.weight")?,
825 in_proj_b: t("in_proj_b.weight")?,
826 conv1d: f("conv1d.weight")?,
827 a_log: f("A_log")?,
828 dt_bias: f("dt_bias")?,
829 norm: f("norm.weight")?,
830 out_proj: t("out_proj.weight")?,
831 }));
832 }
833 let t = |suffix: &str| {
834 load_matrix(model, &format!("{prefix}vmf_attn.{suffix}"), force_f32, ov)
835 };
836 let a_log = load_f32(model, &format!("{prefix}vmf_attn.A_log"), ov).map_err(err)?;
837 let k_gate = if model
841 .tensor(&format!("{prefix}vmf_attn.k_gate.weight"))
842 .is_some()
843 {
844 Some((
845 t("k_gate.weight")?,
846 load_f32(model, &format!("{prefix}vmf_attn.k_gate.bias"), ov).map_err(err)?,
847 ))
848 } else {
849 None
850 };
851 Ok(AttnKind::Linear(VmfPhaseWeights {
852 thq: t("thq.weight")?,
853 thk: t("thk.weight")?,
854 v_proj: t("v_proj.weight")?,
855 out_proj: t("out_proj.weight")?,
856 decay: a_log.iter().map(|&a| (-(a as f64).exp()).exp()).collect(),
857 k_gate,
858 }))
859 };
860
861 let load_short_conv = |prefix: &str| -> Result<AttnKind, CmfError> {
865 let t = |suffix: &str| {
866 load_matrix(
867 model,
868 &format!("{prefix}short_conv.{suffix}"),
869 force_f32,
870 ov,
871 )
872 };
873 Ok(AttnKind::ShortConv(ShortConvWeights {
874 in_proj: t("in_proj.weight")?,
875 conv: load_f32(model, &format!("{prefix}short_conv.conv.weight"), ov)
876 .map_err(err)?,
877 out_proj: t("out_proj.weight")?,
878 }))
879 };
880
881 let load_kda = |prefix: &str| -> Result<AttnKind, CmfError> {
885 let t = |suffix: &str| {
886 load_matrix(model, &format!("{prefix}kda_attn.{suffix}"), force_f32, ov)
887 };
888 let f = |suffix: &str| {
889 load_f32(model, &format!("{prefix}kda_attn.{suffix}"), ov).map_err(err)
890 };
891 let gate = if model
892 .tensor(&format!("{prefix}kda_attn.g_proj.weight"))
893 .is_some()
894 {
895 crate::linear_core::KdaOutGate::Full(t("g_proj.weight")?)
896 } else {
897 crate::linear_core::KdaOutGate::LowRank(
898 t("g_a_proj.weight")?,
899 t("g_b_proj.weight")?,
900 )
901 };
902 Ok(AttnKind::Kda(Box::new(crate::linear_core::KdaWeights {
903 q_proj: t("q_proj.weight")?,
904 k_proj: t("k_proj.weight")?,
905 v_proj: t("v_proj.weight")?,
906 conv_q: f("q_conv1d.weight")?,
907 conv_k: f("k_conv1d.weight")?,
908 conv_v: f("v_conv1d.weight")?,
909 f_a: t("f_a_proj.weight")?,
910 f_b: t("f_b_proj.weight")?,
911 dt_bias: f("dt_bias")?,
912 a_log: f("A_log")?,
913 b_proj: t("b_proj.weight")?,
914 gate,
915 o_norm: f("o_norm.weight")?,
916 o_proj: t("o_proj.weight")?,
917 gate_lower_bound: arch.kda_gate_lower_bound.map(|v| v as f32),
918 })))
919 };
920
921 fn anyhow_like(ok: bool) -> Result<(), ()> {
922 if ok { Ok(()) } else { Err(()) }
923 }
924 let mut layers = Vec::with_capacity(arch.num_layers);
925 let is_g3n = arch.g3n.is_some();
926 let owns_its_layers = is_g3n || arch.arch_name == "deepseek_v4";
931 for li in 0..(if owns_its_layers { 0 } else { arch.num_layers }) {
932 let prefix = format!("model.layers.{li}.");
933 let attn = match arch.layer_types.get(li) {
934 Some(LayerType::LinearAttention) => load_linear_attn(&prefix)?,
935 Some(LayerType::Kda) => load_kda(&prefix)?,
936 Some(LayerType::ShortConv) => load_short_conv(&prefix)?,
937 _ => load_full_attn(&prefix, Some(li))?,
938 };
939 let pre_ffn = format!("{prefix}pre_feedforward_layernorm.weight");
943 let sandwich = model.tensor(&pre_ffn).is_some();
944 layers.push(LayerWeights {
945 input_norm: load_f32(model, &format!("{prefix}input_layernorm.weight"), ov)
946 .map_err(err)?,
947 post_norm: if sandwich {
948 load_f32(model, &pre_ffn, ov).map_err(err)?
949 } else {
950 load_f32(
951 model,
952 &format!("{prefix}post_attention_layernorm.weight"),
953 ov,
954 )
955 .map_err(err)?
956 },
957 attn_out_norm: if sandwich {
958 Some(
959 load_f32(
960 model,
961 &format!("{prefix}post_attention_layernorm.weight"),
962 ov,
963 )
964 .map_err(err)?,
965 )
966 } else {
967 None
968 },
969 ffn_out_norm: if sandwich {
970 Some(
971 load_f32(
972 model,
973 &format!("{prefix}post_feedforward_layernorm.weight"),
974 ov,
975 )
976 .map_err(err)?,
977 )
978 } else {
979 None
980 },
981 layer_scale: model
983 .tensor(&format!("{prefix}layer_scalar"))
984 .and_then(|_| {
985 load_f32(model, &format!("{prefix}layer_scalar"), ov)
986 .ok()
987 .and_then(|v| v.first().copied())
988 }),
989 ffn: build_layer_ffn(model, &arch, li, false, ov)?,
991 attn,
992 });
993 }
994
995 let mtp_present = model
1004 .tensor("model.mtp.layers.0.self_attn.q_proj.weight")
1005 .is_some()
1006 || model.tensor("model.mtp.eh_proj.weight").is_some();
1007 let dsv4_mtp = model.tensor("model.mtp.0.main_proj.weight").is_some();
1012 if arch.mtp.is_some() && !mtp_present && !dsv4_mtp {
1013 tracing::info!(
1014 "header declares an MTP head but the file carries none — \
1015 loading without it"
1016 );
1017 }
1018 let mtp = if let Some(cfg) = arch.mtp.as_ref().filter(|_| mtp_present) {
1019 if cfg.num_layers != 1 {
1020 return Err(CmfError::Parse(format!(
1021 "MTP with {} blocks not supported yet (only 1)",
1022 cfg.num_layers
1023 )));
1024 }
1025 let p = "model.mtp.";
1026 let attn = load_full_attn("model.mtp.layers.0.", None)?;
1027 Some(MtpModule {
1028 enorm: load_f32(model, &format!("{p}enorm.weight"), ov).map_err(err)?,
1029 hnorm: load_f32(model, &format!("{p}hnorm.weight"), ov).map_err(err)?,
1030 eh_proj: load_matrix(model, &format!("{p}eh_proj.weight"), false, ov)?,
1031 layer: LayerWeights {
1032 attn_out_norm: None,
1033 ffn_out_norm: None,
1034 layer_scale: None,
1035 input_norm: load_f32(model, &format!("{p}layers.0.input_layernorm.weight"), ov)
1036 .map_err(err)?,
1037 post_norm: load_f32(
1038 model,
1039 &format!("{p}layers.0.post_attention_layernorm.weight"),
1040 ov,
1041 )
1042 .map_err(err)?,
1043 ffn: build_ffn_at(model, &arch, &format!("{p}layers.0."), false, ov)?,
1047 attn,
1048 },
1049 final_norm: load_f32(model, &format!("{p}norm.weight"), ov).map_err(err)?,
1050 kv: LayerKvCache::new(arch.num_kv_heads, arch.head_dim),
1051 })
1052 } else {
1053 None
1054 };
1055
1056 tracing::info!(
1057 "Pipeline loaded: {} | {}L ({} linear) | {:.2}B params | storage: {} | MTP: {}",
1058 arch.arch_name,
1059 arch.num_layers,
1060 arch.layer_types
1061 .iter()
1062 .filter(|t| matches!(t, LayerType::LinearAttention))
1063 .count(),
1064 model.total_param_count() as f64 / 1e9,
1065 if force_f32 {
1066 "f32 (masked)"
1067 } else {
1068 "quantized mmap"
1069 },
1070 if mtp.is_some() { "yes" } else { "no" }
1071 );
1072
1073 let cap = std::env::var("CMF_MAX_SEQ")
1076 .ok()
1077 .and_then(|v| v.parse::<usize>().ok())
1078 .unwrap_or(8192);
1079 let max_seq_len = arch.max_position_embeddings.min(cap);
1080
1081 let total_layers = arch.num_layers * arch.num_loops;
1083
1084 let mut pipeline = Pipeline::new(
1085 tokenizer,
1086 PipelineWeights {
1087 embed_tokens,
1088 layers,
1089 lm_head,
1090 final_norm,
1091 },
1092 arch.hidden_size,
1093 arch.intermediate_size,
1094 arch.num_attention_heads,
1095 arch.num_kv_heads,
1096 arch.head_dim,
1097 total_layers,
1098 arch.num_layers, arch.loop_final_norm,
1100 arch.vocab_size,
1101 arch.rms_norm_eps,
1102 arch.rope_theta as f32,
1103 arch.norm_style,
1104 max_seq_len,
1105 sampler_config,
1106 );
1107 let rotary = ((arch.head_dim as f32 * arch.partial_rotary_factor) as usize).max(2);
1108 pipeline.set_rotary(rotary, arch.rope_theta as f32);
1109 pipeline.attention_heads_per_layer = arch.attention_heads_per_layer.clone();
1110 if let Some(yarn) = &arch.yarn {
1111 pipeline.inv_freq = std::sync::Arc::new(crate::attention::yarn_inv_freq(
1112 rotary,
1113 arch.rope_theta as f32,
1114 yarn.factor,
1115 yarn.original_max_position_embeddings,
1116 yarn.beta_fast,
1117 yarn.beta_slow,
1118 ));
1119 pipeline.rope_scale = yarn.attention_factor;
1120 }
1121 pipeline.embed_multiplier = arch.embed_multiplier;
1125 pipeline.logit_multiplier = arch.logit_multiplier;
1126 if let Some(qpas) = arch.query_pre_attn_scalar {
1127 pipeline.attn_scale = 1.0 / (qpas as f32).sqrt();
1128 }
1129 if let (Some(w), Some(p)) = (arch.sliding_window, arch.sliding_window_pattern) {
1130 pipeline.swa = Some((w, p));
1131 if let Some(base) = arch.rope_local_base_freq {
1132 pipeline.inv_freq_local = Some(std::sync::Arc::new(
1133 crate::attention::rope_inv_freq(rotary, base as f32),
1134 ));
1135 }
1136 }
1137 let explicit_sliding: Vec<bool> = arch
1138 .layer_types
1139 .iter()
1140 .map(|t| matches!(t, cortiq_core::LayerType::SlidingAttention))
1141 .collect();
1142 if explicit_sliding.iter().any(|&v| v) {
1143 pipeline.sliding_layers = Some(explicit_sliding);
1144 if let Some(w) = arch.sliding_window {
1145 pipeline.swa = Some((w, usize::MAX));
1146 }
1147 let local_rotary = ((arch.head_dim as f32
1148 * arch
1149 .local_partial_rotary_factor
1150 .unwrap_or(arch.partial_rotary_factor))
1151 as usize)
1152 .max(2);
1153 pipeline.rotary_dim_local = Some(local_rotary);
1154 if let Some(base) = arch.rope_local_base_freq {
1155 pipeline.inv_freq_local = Some(std::sync::Arc::new(
1156 crate::attention::rope_inv_freq(local_rotary, base as f32),
1157 ));
1158 }
1159 }
1160 if let (Some(ghd), Some(gkv)) = (arch.global_head_dim, arch.num_global_kv_heads) {
1164 pipeline.global_attn = Some((ghd, gkv));
1165 let prf = arch.global_partial_rotary_factor.unwrap_or(1.0);
1166 let half = ghd / 2;
1167 let ra = (((prf * ghd as f32) as usize) / 2).min(half);
1168 let mut f = vec![0.0f32; half];
1169 for (i, slot) in f.iter_mut().enumerate().take(ra) {
1170 *slot = 1.0 / (arch.rope_theta as f32).powf(2.0 * i as f32 / ghd as f32);
1171 }
1172 pipeline.inv_freq_global = Some(std::sync::Arc::new(f));
1173 let global_at = |li: usize| -> bool {
1178 match &pipeline.sliding_layers {
1179 Some(map) => !map.get(li).copied().unwrap_or(false),
1180 None => pipeline
1181 .swa
1182 .map(|(_, p)| p > 0 && p != usize::MAX && (li + 1) % p == 0)
1183 .unwrap_or(false),
1184 }
1185 };
1186 for li in 0..arch.num_layers {
1187 if global_at(li) {
1188 pipeline.kv_cache.layers[li] = crate::kv_cache::LayerKvCache::new(gkv, ghd);
1189 }
1190 }
1191 }
1192 if let Some(mla) = arch.mla.as_ref() {
1195 let hd = mla.qk_rope_head_dim + mla.qk_nope_head_dim;
1196 pipeline.head_dim = hd;
1197 pipeline.num_kv_heads = arch.num_attention_heads;
1198 pipeline.rotary_dim = mla.qk_rope_head_dim;
1199 let half = mla.qk_rope_head_dim / 2;
1200 let mut f = vec![0.0f32; half];
1201 for (i, slot) in f.iter_mut().enumerate() {
1202 *slot = 1.0
1203 / (arch.rope_theta as f32).powf(2.0 * i as f32 / mla.qk_rope_head_dim as f32);
1204 }
1205 pipeline.inv_freq = std::sync::Arc::new(f);
1206 for li in 0..arch.num_layers {
1207 pipeline.kv_cache.layers[li] =
1208 crate::kv_cache::LayerKvCache::new(arch.num_attention_heads, hd);
1209 }
1210 }
1211 if let Some(fac) = &arch.rope_freq_factors {
1215 let mut f = pipeline.inv_freq.as_ref().clone();
1216 for (i, v) in f.iter_mut().enumerate() {
1217 if let Some(&d) = fac.get(i) {
1218 *v /= d as f32;
1219 }
1220 }
1221 pipeline.inv_freq = std::sync::Arc::new(f);
1222 }
1223 pipeline.attn_v_norm = arch.attn_v_norm;
1224 pipeline.final_softcap = arch.final_logit_softcapping.map(|c| c as f32);
1225 pipeline.attn_softcap = arch.attn_logit_softcapping.unwrap_or(0.0) as f32;
1226 pipeline.vmf_cfg = vmf_cfg;
1227 pipeline.gdn_cfg = gdn_cfg;
1228 pipeline.kda_cfg = kda_cfg;
1229 if let Some(gc) = arch.g3n.as_ref() {
1230 use crate::g3n::{G3nAltUp, G3nGlobals, G3nLaurel, G3nLayer};
1231 anyhow_like(gc.altup_num_inputs == crate::g3n::ALTUP_N).map_err(|_| {
1232 CmfError::Parse(format!(
1233 "g3n: altup_num_inputs {} != supported {}",
1234 gc.altup_num_inputs,
1235 crate::g3n::ALTUP_N
1236 ))
1237 })?;
1238 let t = |name: &str| load_matrix(model, name, force_f32, ov);
1239 let f = |name: &str| load_f32(model, name, ov).map_err(err);
1240 let mut altup_proj = Vec::new();
1241 let mut altup_unembed = Vec::new();
1242 for i in 0..crate::g3n::ALTUP_N - 1 {
1243 altup_proj.push(t(&format!("model.altup_projections.{i}.weight"))?);
1244 altup_unembed.push(t(&format!("model.altup_unembed_projections.{i}.weight"))?);
1245 }
1246 let first_shared = arch.num_layers.saturating_sub(gc.num_kv_shared_layers);
1247 let sliding_of = |li: usize| {
1248 matches!(
1249 arch.layer_types.get(li),
1250 Some(cortiq_core::LayerType::SlidingAttention)
1251 )
1252 };
1253 let mut g3n_layers = Vec::with_capacity(arch.num_layers);
1254 for li in 0..arch.num_layers {
1255 let pfx = format!("model.layers.{li}.");
1256 let shared = li >= first_shared && first_shared > 0;
1257 let share_src = if shared {
1258 let want = sliding_of(li);
1259 (0..first_shared).rev().find(|&j| sliding_of(j) == want)
1260 } else {
1261 None
1262 };
1263 g3n_layers.push(G3nLayer {
1264 altup: G3nAltUp {
1265 router_norm: f(&format!("{pfx}altup.router_norm.weight"))?,
1266 modality_router: t(&format!("{pfx}altup.modality_router.weight"))?,
1267 prediction_coefs: t(&format!("{pfx}altup.prediction_coefs.weight"))?,
1268 correction_coefs: t(&format!("{pfx}altup.correction_coefs.weight"))?,
1269 correct_output_scale: f(&format!("{pfx}altup.correct_output_scale"))?,
1270 },
1271 laurel: G3nLaurel {
1272 left: t(&format!("{pfx}laurel.linear_left.weight"))?,
1273 right: t(&format!("{pfx}laurel.linear_right.weight"))?,
1274 post_norm: f(&format!("{pfx}laurel.post_laurel_norm.weight"))?,
1275 },
1276 input_norm: f(&format!("{pfx}input_layernorm.weight"))?,
1277 post_attn_norm: f(&format!("{pfx}post_attention_layernorm.weight"))?,
1278 pre_ffw_norm: f(&format!("{pfx}pre_feedforward_layernorm.weight"))?,
1279 post_ffw_norm: f(&format!("{pfx}post_feedforward_layernorm.weight"))?,
1280 wq: t(&format!("{pfx}self_attn.q_proj.weight"))?,
1281 wk: if shared {
1282 None
1283 } else {
1284 Some(t(&format!("{pfx}self_attn.k_proj.weight"))?)
1285 },
1286 wv: if shared {
1287 None
1288 } else {
1289 Some(t(&format!("{pfx}self_attn.v_proj.weight"))?)
1290 },
1291 wo: t(&format!("{pfx}self_attn.o_proj.weight"))?,
1292 q_norm: f(&format!("{pfx}self_attn.q_norm.weight"))?,
1293 k_norm: if shared {
1294 None
1295 } else {
1296 Some(f(&format!("{pfx}self_attn.k_norm.weight"))?)
1297 },
1298 kv_share_src: share_src,
1299 sliding: sliding_of(li),
1300 gate: t(&format!("{pfx}mlp.gate_proj.weight"))?,
1301 up: t(&format!("{pfx}mlp.up_proj.weight"))?,
1302 down: t(&format!("{pfx}mlp.down_proj.weight"))?,
1303 sparsity: gc.activation_sparsity.get(li).copied().unwrap_or(0.0),
1304 ple_gate: t(&format!("{pfx}per_layer_input_gate.weight"))?,
1305 ple_proj: t(&format!("{pfx}per_layer_projection.weight"))?,
1306 post_ple_norm: f(&format!("{pfx}post_per_layer_input_norm.weight"))?,
1307 });
1308 }
1309 let hd = arch.head_dim;
1310 let globals = G3nGlobals {
1311 altup_proj,
1312 altup_unembed,
1313 ple_embed: t("model.embed_tokens_per_layer.weight")?,
1314 ple_model_proj: t("model.per_layer_model_projection.weight")?,
1315 ple_norm: f("model.per_layer_projection_norm.weight")?,
1316 ple_vocab: gc.ple_vocab,
1317 ple_dim: gc.ple_dim,
1318 num_layers: arch.num_layers,
1319 hidden: arch.hidden_size,
1320 rms_eps: arch.rms_norm_eps,
1321 inv_freq_local: crate::attention::rope_inv_freq(
1322 hd,
1323 arch.rope_local_base_freq.unwrap_or(10_000.0) as f32,
1324 ),
1325 inv_freq_global: crate::attention::rope_inv_freq(hd, arch.rope_theta as f32),
1326 window: arch.sliding_window.unwrap_or(512),
1327 };
1328 pipeline.g3n = Some(Box::new((globals, g3n_layers)));
1329 }
1330 if arch.arch_name == "deepseek_v4" {
1335 let moe = arch
1336 .moe
1337 .as_ref()
1338 .ok_or_else(|| CmfError::Parse("deepseek_v4: no moe config".into()))?;
1339 let cfg = crate::dsv4::Dsv4Cfg {
1340 dim: arch.hidden_size,
1341 n_heads: arch.num_attention_heads,
1342 head_dim: arch.head_dim,
1343 rope_head_dim: if arch.partial_rotary_factor < 1.0 {
1350 (((arch.head_dim as f32 * arch.partial_rotary_factor) as usize) & !1)
1351 .clamp(2, arch.head_dim)
1352 } else {
1353 64.min(arch.head_dim)
1354 },
1355 q_lora_rank: 0,
1359 o_lora_rank: 0,
1360 o_groups: 8,
1367 hc_mult: 4,
1368 hc_sinkhorn_iters: 20,
1369 hc_eps: 1e-6,
1370 norm_eps: arch.rms_norm_eps as f32,
1371 n_routed_experts: moe.num_experts,
1372 top_k: moe.top_k,
1373 moe_inter: moe.moe_intermediate_size,
1374 route_scale: moe.routed_scaling_factor.unwrap_or(1.0) as f32,
1375 swiglu_limit: 10.0,
1381 window: arch.sliding_window.unwrap_or(128),
1382 index_topk: 512,
1383 vocab: arch.vocab_size,
1384 };
1385 let (g, dl) = crate::dsv4::load(model, &cfg, arch.num_layers)
1386 .map_err(|e| CmfError::Parse(format!("deepseek_v4: {e}")))?;
1387 let mut cfg = cfg;
1392 if let Some(l0) = dl.first() {
1393 cfg.q_lora_rank = l0.wq_a.rows();
1394 let attn_width = arch.num_attention_heads * arch.head_dim;
1395 if l0.wo_a.cols() > 0 && attn_width % l0.wo_a.cols() == 0 {
1396 cfg.o_groups = (attn_width / l0.wo_a.cols()).max(1);
1397 }
1398 cfg.o_lora_rank = l0.wo_b.cols() / cfg.o_groups.max(1);
1399 cfg.hc_mult = (l0.hc_attn_fn.len() / l0.hc_attn_base.len().max(1)) / cfg.dim.max(1);
1400 if cfg.hc_mult == 0 {
1401 cfg.hc_mult = 4;
1402 }
1403 }
1404 let (yf, yo, ybf, ybs) = match &arch.yarn {
1417 Some(y) => (
1418 y.factor,
1419 y.original_max_position_embeddings,
1420 y.beta_fast,
1421 y.beta_slow,
1422 ),
1423 None => {
1424 tracing::warn!(
1425 "deepseek_v4: the header carries no YaRN profile — \
1426 falling back to the release's (factor 16, original \
1427 65536, beta 32/1). Re-converting with a build that \
1428 reads rope_scaling.type would make this exact."
1429 );
1430 (16.0, 65536, 32.0, 1.0)
1431 }
1432 };
1433 pipeline.inv_freq = std::sync::Arc::new(crate::attention::yarn_inv_freq(
1434 cfg.rope_head_dim,
1435 arch.rope_theta as f32,
1436 yf,
1437 yo,
1438 ybf,
1439 ybs,
1440 ));
1441 if let Ok(stats) = std::env::var("CMF_MOE_PIN") {
1446 let cover = std::env::var("CMF_MOE_PIN_COVER")
1447 .ok()
1448 .and_then(|v| v.parse::<f64>().ok())
1449 .filter(|&c| c > 0.0 && c <= 1.0)
1450 .unwrap_or(0.95);
1451 let hot = crate::pin::hot_experts(&stats, cover);
1452 let mut names: Vec<String> = Vec::new();
1453 for e in &model.tensors {
1454 let is_expert = e.name.contains(".mlp.experts.");
1455 if !is_expert {
1456 names.push(e.name.clone()); }
1458 }
1459 let mut kept_experts = 0usize;
1460 if let Some(hot) = &hot {
1461 for (li, experts) in hot {
1462 for e in experts {
1463 for w in ["gate_proj", "up_proj", "down_proj"] {
1464 names.push(format!("model.layers.{li}.mlp.experts.{e}.{w}.weight"));
1465 }
1466 kept_experts += 1;
1467 }
1468 }
1469 }
1470 let r = crate::pin::pin_tensors(model, &names);
1471 tracing::info!(
1472 "закреплено {:.1} ГБ ({} тензоров, горячих экспертов {kept_experts}, покрытие {cover}); лимит {}",
1473 r.bytes as f64 / 1e9,
1474 r.tensors,
1475 r.limit
1476 .map(|l| format!("{:.1} ГБ", l as f64 / 1e9))
1477 .unwrap_or_else(|| "неизвестен".into())
1478 );
1479 if r.skipped > 0 {
1480 tracing::warn!("не закреплено тензоров: {}", r.skipped);
1481 }
1482 }
1483 let st = crate::dsv4::Dsv4State::new(arch.num_layers);
1484 let depth = std::env::var("CMF_DSV4_MTP_DEPTH")
1488 .ok()
1489 .and_then(|v| v.parse::<usize>().ok())
1490 .unwrap_or(3);
1491 pipeline.dsv4_mtp = crate::dsv4::load_mtp(model, &cfg, depth);
1492 crate::dsv4::dspark_reserve_note(&pipeline.dsv4_mtp, &cfg, &dl);
1494 pipeline.dsv4 = Some(Box::new((g, dl, cfg, st)));
1495 }
1496 pipeline.short_conv_cfg = short_conv_cfg;
1497 pipeline.mtp = mtp;
1498 pipeline.install_dynamic_routing(model, false);
1499 match ov {
1503 Overlay::One(sid) => {
1504 pipeline.dyn_active = model.header.skills.iter().position(|s| &s.id == sid);
1505 }
1506 Overlay::Blend(_) => pipeline.dyn_blend_loaded = true,
1507 Overlay::None => {}
1508 }
1509 if let Some(c) = &model.header.calibration {
1512 pipeline.set_calib_temp(c.temperature);
1513 }
1514 let o1 = match crate::nystrom::o1_from_env() {
1520 crate::nystrom::O1Env::Off => None,
1521 crate::nystrom::O1Env::On(cfg) => Some(cfg),
1522 crate::nystrom::O1Env::Unset => model
1523 .header
1524 .provenance
1525 .as_ref()
1526 .and_then(|p| p.get("o1_attn"))
1527 .and_then(crate::nystrom::O1Cfg::from_json),
1528 };
1529 if o1.is_some() {
1530 if pipeline.attn_softcap > 0.0 {
1531 return Err(CmfError::Parse(
1532 "--o1 with attention-logit soft-capping (Gemma-2) is not supported: \
1533 the streaming operator has no capped-score form"
1534 .into(),
1535 ));
1536 }
1537 pipeline.set_o1(o1);
1538 }
1539 Ok(pipeline)
1540 }
1541
1542 pub(crate) fn install_dynamic_routing(&mut self, model: &Arc<CmfModel>, force_f32: bool) {
1547 self.model = Some(model.clone());
1548 self.dyn_force_f32 = force_f32;
1549 let mut per_skill = Vec::with_capacity(model.header.skills.len());
1550 for sk in &model.header.skills {
1551 let mut ffn_layers = std::collections::BTreeSet::new();
1552 let mut non_ffn = false;
1553 let prefix = format!("skill.{}.", sk.id);
1554 for t in model.skill_tensors(&sk.id) {
1555 let rel = &t.name[prefix.len()..]; let toks: Vec<&str> = rel.split('.').collect();
1557 if toks.len() >= 5 && toks[0] == "model" && toks[1] == "layers" && toks[3] == "mlp"
1558 {
1559 if let Ok(li) = toks[2].parse::<usize>() {
1560 ffn_layers.insert(li);
1561 continue;
1562 }
1563 }
1564 non_ffn = true; }
1566 if non_ffn {
1567 tracing::warn!(
1568 "skill '{}' replaces non-FFN tensors — excluded from dynamic \
1569 routing (static overlay still works)",
1570 sk.id
1571 );
1572 per_skill.push(None);
1573 } else {
1574 per_skill.push(Some(ffn_layers.into_iter().collect::<Vec<_>>()));
1575 }
1576 }
1577 self.dyn_skill_layers = per_skill;
1578 }
1579
1580 pub fn set_active_skill(&mut self, idx: Option<usize>) -> Result<(), CmfError> {
1587 self.kv_cache.clear();
1589 self.kv_history.clear();
1590 if self.dyn_active == idx {
1591 return Ok(());
1592 }
1593 let model = self.model.clone().ok_or_else(|| {
1594 CmfError::Parse("dynamic routing needs a model-backed pipeline".into())
1595 })?;
1596 let mut union: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
1597 if let Some(old) = self.dyn_active {
1598 if let Some(Some(ls)) = self.dyn_skill_layers.get(old) {
1599 union.extend(ls.iter().copied());
1600 }
1601 }
1602 let new_id: Option<String> = match idx {
1603 Some(n) => match self.dyn_skill_layers.get(n) {
1604 Some(Some(ls)) => {
1605 union.extend(ls.iter().copied());
1606 Some(model.header.skills[n].id.clone())
1607 }
1608 _ => {
1609 return Err(CmfError::Parse(format!(
1610 "skill index {n} not dynamic-eligible"
1611 )));
1612 }
1613 },
1614 None => None,
1615 };
1616 let ov = match &new_id {
1617 Some(s) => Overlay::One(s),
1618 None => Overlay::None,
1619 };
1620 let arch = model.arch();
1621 for li in union {
1622 self.weights.layers[li].ffn =
1623 build_layer_ffn(&model, arch, li, self.dyn_force_f32, &ov)?;
1624 }
1625 self.dyn_active = idx;
1626 Ok(())
1627 }
1628}