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::{GdnCfg, GdnWeights, VmfPhaseCfg, VmfPhaseWeights};
16use crate::pipeline::{
17    AttnKind, DenseFfn, FfnKind, LayerWeights, MoeFfn, MtpModule, Pipeline, PipelineWeights,
18};
19use crate::qtensor::QTensor;
20use crate::sampler::SamplerConfig;
21use crate::tokenizer::Tokenizer;
22use cortiq_core::quant::dequant_tensor;
23use cortiq_core::{CmfError, CmfModel, LayerType, ModelArch};
24use std::sync::Arc;
25
26/// Tensor source selector (spec §9): backbone, one skill's overlay, or
27/// a soft superposition of top-m skills (claim 14 working tensors).
28pub enum Overlay<'a> {
29    None,
30    One(&'a str),
31    /// (skill_id, weight); weights sum to 1 (softmax(−E/T) upstream).
32    Blend(&'a [(String, f32)]),
33}
34
35impl Overlay<'_> {
36    fn blend_touches(&self, model: &CmfModel, name: &str) -> bool {
37        match self {
38            Overlay::Blend(list) => list
39                .iter()
40                .any(|(sid, _)| model.tensor(&format!("skill.{sid}.{name}")).is_some()),
41            _ => false,
42        }
43    }
44}
45
46fn dequant_by_name(model: &CmfModel, name: &str) -> Result<Vec<f32>, String> {
47    let entry = model
48        .tensor(name)
49        .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
50    let mut out = vec![0.0f32; entry.n_elems()];
51    dequant_tensor(entry, model.entry_bytes(entry), &mut out)?;
52    Ok(out)
53}
54
55/// Weighted working tensor (claim 14): Σ wᵢ·Tᵢ, where Tᵢ is the
56/// skill's replacement when present, else the backbone tensor.
57fn blend_f32(model: &CmfModel, name: &str, list: &[(String, f32)]) -> Result<Vec<f32>, String> {
58    let mut acc: Option<Vec<f32>> = None;
59    for (sid, w) in list {
60        let sname = format!("skill.{sid}.{name}");
61        let src = if model.tensor(&sname).is_some() { &sname } else { name };
62        let t = dequant_by_name(model, src)?;
63        match &mut acc {
64            None => {
65                let mut t = t;
66                for v in t.iter_mut() {
67                    *v *= w;
68                }
69                acc = Some(t);
70            }
71            Some(a) => {
72                for (av, tv) in a.iter_mut().zip(&t) {
73                    *av += w * tv;
74                }
75            }
76        }
77    }
78    acc.ok_or_else(|| "empty blend".into())
79}
80
81/// Dequantize a tensor fully into f32 (norms, masked models).
82fn load_f32(model: &CmfModel, name: &str, ov: &Overlay) -> Result<Vec<f32>, String> {
83    if ov.blend_touches(model, name) {
84        if let Overlay::Blend(list) = ov {
85            return blend_f32(model, name, list);
86        }
87    }
88    let skill = match ov {
89        Overlay::One(s) => Some(*s),
90        _ => None,
91    };
92    let entry = model
93        .resolve_tensor(name, skill)
94        .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
95    let bytes = model.entry_bytes(entry);
96    let mut out = vec![0.0f32; entry.n_elems()];
97    dequant_tensor(entry, bytes, &mut out)?;
98    Ok(out)
99}
100
101/// Build one layer's FFN (dense or MoE) under a given overlay. Shared
102/// by the static loader AND dynamic per-token skill switching
103/// (`Pipeline::set_active_skill`): switching skills = rebuilding the
104/// FFN of the touched layers, cheap because Mapped tensors are just
105/// re-resolved mmap pointers (no dequant, no copy).
106pub(crate) fn build_layer_ffn(
107    model: &Arc<CmfModel>,
108    arch: &ModelArch,
109    li: usize,
110    force_f32: bool,
111    ov: &Overlay,
112) -> Result<FfnKind, CmfError> {
113    let prefix = format!("model.layers.{li}.");
114    let load_dense = |p: &str| -> Result<DenseFfn, CmfError> {
115        Ok(DenseFfn {
116            gate_proj: load_matrix(model, &format!("{p}gate_proj.weight"), force_f32, ov)?,
117            up_proj: load_matrix(model, &format!("{p}up_proj.weight"), force_f32, ov)?,
118            down_proj: load_matrix(model, &format!("{p}down_proj.weight"), force_f32, ov)?,
119        })
120    };
121    let router_name = format!("{prefix}mlp.gate.weight");
122    if model.tensor(&router_name).is_none() {
123        return Ok(FfnKind::Dense(load_dense(&format!("{prefix}mlp."))?));
124    }
125    let cfg = arch.moe.as_ref().ok_or_else(|| {
126        CmfError::Parse(format!("{router_name} present but header has no arch.moe block"))
127    })?;
128    let experts = (0..cfg.num_experts)
129        .map(|e| load_dense(&format!("{prefix}mlp.experts.{e}.")))
130        .collect::<Result<Vec<_>, _>>()?;
131    let shared = if model
132        .tensor(&format!("{prefix}mlp.shared_expert.gate_proj.weight"))
133        .is_some()
134    {
135        Some((
136            load_dense(&format!("{prefix}mlp.shared_expert."))?,
137            load_matrix(model, &format!("{prefix}mlp.shared_expert_gate.weight"), force_f32, ov)?,
138        ))
139    } else {
140        None
141    };
142    Ok(FfnKind::Moe(MoeFfn {
143        router: load_matrix(model, &router_name, force_f32, ov)?,
144        experts,
145        top_k: cfg.top_k,
146        norm_topk_prob: cfg.norm_topk_prob,
147        shared,
148        stats: std::cell::RefCell::new(Vec::new()),
149    }))
150}
151
152fn load_matrix(
153    model: &Arc<CmfModel>,
154    name: &str,
155    force_f32: bool,
156    ov: &Overlay,
157) -> Result<QTensor, CmfError> {
158    // Claim 14: a blended working tensor is materialized in f32 and
159    // held resident (the overlay-cache slot); single skills stay
160    // zero-copy pointers into the mmap.
161    if ov.blend_touches(model, name) {
162        if let Overlay::Blend(list) = ov {
163            let entry = model
164                .tensor(name)
165                .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
166            let data = blend_f32(model, name, list)
167                .map_err(|e| CmfError::Parse(format!("blend: {e}")))?;
168            return Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]));
169        }
170    }
171    let skill = match ov {
172        Overlay::One(s) => Some(*s),
173        _ => None,
174    };
175    // Tensor-source indirection (spec §9): the skill's replacement is
176    // read in place of the backbone tensor — either/or, never a sum.
177    let name: &str = &match skill {
178        Some(sid) if model.tensor(&format!("skill.{sid}.{name}")).is_some() => {
179            format!("skill.{sid}.{name}")
180        }
181        _ => name.to_string(),
182    };
183    let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
184    if force_f32 {
185        let entry = model
186            .tensor(name)
187            .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
188        if entry.shape.len() != 2 {
189            return Err(err(format!("'{name}' is not 2-D")));
190        }
191        let data = load_f32(model, name, &Overlay::None).map_err(err)?;
192        Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]))
193    } else {
194        QTensor::from_model(model, name).map_err(err)
195    }
196}
197
198impl Pipeline {
199    /// Build a runnable pipeline from an opened CMF model.
200    pub fn from_model(model: &Arc<CmfModel>, sampler_config: SamplerConfig) -> Result<Self, CmfError> {
201        Self::from_model_with_skill(model, sampler_config, None)
202    }
203
204    /// Same, with a skill overlaid (spec §9): every layer tensor is
205    /// resolved through tensor-source indirection — the skill's
206    /// full-shape replacement is read in place of the backbone tensor.
207    /// No per-skill model is ever assembled: Mapped tensors are
208    /// pointers into the one shared mmap.
209    pub fn from_model_with_skill(
210        model: &Arc<CmfModel>,
211        sampler_config: SamplerConfig,
212        skill: Option<&str>,
213    ) -> Result<Self, CmfError> {
214        match skill {
215            Some(s) => Self::from_model_with_overlay(model, sampler_config, &Overlay::One(s)),
216            None => Self::from_model_with_overlay(model, sampler_config, &Overlay::None),
217        }
218    }
219
220    /// Soft superposition (claim 14): working tensors accumulated from
221    /// the given (skill, weight) list — softmax(−E/T) upstream.
222    pub fn from_model_with_blend(
223        model: &Arc<CmfModel>,
224        sampler_config: SamplerConfig,
225        blend: &[(String, f32)],
226    ) -> Result<Self, CmfError> {
227        Self::from_model_with_overlay(model, sampler_config, &Overlay::Blend(blend))
228    }
229
230    fn from_model_with_overlay(
231        model: &Arc<CmfModel>,
232        sampler_config: SamplerConfig,
233        ov: &Overlay,
234    ) -> Result<Self, CmfError> {
235        let skill = match ov {
236            Overlay::One(s) => Some(*s),
237            _ => None,
238        };
239        if let Some(sid) = skill {
240            let known = model.header.skills.iter().any(|s| s.id == sid)
241                || model.skill_tensors(sid).next().is_some();
242            if !known {
243                return Err(CmfError::Parse(format!(
244                    "skill '{sid}' not in this container (header.skills: {:?})",
245                    model.header.skills.iter().map(|s| &s.id).collect::<Vec<_>>()
246                )));
247            }
248            tracing::info!(
249                "skill '{sid}': {} replacement tensors overlaid",
250                model.skill_tensors(sid).count()
251            );
252        }
253        let arch = model.arch().clone();
254        let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
255
256        // Masks × quantized mmap: only ATTENTION keeps f32 (the head-mask
257        // path needs f32 slices). FFN masks now run sparse directly on the
258        // quant bytes (sparse_ffn_quant), and embed/lm_head are never
259        // masked — so a masked model runs at quantized RSS, not the old
260        // whole-model-f32 blowup.
261        let masks_present = !model.masks.masks.is_empty();
262        let force_f32 = masks_present; // attention only (head masks)
263
264        // ── Tokenizer: embedded → sidecar → byte-level fallback ──
265        let mut tokenizer = if let Some(vocab_bytes) = &model.vocab {
266            Tokenizer::from_bytes(vocab_bytes)
267                .map_err(|e| CmfError::Parse(format!("embedded tokenizer: {e}")))?
268        } else {
269            let sidecar = model.path.with_file_name("tokenizer.json");
270            if sidecar.exists() {
271                Tokenizer::from_file(&sidecar)
272                    .map_err(|e| CmfError::Parse(format!("sidecar tokenizer: {e}")))?
273            } else {
274                tracing::warn!("no tokenizer in file or sidecar — using byte-level fallback");
275                Tokenizer::byte_level()
276            }
277        };
278        // Chat/eos bundle (spec §6.1): the FILE defines chat behavior.
279        if let Some(tc) = &model.header.tokenizer_config {
280            tokenizer.chat_template = tc.chat_template.clone();
281            tokenizer.extra_eos.extend(tc.eos_token_ids.iter().copied());
282            if tokenizer.bos_token_id.is_none() {
283                tokenizer.bos_token_id = tc.bos_token_id;
284            }
285            tracing::info!(
286                "chat bundle: template {} chars, {} stop ids",
287                tc.chat_template.as_deref().map(str::len).unwrap_or(0),
288                tc.eos_token_ids.len()
289            );
290        }
291
292        // ── Top-level weights (never masked → always quantized) ──
293        let embed_tokens = load_matrix(model, "model.embed_tokens.weight", false, ov)?;
294        let final_norm = load_f32(model, "model.norm.weight", ov).map_err(err)?;
295        let lm_head = if model.tensor("lm_head.weight").is_some() {
296            load_matrix(model, "lm_head.weight", false, ov)?
297        } else if arch.tie_word_embeddings {
298            // Tied: reuse the embedding matrix (re-open, cheap for Mapped).
299            load_matrix(model, "model.embed_tokens.weight", false, ov)?
300        } else {
301            return Err(CmfError::MissingTensor(
302                "lm_head.weight (and tie_word_embeddings is false)".into(),
303            ));
304        };
305
306        // ── Linear-core geometry (required if any linear layer exists) ──
307        let has_linear = arch
308            .layer_types
309            .iter()
310            .any(|t| matches!(t, LayerType::LinearAttention));
311        let mut vmf_cfg = None;
312        let mut gdn_cfg = None;
313        if has_linear {
314            let lc = arch.linear_core.as_ref().ok_or_else(|| {
315                CmfError::Parse(
316                    "model has LinearAttention layers but no arch.linear_core — \
317                     reconvert with the current converter"
318                        .into(),
319                )
320            })?;
321            let need = |v: Option<usize>, name: &str| {
322                v.ok_or_else(|| CmfError::Parse(format!("linear core needs arch.{name}")))
323            };
324            match lc.kind.as_str() {
325                "vmf_phase" => {
326                    vmf_cfg = Some(VmfPhaseCfg {
327                        num_heads: lc.num_heads,
328                        nphase: need(lc.nphase, "linear_core.nphase")?,
329                        value_head_dim: lc.value_head_dim,
330                        hidden_size: arch.hidden_size,
331                        // θ-mass (η′): default 0 (massless); CMF_PHASE_MASS
332                        // widens the phase kernel for folded-unhealed models.
333                        phase_mass: std::env::var("CMF_PHASE_MASS")
334                            .ok()
335                            .and_then(|v| v.parse().ok())
336                            .unwrap_or(0.0),
337                    });
338                }
339                "gated_delta_net" => {
340                    gdn_cfg = Some(GdnCfg {
341                        num_v_heads: lc.num_heads,
342                        num_k_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
343                        key_head_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
344                        value_head_dim: lc.value_head_dim,
345                        conv_kernel: need(
346                            arch.linear_conv_kernel_dim,
347                            "linear_conv_kernel_dim",
348                        )?,
349                        hidden_size: arch.hidden_size,
350                        rms_eps: arch.rms_norm_eps as f64,
351                    });
352                }
353                other => {
354                    return Err(CmfError::Parse(format!(
355                        "unknown linear core '{other}' (this runtime executes: \
356                         gated_delta_net, vmf_phase)"
357                    )));
358                }
359            }
360        }
361
362        // ── Layers ──
363        let load_full_attn = |prefix: &str| -> Result<AttnKind, CmfError> {
364            let t = |suffix: &str| load_matrix(model, &format!("{prefix}{suffix}"), force_f32, ov);
365            let n = |suffix: &str| -> Option<Vec<f32>> {
366                model
367                    .tensor(&format!("{prefix}{suffix}"))
368                    .and_then(|_| load_f32(model, &format!("{prefix}{suffix}"), ov).ok())
369            };
370            let wq = t("self_attn.q_proj.weight")?;
371            // Qwen3.5 output gate: q_proj rows = 2·nh·hd (per-head [q; gate]).
372            let output_gate = wq.rows() == 2 * arch.num_attention_heads * arch.head_dim;
373            // Qwen2-family projection biases (by tensor presence).
374            let bias = match (
375                n("self_attn.q_proj.bias"),
376                n("self_attn.k_proj.bias"),
377                n("self_attn.v_proj.bias"),
378            ) {
379                (Some(a), Some(b), Some(c)) => Some((a, b, c)),
380                _ => None,
381            };
382            Ok(AttnKind::Full {
383                wq,
384                wk: t("self_attn.k_proj.weight")?,
385                wv: t("self_attn.v_proj.weight")?,
386                wo: t("self_attn.o_proj.weight")?,
387                q_norm: n("self_attn.q_norm.weight"),
388                k_norm: n("self_attn.k_norm.weight"),
389                output_gate,
390                bias,
391            })
392        };
393
394        let load_linear_attn = |prefix: &str| -> Result<AttnKind, CmfError> {
395            if gdn_cfg.is_some() {
396                // Faithful vendor operator: tensor names 1:1 with the source.
397                let t = |suffix: &str| {
398                    load_matrix(model, &format!("{prefix}linear_attn.{suffix}"), force_f32, ov)
399                };
400                let f = |suffix: &str| {
401                    load_f32(model, &format!("{prefix}linear_attn.{suffix}"), ov).map_err(err)
402                };
403                return Ok(AttnKind::LinearGdn(GdnWeights {
404                    in_proj_qkv: t("in_proj_qkv.weight")?,
405                    in_proj_z: t("in_proj_z.weight")?,
406                    in_proj_a: t("in_proj_a.weight")?,
407                    in_proj_b: t("in_proj_b.weight")?,
408                    conv1d: f("conv1d.weight")?,
409                    a_log: f("A_log")?,
410                    dt_bias: f("dt_bias")?,
411                    norm: f("norm.weight")?,
412                    out_proj: t("out_proj.weight")?,
413                }));
414            }
415            let t = |suffix: &str| load_matrix(model, &format!("{prefix}vmf_attn.{suffix}"), force_f32, ov);
416            let a_log = load_f32(model, &format!("{prefix}vmf_attn.A_log"), ov).map_err(err)?;
417            Ok(AttnKind::Linear(VmfPhaseWeights {
418                thq: t("thq.weight")?,
419                thk: t("thk.weight")?,
420                v_proj: t("v_proj.weight")?,
421                out_proj: t("out_proj.weight")?,
422                decay: a_log
423                    .iter()
424                    .map(|&a| (-(a as f64).exp()).exp())
425                    .collect(),
426            }))
427        };
428
429        let mut layers = Vec::with_capacity(arch.num_layers);
430        for li in 0..arch.num_layers {
431            let prefix = format!("model.layers.{li}.");
432            let attn = match arch.layer_types.get(li) {
433                Some(LayerType::LinearAttention) => load_linear_attn(&prefix)?,
434                _ => load_full_attn(&prefix)?,
435            };
436            layers.push(LayerWeights {
437                input_norm: load_f32(model, &format!("{prefix}input_layernorm.weight"), ov).map_err(err)?,
438                post_norm: load_f32(model, &format!("{prefix}post_attention_layernorm.weight"), ov)
439                    .map_err(err)?,
440                // FFN always quantized — masks run sparse on quant bytes.
441                ffn: build_layer_ffn(model, &arch, li, false, ov)?,
442                attn,
443            });
444        }
445
446        // ── MTP head (optional, spec §2.1) ──
447        let mtp = if let Some(cfg) = &arch.mtp {
448            if cfg.num_layers != 1 {
449                return Err(CmfError::Parse(format!(
450                    "MTP with {} blocks not supported yet (only 1)",
451                    cfg.num_layers
452                )));
453            }
454            let p = "model.mtp.";
455            let attn = load_full_attn("model.mtp.layers.0.")?;
456            Some(MtpModule {
457                enorm: load_f32(model, &format!("{p}enorm.weight"), ov).map_err(err)?,
458                hnorm: load_f32(model, &format!("{p}hnorm.weight"), ov).map_err(err)?,
459                eh_proj: load_matrix(model, &format!("{p}eh_proj.weight"), false, ov)?,
460                layer: LayerWeights {
461                    input_norm: load_f32(model, &format!("{p}layers.0.input_layernorm.weight"), ov)
462                        .map_err(err)?,
463                    post_norm: load_f32(
464                        model,
465                        &format!("{p}layers.0.post_attention_layernorm.weight"),
466                        ov,
467                    )
468                    .map_err(err)?,
469                    ffn: FfnKind::Dense(DenseFfn {
470                        gate_proj: load_matrix(model, &format!("{p}layers.0.mlp.gate_proj.weight"), false, ov)?,
471                        up_proj: load_matrix(model, &format!("{p}layers.0.mlp.up_proj.weight"), false, ov)?,
472                        down_proj: load_matrix(model, &format!("{p}layers.0.mlp.down_proj.weight"), false, ov)?,
473                    }),
474                    attn,
475                },
476                final_norm: load_f32(model, &format!("{p}norm.weight"), ov).map_err(err)?,
477                kv: LayerKvCache::new(arch.num_kv_heads, arch.head_dim),
478            })
479        } else {
480            None
481        };
482
483        tracing::info!(
484            "Pipeline loaded: {} | {}L ({} linear) | {:.2}B params | storage: {} | MTP: {}",
485            arch.arch_name,
486            arch.num_layers,
487            arch.layer_types
488                .iter()
489                .filter(|t| matches!(t, LayerType::LinearAttention))
490                .count(),
491            model.total_param_count() as f64 / 1e9,
492            if force_f32 { "f32 (masked)" } else { "quantized mmap" },
493            if mtp.is_some() { "yes" } else { "no" }
494        );
495
496        // KV window: the descriptor's max, capped for dev-box safety;
497        // CMF_MAX_SEQ overrides the cap (long-context runs).
498        let cap = std::env::var("CMF_MAX_SEQ")
499            .ok()
500            .and_then(|v| v.parse::<usize>().ok())
501            .unwrap_or(8192);
502        let max_seq_len = arch.max_position_embeddings.min(cap);
503
504        let mut pipeline = Pipeline::new(
505            tokenizer,
506            PipelineWeights {
507                embed_tokens,
508                layers,
509                lm_head,
510                final_norm,
511            },
512            arch.hidden_size,
513            arch.intermediate_size,
514            arch.num_attention_heads,
515            arch.num_kv_heads,
516            arch.head_dim,
517            arch.num_layers,
518            arch.vocab_size,
519            arch.rms_norm_eps,
520            arch.rope_theta as f32,
521            arch.norm_style,
522            max_seq_len,
523            sampler_config,
524        );
525        let rotary = ((arch.head_dim as f32 * arch.partial_rotary_factor) as usize).max(2);
526        pipeline.set_rotary(rotary, arch.rope_theta as f32);
527        pipeline.vmf_cfg = vmf_cfg;
528        pipeline.gdn_cfg = gdn_cfg;
529        pipeline.mtp = mtp;
530        pipeline.install_dynamic_routing(model, false);
531        // Record the load-time overlay so a later set_active_skill(None)
532        // correctly reverts it (the union-diff assumes dyn_active mirrors
533        // the live overlay). Blend loads have no single index to revert.
534        match ov {
535            Overlay::One(sid) => {
536                pipeline.dyn_active = model.header.skills.iter().position(|s| &s.id == sid);
537            }
538            Overlay::Blend(_) => pipeline.dyn_blend_loaded = true,
539            Overlay::None => {}
540        }
541        // B1: apply the measured confidence-calibration temperature, if the
542        // file carries one (softmax(logits / T) for reported Born mass).
543        if let Some(c) = &model.header.calibration {
544            pipeline.set_calib_temp(c.temperature);
545        }
546        Ok(pipeline)
547    }
548
549    /// Record per-skill dynamic-routing metadata: which FFN layers each
550    /// skill actually replaces (derived from the tensors present, not
551    /// the meta `layers` field), and whether the skill is eligible for
552    /// cheap dynamic switching (FFN-only). Called once at load.
553    pub(crate) fn install_dynamic_routing(
554        &mut self,
555        model: &Arc<CmfModel>,
556        force_f32: bool,
557    ) {
558        self.model = Some(model.clone());
559        self.dyn_force_f32 = force_f32;
560        let mut per_skill = Vec::with_capacity(model.header.skills.len());
561        for sk in &model.header.skills {
562            let mut ffn_layers = std::collections::BTreeSet::new();
563            let mut non_ffn = false;
564            let prefix = format!("skill.{}.", sk.id);
565            for t in model.skill_tensors(&sk.id) {
566                let rel = &t.name[prefix.len()..]; // e.g. model.layers.20.mlp.down_proj.weight
567                let toks: Vec<&str> = rel.split('.').collect();
568                if toks.len() >= 5
569                    && toks[0] == "model"
570                    && toks[1] == "layers"
571                    && toks[3] == "mlp"
572                {
573                    if let Ok(li) = toks[2].parse::<usize>() {
574                        ffn_layers.insert(li);
575                        continue;
576                    }
577                }
578                non_ffn = true; // replaces attention / embed / lm_head
579            }
580            if non_ffn {
581                tracing::warn!(
582                    "skill '{}' replaces non-FFN tensors — excluded from dynamic \
583                     routing (static overlay still works)",
584                    sk.id
585                );
586                per_skill.push(None);
587            } else {
588                per_skill.push(Some(ffn_layers.into_iter().collect::<Vec<_>>()));
589            }
590        }
591        self.dyn_skill_layers = per_skill;
592    }
593
594    /// Switch the overlaid skill for subsequent forwards (dynamic
595    /// routing). `idx` = index into model.header.skills; None = backbone.
596    /// Rebuilds the FFN of the union of the old and new skill's touched
597    /// layers with the new overlay — tensor-source indirection made
598    /// dynamic. Cheap: Mapped tensors are re-resolved mmap pointers.
599    /// Result is bit-identical to loading the pipeline with that skill.
600    pub fn set_active_skill(&mut self, idx: Option<usize>) -> Result<(), CmfError> {
601        if self.dyn_active == idx {
602            return Ok(());
603        }
604        let model = self
605            .model
606            .clone()
607            .ok_or_else(|| CmfError::Parse("dynamic routing needs a model-backed pipeline".into()))?;
608        let mut union: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
609        if let Some(old) = self.dyn_active {
610            if let Some(Some(ls)) = self.dyn_skill_layers.get(old) {
611                union.extend(ls.iter().copied());
612            }
613        }
614        let new_id: Option<String> = match idx {
615            Some(n) => match self.dyn_skill_layers.get(n) {
616                Some(Some(ls)) => {
617                    union.extend(ls.iter().copied());
618                    Some(model.header.skills[n].id.clone())
619                }
620                _ => {
621                    return Err(CmfError::Parse(format!(
622                        "skill index {n} not dynamic-eligible"
623                    )))
624                }
625            },
626            None => None,
627        };
628        let ov = match &new_id {
629            Some(s) => Overlay::One(s),
630            None => Overlay::None,
631        };
632        let arch = model.arch();
633        for li in union {
634            self.weights.layers[li].ffn =
635                build_layer_ffn(&model, arch, li, self.dyn_force_f32, &ov)?;
636        }
637        self.dyn_active = idx;
638        Ok(())
639    }
640}