cortiq_engine/

loader.rs

//! Weight loader: CMF tensor directory → Pipeline.
//!
//! Storage rule: models WITH task masks are dequantized to f32 (masked
//! execution needs f32 row access; skill files are small by design).
//! Models without masks keep quantized matrices zero-copy from the mmap
//! (`QTensor::Mapped`) — this is what lets a 15B file run in a few GB
//! of RSS instead of 60 GB of f32.
//!
//! Layer kinds come from `arch.layer_types`: FullAttention loads
//! `self_attn.*` (with auto-detected Qwen3.5 extras: per-head qk-norm by
//! tensor presence, output gate by q_proj row count); LinearAttention
//! loads the canonical core `vmf_attn.*` (folded at convert time).

use crate::kv_cache::LayerKvCache;
use crate::linear_core::{
    GdnCfg, GdnWeights, ShortConvCfg, ShortConvWeights, VmfPhaseCfg, VmfPhaseWeights,
};
use crate::pipeline::{
    AttnKind, DenseFfn, FfnKind, LayerWeights, MoeFfn, MtpModule, Pipeline, PipelineWeights,
};
use crate::qtensor::QTensor;
use crate::sampler::SamplerConfig;
use crate::tokenizer::Tokenizer;
use cortiq_core::quant::dequant_tensor;
use cortiq_core::{CmfError, CmfModel, LayerType, ModelArch};
use std::sync::Arc;

/// Tensor source selector (spec §9): backbone, one skill's overlay, or
/// a soft superposition of top-m skills (claim 14 working tensors).
pub enum Overlay<'a> {
    None,
    One(&'a str),
    /// (skill_id, weight); weights sum to 1 (softmax(−E/T) upstream).
    Blend(&'a [(String, f32)]),
}

impl Overlay<'_> {
    fn blend_touches(&self, model: &CmfModel, name: &str) -> bool {
        match self {
            Overlay::Blend(list) => list
                .iter()
                .any(|(sid, _)| model.tensor(&format!("skill.{sid}.{name}")).is_some()),
            _ => false,
        }
    }
}

fn dequant_by_name(model: &CmfModel, name: &str) -> Result<Vec<f32>, String> {
    let entry = model
        .tensor(name)
        .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
    let mut out = vec![0.0f32; entry.n_elems()];
    dequant_tensor(entry, model.entry_bytes(entry), &mut out)?;
    Ok(out)
}

/// Weighted working tensor (claim 14): Σ wᵢ·Tᵢ, where Tᵢ is the
/// skill's replacement when present, else the backbone tensor.
fn blend_f32(model: &CmfModel, name: &str, list: &[(String, f32)]) -> Result<Vec<f32>, String> {
    let mut acc: Option<Vec<f32>> = None;
    for (sid, w) in list {
        let sname = format!("skill.{sid}.{name}");
        let src = if model.tensor(&sname).is_some() {
            &sname
        } else {
            name
        };
        let t = dequant_by_name(model, src)?;
        match &mut acc {
            None => {
                let mut t = t;
                for v in t.iter_mut() {
                    *v *= w;
                }
                acc = Some(t);
            }
            Some(a) => {
                for (av, tv) in a.iter_mut().zip(&t) {
                    *av += w * tv;
                }
            }
        }
    }
    acc.ok_or_else(|| "empty blend".into())
}

/// Dequantize a tensor fully into f32 (norms, masked models).
fn load_f32(model: &CmfModel, name: &str, ov: &Overlay) -> Result<Vec<f32>, String> {
    if ov.blend_touches(model, name) {
        if let Overlay::Blend(list) = ov {
            return blend_f32(model, name, list);
        }
    }
    let skill = match ov {
        Overlay::One(s) => Some(*s),
        _ => None,
    };
    let entry = model
        .resolve_tensor(name, skill)
        .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
    let bytes = model.entry_bytes(entry);
    let mut out = vec![0.0f32; entry.n_elems()];
    dequant_tensor(entry, bytes, &mut out)?;
    Ok(out)
}

/// Build one layer's FFN (dense or MoE) under a given overlay. Shared
/// by the static loader AND dynamic per-token skill switching
/// (`Pipeline::set_active_skill`): switching skills = rebuilding the
/// FFN of the touched layers, cheap because Mapped tensors are just
/// re-resolved mmap pointers (no dequant, no copy).
pub(crate) fn build_layer_ffn(
    model: &Arc<CmfModel>,
    arch: &ModelArch,
    li: usize,
    force_f32: bool,
    ov: &Overlay,
) -> Result<FfnKind, CmfError> {
    let prefix = format!("model.layers.{li}.");
    let load_dense = |p: &str| -> Result<DenseFfn, CmfError> {
        let gate_proj = load_matrix(model, &format!("{p}gate_proj.weight"), force_f32, ov)?;
        let up_proj = load_matrix(model, &format!("{p}up_proj.weight"), force_f32, ov)?;
        let down_proj = load_matrix(model, &format!("{p}down_proj.weight"), force_f32, ov)?;
        // FFN triple invariant (holds for dense and each MoE expert;
        // enforced loudly so a malformed defrag/repack — spec §11 — fails
        // at load instead of silently mis-computing). inter' is per-layer.
        let inter = gate_proj.rows();
        if up_proj.rows() != inter || down_proj.cols() != inter {
            return Err(CmfError::Parse(format!(
                "{p}: FFN dims disagree (gate.rows={inter}, up.rows={}, \
                 down.cols={}); all three must equal inter'",
                up_proj.rows(),
                down_proj.cols()
            )));
        }
        if down_proj.rows() != arch.hidden_size {
            return Err(CmfError::Parse(format!(
                "{p}: down_proj.rows={} != hidden_size={}",
                down_proj.rows(),
                arch.hidden_size
            )));
        }
        Ok(DenseFfn {
            gate_proj,
            up_proj,
            down_proj,
            act: crate::pipeline::Act::from_arch(&arch.hidden_act),
        })
    };
    let router_name = format!("{prefix}mlp.gate.weight");
    if model.tensor(&router_name).is_none() {
        return Ok(FfnKind::Dense(load_dense(&format!("{prefix}mlp."))?));
    }
    let cfg = arch.moe.as_ref().ok_or_else(|| {
        CmfError::Parse(format!(
            "{router_name} present but header has no arch.moe block"
        ))
    })?;
    // Experts enumerate by TENSOR PRESENCE up to the header count — a
    // moe-defrag'd specialist keeps a per-layer contiguous prefix of
    // renumbered experts (fewer than arch.moe.num_experts), with the
    // router rows sliced to match.
    let mut experts = Vec::new();
    for e in 0..cfg.num_experts {
        if model
            .tensor(&format!("{prefix}mlp.experts.{e}.gate_proj.weight"))
            .is_none()
        {
            break;
        }
        experts.push(load_dense(&format!("{prefix}mlp.experts.{e}."))?);
    }
    if experts.is_empty() {
        return Err(CmfError::Parse(format!(
            "{prefix}: router present but no expert tensors"
        )));
    }
    let shared = if model
        .tensor(&format!("{prefix}mlp.shared_expert.gate_proj.weight"))
        .is_some()
    {
        let gate_name = format!("{prefix}mlp.shared_expert_gate.weight");
        Some((
            load_dense(&format!("{prefix}mlp.shared_expert."))?,
            if model.tensor(&gate_name).is_some() {
                Some(load_matrix(model, &gate_name, force_f32, ov)?)
            } else {
                None
            },
        ))
    } else {
        None
    };
    // LFM2-MoE selection bias (`mlp.expert_bias`): present iff the model
    // routes with a bias; loaded by tensor presence.
    let bias_name = format!("{prefix}mlp.expert_bias");
    let expert_bias = if model.tensor(&bias_name).is_some() {
        Some(load_f32(model, &bias_name, ov).map_err(CmfError::Parse)?)
    } else {
        None
    };
    // CMF_MOE_TOPK=N (opt-in): route to fewer experts than the header
    // asks. MoE decode is memory-bound — every selected expert streams
    // its three matrices per token — so halving k halves that traffic;
    // the renormalized top-k keeps the mixture a proper average.
    // Quality is the experiment — measure ppl before trusting.
    let top_k = std::env::var("CMF_MOE_TOPK")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .filter(|&k| k >= 1 && k <= cfg.top_k)
        .inspect(|k| tracing::info!("MoE top_k override: {} (header {})", k, cfg.top_k))
        .unwrap_or(cfg.top_k);
    // CMF_MOE_TAU=0.x (opt-in): adaptive routing — see MoeFfn::route_tau.
    let route_tau = std::env::var("CMF_MOE_TAU")
        .ok()
        .and_then(|v| v.parse::<f32>().ok())
        .filter(|&t| t > 0.0 && t < 1.0)
        .inspect(|t| tracing::info!("MoE adaptive routing: tau {t}"));
    let mask = moe_task_mask(&prefix, experts.len());
    let router = load_matrix(model, &router_name, force_f32, ov)?;
    if router.rows() != experts.len() {
        return Err(CmfError::Parse(format!(
            "{router_name}: {} rows != {} experts",
            router.rows(),
            experts.len()
        )));
    }
    let top_k = top_k.min(experts.len());
    Ok(FfnKind::Moe(MoeFfn {
        router,
        experts,
        top_k,
        route_tau,
        norm_topk_prob: cfg.norm_topk_prob,
        router_sigmoid: cfg.router_sigmoid,
        expert_bias,
        routed_scaling: cfg.routed_scaling_factor.unwrap_or(1.0),
        shared,
        stats: std::cell::RefCell::new(Vec::new()),
        mask,
    }))
}

/// Task mask over routed experts (opt-in, experimental): DTG-MA applied
/// to MoE. `CMF_MOE_MASK=<stats.json>` points at a claim-12 B-field dump
/// (`CMF_MOE_STATS` output — per-layer expert-selection counts from a
/// task-representative run); `CMF_MOE_MASK_COVER` (default 0.9) keeps,
/// per layer, the smallest top set of experts reaching that fraction of
/// the recorded routing mass. Selection then happens over the allowed
/// set only (softmax renormalizes). Gate any real use on a ppl A/B.
fn moe_task_mask(prefix: &str, ne: usize) -> Option<Vec<bool>> {
    use std::sync::OnceLock;
    static CFG: OnceLock<Option<(std::collections::HashMap<usize, Vec<u64>>, f64)>> =
        OnceLock::new();
    let cfg = CFG.get_or_init(|| {
        let path = std::env::var("CMF_MOE_MASK").ok()?;
        let cover = std::env::var("CMF_MOE_MASK_COVER")
            .ok()
            .and_then(|v| v.parse::<f64>().ok())
            .filter(|&c| c > 0.0 && c <= 1.0)
            .unwrap_or(0.9);
        let text = std::fs::read_to_string(&path)
            .map_err(|e| tracing::warn!("CMF_MOE_MASK: cannot read {path}: {e}"))
            .ok()?;
        let map: std::collections::HashMap<String, Vec<u64>> =
            serde_json::from_str(&text)
                .map_err(|e| tracing::warn!("CMF_MOE_MASK: bad JSON in {path}: {e}"))
                .ok()?;
        tracing::info!("MoE task mask: {path}, cover {cover}");
        Some((
            map.into_iter()
                .filter_map(|(k, v)| Some((k.parse::<usize>().ok()?, v)))
                .collect(),
            cover,
        ))
    });
    let (stats, cover) = cfg.as_ref()?;
    // The layer index rides in the tensor prefix ("model.layers.N.").
    let li: usize = prefix
        .split("layers.")
        .nth(1)?
        .split('.')
        .next()?
        .parse()
        .ok()?;
    let counts = stats.get(&li)?;
    if counts.len() != ne {
        tracing::warn!("CMF_MOE_MASK: layer {li} has {} counts, model has {ne} experts — skipped", counts.len());
        return None;
    }
    let total: u64 = counts.iter().sum();
    if total == 0 {
        return None;
    }
    let mut order: Vec<usize> = (0..ne).collect();
    order.sort_unstable_by_key(|&e| std::cmp::Reverse(counts[e]));
    let mut mask = vec![false; ne];
    let mut acc = 0u64;
    let mut kept = 0usize;
    for &e in &order {
        mask[e] = true;
        acc += counts[e];
        kept += 1;
        if (acc as f64) >= cover * (total as f64) {
            break;
        }
    }
    tracing::info!("MoE task mask L{li}: {kept}/{ne} experts for {:.0}% mass", cover * 100.0);
    Some(mask)
}

fn load_matrix(
    model: &Arc<CmfModel>,
    name: &str,
    force_f32: bool,
    ov: &Overlay,
) -> Result<QTensor, CmfError> {
    // Claim 14: a blended working tensor is materialized in f32 and
    // held resident (the overlay-cache slot); single skills stay
    // zero-copy pointers into the mmap.
    if ov.blend_touches(model, name) {
        if let Overlay::Blend(list) = ov {
            let entry = model
                .tensor(name)
                .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
            let data =
                blend_f32(model, name, list).map_err(|e| CmfError::Parse(format!("blend: {e}")))?;
            return Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]));
        }
    }
    let skill = match ov {
        Overlay::One(s) => Some(*s),
        _ => None,
    };
    // Tensor-source indirection (spec §9): the skill's replacement is
    // read in place of the backbone tensor — either/or, never a sum.
    let name: &str = &match skill {
        Some(sid) if model.tensor(&format!("skill.{sid}.{name}")).is_some() => {
            format!("skill.{sid}.{name}")
        }
        _ => name.to_string(),
    };
    let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
    if force_f32 {
        let entry = model
            .tensor(name)
            .ok_or_else(|| CmfError::MissingTensor(name.to_string()))?;
        if entry.shape.len() != 2 {
            return Err(err(format!("'{name}' is not 2-D")));
        }
        let data = load_f32(model, name, &Overlay::None).map_err(err)?;
        Ok(QTensor::from_f32(data, entry.shape[0], entry.shape[1]))
    } else {
        QTensor::from_model(model, name).map_err(err)
    }
}

impl Pipeline {
    /// Build a runnable pipeline from an opened CMF model.
    pub fn from_model(
        model: &Arc<CmfModel>,
        sampler_config: SamplerConfig,
    ) -> Result<Self, CmfError> {
        Self::from_model_with_skill(model, sampler_config, None)
    }

    /// Same, with a skill overlaid (spec §9): every layer tensor is
    /// resolved through tensor-source indirection — the skill's
    /// full-shape replacement is read in place of the backbone tensor.
    /// No per-skill model is ever assembled: Mapped tensors are
    /// pointers into the one shared mmap.
    pub fn from_model_with_skill(
        model: &Arc<CmfModel>,
        sampler_config: SamplerConfig,
        skill: Option<&str>,
    ) -> Result<Self, CmfError> {
        match skill {
            Some(s) => Self::from_model_with_overlay(model, sampler_config, &Overlay::One(s)),
            None => Self::from_model_with_overlay(model, sampler_config, &Overlay::None),
        }
    }

    /// Soft superposition (claim 14): working tensors accumulated from
    /// the given (skill, weight) list — softmax(−E/T) upstream.
    pub fn from_model_with_blend(
        model: &Arc<CmfModel>,
        sampler_config: SamplerConfig,
        blend: &[(String, f32)],
    ) -> Result<Self, CmfError> {
        Self::from_model_with_overlay(model, sampler_config, &Overlay::Blend(blend))
    }

    fn from_model_with_overlay(
        model: &Arc<CmfModel>,
        sampler_config: SamplerConfig,
        ov: &Overlay,
    ) -> Result<Self, CmfError> {
        let skill = match ov {
            Overlay::One(s) => Some(*s),
            _ => None,
        };
        if let Some(sid) = skill {
            let known = model.header.skills.iter().any(|s| s.id == sid)
                || model.skill_tensors(sid).next().is_some();
            if !known {
                return Err(CmfError::Parse(format!(
                    "skill '{sid}' not in this container (header.skills: {:?})",
                    model
                        .header
                        .skills
                        .iter()
                        .map(|s| &s.id)
                        .collect::<Vec<_>>()
                )));
            }
            tracing::info!(
                "skill '{sid}': {} replacement tensors overlaid",
                model.skill_tensors(sid).count()
            );
        }
        let arch = model.arch().clone();
        let err = |e: String| CmfError::Parse(format!("weight loading: {e}"));
        if let Some(heads) = &arch.attention_heads_per_layer {
            if heads.len() != arch.num_layers {
                return Err(CmfError::Parse(format!(
                    "arch.attention_heads_per_layer has {} entries, expected {}",
                    heads.len(),
                    arch.num_layers
                )));
            }
            if let Some((li, &nh)) = heads
                .iter()
                .enumerate()
                .find(|(_, nh)| **nh == 0 || **nh % arch.num_kv_heads != 0)
            {
                return Err(CmfError::Parse(format!(
                    "layer {li} has {nh} Q heads, which must be nonzero and divisible by {} KV heads",
                    arch.num_kv_heads
                )));
            }
        }
        if arch
            .layer_types
            .iter()
            .any(|t| matches!(t, LayerType::SlidingAttention))
            && arch.sliding_window.is_none()
        {
            return Err(CmfError::Parse(
                "model has SlidingAttention layers but no arch.sliding_window".into(),
            ));
        }

        // Masks × quantized mmap: only ATTENTION keeps f32 (the head-mask
        // path needs f32 slices). FFN masks now run sparse directly on the
        // quant bytes (sparse_ffn_quant), and embed/lm_head are never
        // masked — so a masked model runs at quantized RSS, not the old
        // whole-model-f32 blowup.
        let masks_present = !model.masks.masks.is_empty();
        let force_f32 = masks_present; // attention only (head masks)

        // ── Tokenizer: embedded → sidecar → byte-level fallback ──
        let mut tokenizer = if let Some(vocab_bytes) = &model.vocab {
            Tokenizer::from_bytes(vocab_bytes)
                .map_err(|e| CmfError::Parse(format!("embedded tokenizer: {e}")))?
        } else {
            let sidecar = model.path.with_file_name("tokenizer.json");
            if sidecar.exists() {
                Tokenizer::from_file(&sidecar)
                    .map_err(|e| CmfError::Parse(format!("sidecar tokenizer: {e}")))?
            } else {
                tracing::warn!("no tokenizer in file or sidecar — using byte-level fallback");
                Tokenizer::byte_level()
            }
        };
        // Chat/eos bundle (spec §6.1): the FILE defines chat behavior.
        if let Some(tc) = &model.header.tokenizer_config {
            tokenizer.chat_template = tc.chat_template.clone();
            tokenizer.extra_eos.extend(tc.eos_token_ids.iter().copied());
            if tokenizer.bos_token_id.is_none() {
                tokenizer.bos_token_id = tc.bos_token_id;
            }
            tracing::info!(
                "chat bundle: template {} chars, {} stop ids",
                tc.chat_template.as_deref().map(str::len).unwrap_or(0),
                tc.eos_token_ids.len()
            );
        }
        // Gemma's contract requires <bos> at sequence start, but its
        // tokenizer.json post-processor does not add it (the chat
        // template does). Raw prompts need it too — word salad without.
        if arch.arch_name.to_lowercase().contains("gemma") && tokenizer.bos_token_id.is_some() {
            tokenizer.add_bos = true;
        }

        // ── Top-level weights (never masked → always quantized) ──
        let embed_tokens = load_matrix(model, "model.embed_tokens.weight", false, ov)?;
        let final_norm = load_f32(model, "model.norm.weight", ov).map_err(err)?;
        let lm_head = if model.tensor("lm_head.weight").is_some() {
            load_matrix(model, "lm_head.weight", false, ov)?
        } else if arch.tie_word_embeddings {
            // Tied: reuse the embedding matrix (re-open, cheap for Mapped).
            load_matrix(model, "model.embed_tokens.weight", false, ov)?
        } else {
            return Err(CmfError::MissingTensor(
                "lm_head.weight (and tie_word_embeddings is false)".into(),
            ));
        };

        // ── Linear-core geometry (required if any linear layer exists) ──
        let has_linear = arch
            .layer_types
            .iter()
            .any(|t| matches!(t, LayerType::LinearAttention));
        let mut vmf_cfg = None;
        let mut gdn_cfg = None;
        if has_linear {
            let lc = arch.linear_core.as_ref().ok_or_else(|| {
                CmfError::Parse(
                    "model has LinearAttention layers but no arch.linear_core — \
                     reconvert with the current converter"
                        .into(),
                )
            })?;
            let need = |v: Option<usize>, name: &str| {
                v.ok_or_else(|| CmfError::Parse(format!("linear core needs arch.{name}")))
            };
            match lc.kind.as_str() {
                "vmf_phase" => {
                    vmf_cfg = Some(VmfPhaseCfg {
                        num_heads: lc.num_heads,
                        nphase: need(lc.nphase, "linear_core.nphase")?,
                        value_head_dim: lc.value_head_dim,
                        hidden_size: arch.hidden_size,
                        // θ-mass (η′): default 0 (massless); CMF_PHASE_MASS
                        // widens the phase kernel for folded-unhealed models.
                        phase_mass: std::env::var("CMF_PHASE_MASS")
                            .ok()
                            .and_then(|v| v.parse().ok())
                            .unwrap_or(0.0),
                    });
                }
                "gated_delta_net" => {
                    gdn_cfg = Some(GdnCfg {
                        num_v_heads: lc.num_heads,
                        num_k_heads: need(arch.linear_num_key_heads, "linear_num_key_heads")?,
                        key_head_dim: need(arch.linear_key_head_dim, "linear_key_head_dim")?,
                        value_head_dim: lc.value_head_dim,
                        conv_kernel: need(arch.linear_conv_kernel_dim, "linear_conv_kernel_dim")?,
                        hidden_size: arch.hidden_size,
                        rms_eps: arch.rms_norm_eps,
                    });
                }
                other => {
                    return Err(CmfError::Parse(format!(
                        "unknown linear core '{other}' (this runtime executes: \
                         gated_delta_net, vmf_phase)"
                    )));
                }
            }
        }

        // ── Short-convolution geometry (LFM2 conv mixer layers) ──
        let has_short_conv = arch
            .layer_types
            .iter()
            .any(|t| matches!(t, LayerType::ShortConv));
        let short_conv_cfg = if has_short_conv {
            Some(ShortConvCfg {
                hidden_size: arch.hidden_size,
                kernel: arch.linear_conv_kernel_dim.ok_or_else(|| {
                    CmfError::Parse(
                        "model has ShortConv layers but no arch.linear_conv_kernel_dim — \
                         reconvert with the current converter"
                            .into(),
                    )
                })?,
            })
        } else {
            None
        };

        // ── Layers ──
        let load_full_attn = |prefix: &str, layer: Option<usize>| -> Result<AttnKind, CmfError> {
            let t = |suffix: &str| load_matrix(model, &format!("{prefix}{suffix}"), force_f32, ov);
            let n = |suffix: &str| -> Option<Vec<f32>> {
                model
                    .tensor(&format!("{prefix}{suffix}"))
                    .and_then(|_| load_f32(model, &format!("{prefix}{suffix}"), ov).ok())
            };
            let wq = t("self_attn.q_proj.weight")?;
            let nh = layer
                .and_then(|li| {
                    arch.attention_heads_per_layer
                        .as_ref()
                        .and_then(|v| v.get(li).copied())
                })
                .unwrap_or(arch.num_attention_heads);
            // Qwen3.5 output gate: q_proj rows = 2·nh·hd (per-head [q; gate]).
            // Gemma-4 global layers legitimately have nh·global_head_dim
            // rows (which can equal 2·nh·hd) — never gated.
            let output_gate = arch.global_head_dim.is_none() && wq.rows() == 2 * nh * arch.head_dim;
            if !output_gate && wq.rows() != nh * arch.head_dim {
                return Err(CmfError::Parse(format!(
                    "{prefix}self_attn.q_proj.weight rows={} != heads({nh}) * head_dim({})",
                    wq.rows(),
                    arch.head_dim
                )));
            }
            let gate_name = format!("{prefix}self_attn.g_proj.weight");
            let softplus_gate = if model.tensor(&gate_name).is_some() {
                let gate = load_matrix(model, &gate_name, force_f32, ov)?;
                if gate.cols() != arch.hidden_size {
                    return Err(CmfError::Parse(format!(
                        "{gate_name} cols={} != hidden_size ({})",
                        gate.cols(),
                        arch.hidden_size
                    )));
                }
                let per_head = if gate.rows() == nh {
                    true
                } else if gate.rows() == nh * arch.head_dim {
                    false
                } else {
                    return Err(CmfError::Parse(format!(
                        "{gate_name} rows={} must equal heads ({nh}) or heads*head_dim ({})",
                        gate.rows(),
                        nh * arch.head_dim
                    )));
                };
                Some((gate, per_head))
            } else {
                None
            };
            // Qwen2-family projection biases (by tensor presence).
            let bias = match (
                n("self_attn.q_proj.bias"),
                n("self_attn.k_proj.bias"),
                n("self_attn.v_proj.bias"),
            ) {
                (Some(a), Some(b), Some(c)) => Some((a, b, c)),
                _ => None,
            };
            Ok(AttnKind::Full {
                wq,
                wk: t("self_attn.k_proj.weight")?,
                wv: t("self_attn.v_proj.weight")?,
                wo: t("self_attn.o_proj.weight")?,
                q_norm: n("self_attn.q_norm.weight"),
                k_norm: n("self_attn.k_norm.weight"),
                output_gate,
                softplus_gate,
                bias,
            })
        };

        let load_linear_attn = |prefix: &str| -> Result<AttnKind, CmfError> {
            if gdn_cfg.is_some() {
                // Faithful vendor operator: tensor names 1:1 with the source.
                let t = |suffix: &str| {
                    load_matrix(
                        model,
                        &format!("{prefix}linear_attn.{suffix}"),
                        force_f32,
                        ov,
                    )
                };
                let f = |suffix: &str| {
                    load_f32(model, &format!("{prefix}linear_attn.{suffix}"), ov).map_err(err)
                };
                return Ok(AttnKind::LinearGdn(GdnWeights {
                    in_proj_qkv: t("in_proj_qkv.weight")?,
                    in_proj_z: t("in_proj_z.weight")?,
                    in_proj_a: t("in_proj_a.weight")?,
                    in_proj_b: t("in_proj_b.weight")?,
                    conv1d: f("conv1d.weight")?,
                    a_log: f("A_log")?,
                    dt_bias: f("dt_bias")?,
                    norm: f("norm.weight")?,
                    out_proj: t("out_proj.weight")?,
                }));
            }
            let t = |suffix: &str| {
                load_matrix(model, &format!("{prefix}vmf_attn.{suffix}"), force_f32, ov)
            };
            let a_log = load_f32(model, &format!("{prefix}vmf_attn.A_log"), ov).map_err(err)?;
            // Selective-write gate κ (hybrid_k core): optional by tensor
            // presence — files without it run the classic phase kernel
            // bit-identically.
            let k_gate = if model
                .tensor(&format!("{prefix}vmf_attn.k_gate.weight"))
                .is_some()
            {
                Some((
                    t("k_gate.weight")?,
                    load_f32(model, &format!("{prefix}vmf_attn.k_gate.bias"), ov).map_err(err)?,
                ))
            } else {
                None
            };
            Ok(AttnKind::Linear(VmfPhaseWeights {
                thq: t("thq.weight")?,
                thk: t("thk.weight")?,
                v_proj: t("v_proj.weight")?,
                out_proj: t("out_proj.weight")?,
                decay: a_log.iter().map(|&a| (-(a as f64).exp()).exp()).collect(),
                k_gate,
            }))
        };

        // LFM2 short-conv mixer: in_proj [3·hidden, hidden], a depthwise
        // conv (stored f16 as `[hidden, 1, kernel]` → flattened taps), and
        // out_proj [hidden, hidden]. Names canonicalized at convert time.
        let load_short_conv = |prefix: &str| -> Result<AttnKind, CmfError> {
            let t = |suffix: &str| {
                load_matrix(
                    model,
                    &format!("{prefix}short_conv.{suffix}"),
                    force_f32,
                    ov,
                )
            };
            Ok(AttnKind::ShortConv(ShortConvWeights {
                in_proj: t("in_proj.weight")?,
                conv: load_f32(model, &format!("{prefix}short_conv.conv.weight"), ov)
                    .map_err(err)?,
                out_proj: t("out_proj.weight")?,
            }))
        };

        let mut layers = Vec::with_capacity(arch.num_layers);
        for li in 0..arch.num_layers {
            let prefix = format!("model.layers.{li}.");
            let attn = match arch.layer_types.get(li) {
                Some(LayerType::LinearAttention) => load_linear_attn(&prefix)?,
                Some(LayerType::ShortConv) => load_short_conv(&prefix)?,
                _ => load_full_attn(&prefix, Some(li))?,
            };
            // Gemma-2/3 sandwich: `pre_feedforward_layernorm` present →
            // it is the pre-FFN norm, and post_attention/post_feedforward
            // норms apply to the branch OUTPUTS before their residuals.
            let pre_ffn = format!("{prefix}pre_feedforward_layernorm.weight");
            let sandwich = model.tensor(&pre_ffn).is_some();
            layers.push(LayerWeights {
                input_norm: load_f32(model, &format!("{prefix}input_layernorm.weight"), ov)
                    .map_err(err)?,
                post_norm: if sandwich {
                    load_f32(model, &pre_ffn, ov).map_err(err)?
                } else {
                    load_f32(
                        model,
                        &format!("{prefix}post_attention_layernorm.weight"),
                        ov,
                    )
                    .map_err(err)?
                },
                attn_out_norm: if sandwich {
                    Some(
                        load_f32(
                            model,
                            &format!("{prefix}post_attention_layernorm.weight"),
                            ov,
                        )
                        .map_err(err)?,
                    )
                } else {
                    None
                },
                ffn_out_norm: if sandwich {
                    Some(
                        load_f32(
                            model,
                            &format!("{prefix}post_feedforward_layernorm.weight"),
                            ov,
                        )
                        .map_err(err)?,
                    )
                } else {
                    None
                },
                // Gemma-4: learned scalar multiplying the layer output.
                layer_scale: model
                    .tensor(&format!("{prefix}layer_scalar"))
                    .and_then(|_| {
                        load_f32(model, &format!("{prefix}layer_scalar"), ov)
                            .ok()
                            .and_then(|v| v.first().copied())
                    }),
                // FFN always quantized — masks run sparse on quant bytes.
                ffn: build_layer_ffn(model, &arch, li, false, ov)?,
                attn,
            });
        }

        // ── MTP head (optional, spec §2.1) ──
        let mtp = if let Some(cfg) = &arch.mtp {
            if cfg.num_layers != 1 {
                return Err(CmfError::Parse(format!(
                    "MTP with {} blocks not supported yet (only 1)",
                    cfg.num_layers
                )));
            }
            let p = "model.mtp.";
            let attn = load_full_attn("model.mtp.layers.0.", None)?;
            Some(MtpModule {
                enorm: load_f32(model, &format!("{p}enorm.weight"), ov).map_err(err)?,
                hnorm: load_f32(model, &format!("{p}hnorm.weight"), ov).map_err(err)?,
                eh_proj: load_matrix(model, &format!("{p}eh_proj.weight"), false, ov)?,
                layer: LayerWeights {
                    attn_out_norm: None,
                    ffn_out_norm: None,
                    layer_scale: None,
                    input_norm: load_f32(model, &format!("{p}layers.0.input_layernorm.weight"), ov)
                        .map_err(err)?,
                    post_norm: load_f32(
                        model,
                        &format!("{p}layers.0.post_attention_layernorm.weight"),
                        ov,
                    )
                    .map_err(err)?,
                    ffn: FfnKind::Dense(DenseFfn {
                        gate_proj: load_matrix(
                            model,
                            &format!("{p}layers.0.mlp.gate_proj.weight"),
                            false,
                            ov,
                        )?,
                        up_proj: load_matrix(
                            model,
                            &format!("{p}layers.0.mlp.up_proj.weight"),
                            false,
                            ov,
                        )?,
                        down_proj: load_matrix(
                            model,
                            &format!("{p}layers.0.mlp.down_proj.weight"),
                            false,
                            ov,
                        )?,
                        act: crate::pipeline::Act::from_arch(&arch.hidden_act),
                    }),
                    attn,
                },
                final_norm: load_f32(model, &format!("{p}norm.weight"), ov).map_err(err)?,
                kv: LayerKvCache::new(arch.num_kv_heads, arch.head_dim),
            })
        } else {
            None
        };

        tracing::info!(
            "Pipeline loaded: {} | {}L ({} linear) | {:.2}B params | storage: {} | MTP: {}",
            arch.arch_name,
            arch.num_layers,
            arch.layer_types
                .iter()
                .filter(|t| matches!(t, LayerType::LinearAttention))
                .count(),
            model.total_param_count() as f64 / 1e9,
            if force_f32 {
                "f32 (masked)"
            } else {
                "quantized mmap"
            },
            if mtp.is_some() { "yes" } else { "no" }
        );

        // KV window: the descriptor's max, capped for dev-box safety;
        // CMF_MAX_SEQ overrides the cap (long-context runs).
        let cap = std::env::var("CMF_MAX_SEQ")
            .ok()
            .and_then(|v| v.parse::<usize>().ok())
            .unwrap_or(8192);
        let max_seq_len = arch.max_position_embeddings.min(cap);

        // Looped Transformer: total virtual layers = physical × num_loops.
        let total_layers = arch.num_layers * arch.num_loops;

        let mut pipeline = Pipeline::new(
            tokenizer,
            PipelineWeights {
                embed_tokens,
                layers,
                lm_head,
                final_norm,
            },
            arch.hidden_size,
            arch.intermediate_size,
            arch.num_attention_heads,
            arch.num_kv_heads,
            arch.head_dim,
            total_layers,
            arch.num_layers, // physical layers in weights
            arch.loop_final_norm,
            arch.vocab_size,
            arch.rms_norm_eps,
            arch.rope_theta as f32,
            arch.norm_style,
            max_seq_len,
            sampler_config,
        );
        let rotary = ((arch.head_dim as f32 * arch.partial_rotary_factor) as usize).max(2);
        pipeline.set_rotary(rotary, arch.rope_theta as f32);
        pipeline.attention_heads_per_layer = arch.attention_heads_per_layer.clone();
        if let Some(yarn) = &arch.yarn {
            pipeline.inv_freq = std::sync::Arc::new(crate::attention::yarn_inv_freq(
                rotary,
                arch.rope_theta as f32,
                yarn.factor,
                yarn.original_max_position_embeddings,
                yarn.beta_fast,
                yarn.beta_slow,
            ));
            pipeline.rope_scale = yarn.attention_factor;
        }
        // Gemma-family extras: embedding scale, attention-scale
        // override, and (Gemma-3) sliding-window layers with their own
        // local RoPE base.
        pipeline.embed_multiplier = arch.embed_multiplier;
        if let Some(qpas) = arch.query_pre_attn_scalar {
            pipeline.attn_scale = 1.0 / (qpas as f32).sqrt();
        }
        if let (Some(w), Some(p)) = (arch.sliding_window, arch.sliding_window_pattern) {
            pipeline.swa = Some((w, p));
            if let Some(base) = arch.rope_local_base_freq {
                pipeline.inv_freq_local = Some(std::sync::Arc::new(
                    crate::attention::rope_inv_freq(rotary, base as f32),
                ));
            }
        }
        let explicit_sliding: Vec<bool> = arch
            .layer_types
            .iter()
            .map(|t| matches!(t, cortiq_core::LayerType::SlidingAttention))
            .collect();
        if explicit_sliding.iter().any(|&v| v) {
            pipeline.sliding_layers = Some(explicit_sliding);
            if let Some(w) = arch.sliding_window {
                pipeline.swa = Some((w, usize::MAX));
            }
            let local_rotary = ((arch.head_dim as f32
                * arch
                    .local_partial_rotary_factor
                    .unwrap_or(arch.partial_rotary_factor))
                as usize)
                .max(2);
            pipeline.rotary_dim_local = Some(local_rotary);
            if let Some(base) = arch.rope_local_base_freq {
                pipeline.inv_freq_local = Some(std::sync::Arc::new(
                    crate::attention::rope_inv_freq(local_rotary, base as f32),
                ));
            }
        }
        // Gemma-4: global layers run their own geometry (MQA at
        // global_head_dim) with a proportional RoPE — the first
        // factor·head_dim dims rotate, the zero-padded tail is identity.
        if let (Some(ghd), Some(gkv)) = (arch.global_head_dim, arch.num_global_kv_heads) {
            pipeline.global_attn = Some((ghd, gkv));
            let prf = arch.global_partial_rotary_factor.unwrap_or(1.0);
            let half = ghd / 2;
            let ra = (((prf * ghd as f32) as usize) / 2).min(half);
            let mut f = vec![0.0f32; half];
            for (i, slot) in f.iter_mut().enumerate().take(ra) {
                *slot = 1.0 / (arch.rope_theta as f32).powf(2.0 * i as f32 / ghd as f32);
            }
            pipeline.inv_freq_global = Some(std::sync::Arc::new(f));
            // Re-shape the global layers' KV storage to their geometry.
            if let Some((_, p)) = pipeline.swa {
                for li in 0..arch.num_layers {
                    if (li + 1) % p.max(1) == 0 {
                        pipeline.kv_cache.layers[li] = crate::kv_cache::LayerKvCache::new(gkv, ghd);
                    }
                }
            }
        }
        pipeline.attn_v_norm = arch.attn_v_norm;
        pipeline.final_softcap = arch.final_logit_softcapping.map(|c| c as f32);
        pipeline.vmf_cfg = vmf_cfg;
        pipeline.gdn_cfg = gdn_cfg;
        pipeline.short_conv_cfg = short_conv_cfg;
        pipeline.mtp = mtp;
        pipeline.install_dynamic_routing(model, false);
        // Record the load-time overlay so a later set_active_skill(None)
        // correctly reverts it (the union-diff assumes dyn_active mirrors
        // the live overlay). Blend loads have no single index to revert.
        match ov {
            Overlay::One(sid) => {
                pipeline.dyn_active = model.header.skills.iter().position(|s| &s.id == sid);
            }
            Overlay::Blend(_) => pipeline.dyn_blend_loaded = true,
            Overlay::None => {}
        }
        // B1: apply the measured confidence-calibration temperature, if the
        // file carries one (softmax(logits / T) for reported Born mass).
        if let Some(c) = &model.header.calibration {
            pipeline.set_calib_temp(c.temperature);
        }
        // O(1) Nyström attention (runtime-level, no format change):
        // env CMF_O1 decides; unset falls through to the converter hint
        // in header.provenance.o1_attn (`cortiq convert --o1`), and
        // CMF_O1=off force-disables even the hint. CLI flags override
        // later via set_o1().
        let o1 = match crate::nystrom::o1_from_env() {
            crate::nystrom::O1Env::Off => None,
            crate::nystrom::O1Env::On(cfg) => Some(cfg),
            crate::nystrom::O1Env::Unset => model
                .header
                .provenance
                .as_ref()
                .and_then(|p| p.get("o1_attn"))
                .and_then(crate::nystrom::O1Cfg::from_json),
        };
        if o1.is_some() {
            pipeline.set_o1(o1);
        }
        Ok(pipeline)
    }

    /// Record per-skill dynamic-routing metadata: which FFN layers each
    /// skill actually replaces (derived from the tensors present, not
    /// the meta `layers` field), and whether the skill is eligible for
    /// cheap dynamic switching (FFN-only). Called once at load.
    pub(crate) fn install_dynamic_routing(&mut self, model: &Arc<CmfModel>, force_f32: bool) {
        self.model = Some(model.clone());
        self.dyn_force_f32 = force_f32;
        let mut per_skill = Vec::with_capacity(model.header.skills.len());
        for sk in &model.header.skills {
            let mut ffn_layers = std::collections::BTreeSet::new();
            let mut non_ffn = false;
            let prefix = format!("skill.{}.", sk.id);
            for t in model.skill_tensors(&sk.id) {
                let rel = &t.name[prefix.len()..]; // e.g. model.layers.20.mlp.down_proj.weight
                let toks: Vec<&str> = rel.split('.').collect();
                if toks.len() >= 5 && toks[0] == "model" && toks[1] == "layers" && toks[3] == "mlp"
                {
                    if let Ok(li) = toks[2].parse::<usize>() {
                        ffn_layers.insert(li);
                        continue;
                    }
                }
                non_ffn = true; // replaces attention / embed / lm_head
            }
            if non_ffn {
                tracing::warn!(
                    "skill '{}' replaces non-FFN tensors — excluded from dynamic \
                     routing (static overlay still works)",
                    sk.id
                );
                per_skill.push(None);
            } else {
                per_skill.push(Some(ffn_layers.into_iter().collect::<Vec<_>>()));
            }
        }
        self.dyn_skill_layers = per_skill;
    }

    /// Switch the overlaid skill for subsequent forwards (dynamic
    /// routing). `idx` = index into model.header.skills; None = backbone.
    /// Rebuilds the FFN of the union of the old and new skill's touched
    /// layers with the new overlay — tensor-source indirection made
    /// dynamic. Cheap: Mapped tensors are re-resolved mmap pointers.
    /// Result is bit-identical to loading the pipeline with that skill.
    pub fn set_active_skill(&mut self, idx: Option<usize>) -> Result<(), CmfError> {
        if self.dyn_active == idx {
            return Ok(());
        }
        let model = self.model.clone().ok_or_else(|| {
            CmfError::Parse("dynamic routing needs a model-backed pipeline".into())
        })?;
        let mut union: std::collections::BTreeSet<usize> = std::collections::BTreeSet::new();
        if let Some(old) = self.dyn_active {
            if let Some(Some(ls)) = self.dyn_skill_layers.get(old) {
                union.extend(ls.iter().copied());
            }
        }
        let new_id: Option<String> = match idx {
            Some(n) => match self.dyn_skill_layers.get(n) {
                Some(Some(ls)) => {
                    union.extend(ls.iter().copied());
                    Some(model.header.skills[n].id.clone())
                }
                _ => {
                    return Err(CmfError::Parse(format!(
                        "skill index {n} not dynamic-eligible"
                    )));
                }
            },
            None => None,
        };
        let ov = match &new_id {
            Some(s) => Overlay::One(s),
            None => Overlay::None,
        };
        let arch = model.arch();
        for li in union {
            self.weights.layers[li].ffn =
                build_layer_ffn(&model, arch, li, self.dyn_force_f32, &ov)?;
        }
        self.dyn_active = idx;
        Ok(())
    }
}