Skip to main content

cortiq_engine/
loader.rs

1//! Weight loader: CMF tensor directory → Pipeline.
2//!
3//! Storage rule: models WITH task masks are dequantized to f32 (masked
4//! execution needs f32 row access; skill files are small by design).
5//! Models without masks keep quantized matrices zero-copy from the mmap
6//! (`QTensor::Mapped`) — this is what lets a 15B file run in a few GB
7//! of RSS instead of 60 GB of f32.
8//!
9//! Layer kinds come from `arch.layer_types`: FullAttention loads
10//! `self_attn.*` (with auto-detected Qwen3.5 extras: per-head qk-norm by
11//! tensor presence, output gate by q_proj row count); LinearAttention
12//! loads the canonical core `vmf_attn.*` (folded at convert time).
13
14use 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
28/// Tensor source selector (spec §9): backbone, one skill's overlay, or
29/// a soft superposition of top-m skills (claim 14 working tensors).
30pub enum Overlay<'a> {
31    None,
32    One(&'a str),
33    /// (skill_id, weight); weights sum to 1 (softmax(−E/T) upstream).
34    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
57/// Weighted working tensor (claim 14): Σ wᵢ·Tᵢ, where Tᵢ is the
58/// skill's replacement when present, else the backbone tensor.
59fn 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
87/// Dequantize a tensor fully into f32 (norms, masked models).
88fn 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
107/// Build one layer's FFN (dense or MoE) under a given overlay. Shared
108/// by the static loader AND dynamic per-token skill switching
109/// (`Pipeline::set_active_skill`): switching skills = rebuilding the
110/// FFN of the touched layers, cheap because Mapped tensors are just
111/// re-resolved mmap pointers (no dequant, no copy).
112pub(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    let prefix = format!("model.layers.{li}.");
120    let load_dense = |p: &str| -> Result<DenseFfn, CmfError> {
121        let gate_proj = load_matrix(model, &format!("{p}gate_proj.weight"), force_f32, ov)?;
122        let up_proj = load_matrix(model, &format!("{p}up_proj.weight"), force_f32, ov)?;
123        let down_proj = load_matrix(model, &format!("{p}down_proj.weight"), force_f32, ov)?;
124        // FFN triple invariant (holds for dense and each MoE expert;
125        // enforced loudly so a malformed defrag/repack — spec §11 — fails
126        // at load instead of silently mis-computing). inter' is per-layer.
127        let inter = gate_proj.rows();
128        if up_proj.rows() != inter || down_proj.cols() != inter {
129            return Err(CmfError::Parse(format!(
130                "{p}: FFN dims disagree (gate.rows={inter}, up.rows={}, \
131                 down.cols={}); all three must equal inter'",
132                up_proj.rows(),
133                down_proj.cols()
134            )));
135        }
136        if down_proj.rows() != arch.hidden_size {
137            return Err(CmfError::Parse(format!(
138                "{p}: down_proj.rows={} != hidden_size={}",
139                down_proj.rows(),
140                arch.hidden_size
141            )));
142        }
143        Ok(DenseFfn {
144            gate_proj,
145            up_proj,
146            down_proj,
147            act: crate::pipeline::Act::from_arch_full(arch),
148        })
149    };
150    let router_name = format!("{prefix}mlp.gate.weight");
151    if model.tensor(&router_name).is_none() {
152        return Ok(FfnKind::Dense(load_dense(&format!("{prefix}mlp."))?));
153    }
154    let cfg = arch.moe.as_ref().ok_or_else(|| {
155        CmfError::Parse(format!(
156            "{router_name} present but header has no arch.moe block"
157        ))
158    })?;
159    // Experts enumerate by TENSOR PRESENCE up to the header count — a
160    // moe-defrag'd specialist keeps a per-layer contiguous prefix of
161    // renumbered experts (fewer than arch.moe.num_experts), with the
162    // router rows sliced to match.
163    let mut experts = Vec::new();
164    for e in 0..cfg.num_experts {
165        if model
166            .tensor(&format!("{prefix}mlp.experts.{e}.gate_proj.weight"))
167            .is_none()
168        {
169            break;
170        }
171        experts.push(load_dense(&format!("{prefix}mlp.experts.{e}."))?);
172    }
173    if experts.is_empty() {
174        return Err(CmfError::Parse(format!(
175            "{prefix}: router present but no expert tensors"
176        )));
177    }
178    let shared = if model
179        .tensor(&format!("{prefix}mlp.shared_expert.gate_proj.weight"))
180        .is_some()
181    {
182        let gate_name = format!("{prefix}mlp.shared_expert_gate.weight");
183        Some((
184            load_dense(&format!("{prefix}mlp.shared_expert."))?,
185            if model.tensor(&gate_name).is_some() {
186                Some(load_matrix(model, &gate_name, force_f32, ov)?)
187            } else {
188                None
189            },
190        ))
191    } else {
192        None
193    };
194    // LFM2-MoE selection bias (`mlp.expert_bias`): present iff the model
195    // routes with a bias; loaded by tensor presence.
196    let bias_name = format!("{prefix}mlp.expert_bias");
197    let expert_bias = if model.tensor(&bias_name).is_some() {
198        Some(load_f32(model, &bias_name, ov).map_err(CmfError::Parse)?)
199    } else {
200        None
201    };
202    // CMF_MOE_TOPK=N (opt-in): route to fewer experts than the header
203    // asks. MoE decode is memory-bound — every selected expert streams
204    // its three matrices per token — so halving k halves that traffic;
205    // the renormalized top-k keeps the mixture a proper average.
206    // Quality is the experiment — measure ppl before trusting.
207    let top_k = std::env::var("CMF_MOE_TOPK")
208        .ok()
209        .and_then(|v| v.parse::<usize>().ok())
210        .filter(|&k| k >= 1 && k <= cfg.top_k)
211        .inspect(|k| tracing::info!("MoE top_k override: {} (header {})", k, cfg.top_k))
212        .unwrap_or(cfg.top_k);
213    // CMF_MOE_TAU=0.x (opt-in): adaptive routing — see MoeFfn::route_tau.
214    let route_tau = std::env::var("CMF_MOE_TAU")
215        .ok()
216        .and_then(|v| v.parse::<f32>().ok())
217        .filter(|&t| t > 0.0 && t < 1.0)
218        .inspect(|t| tracing::info!("MoE adaptive routing: tau {t}"));
219    let mask = moe_task_mask(&prefix, experts.len());
220    let router = load_matrix(model, &router_name, force_f32, ov)?;
221    if router.rows() != experts.len() {
222        return Err(CmfError::Parse(format!(
223            "{router_name}: {} rows != {} experts",
224            router.rows(),
225            experts.len()
226        )));
227    }
228    let top_k = top_k.min(experts.len());
229    // Gemma-4: per-expert weight scale after the top-k renorm; its
230    // presence also marks the scale-less-rms router input (the folded
231    // router gain — see the converter).
232    let pes_name = format!("{prefix}mlp.per_expert_scale");
233    let per_expert_scale = if model.tensor(&pes_name).is_some() {
234        Some(load_f32(model, &pes_name, ov).map_err(CmfError::Parse)?)
235    } else {
236        None
237    };
238    let router_input_norm = per_expert_scale.is_some();
239    let moe = MoeFfn {
240        router,
241        experts,
242        top_k,
243        route_tau,
244        norm_topk_prob: cfg.norm_topk_prob,
245        router_sigmoid: cfg.router_sigmoid,
246        expert_bias,
247        routed_scaling: cfg.routed_scaling_factor.unwrap_or(1.0),
248        shared,
249        stats: std::cell::RefCell::new(Vec::new()),
250        mask,
251        per_expert_scale,
252        router_input_norm,
253    };
254    // Gemma-4 dual-branch layer: a dense MLP coexists with the routed
255    // experts, each branch inside its own norm sandwich.
256    if model
257        .tensor(&format!("{prefix}mlp.gate_proj.weight"))
258        .is_some()
259    {
260        let norm = |suffix: &str| -> Result<Vec<f32>, CmfError> {
261            load_f32(model, &format!("{prefix}{suffix}.weight"), ov).map_err(CmfError::Parse)
262        };
263        return Ok(FfnKind::DenseMoe(Box::new(
264            crate::pipeline::DenseMoeFfn {
265                dense: load_dense(&format!("{prefix}mlp."))?,
266                moe,
267                post_norm_1: norm("post_feedforward_layernorm_1")?,
268                pre_norm_2: norm("pre_feedforward_layernorm_2")?,
269                post_norm_2: norm("post_feedforward_layernorm_2")?,
270            },
271        )));
272    }
273    Ok(FfnKind::Moe(moe))
274}
275
276/// Task mask over routed experts (opt-in, experimental): DTG-MA applied
277/// to MoE. `CMF_MOE_MASK=<stats.json>` points at a claim-12 B-field dump
278/// (`CMF_MOE_STATS` output — per-layer expert-selection counts from a
279/// task-representative run); `CMF_MOE_MASK_COVER` (default 0.9) keeps,
280/// per layer, the smallest top set of experts reaching that fraction of
281/// the recorded routing mass. Selection then happens over the allowed
282/// set only (softmax renormalizes). Gate any real use on a ppl A/B.
283fn moe_task_mask(prefix: &str, ne: usize) -> Option<Vec<bool>> {
284    use std::sync::OnceLock;
285    static CFG: OnceLock<Option<(std::collections::HashMap<usize, Vec<u64>>, f64)>> =
286        OnceLock::new();
287    let cfg = CFG.get_or_init(|| {
288        let path = std::env::var("CMF_MOE_MASK").ok()?;
289        let cover = std::env::var("CMF_MOE_MASK_COVER")
290            .ok()
291            .and_then(|v| v.parse::<f64>().ok())
292            .filter(|&c| c > 0.0 && c <= 1.0)
293            .unwrap_or(0.9);
294        let text = std::fs::read_to_string(&path)
295            .map_err(|e| tracing::warn!("CMF_MOE_MASK: cannot read {path}: {e}"))
296            .ok()?;
297        let map: std::collections::HashMap<String, Vec<u64>> =
298            serde_json::from_str(&text)
299                .map_err(|e| tracing::warn!("CMF_MOE_MASK: bad JSON in {path}: {e}"))
300                .ok()?;
301        tracing::info!("MoE task mask: {path}, cover {cover}");
302        Some((
303            map.into_iter()
304                .filter_map(|(k, v)| Some((k.parse::<usize>().ok()?, v)))
305                .collect(),
306            cover,
307        ))
308    });
309    let (stats, cover) = cfg.as_ref()?;
310    // The layer index rides in the tensor prefix ("model.layers.N.").
311    let li: usize = prefix
312        .split("layers.")
313        .nth(1)?
314        .split('.')
315        .next()?
316        .parse()
317        .ok()?;
318    let counts = stats.get(&li)?;
319    if counts.len() != ne {
320        tracing::warn!("CMF_MOE_MASK: layer {li} has {} counts, model has {ne} experts — skipped", counts.len());
321        return None;
322    }
323    let total: u64 = counts.iter().sum();
324    if total == 0 {
325        return None;
326    }
327    let mut order: Vec<usize> = (0..ne).collect();
328    order.sort_unstable_by_key(|&e| std::cmp::Reverse(counts[e]));
329    let mut mask = vec![false; ne];
330    let mut acc = 0u64;
331    let mut kept = 0usize;
332    for &e in &order {
333        mask[e] = true;
334        acc += counts[e];
335        kept += 1;
336        if (acc as f64) >= cover * (total as f64) {
337            break;
338        }
339    }
340    tracing::info!("MoE task mask L{li}: {kept}/{ne} experts for {:.0}% mass", cover * 100.0);
341    Some(mask)
342}
343
344fn load_matrix(
345    model: &Arc<CmfModel>,
346    name: &str,
347    force_f32: bool,
348    ov: &Overlay,
349) -> Result<QTensor, CmfError> {
350    // Claim 14: a blended working tensor is materialized in f32 and
351    // held resident (the overlay-cache slot); single skills stay
352    // zero-copy pointers into the mmap.
353    if ov.blend_touches(model, name) {
354        if let Overlay::Blend(list) = ov {
355            let entry = model
356                .tensor(name)
357                .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
358            let data =
359                blend_f32(model, name, list).map_err(|e| CmfError::Parse(format!("blend: {e}")))?;
360            return Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]));
361        }
362    }
363    let skill = match ov {
364        Overlay::One(s) => Some(*s),
365        _ => None,
366    };
367    // Tensor-source indirection (spec §9): the skill's replacement is
368    // read in place of the backbone tensor — either/or, never a sum.
369    let name: &str = &match skill {
370        Some(sid) if model.tensor(&format!("skill.{sid}.{name}")).is_some() => {
371            format!("skill.{sid}.{name}")
372        }
373        _ => name.to_string(),
374    };
375    let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
376    if force_f32 {
377        let entry = model
378            .tensor(name)
379            .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
380        if entry.shape.len() != 2 {
381            return Err(err(format!("'{name}' is not 2-D")));
382        }
383        let data = load_f32(model, name, &Overlay::None).map_err(err)?;
384        Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]))
385    } else {
386        QTensor::from_model(model, name).map_err(err)
387    }
388}
389
390impl Pipeline {
391    /// Build a runnable pipeline from an opened CMF model.
392    pub fn from_model(
393        model: &Arc<CmfModel>,
394        sampler_config: SamplerConfig,
395    ) -> Result<Self, CmfError> {
396        Self::from_model_with_skill(model, sampler_config, None)
397    }
398
399    /// Same, with a skill overlaid (spec §9): every layer tensor is
400    /// resolved through tensor-source indirection — the skill's
401    /// full-shape replacement is read in place of the backbone tensor.
402    /// No per-skill model is ever assembled: Mapped tensors are
403    /// pointers into the one shared mmap.
404    pub fn from_model_with_skill(
405        model: &Arc<CmfModel>,
406        sampler_config: SamplerConfig,
407        skill: Option<&str>,
408    ) -> Result<Self, CmfError> {
409        match skill {
410            Some(s) => Self::from_model_with_overlay(model, sampler_config, &Overlay::One(s)),
411            None => Self::from_model_with_overlay(model, sampler_config, &Overlay::None),
412        }
413    }
414
415    /// Soft superposition (claim 14): working tensors accumulated from
416    /// the given (skill, weight) list — softmax(−E/T) upstream.
417    pub fn from_model_with_blend(
418        model: &Arc<CmfModel>,
419        sampler_config: SamplerConfig,
420        blend: &[(String, f32)],
421    ) -> Result<Self, CmfError> {
422        Self::from_model_with_overlay(model, sampler_config, &Overlay::Blend(blend))
423    }
424
425    fn from_model_with_overlay(
426        model: &Arc<CmfModel>,
427        sampler_config: SamplerConfig,
428        ov: &Overlay,
429    ) -> Result<Self, CmfError> {
430        let skill = match ov {
431            Overlay::One(s) => Some(*s),
432            _ => None,
433        };
434        if let Some(sid) = skill {
435            let known = model.header.skills.iter().any(|s| s.id == sid)
436                || model.skill_tensors(sid).next().is_some();
437            if !known {
438                return Err(CmfError::Parse(format!(
439                    "skill '{sid}' not in this container (header.skills: {:?})",
440                    model
441                        .header
442                        .skills
443                        .iter()
444                        .map(|s| &s.id)
445                        .collect::<Vec<_>>()
446                )));
447            }
448            tracing::info!(
449                "skill '{sid}': {} replacement tensors overlaid",
450                model.skill_tensors(sid).count()
451            );
452        }
453        let arch = model.arch().clone();
454        let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
455        if let Some(heads) = &arch.attention_heads_per_layer {
456            if heads.len() != arch.num_layers {
457                return Err(CmfError::Parse(format!(
458                    "arch.attention_heads_per_layer has {} entries, expected {}",
459                    heads.len(),
460                    arch.num_layers
461                )));
462            }
463            if let Some((li, &nh)) = heads
464                .iter()
465                .enumerate()
466                .find(|(_, nh)| **nh == 0 || **nh % arch.num_kv_heads != 0)
467            {
468                return Err(CmfError::Parse(format!(
469                    "layer {li} has {nh} Q heads, which must be nonzero and divisible by {} KV heads",
470                    arch.num_kv_heads
471                )));
472            }
473        }
474        if arch
475            .layer_types
476            .iter()
477            .any(|t| matches!(t, LayerType::SlidingAttention))
478            && arch.sliding_window.is_none()
479        {
480            return Err(CmfError::Parse(
481                "model has SlidingAttention layers but no arch.sliding_window".into(),
482            ));
483        }
484
485        // Masks × quantized mmap: only ATTENTION keeps f32 (the head-mask
486        // path needs f32 slices). FFN masks now run sparse directly on the
487        // quant bytes (sparse_ffn_quant), and embed/lm_head are never
488        // masked — so a masked model runs at quantized RSS, not the old
489        // whole-model-f32 blowup.
490        let masks_present = !model.masks.masks.is_empty();
491        let force_f32 = masks_present; // attention only (head masks)
492
493        // ── Tokenizer: embedded → sidecar → byte-level fallback ──
494        let mut tokenizer = if let Some(vocab_bytes) = &model.vocab {
495            Tokenizer::from_bytes(vocab_bytes)
496                .map_err(|e| CmfError::Parse(format!("embedded tokenizer: {e}")))?
497        } else {
498            let sidecar = model.path.with_file_name("tokenizer.json");
499            if sidecar.exists() {
500                Tokenizer::from_file(&sidecar)
501                    .map_err(|e| CmfError::Parse(format!("sidecar tokenizer: {e}")))?
502            } else {
503                tracing::warn!("no tokenizer in file or sidecar — using byte-level fallback");
504                Tokenizer::byte_level()
505            }
506        };
507        // Chat/eos bundle (spec §6.1): the FILE defines chat behavior.
508        if let Some(tc) = &model.header.tokenizer_config {
509            tokenizer.chat_template = tc.chat_template.clone();
510            tokenizer.extra_eos.extend(tc.eos_token_ids.iter().copied());
511            if tokenizer.bos_token_id.is_none() {
512                tokenizer.bos_token_id = tc.bos_token_id;
513            }
514            tracing::info!(
515                "chat bundle: template {} chars, {} stop ids",
516                tc.chat_template.as_deref().map(str::len).unwrap_or(0),
517                tc.eos_token_ids.len()
518            );
519        }
520        // Gemma's contract requires <bos> at sequence start, but its
521        // tokenizer.json post-processor does not add it (the chat
522        // template does). Raw prompts need it too — word salad without.
523        if arch.arch_name.to_lowercase().contains("gemma") && tokenizer.bos_token_id.is_some() {
524            tokenizer.add_bos = true;
525        }
526
527        // ── Top-level weights (never masked → always quantized) ──
528        let embed_tokens = load_matrix(model, "model.embed_tokens.weight", false, ov)?;
529        let final_norm = load_f32(model, "model.norm.weight", ov).map_err(err)?;
530        let lm_head = if model.tensor("lm_head.weight").is_some() {
531            load_matrix(model, "lm_head.weight", false, ov)?
532        } else if arch.tie_word_embeddings {
533            // Tied: reuse the embedding matrix (re-open, cheap for Mapped).
534            load_matrix(model, "model.embed_tokens.weight", false, ov)?
535        } else {
536            return Err(CmfError::MissingTensor(
537                "lm_head.weight (and tie_word_embeddings is false)".into(),
538            ));
539        };
540
541        // ── Linear-core geometry (required if any linear layer exists) ──
542        let has_linear = arch
543            .layer_types
544            .iter()
545            .any(|t| matches!(t, LayerType::LinearAttention));
546        let mut vmf_cfg = None;
547        let mut gdn_cfg = None;
548        if has_linear {
549            let lc = arch.linear_core.as_ref().ok_or_else(|| {
550                CmfError::Parse(
551                    "model has LinearAttention layers but no arch.linear_core — \
552                     reconvert with the current converter"
553                        .into(),
554                )
555            })?;
556            let need = |v: Option<usize>, name: &str| {
557                v.ok_or_else(|| CmfError::Parse(format!("linear core needs arch.{name}")))
558            };
559            match lc.kind.as_str() {
560                "vmf_phase" => {
561                    vmf_cfg = Some(VmfPhaseCfg {
562                        num_heads: lc.num_heads,
563                        nphase: need(lc.nphase, "linear_core.nphase")?,
564                        value_head_dim: lc.value_head_dim,
565                        hidden_size: arch.hidden_size,
566                        // θ-mass (η′): default 0 (massless); CMF_PHASE_MASS
567                        // widens the phase kernel for folded-unhealed models.
568                        phase_mass: std::env::var("CMF_PHASE_MASS")
569                            .ok()
570                            .and_then(|v| v.parse().ok())
571                            .unwrap_or(0.0),
572                    });
573                }
574                "gated_delta_net" => {
575                    gdn_cfg = Some(GdnCfg {
576                        num_v_heads: lc.num_heads,
577                        num_k_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
578                        key_head_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
579                        value_head_dim: lc.value_head_dim,
580                        conv_kernel: need(arch.linear_conv_kernel_dim, "linear_conv_kernel_dim")?,
581                        hidden_size: arch.hidden_size,
582                        rms_eps: arch.rms_norm_eps,
583                    });
584                }
585                other => {
586                    return Err(CmfError::Parse(format!(
587                        "unknown linear core '{other}' (this runtime executes: \
588                         gated_delta_net, vmf_phase)"
589                    )));
590                }
591            }
592        }
593
594        // ── KDA geometry (Kimi Linear / Kimi-K3 delta-attention layers) ──
595        let has_kda = arch
596            .layer_types
597            .iter()
598            .any(|t| matches!(t, LayerType::Kda));
599        let kda_cfg = if has_kda {
600            let need = |v: Option<usize>, name: &str| {
601                v.ok_or_else(|| CmfError::Parse(format!("KDA core needs arch.{name}")))
602            };
603            Some(crate::linear_core::KdaCfg {
604                num_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
605                head_k_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
606                head_v_dim: need(arch.linear_value_head_dim, "linear_value_head_dim")?,
607                conv_kernel: need(arch.linear_conv_kernel_dim, "linear_conv_kernel_dim")?,
608                hidden_size: arch.hidden_size,
609                rms_eps: arch.rms_norm_eps,
610            })
611        } else {
612            None
613        };
614
615        // ── Short-convolution geometry (LFM2 conv mixer layers) ──
616        let has_short_conv = arch
617            .layer_types
618            .iter()
619            .any(|t| matches!(t, LayerType::ShortConv));
620        let short_conv_cfg = if has_short_conv {
621            Some(ShortConvCfg {
622                hidden_size: arch.hidden_size,
623                kernel: arch.linear_conv_kernel_dim.ok_or_else(|| {
624                    CmfError::Parse(
625                        "model has ShortConv layers but no arch.linear_conv_kernel_dim — \
626                         reconvert with the current converter"
627                            .into(),
628                    )
629                })?,
630            })
631        } else {
632            None
633        };
634
635        // ── Layers ──
636        let load_full_attn = |prefix: &str, layer: Option<usize>| -> Result<AttnKind, CmfError> {
637            let t = |suffix: &str| load_matrix(model, &format!("{prefix}{suffix}"), force_f32, ov);
638            let n = |suffix: &str| -> Option<Vec<f32>> {
639                model
640                    .tensor(&format!("{prefix}{suffix}"))
641                    .and_then(|_| load_f32(model, &format!("{prefix}{suffix}"), ov).ok())
642            };
643            // DeepSeek-V2 MLA: the latent projections replace the k/v pair.
644            if let Some(mla) = arch.mla.as_ref() {
645                // Compressed q (K3/V3): q_a → rms → q_b; direct otherwise.
646                let (q_proj, q_a, q_a_norm) = if mla.q_lora_rank.is_some() {
647                    (
648                        t("self_attn.q_b_proj.weight")?,
649                        Some(t("self_attn.q_a_proj.weight")?),
650                        Some(n("self_attn.q_a_layernorm.weight").ok_or_else(|| {
651                            CmfError::Parse(format!("{prefix}: MLA needs q_a_layernorm"))
652                        })?),
653                    )
654                } else {
655                    (t("self_attn.q_proj.weight")?, None, None)
656                };
657                let hd = mla.qk_rope_head_dim + mla.qk_nope_head_dim;
658                let nh = q_proj.rows() / hd;
659                // YaRN mscale²: DeepSeek corrects the softmax scale by
660                // (0.1·mscale_all_dim·ln(factor)+1)².
661                let mut scale = 1.0 / (hd as f32).sqrt();
662                if let Some(y) = arch.yarn.as_ref() {
663                    if let Some(m) = y.mscale_all_dim.filter(|&m| m > 0.0) {
664                        let ms = 0.1 * m * y.factor.ln() + 1.0;
665                        scale *= ms * ms;
666                    }
667                }
668                return Ok(AttnKind::Mla(Box::new(crate::pipeline::MlaWeights {
669                    q_proj,
670                    q_a,
671                    q_a_norm,
672                    kv_a: t("self_attn.kv_a_proj_with_mqa.weight")?,
673                    kv_a_norm: n("self_attn.kv_a_layernorm.weight").ok_or_else(|| {
674                        CmfError::Parse(format!("{prefix}: MLA needs kv_a_layernorm"))
675                    })?,
676                    kv_b: t("self_attn.kv_b_proj.weight")?,
677                    o_proj: t("self_attn.o_proj.weight")?,
678                    nh,
679                    qk_rope: mla.qk_rope_head_dim,
680                    qk_nope: mla.qk_nope_head_dim,
681                    v_dim: mla.v_head_dim,
682                    lora: mla.kv_lora_rank,
683                    scale,
684                    nope: mla.nope,
685                })));
686            }
687            let wq = t("self_attn.q_proj.weight")?;
688            let nh = layer
689                .and_then(|li| {
690                    arch.attention_heads_per_layer
691                        .as_ref()
692                        .and_then(|v| v.get(li).copied())
693                })
694                .unwrap_or(arch.num_attention_heads);
695            // Qwen3.5 output gate: q_proj rows = 2·nh·hd (per-head [q; gate]).
696            // Gemma-4 global layers legitimately have nh·global_head_dim
697            // rows (which can equal 2·nh·hd) — never gated.
698            let output_gate = arch.global_head_dim.is_none() && wq.rows() == 2 * nh * arch.head_dim;
699            // Gemma-4 global layers run MQA at global_head_dim — their
700            // q_proj legitimately carries nh·ghd rows.
701            let is_global_layer = arch.global_head_dim.is_some()
702                && layer.is_some_and(|li| {
703                    arch.sliding_window_pattern
704                        .is_some_and(|p| p > 0 && (li + 1) % p == 0)
705                });
706            let expect = if is_global_layer {
707                nh * arch.global_head_dim.unwrap_or(arch.head_dim)
708            } else {
709                nh * arch.head_dim
710            };
711            if !output_gate && wq.rows() != expect {
712                return Err(CmfError::Parse(format!(
713                    "{prefix}self_attn.q_proj.weight rows={} != heads({nh}) * head_dim({})",
714                    wq.rows(),
715                    expect / nh.max(1)
716                )));
717            }
718            let gate_name = format!("{prefix}self_attn.g_proj.weight");
719            let softplus_gate = if model.tensor(&gate_name).is_some() {
720                let gate = load_matrix(model, &gate_name, force_f32, ov)?;
721                if gate.cols() != arch.hidden_size {
722                    return Err(CmfError::Parse(format!(
723                        "{gate_name} cols={} != hidden_size ({})",
724                        gate.cols(),
725                        arch.hidden_size
726                    )));
727                }
728                let per_head = if gate.rows() == nh {
729                    true
730                } else if gate.rows() == nh * arch.head_dim {
731                    false
732                } else {
733                    return Err(CmfError::Parse(format!(
734                        "{gate_name} rows={} must equal heads ({nh}) or heads*head_dim ({})",
735                        gate.rows(),
736                        nh * arch.head_dim
737                    )));
738                };
739                Some((gate, per_head))
740            } else {
741                None
742            };
743            // Qwen2-family projection biases (by tensor presence).
744            let bias = match (
745                n("self_attn.q_proj.bias"),
746                n("self_attn.k_proj.bias"),
747                n("self_attn.v_proj.bias"),
748            ) {
749                (Some(a), Some(b), Some(c)) => Some((a, b, c)),
750                _ => None,
751            };
752            Ok(AttnKind::Full {
753                wq,
754                wk: t("self_attn.k_proj.weight")?,
755                wv: t("self_attn.v_proj.weight")?,
756                wo: t("self_attn.o_proj.weight")?,
757                q_norm: n("self_attn.q_norm.weight"),
758                k_norm: n("self_attn.k_norm.weight"),
759                output_gate,
760                softplus_gate,
761                bias,
762            })
763        };
764
765        let load_linear_attn = |prefix: &str| -> Result<AttnKind, CmfError> {
766            if gdn_cfg.is_some() {
767                // Faithful vendor operator: tensor names 1:1 with the source.
768                let t = |suffix: &str| {
769                    load_matrix(
770                        model,
771                        &format!("{prefix}linear_attn.{suffix}"),
772                        force_f32,
773                        ov,
774                    )
775                };
776                let f = |suffix: &str| {
777                    load_f32(model, &format!("{prefix}linear_attn.{suffix}"), ov).map_err(err)
778                };
779                return Ok(AttnKind::LinearGdn(GdnWeights {
780                    in_proj_qkv: t("in_proj_qkv.weight")?,
781                    in_proj_z: t("in_proj_z.weight")?,
782                    in_proj_a: t("in_proj_a.weight")?,
783                    in_proj_b: t("in_proj_b.weight")?,
784                    conv1d: f("conv1d.weight")?,
785                    a_log: f("A_log")?,
786                    dt_bias: f("dt_bias")?,
787                    norm: f("norm.weight")?,
788                    out_proj: t("out_proj.weight")?,
789                }));
790            }
791            let t = |suffix: &str| {
792                load_matrix(model, &format!("{prefix}vmf_attn.{suffix}"), force_f32, ov)
793            };
794            let a_log = load_f32(model, &format!("{prefix}vmf_attn.A_log"), ov).map_err(err)?;
795            // Selective-write gate κ (hybrid_k core): optional by tensor
796            // presence — files without it run the classic phase kernel
797            // bit-identically.
798            let k_gate = if model
799                .tensor(&format!("{prefix}vmf_attn.k_gate.weight"))
800                .is_some()
801            {
802                Some((
803                    t("k_gate.weight")?,
804                    load_f32(model, &format!("{prefix}vmf_attn.k_gate.bias"), ov).map_err(err)?,
805                ))
806            } else {
807                None
808            };
809            Ok(AttnKind::Linear(VmfPhaseWeights {
810                thq: t("thq.weight")?,
811                thk: t("thk.weight")?,
812                v_proj: t("v_proj.weight")?,
813                out_proj: t("out_proj.weight")?,
814                decay: a_log.iter().map(|&a| (-(a as f64).exp()).exp()).collect(),
815                k_gate,
816            }))
817        };
818
819        // LFM2 short-conv mixer: in_proj [3·hidden, hidden], a depthwise
820        // conv (stored f16 as `[hidden, 1, kernel]` → flattened taps), and
821        // out_proj [hidden, hidden]. Names canonicalized at convert time.
822        let load_short_conv = |prefix: &str| -> Result<AttnKind, CmfError> {
823            let t = |suffix: &str| {
824                load_matrix(
825                    model,
826                    &format!("{prefix}short_conv.{suffix}"),
827                    force_f32,
828                    ov,
829                )
830            };
831            Ok(AttnKind::ShortConv(ShortConvWeights {
832                in_proj: t("in_proj.weight")?,
833                conv: load_f32(model, &format!("{prefix}short_conv.conv.weight"), ov)
834                    .map_err(err)?,
835                out_proj: t("out_proj.weight")?,
836            }))
837        };
838
839        // KDA layer (Kimi Linear / Kimi-K3): faithful vendor tensors under
840        // the `kda_attn.` canonical prefix. The output gate is full-rank
841        // (g_proj, K3) or low-rank (g_a/g_b, Kimi-Linear-48B) by presence.
842        let load_kda = |prefix: &str| -> Result<AttnKind, CmfError> {
843            let t = |suffix: &str| {
844                load_matrix(model, &format!("{prefix}kda_attn.{suffix}"), force_f32, ov)
845            };
846            let f = |suffix: &str| {
847                load_f32(model, &format!("{prefix}kda_attn.{suffix}"), ov).map_err(err)
848            };
849            let gate = if model
850                .tensor(&format!("{prefix}kda_attn.g_proj.weight"))
851                .is_some()
852            {
853                crate::linear_core::KdaOutGate::Full(t("g_proj.weight")?)
854            } else {
855                crate::linear_core::KdaOutGate::LowRank(
856                    t("g_a_proj.weight")?,
857                    t("g_b_proj.weight")?,
858                )
859            };
860            Ok(AttnKind::Kda(Box::new(crate::linear_core::KdaWeights {
861                q_proj: t("q_proj.weight")?,
862                k_proj: t("k_proj.weight")?,
863                v_proj: t("v_proj.weight")?,
864                conv_q: f("q_conv1d.weight")?,
865                conv_k: f("k_conv1d.weight")?,
866                conv_v: f("v_conv1d.weight")?,
867                f_a: t("f_a_proj.weight")?,
868                f_b: t("f_b_proj.weight")?,
869                dt_bias: f("dt_bias")?,
870                a_log: f("A_log")?,
871                b_proj: t("b_proj.weight")?,
872                gate,
873                o_norm: f("o_norm.weight")?,
874                o_proj: t("o_proj.weight")?,
875                gate_lower_bound: arch.kda_gate_lower_bound.map(|v| v as f32),
876            })))
877        };
878
879        let mut layers = Vec::with_capacity(arch.num_layers);
880        for li in 0..arch.num_layers {
881            let prefix = format!("model.layers.{li}.");
882            let attn = match arch.layer_types.get(li) {
883                Some(LayerType::LinearAttention) => load_linear_attn(&prefix)?,
884                Some(LayerType::Kda) => load_kda(&prefix)?,
885                Some(LayerType::ShortConv) => load_short_conv(&prefix)?,
886                _ => load_full_attn(&prefix, Some(li))?,
887            };
888            // Gemma-2/3 sandwich: `pre_feedforward_layernorm` present →
889            // it is the pre-FFN norm, and post_attention/post_feedforward
890            // норms apply to the branch OUTPUTS before their residuals.
891            let pre_ffn = format!("{prefix}pre_feedforward_layernorm.weight");
892            let sandwich = model.tensor(&pre_ffn).is_some();
893            layers.push(LayerWeights {
894                input_norm: load_f32(model, &format!("{prefix}input_layernorm.weight"), ov)
895                    .map_err(err)?,
896                post_norm: if sandwich {
897                    load_f32(model, &pre_ffn, ov).map_err(err)?
898                } else {
899                    load_f32(
900                        model,
901                        &format!("{prefix}post_attention_layernorm.weight"),
902                        ov,
903                    )
904                    .map_err(err)?
905                },
906                attn_out_norm: if sandwich {
907                    Some(
908                        load_f32(
909                            model,
910                            &format!("{prefix}post_attention_layernorm.weight"),
911                            ov,
912                        )
913                        .map_err(err)?,
914                    )
915                } else {
916                    None
917                },
918                ffn_out_norm: if sandwich {
919                    Some(
920                        load_f32(
921                            model,
922                            &format!("{prefix}post_feedforward_layernorm.weight"),
923                            ov,
924                        )
925                        .map_err(err)?,
926                    )
927                } else {
928                    None
929                },
930                // Gemma-4: learned scalar multiplying the layer output.
931                layer_scale: model
932                    .tensor(&format!("{prefix}layer_scalar"))
933                    .and_then(|_| {
934                        load_f32(model, &format!("{prefix}layer_scalar"), ov)
935                            .ok()
936                            .and_then(|v| v.first().copied())
937                    }),
938                // FFN always quantized — masks run sparse on quant bytes.
939                ffn: build_layer_ffn(model, &arch, li, false, ov)?,
940                attn,
941            });
942        }
943
944        // ── MTP head (optional, spec §2.1) ──
945        let mtp = if let Some(cfg) = &arch.mtp {
946            if cfg.num_layers != 1 {
947                return Err(CmfError::Parse(format!(
948                    "MTP with {} blocks not supported yet (only 1)",
949                    cfg.num_layers
950                )));
951            }
952            let p = "model.mtp.";
953            let attn = load_full_attn("model.mtp.layers.0.", None)?;
954            Some(MtpModule {
955                enorm: load_f32(model, &format!("{p}enorm.weight"), ov).map_err(err)?,
956                hnorm: load_f32(model, &format!("{p}hnorm.weight"), ov).map_err(err)?,
957                eh_proj: load_matrix(model, &format!("{p}eh_proj.weight"), false, ov)?,
958                layer: LayerWeights {
959                    attn_out_norm: None,
960                    ffn_out_norm: None,
961                    layer_scale: None,
962                    input_norm: load_f32(model, &format!("{p}layers.0.input_layernorm.weight"), ov)
963                        .map_err(err)?,
964                    post_norm: load_f32(
965                        model,
966                        &format!("{p}layers.0.post_attention_layernorm.weight"),
967                        ov,
968                    )
969                    .map_err(err)?,
970                    ffn: FfnKind::Dense(DenseFfn {
971                        gate_proj: load_matrix(
972                            model,
973                            &format!("{p}layers.0.mlp.gate_proj.weight"),
974                            false,
975                            ov,
976                        )?,
977                        up_proj: load_matrix(
978                            model,
979                            &format!("{p}layers.0.mlp.up_proj.weight"),
980                            false,
981                            ov,
982                        )?,
983                        down_proj: load_matrix(
984                            model,
985                            &format!("{p}layers.0.mlp.down_proj.weight"),
986                            false,
987                            ov,
988                        )?,
989                        act: crate::pipeline::Act::from_arch_full(&arch),
990                    }),
991                    attn,
992                },
993                final_norm: load_f32(model, &format!("{p}norm.weight"), ov).map_err(err)?,
994                kv: LayerKvCache::new(arch.num_kv_heads, arch.head_dim),
995            })
996        } else {
997            None
998        };
999
1000        tracing::info!(
1001            "Pipeline loaded: {} | {}L ({} linear) | {:.2}B params | storage: {} | MTP: {}",
1002            arch.arch_name,
1003            arch.num_layers,
1004            arch.layer_types
1005                .iter()
1006                .filter(|t| matches!(t, LayerType::LinearAttention))
1007                .count(),
1008            model.total_param_count() as f64 / 1e9,
1009            if force_f32 {
1010                "f32 (masked)"
1011            } else {
1012                "quantized mmap"
1013            },
1014            if mtp.is_some() { "yes" } else { "no" }
1015        );
1016
1017        // KV window: the descriptor's max, capped for dev-box safety;
1018        // CMF_MAX_SEQ overrides the cap (long-context runs).
1019        let cap = std::env::var("CMF_MAX_SEQ")
1020            .ok()
1021            .and_then(|v| v.parse::<usize>().ok())
1022            .unwrap_or(8192);
1023        let max_seq_len = arch.max_position_embeddings.min(cap);
1024
1025        // Looped Transformer: total virtual layers = physical × num_loops.
1026        let total_layers = arch.num_layers * arch.num_loops;
1027
1028        let mut pipeline = Pipeline::new(
1029            tokenizer,
1030            PipelineWeights {
1031                embed_tokens,
1032                layers,
1033                lm_head,
1034                final_norm,
1035            },
1036            arch.hidden_size,
1037            arch.intermediate_size,
1038            arch.num_attention_heads,
1039            arch.num_kv_heads,
1040            arch.head_dim,
1041            total_layers,
1042            arch.num_layers, // physical layers in weights
1043            arch.loop_final_norm,
1044            arch.vocab_size,
1045            arch.rms_norm_eps,
1046            arch.rope_theta as f32,
1047            arch.norm_style,
1048            max_seq_len,
1049            sampler_config,
1050        );
1051        let rotary = ((arch.head_dim as f32 * arch.partial_rotary_factor) as usize).max(2);
1052        pipeline.set_rotary(rotary, arch.rope_theta as f32);
1053        pipeline.attention_heads_per_layer = arch.attention_heads_per_layer.clone();
1054        if let Some(yarn) = &arch.yarn {
1055            pipeline.inv_freq = std::sync::Arc::new(crate::attention::yarn_inv_freq(
1056                rotary,
1057                arch.rope_theta as f32,
1058                yarn.factor,
1059                yarn.original_max_position_embeddings,
1060                yarn.beta_fast,
1061                yarn.beta_slow,
1062            ));
1063            pipeline.rope_scale = yarn.attention_factor;
1064        }
1065        // Gemma-family extras: embedding scale, attention-scale
1066        // override, and (Gemma-3) sliding-window layers with their own
1067        // local RoPE base.
1068        pipeline.embed_multiplier = arch.embed_multiplier;
1069        pipeline.logit_multiplier = arch.logit_multiplier;
1070        if let Some(qpas) = arch.query_pre_attn_scalar {
1071            pipeline.attn_scale = 1.0 / (qpas as f32).sqrt();
1072        }
1073        if let (Some(w), Some(p)) = (arch.sliding_window, arch.sliding_window_pattern) {
1074            pipeline.swa = Some((w, p));
1075            if let Some(base) = arch.rope_local_base_freq {
1076                pipeline.inv_freq_local = Some(std::sync::Arc::new(
1077                    crate::attention::rope_inv_freq(rotary, base as f32),
1078                ));
1079            }
1080        }
1081        let explicit_sliding: Vec<bool> = arch
1082            .layer_types
1083            .iter()
1084            .map(|t| matches!(t, cortiq_core::LayerType::SlidingAttention))
1085            .collect();
1086        if explicit_sliding.iter().any(|&v| v) {
1087            pipeline.sliding_layers = Some(explicit_sliding);
1088            if let Some(w) = arch.sliding_window {
1089                pipeline.swa = Some((w, usize::MAX));
1090            }
1091            let local_rotary = ((arch.head_dim as f32
1092                * arch
1093                    .local_partial_rotary_factor
1094                    .unwrap_or(arch.partial_rotary_factor))
1095                as usize)
1096                .max(2);
1097            pipeline.rotary_dim_local = Some(local_rotary);
1098            if let Some(base) = arch.rope_local_base_freq {
1099                pipeline.inv_freq_local = Some(std::sync::Arc::new(
1100                    crate::attention::rope_inv_freq(local_rotary, base as f32),
1101                ));
1102            }
1103        }
1104        // Gemma-4: global layers run their own geometry (MQA at
1105        // global_head_dim) with a proportional RoPE — the first
1106        // factor·head_dim dims rotate, the zero-padded tail is identity.
1107        if let (Some(ghd), Some(gkv)) = (arch.global_head_dim, arch.num_global_kv_heads) {
1108            pipeline.global_attn = Some((ghd, gkv));
1109            let prf = arch.global_partial_rotary_factor.unwrap_or(1.0);
1110            let half = ghd / 2;
1111            let ra = (((prf * ghd as f32) as usize) / 2).min(half);
1112            let mut f = vec![0.0f32; half];
1113            for (i, slot) in f.iter_mut().enumerate().take(ra) {
1114                *slot = 1.0 / (arch.rope_theta as f32).powf(2.0 * i as f32 / ghd as f32);
1115            }
1116            pipeline.inv_freq_global = Some(std::sync::Arc::new(f));
1117            // Re-shape the global layers' KV storage to their geometry.
1118            // An explicit layer_types map wins over the numeric pattern
1119            // (explicit tags set swa's pattern to usize::MAX, which
1120            // would otherwise leave every global cache mis-shaped).
1121            let global_at = |li: usize| -> bool {
1122                match &pipeline.sliding_layers {
1123                    Some(map) => !map.get(li).copied().unwrap_or(false),
1124                    None => pipeline
1125                        .swa
1126                        .map(|(_, p)| p > 0 && p != usize::MAX && (li + 1) % p == 0)
1127                        .unwrap_or(false),
1128                }
1129            };
1130            for li in 0..arch.num_layers {
1131                if global_at(li) {
1132                    pipeline.kv_cache.layers[li] = crate::kv_cache::LayerKvCache::new(gkv, ghd);
1133                }
1134            }
1135        }
1136        // MLA (DeepSeek-V2): the expand-to-MHA cache holds nh heads of
1137        // rope+nope dims; rotary covers the rope prefix.
1138        if let Some(mla) = arch.mla.as_ref() {
1139            let hd = mla.qk_rope_head_dim + mla.qk_nope_head_dim;
1140            pipeline.head_dim = hd;
1141            pipeline.num_kv_heads = arch.num_attention_heads;
1142            pipeline.rotary_dim = mla.qk_rope_head_dim;
1143            let half = mla.qk_rope_head_dim / 2;
1144            let mut f = vec![0.0f32; half];
1145            for (i, slot) in f.iter_mut().enumerate() {
1146                *slot = 1.0
1147                    / (arch.rope_theta as f32)
1148                        .powf(2.0 * i as f32 / mla.qk_rope_head_dim as f32);
1149            }
1150            pipeline.inv_freq = std::sync::Arc::new(f);
1151            for li in 0..arch.num_layers {
1152                pipeline.kv_cache.layers[li] =
1153                    crate::kv_cache::LayerKvCache::new(arch.num_attention_heads, hd);
1154            }
1155        }
1156        // Per-frequency rope divisors (MiniCPM3 longrope short_factor):
1157        // served at the native window with the trained per-dim factors.
1158        // Applied after every inv_freq build (plain, YaRN, MLA).
1159        if let Some(fac) = &arch.rope_freq_factors {
1160            let mut f = pipeline.inv_freq.as_ref().clone();
1161            for (i, v) in f.iter_mut().enumerate() {
1162                if let Some(&d) = fac.get(i) {
1163                    *v /= d as f32;
1164                }
1165            }
1166            pipeline.inv_freq = std::sync::Arc::new(f);
1167        }
1168        pipeline.attn_v_norm = arch.attn_v_norm;
1169        pipeline.final_softcap = arch.final_logit_softcapping.map(|c| c as f32);
1170        pipeline.attn_softcap = arch.attn_logit_softcapping.unwrap_or(0.0) as f32;
1171        pipeline.vmf_cfg = vmf_cfg;
1172        pipeline.gdn_cfg = gdn_cfg;
1173        pipeline.kda_cfg = kda_cfg;
1174        pipeline.short_conv_cfg = short_conv_cfg;
1175        pipeline.mtp = mtp;
1176        pipeline.install_dynamic_routing(model, false);
1177        // Record the load-time overlay so a later set_active_skill(None)
1178        // correctly reverts it (the union-diff assumes dyn_active mirrors
1179        // the live overlay). Blend loads have no single index to revert.
1180        match ov {
1181            Overlay::One(sid) => {
1182                pipeline.dyn_active = model.header.skills.iter().position(|s| &s.id == sid);
1183            }
1184            Overlay::Blend(_) => pipeline.dyn_blend_loaded = true,
1185            Overlay::None => {}
1186        }
1187        // B1: apply the measured confidence-calibration temperature, if the
1188        // file carries one (softmax(logits / T) for reported Born mass).
1189        if let Some(c) = &model.header.calibration {
1190            pipeline.set_calib_temp(c.temperature);
1191        }
1192        // O(1) Nyström attention (runtime-level, no format change):
1193        // env CMF_O1 decides; unset falls through to the converter hint
1194        // in header.provenance.o1_attn (`cortiq convert --o1`), and
1195        // CMF_O1=off force-disables even the hint. CLI flags override
1196        // later via set_o1().
1197        let o1 = match crate::nystrom::o1_from_env() {
1198            crate::nystrom::O1Env::Off => None,
1199            crate::nystrom::O1Env::On(cfg) => Some(cfg),
1200            crate::nystrom::O1Env::Unset => model
1201                .header
1202                .provenance
1203                .as_ref()
1204                .and_then(|p| p.get("o1_attn"))
1205                .and_then(crate::nystrom::O1Cfg::from_json),
1206        };
1207        if o1.is_some() {
1208            if pipeline.attn_softcap > 0.0 {
1209                return Err(CmfError::Parse(
1210                    "--o1 with attention-logit soft-capping (Gemma-2) is not supported: \
1211                     the streaming operator has no capped-score form"
1212                        .into(),
1213                ));
1214            }
1215            pipeline.set_o1(o1);
1216        }
1217        Ok(pipeline)
1218    }
1219
1220    /// Record per-skill dynamic-routing metadata: which FFN layers each
1221    /// skill actually replaces (derived from the tensors present, not
1222    /// the meta `layers` field), and whether the skill is eligible for
1223    /// cheap dynamic switching (FFN-only). Called once at load.
1224    pub(crate) fn install_dynamic_routing(&mut self, model: &Arc<CmfModel>, force_f32: bool) {
1225        self.model = Some(model.clone());
1226        self.dyn_force_f32 = force_f32;
1227        let mut per_skill = Vec::with_capacity(model.header.skills.len());
1228        for sk in &model.header.skills {
1229            let mut ffn_layers = std::collections::BTreeSet::new();
1230            let mut non_ffn = false;
1231            let prefix = format!("skill.{}.", sk.id);
1232            for t in model.skill_tensors(&sk.id) {
1233                let rel = &t.name[prefix.len()..]; // e.g. model.layers.20.mlp.down_proj.weight
1234                let toks: Vec<&str> = rel.split('.').collect();
1235                if toks.len() >= 5 && toks[0] == "model" && toks[1] == "layers" && toks[3] == "mlp"
1236                {
1237                    if let Ok(li) = toks[2].parse::<usize>() {
1238                        ffn_layers.insert(li);
1239                        continue;
1240                    }
1241                }
1242                non_ffn = true; // replaces attention / embed / lm_head
1243            }
1244            if non_ffn {
1245                tracing::warn!(
1246                    "skill '{}' replaces non-FFN tensors — excluded from dynamic \
1247                     routing (static overlay still works)",
1248                    sk.id
1249                );
1250                per_skill.push(None);
1251            } else {
1252                per_skill.push(Some(ffn_layers.into_iter().collect::<Vec<_>>()));
1253            }
1254        }
1255        self.dyn_skill_layers = per_skill;
1256    }
1257
1258    /// Switch the overlaid skill for subsequent forwards (dynamic
1259    /// routing). `idx` = index into model.header.skills; None = backbone.
1260    /// Rebuilds the FFN of the union of the old and new skill's touched
1261    /// layers with the new overlay — tensor-source indirection made
1262    /// dynamic. Cheap: Mapped tensors are re-resolved mmap pointers.
1263    /// Result is bit-identical to loading the pipeline with that skill.
1264    pub fn set_active_skill(&mut self, idx: Option<usize>) -> Result<(), CmfError> {
1265        if self.dyn_active == idx {
1266            return Ok(());
1267        }
1268        let model = self.model.clone().ok_or_else(|| {
1269            CmfError::Parse("dynamic routing needs a model-backed pipeline".into())
1270        })?;
1271        let mut union: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
1272        if let Some(old) = self.dyn_active {
1273            if let Some(Some(ls)) = self.dyn_skill_layers.get(old) {
1274                union.extend(ls.iter().copied());
1275            }
1276        }
1277        let new_id: Option<String> = match idx {
1278            Some(n) => match self.dyn_skill_layers.get(n) {
1279                Some(Some(ls)) => {
1280                    union.extend(ls.iter().copied());
1281                    Some(model.header.skills[n].id.clone())
1282                }
1283                _ => {
1284                    return Err(CmfError::Parse(format!(
1285                        "skill index {n} not dynamic-eligible"
1286                    )));
1287                }
1288            },
1289            None => None,
1290        };
1291        let ov = match &new_id {
1292            Some(s) => Overlay::One(s),
1293            None => Overlay::None,
1294        };
1295        let arch = model.arch();
1296        for li in union {
1297            self.weights.layers[li].ffn =
1298                build_layer_ffn(&model, arch, li, self.dyn_force_f32, &ov)?;
1299        }
1300        self.dyn_active = idx;
1301        Ok(())
1302    }
1303}