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 from_model_with_overlay(
445 model: &Arc<CmfModel>,
446 sampler_config: SamplerConfig,
447 ov: &Overlay,
448 ) -> Result<Self, CmfError> {
449 let skill = match ov {
450 Overlay::One(s) => Some(*s),
451 _ => None,
452 };
453 if let Some(sid) = skill {
454 let known = model.header.skills.iter().any(|s| s.id == sid)
455 || model.skill_tensors(sid).next().is_some();
456 if !known {
457 return Err(CmfError::Parse(format!(
458 "skill '{sid}' not in this container (header.skills: {:?})",
459 model
460 .header
461 .skills
462 .iter()
463 .map(|s| &s.id)
464 .collect::<Vec<_>>()
465 )));
466 }
467 tracing::info!(
468 "skill '{sid}': {} replacement tensors overlaid",
469 model.skill_tensors(sid).count()
470 );
471 }
472 let arch = model.arch().clone();
473 let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
474 if let Some(heads) = &arch.attention_heads_per_layer {
475 if heads.len() != arch.num_layers {
476 return Err(CmfError::Parse(format!(
477 "arch.attention_heads_per_layer has {} entries, expected {}",
478 heads.len(),
479 arch.num_layers
480 )));
481 }
482 if let Some((li, &nh)) = heads
483 .iter()
484 .enumerate()
485 .find(|(_, nh)| **nh == 0 || **nh % arch.num_kv_heads != 0)
486 {
487 return Err(CmfError::Parse(format!(
488 "layer {li} has {nh} Q heads, which must be nonzero and divisible by {} KV heads",
489 arch.num_kv_heads
490 )));
491 }
492 }
493 if arch
494 .layer_types
495 .iter()
496 .any(|t| matches!(t, LayerType::SlidingAttention))
497 && arch.sliding_window.is_none()
498 {
499 return Err(CmfError::Parse(
500 "model has SlidingAttention layers but no arch.sliding_window".into(),
501 ));
502 }
503
504 let heads_masked = model.masks.masks.iter().any(|m| {
514 m.head_masks.iter().any(|row| {
515 let mut bits = 0usize;
516 for &b in row.iter() {
517 bits += b.count_ones() as usize;
518 }
519 !row.is_empty() && bits < arch.num_attention_heads
520 })
521 });
522 let force_f32 = heads_masked; let mut tokenizer = if let Some(vocab_bytes) = &model.vocab {
526 Tokenizer::from_bytes(vocab_bytes)
527 .map_err(|e| CmfError::Parse(format!("embedded tokenizer: {e}")))?
528 } else {
529 let sidecar = model.path.with_file_name("tokenizer.json");
530 if sidecar.exists() {
531 Tokenizer::from_file(&sidecar)
532 .map_err(|e| CmfError::Parse(format!("sidecar tokenizer: {e}")))?
533 } else {
534 tracing::warn!("no tokenizer in file or sidecar — using byte-level fallback");
535 Tokenizer::byte_level()
536 }
537 };
538 if let Some(tc) = &model.header.tokenizer_config {
540 tokenizer.chat_template = tc.chat_template.clone();
541 tokenizer.extra_eos.extend(tc.eos_token_ids.iter().copied());
542 if tokenizer.bos_token_id.is_none() {
543 tokenizer.bos_token_id = tc.bos_token_id;
544 }
545 tracing::info!(
546 "chat bundle: template {} chars, {} stop ids",
547 tc.chat_template.as_deref().map(str::len).unwrap_or(0),
548 tc.eos_token_ids.len()
549 );
550 }
551 if arch.arch_name.to_lowercase().contains("gemma") && tokenizer.bos_token_id.is_some() {
555 tokenizer.add_bos = true;
556 }
557
558 let embed_tokens = load_matrix(model, "model.embed_tokens.weight", false, ov)?;
560 let final_norm = load_f32(model, "model.norm.weight", ov).map_err(err)?;
561 let lm_head = if model.tensor("lm_head.weight").is_some() {
562 load_matrix(model, "lm_head.weight", false, ov)?
563 } else if arch.tie_word_embeddings {
564 load_matrix(model, "model.embed_tokens.weight", false, ov)?
566 } else {
567 return Err(CmfError::MissingTensor(
568 "lm_head.weight (and tie_word_embeddings is false)".into(),
569 ));
570 };
571
572 let has_linear = arch
574 .layer_types
575 .iter()
576 .any(|t| matches!(t, LayerType::LinearAttention));
577 let mut vmf_cfg = None;
578 let mut gdn_cfg = None;
579 if has_linear {
580 let lc = arch.linear_core.as_ref().ok_or_else(|| {
581 CmfError::Parse(
582 "model has LinearAttention layers but no arch.linear_core — \
583 reconvert with the current converter"
584 .into(),
585 )
586 })?;
587 let need = |v: Option<usize>, name: &str| {
588 v.ok_or_else(|| CmfError::Parse(format!("linear core needs arch.{name}")))
589 };
590 match lc.kind.as_str() {
591 "vmf_phase" => {
592 vmf_cfg = Some(VmfPhaseCfg {
593 num_heads: lc.num_heads,
594 nphase: need(lc.nphase, "linear_core.nphase")?,
595 value_head_dim: lc.value_head_dim,
596 hidden_size: arch.hidden_size,
597 phase_mass: std::env::var("CMF_PHASE_MASS")
600 .ok()
601 .and_then(|v| v.parse().ok())
602 .unwrap_or(0.0),
603 });
604 }
605 "gated_delta_net" => {
606 gdn_cfg = Some(GdnCfg {
607 num_v_heads: lc.num_heads,
608 num_k_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
609 key_head_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
610 value_head_dim: lc.value_head_dim,
611 conv_kernel: need(arch.linear_conv_kernel_dim, "linear_conv_kernel_dim")?,
612 hidden_size: arch.hidden_size,
613 rms_eps: arch.rms_norm_eps,
614 });
615 }
616 other => {
617 return Err(CmfError::Parse(format!(
618 "unknown linear core '{other}' (this runtime executes: \
619 gated_delta_net, vmf_phase)"
620 )));
621 }
622 }
623 }
624
625 let has_kda = arch.layer_types.iter().any(|t| matches!(t, LayerType::Kda));
627 let kda_cfg = if has_kda {
628 let need = |v: Option<usize>, name: &str| {
629 v.ok_or_else(|| CmfError::Parse(format!("KDA core needs arch.{name}")))
630 };
631 Some(crate::linear_core::KdaCfg {
632 num_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
633 head_k_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
634 head_v_dim: need(arch.linear_value_head_dim, "linear_value_head_dim")?,
635 conv_kernel: need(arch.linear_conv_kernel_dim, "linear_conv_kernel_dim")?,
636 hidden_size: arch.hidden_size,
637 rms_eps: arch.rms_norm_eps,
638 })
639 } else {
640 None
641 };
642
643 let has_short_conv = arch
645 .layer_types
646 .iter()
647 .any(|t| matches!(t, LayerType::ShortConv));
648 let short_conv_cfg = if has_short_conv {
649 Some(ShortConvCfg {
650 hidden_size: arch.hidden_size,
651 kernel: arch.linear_conv_kernel_dim.ok_or_else(|| {
652 CmfError::Parse(
653 "model has ShortConv layers but no arch.linear_conv_kernel_dim — \
654 reconvert with the current converter"
655 .into(),
656 )
657 })?,
658 })
659 } else {
660 None
661 };
662
663 let load_full_attn = |prefix: &str, layer: Option<usize>| -> Result<AttnKind, CmfError> {
665 let t = |suffix: &str| load_matrix(model, &format!("{prefix}{suffix}"), force_f32, ov);
666 let n = |suffix: &str| -> Option<Vec<f32>> {
667 model
668 .tensor(&format!("{prefix}{suffix}"))
669 .and_then(|_| load_f32(model, &format!("{prefix}{suffix}"), ov).ok())
670 };
671 if let Some(mla) = arch.mla.as_ref() {
673 let (q_proj, q_a, q_a_norm) = if mla.q_lora_rank.is_some() {
675 (
676 t("self_attn.q_b_proj.weight")?,
677 Some(t("self_attn.q_a_proj.weight")?),
678 Some(n("self_attn.q_a_layernorm.weight").ok_or_else(|| {
679 CmfError::Parse(format!("{prefix}: MLA needs q_a_layernorm"))
680 })?),
681 )
682 } else {
683 (t("self_attn.q_proj.weight")?, None, None)
684 };
685 let hd = mla.qk_rope_head_dim + mla.qk_nope_head_dim;
686 let nh = q_proj.rows() / hd;
687 let mut scale = 1.0 / (hd as f32).sqrt();
690 if let Some(y) = arch.yarn.as_ref() {
691 if let Some(m) = y.mscale_all_dim.filter(|&m| m > 0.0) {
692 let ms = 0.1 * m * y.factor.ln() + 1.0;
693 scale *= ms * ms;
694 }
695 }
696 return Ok(AttnKind::Mla(Box::new(crate::pipeline::MlaWeights {
697 q_proj,
698 q_a,
699 q_a_norm,
700 kv_a: t("self_attn.kv_a_proj_with_mqa.weight")?,
701 kv_a_norm: n("self_attn.kv_a_layernorm.weight").ok_or_else(|| {
702 CmfError::Parse(format!("{prefix}: MLA needs kv_a_layernorm"))
703 })?,
704 kv_b: t("self_attn.kv_b_proj.weight")?,
705 o_proj: t("self_attn.o_proj.weight")?,
706 nh,
707 qk_rope: mla.qk_rope_head_dim,
708 qk_nope: mla.qk_nope_head_dim,
709 v_dim: mla.v_head_dim,
710 lora: mla.kv_lora_rank,
711 scale,
712 nope: mla.nope,
713 })));
714 }
715 let wq = t("self_attn.q_proj.weight")?;
716 let nh = layer
717 .and_then(|li| {
718 arch.attention_heads_per_layer
719 .as_ref()
720 .and_then(|v| v.get(li).copied())
721 })
722 .unwrap_or(arch.num_attention_heads);
723 let output_gate = arch.global_head_dim.is_none() && wq.rows() == 2 * nh * arch.head_dim;
727 let is_global_layer = arch.global_head_dim.is_some()
730 && layer.is_some_and(|li| {
731 arch.sliding_window_pattern
732 .is_some_and(|p| p > 0 && (li + 1) % p == 0)
733 });
734 let expect = if is_global_layer {
735 nh * arch.global_head_dim.unwrap_or(arch.head_dim)
736 } else {
737 nh * arch.head_dim
738 };
739 if !output_gate && wq.rows() != expect {
740 return Err(CmfError::Parse(format!(
741 "{prefix}self_attn.q_proj.weight rows={} != heads({nh}) * head_dim({})",
742 wq.rows(),
743 expect / nh.max(1)
744 )));
745 }
746 let gate_name = format!("{prefix}self_attn.g_proj.weight");
747 let softplus_gate = if model.tensor(&gate_name).is_some() {
748 let gate = load_matrix(model, &gate_name, force_f32, ov)?;
749 if gate.cols() != arch.hidden_size {
750 return Err(CmfError::Parse(format!(
751 "{gate_name} cols={} != hidden_size ({})",
752 gate.cols(),
753 arch.hidden_size
754 )));
755 }
756 let per_head = if gate.rows() == nh {
757 true
758 } else if gate.rows() == nh * arch.head_dim {
759 false
760 } else {
761 return Err(CmfError::Parse(format!(
762 "{gate_name} rows={} must equal heads ({nh}) or heads*head_dim ({})",
763 gate.rows(),
764 nh * arch.head_dim
765 )));
766 };
767 Some((gate, per_head))
768 } else {
769 None
770 };
771 let bias = match (
773 n("self_attn.q_proj.bias"),
774 n("self_attn.k_proj.bias"),
775 n("self_attn.v_proj.bias"),
776 ) {
777 (Some(a), Some(b), Some(c)) => Some((a, b, c)),
778 _ => None,
779 };
780 Ok(AttnKind::Full {
781 wq,
782 wk: t("self_attn.k_proj.weight")?,
783 wv: t("self_attn.v_proj.weight")?,
784 wo: t("self_attn.o_proj.weight")?,
785 q_norm: n("self_attn.q_norm.weight"),
786 k_norm: n("self_attn.k_norm.weight"),
787 output_gate,
788 softplus_gate,
789 bias,
790 })
791 };
792
793 let load_linear_attn = |prefix: &str| -> Result<AttnKind, CmfError> {
794 if gdn_cfg.is_some() {
795 let t = |suffix: &str| {
797 load_matrix(
798 model,
799 &format!("{prefix}linear_attn.{suffix}"),
800 force_f32,
801 ov,
802 )
803 };
804 let f = |suffix: &str| {
805 load_f32(model, &format!("{prefix}linear_attn.{suffix}"), ov).map_err(err)
806 };
807 return Ok(AttnKind::LinearGdn(GdnWeights {
808 in_proj_qkv: t("in_proj_qkv.weight")?,
809 in_proj_z: t("in_proj_z.weight")?,
810 in_proj_a: t("in_proj_a.weight")?,
811 in_proj_b: t("in_proj_b.weight")?,
812 conv1d: f("conv1d.weight")?,
813 a_log: f("A_log")?,
814 dt_bias: f("dt_bias")?,
815 norm: f("norm.weight")?,
816 out_proj: t("out_proj.weight")?,
817 }));
818 }
819 let t = |suffix: &str| {
820 load_matrix(model, &format!("{prefix}vmf_attn.{suffix}"), force_f32, ov)
821 };
822 let a_log = load_f32(model, &format!("{prefix}vmf_attn.A_log"), ov).map_err(err)?;
823 let k_gate = if model
827 .tensor(&format!("{prefix}vmf_attn.k_gate.weight"))
828 .is_some()
829 {
830 Some((
831 t("k_gate.weight")?,
832 load_f32(model, &format!("{prefix}vmf_attn.k_gate.bias"), ov).map_err(err)?,
833 ))
834 } else {
835 None
836 };
837 Ok(AttnKind::Linear(VmfPhaseWeights {
838 thq: t("thq.weight")?,
839 thk: t("thk.weight")?,
840 v_proj: t("v_proj.weight")?,
841 out_proj: t("out_proj.weight")?,
842 decay: a_log.iter().map(|&a| (-(a as f64).exp()).exp()).collect(),
843 k_gate,
844 }))
845 };
846
847 let load_short_conv = |prefix: &str| -> Result<AttnKind, CmfError> {
851 let t = |suffix: &str| {
852 load_matrix(
853 model,
854 &format!("{prefix}short_conv.{suffix}"),
855 force_f32,
856 ov,
857 )
858 };
859 Ok(AttnKind::ShortConv(ShortConvWeights {
860 in_proj: t("in_proj.weight")?,
861 conv: load_f32(model, &format!("{prefix}short_conv.conv.weight"), ov)
862 .map_err(err)?,
863 out_proj: t("out_proj.weight")?,
864 }))
865 };
866
867 let load_kda = |prefix: &str| -> Result<AttnKind, CmfError> {
871 let t = |suffix: &str| {
872 load_matrix(model, &format!("{prefix}kda_attn.{suffix}"), force_f32, ov)
873 };
874 let f = |suffix: &str| {
875 load_f32(model, &format!("{prefix}kda_attn.{suffix}"), ov).map_err(err)
876 };
877 let gate = if model
878 .tensor(&format!("{prefix}kda_attn.g_proj.weight"))
879 .is_some()
880 {
881 crate::linear_core::KdaOutGate::Full(t("g_proj.weight")?)
882 } else {
883 crate::linear_core::KdaOutGate::LowRank(
884 t("g_a_proj.weight")?,
885 t("g_b_proj.weight")?,
886 )
887 };
888 Ok(AttnKind::Kda(Box::new(crate::linear_core::KdaWeights {
889 q_proj: t("q_proj.weight")?,
890 k_proj: t("k_proj.weight")?,
891 v_proj: t("v_proj.weight")?,
892 conv_q: f("q_conv1d.weight")?,
893 conv_k: f("k_conv1d.weight")?,
894 conv_v: f("v_conv1d.weight")?,
895 f_a: t("f_a_proj.weight")?,
896 f_b: t("f_b_proj.weight")?,
897 dt_bias: f("dt_bias")?,
898 a_log: f("A_log")?,
899 b_proj: t("b_proj.weight")?,
900 gate,
901 o_norm: f("o_norm.weight")?,
902 o_proj: t("o_proj.weight")?,
903 gate_lower_bound: arch.kda_gate_lower_bound.map(|v| v as f32),
904 })))
905 };
906
907 fn anyhow_like(ok: bool) -> Result<(), ()> {
908 if ok { Ok(()) } else { Err(()) }
909 }
910 let mut layers = Vec::with_capacity(arch.num_layers);
911 let is_g3n = arch.g3n.is_some();
912 let owns_its_layers = is_g3n || arch.arch_name == "deepseek_v4";
917 for li in 0..(if owns_its_layers { 0 } else { arch.num_layers }) {
918 let prefix = format!("model.layers.{li}.");
919 let attn = match arch.layer_types.get(li) {
920 Some(LayerType::LinearAttention) => load_linear_attn(&prefix)?,
921 Some(LayerType::Kda) => load_kda(&prefix)?,
922 Some(LayerType::ShortConv) => load_short_conv(&prefix)?,
923 _ => load_full_attn(&prefix, Some(li))?,
924 };
925 let pre_ffn = format!("{prefix}pre_feedforward_layernorm.weight");
929 let sandwich = model.tensor(&pre_ffn).is_some();
930 layers.push(LayerWeights {
931 input_norm: load_f32(model, &format!("{prefix}input_layernorm.weight"), ov)
932 .map_err(err)?,
933 post_norm: if sandwich {
934 load_f32(model, &pre_ffn, ov).map_err(err)?
935 } else {
936 load_f32(
937 model,
938 &format!("{prefix}post_attention_layernorm.weight"),
939 ov,
940 )
941 .map_err(err)?
942 },
943 attn_out_norm: if sandwich {
944 Some(
945 load_f32(
946 model,
947 &format!("{prefix}post_attention_layernorm.weight"),
948 ov,
949 )
950 .map_err(err)?,
951 )
952 } else {
953 None
954 },
955 ffn_out_norm: if sandwich {
956 Some(
957 load_f32(
958 model,
959 &format!("{prefix}post_feedforward_layernorm.weight"),
960 ov,
961 )
962 .map_err(err)?,
963 )
964 } else {
965 None
966 },
967 layer_scale: model
969 .tensor(&format!("{prefix}layer_scalar"))
970 .and_then(|_| {
971 load_f32(model, &format!("{prefix}layer_scalar"), ov)
972 .ok()
973 .and_then(|v| v.first().copied())
974 }),
975 ffn: build_layer_ffn(model, &arch, li, false, ov)?,
977 attn,
978 });
979 }
980
981 let mtp_present = model
990 .tensor("model.mtp.layers.0.self_attn.q_proj.weight")
991 .is_some()
992 || model.tensor("model.mtp.eh_proj.weight").is_some();
993 let dsv4_mtp = model.tensor("model.mtp.0.main_proj.weight").is_some();
998 if arch.mtp.is_some() && !mtp_present && !dsv4_mtp {
999 tracing::info!(
1000 "header declares an MTP head but the file carries none — \
1001 loading without it"
1002 );
1003 }
1004 let mtp = if let Some(cfg) = arch.mtp.as_ref().filter(|_| mtp_present) {
1005 if cfg.num_layers != 1 {
1006 return Err(CmfError::Parse(format!(
1007 "MTP with {} blocks not supported yet (only 1)",
1008 cfg.num_layers
1009 )));
1010 }
1011 let p = "model.mtp.";
1012 let attn = load_full_attn("model.mtp.layers.0.", None)?;
1013 Some(MtpModule {
1014 enorm: load_f32(model, &format!("{p}enorm.weight"), ov).map_err(err)?,
1015 hnorm: load_f32(model, &format!("{p}hnorm.weight"), ov).map_err(err)?,
1016 eh_proj: load_matrix(model, &format!("{p}eh_proj.weight"), false, ov)?,
1017 layer: LayerWeights {
1018 attn_out_norm: None,
1019 ffn_out_norm: None,
1020 layer_scale: None,
1021 input_norm: load_f32(model, &format!("{p}layers.0.input_layernorm.weight"), ov)
1022 .map_err(err)?,
1023 post_norm: load_f32(
1024 model,
1025 &format!("{p}layers.0.post_attention_layernorm.weight"),
1026 ov,
1027 )
1028 .map_err(err)?,
1029 ffn: build_ffn_at(model, &arch, &format!("{p}layers.0."), false, ov)?,
1033 attn,
1034 },
1035 final_norm: load_f32(model, &format!("{p}norm.weight"), ov).map_err(err)?,
1036 kv: LayerKvCache::new(arch.num_kv_heads, arch.head_dim),
1037 })
1038 } else {
1039 None
1040 };
1041
1042 tracing::info!(
1043 "Pipeline loaded: {} | {}L ({} linear) | {:.2}B params | storage: {} | MTP: {}",
1044 arch.arch_name,
1045 arch.num_layers,
1046 arch.layer_types
1047 .iter()
1048 .filter(|t| matches!(t, LayerType::LinearAttention))
1049 .count(),
1050 model.total_param_count() as f64 / 1e9,
1051 if force_f32 {
1052 "f32 (masked)"
1053 } else {
1054 "quantized mmap"
1055 },
1056 if mtp.is_some() { "yes" } else { "no" }
1057 );
1058
1059 let cap = std::env::var("CMF_MAX_SEQ")
1062 .ok()
1063 .and_then(|v| v.parse::<usize>().ok())
1064 .unwrap_or(8192);
1065 let max_seq_len = arch.max_position_embeddings.min(cap);
1066
1067 let total_layers = arch.num_layers * arch.num_loops;
1069
1070 let mut pipeline = Pipeline::new(
1071 tokenizer,
1072 PipelineWeights {
1073 embed_tokens,
1074 layers,
1075 lm_head,
1076 final_norm,
1077 },
1078 arch.hidden_size,
1079 arch.intermediate_size,
1080 arch.num_attention_heads,
1081 arch.num_kv_heads,
1082 arch.head_dim,
1083 total_layers,
1084 arch.num_layers, arch.loop_final_norm,
1086 arch.vocab_size,
1087 arch.rms_norm_eps,
1088 arch.rope_theta as f32,
1089 arch.norm_style,
1090 max_seq_len,
1091 sampler_config,
1092 );
1093 let rotary = ((arch.head_dim as f32 * arch.partial_rotary_factor) as usize).max(2);
1094 pipeline.set_rotary(rotary, arch.rope_theta as f32);
1095 pipeline.attention_heads_per_layer = arch.attention_heads_per_layer.clone();
1096 if let Some(yarn) = &arch.yarn {
1097 pipeline.inv_freq = std::sync::Arc::new(crate::attention::yarn_inv_freq(
1098 rotary,
1099 arch.rope_theta as f32,
1100 yarn.factor,
1101 yarn.original_max_position_embeddings,
1102 yarn.beta_fast,
1103 yarn.beta_slow,
1104 ));
1105 pipeline.rope_scale = yarn.attention_factor;
1106 }
1107 pipeline.embed_multiplier = arch.embed_multiplier;
1111 pipeline.logit_multiplier = arch.logit_multiplier;
1112 if let Some(qpas) = arch.query_pre_attn_scalar {
1113 pipeline.attn_scale = 1.0 / (qpas as f32).sqrt();
1114 }
1115 if let (Some(w), Some(p)) = (arch.sliding_window, arch.sliding_window_pattern) {
1116 pipeline.swa = Some((w, p));
1117 if let Some(base) = arch.rope_local_base_freq {
1118 pipeline.inv_freq_local = Some(std::sync::Arc::new(
1119 crate::attention::rope_inv_freq(rotary, base as f32),
1120 ));
1121 }
1122 }
1123 let explicit_sliding: Vec<bool> = arch
1124 .layer_types
1125 .iter()
1126 .map(|t| matches!(t, cortiq_core::LayerType::SlidingAttention))
1127 .collect();
1128 if explicit_sliding.iter().any(|&v| v) {
1129 pipeline.sliding_layers = Some(explicit_sliding);
1130 if let Some(w) = arch.sliding_window {
1131 pipeline.swa = Some((w, usize::MAX));
1132 }
1133 let local_rotary = ((arch.head_dim as f32
1134 * arch
1135 .local_partial_rotary_factor
1136 .unwrap_or(arch.partial_rotary_factor))
1137 as usize)
1138 .max(2);
1139 pipeline.rotary_dim_local = Some(local_rotary);
1140 if let Some(base) = arch.rope_local_base_freq {
1141 pipeline.inv_freq_local = Some(std::sync::Arc::new(
1142 crate::attention::rope_inv_freq(local_rotary, base as f32),
1143 ));
1144 }
1145 }
1146 if let (Some(ghd), Some(gkv)) = (arch.global_head_dim, arch.num_global_kv_heads) {
1150 pipeline.global_attn = Some((ghd, gkv));
1151 let prf = arch.global_partial_rotary_factor.unwrap_or(1.0);
1152 let half = ghd / 2;
1153 let ra = (((prf * ghd as f32) as usize) / 2).min(half);
1154 let mut f = vec![0.0f32; half];
1155 for (i, slot) in f.iter_mut().enumerate().take(ra) {
1156 *slot = 1.0 / (arch.rope_theta as f32).powf(2.0 * i as f32 / ghd as f32);
1157 }
1158 pipeline.inv_freq_global = Some(std::sync::Arc::new(f));
1159 let global_at = |li: usize| -> bool {
1164 match &pipeline.sliding_layers {
1165 Some(map) => !map.get(li).copied().unwrap_or(false),
1166 None => pipeline
1167 .swa
1168 .map(|(_, p)| p > 0 && p != usize::MAX && (li + 1) % p == 0)
1169 .unwrap_or(false),
1170 }
1171 };
1172 for li in 0..arch.num_layers {
1173 if global_at(li) {
1174 pipeline.kv_cache.layers[li] = crate::kv_cache::LayerKvCache::new(gkv, ghd);
1175 }
1176 }
1177 }
1178 if let Some(mla) = arch.mla.as_ref() {
1181 let hd = mla.qk_rope_head_dim + mla.qk_nope_head_dim;
1182 pipeline.head_dim = hd;
1183 pipeline.num_kv_heads = arch.num_attention_heads;
1184 pipeline.rotary_dim = mla.qk_rope_head_dim;
1185 let half = mla.qk_rope_head_dim / 2;
1186 let mut f = vec![0.0f32; half];
1187 for (i, slot) in f.iter_mut().enumerate() {
1188 *slot = 1.0
1189 / (arch.rope_theta as f32).powf(2.0 * i as f32 / mla.qk_rope_head_dim as f32);
1190 }
1191 pipeline.inv_freq = std::sync::Arc::new(f);
1192 for li in 0..arch.num_layers {
1193 pipeline.kv_cache.layers[li] =
1194 crate::kv_cache::LayerKvCache::new(arch.num_attention_heads, hd);
1195 }
1196 }
1197 if let Some(fac) = &arch.rope_freq_factors {
1201 let mut f = pipeline.inv_freq.as_ref().clone();
1202 for (i, v) in f.iter_mut().enumerate() {
1203 if let Some(&d) = fac.get(i) {
1204 *v /= d as f32;
1205 }
1206 }
1207 pipeline.inv_freq = std::sync::Arc::new(f);
1208 }
1209 pipeline.attn_v_norm = arch.attn_v_norm;
1210 pipeline.final_softcap = arch.final_logit_softcapping.map(|c| c as f32);
1211 pipeline.attn_softcap = arch.attn_logit_softcapping.unwrap_or(0.0) as f32;
1212 pipeline.vmf_cfg = vmf_cfg;
1213 pipeline.gdn_cfg = gdn_cfg;
1214 pipeline.kda_cfg = kda_cfg;
1215 if let Some(gc) = arch.g3n.as_ref() {
1216 use crate::g3n::{G3nAltUp, G3nGlobals, G3nLaurel, G3nLayer};
1217 anyhow_like(gc.altup_num_inputs == crate::g3n::ALTUP_N).map_err(|_| {
1218 CmfError::Parse(format!(
1219 "g3n: altup_num_inputs {} != supported {}",
1220 gc.altup_num_inputs,
1221 crate::g3n::ALTUP_N
1222 ))
1223 })?;
1224 let t = |name: &str| load_matrix(model, name, force_f32, ov);
1225 let f = |name: &str| load_f32(model, name, ov).map_err(err);
1226 let mut altup_proj = Vec::new();
1227 let mut altup_unembed = Vec::new();
1228 for i in 0..crate::g3n::ALTUP_N - 1 {
1229 altup_proj.push(t(&format!("model.altup_projections.{i}.weight"))?);
1230 altup_unembed.push(t(&format!("model.altup_unembed_projections.{i}.weight"))?);
1231 }
1232 let first_shared = arch.num_layers.saturating_sub(gc.num_kv_shared_layers);
1233 let sliding_of = |li: usize| {
1234 matches!(
1235 arch.layer_types.get(li),
1236 Some(cortiq_core::LayerType::SlidingAttention)
1237 )
1238 };
1239 let mut g3n_layers = Vec::with_capacity(arch.num_layers);
1240 for li in 0..arch.num_layers {
1241 let pfx = format!("model.layers.{li}.");
1242 let shared = li >= first_shared && first_shared > 0;
1243 let share_src = if shared {
1244 let want = sliding_of(li);
1245 (0..first_shared).rev().find(|&j| sliding_of(j) == want)
1246 } else {
1247 None
1248 };
1249 g3n_layers.push(G3nLayer {
1250 altup: G3nAltUp {
1251 router_norm: f(&format!("{pfx}altup.router_norm.weight"))?,
1252 modality_router: t(&format!("{pfx}altup.modality_router.weight"))?,
1253 prediction_coefs: t(&format!("{pfx}altup.prediction_coefs.weight"))?,
1254 correction_coefs: t(&format!("{pfx}altup.correction_coefs.weight"))?,
1255 correct_output_scale: f(&format!("{pfx}altup.correct_output_scale"))?,
1256 },
1257 laurel: G3nLaurel {
1258 left: t(&format!("{pfx}laurel.linear_left.weight"))?,
1259 right: t(&format!("{pfx}laurel.linear_right.weight"))?,
1260 post_norm: f(&format!("{pfx}laurel.post_laurel_norm.weight"))?,
1261 },
1262 input_norm: f(&format!("{pfx}input_layernorm.weight"))?,
1263 post_attn_norm: f(&format!("{pfx}post_attention_layernorm.weight"))?,
1264 pre_ffw_norm: f(&format!("{pfx}pre_feedforward_layernorm.weight"))?,
1265 post_ffw_norm: f(&format!("{pfx}post_feedforward_layernorm.weight"))?,
1266 wq: t(&format!("{pfx}self_attn.q_proj.weight"))?,
1267 wk: if shared {
1268 None
1269 } else {
1270 Some(t(&format!("{pfx}self_attn.k_proj.weight"))?)
1271 },
1272 wv: if shared {
1273 None
1274 } else {
1275 Some(t(&format!("{pfx}self_attn.v_proj.weight"))?)
1276 },
1277 wo: t(&format!("{pfx}self_attn.o_proj.weight"))?,
1278 q_norm: f(&format!("{pfx}self_attn.q_norm.weight"))?,
1279 k_norm: if shared {
1280 None
1281 } else {
1282 Some(f(&format!("{pfx}self_attn.k_norm.weight"))?)
1283 },
1284 kv_share_src: share_src,
1285 sliding: sliding_of(li),
1286 gate: t(&format!("{pfx}mlp.gate_proj.weight"))?,
1287 up: t(&format!("{pfx}mlp.up_proj.weight"))?,
1288 down: t(&format!("{pfx}mlp.down_proj.weight"))?,
1289 sparsity: gc.activation_sparsity.get(li).copied().unwrap_or(0.0),
1290 ple_gate: t(&format!("{pfx}per_layer_input_gate.weight"))?,
1291 ple_proj: t(&format!("{pfx}per_layer_projection.weight"))?,
1292 post_ple_norm: f(&format!("{pfx}post_per_layer_input_norm.weight"))?,
1293 });
1294 }
1295 let hd = arch.head_dim;
1296 let globals = G3nGlobals {
1297 altup_proj,
1298 altup_unembed,
1299 ple_embed: t("model.embed_tokens_per_layer.weight")?,
1300 ple_model_proj: t("model.per_layer_model_projection.weight")?,
1301 ple_norm: f("model.per_layer_projection_norm.weight")?,
1302 ple_vocab: gc.ple_vocab,
1303 ple_dim: gc.ple_dim,
1304 num_layers: arch.num_layers,
1305 hidden: arch.hidden_size,
1306 rms_eps: arch.rms_norm_eps,
1307 inv_freq_local: crate::attention::rope_inv_freq(
1308 hd,
1309 arch.rope_local_base_freq.unwrap_or(10_000.0) as f32,
1310 ),
1311 inv_freq_global: crate::attention::rope_inv_freq(hd, arch.rope_theta as f32),
1312 window: arch.sliding_window.unwrap_or(512),
1313 };
1314 pipeline.g3n = Some(Box::new((globals, g3n_layers)));
1315 }
1316 if arch.arch_name == "deepseek_v4" {
1321 let moe = arch
1322 .moe
1323 .as_ref()
1324 .ok_or_else(|| CmfError::Parse("deepseek_v4: no moe config".into()))?;
1325 let cfg = crate::dsv4::Dsv4Cfg {
1326 dim: arch.hidden_size,
1327 n_heads: arch.num_attention_heads,
1328 head_dim: arch.head_dim,
1329 rope_head_dim: if arch.partial_rotary_factor < 1.0 {
1336 (((arch.head_dim as f32 * arch.partial_rotary_factor) as usize) & !1)
1337 .clamp(2, arch.head_dim)
1338 } else {
1339 64.min(arch.head_dim)
1340 },
1341 q_lora_rank: 0,
1345 o_lora_rank: 0,
1346 o_groups: 8,
1353 hc_mult: 4,
1354 hc_sinkhorn_iters: 20,
1355 hc_eps: 1e-6,
1356 norm_eps: arch.rms_norm_eps as f32,
1357 n_routed_experts: moe.num_experts,
1358 top_k: moe.top_k,
1359 moe_inter: moe.moe_intermediate_size,
1360 route_scale: moe.routed_scaling_factor.unwrap_or(1.0) as f32,
1361 swiglu_limit: 10.0,
1367 window: arch.sliding_window.unwrap_or(128),
1368 index_topk: 512,
1369 vocab: arch.vocab_size,
1370 };
1371 let (g, dl) = crate::dsv4::load(model, &cfg, arch.num_layers)
1372 .map_err(|e| CmfError::Parse(format!("deepseek_v4: {e}")))?;
1373 let mut cfg = cfg;
1378 if let Some(l0) = dl.first() {
1379 cfg.q_lora_rank = l0.wq_a.rows();
1380 let attn_width = arch.num_attention_heads * arch.head_dim;
1381 if l0.wo_a.cols() > 0 && attn_width % l0.wo_a.cols() == 0 {
1382 cfg.o_groups = (attn_width / l0.wo_a.cols()).max(1);
1383 }
1384 cfg.o_lora_rank = l0.wo_b.cols() / cfg.o_groups.max(1);
1385 cfg.hc_mult = (l0.hc_attn_fn.len() / l0.hc_attn_base.len().max(1)) / cfg.dim.max(1);
1386 if cfg.hc_mult == 0 {
1387 cfg.hc_mult = 4;
1388 }
1389 }
1390 let (yf, yo, ybf, ybs) = match &arch.yarn {
1403 Some(y) => (
1404 y.factor,
1405 y.original_max_position_embeddings,
1406 y.beta_fast,
1407 y.beta_slow,
1408 ),
1409 None => {
1410 tracing::warn!(
1411 "deepseek_v4: the header carries no YaRN profile — \
1412 falling back to the release's (factor 16, original \
1413 65536, beta 32/1). Re-converting with a build that \
1414 reads rope_scaling.type would make this exact."
1415 );
1416 (16.0, 65536, 32.0, 1.0)
1417 }
1418 };
1419 pipeline.inv_freq = std::sync::Arc::new(crate::attention::yarn_inv_freq(
1420 cfg.rope_head_dim,
1421 arch.rope_theta as f32,
1422 yf,
1423 yo,
1424 ybf,
1425 ybs,
1426 ));
1427 if let Ok(stats) = std::env::var("CMF_MOE_PIN") {
1432 let cover = std::env::var("CMF_MOE_PIN_COVER")
1433 .ok()
1434 .and_then(|v| v.parse::<f64>().ok())
1435 .filter(|&c| c > 0.0 && c <= 1.0)
1436 .unwrap_or(0.95);
1437 let hot = crate::pin::hot_experts(&stats, cover);
1438 let mut names: Vec<String> = Vec::new();
1439 for e in &model.tensors {
1440 let is_expert = e.name.contains(".mlp.experts.");
1441 if !is_expert {
1442 names.push(e.name.clone()); }
1444 }
1445 let mut kept_experts = 0usize;
1446 if let Some(hot) = &hot {
1447 for (li, experts) in hot {
1448 for e in experts {
1449 for w in ["gate_proj", "up_proj", "down_proj"] {
1450 names.push(format!("model.layers.{li}.mlp.experts.{e}.{w}.weight"));
1451 }
1452 kept_experts += 1;
1453 }
1454 }
1455 }
1456 let r = crate::pin::pin_tensors(model, &names);
1457 tracing::info!(
1458 "закреплено {:.1} ГБ ({} тензоров, горячих экспертов {kept_experts}, покрытие {cover}); лимит {}",
1459 r.bytes as f64 / 1e9,
1460 r.tensors,
1461 r.limit
1462 .map(|l| format!("{:.1} ГБ", l as f64 / 1e9))
1463 .unwrap_or_else(|| "неизвестен".into())
1464 );
1465 if r.skipped > 0 {
1466 tracing::warn!("не закреплено тензоров: {}", r.skipped);
1467 }
1468 }
1469 let st = crate::dsv4::Dsv4State::new(arch.num_layers);
1470 let depth = std::env::var("CMF_DSV4_MTP_DEPTH")
1474 .ok()
1475 .and_then(|v| v.parse::<usize>().ok())
1476 .unwrap_or(3);
1477 pipeline.dsv4_mtp = crate::dsv4::load_mtp(model, &cfg, depth);
1478 crate::dsv4::dspark_reserve_note(&pipeline.dsv4_mtp, &cfg, &dl);
1480 pipeline.dsv4 = Some(Box::new((g, dl, cfg, st)));
1481 }
1482 pipeline.short_conv_cfg = short_conv_cfg;
1483 pipeline.mtp = mtp;
1484 pipeline.install_dynamic_routing(model, false);
1485 match ov {
1489 Overlay::One(sid) => {
1490 pipeline.dyn_active = model.header.skills.iter().position(|s| &s.id == sid);
1491 }
1492 Overlay::Blend(_) => pipeline.dyn_blend_loaded = true,
1493 Overlay::None => {}
1494 }
1495 if let Some(c) = &model.header.calibration {
1498 pipeline.set_calib_temp(c.temperature);
1499 }
1500 let o1 = match crate::nystrom::o1_from_env() {
1506 crate::nystrom::O1Env::Off => None,
1507 crate::nystrom::O1Env::On(cfg) => Some(cfg),
1508 crate::nystrom::O1Env::Unset => model
1509 .header
1510 .provenance
1511 .as_ref()
1512 .and_then(|p| p.get("o1_attn"))
1513 .and_then(crate::nystrom::O1Cfg::from_json),
1514 };
1515 if o1.is_some() {
1516 if pipeline.attn_softcap > 0.0 {
1517 return Err(CmfError::Parse(
1518 "--o1 with attention-logit soft-capping (Gemma-2) is not supported: \
1519 the streaming operator has no capped-score form"
1520 .into(),
1521 ));
1522 }
1523 pipeline.set_o1(o1);
1524 }
1525 Ok(pipeline)
1526 }
1527
1528 pub(crate) fn install_dynamic_routing(&mut self, model: &Arc<CmfModel>, force_f32: bool) {
1533 self.model = Some(model.clone());
1534 self.dyn_force_f32 = force_f32;
1535 let mut per_skill = Vec::with_capacity(model.header.skills.len());
1536 for sk in &model.header.skills {
1537 let mut ffn_layers = std::collections::BTreeSet::new();
1538 let mut non_ffn = false;
1539 let prefix = format!("skill.{}.", sk.id);
1540 for t in model.skill_tensors(&sk.id) {
1541 let rel = &t.name[prefix.len()..]; let toks: Vec<&str> = rel.split('.').collect();
1543 if toks.len() >= 5 && toks[0] == "model" && toks[1] == "layers" && toks[3] == "mlp"
1544 {
1545 if let Ok(li) = toks[2].parse::<usize>() {
1546 ffn_layers.insert(li);
1547 continue;
1548 }
1549 }
1550 non_ffn = true; }
1552 if non_ffn {
1553 tracing::warn!(
1554 "skill '{}' replaces non-FFN tensors — excluded from dynamic \
1555 routing (static overlay still works)",
1556 sk.id
1557 );
1558 per_skill.push(None);
1559 } else {
1560 per_skill.push(Some(ffn_layers.into_iter().collect::<Vec<_>>()));
1561 }
1562 }
1563 self.dyn_skill_layers = per_skill;
1564 }
1565
1566 pub fn set_active_skill(&mut self, idx: Option<usize>) -> Result<(), CmfError> {
1573 self.kv_cache.clear();
1575 self.kv_history.clear();
1576 if self.dyn_active == idx {
1577 return Ok(());
1578 }
1579 let model = self.model.clone().ok_or_else(|| {
1580 CmfError::Parse("dynamic routing needs a model-backed pipeline".into())
1581 })?;
1582 let mut union: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
1583 if let Some(old) = self.dyn_active {
1584 if let Some(Some(ls)) = self.dyn_skill_layers.get(old) {
1585 union.extend(ls.iter().copied());
1586 }
1587 }
1588 let new_id: Option<String> = match idx {
1589 Some(n) => match self.dyn_skill_layers.get(n) {
1590 Some(Some(ls)) => {
1591 union.extend(ls.iter().copied());
1592 Some(model.header.skills[n].id.clone())
1593 }
1594 _ => {
1595 return Err(CmfError::Parse(format!(
1596 "skill index {n} not dynamic-eligible"
1597 )));
1598 }
1599 },
1600 None => None,
1601 };
1602 let ov = match &new_id {
1603 Some(s) => Overlay::One(s),
1604 None => Overlay::None,
1605 };
1606 let arch = model.arch();
1607 for li in union {
1608 self.weights.layers[li].ffn =
1609 build_layer_ffn(&model, arch, li, self.dyn_force_f32, &ov)?;
1610 }
1611 self.dyn_active = idx;
1612 Ok(())
1613 }
1614}