inferencelayer 0.2.1

//! CPU encoder runtime (feature `encoder-cpu`, native-only): the no-GPU fallback AND the parity
//! oracle every encoder GPU kernel is verified against.
//!
//! Design: op-for-op clarity in the `reference.rs` spirit, but batched and fast — every linear
//! projection runs as ONE `gemm` over the whole concatenated token matrix (ragged batching means
//! the token dimension simply concatenates), attention/norm/activation stages parallelize over
//! rows with rayon inside an encoder-owned thread pool (never the global pool the host process
//! may be using). `gemm` is the same pure-Rust backend candle-CPU uses, which is what makes the
//! "as fast as embed_anything on CPU" bar structurally reachable without any candle dependency.
//!
//! The executor branches on the DECLARATIVE config only (post-LN vs prenorm residual placement,
//! learned vs rotary positions, dense vs GLU feed-forward, per-layer windows) — never on
//! `EncArch`; the architecture tag selects a tensor-name table at load and nothing else.

use crate::encoder_weights::{
    Act, EncArch, EncBatch, EncoderConfig, MaskKind, MlpKind, NormKind, PosKind, RelAttn,
};
use crate::pooling::{self, EmbedOut, Pooling};
use crate::weights::LazySt;
use anyhow::{Context, Result};
use rayon::prelude::*;
use std::path::Path;

/// One norm: weight + optional bias (bias-free for ModernBERT).
struct CpuNorm {
    w: Vec<f32>,
    b: Option<Vec<f32>>,
}

/// One linear projection, HF Linear layout (`[n, k]` row-major), optional bias.
#[derive(Default)]
struct CpuLinear {
    w: Vec<f32>,
    b: Option<Vec<f32>>,
    n: usize,
    k: usize,
    /// `w` repacked into the microkernel's panel layout, built on first use and reused for the life
    /// of the model. This is the whole point of `cpu_gemm`: the weight is immutable, so the pack
    /// belongs at load time, not in every call (which is what the `gemm` crate does — see
    /// [`crate::cpu_gemm`]). `OnceLock` rather than eager packing so that the many loaders which
    /// build a `CpuLinear` for a *tensor they then never run* (shape probes, unused heads) pay
    /// nothing, and so the 18 construction sites need no edit.
    packed: std::sync::OnceLock<crate::cpu_gemm::PackedWeight>,
    /// INT8 twin (`OSFKB_CPU_I8=1` on an i8mm host): per-row-quantized, smmla-panel-packed —
    /// the llama.cpp Arm path, opt-in behind its own cosine gate (`tests/cpu_i8.rs`).
    packed_i8: std::sync::OnceLock<crate::cpu_gemm::PackedWeightI8>,
    /// F16-storage twin (`OSFKB_CPU_F16=1` on a FEAT_FHM host): the widening-FMLAL kernel —
    /// same MAC rate as FMA from half the streamed bytes, which is what the measured multicore
    /// wall is made of. Opt-in behind its own cosine gate (`tests/cpu_f16.rs`).
    packed_f16: std::sync::OnceLock<crate::cpu_gemm::PackedWeightF16>,
}

/// One encoder layer. The two residual placements read these fields differently:
/// post-LN (BERT): `attn_norm` is applied AFTER the attention residual, `mlp_norm` after the
/// MLP residual. Prenorm (ModernBERT): `attn_norm` normalizes the attention INPUT (`None` =
/// layer-0 skip), `mlp_norm` the MLP input.
/// The layer's mixing operator: attention, or the LFM2 gated short-conv.
// `Attn` holds six `CpuLinear`s to `Conv`'s three, so it is the larger variant by a few hundred
// bytes. One of these exists per layer (tens per model, built once at load), so boxing to equalize
// them would trade a pointer chase in the hot loop for memory nobody misses.
#[allow(clippy::large_enum_variant)]
enum CpuOp {
    Attn {
        q: CpuLinear,
        k: CpuLinear,
        v: CpuLinear,
        o: CpuLinear,
        /// Per-head RMS weights for q/k before rotary (Qwen3/LFM2; `[head_dim]` each).
        q_norm: Option<Vec<f32>>,
        k_norm: Option<Vec<f32>>,
    },
    /// LFM2 gated short-conv (stateless full-sequence form of the decoder's ring buffer):
    /// `in_proj` → per-channel (B, C, x); `bx = B⊙x`; depthwise causal conv over the last
    /// `conv_l` taps (zero history before the sequence start); output `C ⊙ (bx ⊛ w)` → `out_proj`.
    Conv {
        in_proj: CpuLinear,
        /// Depthwise taps `[hidden · conv_l]`, oldest-first (tap `l-1` is the current token).
        conv_w: Vec<f32>,
        out_proj: CpuLinear,
    },
}

struct CpuLayer {
    attn_norm: Option<CpuNorm>,
    op: CpuOp,
    mlp_norm: CpuNorm,
    /// Dense `up` ([I, H]) or fused GLU `Wi` ([2I, H]).
    up: CpuLinear,
    down: CpuLinear,
    /// DeBERTa: this layer's relative-position projections (`None` on every other architecture).
    rel: Option<CpuRelPos>,
}

/// DeBERTa's relative-embedding table projected through ONE layer's own q/k projections
/// (`share_att_key: true`). Both are `[2·span, n_heads · head_dim]`, computed once at load — they
/// depend only on the weights, never on the input, so recomputing them per forward would be pure
/// waste (~0.3 GMAC per layer).
struct CpuRelPos {
    /// `key_proj(LayerNorm(rel_embeddings))` — the content→position (c2p) keys.
    pos_k: Vec<f32>,
    /// `query_proj(LayerNorm(rel_embeddings))` — the position→content (p2c) queries.
    pos_q: Vec<f32>,
}

/// `BertForSequenceClassification` scoring head (sentence-transformers CrossEncoder): the pooler
/// dense (`h×h`, tanh applied to the CLS state) then the classifier (`n_labels×h`).
struct CpuCrossHead {
    pooler: CpuLinear,
    classifier: CpuLinear,
}

/// All encoder weights resident in RAM as f32.
struct CpuWeights {
    word: Vec<f32>,
    /// Learned absolute position table (None for rotary models).
    pos: Option<Vec<f32>>,
    ttype: Option<Vec<f32>>,
    /// Embedding-stage norm (BERT/ModernBERT); rotary decoders (Qwen3) have none — the first
    /// layer's own input norm is the first normalization the stream sees.
    emb_norm: Option<CpuNorm>,
    layers: Vec<CpuLayer>,
    /// Trailing norm before pooling (ModernBERT `final_norm`; None for post-LN BERT).
    final_norm: Option<CpuNorm>,
    /// Late-interaction per-token projection (pylate `1_Dense`; ColBERT).
    proj: Option<CpuLinear>,
    /// Post-pooling projection (SigLIP text `head`).
    pooled_head: Option<CpuLinear>,
    /// SigLIP vision extras (patch embedding + MAP attention-pooling head).
    vision: Option<CpuVision>,
    /// `*ForSequenceClassification` scoring head (CrossEncoder); None for embedders.
    cross_encoder: Option<CpuCrossHead>,
}

/// SigLIP vision-tower extras: patch embedding and the MAP attention-pooling head
/// (torch `MultiheadAttention` layout: packed `in_proj` q|k|v with biases).
struct CpuVision {
    /// Patch embedding as a linear over im2row rows `[3·p²] → [hidden]` (conv ≡ GEMM).
    patch: CpuLinear,
    /// Learned per-patch position table `[n_patches · hidden]`.
    pos: Vec<f32>,
    /// Learned probe query `[hidden]`.
    probe: Vec<f32>,
    /// Packed q|k|v projection `[3H, H]` + bias `[3H]`.
    in_proj_w: Vec<f32>,
    in_proj_b: Vec<f32>,
    out_proj: CpuLinear,
    head_ln: CpuNorm,
    head_fc1: CpuLinear,
    head_fc2: CpuLinear,
}

/// The CPU encoder: loads a checkpoint directory (`config.json` + `model.safetensors`) and
/// embeds ragged token batches. Deterministic per input (fixed reduction orders), so it serves
/// as the GPU kernels' tolerance-gated oracle.
pub struct CpuEncoder {
    cfg: EncoderConfig,
    w: CpuWeights,
    pool: rayon::ThreadPool,
}

impl CpuEncoder {
    /// Load from a checkpoint dir. Refuses (loudly) architectures whose CPU arm hasn't landed —
    /// "recognized" and "supported" stay honest.
    pub fn load(dir: &Path) -> Result<Self> {
        // SigLIP checkpoints are a JOINT text+vision config (often without `architectures`);
        // route them to the tower loaders before the generic dir-level resolution.
        if crate::encoder_weights::siglip_configs_from_dir(dir).is_ok() {
            return Self::load_siglip_text(dir);
        }
        let cfg = crate::encoder_weights::encoder_config_from_dir(dir)?;
        let st = LazySt::open(dir)?;
        let w = match cfg.arch {
            EncArch::Bert | EncArch::XlmRoberta => load_bert_family(&st, &cfg)?,
            EncArch::DebertaV2 => load_deberta_v2(&st, &cfg)?,
            EncArch::ModernBert => load_modernbert(&st, &cfg)?,
            EncArch::Qwen3Embed => load_qwen3_embed(&st, &cfg)?,
            EncArch::Lfm2Colbert => load_lfm2_colbert(&st, &cfg, dir)?,
            EncArch::NomicBert => load_nomic_bert(&st, &cfg)?,
            // Config parses (detect/validate work); the tensor table awaits verification against
            // a real checkpoint — loading must fail loudly, never guess.
            other => anyhow::bail!(
                "CPU encoder tensor table for {other:?} pending verification against a real \
                 checkpoint"
            ),
        };
        let pool = rayon::ThreadPoolBuilder::new()
            .thread_name(|i| format!("lfm2-enc-cpu-{i}"))
            // `OSFKB_ENC_THREADS`: pool width knob. Default = all logical cores; on big.LITTLE
            // hosts (M-series: 12P+4E) the last wave of a parallel GEMM runs at E-core speed,
            // and capping to the performance-core count is a measurable A/B (the 54%-of-roofline
            // question). Env-tuned rather than auto: core topology detection is not portable.
            .num_threads(
                std::env::var("OSFKB_ENC_THREADS")
                    .ok()
                    .and_then(|v| v.parse().ok())
                    .unwrap_or(0),
            )
            .build()
            .context("encoder thread pool")?;
        Ok(Self { cfg, w, pool })
    }

    /// Load the SigLIP TEXT tower from a joint checkpoint.
    pub fn load_siglip_text(dir: &Path) -> Result<Self> {
        let (cfg, _) = crate::encoder_weights::siglip_configs_from_dir(dir)?;
        let st = LazySt::open(dir)?;
        let w = load_siglip_text_weights(&st, &cfg)?;
        let pool = rayon::ThreadPoolBuilder::new()
            .thread_name(|i| format!("lfm2-enc-cpu-{i}"))
            .build()
            .context("encoder thread pool")?;
        Ok(Self { cfg, w, pool })
    }

    /// Load the SigLIP VISION tower from a joint checkpoint.
    pub fn load_siglip_vision(dir: &Path) -> Result<Self> {
        let (_, spec) = crate::encoder_weights::siglip_configs_from_dir(dir)?;
        let st = LazySt::open(dir)?;
        let w = load_siglip_vision_weights(&st, &spec)?;
        let pool = rayon::ThreadPoolBuilder::new()
            .thread_name(|i| format!("lfm2-enc-cpu-{i}"))
            .build()
            .context("encoder thread pool")?;
        Ok(Self {
            cfg: spec.config,
            w,
            pool,
        })
    }

    /// Patch-embed one image's normalized pixels (`[3 · image²]` CHW) into the initial
    /// `[n_patches, hidden]` residual stream: im2row patch convolution + learned positions. This is
    /// the pre-transformer vision stage, factored out so the GPU vision tower reuses the EXACT same
    /// patch embedding as the CPU oracle (only the transformer layers differ by device).
    pub(crate) fn patch_embed_rows(
        &self,
        pixels: &[f32],
        image_size: usize,
        patch_size: usize,
    ) -> Result<Vec<f32>> {
        let vis = self
            .w
            .vision
            .as_ref()
            .context("patch_embed_rows requires the SigLIP vision tower")?;
        let h = self.cfg.hidden;
        anyhow::ensure!(
            pixels.len() == 3 * image_size * image_size,
            "pixel buffer {} != 3·{image_size}²",
            pixels.len()
        );
        let grid = image_size / patch_size;
        let n = grid * grid;
        anyhow::ensure!(n == self.cfg.max_pos, "patch count mismatch");
        self.pool.install(|| {
            // im2row: rows of [3·p²] in the conv-weight layout (c-major, then ky, kx).
            let mut rows = vec![0f32; n * 3 * patch_size * patch_size];
            let ps2 = patch_size * patch_size;
            for py in 0..grid {
                for px in 0..grid {
                    let r = py * grid + px;
                    for c in 0..3 {
                        for ky in 0..patch_size {
                            for kx in 0..patch_size {
                                let src = c * image_size * image_size
                                    + (py * patch_size + ky) * image_size
                                    + px * patch_size
                                    + kx;
                                rows[r * 3 * ps2 + c * ps2 + ky * patch_size + kx] = pixels[src];
                            }
                        }
                    }
                }
            }
            let mut cur = vec![0f32; n * h];
            linear_from(&mut cur, &rows, &vis.patch);
            for (i, row) in cur.chunks_exact_mut(h).enumerate() {
                for (r, p) in row.iter_mut().zip(&vis.pos[i * h..(i + 1) * h]) {
                    *r += p;
                }
            }
            Ok(cur)
        })
    }

    /// The SigLIP MAP (multihead-attention-pooling) head over the post-transformer `[n, hidden]`
    /// hidden states → the single pooled, L2-normalized image embedding. Factored out so the GPU
    /// vision tower reuses the EXACT same head as the CPU oracle (packed torch MultiheadAttention
    /// layout: a learned probe attends the patches, then a residual LN+MLP block).
    pub(crate) fn map_head(&self, hidden: &[f32]) -> Result<Vec<f32>> {
        let vis = self
            .w
            .vision
            .as_ref()
            .context("map_head requires the SigLIP vision tower")?;
        let h = self.cfg.hidden;
        anyhow::ensure!(
            !hidden.is_empty() && hidden.len().is_multiple_of(h),
            "hidden {} not a positive multiple of hidden {h}",
            hidden.len()
        );
        let n = hidden.len() / h;
        self.pool.install(|| {
            let nh = self.cfg.n_heads;
            let hd = self.cfg.head_dim;
            let mut q = vec![0f32; h];
            // q = Wq·probe + bq (rows 0..h of in_proj).
            for (o, qv) in q.iter_mut().enumerate() {
                let mut acc = vis.in_proj_b[o];
                for k in 0..h {
                    acc += vis.in_proj_w[o * h + k] * vis.probe[k];
                }
                *qv = acc;
            }
            let mut kbuf = vec![0f32; n * h];
            let mut vbuf = vec![0f32; n * h];
            for (i, row) in hidden.chunks_exact(h).enumerate() {
                for o in 0..h {
                    let (mut ka, mut va) = (vis.in_proj_b[h + o], vis.in_proj_b[2 * h + o]);
                    for (k, &rk) in row.iter().enumerate() {
                        ka += vis.in_proj_w[(h + o) * h + k] * rk;
                        va += vis.in_proj_w[(2 * h + o) * h + k] * rk;
                    }
                    kbuf[i * h + o] = ka;
                    vbuf[i * h + o] = va;
                }
            }
            let scale = 1.0 / (hd as f32).sqrt();
            let mut ctx = vec![0f32; h];
            for head in 0..nh {
                let qo = head * hd;
                let mut scores: Vec<f32> = (0..n)
                    .map(|i| {
                        (0..hd)
                            .map(|d| q[qo + d] * kbuf[i * h + qo + d])
                            .sum::<f32>()
                            * scale
                    })
                    .collect();
                let mx = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
                let mut denom = 0f32;
                for sc in &mut scores {
                    *sc = (*sc - mx).exp();
                    denom += *sc;
                }
                for (i, sc) in scores.iter().enumerate() {
                    let p = sc / denom;
                    for d in 0..hd {
                        ctx[qo + d] += p * vbuf[i * h + qo + d];
                    }
                }
            }
            let mut attn_out = vec![0f32; h];
            linear_from(&mut attn_out, &ctx, &vis.out_proj);
            // residual = attn_out; ln; mlp; residual add; the pooled embedding.
            let mut normed = attn_out.clone();
            layer_norm_rows(&mut normed, h, &vis.head_ln, self.cfg.eps, false);
            let mut mid = vec![0f32; vis.head_fc1.n];
            linear_from(&mut mid, &normed, &vis.head_fc1);
            let act = match self.cfg.mlp_kind {
                MlpKind::Dense { act, .. } => act,
                MlpKind::Glu { act } => act,
            };
            activate(&mut mid, act);
            let mut out = vec![0f32; h];
            linear_from(&mut out, &mid, &vis.head_fc2);
            for (o, r) in out.iter_mut().zip(&attn_out) {
                *o += r;
            }
            pooling::l2_normalize(&mut out);
            Ok(out)
        })
    }

    /// Embed ONE image's normalized pixels (`[3 · image² ]` CHW, already resized/rescaled/
    /// normalized): im2row patch embedding + positions → transformer → post-LN → MAP head.
    /// Vision tower only.
    pub fn encode_image(
        &self,
        pixels: &[f32],
        image_size: usize,
        patch_size: usize,
    ) -> Result<Vec<f32>> {
        let cur = self.patch_embed_rows(pixels, image_size, patch_size)?;
        let n = cur.len() / self.cfg.hidden;
        // Transformer over one ragged sequence of n patches.
        let batch = EncBatch {
            tokens: vec![0; n], // unused: the stream is already materialized
            seq_starts: vec![0, n as u32],
            valid: None,
            type_ids: None,
        };
        let hidden = self.pool.install(|| self.run_layers(cur, &batch))?;
        self.map_head(&hidden)
    }

    /// The parsed configuration.
    pub fn config(&self) -> &EncoderConfig {
        &self.cfg
    }

    /// Embed a ragged batch. Sequence length (plus any learned-position offset) must fit
    /// `max_pos`; empty sequences are refused (a mean over zero rows has no value).
    pub fn encode(&self, batch: &EncBatch) -> Result<EmbedOut> {
        let cfg = &self.cfg;
        anyhow::ensure!(batch.n_seqs() > 0, "empty batch");
        let offset = match cfg.pos_kind {
            PosKind::Learned { offset } => offset,
            PosKind::Rope { .. } => 0,
        };
        for w in batch.seq_starts.windows(2) {
            let len = (w[1] - w[0]) as usize;
            anyhow::ensure!(len > 0, "empty sequence in batch");
            anyhow::ensure!(
                len + offset <= cfg.max_pos,
                "sequence of {len} tokens exceeds max positions {} (offset {offset})",
                cfg.max_pos
            );
        }
        self.pool.install(|| self.encode_inner(batch))
    }

    fn encode_inner(&self, batch: &EncBatch) -> Result<EmbedOut> {
        let cur = self.forward_hidden_inner(batch)?;
        // Pooling + L2 normalize — shared semantics from `pooling`.
        self.pool_outputs(cur, batch)
    }

    /// The raw pre-pooling `[total_tokens, hidden]` hidden-state stream (row-major, ragged-
    /// concatenated) — the encoder's `last_hidden_state`. Every pooling mode is a reduction of
    /// this: `Cls` L2-normalizes each sequence's FIRST row, `LastToken` its last, `Mean` the
    /// row-average. Exposed so parity gates and callers can inspect per-token states directly.
    /// Runs inside the encoder-owned rayon pool.
    pub fn forward_hidden(&self, batch: &EncBatch) -> Result<Vec<f32>> {
        let cfg = &self.cfg;
        anyhow::ensure!(batch.n_seqs() > 0, "empty batch");
        let offset = match cfg.pos_kind {
            PosKind::Learned { offset } => offset,
            PosKind::Rope { .. } => 0,
        };
        for w in batch.seq_starts.windows(2) {
            let len = (w[1] - w[0]) as usize;
            anyhow::ensure!(len > 0, "empty sequence in batch");
            anyhow::ensure!(
                len + offset <= cfg.max_pos,
                "sequence of {len} tokens exceeds max positions {} (offset {offset})",
                cfg.max_pos
            );
        }
        self.pool.install(|| self.forward_hidden_inner(batch))
    }

    /// Embedding stage + transformer layers → the pre-pooling `[T, hidden]` stream. Assumes a
    /// validated batch (via `encode`/`forward_hidden`) and runs inside the encoder's rayon pool.
    fn forward_hidden_inner(&self, batch: &EncBatch) -> Result<Vec<f32>> {
        let cfg = &self.cfg;
        let (h, t) = (cfg.hidden, batch.tokens.len());
        let seq_of = row_to_seq(&batch.seq_starts, t);
        if let Some(v) = &batch.valid {
            anyhow::ensure!(
                v.len() == t,
                "validity mask length {} != tokens {t}",
                v.len()
            );
        }
        if let Some(ty) = &batch.type_ids {
            anyhow::ensure!(ty.len() == t, "type_ids length {} != tokens {t}", ty.len());
        }

        // Embedding stage: word (+ learned position, + per-segment token-type), then embedding norm.
        let mut cur = vec![0f32; t * h];
        cur.par_chunks_mut(h).enumerate().for_each(|(i, row)| {
            let tok = batch.tokens[i] as usize;
            row.copy_from_slice(&self.w.word[tok * h..(tok + 1) * h]);
            if let (Some(pos_table), PosKind::Learned { offset }) = (&self.w.pos, cfg.pos_kind) {
                let local = i - batch.seq_starts[seq_of[i]] as usize;
                for (r, p) in row
                    .iter_mut()
                    .zip(&pos_table[(offset + local) * h..(offset + local + 1) * h])
                {
                    *r += p;
                }
            }
            if let Some(tt) = &self.w.ttype {
                // Segment embedding: type row 0 (single text) or 1 (CrossEncoder second text).
                let ty = batch.type_ids.as_ref().map(|t| t[i] as usize).unwrap_or(0);
                for (r, v) in row.iter_mut().zip(&tt[ty * h..(ty + 1) * h]) {
                    *r += v;
                }
            }
        });
        if let Some(n) = &self.w.emb_norm {
            layer_norm_rows(
                &mut cur,
                h,
                n,
                cfg.eps,
                matches!(cfg.norm_kind, NormKind::RmsNorm),
            );
        }

        self.run_layers(cur, batch)
    }

    /// Transformer layers + trailing norm over a pre-built `[T, hidden]` stream (token
    /// embeddings, or vision patch embeddings — the layer stack is input-agnostic).
    fn run_layers(&self, mut cur: Vec<f32>, batch: &EncBatch) -> Result<Vec<f32>> {
        let cfg = &self.cfg;
        let (h, t) = (cfg.hidden, batch.seq_starts[batch.n_seqs()] as usize);
        let seq_of = row_to_seq(&batch.seq_starts, t);
        let valid = batch.valid.as_deref();
        let hd = cfg.head_dim;
        let nkv = cfg.n_kv_heads;
        let scale = 1.0 / (hd as f32).sqrt();
        let causal = cfg.attn_mask == MaskKind::Causal;
        let rms = matches!(cfg.norm_kind, NormKind::RmsNorm);
        let mut up_width = match cfg.mlp_kind {
            MlpKind::Dense { .. } => cfg.intermediate,
            MlpKind::Glu { .. } => 2 * cfg.intermediate,
        };
        if cfg.layer_is_attn.iter().any(|a| !a) {
            // Conv layers stage their [T, 3H] in_proj output in `mid`.
            up_width = up_width.max(3 * h);
        }
        // q spans all heads; k/v span the (possibly fewer) kv heads — GQA geometry.
        let qh = cfg.n_heads * hd;
        let kvh = nkv * hd;
        let mut attn_in = vec![0f32; t * h];
        let mut q = vec![0f32; t * qh];
        let mut k = vec![0f32; t * kvh];
        let mut v = vec![0f32; t * kvh];
        let mut ctx = vec![0f32; t * qh];
        let mut attn_out = vec![0f32; t * h];
        let mut mid = vec![0f32; t * up_width];
        let mut glu_h = vec![0f32; t * cfg.intermediate];
        let mut ffn_out = vec![0f32; t * h];
        let mut prof = StageProf::default();
        let pon = StageProf::on();
        let now = std::time::Instant::now;
        for (li, layer) in self.w.layers.iter().enumerate() {
            let window = cfg.layer_window[li];
            // Attention input: prenorm normalizes (layer 0 may skip), post-LN feeds `cur` raw.
            if cfg.prenorm {
                attn_in.copy_from_slice(&cur);
                if let Some(n) = &layer.attn_norm {
                    layer_norm_rows(&mut attn_in, h, n, cfg.eps, rms);
                }
            } else {
                attn_in.copy_from_slice(&cur);
            }
            match &layer.op {
                CpuOp::Attn {
                    q: qw,
                    k: kw,
                    v: vw,
                    o: ow,
                    q_norm,
                    k_norm,
                } => {
                    let _t = now();
                    linear(&mut q, &attn_in, qw);
                    linear(&mut k, &attn_in, kw);
                    linear(&mut v, &attn_in, vw);
                    prof.tick(pon, |p| &mut p.qkv, _t);
                    if let Some(w) = q_norm {
                        rms_norm_heads(&mut q, cfg.n_heads, hd, w, cfg.eps);
                    }
                    if let Some(w) = k_norm {
                        rms_norm_heads(&mut k, nkv, hd, w, cfg.eps);
                    }
                    if let PosKind::Rope { theta, local_theta } = cfg.pos_kind {
                        let th = if window > 0 { local_theta } else { theta };
                        apply_rope(&mut q, &batch.seq_starts, &seq_of, cfg.n_heads, hd, th);
                        apply_rope(&mut k, &batch.seq_starts, &seq_of, nkv, hd, th);
                    }
                    let _t = now();
                    match (&layer.rel, cfg.rel_attn) {
                        // DeBERTa: position lives in the SCORE, not the embedding.
                        (Some(rel), Some(spec)) => attention_disentangled(
                            &mut ctx,
                            &q,
                            &k,
                            &v,
                            &batch.seq_starts,
                            &seq_of,
                            cfg.n_heads,
                            hd,
                            rel,
                            spec,
                            valid,
                        ),
                        _ => attention_ragged(
                            &mut ctx,
                            &q,
                            &k,
                            &v,
                            &batch.seq_starts,
                            &seq_of,
                            cfg.n_heads,
                            nkv,
                            hd,
                            scale,
                            window,
                            causal,
                            valid,
                        ),
                    }
                    prof.tick(pon, |p| &mut p.attn, _t);
                    let _t = now();
                    linear(&mut attn_out, &ctx, ow);
                    prof.tick(pon, |p| &mut p.o, _t);
                }
                CpuOp::Conv {
                    in_proj,
                    conv_w,
                    out_proj,
                } => {
                    // bcx spans [T, 3H]; `mid` is the widest scratch (≥ 3H whenever a conv
                    // layer exists — enforced at load).
                    linear(&mut mid, &attn_in, in_proj);
                    conv_ragged(
                        &mut ctx,
                        &mid,
                        conv_w,
                        &batch.seq_starts,
                        &seq_of,
                        h,
                        cfg.conv_l,
                        valid,
                    );
                    linear(&mut attn_out, &ctx[..t * h], out_proj);
                }
            }
            let _t = now();
            if cfg.prenorm {
                add_into(&mut attn_out, &cur);
                // attn_out currently holds x + attn(x); it becomes the new residual stream.
                cur.copy_from_slice(&attn_out);
                attn_in.copy_from_slice(&cur);
                layer_norm_rows(&mut attn_in, h, &layer.mlp_norm, cfg.eps, rms);
            } else {
                let n = layer.attn_norm.as_ref().context("post-LN attn norm")?;
                layer_norm_rows_add(&mut attn_out, &cur, h, n, cfg.eps, rms);
                attn_in.copy_from_slice(&attn_out);
            }

            prof.tick(pon, |p| &mut p.norm, _t);

            let _t = now();
            let mlp_input = &attn_in;
            linear(&mut mid, mlp_input, &layer.up);
            let down_input: &[f32] = match cfg.mlp_kind {
                MlpKind::Dense { act, .. } => {
                    activate(&mut mid, act);
                    &mid
                }
                MlpKind::Glu { act } => {
                    glu_split(&mid, &mut glu_h, cfg.intermediate, act);
                    &glu_h
                }
            };
            linear_from(&mut ffn_out, down_input, &layer.down);
            prof.tick(pon, |p| &mut p.mlp, _t);
            let _t = now();
            if cfg.prenorm {
                add_into(&mut ffn_out, &cur);
                cur.copy_from_slice(&ffn_out);
            } else {
                layer_norm_rows_add(&mut ffn_out, &attn_in, h, &layer.mlp_norm, cfg.eps, rms);
                cur.copy_from_slice(&ffn_out);
            }
            prof.tick(pon, |p| &mut p.norm, _t);
        }
        if let Some(n) = &self.w.final_norm {
            layer_norm_rows(&mut cur, h, n, cfg.eps, rms);
        }
        prof.report();
        Ok(cur)
    }

    /// Reduce a `[T, hidden]` stream to the configured embedding output.
    fn pool_outputs(&self, cur: Vec<f32>, batch: &EncBatch) -> Result<EmbedOut> {
        let cfg = &self.cfg;
        let (h, t) = (cfg.hidden, batch.seq_starts[batch.n_seqs()] as usize);
        match cfg.pooling {
            Pooling::Mean => {
                let mut pooled = pooling::mean_pool(&cur, h, &batch.seq_starts);
                for e in &mut pooled {
                    pooling::l2_normalize(e);
                }
                Ok(EmbedOut::Pooled(pooled))
            }
            Pooling::Cls => {
                // The first token's hidden state ([CLS] for BERT-family sentence-transformers),
                // then L2 — no post-pooling head (the SigLIP `pooled_head` is a LastToken/EOS
                // projection; CLS embedders never ship one, and the GPU `cls_pool` twin applies
                // none, so keeping this head-free guarantees CPU/GPU parity by construction).
                let mut pooled: Vec<Vec<f32>> = batch
                    .seq_starts
                    .windows(2)
                    .map(|w| {
                        let first = w[0] as usize;
                        cur[first * h..(first + 1) * h].to_vec()
                    })
                    .collect();
                for e in &mut pooled {
                    pooling::l2_normalize(e);
                }
                Ok(EmbedOut::Pooled(pooled))
            }
            Pooling::LastToken => {
                let mut pooled: Vec<Vec<f32>> = batch
                    .seq_starts
                    .windows(2)
                    .map(|w| {
                        let last = (w[1] - 1) as usize;
                        cur[last * h..(last + 1) * h].to_vec()
                    })
                    .collect();
                if let Some(head) = &self.w.pooled_head {
                    pooled = pooled
                        .into_iter()
                        .map(|v| {
                            let mut out = vec![0f32; head.n];
                            linear_from(&mut out, &v, head);
                            out
                        })
                        .collect();
                }
                for e in &mut pooled {
                    pooling::l2_normalize(e);
                }
                Ok(EmbedOut::Pooled(pooled))
            }
            Pooling::PerToken { dim } => {
                let proj = self
                    .w
                    .proj
                    .as_ref()
                    .context("PerToken pooling requires the Dense projection")?;
                anyhow::ensure!(proj.n == dim, "projection rows {} != dim {dim}", proj.n);
                let mut projected = vec![0f32; t * dim];
                linear(&mut projected, &cur, proj);
                let mut out: Vec<Vec<Vec<f32>>> = Vec::with_capacity(batch.n_seqs());
                for w in batch.seq_starts.windows(2) {
                    let mut seq_vecs = Vec::with_capacity((w[1] - w[0]) as usize);
                    for i in w[0] as usize..w[1] as usize {
                        let mut v = projected[i * dim..(i + 1) * dim].to_vec();
                        pooling::l2_normalize(&mut v);
                        seq_vecs.push(v);
                    }
                    out.push(seq_vecs);
                }
                Ok(EmbedOut::PerToken(out))
            }
            Pooling::CrossEncoder { n_labels } => {
                // BertForSequenceClassification score: CLS (first token) → pooler.dense + tanh →
                // classifier → `n_labels` raw logits. No L2 normalization — it is a score, not an
                // embedding. One `[n_labels]` vector per input pair.
                let head = self
                    .w
                    .cross_encoder
                    .as_ref()
                    .context("CrossEncoder pooling requires the scoring head")?;
                let mut out: Vec<Vec<f32>> = Vec::with_capacity(batch.n_seqs());
                for w in batch.seq_starts.windows(2) {
                    let cls = &cur[w[0] as usize * h..(w[0] as usize + 1) * h];
                    let mut pooled = vec![0f32; head.pooler.n];
                    linear_from(&mut pooled, cls, &head.pooler);
                    for x in &mut pooled {
                        *x = x.tanh();
                    }
                    let mut logits = vec![0f32; n_labels];
                    linear_from(&mut logits, &pooled, &head.classifier);
                    out.push(logits);
                }
                Ok(EmbedOut::Pooled(out))
            }
            other => anyhow::bail!("pooling {other:?} lands with its phase"),
        }
    }
}

/// Row index → sequence index lookup (ragged batches; rows of one sequence are contiguous).
fn row_to_seq(seq_starts: &[u32], t: usize) -> Vec<usize> {
    let mut map = vec![0usize; t];
    for (s, w) in seq_starts.windows(2).enumerate() {
        for r in w[0]..w[1] {
            map[r as usize] = s;
        }
    }
    map
}

/// `y = x · Wᵀ (+ b)` over the whole token matrix.
fn linear(y: &mut [f32], x: &[f32], l: &CpuLinear) {
    linear_from(y, x, l);
}

/// `linear` with an explicit input slice (used where the input buffer differs from the layer's
/// nominal stream).
///
/// On aarch64 this runs [`crate::cpu_gemm`], whose weight is packed ONCE (lazily, on first use) and
/// reused forever. Elsewhere it stays on the `gemm` crate — which repacks the weight on every call,
/// and is why the packed path exists at all. Both are covered by the same parity gates.
// The `return` is the aarch64 arm of a cfg-gated pair; clippy only ever sees one arm.
#[allow(clippy::needless_return)]
fn linear_from(y: &mut [f32], x: &[f32], l: &CpuLinear) {
    let (n, k) = (l.n, l.k);
    let m = x.len() / k;
    assert_eq!(x.len(), m * k, "lhs shape");
    assert_eq!(l.w.len(), n * k, "weight shape");
    let y = &mut y[..m * n];

    #[cfg(target_arch = "aarch64")]
    {
        // INT8 opt-in (i8mm hosts): ~4× the MAC rate of the f32 FMA path. Dynamic per-row
        // activation quant — NOT parity-class, gated by tests/cpu_i8.rs's cosine contract.
        if std::env::var("OSFKB_CPU_I8").ok().as_deref() == Some("1")
            && crate::cpu_gemm::i8mm_available()
        {
            let packed = l
                .packed_i8
                .get_or_init(|| crate::cpu_gemm::PackedWeightI8::new(&l.w, n, k));
            crate::cpu_gemm::gemm_i8(y, x, packed, m, l.b.as_deref());
            return;
        }
        // F16-storage opt-in (FEAT_FHM hosts): FMLAL widens f16 pairs into f32 accumulators —
        // half the W+A traffic on a GEMM whose multicore ceiling IS that traffic (prototype:
        // 1.30–1.40× at the bandwidth-bound shapes, neutral where W is cache-resident). NOT
        // parity-class: f16 storage rounding, gated by tests/cpu_f16.rs's cosine contract.
        // Listed after int8 so an explicit OSFKB_CPU_I8=1 keeps winning if both are set.
        if std::env::var("OSFKB_CPU_F16").ok().as_deref() == Some("1")
            && crate::cpu_gemm::fhm_available()
        {
            let packed = l
                .packed_f16
                .get_or_init(|| crate::cpu_gemm::PackedWeightF16::new(&l.w, n, k));
            crate::cpu_gemm::gemm_f16(y, x, packed, m, l.b.as_deref());
            return;
        }
        let packed = l
            .packed
            .get_or_init(|| crate::cpu_gemm::PackedWeight::new(&l.w, n, k));
        crate::cpu_gemm::gemm_packed(y, x, packed, m, l.b.as_deref());
        return;
    }

    #[cfg(not(target_arch = "aarch64"))]
    {
        // SAFETY: the asserts above establish exactly the bounds this dereferences.
        unsafe {
            gemm::gemm(
                m,
                n,
                k,
                y.as_mut_ptr(),
                1,
                n as isize,
                false,
                x.as_ptr(),
                1,
                k as isize,
                l.w.as_ptr(),
                k as isize,
                1,
                0.0f32,
                1.0f32,
                false,
                false,
                false,
                gemm::Parallelism::Rayon(rayon::current_num_threads()),
            );
        }
        if let Some(b) = &l.b {
            y.par_chunks_mut(n).for_each(|row| {
                for (r, bi) in row.iter_mut().zip(b) {
                    *r += bi;
                }
            });
        }
    }
}

/// LFM2 gated short-conv over a ragged batch, stateless full-sequence form: for token row `i`
/// in sequence `s`, `y[i,c] = C(i,c) · Σ_k w[c,k] · bx(i − (l−1−k), c)` with `bx(j,c) =
/// B(j,c)·x(j,c)` and zero history before the sequence start (state never leaks across
/// sequences). `bcx` is the `[T, 3H]` in_proj output laid out `B | C | x` per token.
#[allow(clippy::too_many_arguments)]
fn conv_ragged(
    y: &mut [f32],
    bcx: &[f32],
    conv_w: &[f32],
    seq_starts: &[u32],
    seq_of: &[usize],
    h: usize,
    l: usize,
    valid: Option<&[u32]>,
) {
    let w3 = 3 * h;
    y[..seq_of.len() * h]
        .par_chunks_mut(h)
        .enumerate()
        .for_each(|(i, out)| {
            // HF zeroes the conv-block INPUT at padded positions before in_proj; with bias-free
            // projections that equals: invalid position ⇒ its Bx taps are 0 AND its own C gate
            // is 0 (output row zero).
            if valid.is_some_and(|v| v[i] == 0) {
                out.fill(0.0);
                return;
            }
            let a = seq_starts[seq_of[i]] as usize;
            for c in 0..h {
                let cc = bcx[i * w3 + h + c];
                let mut acc = 0.0f32;
                for k in 0..l {
                    let back = l - 1 - k;
                    if i >= a + back {
                        let j = i - back;
                        if valid.is_some_and(|v| v[j] == 0) {
                            continue;
                        }
                        let bx = bcx[j * w3 + c] * bcx[j * w3 + 2 * h + c];
                        acc += conv_w[c * l + k] * bx;
                    }
                }
                out[c] = cc * acc;
            }
        });
}

/// Rotary embedding in place, HF `rotate_half` convention: pairs `(j, j + hd/2)`, frequency
/// `theta^(-2j/hd)`, angle `p·freq` with `p` LOCAL to the sequence. Mask-agnostic — the same
/// rotation serves bidirectional and causal attention.
fn apply_rope(
    x: &mut [f32],
    seq_starts: &[u32],
    seq_of: &[usize],
    n_heads: usize,
    hd: usize,
    theta: f32,
) {
    let h = n_heads * hd;
    let half = hd / 2;
    x.par_chunks_mut(h).enumerate().for_each(|(i, row)| {
        let p = (i - seq_starts[seq_of[i]] as usize) as f32;
        for head in 0..n_heads {
            let base = head * hd;
            for j in 0..half {
                let freq = theta.powf(-2.0 * j as f32 / hd as f32);
                let (sin, cos) = (p * freq).sin_cos();
                let a = row[base + j];
                let b = row[base + j + half];
                row[base + j] = a * cos - b * sin;
                row[base + j + half] = a * sin + b * cos;
            }
        }
    });
}

/// Dense attention over a ragged batch: bidirectional or `causal` (keys `≤ i`), with `window > 0`
/// restricting keys to `|i − j| ≤ window/2` (ModernBERT local span semantics), and GQA — q spans
/// `n_heads`, k/v span `n_kv_heads` (each kv head serves `n_heads / n_kv_heads` q heads).
/// Row-parallel, sequential accumulation per output row → deterministic.
#[allow(clippy::too_many_arguments)]
fn attention_ragged(
    ctx: &mut [f32],
    q: &[f32],
    k: &[f32],
    v: &[f32],
    seq_starts: &[u32],
    seq_of: &[usize],
    n_heads: usize,
    n_kv_heads: usize,
    hd: usize,
    scale: f32,
    window: u32,
    causal: bool,
    valid: Option<&[u32]>,
) {
    let qh = n_heads * hd;
    let kvh = n_kv_heads * hd;
    let group = n_heads / n_kv_heads;
    ctx.par_chunks_mut(qh).enumerate().for_each(|(i, out)| {
        let s = seq_of[i];
        let (a, b) = (seq_starts[s] as usize, seq_starts[s + 1] as usize);
        let (mut ja, mut jb) = (a, b);
        if causal {
            jb = jb.min(i + 1);
        }
        if window > 0 {
            let half_w = (window / 2) as usize;
            ja = ja.max(i.saturating_sub(half_w));
            jb = jb.min(i + half_w + 1);
        }
        let mut scores = vec![0f32; jb - ja];
        for head in 0..n_heads {
            let kv_head = head / group;
            let qo = i * qh + head * hd;
            let qi = &q[qo..qo + hd];
            let mut mx = f32::NEG_INFINITY;
            for (jj, sc) in scores.iter_mut().enumerate() {
                // Invalid tokens (ColBERT expansion pads) are never attention KEYS.
                if valid.is_some_and(|v| v[ja + jj] == 0) {
                    *sc = f32::NEG_INFINITY;
                    continue;
                }
                let ko = (ja + jj) * kvh + kv_head * hd;
                *sc = dot(qi, &k[ko..ko + hd]) * scale;
                mx = mx.max(*sc);
            }
            // Vectorized softmax middle: masked (−inf) scores stay EXACTLY zero, so the
            // degenerate all-masked row still lands in the denom == 0 arm below.
            let denom = crate::simd_math::exp_sub_sum(&mut scores, mx);
            // All keys invalid (can only happen with a degenerate mask): zero output, no NaN.
            let inv = if denom > 0.0 { 1.0 / denom } else { 0.0 };
            let dst = &mut out[head * hd..(head + 1) * hd];
            dst.fill(0.0);
            for (jj, sc) in scores.iter().enumerate() {
                let vo = (ja + jj) * kvh + kv_head * hd;
                axpy(dst, &v[vo..vo + hd], sc * inv);
            }
        }
    });
}

/// DeBERTa disentangled attention — the content term plus `c2p` and `p2c`, all sharing one scale.
///
/// For query `i` and key `j` (positions are LOCAL to their sequence, so a ragged batch behaves
/// exactly like the sequences run solo):
///
/// ```text
/// score = ( q_i·k_j  +  q_i·Kr[m]  +  k_j·Qr[m] ) / sqrt(head_dim · scale_factor)
///   m = clamp(bucket(i - j) + span, 0, 2·span - 1)
/// ```
///
/// The two relative terms share the index `m` because the bucket function is odd — see
/// [`RelAttn`](crate::encoder_weights::RelAttn), which derives that from HF's asymmetric-looking
/// `c2p_pos` / `-p2c_pos` gathers. `Kr`/`Qr` are this layer's [`CpuRelPos`].
///
/// No GQA (`n_kv_heads == n_heads` on every DeBERTa) and no windowing — full bidirectional
/// attention within each sequence.
#[allow(clippy::too_many_arguments)]
/// Shortest sequence for which the GEMM form of [`attention_disentangled`] is worth its overheads.
///
/// Below it the scalar twin wins and the GEMM form is a real regression, for two reasons that both
/// vanish with length: the position LUTs span the *reachable bucket range*, which is wider than a
/// short sequence (so the restructure adds more work than it saves), and there are too few row
/// blocks to fill the cores. Above it the c2c GEMM's reuse dominates and the win grows with T²:
///
/// ```text
///   T:       32     64    128    192    256    384    512
///   speedup: 0.67x  0.80x  1.01x  1.09x  1.47x  1.93x  2.46x
/// ```
///
/// Measured, not guessed — `gemm_form_beats_the_scalar_kernel` reprints that row and asserts the win
/// above this cutover, so a future kernel change that moves the crossover fails the test.
const GEMM_ATTN_MIN_T: usize = 192;

/// DeBERTa disentangled attention, GEMM form.
///
/// The scalar twin below ([`attention_disentangled_scalar`]) computes the same thing and is the
/// reference this is tested against. It loses because it has **no reuse**: every (query, key, head)
/// re-streams a 64-float key from memory for a 2-FLOP-per-8-byte dot, which measures ~16 GF/s
/// against a ~120 GF/s NEON roofline — bandwidth-bound, not ALU-bound. Making the dot itself faster
/// cannot fix that; only reuse can.
///
/// So the content term becomes a real GEMM over a block of query rows, where one streamed K tile
/// serves `RBLK` queries instead of one. The position terms stay gathers, but from a per-head
/// `[T, buckets]` LUT built by GEMM once — ~1 MB, L2-resident, and the same slab serves every query.
///
/// What this is NOT: an attempt to save FLOPs by hoisting the position terms. That was the first
/// design, and it was wrong — hoisting c2p/p2c costs `T × 2·span` where computing them inline costs
/// `T × T`, so for every sequence this encoder actually serves (T ≤ 512 = 2·span) it does MORE work,
/// not less, and it duly shipped a 25% regression. The win here is memory hierarchy, not arithmetic.
#[allow(clippy::too_many_arguments)]
fn attention_disentangled(
    ctx: &mut [f32],
    q: &[f32],
    k: &[f32],
    v: &[f32],
    seq_starts: &[u32],
    seq_of: &[usize],
    n_heads: usize,
    hd: usize,
    rel: &CpuRelPos,
    spec: RelAttn,
    valid: Option<&[u32]>,
) {
    let longest = seq_starts
        .windows(2)
        .map(|w| (w[1] - w[0]) as usize)
        .max()
        .unwrap_or(0);
    if longest < GEMM_ATTN_MIN_T || std::env::var("OSFKB_ENC_ATTN_SCALAR").is_ok_and(|v| v != "0") {
        return attention_disentangled_scalar(
            ctx, q, k, v, seq_starts, seq_of, n_heads, hd, rel, spec, valid,
        );
    }
    attention_disentangled_gemm(ctx, q, k, v, seq_starts, n_heads, hd, rel, spec, valid)
}

/// The GEMM form itself, reached through [`attention_disentangled`] above the length cutover.
#[allow(clippy::too_many_arguments)]
fn attention_disentangled_gemm(
    ctx: &mut [f32],
    q: &[f32],
    k: &[f32],
    v: &[f32],
    seq_starts: &[u32],
    n_heads: usize,
    hd: usize,
    rel: &CpuRelPos,
    spec: RelAttn,
    valid: Option<&[u32]>,
) {
    /// Query rows per task. Sized so the tasks (`heads × T/RBLK`) outnumber the cores while each
    /// still amortizes a streamed K tile over enough rows to leave the bandwidth wall.
    const RBLK: usize = 32;

    let qh = n_heads * hd;
    let scale = 1.0 / ((hd * spec.scale_factor()) as f32).sqrt();

    for s in 0..seq_starts.len().saturating_sub(1) {
        let (a, b) = (seq_starts[s] as usize, seq_starts[s + 1] as usize);
        let tn = b - a;
        if tn == 0 {
            continue;
        }
        // `index()` is a function of (i - j) alone, so the whole bucket map is one small vector indexed
        // by `i - j + tn - 1` — L1-resident, versus a T² table.
        let bidx: Vec<u32> = (0..2 * tn - 1)
            .map(|d| spec.index(d, tn - 1) as u32)
            .collect();
        let (m_lo, m_hi) = (
            *bidx.iter().min().unwrap() as usize,
            *bidx.iter().max().unwrap() as usize,
        );
        let bn = m_hi - m_lo + 1;

        // p2c LUT: [head][key j, bucket] = k_j · pos_q[bucket]. Depends on the KEY, so it cannot be
        // folded into the per-row-block GEMM — build it once per head and let every block gather.
        let mut p2c = vec![0f32; if spec.p2c { n_heads * tn * bn } else { 0 }];
        if spec.p2c {
            p2c.par_chunks_mut(tn * bn)
                .enumerate()
                .for_each(|(h, slab)| {
                    // SAFETY: disjoint per-head output slab; inputs are read-only. Strides map the
                    // per-head views of the [T, qh] / [buckets, qh] tables (element[r][c] = r·rs + c·cs).
                    unsafe {
                        gemm::gemm(
                            tn,
                            bn,
                            hd,
                            slab.as_mut_ptr(),
                            1,
                            bn as isize,
                            false,
                            k.as_ptr().add(a * qh + h * hd),
                            1,
                            qh as isize,
                            rel.pos_q.as_ptr().add(m_lo * qh + h * hd),
                            qh as isize,
                            1,
                            0.0,
                            1.0,
                            false,
                            false,
                            false,
                            gemm::Parallelism::None,
                        );
                    }
                });
        }

        let seq_ctx = &mut ctx[a * qh..b * qh];
        seq_ctx
            .par_chunks_mut(RBLK * qh)
            .enumerate()
            .for_each(|(blk, out)| {
                let r0 = blk * RBLK;
                let rb = out.len() / qh; // the last block is ragged
                let mut scores = vec![0f32; rb * tn];
                let mut c2p = vec![0f32; if spec.c2p { rb * bn } else { 0 }];
                for h in 0..n_heads {
                    // SAFETY: `scores`/`c2p` are this task's own buffers; `out` is a disjoint row
                    // block; q/k/v/rel are read-only. m/n/k and strides are checked against the
                    // scalar twin by `gemm_form_matches_the_scalar_reference`.
                    unsafe {
                        // content → content: Q_blk [rb, hd] × K^T [hd, tn]
                        gemm::gemm(
                            rb,
                            tn,
                            hd,
                            scores.as_mut_ptr(),
                            1,
                            tn as isize,
                            false,
                            q.as_ptr().add((a + r0) * qh + h * hd),
                            1,
                            qh as isize,
                            k.as_ptr().add(a * qh + h * hd),
                            qh as isize,
                            1,
                            0.0,
                            1.0,
                            false,
                            false,
                            false,
                            gemm::Parallelism::None,
                        );
                        if spec.c2p {
                            // content query → position key: Q_blk [rb, hd] × pos_k^T [hd, bn]
                            gemm::gemm(
                                rb,
                                bn,
                                hd,
                                c2p.as_mut_ptr(),
                                1,
                                bn as isize,
                                false,
                                q.as_ptr().add((a + r0) * qh + h * hd),
                                1,
                                qh as isize,
                                rel.pos_k.as_ptr().add(m_lo * qh + h * hd),
                                qh as isize,
                                1,
                                0.0,
                                1.0,
                                false,
                                false,
                                false,
                                gemm::Parallelism::None,
                            );
                        }
                    }
                    for ii in 0..rb {
                        let i = r0 + ii;
                        let row = &mut scores[ii * tn..(ii + 1) * tn];
                        let mut mx = f32::NEG_INFINITY;
                        for (j, sc) in row.iter_mut().enumerate() {
                            if valid.is_some_and(|v| v[a + j] == 0) {
                                *sc = f32::NEG_INFINITY;
                                continue;
                            }
                            let m = bidx[i + tn - 1 - j] as usize - m_lo;
                            let mut acc = *sc;
                            if spec.c2p {
                                acc += c2p[ii * bn + m];
                            }
                            if spec.p2c {
                                acc += p2c[h * tn * bn + j * bn + m];
                            }
                            *sc = acc * scale;
                            mx = mx.max(*sc);
                        }
                        let denom = crate::simd_math::exp_sub_sum(row, mx);
                        let inv = if denom > 0.0 { 1.0 / denom } else { 0.0 };
                        for sc in row.iter_mut() {
                            *sc *= inv;
                        }
                    }
                    // P [rb, tn] × V [tn, hd] straight into this block's head slice of ctx.
                    // SAFETY: as above — `out` rows are this task's alone.
                    unsafe {
                        gemm::gemm(
                            rb,
                            hd,
                            tn,
                            out.as_mut_ptr().add(h * hd),
                            1,
                            qh as isize,
                            false,
                            scores.as_ptr(),
                            1,
                            tn as isize,
                            v.as_ptr().add(a * qh + h * hd),
                            1,
                            qh as isize,
                            0.0,
                            1.0,
                            false,
                            false,
                            false,
                            gemm::Parallelism::None,
                        );
                    }
                }
            });
    }
}

/// The pre-GEMM scalar kernel: kept as the reference oracle for the fast path (and as an A/B switch
/// via `OSFKB_ENC_ATTN_SCALAR=1`, which is how the fast path's speedup is measured on a box whose
/// load average swings 2.8× between runs — the arms have to run in the same process to be comparable).
#[allow(clippy::too_many_arguments)]
fn attention_disentangled_scalar(
    ctx: &mut [f32],
    q: &[f32],
    k: &[f32],
    v: &[f32],
    seq_starts: &[u32],
    seq_of: &[usize],
    n_heads: usize,
    hd: usize,
    rel: &CpuRelPos,
    spec: RelAttn,
    valid: Option<&[u32]>,
) {
    let qh = n_heads * hd;
    let scale = 1.0 / ((hd * spec.scale_factor()) as f32).sqrt();
    ctx.par_chunks_mut(qh).enumerate().for_each(|(i, out)| {
        let s = seq_of[i];
        let (a, b) = (seq_starts[s] as usize, seq_starts[s + 1] as usize);
        // Relative distance is measured within the sequence, not across the packed batch.
        let qi_local = i - a;
        let mut scores = vec![0f32; b - a];
        // Row of the relative table each key uses — identical for every head, so hoist it.
        let rows: Vec<usize> = (a..b).map(|j| spec.index(qi_local, j - a)).collect();
        for head in 0..n_heads {
            let qo = i * qh + head * hd;
            let qi = &q[qo..qo + hd];
            let mut mx = f32::NEG_INFINITY;
            for (jj, sc) in scores.iter_mut().enumerate() {
                if valid.is_some_and(|v| v[a + jj] == 0) {
                    *sc = f32::NEG_INFINITY;
                    continue;
                }
                let ko = (a + jj) * qh + head * hd;
                let kj = &k[ko..ko + hd];
                let ro = rows[jj] * qh + head * hd;

                // content → content
                let mut acc = dot(qi, kj);
                if spec.c2p {
                    // content query → position key
                    acc += dot(qi, &rel.pos_k[ro..ro + hd]);
                }
                if spec.p2c {
                    // position query → content key
                    acc += dot(kj, &rel.pos_q[ro..ro + hd]);
                }
                *sc = acc * scale;
                mx = mx.max(*sc);
            }
            let mut denom = 0f32;
            for sc in scores.iter_mut() {
                *sc = (*sc - mx).exp();
                denom += *sc;
            }
            let inv = if denom > 0.0 { 1.0 / denom } else { 0.0 };
            let dst = &mut out[head * hd..(head + 1) * hd];
            dst.fill(0.0);
            for (jj, sc) in scores.iter().enumerate() {
                let vo = (a + jj) * qh + head * hd;
                axpy(dst, &v[vo..vo + hd], sc * inv);
            }
        }
    });
}

/// Row-wise norm in place: LayerNorm `(x - mean)/sqrt(var + eps)·w (+ b)` (biased variance) or,
/// with `rms`, RMSNorm `x/sqrt(mean(x²) + eps)·w` (no centering, no bias).
/// Fused `x += res` + LayerNorm — the post-LN residual pair (`add_into` then
/// [`layer_norm_rows`]) in ONE parallel sweep structure. The add folds into the mean pass, so
/// `res` is read once and `x` sees three passes instead of the pair's five; the accumulation
/// ORDER is byte-for-byte the unfused pair's (add store, then the same sequential mean / two-pass
/// variance / normalize), so the fusion is bitwise-neutral — required, the CPU path is the oracle.
fn layer_norm_rows_add(x: &mut [f32], res: &[f32], h: usize, norm: &CpuNorm, eps: f32, rms: bool) {
    x.par_chunks_mut(h)
        .zip(res.par_chunks(h))
        .for_each(|(row, rrow)| {
            let mut sum = 0f32;
            for (v, r) in row.iter_mut().zip(rrow) {
                *v += r;
                sum += *v;
            }
            let mean: f32 = if rms { 0.0 } else { sum / h as f32 };
            let var: f32 = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / h as f32;
            let inv = 1.0 / (var + eps).sqrt();
            for (i, r) in row.iter_mut().enumerate() {
                *r = (*r - mean) * inv * norm.w[i];
                if let Some(b) = &norm.b {
                    *r += b[i];
                }
            }
        });
}

fn layer_norm_rows(x: &mut [f32], h: usize, norm: &CpuNorm, eps: f32, rms: bool) {
    x.par_chunks_mut(h).for_each(|row| {
        let mean: f32 = if rms {
            0.0
        } else {
            row.iter().sum::<f32>() / h as f32
        };
        let var: f32 = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / h as f32;
        let inv = 1.0 / (var + eps).sqrt();
        for (i, r) in row.iter_mut().enumerate() {
            *r = (*r - mean) * inv * norm.w[i];
            if let Some(b) = &norm.b {
                *r += b[i];
            }
        }
    });
}

/// Per-head RMS normalization over `[T, n_heads·hd]` with a shared `[hd]` weight (Qwen3 qk-norm,
/// applied BEFORE rotary).
fn rms_norm_heads(x: &mut [f32], n_heads: usize, hd: usize, w: &[f32], eps: f32) {
    let h = n_heads * hd;
    x.par_chunks_mut(h).for_each(|row| {
        for head in 0..n_heads {
            let s = &mut row[head * hd..(head + 1) * hd];
            let ms: f32 = s.iter().map(|v| v * v).sum::<f32>() / hd as f32;
            let inv = 1.0 / (ms + eps).sqrt();
            for (v, wi) in s.iter_mut().zip(w) {
                *v *= inv * wi;
            }
        }
    });
}

/// Elementwise `dst += src`.
/// Per-stage wall-clock inside the encoder layer loop, enabled by `OSFKB_ENC_PROFILE=1`.
///
/// Exists because the CPU optimisation order is only defensible if the split between *attention*
/// and the *projection/FFN GEMMs* is measured. The first pass of this campaign guessed attention
/// was the lever, restructured it, and shipped a 25% REGRESSION — the arithmetic that justified the
/// restructure was only true for sequences far longer than any this encoder serves.
#[derive(Default)]
struct StageProf {
    qkv: f64,
    attn: f64,
    o: f64,
    mlp: f64,
    norm: f64,
}

impl StageProf {
    fn on() -> bool {
        std::env::var("OSFKB_ENC_PROFILE").is_ok_and(|v| v != "0")
    }
    #[inline]
    fn tick(&mut self, on: bool, slot: fn(&mut Self) -> &mut f64, t: std::time::Instant) {
        if on {
            *slot(self) += t.elapsed().as_secs_f64();
        }
    }
    fn report(&self) {
        let tot = self.qkv + self.attn + self.o + self.mlp + self.norm;
        if tot <= 0.0 {
            return;
        }
        let pct = |x: f64| 100.0 * x / tot;
        eprintln!(
            "  [enc cpu] total {:6.1} ms | qkv {:5.1} ({:4.1}%) attn {:5.1} ({:4.1}%) \
o {:5.1} ({:4.1}%) mlp {:5.1} ({:4.1}%) norm+res {:5.1} ({:4.1}%)",
            tot * 1e3,
            self.qkv * 1e3,
            pct(self.qkv),
            self.attn * 1e3,
            pct(self.attn),
            self.o * 1e3,
            pct(self.o),
            self.mlp * 1e3,
            pct(self.mlp),
            self.norm * 1e3,
            pct(self.norm),
        );
    }
}

/// The attention kernels run one of these per (query, key, head) — at 512 tokens × 12 heads that is
/// ~1.2 G MACs of exactly this shape, so it is where the encoder's CPU time actually goes. See
/// [`crate::simd`] for why the idiomatic `.sum()` spelling gives up ~4× before it starts.
#[inline]
pub(crate) fn dot(a: &[f32], b: &[f32]) -> f32 {
    crate::simd::dot(a, b)
}

/// `dst += p · src` — the attention value accumulation, same vectorization argument as [`dot`].
#[inline]
/// `dst += p · src`. On aarch64 this is `vfmaq_n_f32` — safe Rust never emits the fused
/// multiply-add (no FP contraction), and the attention P·V accumulate lives on this function.
/// FMA vs mul+add moves last bits only (tolerance-class, same as `simd::dot`).
fn axpy(dst: &mut [f32], src: &[f32], p: f32) {
    debug_assert_eq!(dst.len(), src.len());
    #[cfg(target_arch = "aarch64")]
    {
        use std::arch::aarch64::*;
        let n4 = dst.len() / 4 * 4;
        // SAFETY: 4-lane loads/stores strictly below n4 ≤ len on both slices.
        unsafe {
            let (pd, ps) = (dst.as_mut_ptr(), src.as_ptr());
            let mut i = 0;
            while i < n4 {
                vst1q_f32(
                    pd.add(i),
                    vfmaq_n_f32(vld1q_f32(pd.add(i)), vld1q_f32(ps.add(i)), p),
                );
                i += 4;
            }
        }
        for (d, s) in dst[n4..].iter_mut().zip(&src[n4..]) {
            *d += p * s;
        }
        return;
    }
    #[cfg(not(target_arch = "aarch64"))]
    {
        let mut cd = dst.chunks_exact_mut(4);
        let mut cs = src.chunks_exact(4);
        for (d, s) in cd.by_ref().zip(cs.by_ref()) {
            for l in 0..4 {
                d[l] += p * s[l];
            }
        }
        for (d, s) in cd.into_remainder().iter_mut().zip(cs.remainder()) {
            *d += p * s;
        }
    }
}

fn add_into(dst: &mut [f32], src: &[f32]) {
    dst.par_chunks_mut(4096)
        .zip(src.par_chunks(4096))
        .for_each(|(d, s)| {
            for (di, si) in d.iter_mut().zip(s) {
                *di += si;
            }
        });
}

/// Map the config's activation onto the vectorized module's kind (1:1 by construction).
fn act_kind(act: Act) -> crate::simd_math::ActKind {
    use crate::simd_math::ActKind as K;
    match act {
        Act::GeluErf => K::GeluErf,
        Act::GeluTanh => K::GeluTanh,
        Act::Silu => K::Silu,
        Act::Tanh => K::Tanh,
        Act::Relu => K::Relu,
    }
}

/// Elementwise activation in place, vectorized (`simd_math`). Erf-GELU now runs the SAME A&S
/// 7.1.26 approximation the GPU's `ACT_FNS` uses (|ε| ≤ 1.5e-7 — far inside the 1e-4 parity
/// gates); the old scalar-libm pass measured 6.1 ns/elem, whole milliseconds of an erf-GELU
/// encode.
fn activate(x: &mut [f32], act: Act) {
    let k = act_kind(act);
    x.par_chunks_mut(4096)
        .for_each(|chunk| crate::simd_math::act_slice(chunk, k));
}

/// GLU split: `mid` is `[T, 2I]` per token as `input | gate`; `out[t] = act(input) ⊙ gate`
/// (the HF ModernBERT `chunk(2)` convention; nomic's `silu(fc11)·fc12` concatenates to the
/// identical layout with `fc11` first). Vectorized as copy → `act_slice` → elementwise multiply.
fn glu_split(mid: &[f32], out: &mut [f32], i_width: usize, act: Act) {
    let k = act_kind(act);
    out.par_chunks_mut(i_width)
        .zip(mid.par_chunks(2 * i_width))
        .for_each(|(o, m)| {
            o.copy_from_slice(&m[..i_width]);
            crate::simd_math::act_slice(o, k);
            for (oj, g) in o.iter_mut().zip(&m[i_width..]) {
                *oj *= g;
            }
        });
}

/// BERT/XLM-R tensor-name table → f32 weights. Bare sentence-transformers checkpoints name
/// tensors `embeddings.*`/`encoder.*`; task-headed exports prefix `bert.`/`roberta.` — both load.
fn load_bert_family(st: &LazySt, cfg: &EncoderConfig) -> Result<CpuWeights> {
    let prefix = ["", "bert.", "roberta."]
        .into_iter()
        .find(|p| st.has(&format!("{p}embeddings.word_embeddings.weight")))
        .context("word embeddings not found under known prefixes ('', 'bert.', 'roberta.')")?;
    let get = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(&format!("{prefix}{name}")) };
    let word = get("embeddings.word_embeddings.weight")?;
    anyhow::ensure!(
        word.len() == cfg.vocab * cfg.hidden,
        "word embedding shape {} != vocab {} × hidden {}",
        word.len(),
        cfg.vocab,
        cfg.hidden
    );
    let ttype = if cfg.type_vocab > 0 {
        Some(get("embeddings.token_type_embeddings.weight")?)
    } else {
        None
    };
    anyhow::ensure!(
        matches!(cfg.norm_kind, NormKind::LayerNorm { bias: true }),
        "BERT-family norms are biased LayerNorm"
    );
    let lin = |w: Vec<f32>, b: Vec<f32>, n: usize, k: usize| CpuLinear {
        w,
        b: Some(b),
        n,
        k,
        ..Default::default()
    };
    let (h, im) = (cfg.hidden, cfg.intermediate);
    // jina-reranker-v2 ships a "jina-flash" XLM-RoBERTa: identical math (absolute positions,
    // post-LayerNorm, dense GELU) but a FUSED QKV (`encoder.layers.N.mixer.Wqkv`, `[3h,h]`) and
    // renamed tensors (`mixer.out_proj`, `mlp.fc1/fc2`, `norm1/norm2`, `emb_ln`). Detect it and
    // remap, splitting Wqkv row-wise into q|k|v. Everything downstream is the standard bert family.
    let jina = st.has(&format!("{prefix}encoder.layers.0.mixer.Wqkv.weight"));
    let mut layers = Vec::with_capacity(cfg.n_layers);
    for i in 0..cfg.n_layers {
        let (attn_norm, op, mlp_norm, up, down) = if jina {
            let lp = format!("encoder.layers.{i}");
            // Wqkv `[3h,h]` row-major (PyTorch `[out,in]`): q=rows[0,h) k=[h,2h) v=[2h,3h); bias `[3h]`.
            let wqkv = get(&format!("{lp}.mixer.Wqkv.weight"))?;
            let bqkv = get(&format!("{lp}.mixer.Wqkv.bias"))?;
            let w = |r: usize| wqkv[r * h * h..(r + 1) * h * h].to_vec();
            let b = |r: usize| bqkv[r * h..(r + 1) * h].to_vec();
            (
                Some(CpuNorm {
                    w: get(&format!("{lp}.norm1.weight"))?,
                    b: Some(get(&format!("{lp}.norm1.bias"))?),
                }),
                CpuOp::Attn {
                    q: lin(w(0), b(0), h, h),
                    k: lin(w(1), b(1), h, h),
                    v: lin(w(2), b(2), h, h),
                    o: lin(
                        get(&format!("{lp}.mixer.out_proj.weight"))?,
                        get(&format!("{lp}.mixer.out_proj.bias"))?,
                        h,
                        h,
                    ),
                    q_norm: None,
                    k_norm: None,
                },
                CpuNorm {
                    w: get(&format!("{lp}.norm2.weight"))?,
                    b: Some(get(&format!("{lp}.norm2.bias"))?),
                },
                lin(
                    get(&format!("{lp}.mlp.fc1.weight"))?,
                    get(&format!("{lp}.mlp.fc1.bias"))?,
                    im,
                    h,
                ),
                lin(
                    get(&format!("{lp}.mlp.fc2.weight"))?,
                    get(&format!("{lp}.mlp.fc2.bias"))?,
                    h,
                    im,
                ),
            )
        } else {
            let p = format!("encoder.layer.{i}");
            (
                Some(CpuNorm {
                    w: get(&format!("{p}.attention.output.LayerNorm.weight"))?,
                    b: Some(get(&format!("{p}.attention.output.LayerNorm.bias"))?),
                }),
                CpuOp::Attn {
                    q: lin(
                        get(&format!("{p}.attention.self.query.weight"))?,
                        get(&format!("{p}.attention.self.query.bias"))?,
                        h,
                        h,
                    ),
                    k: lin(
                        get(&format!("{p}.attention.self.key.weight"))?,
                        get(&format!("{p}.attention.self.key.bias"))?,
                        h,
                        h,
                    ),
                    v: lin(
                        get(&format!("{p}.attention.self.value.weight"))?,
                        get(&format!("{p}.attention.self.value.bias"))?,
                        h,
                        h,
                    ),
                    o: lin(
                        get(&format!("{p}.attention.output.dense.weight"))?,
                        get(&format!("{p}.attention.output.dense.bias"))?,
                        h,
                        h,
                    ),
                    q_norm: None,
                    k_norm: None,
                },
                CpuNorm {
                    w: get(&format!("{p}.output.LayerNorm.weight"))?,
                    b: Some(get(&format!("{p}.output.LayerNorm.bias"))?),
                },
                lin(
                    get(&format!("{p}.intermediate.dense.weight"))?,
                    get(&format!("{p}.intermediate.dense.bias"))?,
                    im,
                    h,
                ),
                lin(
                    get(&format!("{p}.output.dense.weight"))?,
                    get(&format!("{p}.output.dense.bias"))?,
                    h,
                    im,
                ),
            )
        };
        layers.push(CpuLayer {
            rel: None,
            attn_norm,
            op,
            mlp_norm,
            up,
            down,
        });
    }
    // CrossEncoder (`*ForSequenceClassification`): load the scoring head. Two layouts, same
    // linear→tanh→linear math (GPU twin in encoder.rs): BERT — pooler under the body prefix
    // (`bert.pooler.dense.*`) + a top-level `classifier.*`; RoBERTa/XLM-R — `classifier.dense`
    // → tanh → `classifier.out_proj`, no pooler (bge-reranker-v2-m3, jina-reranker-v2).
    let cross_encoder = if let Pooling::CrossEncoder { n_labels } = cfg.pooling {
        let head = if st.has(&format!("{prefix}pooler.dense.weight")) {
            CpuCrossHead {
                pooler: lin(get("pooler.dense.weight")?, get("pooler.dense.bias")?, h, h),
                classifier: lin(
                    st.tensor_f32("classifier.weight")?,
                    st.tensor_f32("classifier.bias")?,
                    n_labels,
                    h,
                ),
            }
        } else {
            CpuCrossHead {
                pooler: lin(
                    st.tensor_f32("classifier.dense.weight")?,
                    st.tensor_f32("classifier.dense.bias")?,
                    h,
                    h,
                ),
                classifier: lin(
                    st.tensor_f32("classifier.out_proj.weight")?,
                    st.tensor_f32("classifier.out_proj.bias")?,
                    n_labels,
                    h,
                ),
            }
        };
        Some(head)
    } else {
        None
    };
    Ok(CpuWeights {
        word,
        pos: Some(get("embeddings.position_embeddings.weight")?),
        ttype,
        // jina-flash renames the embeddings LayerNorm to `emb_ln`.
        emb_norm: Some(if jina {
            CpuNorm {
                w: get("emb_ln.weight")?,
                b: Some(get("emb_ln.bias")?),
            }
        } else {
            CpuNorm {
                w: get("embeddings.LayerNorm.weight")?,
                b: Some(get("embeddings.LayerNorm.bias")?),
            }
        }),
        layers,
        final_norm: None,
        proj: None,
        pooled_head: None,
        vision: None,
        cross_encoder,
    })
}

/// DeBERTa-v2/v3 tensor table. A BERT body with three differences that all show up here:
///
/// 1. **No `embeddings.position_embeddings`** (`position_biased_input: false`) — `pos` is `None`
///    and the embedding stage adds nothing. Position enters only through attention.
/// 2. **`encoder.rel_embeddings`** `[2·span, hidden]`, LayerNorm'd by **`encoder.LayerNorm`**
///    (`norm_rel_ebd: layer_norm`) and then projected through each layer's OWN `key_proj` /
///    `query_proj` (`share_att_key: true` — there are no `pos_*_proj` tensors). Both projections
///    are precomputed per layer into [`CpuRelPos`]: they depend only on weights.
/// 3. Projections are named `query_proj`/`key_proj`/`value_proj`, not `query`/`key`/`value`.
fn load_deberta_v2(st: &LazySt, cfg: &EncoderConfig) -> Result<CpuWeights> {
    // The last prefix is GLiNER's: its checkpoint nests the whole backbone under the token
    // representation layer, so one `model.safetensors` carries backbone AND head.
    let prefix = ["", "deberta.", "token_rep_layer.bert_layer.model."]
        .into_iter()
        .find(|p| st.has(&format!("{p}embeddings.word_embeddings.weight")))
        .context(
            "word embeddings not found under known prefixes ('', 'deberta.', \
             'token_rep_layer.bert_layer.model.')",
        )?;
    let get = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(&format!("{prefix}{name}")) };
    let spec = cfg.rel_attn.context("DeBERTa config without rel_attn")?;
    let (h, im) = (cfg.hidden, cfg.intermediate);

    let word = get("embeddings.word_embeddings.weight")?;
    anyhow::ensure!(
        word.len() == cfg.vocab * cfg.hidden,
        "word embedding shape {} != vocab {} × hidden {}",
        word.len(),
        cfg.vocab,
        cfg.hidden
    );
    let lin = |w: Vec<f32>, b: Vec<f32>, n: usize, k: usize| CpuLinear {
        w,
        b: Some(b),
        n,
        k,
        ..Default::default()
    };

    // The relative table, normalized ONCE (HF `get_rel_embedding`), then shared by every layer.
    // HF slices `[0 : 2·span]`; a checkpoint whose table is exactly that size (the usual case)
    // is unaffected, but slice anyway so an over-long table cannot silently shift the indices.
    let rows = 2 * spec.span;
    let mut rel_emb = get("encoder.rel_embeddings.weight")?;
    anyhow::ensure!(
        rel_emb.len() >= rows * h,
        "rel_embeddings has {} rows, need 2·span = {rows}",
        rel_emb.len() / h.max(1)
    );
    rel_emb.truncate(rows * h);
    let rel_norm = CpuNorm {
        w: get("encoder.LayerNorm.weight")?,
        b: Some(get("encoder.LayerNorm.bias")?),
    };
    layer_norm_rows(&mut rel_emb, h, &rel_norm, cfg.eps, false);

    let mut layers = Vec::with_capacity(cfg.n_layers);
    for i in 0..cfg.n_layers {
        let p = format!("encoder.layer.{i}");
        let qw = lin(
            get(&format!("{p}.attention.self.query_proj.weight"))?,
            get(&format!("{p}.attention.self.query_proj.bias"))?,
            h,
            h,
        );
        let kw = lin(
            get(&format!("{p}.attention.self.key_proj.weight"))?,
            get(&format!("{p}.attention.self.key_proj.bias"))?,
            h,
            h,
        );
        // share_att_key: the position projections ARE the content projections (bias included, as
        // in HF — `self.key_proj(rel_embeddings)` goes through the full nn.Linear).
        let mut pos_k = vec![0f32; rows * h];
        let mut pos_q = vec![0f32; rows * h];
        linear(&mut pos_k, &rel_emb, &kw);
        linear(&mut pos_q, &rel_emb, &qw);

        layers.push(CpuLayer {
            rel: Some(CpuRelPos { pos_k, pos_q }),
            attn_norm: Some(CpuNorm {
                w: get(&format!("{p}.attention.output.LayerNorm.weight"))?,
                b: Some(get(&format!("{p}.attention.output.LayerNorm.bias"))?),
            }),
            op: CpuOp::Attn {
                q: qw,
                k: kw,
                v: lin(
                    get(&format!("{p}.attention.self.value_proj.weight"))?,
                    get(&format!("{p}.attention.self.value_proj.bias"))?,
                    h,
                    h,
                ),
                o: lin(
                    get(&format!("{p}.attention.output.dense.weight"))?,
                    get(&format!("{p}.attention.output.dense.bias"))?,
                    h,
                    h,
                ),
                q_norm: None,
                k_norm: None,
            },
            mlp_norm: CpuNorm {
                w: get(&format!("{p}.output.LayerNorm.weight"))?,
                b: Some(get(&format!("{p}.output.LayerNorm.bias"))?),
            },
            up: lin(
                get(&format!("{p}.intermediate.dense.weight"))?,
                get(&format!("{p}.intermediate.dense.bias"))?,
                im,
                h,
            ),
            down: lin(
                get(&format!("{p}.output.dense.weight"))?,
                get(&format!("{p}.output.dense.bias"))?,
                h,
                im,
            ),
        });
    }

    Ok(CpuWeights {
        word,
        pos: None, // position_biased_input: false — there is no absolute table
        ttype: if cfg.type_vocab > 0 {
            Some(get("embeddings.token_type_embeddings.weight")?)
        } else {
            None
        },
        emb_norm: Some(CpuNorm {
            w: get("embeddings.LayerNorm.weight")?,
            b: Some(get("embeddings.LayerNorm.bias")?),
        }),
        layers,
        final_norm: None,
        proj: None,
        pooled_head: None,
        vision: None,
        cross_encoder: None,
    })
}

/// ModernBERT tensor table (HF `modeling_modernbert.py` names, `model.` prefix optional):
/// `embeddings.tok_embeddings` + bias-free `embeddings.norm`, per layer fused `attn.Wqkv`
/// (sliced to q|k|v at load so the executor's separate projections apply unchanged), `attn.Wo`,
/// `attn_norm` (ABSENT on layer 0), `mlp_norm`, GeGLU `mlp.Wi`/`mlp.Wo`, trailing `final_norm`.
fn load_modernbert(st: &LazySt, cfg: &EncoderConfig) -> Result<CpuWeights> {
    let prefix = ["", "model."]
        .into_iter()
        .find(|p| st.has(&format!("{p}embeddings.tok_embeddings.weight")))
        .context("ModernBERT tok_embeddings not found under known prefixes ('', 'model.')")?;
    let get = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(&format!("{prefix}{name}")) };
    let has = |name: &str| st.has(&format!("{prefix}{name}"));
    let word = get("embeddings.tok_embeddings.weight")?;
    anyhow::ensure!(
        word.len() == cfg.vocab * cfg.hidden,
        "tok embedding shape {} != vocab {} × hidden {}",
        word.len(),
        cfg.vocab,
        cfg.hidden
    );
    anyhow::ensure!(!cfg.qkv_bias, "ModernBERT is bias-free");
    let norm_of = |name: &str| -> Result<CpuNorm> {
        Ok(CpuNorm {
            w: get(name)?,
            b: if has(&format!("{}.bias", name.trim_end_matches(".weight"))) {
                Some(get(&format!("{}.bias", name.trim_end_matches(".weight")))?)
            } else {
                None
            },
        })
    };
    let (h, im) = (cfg.hidden, cfg.intermediate);
    let mut layers = Vec::with_capacity(cfg.n_layers);
    for i in 0..cfg.n_layers {
        let p = format!("layers.{i}");
        // Fused Wqkv [3H, H] → three [H, H] row slices (per-token layout q|k|v).
        let wqkv = get(&format!("{p}.attn.Wqkv.weight"))?;
        anyhow::ensure!(
            wqkv.len() == 3 * h * h,
            "{p}.attn.Wqkv shape {} != 3·{h}·{h}",
            wqkv.len()
        );
        let slice = |r: usize| CpuLinear {
            w: wqkv[r * h * h..(r + 1) * h * h].to_vec(),
            b: None,
            n: h,
            k: h,
            ..Default::default()
        };
        layers.push(CpuLayer {
            rel: None,
            attn_norm: if i == 0 && cfg.skip_first_attn_norm {
                anyhow::ensure!(
                    !has(&format!("{p}.attn_norm.weight")),
                    "layer 0 attn_norm present but config says skip — checkpoint mismatch"
                );
                None
            } else {
                Some(norm_of(&format!("{p}.attn_norm.weight"))?)
            },
            op: CpuOp::Attn {
                q: slice(0),
                k: slice(1),
                v: slice(2),
                o: CpuLinear {
                    w: get(&format!("{p}.attn.Wo.weight"))?,
                    b: None,
                    n: h,
                    k: h,
                    ..Default::default()
                },
                q_norm: None,
                k_norm: None,
            },
            mlp_norm: norm_of(&format!("{p}.mlp_norm.weight"))?,
            up: CpuLinear {
                w: get(&format!("{p}.mlp.Wi.weight"))?,
                b: None,
                n: 2 * im,
                k: h,
                ..Default::default()
            },
            down: CpuLinear {
                w: get(&format!("{p}.mlp.Wo.weight"))?,
                b: None,
                n: h,
                k: im,
                ..Default::default()
            },
        });
    }
    Ok(CpuWeights {
        word,
        pos: None,
        ttype: None,
        emb_norm: Some(norm_of("embeddings.norm.weight")?),
        layers,
        final_norm: Some(norm_of("final_norm.weight")?),
        proj: None,
        pooled_head: None,
        vision: None,
        cross_encoder: None,
    })
}

/// Qwen3-Embedding tensor table (HF Qwen3ForCausalLM names, `model.` prefix): separate
/// q/k/v/o projections (GQA shapes), per-head `q_norm`/`k_norm`, RMS `input_layernorm` /
/// `post_attention_layernorm` / trailing `model.norm`, and the SiLU-GLU MLP whose
/// `gate_proj`/`up_proj` concatenate into the executor's fused `[2I, H]` layout
/// (`silu(gate) ⊙ up` = act(first half) ⊙ second half). The `lm_head` is never read.
fn load_qwen3_embed(st: &LazySt, cfg: &EncoderConfig) -> Result<CpuWeights> {
    let get = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(&format!("model.{name}")) };
    let word = get("embed_tokens.weight")?;
    anyhow::ensure!(
        word.len() == cfg.vocab * cfg.hidden,
        "embed_tokens shape {} != vocab {} × hidden {}",
        word.len(),
        cfg.vocab,
        cfg.hidden
    );
    let (h, hd, im) = (cfg.hidden, cfg.head_dim, cfg.intermediate);
    let (qh, kvh) = (cfg.n_heads * hd, cfg.n_kv_heads * hd);
    let lin = |w: Vec<f32>, n: usize, k: usize| -> Result<CpuLinear> {
        anyhow::ensure!(w.len() == n * k, "linear shape {} != {n}×{k}", w.len());
        Ok(CpuLinear {
            w,
            b: None,
            n,
            k,
            ..Default::default()
        })
    };
    let mut layers = Vec::with_capacity(cfg.n_layers);
    for i in 0..cfg.n_layers {
        let p = format!("layers.{i}");
        let gate = get(&format!("{p}.mlp.gate_proj.weight"))?;
        let up = get(&format!("{p}.mlp.up_proj.weight"))?;
        anyhow::ensure!(
            gate.len() == im * h && up.len() == im * h,
            "{p} gate/up shapes {}/{} != {im}×{h}",
            gate.len(),
            up.len()
        );
        let mut wi = gate;
        wi.extend_from_slice(&up);
        layers.push(CpuLayer {
            rel: None,
            attn_norm: Some(CpuNorm {
                w: get(&format!("{p}.input_layernorm.weight"))?,
                b: None,
            }),
            op: CpuOp::Attn {
                q: lin(get(&format!("{p}.self_attn.q_proj.weight"))?, qh, h)?,
                k: lin(get(&format!("{p}.self_attn.k_proj.weight"))?, kvh, h)?,
                v: lin(get(&format!("{p}.self_attn.v_proj.weight"))?, kvh, h)?,
                o: lin(get(&format!("{p}.self_attn.o_proj.weight"))?, h, qh)?,
                q_norm: Some(get(&format!("{p}.self_attn.q_norm.weight"))?),
                k_norm: Some(get(&format!("{p}.self_attn.k_norm.weight"))?),
            },
            mlp_norm: CpuNorm {
                w: get(&format!("{p}.post_attention_layernorm.weight"))?,
                b: None,
            },
            up: CpuLinear {
                w: wi,
                b: None,
                n: 2 * im,
                k: h,
                ..Default::default()
            },
            down: lin(get(&format!("{p}.mlp.down_proj.weight"))?, h, im)?,
        });
    }
    Ok(CpuWeights {
        word,
        pos: None,
        ttype: None,
        emb_norm: None,
        layers,
        final_norm: Some(CpuNorm {
            w: get("norm.weight")?,
            b: None,
        }),
        proj: None,
        pooled_head: None,
        vision: None,
        cross_encoder: None,
    })
}

/// LFM2-ColBERT tensor table (verified against LiquidAI/LFM2-ColBERT-350M): bare `Lfm2Model`
/// names — `embed_tokens`, per layer `operator_norm` + (`conv.{in_proj,conv,out_proj}` |
/// `self_attn.{q,k,v,out}_proj` + `q/k_layernorm`) + `ffn_norm` + llama-SwiGLU
/// `feed_forward.{w1,w2,w3}` (w1|w3 fused into the act(first)⊙second layout), trailing
/// `embedding_norm` (LFM2's FINAL norm), plus the pylate `1_Dense/model.safetensors`
/// per-token projection.
fn load_lfm2_colbert(st: &LazySt, cfg: &EncoderConfig, dir: &Path) -> Result<CpuWeights> {
    let get = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(name) };
    let word = get("embed_tokens.weight")?;
    anyhow::ensure!(
        word.len() == cfg.vocab * cfg.hidden,
        "embed_tokens shape {} != vocab {} × hidden {}",
        word.len(),
        cfg.vocab,
        cfg.hidden
    );
    let (h, hd, im, l) = (cfg.hidden, cfg.head_dim, cfg.intermediate, cfg.conv_l);
    let (qh, kvh) = (cfg.n_heads * hd, cfg.n_kv_heads * hd);
    let lin = |w: Vec<f32>, n: usize, k: usize| -> Result<CpuLinear> {
        anyhow::ensure!(w.len() == n * k, "linear shape {} != {n}×{k}", w.len());
        Ok(CpuLinear {
            w,
            b: None,
            n,
            k,
            ..Default::default()
        })
    };
    let mut layers = Vec::with_capacity(cfg.n_layers);
    for i in 0..cfg.n_layers {
        let p = format!("layers.{i}");
        let op = if cfg.layer_is_attn[i] {
            CpuOp::Attn {
                q: lin(get(&format!("{p}.self_attn.q_proj.weight"))?, qh, h)?,
                k: lin(get(&format!("{p}.self_attn.k_proj.weight"))?, kvh, h)?,
                v: lin(get(&format!("{p}.self_attn.v_proj.weight"))?, kvh, h)?,
                o: lin(get(&format!("{p}.self_attn.out_proj.weight"))?, h, qh)?,
                q_norm: Some(get(&format!("{p}.self_attn.q_layernorm.weight"))?),
                k_norm: Some(get(&format!("{p}.self_attn.k_layernorm.weight"))?),
            }
        } else {
            let conv_w = get(&format!("{p}.conv.conv.weight"))?;
            anyhow::ensure!(
                conv_w.len() == h * l,
                "{p} conv taps {} != hidden {h} × conv_l {l}",
                conv_w.len()
            );
            CpuOp::Conv {
                in_proj: lin(get(&format!("{p}.conv.in_proj.weight"))?, 3 * h, h)?,
                conv_w,
                out_proj: lin(get(&format!("{p}.conv.out_proj.weight"))?, h, h)?,
            }
        };
        let gate = get(&format!("{p}.feed_forward.w1.weight"))?;
        let up = get(&format!("{p}.feed_forward.w3.weight"))?;
        anyhow::ensure!(
            gate.len() == im * h && up.len() == im * h,
            "{p} w1/w3 shapes {}/{} != {im}×{h}",
            gate.len(),
            up.len()
        );
        let mut wi = gate;
        wi.extend_from_slice(&up);
        layers.push(CpuLayer {
            rel: None,
            attn_norm: Some(CpuNorm {
                w: get(&format!("{p}.operator_norm.weight"))?,
                b: None,
            }),
            op,
            mlp_norm: CpuNorm {
                w: get(&format!("{p}.ffn_norm.weight"))?,
                b: None,
            },
            up: CpuLinear {
                w: wi,
                b: None,
                n: 2 * im,
                k: h,
                ..Default::default()
            },
            down: lin(get(&format!("{p}.feed_forward.w2.weight"))?, h, im)?,
        });
    }
    // The pylate Dense module lives in its own safetensors under 1_Dense/.
    let dense_st = LazySt::open(&dir.join("1_Dense"))?;
    let dim = match cfg.pooling {
        Pooling::PerToken { dim } => dim,
        other => anyhow::bail!("ColBERT config must use PerToken pooling, got {other:?}"),
    };
    let proj = lin(dense_st.tensor_f32("linear.weight")?, dim, h)?;
    Ok(CpuWeights {
        word,
        pos: None,
        ttype: None,
        emb_norm: None,
        layers,
        // LFM2 quirk: `embedding_norm` is the FINAL norm (matches the decoder engine's use).
        final_norm: Some(CpuNorm {
            w: get("embedding_norm.weight")?,
            b: None,
        }),
        proj: Some(proj),
        pooled_head: None,
        vision: None,
        cross_encoder: None,
    })
}

/// SigLIP TEXT tower tensor table (HF `SiglipTextModel` names under `text_model.`): token +
/// position embeddings (no embedding-stage norm), per layer biased `layer_norm1`/`layer_norm2`
/// (pre-LN), biased q/k/v/out projections, biased `mlp.fc1`/`fc2`, trailing `final_layer_norm`,
/// and the post-pooling `head` projection.
fn load_siglip_text_weights(st: &LazySt, cfg: &EncoderConfig) -> Result<CpuWeights> {
    let get = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(&format!("text_model.{name}")) };
    let word = get("embeddings.token_embedding.weight")?;
    anyhow::ensure!(
        word.len() == cfg.vocab * cfg.hidden,
        "token embedding shape {} != vocab {} × hidden {}",
        word.len(),
        cfg.vocab,
        cfg.hidden
    );
    let (h, im) = (cfg.hidden, cfg.intermediate);
    let lin = |w: Vec<f32>, b: Vec<f32>, n: usize, k: usize| -> Result<CpuLinear> {
        anyhow::ensure!(w.len() == n * k && b.len() == n, "linear shape");
        Ok(CpuLinear {
            w,
            b: Some(b),
            n,
            k,
            ..Default::default()
        })
    };
    let mut layers = Vec::with_capacity(cfg.n_layers);
    for i in 0..cfg.n_layers {
        let p = format!("encoder.layers.{i}");
        layers.push(CpuLayer {
            rel: None,
            attn_norm: Some(CpuNorm {
                w: get(&format!("{p}.layer_norm1.weight"))?,
                b: Some(get(&format!("{p}.layer_norm1.bias"))?),
            }),
            op: CpuOp::Attn {
                q: lin(
                    get(&format!("{p}.self_attn.q_proj.weight"))?,
                    get(&format!("{p}.self_attn.q_proj.bias"))?,
                    h,
                    h,
                )?,
                k: lin(
                    get(&format!("{p}.self_attn.k_proj.weight"))?,
                    get(&format!("{p}.self_attn.k_proj.bias"))?,
                    h,
                    h,
                )?,
                v: lin(
                    get(&format!("{p}.self_attn.v_proj.weight"))?,
                    get(&format!("{p}.self_attn.v_proj.bias"))?,
                    h,
                    h,
                )?,
                o: lin(
                    get(&format!("{p}.self_attn.out_proj.weight"))?,
                    get(&format!("{p}.self_attn.out_proj.bias"))?,
                    h,
                    h,
                )?,
                q_norm: None,
                k_norm: None,
            },
            mlp_norm: CpuNorm {
                w: get(&format!("{p}.layer_norm2.weight"))?,
                b: Some(get(&format!("{p}.layer_norm2.bias"))?),
            },
            up: lin(
                get(&format!("{p}.mlp.fc1.weight"))?,
                get(&format!("{p}.mlp.fc1.bias"))?,
                im,
                h,
            )?,
            down: lin(
                get(&format!("{p}.mlp.fc2.weight"))?,
                get(&format!("{p}.mlp.fc2.bias"))?,
                h,
                im,
            )?,
        });
    }
    Ok(CpuWeights {
        word,
        pos: Some(get("embeddings.position_embedding.weight")?),
        ttype: None,
        emb_norm: None,
        layers,
        final_norm: Some(CpuNorm {
            w: get("final_layer_norm.weight")?,
            b: Some(get("final_layer_norm.bias")?),
        }),
        proj: None,
        pooled_head: Some(lin(get("head.weight")?, get("head.bias")?, h, h)?),
        vision: None,
        cross_encoder: None,
    })
}

/// SigLIP VISION tower tensor table (under `vision_model.`): conv patch embedding (flattened to
/// the im2row GEMM layout), per-patch positions, the same biased pre-LN layer stack, trailing
/// `post_layernorm`, and the MAP attention-pooling head (torch packed `in_proj`).
fn load_siglip_vision_weights(
    st: &LazySt,
    spec: &crate::encoder_weights::SiglipVisionSpec,
) -> Result<CpuWeights> {
    let cfg = &spec.config;
    let get = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(&format!("vision_model.{name}")) };
    let (h, im) = (cfg.hidden, cfg.intermediate);
    let lin = |w: Vec<f32>, b: Vec<f32>, n: usize, k: usize| -> Result<CpuLinear> {
        anyhow::ensure!(w.len() == n * k && b.len() == n, "linear shape");
        Ok(CpuLinear {
            w,
            b: Some(b),
            n,
            k,
            ..Default::default()
        })
    };
    let mut layers = Vec::with_capacity(cfg.n_layers);
    for i in 0..cfg.n_layers {
        let p = format!("encoder.layers.{i}");
        layers.push(CpuLayer {
            rel: None,
            attn_norm: Some(CpuNorm {
                w: get(&format!("{p}.layer_norm1.weight"))?,
                b: Some(get(&format!("{p}.layer_norm1.bias"))?),
            }),
            op: CpuOp::Attn {
                q: lin(
                    get(&format!("{p}.self_attn.q_proj.weight"))?,
                    get(&format!("{p}.self_attn.q_proj.bias"))?,
                    h,
                    h,
                )?,
                k: lin(
                    get(&format!("{p}.self_attn.k_proj.weight"))?,
                    get(&format!("{p}.self_attn.k_proj.bias"))?,
                    h,
                    h,
                )?,
                v: lin(
                    get(&format!("{p}.self_attn.v_proj.weight"))?,
                    get(&format!("{p}.self_attn.v_proj.bias"))?,
                    h,
                    h,
                )?,
                o: lin(
                    get(&format!("{p}.self_attn.out_proj.weight"))?,
                    get(&format!("{p}.self_attn.out_proj.bias"))?,
                    h,
                    h,
                )?,
                q_norm: None,
                k_norm: None,
            },
            mlp_norm: CpuNorm {
                w: get(&format!("{p}.layer_norm2.weight"))?,
                b: Some(get(&format!("{p}.layer_norm2.bias"))?),
            },
            up: lin(
                get(&format!("{p}.mlp.fc1.weight"))?,
                get(&format!("{p}.mlp.fc1.bias"))?,
                im,
                h,
            )?,
            down: lin(
                get(&format!("{p}.mlp.fc2.weight"))?,
                get(&format!("{p}.mlp.fc2.bias"))?,
                h,
                im,
            )?,
        });
    }
    let ps2 = spec.patch_size * spec.patch_size;
    let patch_w = get("embeddings.patch_embedding.weight")?; // [H, 3, p, p] row-major
    anyhow::ensure!(
        patch_w.len() == h * 3 * ps2,
        "patch embedding shape {} != {h}×3×{ps2}",
        patch_w.len()
    );
    let in_proj_w = get("head.attention.in_proj_weight")?;
    anyhow::ensure!(in_proj_w.len() == 3 * h * h, "in_proj shape");
    Ok(CpuWeights {
        word: Vec::new(),
        pos: None,
        ttype: None,
        emb_norm: None,
        layers,
        final_norm: Some(CpuNorm {
            w: get("post_layernorm.weight")?,
            b: Some(get("post_layernorm.bias")?),
        }),
        proj: None,
        pooled_head: None,
        vision: Some(CpuVision {
            patch: lin(patch_w, get("embeddings.patch_embedding.bias")?, h, 3 * ps2)?,
            pos: get("embeddings.position_embedding.weight")?,
            probe: get("head.probe")?,
            in_proj_w,
            in_proj_b: get("head.attention.in_proj_bias")?,
            out_proj: lin(
                get("head.attention.out_proj.weight")?,
                get("head.attention.out_proj.bias")?,
                h,
                h,
            )?,
            head_ln: CpuNorm {
                w: get("head.layernorm.weight")?,
                b: Some(get("head.layernorm.bias")?),
            },
            head_fc1: lin(
                get("head.mlp.fc1.weight")?,
                get("head.mlp.fc1.bias")?,
                im,
                h,
            )?,
            head_fc2: lin(
                get("head.mlp.fc2.weight")?,
                get("head.mlp.fc2.bias")?,
                h,
                im,
            )?,
        }),
        cross_encoder: None,
    })
}

/// nomic-bert tensor table (verified against nomic-ai/nomic-embed-text-v1.5): word + token-type
/// embeddings with biased `emb_ln` (rotary — no position table), POST-norm layers (`norm1` after
/// the attention residual, `norm2` after the MLP residual, both biased), fused bias-free
/// `attn.Wqkv` (sliced q|k|v), `attn.out_proj`, SwiGLU `mlp.fc11` (gate) | `mlp.fc12` fused into
/// the executor's act(first)⊙second layout, `mlp.fc2` down. No trailing norm (post-norm).
fn load_nomic_bert(st: &LazySt, cfg: &EncoderConfig) -> Result<CpuWeights> {
    let get = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(name) };
    let word = get("embeddings.word_embeddings.weight")?;
    anyhow::ensure!(
        word.len() == cfg.vocab * cfg.hidden,
        "word embedding shape {} != vocab {} × hidden {}",
        word.len(),
        cfg.vocab,
        cfg.hidden
    );
    let (h, im) = (cfg.hidden, cfg.intermediate);
    let mut layers = Vec::with_capacity(cfg.n_layers);
    for i in 0..cfg.n_layers {
        let p = format!("encoder.layers.{i}");
        let wqkv = get(&format!("{p}.attn.Wqkv.weight"))?;
        anyhow::ensure!(
            wqkv.len() == 3 * h * h,
            "{p}.attn.Wqkv shape {} != 3·{h}·{h}",
            wqkv.len()
        );
        let slice = |r: usize| CpuLinear {
            w: wqkv[r * h * h..(r + 1) * h * h].to_vec(),
            b: None,
            n: h,
            k: h,
            ..Default::default()
        };
        // Verified nomic GLU order: `y = fc11(x) · act(fc12(x))` — fc12 is the GATE. The
        // executor's fused layout is act(first) ⊙ second, so fc12 concatenates FIRST.
        let gate = get(&format!("{p}.mlp.fc12.weight"))?;
        let lin_half = get(&format!("{p}.mlp.fc11.weight"))?;
        anyhow::ensure!(
            gate.len() == im * h && lin_half.len() == im * h,
            "{p} fc11/fc12 shapes {}/{} != {im}×{h}",
            lin_half.len(),
            gate.len()
        );
        let mut wi = gate;
        wi.extend_from_slice(&lin_half);
        layers.push(CpuLayer {
            rel: None,
            // Post-norm: `norm1` is the POST-attention norm (executor post-LN semantics).
            attn_norm: Some(CpuNorm {
                w: get(&format!("{p}.norm1.weight"))?,
                b: Some(get(&format!("{p}.norm1.bias"))?),
            }),
            op: CpuOp::Attn {
                q: slice(0),
                k: slice(1),
                v: slice(2),
                o: CpuLinear {
                    w: get(&format!("{p}.attn.out_proj.weight"))?,
                    b: None,
                    n: h,
                    k: h,
                    ..Default::default()
                },
                q_norm: None,
                k_norm: None,
            },
            mlp_norm: CpuNorm {
                w: get(&format!("{p}.norm2.weight"))?,
                b: Some(get(&format!("{p}.norm2.bias"))?),
            },
            up: CpuLinear {
                w: wi,
                b: None,
                n: 2 * im,
                k: h,
                ..Default::default()
            },
            down: CpuLinear {
                w: get(&format!("{p}.mlp.fc2.weight"))?,
                b: None,
                n: h,
                k: im,
                ..Default::default()
            },
        });
    }
    Ok(CpuWeights {
        word,
        pos: None,
        ttype: Some(get("embeddings.token_type_embeddings.weight")?),
        emb_norm: Some(CpuNorm {
            w: get("emb_ln.weight")?,
            b: Some(get("emb_ln.bias")?),
        }),
        layers,
        final_norm: None,
        proj: None,
        pooled_head: None,
        vision: None,
        cross_encoder: None,
    })
}

#[cfg(test)]
mod dot_bench {
    use super::dot;

    /// Is the four-accumulator [`dot`] actually faster than the idiomatic `.sum()`, at the shape
    /// the encoder attention runs (hd=64)?
    ///
    /// Both arms run INTERLEAVED in one process, min-of-N. Comparing two separate binaries across
    /// two runs on this box is not a measurement — the load average swings 2.8× within minutes and
    /// has already sent this campaign down a blind alley once. Alternating the arms is what makes
    /// the delta attributable to the code rather than to the neighbours.
    #[test]
    #[ignore = "perf probe: cargo test --release -p inferencelayer dot_bench -- --ignored --nocapture"]
    fn four_accumulators_beat_dot_sum() {
        const HD: usize = 64;
        const N: usize = 4096; // one query row's worth of keys, ×heads
        let a: Vec<f32> = (0..HD * N).map(|i| (i % 17) as f32 * 0.01 - 0.08).collect();
        let b: Vec<f32> = (0..HD * N).map(|i| (i % 23) as f32 * 0.01 - 0.11).collect();

        let scalar = |x: &[f32], y: &[f32]| -> f32 { x.iter().zip(y).map(|(p, q)| p * q).sum() };
        let (mut fast, mut slow) = (f64::MAX, f64::MAX);
        let (mut fsum, mut ssum) = (0f32, 0f32);
        for _ in 0..200 {
            let t = std::time::Instant::now();
            let mut s = 0f32;
            for i in 0..N {
                s += dot(&a[i * HD..(i + 1) * HD], &b[i * HD..(i + 1) * HD]);
            }
            fast = fast.min(t.elapsed().as_secs_f64());
            fsum = s;

            let t = std::time::Instant::now();
            let mut s = 0f32;
            for i in 0..N {
                s += scalar(&a[i * HD..(i + 1) * HD], &b[i * HD..(i + 1) * HD]);
            }
            slow = slow.min(t.elapsed().as_secs_f64());
            ssum = s;
        }
        let gf = |t: f64| (2.0 * (HD * N) as f64) / t / 1e9;
        eprintln!(
            "\n  dot(hd=64) x{N}:  4-acc {:.1} GF/s | .sum() {:.1} GF/s  ->  {:.2}x\n",
            gf(fast),
            gf(slow),
            slow / fast
        );
        assert!(
            (fsum - ssum).abs() < 1e-2,
            "arms must agree: {fsum} vs {ssum}"
        );
        assert!(
            fast < slow,
            "4-acc dot ({fast:.6}s) must beat .sum() ({slow:.6}s)"
        );
    }
}

#[cfg(test)]
mod attn_equiv {
    use super::*;

    fn synth(n: usize, f: impl Fn(usize) -> f32) -> Vec<f32> {
        (0..n).map(f).collect()
    }

    /// The GEMM form must reproduce the scalar kernel it replaces — bit-for-bit is not required
    /// (the summation order legitimately differs), but far below the 1e-4 DeBERTa-vs-HF gate is.
    ///
    /// Called directly rather than through `OSFKB_ENC_ATTN_SCALAR`, because that switch is a process
    /// -global env read and the test binary runs its tests on parallel threads.
    #[test]
    fn gemm_form_matches_the_scalar_reference() {
        for (span, max_rel, c2p, p2c, seqs, n_heads, hd) in [
            // Production DeBERTa-v3 geometry (GLiNER's backbone), single sequence.
            (
                256usize,
                512usize,
                true,
                true,
                vec![112usize],
                12usize,
                64usize,
            ),
            // Ragged batch + a block-ragged tail (77 = 2·32 + 13) + log-bucket region well exercised.
            (16, 64, true, true, vec![77, 40], 4, 16),
            // Ablations: each positional term alone.
            (16, 64, true, false, vec![33], 4, 16),
            (16, 64, false, true, vec![33], 4, 16),
        ] {
            let spec = RelAttn {
                span,
                max_rel,
                c2p,
                p2c,
            };
            let qh = n_heads * hd;
            let t: usize = seqs.iter().sum();
            let mut starts = vec![0u32];
            for s in &seqs {
                starts.push(starts.last().unwrap() + *s as u32);
            }
            let seq_of = row_to_seq(&starts, t);

            let q = synth(t * qh, |i| ((i % 37) as f32 - 18.0) * 0.031);
            let k = synth(t * qh, |i| ((i % 41) as f32 - 20.0) * 0.027);
            let v = synth(t * qh, |i| ((i % 29) as f32 - 14.0) * 0.043);
            let rel = CpuRelPos {
                pos_k: synth(2 * span * qh, |i| ((i % 23) as f32 - 11.0) * 0.019),
                pos_q: synth(2 * span * qh, |i| ((i % 31) as f32 - 15.0) * 0.017),
            };
            // Padding in the middle of a sequence: the mask must survive the GEMM restructure.
            let mut valid = vec![1u32; t];
            valid[5] = 0;
            valid[t - 2] = 0;

            let mut fast = vec![0f32; t * qh];
            let mut refr = vec![0f32; t * qh];
            attention_disentangled(
                &mut fast,
                &q,
                &k,
                &v,
                &starts,
                &seq_of,
                n_heads,
                hd,
                &rel,
                spec,
                Some(&valid),
            );
            attention_disentangled_scalar(
                &mut refr,
                &q,
                &k,
                &v,
                &starts,
                &seq_of,
                n_heads,
                hd,
                &rel,
                spec,
                Some(&valid),
            );
            let dmax = fast
                .iter()
                .zip(&refr)
                .map(|(a, b)| (a - b).abs())
                .fold(0f32, f32::max);
            assert!(
                dmax < 1e-5,
                "span={span} c2p={c2p} p2c={p2c} seqs={seqs:?}: max |Δ| {dmax:.3e} vs scalar"
            );
        }
    }

    /// The fast path is the only path in production, so it must also be *deterministic*: it is the
    /// oracle the GPU kernels are gated against, and a gate that moves is not a gate.
    #[test]
    fn gemm_form_is_deterministic() {
        let spec = RelAttn {
            span: 256,
            max_rel: 512,
            c2p: true,
            p2c: true,
        };
        let (n_heads, hd, t) = (12usize, 64usize, 96usize);
        let qh = n_heads * hd;
        let starts = vec![0u32, t as u32];
        let seq_of = row_to_seq(&starts, t);
        let q = synth(t * qh, |i| ((i % 37) as f32 - 18.0) * 0.031);
        let k = synth(t * qh, |i| ((i % 41) as f32 - 20.0) * 0.027);
        let v = synth(t * qh, |i| ((i % 29) as f32 - 14.0) * 0.043);
        let rel = CpuRelPos {
            pos_k: synth(512 * qh, |i| ((i % 23) as f32 - 11.0) * 0.019),
            pos_q: synth(512 * qh, |i| ((i % 31) as f32 - 15.0) * 0.017),
        };
        let run = || {
            let mut o = vec![0f32; t * qh];
            attention_disentangled(
                &mut o, &q, &k, &v, &starts, &seq_of, n_heads, hd, &rel, spec, None,
            );
            o
        };
        assert_eq!(
            run(),
            run(),
            "two runs of the fast path must be bit-identical"
        );
    }
}

#[cfg(test)]
mod attn_bench {
    use super::*;

    /// GEMM form vs the scalar twin at production DeBERTa geometry, INTERLEAVED min-of-N.
    ///
    /// Run through the full model the two arms disagreed (backbone said 1.8×, end-to-end said 1.05×)
    /// because they were measured minutes apart on a box whose load average moves by 2× in that time.
    /// Alternating the arms inside one process is the only comparison this machine can actually
    /// support.
    #[test]
    #[ignore = "perf probe: cargo test --release attn_bench -- --ignored --nocapture"]
    fn gemm_form_beats_the_scalar_kernel() {
        let (n_heads, hd, span) = (12usize, 64usize, 256usize);
        let spec = RelAttn {
            span,
            max_rel: 512,
            c2p: true,
            p2c: true,
        };
        let qh = n_heads * hd;
        for t in [32usize, 64, 128, 192, 256, 384, 512] {
            let starts = vec![0u32, t as u32];
            let seq_of = row_to_seq(&starts, t);
            let g = |n, m, s: f32| -> Vec<f32> {
                (0..n)
                    .map(|i| ((i % m) as f32 - (m / 2) as f32) * s)
                    .collect()
            };
            let (q, k, v) = (
                g(t * qh, 37, 0.031),
                g(t * qh, 41, 0.027),
                g(t * qh, 29, 0.043),
            );
            let rel = CpuRelPos {
                pos_k: g(2 * span * qh, 23, 0.019),
                pos_q: g(2 * span * qh, 31, 0.017),
            };
            let mut out = vec![0f32; t * qh];
            let (mut fast, mut slow) = (f64::MAX, f64::MAX);
            for _ in 0..12 {
                let t0 = std::time::Instant::now();
                attention_disentangled_gemm(
                    &mut out, &q, &k, &v, &starts, n_heads, hd, &rel, spec, None,
                );
                fast = fast.min(t0.elapsed().as_secs_f64());

                let t0 = std::time::Instant::now();
                attention_disentangled_scalar(
                    &mut out, &q, &k, &v, &starts, &seq_of, n_heads, hd, &rel, spec, None,
                );
                slow = slow.min(t0.elapsed().as_secs_f64());
            }
            eprintln!(
                "  T={t:4}  gemm {:6.2} ms | scalar {:6.2} ms  ->  {:.2}x   (×6 layers: {:5.1} -> {:5.1} ms)",
                fast * 1e3,
                slow * 1e3,
                slow / fast,
                slow * 6e3,
                fast * 6e3,
            );
            if t >= GEMM_ATTN_MIN_T {
                assert!(
                    fast < slow,
                    "T={t}: gemm form must beat the scalar kernel above the cutover"
                );
            }
        }
    }
}

#[cfg(test)]
mod gemm_pack_probe {
    /// How much of a `linear()` call is the `gemm` crate REPACKING the weight matrix?
    ///
    /// The pack cost depends on the weight, not on the batch — so it is invisible at large M and
    /// ruinous at small M. This probe reports achieved GF/s across the encoder's real M values with
    /// the SAME weight reused every call (the production pattern: weights are immutable, and a
    /// prepack could be hoisted to load time). A flat GF/s curve would mean packing is free and P5's
    /// premise is wrong; a curve that collapses at small M sizes the prize directly.
    #[test]
    #[ignore = "perf probe: cargo test --release gemm_pack_probe -- --ignored --nocapture"]
    fn pack_cost_dominates_at_small_m() {
        // DeBERTa-v3-small: hidden 768, intermediate 3072 — the qkv/o and up/down shapes.
        for (k, n, tag) in [(768usize, 768usize, "qkv/o "), (768, 3072, "mlp-up")] {
            let w: Vec<f32> = (0..k * n).map(|i| ((i % 19) as f32 - 9.0) * 0.01).collect();
            eprintln!("\n  [{tag}] K={k} N={n}");
            for m in [16usize, 32, 64, 128, 256, 512] {
                let x: Vec<f32> = (0..m * k)
                    .map(|i| ((i % 23) as f32 - 11.0) * 0.01)
                    .collect();
                let mut y = vec![0f32; m * n];
                let mut best = f64::MAX;
                for _ in 0..20 {
                    let t0 = std::time::Instant::now();
                    // SAFETY: row-major [m,k] × [k,n] → [m,n]; all three buffers are ours.
                    unsafe {
                        gemm::gemm(
                            m,
                            n,
                            k,
                            y.as_mut_ptr(),
                            1,
                            n as isize,
                            false,
                            x.as_ptr(),
                            1,
                            k as isize,
                            w.as_ptr(),
                            1,
                            n as isize,
                            0.0,
                            1.0,
                            false,
                            false,
                            false,
                            gemm::Parallelism::Rayon(0),
                        );
                    }
                    best = best.min(t0.elapsed().as_secs_f64());
                }
                let gf = 2.0 * (m * n * k) as f64 / best / 1e9;
                eprintln!("    M={m:4}  {:7.3} ms  {gf:7.1} GF/s", best * 1e3);
            }
        }
    }
}