franken_ocr 0.8.0

//! The 12-layer DeepseekV2/Llama-MHA decoder driver ([SPEC-070..081],
//! PROPOSED_ARCHITECTURE.md §6.7).
//!
//! `embed_tokens` -> per layer `h = x + self_attn(input_layernorm(x))`;
//! `out = h + mlp(post_attention_layernorm(h))` -> final RMSNorm -> `lm_head`
//! GEMV 1280 -> 129280. RoPE = Llama variant (theta 10000, head_dim 128,
//! NEOX-style `rotate_half`, NOT the DeepseekV2 interleave — [SPEC-078]). Layer
//! 0 is dense MLP; layers 1..11 are MoE; all 12 layers use R-SWA attention.
//!
//! ## What lives here vs. the siblings
//!
//! This file is the *driver*: it owns the per-layer pre-norm residual loop
//! ([SPEC-072]), the token-embedding lookup ([SPEC-070]), the final norm, the
//! `lm_head` GEMV ([SPEC-081]), and the **RoPE** math ([SPEC-078]) applied to Q/K
//! before they are handed to attention. The attention kernel itself is
//! [`crate::native_engine::rswa::attention`]; the per-layer MLP/MoE block is
//! [`crate::native_engine::moe::forward`] (layers 1..11) or the dense SwiGLU MLP
//! ([SPEC-075], layer 0). Norms funnel through the [`crate::native_engine::nn`]
//! facade ([`nn::rms_norm`]); GEMMs through [`nn::matmul`].
//!
//! ## Weights plumbing (wired)
//!
//! [`super::weights::Weights`] exposes named-tensor accessors (`mat`/`vec`,
//! widening BF16→f32 at the boundary), so the top-level [`forward`]/[`lm_head`]
//! entrypoints pull the per-layer norm/projection slices, the `model.norm`
//! weight, and the top-level `lm_head.weight` straight out of it. [`forward`] is
//! the **stateless, full-sequence prefill** path (full causal MHA via
//! [`prefill_attention`]); incremental decode threads a separate
//! `&mut [RingCache; 12]` (bd-1gv.17) and is out of scope here. ALL the
//! load-bearing math — token embed, RMSNorm, RoPE, the dense/MoE MLPs, the
//! lm_head GEMV, and the per-layer driver — stays unit-tested as pure functions
//! over explicit `&[f32]` / [`Mat`] slices; the entrypoints are the thin wiring
//! over them.

// Numerical kernel file: many hot loops index parallel stride-arrays by range
// (`a[r*stride..]`, `caches[layer]`/`weights[layer]`), where `clippy::needless_range_loop`
// is a false positive — the index is genuinely needed across several arrays.
#![allow(clippy::needless_range_loop)]

use super::calib;
use super::moe;
use super::nn;
use super::rswa::{self, BatchedRingCache, RingCache};
use super::tensor::{Mat, QInt4, QInt8, WeightLayout};
use super::weights::{DType, Weights};
use crate::error::{FocrError, FocrResult};
use crate::quant::calib as quant_calib;
use crate::simd;
use rayon::prelude::*;

// Clock seam: `std::time::Instant` traps on wasm32-unknown-unknown; `web-time`
// re-exports std's types on native targets, so native behavior is unchanged.
#[cfg(not(target_arch = "wasm32"))]
use std::time::Instant;
#[cfg(target_arch = "wasm32")]
use web_time::Instant;

fn checked_shape_mul(context: &str, lhs: usize, rhs: usize, expression: &str) -> FocrResult<usize> {
    lhs.checked_mul(rhs).ok_or_else(|| {
        FocrError::Other(anyhow::anyhow!(
            "{context}: usize overflow computing {expression} ({lhs} * {rhs})"
        ))
    })
}

fn checked_mat_len(context: &str, x: &Mat) -> FocrResult<usize> {
    let expected = checked_shape_mul(context, x.rows, x.cols, "rows*cols")?;
    if x.data.len() != expected {
        return Err(FocrError::Other(anyhow::anyhow!(
            "{context}: data len {} != rows*cols {} for shape [{}, {}]",
            x.data.len(),
            expected,
            x.rows,
            x.cols
        )));
    }
    Ok(expected)
}

/// Decoder hyperparameters (compile-time shapes — plan §1.1 P2). Mirrors the
/// pinned `config.json` so the forward never reads runtime params.
pub mod config {
    /// Hidden size ([SPEC-010]).
    pub const HIDDEN_SIZE: usize = 1280;
    /// Dense MLP intermediate size ([SPEC-010]).
    pub const INTERMEDIATE_SIZE: usize = 6848;
    /// Decoder layer count ([SPEC-010]).
    pub const NUM_HIDDEN_LAYERS: usize = 12;
    /// Attention heads ([SPEC-010]); MHA (kv heads = heads, no GQA).
    pub const NUM_ATTENTION_HEADS: usize = 10;
    /// KV heads ([SPEC-010]).
    pub const NUM_KEY_VALUE_HEADS: usize = 10;
    /// Per-head dim (`hidden / heads`) ([SPEC-011]).
    pub const HEAD_DIM: usize = 128;
    /// Vocabulary size ([SPEC-010]).
    pub const VOCAB_SIZE: usize = 129280;
    /// RoPE theta ([SPEC-013/078]).
    pub const ROPE_THETA: f32 = 10000.0;
    /// RMSNorm epsilon ([SPEC-013/071]).
    pub const RMS_NORM_EPS: f32 = 1e-6;
    /// First dense layer count (layer 0 dense) ([SPEC-012]).
    pub const FIRST_K_DENSE_REPLACE: usize = 1;
    /// R-SWA window ([SPEC-015/090]).
    pub const RING_WINDOW: usize = 128;
    /// BOS token id ([SPEC-014]).
    pub const BOS_TOKEN_ID: u32 = 0;
    /// EOS token id ([SPEC-014]).
    pub const EOS_TOKEN_ID: u32 = 1;
}

// ── Token embedding ([SPEC-070]) ────────────────────────────────────────────

/// Gather rows of the `embed_tokens` table for a sequence of token ids
/// ([SPEC-070]).
///
/// `table` is the row-major `[vocab, hidden]` embedding matrix
/// (`model.embed_tokens.weight`); `ids` is the prompt/decode token sequence.
/// Returns the `[seq, hidden]` activation [`Mat`] that begins the residual
/// stream — exactly the `inputs_embeds` the connector then masked-scatters
/// vision features into.
///
/// # Errors
/// [`FocrError::Other`] if `table.len() != vocab * hidden`, or any id is
/// `>= vocab`.
pub fn embed_tokens(table: &[f32], vocab: usize, hidden: usize, ids: &[u32]) -> FocrResult<Mat> {
    let expected_table_len = checked_shape_mul("embed_tokens", vocab, hidden, "vocab*hidden")?;
    if table.len() != expected_table_len {
        return Err(FocrError::Other(anyhow::anyhow!(
            "embed_tokens: table len {} != vocab*hidden {}",
            table.len(),
            expected_table_len
        )));
    }
    let seq = ids.len();
    let out_len = checked_shape_mul("embed_tokens", seq, hidden, "seq*hidden")?;
    let mut out = vec![0.0f32; out_len];
    for (t, &id) in ids.iter().enumerate() {
        let row = id as usize;
        if row >= vocab {
            return Err(FocrError::Other(anyhow::anyhow!(
                "embed_tokens: id {row} out of range (vocab {vocab})"
            )));
        }
        let src = &table[row * hidden..(row + 1) * hidden];
        out[t * hidden..(t + 1) * hidden].copy_from_slice(src);
    }
    Ok(Mat::from_vec(seq, hidden, out))
}

// ── RoPE ([SPEC-078], Llama-style NEOX rotate_half) ─────────────────────────

/// Precomputed RoPE phases for the prompt positions: `cos[p][i]`, `sin[p][i]`
/// laid out `[seq, head_dim]` row-major, with the half-dim freqs duplicated into
/// both halves (Llama `cat(freqs, freqs)`), so they can be applied head-by-head
/// with a single index.
///
/// `inv_freq[i] = theta^(-2i/head_dim)` for `i in 0..head_dim/2`; for position
/// `p`, `angle = p * inv_freq[i]`. `cos`/`sin` each duplicate `i` into columns
/// `i` and `i + head_dim/2`.
#[derive(Debug, Clone)]
pub struct RopeTable {
    /// Per-(position, channel) cosine, row-major `[seq, head_dim]`.
    pub cos: Vec<f32>,
    /// Per-(position, channel) sine, row-major `[seq, head_dim]`.
    pub sin: Vec<f32>,
    /// Per-head dimension (here 128).
    pub head_dim: usize,
}

impl RopeTable {
    /// Build the cos/sin tables for the given absolute `position_ids`
    /// ([SPEC-095]: ALWAYS the true logical positions, never ring slots).
    ///
    /// Vanilla un-scaled base-`theta` RoPE (OQ-5: no YARN / NTK / linear
    /// scaling). `head_dim` must be even.
    ///
    /// # Panics
    /// Panics if `head_dim` is not even.
    #[must_use]
    pub fn build(position_ids: &[usize], head_dim: usize, theta: f32) -> Self {
        assert!(
            head_dim.is_multiple_of(2),
            "RopeTable: head_dim must be even"
        );
        let half = head_dim / 2;
        let seq = position_ids.len();
        let table_len = seq.checked_mul(head_dim);
        assert!(
            table_len.is_some(),
            "RopeTable: seq*head_dim overflow ({seq} * {head_dim})"
        );
        let table_len = table_len.unwrap_or(0);
        if seq == 0 {
            return Self {
                cos: Vec::new(),
                sin: Vec::new(),
                head_dim,
            };
        }
        // inv_freq[i] = theta^(-2i/head_dim) = theta^(-(i/half)).  Computed in
        // f64 then narrowed — matches the HF float32 rotary embedding closely.
        let inv_freq: Vec<f32> = (0..half)
            .map(|i| {
                let exponent = (2 * i) as f64 / head_dim as f64;
                (1.0 / (theta as f64).powf(exponent)) as f32
            })
            .collect();
        let mut cos = vec![0.0f32; table_len];
        let mut sin = vec![0.0f32; table_len];
        for (p_idx, &pos) in position_ids.iter().enumerate() {
            let base = p_idx * head_dim;
            for i in 0..half {
                let angle = pos as f32 * inv_freq[i];
                let (s, c) = angle.sin_cos();
                // Llama cat(freqs, freqs): freq i lands in column i and i+half.
                cos[base + i] = c;
                cos[base + half + i] = c;
                sin[base + i] = s;
                sin[base + half + i] = s;
            }
        }
        Self { cos, sin, head_dim }
    }
}

/// Apply Llama-style RoPE in place to a `[seq, num_heads * head_dim]` Q or K
/// activation ([SPEC-078]).
///
/// Each head's `head_dim` block is rotated by `rotate_half`:
/// `out = x * cos + rotate_half(x) * sin`, where
/// `rotate_half([a; b]) = [-b; a]` over the two contiguous halves (NEOX layout,
/// NOT the DeepseekV2 interleave). `rope` must have been built from the SAME
/// positions as the rows of `x` ([SPEC-095]).
///
/// # Errors
/// [`FocrError::Other`] on a shape mismatch (`x.cols % head_dim != 0`,
/// `x.rows != positions`, …).
pub fn apply_rope(x: &mut Mat, rope: &RopeTable) -> FocrResult<()> {
    checked_mat_len("apply_rope x", x)?;
    let head_dim = rope.head_dim;
    if head_dim == 0 || !x.cols.is_multiple_of(head_dim) {
        return Err(FocrError::Other(anyhow::anyhow!(
            "apply_rope: cols {} not a multiple of head_dim {head_dim}",
            x.cols
        )));
    }
    let seq = x.rows;
    if rope.cos.len() != seq * head_dim {
        return Err(FocrError::Other(anyhow::anyhow!(
            "apply_rope: rope built for {} positions, x has {seq} rows",
            rope.cos.len() / head_dim
        )));
    }
    let num_heads = x.cols / head_dim;
    let half = head_dim / 2;
    for t in 0..seq {
        let rope_base = t * head_dim;
        let row = x.row_mut(t);
        for h in 0..num_heads {
            let hb = h * head_dim;
            // rotate_half over the two halves of THIS head's block.
            for i in 0..half {
                let a = row[hb + i]; // first half
                let b = row[hb + half + i]; // second half
                let cos_a = rope.cos[rope_base + i];
                let sin_a = rope.sin[rope_base + i];
                // rotate_half([a,b]) -> first half gets -b, second half gets a.
                row[hb + i] = a * cos_a - b * sin_a;
                row[hb + half + i] = b * cos_a + a * sin_a;
            }
        }
    }
    Ok(())
}

// ── Dense SwiGLU MLP ([SPEC-075], layer 0 & a shared-expert primitive) ───────

/// SwiGLU feed-forward `down(silu(gate(x)) * up(x))` over `[seq, hidden]`
/// activations ([SPEC-075]).
///
/// `gate_w` / `up_w` are row-major `[inter, hidden]` (`F.linear` weight layout,
/// `out_features x in_features`); `down_w` is `[hidden, inter]`. `act_fn = silu`.
/// This is the layer-0 dense MLP and the exact primitive shape the MoE shared /
/// routed experts reuse, so it lives here next to the driver and is shared by
/// the MoE module if needed.
///
/// # Errors
/// [`FocrError::Other`] on any weight-shape mismatch (propagated from the
/// underlying GEMMs).
pub fn dense_mlp(
    x: &Mat,
    gate_w: &[f32],
    up_w: &[f32],
    down_w: &[f32],
    hidden: usize,
    inter: usize,
) -> FocrResult<Mat> {
    if x.cols != hidden {
        return Err(FocrError::Other(anyhow::anyhow!(
            "dense_mlp: x.cols {} != hidden {hidden}",
            x.cols
        )));
    }
    // gate = x @ gate_w^T  -> [seq, inter]
    let mut gate = linear_no_bias(x, gate_w, hidden, inter)?;
    // up   = x @ up_w^T    -> [seq, inter]
    let up = linear_no_bias(x, up_w, hidden, inter)?;
    // silu(gate) * up, elementwise.
    nn::silu(&mut gate);
    for (g, u) in gate.data.iter_mut().zip(up.data.iter()) {
        *g *= *u;
    }
    // down = (silu(gate)*up) @ down_w^T -> [seq, hidden]
    linear_no_bias(&gate, down_w, inter, hidden)
}

/// Bias-free linear `y = x @ w^T` where `w` is the PyTorch `[out, in]` row-major
/// weight (`F.linear` convention) — transposed on the fly into the `[in, out]`
/// layout [`nn::matmul`] contracts over.
///
/// Small/clear: builds the transpose once (the GEMM dominates), keeping the
/// driver readable. The int8 path ([`nn::linear_int8_dynamic`]) is the
/// kill-switch-gated perf swap and slots in here unchanged.
///
/// # Errors
/// [`FocrError::Other`] if `x.cols != in_` or `w.len() != out * in_`.
pub(crate) fn linear_no_bias(x: &Mat, w: &[f32], in_: usize, out: usize) -> FocrResult<Mat> {
    checked_mat_len("linear_no_bias x", x)?;
    if x.cols != in_ {
        return Err(FocrError::Other(anyhow::anyhow!(
            "linear_no_bias: x.cols {} != in {in_}",
            x.cols
        )));
    }
    let expected_weight_len = checked_shape_mul("linear_no_bias", out, in_, "out*in")?;
    if w.len() != expected_weight_len {
        return Err(FocrError::Other(anyhow::anyhow!(
            "linear_no_bias: w len {} != out*in {}",
            w.len(),
            expected_weight_len
        )));
    }
    // Transpose [out, in] -> [in, out] so matmul does [seq,in] x [in,out].
    let mut wt = vec![0.0f32; expected_weight_len];
    for i in 0..in_ {
        let dst = &mut wt[i * out..(i + 1) * out];
        for (o, slot) in dst.iter_mut().enumerate() {
            *slot = w[o * in_ + i];
        }
    }
    let w_mat = Mat::from_vec(in_, out, wt);
    nn::matmul(x, &w_mat)
}

// ── Final norm + lm_head ([SPEC-071], [SPEC-081]) ───────────────────────────

/// Final `model.norm` RMSNorm then `lm_head` projection `[seq, hidden] ->
/// [seq, vocab]` ([SPEC-071]/[SPEC-081]).
///
/// `norm_w` is the `model.norm.weight` (length `hidden`); `head_w` is the
/// row-major `[vocab, hidden]` `lm_head.weight` (non-tied — `tie_word_embeddings
/// = false`). Returns the f32 logits (HF casts to float for sampling; we are
/// already f32).
///
/// # Errors
/// [`FocrError::Other`] on any shape mismatch.
pub fn norm_and_lm_head(
    hidden: &Mat,
    norm_w: &[f32],
    head_w: &[f32],
    vocab: usize,
    eps: f32,
) -> FocrResult<Mat> {
    let normed = nn::rms_norm(hidden, Some(norm_w), eps)?;
    lm_head_proj(&normed, head_w, vocab)
}

/// [`norm_and_lm_head`] with the head weight **already transposed** to a matmul-ready
/// `[hidden, vocab]` [`Mat`] — **bit-identical** logits (same `rms_norm`, same
/// `nn::matmul`), but WITHOUT re-transposing the (large, tied) embedding on every call.
///
/// The naive [`lm_head_proj`]/[`linear_no_bias`] path transposes the whole
/// `[vocab, hidden]` head matrix per call (~0.6 GB for GOT's 151860×1024). The
/// O(n)-per-token KV-cache decode invokes the head once per generated token, so that
/// redundant transpose measured as **~95% of decode wall-clock** (463 s of a 487 s
/// page). Loop callers build the transpose ONCE and pass it here; the one-shot seeding
/// prefill keeps the naive path. Numerically identical, so parity/oracle gates hold.
///
/// # Errors
/// [`FocrError::Other`] on any `rms_norm`/`matmul` shape mismatch (e.g.
/// `head_wt.rows != hidden.cols`).
pub(crate) fn norm_and_lm_head_pretransposed(
    hidden: &Mat,
    norm_w: &[f32],
    head_wt: &Mat,
    eps: f32,
) -> FocrResult<Mat> {
    let normed = nn::rms_norm(hidden, Some(norm_w), eps)?;
    nn::matmul(&normed, head_wt)
}

/// Project (already-normed) hidden states to vocab logits — the bare
/// `lm_head` GEMM `[seq, hidden] x [vocab, hidden]^T -> [seq, vocab]`
/// ([SPEC-081]).
///
/// `head_w` is row-major `[vocab, hidden]`. For the single-token decode step
/// (`seq == 1`) this is a GEMV.
///
/// # Errors
/// [`FocrError::Other`] if `head_w.len() != vocab * hidden`.
pub fn lm_head_proj(hidden: &Mat, head_w: &[f32], vocab: usize) -> FocrResult<Mat> {
    let h = hidden.cols;
    linear_no_bias(hidden, head_w, h, vocab)
}

// ── Per-layer driver ([SPEC-072]) ──────────────────────────────────────────

/// The fp32 weights one decoder layer needs from the driver's point of view.
///
/// Norms + projections are explicit slices so the driver is fully testable
/// before the `.focrq` reader exists; the attention/MLP kernels consume the
/// rest through their own modules. `gate_w`/`up_w`/`down_w` are populated for the
/// **dense** layer-0 path; MoE layers (1..11) route through
/// [`crate::native_engine::moe::forward`] and ignore them.
#[derive(Debug, Clone)]
pub struct LayerWeights<'a> {
    /// `input_layernorm.weight`, length `hidden`.
    pub input_ln: &'a [f32],
    /// `post_attention_layernorm.weight`, length `hidden`.
    pub post_attn_ln: &'a [f32],
    /// `q_proj.weight`, row-major `[num_heads*head_dim, hidden]`.
    pub q_proj: &'a [f32],
    /// `k_proj.weight`, row-major `[num_kv_heads*head_dim, hidden]`.
    pub k_proj: &'a [f32],
    /// `v_proj.weight`, row-major `[num_kv_heads*head_dim, hidden]`.
    pub v_proj: &'a [f32],
    /// `o_proj.weight`, row-major `[hidden, num_heads*head_dim]`.
    pub o_proj: &'a [f32],
    /// Dense `gate_proj.weight` `[inter, hidden]` (layer 0 only).
    pub gate_w: &'a [f32],
    /// Dense `up_proj.weight` `[inter, hidden]` (layer 0 only).
    pub up_w: &'a [f32],
    /// Dense `down_proj.weight` `[hidden, inter]` (layer 0 only).
    pub down_w: &'a [f32],
}

/// Project Q/K/V from the (already input-normed) hidden states and apply RoPE to
/// Q and K ([SPEC-078]/[SPEC-090]).
///
/// Returns `(q, k, v)`, each `[seq, num_heads*head_dim]`, ready to hand to
/// [`crate::native_engine::rswa::attention`]. No QKV bias ([SPEC-090],
/// `attention_bias=false`).
///
/// # Errors
/// [`FocrError::Other`] on a projection / RoPE shape mismatch.
pub fn qkv_with_rope(
    normed: &Mat,
    lw: &LayerWeights<'_>,
    rope: &RopeTable,
    hidden: usize,
    qkv_dim: usize,
) -> FocrResult<(Mat, Mat, Mat)> {
    let mut q = linear_no_bias(normed, lw.q_proj, hidden, qkv_dim)?;
    let mut k = linear_no_bias(normed, lw.k_proj, hidden, qkv_dim)?;
    let v = linear_no_bias(normed, lw.v_proj, hidden, qkv_dim)?;
    apply_rope(&mut q, rope)?;
    apply_rope(&mut k, rope)?;
    Ok((q, k, v))
}

/// Self-attention sub-block for one layer, *given* the raw attention context
/// (the `[seq, num_heads*head_dim]` output of R-SWA, BEFORE `o_proj`).
///
/// Applies `o_proj`: `attn_out = context @ o_proj^T`, returning `[seq, hidden]`.
/// Split out so the driver can call the real [`crate::native_engine::rswa`]
/// kernel for the context and keep the projection here (and so it is unit
/// testable without the ring cache).
///
/// # Errors
/// [`FocrError::Other`] on a shape mismatch.
pub fn attn_output_proj(
    context: &Mat,
    o_proj: &[f32],
    hidden: usize,
    qkv_dim: usize,
) -> FocrResult<Mat> {
    linear_no_bias(context, o_proj, qkv_dim, hidden)
}

/// One full decoder layer, pre-norm residual ([SPEC-072]):
/// `h = x + o_proj(attn(rope(qkv(input_ln(x)))))`;
/// `out = h + mlp(post_attention_ln(h))`.
///
/// `attn_context` is a closure that runs the R-SWA attention kernel over the
/// rope'd `(q, k, v)` and returns the pre-`o_proj` context `[seq,
/// num_heads*head_dim]` — injected so this driver is testable against a pure
/// reference attention (and so the real ring-cache kernel
/// [`crate::native_engine::rswa::attention`] plugs in at the call site without
/// this function owning the per-layer cache). `mlp` is a closure that runs the
/// layer's MLP/MoE block over `[seq, hidden]` (dense for layer 0 via
/// [`dense_mlp`]; [`crate::native_engine::moe::forward`] for layers 1..11).
///
/// # Errors
/// Propagates any sub-step error.
// decoder forward: args are model state + tensors
#[allow(clippy::too_many_arguments)]
pub fn layer_forward<A, M>(
    x: &Mat,
    lw: &LayerWeights<'_>,
    rope: &RopeTable,
    hidden: usize,
    qkv_dim: usize,
    eps: f32,
    attn_context: A,
    mlp: M,
) -> FocrResult<Mat>
where
    A: FnOnce(&Mat, &Mat, &Mat) -> FocrResult<Mat>,
    M: FnOnce(&Mat) -> FocrResult<Mat>,
{
    // --- attention sub-block ---
    let normed = nn::rms_norm(x, Some(lw.input_ln), eps)?;
    let (q, k, v) = qkv_with_rope(&normed, lw, rope, hidden, qkv_dim)?;
    let context = attn_context(&q, &k, &v)?;
    let attn_out = attn_output_proj(&context, lw.o_proj, hidden, qkv_dim)?;
    let h = add_residual(x, &attn_out)?;

    // --- MLP / MoE sub-block ---
    let normed2 = nn::rms_norm(&h, Some(lw.post_attn_ln), eps)?;
    let mlp_out = mlp(&normed2)?;
    add_residual(&h, &mlp_out)
}

/// Elementwise residual add `a + b` for two equal-shaped activation [`Mat`]s.
///
/// # Errors
/// [`FocrError::Other`] if the shapes differ.
pub fn add_residual(a: &Mat, b: &Mat) -> FocrResult<Mat> {
    checked_mat_len("add_residual lhs", a)?;
    checked_mat_len("add_residual rhs", b)?;
    if a.shape() != b.shape() {
        return Err(FocrError::Other(anyhow::anyhow!(
            "add_residual: shape mismatch {:?} vs {:?}",
            a.shape(),
            b.shape()
        )));
    }
    let data = a
        .data
        .iter()
        .zip(b.data.iter())
        .map(|(x, y)| x + y)
        .collect::<Vec<_>>();
    Ok(Mat::from_vec(a.rows, a.cols, data))
}

// ── Prefill self-attention (full causal MHA over the whole sequence) ────────

/// Full causal multi-head self-attention over the prefill, *given* the already
/// RoPE'd `(q, k, v)` ([SPEC-090], baidu `SlidingWindowLlamaAttention` prefill
/// path `_attn_forward`).
///
/// In prefill the **entire** sequence is the R-SWA reference block (no token is
/// ever evicted), so a plain lower-triangular causal mask is exactly the
/// reference math — the ring buffer only ever restricts the *generated* tail
/// during incremental decode. `q`/`k`/`v` are token-major `[seq,
/// num_heads*head_dim]` (head `h` in columns `h*head_dim .. (h+1)*head_dim`).
///
/// Transposes to head-major `[num_heads, seq, head_dim]`, runs [`nn::sdpa`] with
/// `scale = 1/sqrt(head_dim)` and `causal = true`, then repacks the context back
/// to token-major `[seq, num_heads*head_dim]` for `o_proj`. MHA: `num_kv_heads ==
/// num_heads` (`repeat_kv` is a no-op), no QKV bias.
///
/// # Errors
/// [`FocrError::Other`] on a shape mismatch (`q/k/v` cols not `num_heads *
/// head_dim`, or `q/k/v` rows disagree).
pub fn prefill_attention(
    q: &Mat,
    k: &Mat,
    v: &Mat,
    num_heads: usize,
    head_dim: usize,
) -> FocrResult<Mat> {
    checked_mat_len("prefill_attention q", q)?;
    checked_mat_len("prefill_attention k", k)?;
    checked_mat_len("prefill_attention v", v)?;
    let dim = checked_shape_mul(
        "prefill_attention",
        num_heads,
        head_dim,
        "num_heads*head_dim",
    )?;
    let seq = q.rows;
    if q.cols != dim || k.cols != dim || v.cols != dim {
        return Err(FocrError::Other(anyhow::anyhow!(
            "prefill_attention: q/k/v cols ({},{},{}) != num_heads*head_dim {dim}",
            q.cols,
            k.cols,
            v.cols
        )));
    }
    if k.rows != seq || v.rows != seq {
        return Err(FocrError::Other(anyhow::anyhow!(
            "prefill_attention: q/k/v rows disagree ({}, {}, {})",
            seq,
            k.rows,
            v.rows
        )));
    }

    // Token-major [seq, num_heads*head_dim] -> head-major [num_heads, seq, head_dim].
    let span = seq * head_dim;
    let mut qh = vec![0.0f32; num_heads * span];
    let mut kh = vec![0.0f32; num_heads * span];
    let mut vh = vec![0.0f32; num_heads * span];
    for s in 0..seq {
        let (qr, kr, vr) = (q.row(s), k.row(s), v.row(s));
        for h in 0..num_heads {
            let src = h * head_dim;
            let dst = h * span + s * head_dim;
            qh[dst..dst + head_dim].copy_from_slice(&qr[src..src + head_dim]);
            kh[dst..dst + head_dim].copy_from_slice(&kr[src..src + head_dim]);
            vh[dst..dst + head_dim].copy_from_slice(&vr[src..src + head_dim]);
        }
    }

    // num_bh = batch(1) * num_heads; causal lower-triangular mask.
    let scale = 1.0f32 / (head_dim as f32).sqrt();
    let ctx = nn::sdpa(
        &qh, &kh, &vh, num_heads, seq, seq, head_dim, head_dim, scale, true,
    );
    let expected = num_heads * span;
    if ctx.len() != expected {
        return Err(FocrError::Other(anyhow::anyhow!(
            "prefill_attention: sdpa context len {} != expected {expected}",
            ctx.len()
        )));
    }

    // Head-major [num_heads, seq, head_dim] -> token-major [seq, num_heads*head_dim].
    let mut out = Mat::zeros(seq, dim);
    for h in 0..num_heads {
        for s in 0..seq {
            let src = h * span + s * head_dim;
            let dst = s * dim + h * head_dim;
            out.data[dst..dst + head_dim].copy_from_slice(&ctx[src..src + head_dim]);
        }
    }
    Ok(out)
}

// ── Top-level entrypoints (Weights-backed) ──────────────────────────────────

/// Final `model.norm` RMSNorm + `lm_head` projection over the decoder hidden
/// states ([SPEC-071]/[SPEC-081]).
///
/// Pulls `model.norm.weight` (length `hidden`) and the **top-level**
/// `lm_head.weight` (`[vocab, hidden] = [129280, 1280]`, NOT `model.lm_head`;
/// `tie_word_embeddings = false`) and composes the tested
/// [`norm_and_lm_head`]. Returns the f32 logits `[seq, vocab]`; the decode driver
/// feeds only the last hidden row, yielding `[1, vocab]`.
///
/// # Errors
/// [`FocrError::FormatMismatch`] if `model.norm.weight` / `lm_head.weight` are
/// absent, or [`FocrError::Other`] on a shape mismatch.
pub fn lm_head(weights: &Weights, hidden: &Mat) -> FocrResult<Mat> {
    let norm_w = weights.vec("model.norm.weight")?;
    let head = weights.mat("lm_head.weight")?;
    norm_and_lm_head(
        hidden,
        &norm_w,
        &head.data,
        config::VOCAB_SIZE,
        config::RMS_NORM_EPS,
    )
}

/// Run the full 12-layer DeepSeek-V2 MoE decoder over `inputs_embeds`,
/// returning the final `model.norm`-ready hidden states `[seq, hidden]`
/// ([SPEC-072], baidu `DeepseekV2Model.forward`).
///
/// Stateless, full-sequence (no KV cache): correct for **prefill / parity**. In
/// prefill every token is an R-SWA reference token, so attention is plain full
/// causal MHA (the ring cache only constrains incremental decode, which threads
/// a separate `&mut [RingCache; 12]` — out of scope here). Per layer 0..11:
///
/// 1. RMSNorm(`input_layernorm.weight`) -> Q/K/V proj (separate
///    `self_attn.{q,k,v}_proj.weight`) -> RoPE (theta 10000, NEOX `rotate_half`)
///    -> full causal SDPA ([`prefill_attention`]) -> `o_proj` -> residual.
/// 2. RMSNorm(`post_attention_layernorm.weight`) -> MLP: layer 0 dense
///    ([`moe::dense_forward`]), layers 1..11 MoE ([`moe::forward`]) -> residual.
///
/// Reuses the unit-tested [`layer_forward`] / [`qkv_with_rope`] / [`apply_rope`]
/// driver; the per-layer weight slices are pulled fresh from [`Weights`].
///
/// # Errors
/// [`FocrError::FormatMismatch`] on an absent tensor; [`FocrError::Other`] on a
/// shape mismatch or kernel error (propagated from the sub-steps).
pub fn forward(weights: &Weights, inputs_embeds: &Mat) -> FocrResult<Mat> {
    checked_mat_len("decoder::forward inputs_embeds", inputs_embeds)?;
    let hidden = config::HIDDEN_SIZE;
    let qkv_dim = checked_shape_mul(
        "decoder::forward",
        config::NUM_ATTENTION_HEADS,
        config::HEAD_DIM,
        "num_heads*head_dim",
    )?;
    let eps = config::RMS_NORM_EPS;
    if inputs_embeds.cols != hidden {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::forward: inputs_embeds cols {} != hidden {hidden}",
            inputs_embeds.cols
        )));
    }
    let seq = inputs_embeds.rows;

    // Absolute position_ids 0..seq ([SPEC-095]: true logical positions); one
    // RoPE table shared across all 12 layers.
    let positions: Vec<usize> = (0..seq).collect();
    let rope = RopeTable::build(&positions, config::HEAD_DIM, config::ROPE_THETA);

    let mut x = inputs_embeds.clone();
    for layer in 0..config::NUM_HIDDEN_LAYERS {
        let prefix = format!("model.layers.{layer}");
        // Owned per-layer weight Mats/Vecs; LayerWeights borrows their slices.
        let input_ln = weights.vec(&format!("{prefix}.input_layernorm.weight"))?;
        let post_attn_ln = weights.vec(&format!("{prefix}.post_attention_layernorm.weight"))?;
        let q_proj = weights.mat(&format!("{prefix}.self_attn.q_proj.weight"))?;
        let k_proj = weights.mat(&format!("{prefix}.self_attn.k_proj.weight"))?;
        let v_proj = weights.mat(&format!("{prefix}.self_attn.v_proj.weight"))?;
        let o_proj = weights.mat(&format!("{prefix}.self_attn.o_proj.weight"))?;

        let lw = LayerWeights {
            input_ln: &input_ln,
            post_attn_ln: &post_attn_ln,
            q_proj: &q_proj.data,
            k_proj: &k_proj.data,
            v_proj: &v_proj.data,
            o_proj: &o_proj.data,
            // Dense gate/up/down are pulled inside the MLP closure (moe::*), so
            // the driver's vestigial dense slots stay empty here.
            gate_w: &[],
            up_w: &[],
            down_w: &[],
        };

        x = layer_forward(
            &x,
            &lw,
            &rope,
            hidden,
            qkv_dim,
            eps,
            |q, k, v| prefill_attention(q, k, v, config::NUM_ATTENTION_HEADS, config::HEAD_DIM),
            |normed| {
                if layer < config::FIRST_K_DENSE_REPLACE {
                    moe::dense_forward(weights, normed)
                } else {
                    moe::forward(weights, normed, layer)
                }
            },
        )?;
    }
    Ok(x)
}

/// Repack token-major `[seq, num_heads*head_dim]` K/V (the layout
/// [`qkv_with_rope`] returns) into the head-major `[num_heads, seq, head_dim]`
/// flat layout that [`RingCache::record_prefill`] consumes (head `h`'s `seq`
/// rows are contiguous). Pure index gymnastics — no math.
///
/// # Errors
/// [`FocrError::Other`] on a `seq*head_dim`/`num_heads*…` overflow.
fn token_major_to_head_major(
    k: &Mat,
    v: &Mat,
    seq: usize,
    num_heads: usize,
    head_dim: usize,
) -> FocrResult<(Vec<f32>, Vec<f32>)> {
    let span = checked_shape_mul("token_major_to_head_major", seq, head_dim, "seq*head_dim")?;
    let total = checked_shape_mul(
        "token_major_to_head_major",
        num_heads,
        span,
        "num_heads*seq*head_dim",
    )?;
    let mut kh = vec![0.0f32; total];
    let mut vh = vec![0.0f32; total];
    for s in 0..seq {
        let (kr, vr) = (k.row(s), v.row(s));
        for h in 0..num_heads {
            let src = h * head_dim;
            let dst = h * span + s * head_dim;
            kh[dst..dst + head_dim].copy_from_slice(&kr[src..src + head_dim]);
            vh[dst..dst + head_dim].copy_from_slice(&vr[src..src + head_dim]);
        }
    }
    Ok((kh, vh))
}

// ── Conservative mixed-precision decoder cache ──────────────────────────────
//
// Doctrine #2 quantizes only the dense/FFN/expert GEMMs. Attention q/k/v/o and
// `lm_head` stay high precision. This cache mirrors that recipe exactly: it
// widens the relatively small gated/high-precision set once and keeps the 2,148
// FFN/expert matrices in their per-channel int8 form. It therefore avoids the
// former ~10.5 GB all-f32 cache without silently opting gated tensors into int8.

/// Owned conservative-recipe decoder weights, built once and borrowed by every
/// prefill/decode step.
pub struct DecoderWeightCache {
    layers: Vec<CachedLayer>,
    /// Final `model.norm.weight` (RMSNorm before the head).
    final_norm: Vec<f32>,
    /// `lm_head.weight`, widened once to row-major f32 `[vocab, hidden]`.
    lm_head: Vec<f32>,
}

/// One decoder layer's mixed-precision weights.
struct CachedLayer {
    input_ln: Vec<f32>,
    post_attn_ln: Vec<f32>,
    q_proj: Vec<f32>,
    k_proj: Vec<f32>,
    v_proj: Vec<f32>,
    o_proj: Vec<f32>,
    mlp: CachedMlpI8,
}

impl DecoderWeightCache {
    /// Build the fixed validated recipe from [`Weights`]: attention and
    /// `lm_head` high precision, FFN/expert projections int8, router/norms f32.
    ///
    /// # Errors
    /// [`FocrError::FormatMismatch`] if any expected tensor is absent/mis-shaped.
    pub fn build(weights: &Weights) -> FocrResult<Self> {
        let mut layers = Vec::with_capacity(config::NUM_HIDDEN_LAYERS);
        for layer in 0..config::NUM_HIDDEN_LAYERS {
            let prefix = format!("model.layers.{layer}");
            let input_ln = weights.vec(&format!("{prefix}.input_layernorm.weight"))?;
            let post_attn_ln = weights.vec(&format!("{prefix}.post_attention_layernorm.weight"))?;
            let q_proj = weights
                .mat(&format!("{prefix}.self_attn.q_proj.weight"))?
                .data;
            let k_proj = weights
                .mat(&format!("{prefix}.self_attn.k_proj.weight"))?
                .data;
            let v_proj = weights
                .mat(&format!("{prefix}.self_attn.v_proj.weight"))?
                .data;
            let o_proj = weights
                .mat(&format!("{prefix}.self_attn.o_proj.weight"))?
                .data;
            // FFN/expert projections go through `quant_ffn_loaded`: qint8 (or
            // bf16 quantized at load) exactly as before, or PACKED int4 kept
            // packed when the wasm experts-int4 artifact provides it (bd-4l71).
            let mlp = if layer < config::FIRST_K_DENSE_REPLACE {
                let p = format!("{prefix}.mlp");
                let inter = moe::config::DENSE_INTERMEDIATE_SIZE;
                CachedMlpI8::Dense {
                    gate: quant_ffn_loaded(weights, &format!("{p}.gate_proj.weight"), inter)?,
                    up: quant_ffn_loaded(weights, &format!("{p}.up_proj.weight"), inter)?,
                    down: quant_ffn_loaded(
                        weights,
                        &format!("{p}.down_proj.weight"),
                        config::HIDDEN_SIZE,
                    )?,
                }
            } else {
                let p = format!("{prefix}.mlp");
                let gate = weights.mat(&format!("{p}.gate.weight"))?.data;
                let inter = moe::config::MOE_INTERMEDIATE_SIZE;
                let mut experts = Vec::with_capacity(moe::config::N_ROUTED_EXPERTS);
                for e in 0..moe::config::N_ROUTED_EXPERTS {
                    experts.push([
                        quant_ffn_loaded(
                            weights,
                            &format!("{p}.experts.{e}.gate_proj.weight"),
                            inter,
                        )?,
                        quant_ffn_loaded(
                            weights,
                            &format!("{p}.experts.{e}.up_proj.weight"),
                            inter,
                        )?,
                        quant_ffn_loaded(
                            weights,
                            &format!("{p}.experts.{e}.down_proj.weight"),
                            config::HIDDEN_SIZE,
                        )?,
                    ]);
                }
                let shared_inter = moe::config::SHARED_INTERMEDIATE_SIZE;
                let shared = [
                    quant_ffn_loaded(
                        weights,
                        &format!("{p}.shared_experts.gate_proj.weight"),
                        shared_inter,
                    )?,
                    quant_ffn_loaded(
                        weights,
                        &format!("{p}.shared_experts.up_proj.weight"),
                        shared_inter,
                    )?,
                    quant_ffn_loaded(
                        weights,
                        &format!("{p}.shared_experts.down_proj.weight"),
                        config::HIDDEN_SIZE,
                    )?,
                ];
                CachedMlpI8::Moe {
                    gate,
                    experts,
                    shared,
                }
            };
            layers.push(CachedLayer {
                input_ln,
                post_attn_ln,
                q_proj,
                k_proj,
                v_proj,
                o_proj,
                mlp,
            });
        }
        let final_norm = weights.vec("model.norm.weight")?;
        let lm_head = weights.mat("lm_head.weight")?.data;
        Ok(Self {
            layers,
            final_norm,
            lm_head,
        })
    }
}

/// Borrow a cached layer's attention weights as a [`LayerWeights`] (MLP slots
/// stay empty — the MLP is run via [`cached_mlp`]).
fn cached_layer_weights(cl: &CachedLayer) -> LayerWeights<'_> {
    LayerWeights {
        input_ln: &cl.input_ln,
        post_attn_ln: &cl.post_attn_ln,
        q_proj: &cl.q_proj,
        k_proj: &cl.k_proj,
        v_proj: &cl.v_proj,
        o_proj: &cl.o_proj,
        gate_w: &[],
        up_w: &[],
        down_w: &[],
    }
}

/// Run the recipe-approved int8 MLP/MoE for one prefill activation block.
///
/// `calib_layer` is the decoder layer index when the activation-statistics
/// recorder is armed (`FOCR_CALIB_OUT`, bd-50wo stage A) and `None` otherwise;
/// it only ever selects the stat key the recorder files this block's inputs
/// under, and is inert when the recorder is off.
fn cached_mlp(mlp: &CachedMlpI8, normed: &Mat, calib_layer: Option<usize>) -> FocrResult<Mat> {
    prefill_mlp_i8(mlp, normed, calib_layer)
}

// ── Decode phase profiler (FOCR_PROFILE_DECODE) ──────────────────────────────
//
// Accumulates wall-nanoseconds per decode phase across all tokens so the
// throughput bottleneck is attributable (lm_head vs attention vs MoE experts vs
// routing). A handful of `Instant::now()` (~20 ns) + `fetch_add` per token —
// negligible vs the ~23 ms/token it measures, and entirely skipped when off.
pub mod prof {
    use std::sync::OnceLock;
    use std::sync::atomic::{AtomicU64, Ordering};

    /// Final RMSNorm + `lm_head` GEMV.
    pub static LMHEAD_NS: AtomicU64 = AtomicU64::new(0);
    /// Attention sub-block: q/k/v/o projections + RoPE + ring + R-SWA attention.
    pub static ATTN_NS: AtomicU64 = AtomicU64::new(0);
    /// MoE/MLP routed + shared experts (the SwiGLU GEMVs), excluding routing.
    pub static EXPERTS_NS: AtomicU64 = AtomicU64::new(0);
    /// MoE top-k router (softmax over the 64-expert gate).
    pub static ROUTE_NS: AtomicU64 = AtomicU64::new(0);

    pub fn enabled() -> bool {
        static FLAG: OnceLock<bool> = OnceLock::new();
        *FLAG.get_or_init(|| std::env::var_os("FOCR_PROFILE_DECODE").is_some())
    }
    #[inline]
    pub fn add(c: &AtomicU64, ns: u64) {
        c.fetch_add(ns, Ordering::Relaxed);
    }
    /// Reset all counters (call before a timed decode loop).
    pub fn reset() {
        for c in [&LMHEAD_NS, &ATTN_NS, &EXPERTS_NS, &ROUTE_NS] {
            c.store(0, Ordering::Relaxed);
        }
    }
    /// `(lmhead, attn, experts, route)` accumulated milliseconds.
    pub fn snapshot_ms() -> (f64, f64, f64, f64) {
        let ms = |c: &AtomicU64| c.load(Ordering::Relaxed) as f64 / 1e6;
        (ms(&LMHEAD_NS), ms(&ATTN_NS), ms(&EXPERTS_NS), ms(&ROUTE_NS))
    }
}

// ── Bespoke single-token (m=1) decode GEMV — the decode-throughput kernel ────
//
// Decode is one token at a time: every projection is a matrix·vector product
// `y[N] = W[N,K] · x[K]`, with W stored `[out, in]` row-major (output-channel
// major) — so each output row is a CONTIGUOUS `K`-length dot product. The
// generic GEMM path transposed W every call and fell below ft-kernel-cpu's
// threading threshold (m=1), running single-threaded at ~0.4 GFLOP/s. This
// kernel instead: (1) no transpose — dots the native row layout; (2) fans the N
// independent rows across cores with rayon; (3) an 8-wide unrolled inner dot
// that LLVM lowers to NEON/AVX FMA. This is intentionally model-specialized, not
// a general BLAS — exactly the decode hot path and nothing else.

/// 8-wide unrolled f32 dot product (`sum a[i]*b[i]`). The fixed 8-lane reduction
/// order is deterministic and auto-vectorizes to NEON/AVX FMA under `-O3`.
#[inline]
fn dot_f32(a: &[f32], b: &[f32]) -> f32 {
    debug_assert_eq!(a.len(), b.len());
    let n = a.len();
    let mut acc = [0.0f32; 8];
    let chunks = n / 8;
    for c in 0..chunks {
        let base = c * 8;
        // Indexed (not iter) so the 8 independent FMAs vectorize cleanly.
        for (l, slot) in acc.iter_mut().enumerate() {
            *slot += a[base + l] * b[base + l];
        }
    }
    let mut s = ((acc[0] + acc[1]) + (acc[2] + acc[3])) + ((acc[4] + acc[5]) + (acc[6] + acc[7]));
    for i in (chunks * 8)..n {
        s += a[i] * b[i];
    }
    s
}

/// Single-token GEMV: `y[o] = dot(x[0..k], w[o*k .. o*k+k])` for `o in 0..n`,
/// over `w` in `[n, k]` row-major layout. The `n` output rows are independent →
/// fanned across the rayon pool. Returns the `[n]` result.
fn gemv(x: &[f32], w: &[f32], n: usize, k: usize) -> Vec<f32> {
    debug_assert_eq!(x.len(), k);
    debug_assert_eq!(w.len(), n * k);
    let mut y = vec![0.0f32; n];
    // Chunked parallelism keeps per-task work well above rayon's dispatch cost
    // even for the small projections; the huge lm_head row-splits across cores.
    y.par_chunks_mut(64).enumerate().for_each(|(blk, ys)| {
        let base = blk * 64;
        for (j, slot) in ys.iter_mut().enumerate() {
            let o = base + j;
            *slot = dot_f32(x, &w[o * k..o * k + k]);
        }
    });
    y
}

/// SiLU / swish: `x * sigmoid(x) = x / (1 + e^-x)` ([SPEC-075], the SwiGLU gate).
#[inline]
fn silu(x: f32) -> f32 {
    x / (1.0 + (-x).exp())
}

/// Decode the recipe-approved int8 MLP/MoE for one normalized token row.
/// `calib_layer` as in [`cached_mlp`].
fn decode_mlp(mlp: &CachedMlpI8, normed: &Mat, calib_layer: Option<usize>) -> FocrResult<Vec<f32>> {
    decode_mlp_i8(mlp, normed, calib_layer)
}

/// Final RMSNorm + `lm_head` over the decode hidden (`[1, hidden]`), off the
/// cached head weights via the bespoke [`gemv`] — the per-token `[vocab, hidden]`
/// projection was a major decode cost run single-threaded by the GEMM path.
///
/// # Errors
/// [`FocrError::Other`] on a shape mismatch.
pub fn lm_head_cached(wc: &DecoderWeightCache, hidden: &Mat) -> FocrResult<Mat> {
    if hidden.rows != 1 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::lm_head_cached: expected a single decode row, got {} rows",
            hidden.rows
        )));
    }
    let t = prof::enabled().then(Instant::now);
    let normed = nn::rms_norm(hidden, Some(&wc.final_norm), config::RMS_NORM_EPS)?;
    let row = normed.row(0);
    if calib::enabled() {
        calib::record_row(quant_calib::LM_HEAD_IN_KEY, row);
    }
    // FOCR_LMHEAD_SHARD: vocab-tiled head (default OFF ⇒ the monolithic `gemv`).
    // Byte-for-byte identical either way — each logit is an independent dot.
    let logits = if lmhead_shard_enabled() {
        gemv_sharded(
            row,
            &wc.lm_head,
            config::VOCAB_SIZE,
            config::HIDDEN_SIZE,
            lmhead_shard_tiles(),
        )
    } else {
        gemv(row, &wc.lm_head, config::VOCAB_SIZE, config::HIDDEN_SIZE)
    };
    if let Some(t) = t {
        prof::add(&prof::LMHEAD_NS, t.elapsed().as_nanos() as u64);
    }
    Ok(Mat::from_vec(1, config::VOCAB_SIZE, logits))
}

// ── Chunked prefill (FOCR_PREFILL_CHUNK, bd-1azu.9) ──────────────────────────
//
// Chunked prefill co-schedules a slice of the prompt into the decode batch as a
// "mixed prefill/decode forward": instead of running the WHOLE prefill through
// each layer in one monolithic SDPA, it consumes `C` tokens at a time, pushing
// each chunk through all 12 layers (writing its K/V into the per-layer rings)
// before the next chunk. This is the front end of the continuous-batch spine —
// a prefill chunk is processed like a decode step, one pass through the layers,
// so it can later share the batched per-layer GEMMs with in-flight decode rows.
//
// LOSSLESS by construction: chunking is a pure TILING of the SAME causal
// attention. Every per-token op outside attention (RMSNorm, the q/k/v/o
// projections, RoPE at the TRUE absolute position, the dense/MoE MLP, the
// residual adds) is independent ROW-by-ROW, so a chunk's rows are byte-identical
// whether projected alone or as part of the whole sequence. The only cross-token
// op is attention, and for chunk `[c0, c1)` each new query at global position
// `t ∈ [c0, c1)` attends EXACTLY the prior tokens it would monolithically —
// keys `[0, t]` (earlier chunks already written into the running K/V, plus this
// chunk up to `t`) — in the SAME ascending reduction order. The trailing
// (future) keys a monolithic SDPA carries are masked to 0 for row `t` and add
// exact `0.0`, so restricting the key set to `[0, c1)` is byte-for-byte the same
// reduction (see [`chunk_prefill_attention`]). Hence the final hidden AND every
// layer's ring K/V equal the monolithic [`prefill_with_cache`] output exactly.

/// `FOCR_PREFILL_CHUNK`: kill-switch arming chunked prefill. UNSET ⇒ today's
/// monolithic [`prefill_with_cache`] / [`prefill_with_cache_i8`] path, byte-for-
/// byte. When present its value is the chunk size `C` (tokens consumed per
/// pass); present-but-invalid (`empty`/unparseable/`0`) falls back to
/// [`DEFAULT_PREFILL_CHUNK`] so mere presence still arms the lever, exactly like
/// [`BATCH_SIZE_ENV`].
const PREFILL_CHUNK_ENV: &str = "FOCR_PREFILL_CHUNK";

/// Fallback chunk size when [`PREFILL_CHUNK_ENV`] is present but unparseable/`0`.
const DEFAULT_PREFILL_CHUNK: usize = 256;

/// The configured prefill chunk size ([`PREFILL_CHUNK_ENV`], read ONCE into a
/// process-global per doctrine — never re-read per prefill). `None` ⇒ the
/// monolithic path (the unset default); `Some(C)` ⇒ chunk `C > 0` tokens per
/// pass. A chunk `>=` the sequence length collapses to a single chunk, which is
/// itself byte-for-byte the monolithic path.
#[must_use]
pub fn prefill_chunk_size() -> Option<usize> {
    static SIZE: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
    *SIZE.get_or_init(|| {
        // unset ⇒ monolithic prefill, byte-for-byte
        std::env::var_os(PREFILL_CHUNK_ENV)?;
        Some(
            std::env::var(PREFILL_CHUNK_ENV)
                .ok()
                .and_then(|v| v.trim().parse::<usize>().ok())
                .filter(|&n| n > 0)
                .unwrap_or(DEFAULT_PREFILL_CHUNK),
        )
    })
}

/// Tile `[0, seq)` into CONTIGUOUS, gap-free, ascending `[c0, c1)` chunks of at
/// most `chunk` tokens (the last is the `seq % chunk` remainder). `chunk` is
/// clamped to `>= 1`. The ascending coverage is what lets chunk `g`'s attention
/// see exactly the keys earlier chunks already wrote — and preserves token order.
fn prefill_chunk_bounds(seq: usize, chunk: usize) -> Vec<(usize, usize)> {
    let chunk = chunk.max(1);
    let mut bounds = Vec::with_capacity(seq.div_ceil(chunk));
    let mut c0 = 0usize;
    while c0 < seq {
        let c1 = (c0 + chunk).min(seq);
        bounds.push((c0, c1));
        c0 = c1;
    }
    bounds
}

/// Causal self-attention for ONE prefill chunk: the `cs = c1 - c0` new queries
/// `q_chunk` (global positions `[c0, c1)`, already RoPE'd) attend over the
/// running reference keys/values `k_prefix`/`v_prefix` — the FULL `[0, c1)`
/// prefix (earlier chunks ++ this chunk's own K/V) — under the triangular causal
/// mask, returning the `[cs, num_heads*head_dim]` context (pre-`o_proj`).
///
/// BYTE-FOR-BYTE identical to the corresponding rows of the monolithic
/// [`prefill_attention`] over the whole sequence. The trick: front-pad the
/// queries to the `c1` rows the prefix spans (rows `[0, c0)` are scratch and
/// discarded) so [`prefill_attention`]'s top-left causal mask (`limit = row+1`)
/// lands each REAL query row `t` on keys `[0, t]` — exactly its monolithic
/// reach. Each kept row's score dot (over `head_dim`), softmax (ascending over
/// `[0, t]`), and `P·V` (ascending, with the future keys a monolithic pass would
/// carry contributing exact `0.0`) are unchanged, and the scratch rows are
/// row-independent in SDPA, so they never perturb a kept row.
///
/// # Errors
/// [`FocrError::Other`] on a shape mismatch (`q_chunk`/`k_prefix`/`v_prefix`
/// widths disagree, `c0 + cs != c1`, or `c0 > c1`).
pub fn chunk_prefill_attention(
    q_chunk: &Mat,
    k_prefix: &Mat,
    v_prefix: &Mat,
    num_heads: usize,
    head_dim: usize,
    c0: usize,
) -> FocrResult<Mat> {
    checked_mat_len("chunk_prefill_attention q_chunk", q_chunk)?;
    checked_mat_len("chunk_prefill_attention k_prefix", k_prefix)?;
    checked_mat_len("chunk_prefill_attention v_prefix", v_prefix)?;
    let dim = checked_shape_mul(
        "chunk_prefill_attention",
        num_heads,
        head_dim,
        "num_heads*head_dim",
    )?;
    let c1 = k_prefix.rows;
    let cs = q_chunk.rows;
    if q_chunk.cols != dim || k_prefix.cols != dim || v_prefix.cols != dim {
        return Err(FocrError::Other(anyhow::anyhow!(
            "chunk_prefill_attention: q/k/v cols ({},{},{}) != num_heads*head_dim {dim}",
            q_chunk.cols,
            k_prefix.cols,
            v_prefix.cols
        )));
    }
    if v_prefix.rows != c1 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "chunk_prefill_attention: k/v prefix rows disagree ({}, {})",
            c1,
            v_prefix.rows
        )));
    }
    if c0 > c1 || c0 + cs != c1 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "chunk_prefill_attention: chunk [{c0}, {}) of {cs} rows does not fit prefix {c1}",
            c0 + cs
        )));
    }
    // Front-pad the chunk's queries to the prefix length so SDPA's top-left
    // causal mask aligns each real query row to its TRUE global position. The
    // scratch rows `[0, c0)` are discarded (their masked attention over the
    // zero query never feeds back into a kept row).
    let mut q_padded = Mat::zeros(c1, dim);
    q_padded.data[c0 * dim..c1 * dim].copy_from_slice(&q_chunk.data);
    let ctx_full = prefill_attention(&q_padded, k_prefix, v_prefix, num_heads, head_dim)?;
    Ok(Mat::from_vec(
        cs,
        dim,
        ctx_full.data[c0 * dim..c1 * dim].to_vec(),
    ))
}

/// Run the full 12-layer prefill over `inputs_embeds` exactly like [`forward`],
/// but ALSO capture each layer's RoPE'd K/V into a per-layer [`RingCache`] — the
/// R-SWA reference block ([SPEC-091], never evicted). Returns the final
/// `model.norm`-ready hidden `[seq, hidden]` (bit-identical to [`forward`]'s
/// output, same kernels) plus the 12 populated caches for
/// [`decode_step_with_cache`].
///
/// This is the prefill half of the O(n) cached decode: the stateless loop
/// re-runs [`forward`] over the whole growing sequence every step (O(n^2) — the
/// MLP/MoE re-processes all prior tokens); here we pay the prefill ONCE and then
/// extend by a single token per step.
///
/// `FOCR_PREFILL_CHUNK` ([`prefill_chunk_size`]) arms the chunked path
/// ([`prefill_with_cache_chunked`]); unset ⇒ the monolithic loop, byte-for-byte
/// the original.
///
/// # Errors
/// As [`forward`], plus a [`RingCache::record_prefill`] error if `seq` exceeds
/// the cache capacity (it is sized to `seq`, so only an internal inconsistency
/// trips it).
pub fn prefill_with_cache(
    wc: &DecoderWeightCache,
    inputs_embeds: &Mat,
) -> FocrResult<(Mat, Vec<RingCache>)> {
    prefill_with_cache_chunked(wc, inputs_embeds, prefill_chunk_size())
}

/// [`prefill_with_cache`] with the chunk decision made explicit (so the parity
/// test can exercise both schedules in one process without re-reading the
/// kill-switch). `chunk = None` runs the monolithic loop — the exact original
/// path; `chunk = Some(C)` tiles the prefill into `C`-token chunks (lossless, see
/// the module note above [`PREFILL_CHUNK_ENV`]).
///
/// # Errors
/// As [`prefill_with_cache`].
pub fn prefill_with_cache_chunked(
    wc: &DecoderWeightCache,
    inputs_embeds: &Mat,
    chunk: Option<usize>,
) -> FocrResult<(Mat, Vec<RingCache>)> {
    checked_mat_len("decoder::prefill_with_cache inputs_embeds", inputs_embeds)?;
    let hidden = config::HIDDEN_SIZE;
    let num_heads = config::NUM_ATTENTION_HEADS;
    let head_dim = config::HEAD_DIM;
    let qkv_dim = checked_shape_mul(
        "decoder::prefill_with_cache",
        num_heads,
        head_dim,
        "num_heads*head_dim",
    )?;
    let eps = config::RMS_NORM_EPS;
    if inputs_embeds.cols != hidden {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::prefill_with_cache: inputs_embeds cols {} != hidden {hidden}",
            inputs_embeds.cols
        )));
    }
    let seq = inputs_embeds.rows;

    let mut caches: Vec<RingCache> = (0..config::NUM_HIDDEN_LAYERS)
        .map(|_| RingCache::new(seq.max(1)))
        .collect();

    if let Some(chunk) = chunk {
        // ── Chunked schedule: push each `chunk`-token slice through all layers,
        //    growing the per-layer K/V, then seed the rings ONCE at the end. ──
        let mut out = Mat::zeros(seq, hidden);
        let mut k_full: Vec<Mat> = (0..config::NUM_HIDDEN_LAYERS)
            .map(|_| Mat::zeros(seq, qkv_dim))
            .collect();
        let mut v_full: Vec<Mat> = (0..config::NUM_HIDDEN_LAYERS)
            .map(|_| Mat::zeros(seq, qkv_dim))
            .collect();
        for (c0, c1) in prefill_chunk_bounds(seq, chunk) {
            // RoPE over THIS chunk's TRUE absolute positions [c0, c1) ([SPEC-095]).
            let positions: Vec<usize> = (c0..c1).collect();
            let rope = RopeTable::build(&positions, head_dim, config::ROPE_THETA);
            let mut x = Mat::from_vec(
                c1 - c0,
                hidden,
                inputs_embeds.data[c0 * hidden..c1 * hidden].to_vec(),
            );
            for layer in 0..config::NUM_HIDDEN_LAYERS {
                let cl = &wc.layers[layer];
                let lw = cached_layer_weights(cl);
                let normed = nn::rms_norm(&x, Some(lw.input_ln), eps)?;
                if calib::enabled() {
                    calib::record_rows(&quant_calib::attn_in_key(layer), &normed.data, normed.cols);
                }
                let (q, k, v) = qkv_with_rope(&normed, &lw, &rope, hidden, qkv_dim)?;
                // Append this chunk's K/V into the running reference block.
                k_full[layer].data[c0 * qkv_dim..c1 * qkv_dim].copy_from_slice(&k.data);
                v_full[layer].data[c0 * qkv_dim..c1 * qkv_dim].copy_from_slice(&v.data);
                let kpre = Mat::from_vec(c1, qkv_dim, k_full[layer].data[..c1 * qkv_dim].to_vec());
                let vpre = Mat::from_vec(c1, qkv_dim, v_full[layer].data[..c1 * qkv_dim].to_vec());
                let context = chunk_prefill_attention(&q, &kpre, &vpre, num_heads, head_dim, c0)?;
                if calib::enabled() {
                    calib::record_rows(&quant_calib::o_in_key(layer), &context.data, context.cols);
                }
                let attn_out = attn_output_proj(&context, lw.o_proj, hidden, qkv_dim)?;
                let h = add_residual(&x, &attn_out)?;
                let normed2 = nn::rms_norm(&h, Some(lw.post_attn_ln), eps)?;
                let mlp_out = cached_mlp(&cl.mlp, &normed2, calib::enabled().then_some(layer))?;
                x = add_residual(&h, &mlp_out)?;
            }
            out.data[c0 * hidden..c1 * hidden].copy_from_slice(&x.data);
        }
        // Seed each ring with the full accumulated reference block — byte-for-byte
        // the monolithic `record_prefill` (same K/V, same order).
        for layer in 0..config::NUM_HIDDEN_LAYERS {
            let (kh, vh) = token_major_to_head_major(
                &k_full[layer],
                &v_full[layer],
                seq,
                num_heads,
                head_dim,
            )?;
            caches[layer].record_prefill(&kh, &vh, seq)?;
        }
        return Ok((out, caches));
    }

    // ── Monolithic schedule (the unset default): one SDPA over the whole seq. ──
    // Absolute positions 0..seq, one shared RoPE table ([SPEC-095]).
    let positions: Vec<usize> = (0..seq).collect();
    let rope = RopeTable::build(&positions, head_dim, config::ROPE_THETA);

    let mut x = inputs_embeds.clone();
    for layer in 0..config::NUM_HIDDEN_LAYERS {
        let cl = &wc.layers[layer];
        let lw = cached_layer_weights(cl);

        // Attention sub-block — mirrors `layer_forward`, but intercepts (k, v) to
        // seed the ring cache's reference block before the prefill SDPA.
        let normed = nn::rms_norm(&x, Some(lw.input_ln), eps)?;
        if calib::enabled() {
            calib::record_rows(&quant_calib::attn_in_key(layer), &normed.data, normed.cols);
        }
        let (q, k, v) = qkv_with_rope(&normed, &lw, &rope, hidden, qkv_dim)?;
        let (kh, vh) = token_major_to_head_major(&k, &v, seq, num_heads, head_dim)?;
        caches[layer].record_prefill(&kh, &vh, seq)?;
        let context = prefill_attention(&q, &k, &v, num_heads, head_dim)?;
        if calib::enabled() {
            calib::record_rows(&quant_calib::o_in_key(layer), &context.data, context.cols);
        }
        let attn_out = attn_output_proj(&context, lw.o_proj, hidden, qkv_dim)?;
        let h = add_residual(&x, &attn_out)?;

        // MLP / MoE sub-block (dense layer 0, MoE 1..11) off the cached weights.
        let normed2 = nn::rms_norm(&h, Some(lw.post_attn_ln), eps)?;
        let mlp_out = cached_mlp(&cl.mlp, &normed2, calib::enabled().then_some(layer))?;
        x = add_residual(&h, &mlp_out)?;
    }
    Ok((x, caches))
}

/// One incremental decode step over the per-layer [`RingCache`]s
/// ([SPEC-091..095]).
///
/// `token_embed` is the new token's embedding row `[1, hidden]` (from
/// [`embed_tokens`]); `position` is its TRUE absolute position (`prefill_len +
/// generated_so_far`), used for RoPE ([SPEC-095] — always the logical position,
/// never a ring slot). Each layer writes this token's RoPE'd K/V into its ring,
/// then runs R-SWA [`rswa::decode_attention`] over (reference ++ ring). Returns
/// the final `model.norm`-ready hidden `[1, hidden]` for `lm_head`.
///
/// For the first [`rswa::RING_WINDOW`] (128) decode steps this is bit-identical
/// to the stateless [`forward`] over the equivalent full sequence (the ring holds
/// every generated token; no eviction, so the union is the full causal context);
/// beyond that the generated tail is windowed to `W` — the reference model's
/// actual sliding-window behavior, which the stateless full-causal path does NOT
/// reproduce.
///
/// # Errors
/// As [`forward`], plus a [`RingCache::write_decode_step`] error if the caches
/// were not populated by [`prefill_with_cache`] first, or a shape mismatch.
pub fn decode_step_with_cache(
    wc: &DecoderWeightCache,
    caches: &mut [RingCache],
    token_embed: &Mat,
    position: usize,
) -> FocrResult<Mat> {
    let hidden = config::HIDDEN_SIZE;
    let qkv_dim = checked_shape_mul(
        "decoder::decode_step_with_cache",
        config::NUM_ATTENTION_HEADS,
        config::HEAD_DIM,
        "num_heads*head_dim",
    )?;
    let eps = config::RMS_NORM_EPS;
    if token_embed.rows != 1 || token_embed.cols != hidden {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::decode_step_with_cache: token_embed shape [{}, {}] != [1, {hidden}]",
            token_embed.rows,
            token_embed.cols
        )));
    }
    if caches.len() != config::NUM_HIDDEN_LAYERS {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::decode_step_with_cache: {} caches != {} layers",
            caches.len(),
            config::NUM_HIDDEN_LAYERS
        )));
    }

    // RoPE at the single TRUE absolute position (one row, one shared table).
    let rope = RopeTable::build(&[position], config::HEAD_DIM, config::ROPE_THETA);

    let profiling = prof::enabled();
    let mut x = token_embed.clone();
    for layer in 0..config::NUM_HIDDEN_LAYERS {
        let cl = &wc.layers[layer];
        let t_attn = profiling.then(Instant::now);

        // Attention via the bespoke m=1 GEMV: project q/k/v, RoPE q/k at the true
        // `position`, push K/V into the ring (the query attends to itself as the
        // newest ring token — OQ-3), R-SWA decode attention over reference ++
        // ring, then o_proj. seq == 1, so the token-major `[1, qkv_dim]` rows ARE
        // the head-major `[num_heads, head_dim]` flats the ring kernels expect.
        let normed = nn::rms_norm(&x, Some(&cl.input_ln), eps)?;
        let nrow = normed.row(0);
        if calib::enabled() {
            calib::record_row(&quant_calib::attn_in_key(layer), nrow);
        }
        let mut q = Mat::from_vec(1, qkv_dim, gemv(nrow, &cl.q_proj, qkv_dim, hidden));
        let mut k = Mat::from_vec(1, qkv_dim, gemv(nrow, &cl.k_proj, qkv_dim, hidden));
        let v = gemv(nrow, &cl.v_proj, qkv_dim, hidden);
        apply_rope(&mut q, &rope)?;
        apply_rope(&mut k, &rope)?;
        caches[layer].write_decode_step(&k.data, &v)?;
        let context = rswa::decode_attention(&caches[layer], &q.data)?;
        if calib::enabled() {
            calib::record_row(&quant_calib::o_in_key(layer), &context.data);
        }
        let attn_out = Mat::from_vec(1, hidden, gemv(&context.data, &cl.o_proj, hidden, qkv_dim));
        let h = add_residual(&x, &attn_out)?;
        if let Some(t) = t_attn {
            prof::add(&prof::ATTN_NS, t.elapsed().as_nanos() as u64);
        }

        // MLP / MoE via the bespoke GEMV (routed top-k experts + shared).
        let normed2 = nn::rms_norm(&h, Some(&cl.post_attn_ln), eps)?;
        let mlp_out = Mat::from_vec(
            1,
            hidden,
            decode_mlp(&cl.mlp, &normed2, calib::enabled().then_some(layer))?,
        );
        x = add_residual(&h, &mlp_out)?;
    }
    Ok(x)
}

// ╔══════════════════════════════════════════════════════════════════════════╗
// ║  INT8 DECODE ENGINE — per-output-channel symmetric S8S8 dense contractions  ║
// ╚══════════════════════════════════════════════════════════════════════════╝
//
// The f32 cache above dequantizes the whole 12-layer MoE decoder to ~10.5 GB of
// f32 and reads ~1.9 GB/token at decode — memory-bound on the ~120 GB/s M4 bus.
// This int8 engine stores the SAME weights as per-output-channel symmetric int8
// (~2.8 GB, 4x less traffic) and runs the hot GEMV/GEMM through the bespoke
// `simd::igemm_s8s8` backend (measured LLVM autovec by default on Apple, with
// forced SDOT/SMMLA routes; AVX-512 / AVX-VNNI / AVX2 on x86).
//
// Precision: weights quantize as `w_scale[o] = max|W[o,:]|/127` (symmetric, the
// proven `nn::quantize_int8` / `ft_kernel_cpu::quantize_per_output_channel_i8`
// convention); activations quantize dynamically per row (per-tensor for the m=1
// decode GEMV), `a_scale = max|x|/127`, then dequant `acc · a_scale · w_scale[o]`
// in exact i32 — identical math to `nn::linear_int8_dynamic`. The router gate and
// every RMSNorm weight stay f32 (tiny, and top-k selection is precision-sensitive).
// Accuracy is VERIFIED end-to-end against the f32-stateless oracle + baidu CER.

/// Block of int8 GEMV output rows fanned to one rayon task (each a single
/// `simd::igemm_s8s8` dense contraction). Matches the f32 `gemv`'s 64-row blocking.
const I8_GEMV_BLOCK: usize = 64;

/// Dynamically quantize a single activation row `x[k]` to symmetric int8
/// (`a_scale = max|x|/127`, zero-point 0), returning `(xq, a_scale)`. The m=1
/// decode case of `ft_kernel_cpu`'s per-row `DynamicQuantizeLinear`.
#[inline]
fn quantize_row_i8(x: &[f32]) -> (Vec<i8>, f32) {
    let amax = x.iter().fold(0.0f32, |m, &v| m.max(v.abs()));
    let a_scale = if amax > 0.0 { amax / 127.0 } else { 1.0 };
    let xq: Vec<i8> = x
        .iter()
        .map(|&v| (v / a_scale).round_ties_even().clamp(-127.0, 127.0) as i8)
        .collect();
    (xq, a_scale)
}

/// Bespoke single-token (m=1) int8 GEMV — the decode-throughput kernel.
///
/// `y[o] = (Σ_i xq[i]·w[o,i]) · a_scale · w_scale[o]` for `o in 0..n`, with the
/// activation `x[k]` dynamically int8-quantized once, and the pre-quantized
/// weight `qw` in `[n,k]` output-channel-major int8 (the native `[out,in]`
/// checkpoint layout — no transpose). The `n` output rows fan across the rayon
/// pool in `I8_GEMV_BLOCK` chunks, each a single dispatched `simd::igemm_s8s8`
/// call; i32 accumulation is exact (proven `K ≤ 6848 < i32::MAX`, `src/simd`).
fn gemv_i8(x: &[f32], qw: &QInt8) -> Vec<f32> {
    debug_assert_eq!(x.len(), qw.k);
    let (xq, a_scale) = quantize_row_i8(x);
    gemv_i8_prequant(&xq, a_scale, qw)
}

/// Block-parallel int8 GEMV from a PRE-QUANTIZED activation row `(xq, a_scale)` —
/// the body of [`gemv_i8`] hoisted out of the per-call activation quantize so a
/// caller that already produced `(xq, a_scale)` (the fused norm->quant epilogue,
/// [`FUSE_NORM_QUANT_ENV`]) can feed the SAME int8 row to several weights without
/// re-quantizing. BYTE-FOR-BYTE identical to [`gemv_i8`]: each of the `n` output
/// channels is an independent i32 dot dequantized by `a_scale · scale[o]`, so the
/// rayon row-chunking never changes a per-channel value (the same N-independence
/// the 64-row blocking and `fuse_qkv` already rely on).
/// Layout-aware channel-block int8 GEMM: contract `xq [m, k]` against weight
/// rows `[base, base + cnt)` of `qw` into `acc [m, cnt]` (zeroed here).
///
/// `RowMajor` slices the canonical `[n, k]` buffer exactly as before;
/// `SmmlaPanels` (an `--arch aarch64-smmla` artifact kept packed by the
/// loader, bd-2mo.3) slices the offline panel stream — pair-aligned because
/// every caller blocks by the even `I8_GEMV_BLOCK` (or passes base 0) — and
/// feeds the i8mm kernel with ZERO runtime shuffle. Bit-identical either way:
/// the packing is a pure zero-padded permutation and integer accumulation is
/// exact, so each output channel's i32 is the same in both layouts.
fn igemm_i8_block(qw: &QInt8, xq: &[i8], m: usize, base: usize, cnt: usize, acc: &mut [i32]) {
    let k = qw.k;
    match qw.layout {
        WeightLayout::RowMajor => {
            acc.fill(0);
            simd::igemm_s8s8(xq, &qw.w[base * k..(base + cnt) * k], m, k, cnt, acc);
        }
        WeightLayout::SmmlaPanels => {
            debug_assert_eq!(base % 2, 0, "channel blocks must be pair-aligned");
            let kb = k.div_ceil(8);
            let start = (base / 2) * kb * 16;
            let len = cnt.div_ceil(2) * kb * 16;
            simd::igemm_s8s8_packed_b(xq, &qw.w[start..start + len], m, k, cnt, acc);
        }
    }
}

fn gemv_i8_prequant(xq: &[i8], a_scale: f32, qw: &QInt8) -> Vec<f32> {
    let k = qw.k;
    let n = qw.n;
    debug_assert_eq!(xq.len(), k);
    debug_assert_eq!(qw.w.len(), qw.expected_w_len());
    debug_assert_eq!(qw.scales.len(), n);
    let mut y = vec![0.0f32; n];
    y.par_chunks_mut(I8_GEMV_BLOCK)
        .enumerate()
        .for_each(|(blk, ys)| {
            let base = blk * I8_GEMV_BLOCK;
            let cnt = ys.len();
            let mut acc = vec![0i32; cnt];
            igemm_i8_block(qw, xq, 1, base, cnt, &mut acc);
            for (j, slot) in ys.iter_mut().enumerate() {
                *slot = acc[j] as f32 * a_scale * qw.scales[base + j];
            }
        });
    y
}

/// `M=B` batched twin of [`gemv_i8_bias_prequant`] for the DENSE decoder's batch
/// spine (A7.5, bd-3jo6.1.7.5): every stream's activation row arrives ALREADY
/// quantized on its own ties-to-even `(xq, a_scale)` (exactly what the m=1 dense
/// decode step produces), the i32 contraction runs as ONE `M=B`
/// [`simd::igemm_s8s8`] per channel block (M-independent — the weight panel is
/// read once and reused across all `B` rows), and the dequant
/// `acc as f32 * a_scale[r] * scales[o]` plus the optional per-output f32 bias
/// use the SAME operands in the SAME order as [`gemv_i8_prequant`] +
/// [`gemv_i8_bias_prequant`] — so row `r` is BYTE-FOR-BYTE the m=1 result for
/// stream `r` alone (the lossless contract; gated in-module and by the dense
/// bit-identity suite).
pub(crate) fn gemm_i8_bias_prequant_batched(
    rows: &[(&[i8], f32)],
    qw: &QInt8,
    bias: Option<&[f32]>,
) -> Vec<Vec<f32>> {
    let b = rows.len();
    let k = qw.k;
    let n = qw.n;
    debug_assert_eq!(qw.w.len(), qw.expected_w_len());
    debug_assert_eq!(qw.scales.len(), n);
    if b == 0 {
        return Vec::new();
    }
    // Pack the prequantized rows `[b, k]` (no re-quantize — the scales are the
    // streams' own ties-to-even scales).
    let mut xq = vec![0i8; b * k];
    for (r, (row, _)) in rows.iter().enumerate() {
        debug_assert_eq!(row.len(), k);
        xq[r * k..(r + 1) * k].copy_from_slice(row);
    }
    // Channel-major `[n, b]` scratch: disjoint contiguous chunks per channel
    // block, race-free fan-out (the gemv_i8_batched pattern).
    let mut ycm = vec![0.0f32; n * b];
    ycm.par_chunks_mut(I8_GEMV_BLOCK * b)
        .enumerate()
        .for_each(|(blk, ys)| {
            let base = blk * I8_GEMV_BLOCK;
            let cnt = ys.len() / b;
            let mut acc = vec![0i32; b * cnt];
            igemm_i8_block(qw, &xq, b, base, cnt, &mut acc);
            for j in 0..cnt {
                let scale_o = qw.scales[base + j];
                let bias_o = bias.map_or(0.0, |bb| bb[base + j]);
                for (r, &(_, a_scale)) in rows.iter().enumerate() {
                    // Same expression shape as gemv_i8_prequant (+ the bias add
                    // gemv_i8_bias_prequant performs after dequant).
                    ys[j * b + r] = acc[r * cnt + j] as f32 * a_scale * scale_o + bias_o;
                }
            }
        });
    // Transpose channel-major `[n, b]` -> per-stream rows `[b][n]`.
    let mut out: Vec<Vec<f32>> = (0..b).map(|_| vec![0.0f32; n]).collect();
    for o in 0..n {
        let col = o * b;
        for (r, row) in out.iter_mut().enumerate() {
            row[o] = ycm[col + r];
        }
    }
    out
}

/// Shared dynamic activation quantizer for the dense-model decode paths. Uses
/// true division plus ties-to-even, matching [`quantize_row_i8`], the prefill's
/// `nn::linear_int8_dynamic`, and the frozen `.focrq` numeric contract.
#[inline]
pub(crate) fn quantize_row_i8_te(x: &[f32]) -> (Vec<i8>, f32) {
    quantize_row_i8(x)
}

/// n-parallel m=1 int8 GEMV from a PRE-QUANTIZED activation `(xq, a_scale)`, adding
/// an optional per-output f32 `bias` after dequant. Reuses [`gemv_i8_prequant`] (the
/// output channels fan across the rayon pool — the decode kernel that keeps m=1 fast,
/// vs `nn::linear_int8_dynamic` which parallelizes over the m=1 row = single-thread).
/// The fused-qkv decode quantizes the normed row ONCE and feeds q/k/v (one `[3·d, h]`
/// panel + concatenated biases) through a single call.
pub(crate) fn gemv_i8_bias_prequant(
    xq: &[i8],
    a_scale: f32,
    qw: &QInt8,
    bias: Option<&[f32]>,
) -> Vec<f32> {
    let mut y = gemv_i8_prequant(xq, a_scale, qw);
    if let Some(b) = bias {
        debug_assert_eq!(b.len(), y.len());
        for (v, &bb) in y.iter_mut().zip(b) {
            *v += bb;
        }
    }
    y
}

/// `FOCR_FUSE_NORM_QUANT` (bd-1azu.54, Lever 1): fold the decode-row int8
/// activation quantize into the `input_layernorm` RMSNorm epilogue so the
/// normalized row is quantized ONCE as it is produced and the same int8 row feeds
/// q/k/v, instead of `nn::rms_norm` -> f32 `Mat` -> a re-quantize inside each
/// [`gemv_i8`]. DEFAULT OFF — unset reproduces the norm-then-gemv path
/// byte-for-byte. Read ONCE into a process-wide bool.
const FUSE_NORM_QUANT_ENV: &str = "FOCR_FUSE_NORM_QUANT";

fn fuse_norm_quant_enabled() -> bool {
    static FLAG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *FLAG.get_or_init(|| std::env::var_os(FUSE_NORM_QUANT_ENV).is_some())
}

/// Fused RMSNorm + per-row int8 quantize (FOCR_FUSE_NORM_QUANT): returns the SAME
/// `(xq, a_scale)` that `quantize_row_i8(nn::rms_norm(x, weight, eps).row(0))`
/// produces. The normalized values come from the SAME `nn::rms_norm` kernel call,
/// then the quantize (amax/scale/round/clamp of [`quantize_row_i8`]) is applied in
/// the epilogue — so the int8 bytes + scale are byte-for-byte the separate
/// norm-then-quantize, while the int8 GEMV ([`gemv_i8_prequant`]) no longer
/// re-reads an f32 row to re-quantize it. `x` is a single decode row (`rows == 1`).
///
/// # Errors
/// Propagates [`nn::rms_norm`]'s shape error.
fn rms_norm_quant_i8(x: &Mat, weight: Option<&[f32]>, eps: f32) -> FocrResult<(Vec<i8>, f32)> {
    let normed = nn::rms_norm(x, weight, eps)?;
    Ok(quantize_row_i8(normed.row(0)))
}

/// One SwiGLU expert over a single decode row `x[hidden]`, quantized:
/// `down(silu(gate·x) * (up·x))` over [`QLinear`] weights (int8 arm identical
/// to the historical int8-only path; int4 arm consumes packed nibbles).
/// Internally parallel (used for the big dense layer-0 MLP + the shared
/// expert).
fn expert_gemv_i8(
    x: &[f32],
    gate: &QLinear,
    up: &QLinear,
    down: &QLinear,
    calib: Option<(usize, quant_calib::MlpUnit)>,
) -> Vec<f32> {
    if let Some((layer, unit)) = calib.filter(|_| calib::enabled()) {
        calib::record_row(&quant_calib::ffn_in_key(layer, unit), x);
    }
    let g = gemv_q(x, gate);
    let u = gemv_q(x, up);
    let inter = gate.n();
    let mut act = vec![0.0f32; inter];
    for i in 0..inter {
        act[i] = silu(g[i]) * u[i];
    }
    if let Some((layer, unit)) = calib.filter(|_| calib::enabled()) {
        calib::record_row(&quant_calib::down_in_key(layer, unit), &act);
    }
    gemv_q(&act, down)
}

/// SERIAL int8 GEMV from a pre-quantized activation — a single `simd::igemm_s8s8`
/// dense-int8 call over all `n` rows, NO rayon. Used inside the per-expert rayon
/// tasks below, where parallelism is ACROSS experts (one task each) rather than
/// within: 6 routed experts is the right granularity to saturate the pool with
/// ONE dispatch, vs the 18 tiny internally-parallel GEMVs the per-row `gemv_i8`
/// would spawn (measured: MoE experts were 56% of decode and dispatch-bound —
/// dispatch-bound, not compute-bound).
fn gemv_i8_serial(xq: &[i8], a_scale: f32, qw: &QInt8) -> Vec<f32> {
    let k = qw.k;
    let n = qw.n;
    debug_assert_eq!(xq.len(), k);
    let mut acc = vec![0i32; n];
    igemm_i8_block(qw, xq, 1, 0, n, &mut acc);
    let mut y = vec![0.0f32; n];
    for (o, slot) in y.iter_mut().enumerate() {
        *slot = acc[o] as f32 * a_scale * qw.scales[o];
    }
    y
}

/// One SwiGLU expert, fully SERIAL (the input is quantized ONCE and reused for
/// gate+up). The serial twin of [`expert_gemv_i8`] for cross-expert
/// parallelism; dispatches per-tensor between the int8 and packed-int4 arms.
fn expert_gemv_i8_serial(
    x: &[f32],
    gate: &QLinear,
    up: &QLinear,
    down: &QLinear,
    calib: Option<(usize, quant_calib::MlpUnit)>,
) -> Vec<f32> {
    if let Some((layer, unit)) = calib.filter(|_| calib::enabled()) {
        calib::record_row(&quant_calib::ffn_in_key(layer, unit), x);
    }
    let (xq, a_scale) = quantize_row_i8(x);
    let g = gemv_q_serial(&xq, a_scale, gate);
    let u = gemv_q_serial(&xq, a_scale, up);
    let inter = gate.n();
    let mut act = vec![0.0f32; inter];
    for i in 0..inter {
        act[i] = silu(g[i]) * u[i];
    }
    if let Some((layer, unit)) = calib.filter(|_| calib::enabled()) {
        calib::record_row(&quant_calib::down_in_key(layer, unit), &act);
    }
    let (aq, a_scale2) = quantize_row_i8(&act);
    gemv_q_serial(&aq, a_scale2, down)
}

/// Quantize an exact `[out, in_]` row-major f32 weight to per-output-channel
/// symmetric int8.
fn quant_oc(w: &[f32], out: usize, in_: usize, name: &str) -> FocrResult<QInt8> {
    let expected = out.checked_mul(in_).ok_or_else(|| {
        FocrError::FormatMismatch(format!(
            "tensor {name:?}: output/input shape [{out}, {in_}] overflows usize"
        ))
    })?;
    if w.len() != expected {
        return Err(FocrError::FormatMismatch(format!(
            "tensor {name:?}: {} elements != expected output/input shape [{out}, {in_}] ({expected})",
            w.len()
        )));
    }
    Ok(nn::quantize_int8(w, out, in_))
}

/// Resolve `name` to a per-output-channel int8 weight, from EITHER source:
///
/// * a raw bf16/f32 safetensors record → widen the `[out, in]` mat and
///   [`quant_oc`] it at load time (the original path), OR
/// * a pre-quantized `.focrq` record (`QInt8PerChan`, produced by `focr convert`)
///   → read it back verbatim via [`Weights::qint8`], skipping the re-quantize.
///
/// For recipe-approved FFN/expert tensors the two are byte-identical: `focr
/// convert` uses the same [`nn::quantize_int8`] as this fallback. Attention and
/// `lm_head` remain high precision in the default artifact and reach this helper
/// only when the separately gated experimental all-int8 cache is requested.
///
/// # Errors
/// [`FocrError::FormatMismatch`] if `name` is absent or mis-shaped.
pub(crate) fn quant_oc_loaded(weights: &Weights, name: &str, out: usize) -> FocrResult<QInt8> {
    if matches!(
        weights.record(name).map(|rec| rec.dtype),
        Some(DType::QInt8PerChan)
    ) {
        let q = weights.qint8(name)?;
        if q.n != out {
            return Err(FocrError::FormatMismatch(format!(
                "QInt8 tensor {name:?} has {} output rows; expected {out}",
                q.n
            )));
        }
        return Ok(q);
    }

    let record = weights.record(name).ok_or_else(|| {
        FocrError::FormatMismatch(format!("tensor {name:?} not found in weights directory"))
    })?;
    let [rows, in_] = record.shape.as_slice() else {
        return Err(FocrError::FormatMismatch(format!(
            "tensor {name:?} has rank {}; expected 2 ([out, in])",
            record.shape.len()
        )));
    };
    if *rows != out {
        return Err(FocrError::FormatMismatch(format!(
            "tensor {name:?} has {rows} output rows; expected {out}"
        )));
    }
    let mat = weights.mat(name)?;
    quant_oc(&mat.data, out, *in_, name)
}

// ── Packed int4 expert consumption (bd-4l71) ─────────────────────────────────
//
// The wasm experts-int4 recipe (`unlimited-ocr-wasm-experts-int4-attn-int8-
// lmhead-int8-v1`) stores the decoder FFN/expert bulk as `QInt4PerGroup`. The
// paths below consume the PACKED nibbles directly through
// [`simd::int4::igemm_s4s8_packed`] — B is **never** materialized as a dense i8
// buffer (halving resident expert bytes is the entire point of the recipe).
// The activation side reuses the int8 path's [`quantize_row_i8`] (true
// division + ties-to-even), so the int8 and int4 arms share one activation
// contract. Conservative artifacts (qint8 or bf16 FFN tensors) take the
// existing [`QInt8`] arm byte-for-byte unchanged.

/// A quantized decoder FFN/expert GEMM weight: per-output-channel int8 (the
/// conservative recipe / quantize-at-load path) or per-group PACKED int4 (the
/// wasm experts-int4 artifact). The MLP caches hold this so the same
/// prefill/decode drivers consume both storages; which arm a tensor takes is
/// decided once at cache build by [`quant_ffn_loaded`].
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum QLinear {
    /// Per-output-channel symmetric int8 (`[n, k]` + `n` scales).
    I8(QInt8),
    /// Group-quantized packed int4 (`[n, k/2]` nibbles + `n*(k/group)` scales),
    /// consumed packed by [`simd::int4::igemm_s4s8_packed`].
    I4(QInt4),
}

impl QLinear {
    /// Output dimension (rows).
    pub(crate) fn n(&self) -> usize {
        match self {
            QLinear::I8(q) => q.n,
            QLinear::I4(q) => q.n,
        }
    }

    /// Contraction dimension (columns). Exercised by the loader-plumbing tests
    /// (the forward paths read `k` through the concrete arms).
    #[cfg_attr(not(test), allow(dead_code))]
    pub(crate) fn k(&self) -> usize {
        match self {
            QLinear::I8(q) => q.k,
            QLinear::I4(q) => q.k,
        }
    }
}

/// Resolve a decoder FFN/expert tensor `name` to a [`QLinear`]:
///
/// * a `QInt4PerGroup` record (the wasm experts-int4 artifact) → the packed
///   [`QLinear::I4`] arm, verbatim (nibbles stay packed in residency), or
/// * anything else (pre-quantized `QInt8PerChan`, or raw bf16/f32 quantized at
///   load) → the existing [`quant_oc_loaded`] int8 arm, byte-for-byte the
///   conservative behavior.
///
/// # Errors
/// [`FocrError::FormatMismatch`] if `name` is absent, mis-shaped, or its
/// output-row count disagrees with `out`.
pub(crate) fn quant_ffn_loaded(weights: &Weights, name: &str, out: usize) -> FocrResult<QLinear> {
    if matches!(
        weights.record(name).map(|rec| rec.dtype),
        Some(DType::QInt4PerGroup)
    ) {
        let q = weights.qint4(name)?;
        if q.n != out {
            return Err(FocrError::FormatMismatch(format!(
                "QInt4 tensor {name:?} has {} output rows; expected {out}",
                q.n
            )));
        }
        return Ok(QLinear::I4(q));
    }
    Ok(QLinear::I8(quant_oc_loaded(weights, name, out)?))
}

/// Block-parallel m=1 packed-int4 GEMV — the int4 twin of [`gemv_i8`]. The
/// activation row is dynamically int8-quantized ONCE (the same
/// [`quantize_row_i8`] the int8 path uses); the `n` output rows fan across the
/// rayon pool in [`I8_GEMV_BLOCK`] chunks, each one
/// [`simd::int4::igemm_s4s8_packed`] call over that block's CONTIGUOUS packed
/// rows/scales, then the activation scale is applied
/// (`y[o] = a_scale * Σ_g scale[o,g]·(Σ_{k∈g} xq[k]·q4[o,k])`). Each output
/// channel is an independent per-group contraction, so the block partition
/// never changes a value — [`gemv_i4_serial`] is bit-identical (gated in-module).
fn gemv_i4(x: &[f32], q: &QInt4) -> Vec<f32> {
    debug_assert_eq!(x.len(), q.k);
    let (xq, a_scale) = quantize_row_i8(x);
    let kbytes = q.k / 2;
    let groups = q.k / q.group_size;
    let mut y = vec![0.0f32; q.n];
    y.par_chunks_mut(I8_GEMV_BLOCK)
        .enumerate()
        .for_each(|(blk, ys)| {
            let base = blk * I8_GEMV_BLOCK;
            let cnt = ys.len();
            // Scales may live as raw LE bytes in the blob (zero-copy residency,
            // bd-4l71): decode just this block's groups into a task-local
            // scratch — identical values to the owned-f32 storage.
            let mut scale_scratch = Vec::new();
            let scales =
                q.scales
                    .slice_f32(base * groups, (base + cnt) * groups, &mut scale_scratch);
            simd::int4::igemm_s4s8_packed(
                &xq,
                &q.packed[base * kbytes..(base + cnt) * kbytes],
                scales,
                q.group_size,
                1,
                q.k,
                cnt,
                ys,
            );
            for slot in ys.iter_mut() {
                *slot *= a_scale;
            }
        });
    y
}

/// SERIAL packed-int4 GEMV from a pre-quantized activation — the int4 twin of
/// [`gemv_i8_serial`]: one [`simd::int4::igemm_s4s8_packed`] call over all `n`
/// rows, NO rayon (used inside per-expert rayon tasks where parallelism is
/// ACROSS experts). Bit-identical to [`gemv_i4`] for the same `(xq, a_scale)`.
fn gemv_i4_serial(xq: &[i8], a_scale: f32, q: &QInt4) -> Vec<f32> {
    debug_assert_eq!(xq.len(), q.k);
    let mut y = vec![0.0f32; q.n];
    let mut scale_scratch = Vec::new();
    let scales = q.scales.slice_f32(0, q.scales.len(), &mut scale_scratch);
    simd::int4::igemm_s4s8_packed(xq, &q.packed, scales, q.group_size, 1, q.k, q.n, &mut y);
    for slot in y.iter_mut() {
        *slot *= a_scale;
    }
    y
}

/// Storage-dispatching m=1 GEMV over a [`QLinear`] (fresh activation quantize):
/// the int8 arm is exactly [`gemv_i8`], the int4 arm consumes packed nibbles
/// via [`gemv_i4`].
fn gemv_q(x: &[f32], w: &QLinear) -> Vec<f32> {
    match w {
        QLinear::I8(q) => gemv_i8(x, q),
        QLinear::I4(q) => gemv_i4(x, q),
    }
}

/// Storage-dispatching SERIAL m=1 GEMV from a pre-quantized activation row —
/// the [`QLinear`] twin of [`gemv_i8_serial`] / [`gemv_i4_serial`].
fn gemv_q_serial(xq: &[i8], a_scale: f32, w: &QLinear) -> Vec<f32> {
    match w {
        QLinear::I8(q) => gemv_i8_serial(xq, a_scale, q),
        QLinear::I4(q) => gemv_i4_serial(xq, a_scale, q),
    }
}

/// Prefill (m≥1) dynamic-activation packed-int4 linear — the int4 twin of
/// [`nn::linear_int8_dynamic`]. Every activation row is int8-quantized with the
/// SAME [`quantize_row_i8`] rule the int8 prefill kernel applies (true division,
/// ties-to-even), the whole `[m, k] · [n, k]ᵀ` contraction runs as ONE
/// [`simd::int4::igemm_s4s8_packed`] call over the packed nibbles (B never
/// materialized), and each output row is dequantized by its own activation
/// scale. Row `r` of the result is therefore bit-identical to the m=1
/// [`gemv_i4_serial`] over that row alone.
///
/// # Errors
/// [`FocrError::Other`] on a shape mismatch.
fn linear_int4_dynamic(x: &Mat, q: &QInt4) -> FocrResult<Mat> {
    checked_mat_len("linear_int4_dynamic x", x)?;
    if x.cols != q.k {
        return Err(FocrError::Other(anyhow::anyhow!(
            "linear_int4_dynamic: x.cols {} != w.k {}",
            x.cols,
            q.k
        )));
    }
    let (m, k, n) = (x.rows, x.cols, q.n);
    let mut xq = vec![0i8; m * k];
    let mut a_scales = vec![0.0f32; m];
    for r in 0..m {
        let (row_q, scale) = quantize_row_i8(x.row(r));
        xq[r * k..(r + 1) * k].copy_from_slice(&row_q);
        a_scales[r] = scale;
    }
    let mut out = vec![0.0f32; m * n];
    let mut scale_scratch = Vec::new();
    let scales = q.scales.slice_f32(0, q.scales.len(), &mut scale_scratch);
    simd::int4::igemm_s4s8_packed(&xq, &q.packed, scales, q.group_size, m, k, n, &mut out);
    for (r, &scale) in a_scales.iter().enumerate() {
        for value in &mut out[r * n..(r + 1) * n] {
            *value *= scale;
        }
    }
    Ok(Mat::from_vec(m, n, out))
}

/// Storage-dispatching prefill linear over a [`QLinear`]: the int8 arm is
/// exactly [`nn::linear_int8_dynamic`] (no bias — the decoder FFN projections
/// carry none), the int4 arm is [`linear_int4_dynamic`].
///
/// # Errors
/// Shape mismatches from either arm.
fn linear_q_dynamic(x: &Mat, w: &QLinear) -> FocrResult<Mat> {
    match w {
        QLinear::I8(q) => nn::linear_int8_dynamic(x, q, None),
        QLinear::I4(q) => linear_int4_dynamic(x, q),
    }
}

// ── CCD-sharded / L3-tiled lm_head (FOCR_LMHEAD_SHARD, bd-1azu.25) ────────────
//
// The `lm_head` projects the final hidden `[1, 1280]` against the `[129280, 1280]`
// vocab weight to `[1, 129280]` logits. Each logit `o` is a SELF-CONTAINED dot
// `Σ_i x[i]·w[o,i]` (int8: `(Σ_i xq[i]·w[o,i])·a_scale·scale[o]`) — the vocab
// columns never reduce into one another. So splitting the 129280 output columns
// into CONTIGUOUS tiles, computing each tile, and writing it back into its own
// `[.., tile]` column span is BYTE-FOR-BYTE identical to the single monolithic
// GEMV: same single activation quantize, same per-logit i32 contraction
// (N-independent — already relied on by [`gemv_i8`]'s 64-row blocking and
// [`fuse_qkv`]), same per-channel dequant operands in the same order, and the
// ascending tile order preserves the vocab column order — so argmax/sampling are
// unchanged. This is the L3-resident weight-tiling seam (read one vocab tile of
// the 660 MB int8 panel at a time); DEFAULT OFF keeps the exact monolithic path.

/// `FOCR_LMHEAD_SHARD`: kill-switch arming the vocab-tiled (`lm_head`-sharded)
/// projection ([`gemv_i8_sharded`] / [`gemv_sharded`]). DEFAULT OFF — the head
/// stays the single monolithic [`gemv_i8`]/[`gemv`]; armed, the 129280 vocab
/// columns are computed in [`lmhead_shard_tiles`] CONTIGUOUS tiles, byte-for-byte
/// identical. Consulted ONCE per `lm_head` call, never inside the math, exactly
/// like [`QKV_FUSED_ENV`].
const LMHEAD_SHARD_ENV: &str = "FOCR_LMHEAD_SHARD";

/// `FOCR_LMHEAD_SHARD_TILES`: number of CONTIGUOUS vocab tiles the sharded
/// `lm_head` splits its output columns into. Parsed ONCE; defaults to
/// [`DEFAULT_LMHEAD_SHARD_TILES`] when unset, empty, unparseable, or `0`. The tile
/// count never changes a logit value or the column order — only how the columns
/// are grouped into kernel calls.
const LMHEAD_SHARD_TILES_ENV: &str = "FOCR_LMHEAD_SHARD_TILES";

/// Fallback vocab-tile count when [`LMHEAD_SHARD_TILES_ENV`] is unset/invalid.
const DEFAULT_LMHEAD_SHARD_TILES: usize = 16;

/// Whether the vocab-tiled `lm_head` is armed (the [`LMHEAD_SHARD_ENV`]
/// kill-switch, read ONCE into a process-wide bool — never touched inside the
/// per-logit math). The sharded kernels are pure and testable regardless of this
/// flag; only the head's decision to route through them is gated here.
#[must_use]
pub fn lmhead_shard_enabled() -> bool {
    static FLAG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *FLAG.get_or_init(|| std::env::var_os(LMHEAD_SHARD_ENV).is_some())
}

/// The configured contiguous vocab-tile count ([`LMHEAD_SHARD_TILES_ENV`], read
/// ONCE). Always `>= 1`.
#[must_use]
pub fn lmhead_shard_tiles() -> usize {
    static TILES: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
    *TILES.get_or_init(|| {
        std::env::var(LMHEAD_SHARD_TILES_ENV)
            .ok()
            .and_then(|v| v.trim().parse::<usize>().ok())
            .filter(|&n| n > 0)
            .unwrap_or(DEFAULT_LMHEAD_SHARD_TILES)
    })
}

/// Partition `n` output channels (vocab columns) into `tiles` CONTIGUOUS,
/// gap-free, non-overlapping ranges covering `[0, n)` in fixed ascending order
/// (the `n % tiles` remainder is spread one-per-tile across the leading tiles).
/// `tiles` is clamped to `[1, n.max(1)]` so every emitted range is in-bounds.
/// The ascending coverage is what preserves the vocab column order — and hence
/// argmax/sampling — under sharding.
fn vocab_tile_ranges(n: usize, tiles: usize) -> Vec<(usize, usize)> {
    let tiles = tiles.clamp(1, n.max(1));
    let base = n / tiles;
    let rem = n % tiles;
    let mut ranges = Vec::with_capacity(tiles);
    let mut start = 0usize;
    for t in 0..tiles {
        let len = base + usize::from(t < rem);
        let end = start + len;
        ranges.push((start, end));
        start = end;
    }
    debug_assert_eq!(start, n, "vocab_tile_ranges must cover [0, n)");
    ranges
}

/// Vocab-tiled f32 `lm_head` GEMV — the `FOCR_LMHEAD_SHARD` twin of [`gemv`].
/// Splits the `n` output channels into [`vocab_tile_ranges`] CONTIGUOUS tiles and
/// computes each tile's logit slice into its own `y[start..end]` span. BYTE-FOR-
/// BYTE identical to [`gemv`]: every channel `o`'s value is `dot_f32(x, w[o,:])`
/// regardless of which tile (or 64-row rayon chunk) it lands in — tiling only
/// repartitions the column ranges, never a per-logit reduction or the order.
fn gemv_sharded(x: &[f32], w: &[f32], n: usize, k: usize, tiles: usize) -> Vec<f32> {
    debug_assert_eq!(x.len(), k);
    debug_assert_eq!(w.len(), n * k);
    let mut y = vec![0.0f32; n];
    for (start, end) in vocab_tile_ranges(n, tiles) {
        // Same 64-row rayon chunking as `gemv`, confined to this tile's columns.
        y[start..end]
            .par_chunks_mut(64)
            .enumerate()
            .for_each(|(blk, ys)| {
                let base = start + blk * 64;
                for (j, slot) in ys.iter_mut().enumerate() {
                    let o = base + j;
                    *slot = dot_f32(x, &w[o * k..o * k + k]);
                }
            });
    }
    y
}

/// Vocab-tiled int8 `lm_head` GEMV — the `FOCR_LMHEAD_SHARD` twin of [`gemv_i8`].
/// The activation is dynamically int8-quantized ONCE (the same `(xq, a_scale)`
/// [`gemv_i8`] would produce), then the `n` output channels are computed in
/// [`vocab_tile_ranges`] CONTIGUOUS tiles, each tile's slice written into its own
/// `y[start..end]` span. BYTE-FOR-BYTE identical to [`gemv_i8`]: each channel
/// `o`'s i32 dot `Σ_i xq[i]·w[o,i]` is independent of how the columns are grouped
/// into [`simd::igemm_s8s8`] calls (the N-independence [`gemv_i8`] already exploits
/// with its 64-row blocking), and the dequant `acc·a_scale·scales[o]` uses the
/// SAME operands in the SAME left-associative order.
fn gemv_i8_sharded(x: &[f32], qw: &QInt8, tiles: usize) -> Vec<f32> {
    let k = qw.k;
    let n = qw.n;
    debug_assert_eq!(x.len(), k);
    debug_assert_eq!(qw.w.len(), qw.expected_w_len());
    debug_assert_eq!(qw.scales.len(), n);
    let (xq, a_scale) = quantize_row_i8(x);
    let mut y = vec![0.0f32; n];
    for (start, end) in vocab_tile_ranges(n, tiles) {
        // Same `I8_GEMV_BLOCK` rayon chunking as `gemv_i8`, confined to this
        // tile's columns; `base`/`scales` index the ABSOLUTE channel.
        y[start..end]
            .par_chunks_mut(I8_GEMV_BLOCK)
            .enumerate()
            .for_each(|(blk, ys)| {
                let base = start + blk * I8_GEMV_BLOCK;
                let cnt = ys.len();
                let mut acc = vec![0i32; cnt];
                igemm_i8_block(qw, &xq, 1, base, cnt, &mut acc);
                for (j, slot) in ys.iter_mut().enumerate() {
                    *slot = acc[j] as f32 * a_scale * qw.scales[base + j];
                }
            });
    }
    y
}

/// `FOCR_QKV_FUSED`: at cache-build, STACK q/k/v into ONE `[3*qkv_dim, hidden]`
/// int8 weight so each decode token runs a SINGLE block-parallel [`gemv_i8`]
/// (one activation quantize, one rayon dispatch wave) instead of three. The
/// default path (flag unset) is byte-for-byte unchanged.
const QKV_FUSED_ENV: &str = "FOCR_QKV_FUSED";

/// Read [`QKV_FUSED_ENV`] ONCE into a process-wide bool (build-time only; never
/// touched per-token). DEFAULT **ON** since 2026-07-07 (bd-241s): the fused
/// path is the ledgered LOSSLESS win (NEGATIVE_EVIDENCE 2026-06-27 KEPT entry:
/// M4 −8.9% / x86-avx2 −21% decode; byte-identical proven by the unit gate +
/// page_0590 sha + 20-page CER; re-confirmed 2026-07-07: 0.072→0.052 s/tok
/// best-of-3 on page_0009 @8T, outputs byte-identical). `FOCR_QKV_FUSED=0`
/// restores the three-call path — now the kill-switch, no longer the default.
fn qkv_fused_enabled() -> bool {
    static FLAG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *FLAG.get_or_init(|| {
        !matches!(
            std::env::var(QKV_FUSED_ENV)
                .ok()
                .map(|v| v.trim().to_ascii_lowercase())
                .as_deref(),
            Some("0" | "off" | "false" | "no")
        )
    })
}

/// Stack three same-shaped `[n, k]` per-output-channel int8 weights into ONE
/// `[3n, k]` weight: int8 rows concatenated (q then k then v), per-channel
/// `scales` concatenated in the same order. Feeding the result to one
/// block-parallel [`gemv_i8`] yields output rows `[0,n)`, `[n,2n)`, `[2n,3n)`
/// that are BYTE-IDENTICAL to three separate `gemv_i8(x, {q,k,v})` calls: the
/// activation is quantized once to the same `(xq, a_scale)`, and each output row
/// `o` is an independent i32 dot of the SAME `xq` against the SAME `w[o,:]`
/// dequantized by the SAME `scales[o]` — the rayon row-chunking only changes
/// which rows share a task, never a per-row reduction.
fn fuse_qkv(q: &QInt8, k: &QInt8, v: &QInt8) -> QInt8 {
    debug_assert_eq!(q.k, k.k, "fuse_qkv: k contraction dim mismatch");
    debug_assert_eq!(q.k, v.k, "fuse_qkv: k contraction dim mismatch");
    debug_assert_eq!(q.n, k.n, "fuse_qkv: q/k output dim mismatch");
    debug_assert_eq!(q.n, v.n, "fuse_qkv: q/v output dim mismatch");
    let (n, kk) = (q.n, q.k);
    let mut w = Vec::with_capacity(q.w.len() + k.w.len() + v.w.len());
    w.extend_from_slice(&q.w);
    w.extend_from_slice(&k.w);
    w.extend_from_slice(&v.w);
    let mut scales = Vec::with_capacity(n * 3);
    scales.extend_from_slice(&q.scales);
    scales.extend_from_slice(&k.scales);
    scales.extend_from_slice(&v.scales);
    // Offline SMMLA panels concatenate exactly like rows: each segment's row
    // count is even (every registered attention dim), so segment boundaries
    // land on pair boundaries and the concatenated stream IS the panel pack
    // of the concatenated matrix.
    if q.layout == WeightLayout::SmmlaPanels
        && k.layout == WeightLayout::SmmlaPanels
        && v.layout == WeightLayout::SmmlaPanels
        && q.n.is_multiple_of(2)
        && k.n.is_multiple_of(2)
    {
        return QInt8::new_smmla_panels(w, scales, 3 * n, kk);
    }
    debug_assert!(
        q.layout == WeightLayout::RowMajor
            && k.layout == WeightLayout::RowMajor
            && v.layout == WeightLayout::RowMajor,
        "fuse_qkv: mixed or odd-row packed layouts are unreachable from the loader"
    );
    QInt8::new(w, scales, 3 * n, kk)
}

/// A layer's quantized MLP weights — dense (layer 0) or MoE (layers 1..11).
/// Every projection is a [`QLinear`]: per-output-channel int8 for conservative
/// artifacts (byte-for-byte the historical behavior) or per-group PACKED int4
/// for the wasm experts-int4 artifact (bd-4l71). The MoE router `gate` stays
/// f32 (top-k routing is precision-sensitive and the gate is tiny: `[64,1280]`).
enum CachedMlpI8 {
    Dense {
        gate: QLinear,
        up: QLinear,
        down: QLinear,
    },
    Moe {
        gate: Vec<f32>,
        experts: Vec<[QLinear; 3]>,
        shared: [QLinear; 3],
    },
}

/// One decoder layer's int8 weights (attention projections + MLP). Norm weights
/// stay f32.
struct CachedLayerI8 {
    input_ln: Vec<f32>,
    post_attn_ln: Vec<f32>,
    q_proj: QInt8,
    k_proj: QInt8,
    v_proj: QInt8,
    /// Optional fused `[3*qkv_dim, hidden]` stack of `q_proj`/`k_proj`/`v_proj`,
    /// built ONLY when [`QKV_FUSED_ENV`] is set (the separate three are always
    /// retained for prefill). When present, decode runs one block-parallel GEMV
    /// over all `3*qkv_dim` output rows instead of three; byte-identical output.
    qkv: Option<QInt8>,
    o_proj: QInt8,
    mlp: CachedMlpI8,
}

/// The int8 twin of [`DecoderWeightCache`]: every load-bearing GEMM weight stored
/// per-output-channel symmetric int8 (~2.8 GB vs ~10.5 GB f32). Built ONCE; prefill
/// (via [`nn::linear_int8_dynamic`]) and decode (via [`gemv_i8`]) run off it.
pub struct DecoderWeightCacheI8 {
    layers: Vec<CachedLayerI8>,
    final_norm: Vec<f32>,
    lm_head: QInt8,
}

impl DecoderWeightCacheI8 {
    /// Quantize every decoder tensor ONCE from [`Weights`] to per-output-channel
    /// symmetric int8 (the dense layer-0 MLP, the 64 routed + 1 shared expert per
    /// MoE layer, all four attention projections, and the `lm_head`). The router
    /// gate and the RMSNorm weights stay f32.
    ///
    /// # Errors
    /// [`FocrError::FormatMismatch`] if any expected tensor is absent/mis-shaped.
    pub fn build(weights: &Weights) -> FocrResult<Self> {
        let hidden = config::HIDDEN_SIZE;
        let qkv_dim = config::NUM_ATTENTION_HEADS * config::HEAD_DIM;
        let mut layers = Vec::with_capacity(config::NUM_HIDDEN_LAYERS);
        for layer in 0..config::NUM_HIDDEN_LAYERS {
            let prefix = format!("model.layers.{layer}");
            let input_ln = weights.vec(&format!("{prefix}.input_layernorm.weight"))?;
            let post_attn_ln = weights.vec(&format!("{prefix}.post_attention_layernorm.weight"))?;
            let q_proj = quant_oc_loaded(
                weights,
                &format!("{prefix}.self_attn.q_proj.weight"),
                qkv_dim,
            )?;
            let k_proj = quant_oc_loaded(
                weights,
                &format!("{prefix}.self_attn.k_proj.weight"),
                qkv_dim,
            )?;
            let v_proj = quant_oc_loaded(
                weights,
                &format!("{prefix}.self_attn.v_proj.weight"),
                qkv_dim,
            )?;
            let o_proj = quant_oc_loaded(
                weights,
                &format!("{prefix}.self_attn.o_proj.weight"),
                hidden,
            )?;
            // Fused q/k/v stack (FOCR_QKV_FUSED): built ONCE here so decode runs a
            // single block-parallel GEMV. Byte-identical to the 3 separate calls.
            let qkv = qkv_fused_enabled().then(|| fuse_qkv(&q_proj, &k_proj, &v_proj));
            // FFN/expert projections dispatch per-tensor: qint8/bf16 exactly as
            // before, PACKED int4 kept packed for the wasm artifact (bd-4l71).
            let mlp = if layer < config::FIRST_K_DENSE_REPLACE {
                let p = format!("{prefix}.mlp");
                let inter = moe::config::DENSE_INTERMEDIATE_SIZE;
                CachedMlpI8::Dense {
                    gate: quant_ffn_loaded(weights, &format!("{p}.gate_proj.weight"), inter)?,
                    up: quant_ffn_loaded(weights, &format!("{p}.up_proj.weight"), inter)?,
                    down: quant_ffn_loaded(weights, &format!("{p}.down_proj.weight"), hidden)?,
                }
            } else {
                let p = format!("{prefix}.mlp");
                let gate = weights.mat(&format!("{p}.gate.weight"))?.data;
                let inter = moe::config::MOE_INTERMEDIATE_SIZE;
                let mut experts = Vec::with_capacity(moe::config::N_ROUTED_EXPERTS);
                for e in 0..moe::config::N_ROUTED_EXPERTS {
                    experts.push([
                        quant_ffn_loaded(
                            weights,
                            &format!("{p}.experts.{e}.gate_proj.weight"),
                            inter,
                        )?,
                        quant_ffn_loaded(
                            weights,
                            &format!("{p}.experts.{e}.up_proj.weight"),
                            inter,
                        )?,
                        quant_ffn_loaded(
                            weights,
                            &format!("{p}.experts.{e}.down_proj.weight"),
                            hidden,
                        )?,
                    ]);
                }
                let si = moe::config::SHARED_INTERMEDIATE_SIZE;
                let shared = [
                    quant_ffn_loaded(weights, &format!("{p}.shared_experts.gate_proj.weight"), si)?,
                    quant_ffn_loaded(weights, &format!("{p}.shared_experts.up_proj.weight"), si)?,
                    quant_ffn_loaded(
                        weights,
                        &format!("{p}.shared_experts.down_proj.weight"),
                        hidden,
                    )?,
                ];
                CachedMlpI8::Moe {
                    gate,
                    experts,
                    shared,
                }
            };
            layers.push(CachedLayerI8 {
                input_ln,
                post_attn_ln,
                q_proj,
                k_proj,
                v_proj,
                qkv,
                o_proj,
                mlp,
            });
        }
        let final_norm = weights.vec("model.norm.weight")?;
        let lm_head = quant_oc_loaded(weights, "lm_head.weight", config::VOCAB_SIZE)?;
        Ok(Self {
            layers,
            final_norm,
            lm_head,
        })
    }
}

/// One SwiGLU expert over a `[n_tok, hidden]` activation, int8 — `down(silu(gate·x)
/// * (up·x))` via [`nn::linear_int8_dynamic`]. The int8 twin of [`moe::expert_mlp`].
/// Kept int8-typed for the zoo decoders (`decoder_qwen2`); the Unlimited MLP
/// caches route through the storage-dispatching [`expert_mlp_q`].
pub(crate) fn expert_mlp_i8(x: &Mat, gate: &QInt8, up: &QInt8, down: &QInt8) -> FocrResult<Mat> {
    let mut g = nn::linear_int8_dynamic(x, gate, None)?;
    nn::silu(&mut g);
    let u = nn::linear_int8_dynamic(x, up, None)?;
    if g.data.len() != u.data.len() {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::expert_mlp_i8: gate/up shape mismatch ({} vs {})",
            g.data.len(),
            u.data.len()
        )));
    }
    for (a, &b) in g.data.iter_mut().zip(u.data.iter()) {
        *a *= b;
    }
    nn::linear_int8_dynamic(&g, down, None)
}

/// One SwiGLU expert over a `[n_tok, hidden]` activation, storage-dispatching —
/// `down(silu(gate·x) * (up·x))` via [`linear_q_dynamic`] (int8 arm =
/// [`nn::linear_int8_dynamic`] exactly, so it cannot drift from
/// [`expert_mlp_i8`]; int4 arm = packed [`linear_int4_dynamic`]).
pub(crate) fn expert_mlp_q(
    x: &Mat,
    gate: &QLinear,
    up: &QLinear,
    down: &QLinear,
    calib: Option<(usize, quant_calib::MlpUnit)>,
) -> FocrResult<Mat> {
    if let Some((layer, unit)) = calib.filter(|_| calib::enabled()) {
        calib::record_rows(&quant_calib::ffn_in_key(layer, unit), &x.data, x.cols);
    }
    let mut g = linear_q_dynamic(x, gate)?;
    nn::silu(&mut g);
    let u = linear_q_dynamic(x, up)?;
    if g.data.len() != u.data.len() {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::expert_mlp_q: gate/up shape mismatch ({} vs {})",
            g.data.len(),
            u.data.len()
        )));
    }
    for (a, &b) in g.data.iter_mut().zip(u.data.iter()) {
        *a *= b;
    }
    if let Some((layer, unit)) = calib.filter(|_| calib::enabled()) {
        calib::record_rows(&quant_calib::down_in_key(layer, unit), &g.data, g.cols);
    }
    linear_q_dynamic(&g, down)
}

/// Int8 MoE block over `[n_tok, hidden]` — the int8 twin of [`moe::moe_block`]:
/// route top-6 (f32 router gate), gather each selected token into a compact
/// activation, run [`expert_mlp_i8`], scatter back weighted, then add the shared
/// expert at weight 1.0. Mathematically identical structure; only the GEMMs go int8.
#[allow(clippy::needless_range_loop)]
fn moe_block_i8(
    hidden: &Mat,
    gate: &[f32],
    experts: &[[QLinear; 3]],
    shared: &[QLinear; 3],
    calib_layer: Option<usize>,
) -> FocrResult<Mat> {
    let n_tok = hidden.rows;
    let h = hidden.cols;
    let routing = moe::route_default(hidden, gate)?;
    let mut out = Mat::zeros(n_tok, h);
    let contribution_len = n_tok
        .checked_mul(moe::config::NUM_EXPERTS_PER_TOK)
        .and_then(|rows| rows.checked_mul(h))
        .ok_or_else(|| {
            FocrError::Other(anyhow::anyhow!(
                "decoder::moe_block_i8: n_tok*top_k*hidden overflow"
            ))
        })?;
    let mut contributions = vec![0.0f32; contribution_len];
    let mut per_expert: Vec<Vec<(usize, usize, f32)>> =
        vec![Vec::new(); moe::config::N_ROUTED_EXPERTS];
    for t in 0..n_tok {
        for j in 0..moe::config::NUM_EXPERTS_PER_TOK {
            per_expert[routing.indices[t][j]].push((t, j, routing.weights[t][j]));
        }
    }
    for (e, members) in per_expert.iter().enumerate() {
        if members.is_empty() {
            continue;
        }
        let m = members.len();
        let mut sub = Mat::zeros(m, h);
        for (r, &(t, _slot, _w)) in members.iter().enumerate() {
            sub.row_mut(r).copy_from_slice(hidden.row(t));
        }
        let y = expert_mlp_q(
            &sub,
            &experts[e][0],
            &experts[e][1],
            &experts[e][2],
            calib_layer.map(|l| (l, quant_calib::MlpUnit::Expert(e))),
        )?;
        for (r, &(t, slot, w)) in members.iter().enumerate() {
            let yr = y.row(r);
            let base = (t * moe::config::NUM_EXPERTS_PER_TOK + slot) * h;
            let dst = &mut contributions[base..base + h];
            for (value, &expert_value) in dst.iter_mut().zip(yr.iter()) {
                *value = w * expert_value;
            }
        }
    }
    for t in 0..n_tok {
        let token_base = t * moe::config::NUM_EXPERTS_PER_TOK * h;
        let rows: [&[f32]; moe::config::NUM_EXPERTS_PER_TOK] = std::array::from_fn(|slot| {
            let base = token_base + slot * h;
            &contributions[base..base + h]
        });
        moe::combine_routed_rows(rows, &routing.indices[t], out.row_mut(t))?;
    }
    let shared_out = expert_mlp_q(
        hidden,
        &shared[0],
        &shared[1],
        &shared[2],
        calib_layer.map(|l| (l, quant_calib::MlpUnit::Shared)),
    )?;
    for (o, &s) in out.data.iter_mut().zip(shared_out.data.iter()) {
        *o += s;
    }
    Ok(out)
}

/// Run one int8 layer's MLP/MoE over the `post_attention_layernorm`'d hidden
/// (prefill, m>1). Int8 twin of [`cached_mlp`].
fn prefill_mlp_i8(mlp: &CachedMlpI8, normed: &Mat, calib_layer: Option<usize>) -> FocrResult<Mat> {
    match mlp {
        CachedMlpI8::Dense { gate, up, down } => expert_mlp_q(
            normed,
            gate,
            up,
            down,
            calib_layer.map(|l| (l, quant_calib::MlpUnit::Dense)),
        ),
        CachedMlpI8::Moe {
            gate,
            experts,
            shared,
        } => moe_block_i8(normed, gate, experts, shared, calib_layer),
    }
}

/// Int8 decode MLP/MoE over a single `post_attention_layernorm`'d row. Int8 twin
/// of [`decode_mlp`] (route top-k via the f32 gate, weighted [`expert_gemv_i8`]
/// sum, + shared expert).
fn decode_mlp_i8(
    mlp: &CachedMlpI8,
    normed: &Mat,
    calib_layer: Option<usize>,
) -> FocrResult<Vec<f32>> {
    let hidden = config::HIDDEN_SIZE;
    let row = normed.row(0);
    let profiling = prof::enabled();
    match mlp {
        CachedMlpI8::Dense { gate, up, down } => {
            let t = profiling.then(Instant::now);
            let y = expert_gemv_i8(
                row,
                gate,
                up,
                down,
                calib_layer.map(|l| (l, quant_calib::MlpUnit::Dense)),
            );
            if let Some(t) = t {
                prof::add(&prof::EXPERTS_NS, t.elapsed().as_nanos() as u64);
            }
            Ok(y)
        }
        CachedMlpI8::Moe {
            gate,
            experts,
            shared,
        } => {
            let tr = profiling.then(Instant::now);
            let routing = moe::route_default(normed, gate)?;
            if let Some(t) = tr {
                prof::add(&prof::ROUTE_NS, t.elapsed().as_nanos() as u64);
            }
            let te = profiling.then(Instant::now);
            // Fan the 6 routed experts across the pool as FULLY SERIAL tasks (one
            // dispatch, no nested parallelism), CONCURRENTLY with the larger shared
            // expert (inter 1792), which keeps its own internal parallelism via
            // `rayon::join`. Measured best of the schemes tried: the shared expert
            // gets the spare cores while the 6 routed saturate the rest, vs an
            // all-serial fan-out where the shared becomes a single-core long pole.
            let idx = routing.indices[0];
            let wts = routing.weights[0];
            let (partials, s) = rayon::join(
                || {
                    (0..moe::config::NUM_EXPERTS_PER_TOK)
                        .into_par_iter()
                        .map(|j| {
                            let e = idx[j];
                            let w = wts[j];
                            let mut y = expert_gemv_i8_serial(
                                row,
                                &experts[e][0],
                                &experts[e][1],
                                &experts[e][2],
                                calib_layer.map(|l| (l, quant_calib::MlpUnit::Expert(e))),
                            );
                            for v in y.iter_mut() {
                                *v *= w;
                            }
                            y
                        })
                        .collect::<Vec<Vec<f32>>>()
                },
                || {
                    expert_gemv_i8(
                        row,
                        &shared[0],
                        &shared[1],
                        &shared[2],
                        calib_layer.map(|l| (l, quant_calib::MlpUnit::Shared)),
                    )
                },
            );
            let mut out = vec![0.0f32; hidden];
            let rows: [&[f32]; moe::config::NUM_EXPERTS_PER_TOK] =
                std::array::from_fn(|slot| partials[slot].as_slice());
            moe::combine_routed_rows(rows, &idx, &mut out)?;
            for c in 0..hidden {
                out[c] += s[c];
            }
            if let Some(t) = te {
                prof::add(&prof::EXPERTS_NS, t.elapsed().as_nanos() as u64);
            }
            Ok(out)
        }
    }
}

/// Final RMSNorm + int8 `lm_head` over the decode hidden `[1, hidden]`. Int8 twin
/// of [`lm_head_cached`].
///
/// # Errors
/// [`FocrError::Other`] on a shape mismatch.
pub fn lm_head_cached_i8(wc: &DecoderWeightCacheI8, hidden: &Mat) -> FocrResult<Mat> {
    if hidden.rows != 1 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::lm_head_cached_i8: expected a single decode row, got {} rows",
            hidden.rows
        )));
    }
    let t = prof::enabled().then(Instant::now);
    let normed = nn::rms_norm(hidden, Some(&wc.final_norm), config::RMS_NORM_EPS)?;
    let row = normed.row(0);
    // FOCR_LMHEAD_SHARD: vocab-tiled head (default OFF ⇒ the monolithic `gemv_i8`).
    // Byte-for-byte identical either way — each logit is an independent int8 dot.
    let logits = if lmhead_shard_enabled() {
        gemv_i8_sharded(row, &wc.lm_head, lmhead_shard_tiles())
    } else {
        gemv_i8(row, &wc.lm_head)
    };
    if let Some(t) = t {
        prof::add(&prof::LMHEAD_NS, t.elapsed().as_nanos() as u64);
    }
    Ok(Mat::from_vec(1, config::VOCAB_SIZE, logits))
}

// ── ngram-ban into the int8 lm_head epilogue (FOCR_FUSE_NGRAM_LMHEAD, bd-1azu.54) ──
//
// The greedy decode loop produces `[1, vocab]` lm_head logits, then the sampler's
// `masked_sliding_window_logits_if_needed` COPIES the whole 129280-wide row and
// sets every sliding-window no-repeat-ngram-banned token to -inf before argmax.
// Because the ban SET is a pure function of the generated sequence (independent of
// the logit values), it can be folded into the lm_head dequant epilogue: as each
// output channel `o` is dequantized, a banned `o` is written -inf instead of its
// dot product. This produces a logits row BYTE-FOR-BYTE identical to `gemv_i8`
// followed by `masked_sliding_window_logits_if_needed` — non-banned channels keep
// the same `acc·a_scale·scale[o]`, banned channels are -inf in both — so the argmax
// (and the chosen token) is unchanged, with no separate full-row copy/mask pass.

/// `FOCR_FUSE_NGRAM_LMHEAD` (bd-1azu.54, Lever 3): fold the sliding-window
/// no-repeat-ngram ban into the int8 lm_head dequant epilogue (mask as logits are
/// produced) instead of the sampler's separate copy-then-mask pass. DEFAULT OFF —
/// unset reproduces the `lm_head_cached_i8` -> `sampler::decode_step` path
/// byte-for-byte. Read ONCE into a process-wide bool.
const FUSE_NGRAM_LMHEAD_ENV: &str = "FOCR_FUSE_NGRAM_LMHEAD";

/// Read [`FUSE_NGRAM_LMHEAD_ENV`] ONCE into a process-wide bool (consulted by the
/// decode driver, never inside a per-channel loop).
pub fn fuse_ngram_lmhead_enabled() -> bool {
    static FLAG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *FLAG.get_or_init(|| std::env::var_os(FUSE_NGRAM_LMHEAD_ENV).is_some())
}

/// Int8 lm_head GEMV with the no-repeat-ngram ban folded into the dequant epilogue
/// (FOCR_FUSE_NGRAM_LMHEAD). `banned` is the set of vocab ids to mask to -inf
/// (from [`sampler::collect_sliding_window_ngram_bans`]); out-of-range ids are
/// ignored, and a repeated id is idempotent. With `banned` empty this is exactly
/// [`gemv_i8`]. BYTE-FOR-BYTE identical to `gemv_i8(x, qw)` then setting `y[b] =
/// -inf` for each in-range `b`: non-banned channels are the SAME
/// `acc·a_scale·scale[o]`, banned channels are `f32::NEG_INFINITY`.
fn gemv_i8_ngram_masked(x: &[f32], qw: &QInt8, banned: &[u32]) -> Vec<f32> {
    let n = qw.n;
    if banned.is_empty() {
        // No ban this step ⇒ byte-for-byte the unmasked head (matches the sampler's
        // `masked_sliding_window_logits_if_needed` returning `None`).
        return gemv_i8(x, qw);
    }
    let mut ban_mask = vec![false; n];
    for &b in banned {
        let bi = b as usize;
        if bi < n {
            ban_mask[bi] = true;
        }
    }
    let k = qw.k;
    debug_assert_eq!(x.len(), k);
    debug_assert_eq!(qw.w.len(), qw.expected_w_len());
    debug_assert_eq!(qw.scales.len(), n);
    let (xq, a_scale) = quantize_row_i8(x);
    let mut y = vec![0.0f32; n];
    y.par_chunks_mut(I8_GEMV_BLOCK)
        .enumerate()
        .for_each(|(blk, ys)| {
            let base = blk * I8_GEMV_BLOCK;
            let cnt = ys.len();
            let mut acc = vec![0i32; cnt];
            igemm_i8_block(qw, &xq, 1, base, cnt, &mut acc);
            for (j, slot) in ys.iter_mut().enumerate() {
                *slot = if ban_mask[base + j] {
                    f32::NEG_INFINITY
                } else {
                    acc[j] as f32 * a_scale * qw.scales[base + j]
                };
            }
        });
    y
}

/// Final RMSNorm + int8 lm_head with the no-repeat-ngram ban folded into the
/// dequant epilogue — the FOCR_FUSE_NGRAM_LMHEAD twin of [`lm_head_cached_i8`].
/// Returns `[1, vocab]` logits already masked, so the caller argmaxes directly
/// (no separate sampler masking pass; see [`sampler::decode_step_premasked`]). The
/// masked row is byte-for-byte the [`lm_head_cached_i8`] logits with
/// `masked_sliding_window_logits_if_needed` applied — the monolithic head is itself
/// byte-identical to the `FOCR_LMHEAD_SHARD` tiling, so this matches regardless of
/// that flag.
///
/// # Errors
/// [`FocrError::Other`] on a shape mismatch (mirrors [`lm_head_cached_i8`]).
pub fn lm_head_cached_i8_ngram_masked(
    wc: &DecoderWeightCacheI8,
    hidden: &Mat,
    banned: &[u32],
) -> FocrResult<Mat> {
    if hidden.rows != 1 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::lm_head_cached_i8_ngram_masked: expected a single decode row, got {} rows",
            hidden.rows
        )));
    }
    let t = prof::enabled().then(Instant::now);
    let normed = nn::rms_norm(hidden, Some(&wc.final_norm), config::RMS_NORM_EPS)?;
    let row = normed.row(0);
    let logits = gemv_i8_ngram_masked(row, &wc.lm_head, banned);
    if let Some(t) = t {
        prof::add(&prof::LMHEAD_NS, t.elapsed().as_nanos() as u64);
    }
    Ok(Mat::from_vec(1, config::VOCAB_SIZE, logits))
}

/// Int8 prefill: the [`prefill_with_cache`] structure with the GEMMs routed through
/// [`nn::linear_int8_dynamic`] (threaded dense-int8 dispatch). Captures each layer's RoPE'd
/// K/V into the R-SWA ring just like the f32 path. Returns the final
/// `model.norm`-ready hidden `[seq, hidden]` + the 12 populated caches.
///
/// `FOCR_PREFILL_CHUNK` ([`prefill_chunk_size`]) arms the chunked schedule
/// ([`prefill_with_cache_i8_chunked`]); unset ⇒ the monolithic loop, byte-for-byte.
///
/// # Errors
/// As [`prefill_with_cache`].
pub fn prefill_with_cache_i8(
    wc: &DecoderWeightCacheI8,
    inputs_embeds: &Mat,
) -> FocrResult<(Mat, Vec<RingCache>)> {
    prefill_with_cache_i8_chunked(wc, inputs_embeds, prefill_chunk_size())
}

/// [`prefill_with_cache_i8`] with the chunk decision made explicit (the int8 twin
/// of [`prefill_with_cache_chunked`]). `chunk = None` is the exact monolithic
/// int8 path; `chunk = Some(C)` tiles the prefill into `C`-token chunks. The
/// attention is the SAME f32 [`chunk_prefill_attention`] as the f32 path — only
/// the projections/MLP run int8 — and every chunked op is row-independent, so it
/// is byte-for-byte the monolithic int8 prefill.
///
/// # Errors
/// As [`prefill_with_cache_i8`].
pub fn prefill_with_cache_i8_chunked(
    wc: &DecoderWeightCacheI8,
    inputs_embeds: &Mat,
    chunk: Option<usize>,
) -> FocrResult<(Mat, Vec<RingCache>)> {
    checked_mat_len(
        "decoder::prefill_with_cache_i8 inputs_embeds",
        inputs_embeds,
    )?;
    let hidden = config::HIDDEN_SIZE;
    let num_heads = config::NUM_ATTENTION_HEADS;
    let head_dim = config::HEAD_DIM;
    // qkv width = num_heads*head_dim (the int8 linears infer it from `q_proj.n`,
    // so it isn't threaded through; we still need it for the chunked K/V buffers).
    let qkv_dim = checked_shape_mul(
        "decoder::prefill_with_cache_i8",
        num_heads,
        head_dim,
        "num_heads*head_dim",
    )?;
    let eps = config::RMS_NORM_EPS;
    if inputs_embeds.cols != hidden {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::prefill_with_cache_i8: inputs_embeds cols {} != hidden {hidden}",
            inputs_embeds.cols
        )));
    }
    let seq = inputs_embeds.rows;
    let mut caches: Vec<RingCache> = (0..config::NUM_HIDDEN_LAYERS)
        .map(|_| RingCache::new(seq.max(1)))
        .collect();

    if let Some(chunk) = chunk {
        // ── Chunked schedule (int8 projections/MLP; f32 chunked attention). ──
        let mut out = Mat::zeros(seq, hidden);
        let mut k_full: Vec<Mat> = (0..config::NUM_HIDDEN_LAYERS)
            .map(|_| Mat::zeros(seq, qkv_dim))
            .collect();
        let mut v_full: Vec<Mat> = (0..config::NUM_HIDDEN_LAYERS)
            .map(|_| Mat::zeros(seq, qkv_dim))
            .collect();
        for (c0, c1) in prefill_chunk_bounds(seq, chunk) {
            let positions: Vec<usize> = (c0..c1).collect();
            let rope = RopeTable::build(&positions, head_dim, config::ROPE_THETA);
            let mut x = Mat::from_vec(
                c1 - c0,
                hidden,
                inputs_embeds.data[c0 * hidden..c1 * hidden].to_vec(),
            );
            for layer in 0..config::NUM_HIDDEN_LAYERS {
                let cl = &wc.layers[layer];
                let normed = nn::rms_norm(&x, Some(&cl.input_ln), eps)?;
                let mut q = nn::linear_int8_dynamic(&normed, &cl.q_proj, None)?;
                let mut k = nn::linear_int8_dynamic(&normed, &cl.k_proj, None)?;
                let v = nn::linear_int8_dynamic(&normed, &cl.v_proj, None)?;
                apply_rope(&mut q, &rope)?;
                apply_rope(&mut k, &rope)?;
                k_full[layer].data[c0 * qkv_dim..c1 * qkv_dim].copy_from_slice(&k.data);
                v_full[layer].data[c0 * qkv_dim..c1 * qkv_dim].copy_from_slice(&v.data);
                let kpre = Mat::from_vec(c1, qkv_dim, k_full[layer].data[..c1 * qkv_dim].to_vec());
                let vpre = Mat::from_vec(c1, qkv_dim, v_full[layer].data[..c1 * qkv_dim].to_vec());
                let context = chunk_prefill_attention(&q, &kpre, &vpre, num_heads, head_dim, c0)?;
                let attn_out = nn::linear_int8_dynamic(&context, &cl.o_proj, None)?;
                let h = add_residual(&x, &attn_out)?;
                let normed2 = nn::rms_norm(&h, Some(&cl.post_attn_ln), eps)?;
                let mlp_out = prefill_mlp_i8(&cl.mlp, &normed2, calib::enabled().then_some(layer))?;
                x = add_residual(&h, &mlp_out)?;
            }
            out.data[c0 * hidden..c1 * hidden].copy_from_slice(&x.data);
        }
        for layer in 0..config::NUM_HIDDEN_LAYERS {
            let (kh, vh) = token_major_to_head_major(
                &k_full[layer],
                &v_full[layer],
                seq,
                num_heads,
                head_dim,
            )?;
            caches[layer].record_prefill(&kh, &vh, seq)?;
        }
        return Ok((out, caches));
    }

    // ── Monolithic schedule (the unset default). ──
    let positions: Vec<usize> = (0..seq).collect();
    let rope = RopeTable::build(&positions, head_dim, config::ROPE_THETA);

    let mut x = inputs_embeds.clone();
    for layer in 0..config::NUM_HIDDEN_LAYERS {
        let cl = &wc.layers[layer];
        let normed = nn::rms_norm(&x, Some(&cl.input_ln), eps)?;
        let mut q = nn::linear_int8_dynamic(&normed, &cl.q_proj, None)?;
        let mut k = nn::linear_int8_dynamic(&normed, &cl.k_proj, None)?;
        let v = nn::linear_int8_dynamic(&normed, &cl.v_proj, None)?;
        apply_rope(&mut q, &rope)?;
        apply_rope(&mut k, &rope)?;
        let (kh, vh) = token_major_to_head_major(&k, &v, seq, num_heads, head_dim)?;
        caches[layer].record_prefill(&kh, &vh, seq)?;
        let context = prefill_attention(&q, &k, &v, num_heads, head_dim)?;
        let attn_out = nn::linear_int8_dynamic(&context, &cl.o_proj, None)?;
        let h = add_residual(&x, &attn_out)?;
        let normed2 = nn::rms_norm(&h, Some(&cl.post_attn_ln), eps)?;
        let mlp_out = prefill_mlp_i8(&cl.mlp, &normed2, calib::enabled().then_some(layer))?;
        x = add_residual(&h, &mlp_out)?;
    }
    Ok((x, caches))
}

/// One int8 incremental decode step over the per-layer [`RingCache`]s. Int8 twin
/// of [`decode_step_with_cache`] — every projection via the bespoke m=1 [`gemv_i8`]
/// (Apple autovec or forced NEON routes / x86 VNNI), with identical R-SWA and ring contracts.
///
/// # Errors
/// As [`decode_step_with_cache`].
pub fn decode_step_with_cache_i8(
    wc: &DecoderWeightCacheI8,
    caches: &mut [RingCache],
    token_embed: &Mat,
    position: usize,
) -> FocrResult<Mat> {
    let hidden = config::HIDDEN_SIZE;
    let qkv_dim = checked_shape_mul(
        "decoder::decode_step_with_cache_i8",
        config::NUM_ATTENTION_HEADS,
        config::HEAD_DIM,
        "num_heads*head_dim",
    )?;
    let eps = config::RMS_NORM_EPS;
    if token_embed.rows != 1 || token_embed.cols != hidden {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::decode_step_with_cache_i8: token_embed shape [{}, {}] != [1, {hidden}]",
            token_embed.rows,
            token_embed.cols
        )));
    }
    if caches.len() != config::NUM_HIDDEN_LAYERS {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::decode_step_with_cache_i8: {} caches != {} layers",
            caches.len(),
            config::NUM_HIDDEN_LAYERS
        )));
    }
    let rope = RopeTable::build(&[position], config::HEAD_DIM, config::ROPE_THETA);
    let profiling = prof::enabled();
    let mut x = token_embed.clone();
    for layer in 0..config::NUM_HIDDEN_LAYERS {
        let cl = &wc.layers[layer];
        let t_attn = profiling.then(Instant::now);
        let (mut q, mut k, v) = if fuse_norm_quant_enabled() {
            // FOCR_FUSE_NORM_QUANT (Lever 1): quantize the `input_layernorm` output
            // ONCE in the norm epilogue and reuse that int8 row across q/k/v —
            // byte-identical to re-quantizing `normed.row(0)` inside each `gemv_i8`
            // (the quantize is deterministic, so once-then-reuse == thrice).
            let (xq, a_scale) = rms_norm_quant_i8(&x, Some(&cl.input_ln), eps)?;
            if let Some(qkv) = &cl.qkv {
                let out = gemv_i8_prequant(&xq, a_scale, qkv);
                let q = Mat::from_vec(1, qkv_dim, out[0..qkv_dim].to_vec());
                let k = Mat::from_vec(1, qkv_dim, out[qkv_dim..2 * qkv_dim].to_vec());
                let v = out[2 * qkv_dim..3 * qkv_dim].to_vec();
                (q, k, v)
            } else {
                let q = Mat::from_vec(1, qkv_dim, gemv_i8_prequant(&xq, a_scale, &cl.q_proj));
                let k = Mat::from_vec(1, qkv_dim, gemv_i8_prequant(&xq, a_scale, &cl.k_proj));
                let v = gemv_i8_prequant(&xq, a_scale, &cl.v_proj);
                (q, k, v)
            }
        } else {
            let normed = nn::rms_norm(&x, Some(&cl.input_ln), eps)?;
            let nrow = normed.row(0);
            if let Some(qkv) = &cl.qkv {
                // FOCR_QKV_FUSED: ONE quantize of `nrow`, ONE block-parallel GEMV over
                // all 3*qkv_dim output rows, then slice into q/k/v. Byte-identical to
                // the three-call `else` branch (each output row is an independent dot).
                let out = gemv_i8(nrow, qkv);
                let q = Mat::from_vec(1, qkv_dim, out[0..qkv_dim].to_vec());
                let k = Mat::from_vec(1, qkv_dim, out[qkv_dim..2 * qkv_dim].to_vec());
                let v = out[2 * qkv_dim..3 * qkv_dim].to_vec();
                (q, k, v)
            } else {
                let q = Mat::from_vec(1, qkv_dim, gemv_i8(nrow, &cl.q_proj));
                let k = Mat::from_vec(1, qkv_dim, gemv_i8(nrow, &cl.k_proj));
                let v = gemv_i8(nrow, &cl.v_proj);
                (q, k, v)
            }
        };
        apply_rope(&mut q, &rope)?;
        apply_rope(&mut k, &rope)?;
        caches[layer].write_decode_step(&k.data, &v)?;
        let context = rswa::decode_attention(&caches[layer], &q.data)?;
        let attn_out = Mat::from_vec(1, hidden, gemv_i8(&context.data, &cl.o_proj));
        let h = add_residual(&x, &attn_out)?;
        if let Some(t) = t_attn {
            prof::add(&prof::ATTN_NS, t.elapsed().as_nanos() as u64);
        }
        let normed2 = nn::rms_norm(&h, Some(&cl.post_attn_ln), eps)?;
        let mlp_out = Mat::from_vec(
            1,
            hidden,
            decode_mlp_i8(&cl.mlp, &normed2, calib::enabled().then_some(layer))?,
        );
        x = add_residual(&h, &mlp_out)?;
    }
    Ok(x)
}

// ── Phase-6 batched continuous-decode spine (bd-1azu.3) ──────────────────────
//
// One decode step over `B` in-flight page-streams. Per layer, the LINEAR
// projections that share a single weight panel across streams — the fused q/k/v
// stack, `o_proj`, and the final `lm_head` — become ONE `M=B` int8 GEMM that
// reads each weight row once and reuses it across all `B` activation rows (the
// decode-throughput win). This is LOSSLESS by construction: the per-output i32
// contraction `acc[r,o] = Σ_p xq[r,p]·w[o,p]` is independent of `M`, so row `r`
// of the `M=B` GEMM is BYTE-FOR-BYTE the row a standalone `m=1` GEMV produces
// (proven across every tier/shape in `tests/batched_igemm_parity.rs`, bd-1azu.2).
// Everything genuinely per-stream — RoPE at each stream's TRUE absolute position,
// the R-SWA ring write, `decode_attention`, and the MoE top-k dispatch — stays a
// faithful per-stream loop over the existing single-stream kernels, so each
// stream's output equals `decode_step_with_cache_i8` run for that stream alone
// (Doctrine #1: correctness first).

/// `FOCR_BATCH_SPINE`: kill-switch that ARMS the Phase-6 batched continuous-decode
/// spine ([`batched_decode_step_i8`]). DEFAULT OFF — the spine is an additive
/// throughput path; the single-stream [`decode_step_with_cache_i8`] stays the
/// default decode. Consulted ONCE by the (future) batched driver, never inside the
/// math, exactly like [`QKV_FUSED_ENV`].
const BATCH_SPINE_ENV: &str = "FOCR_BATCH_SPINE";

/// `FOCR_BATCH_SIZE`: the maximum number of in-flight page-streams (`B`) the spine
/// admits per batched step. Parsed ONCE; defaults to [`DEFAULT_BATCH_SIZE`] when
/// unset, empty, unparseable, or `0`.
const BATCH_SIZE_ENV: &str = "FOCR_BATCH_SIZE";

/// Fallback in-flight stream cap `B` when [`BATCH_SIZE_ENV`] is unset/invalid.
const DEFAULT_BATCH_SIZE: usize = 8;

/// Whether the batched decode spine is armed (the [`BATCH_SPINE_ENV`] kill-switch,
/// read ONCE into a process-wide bool — never touched per-token). The batched
/// kernels themselves are pure and testable regardless of this flag; only the
/// driver's decision to route through them is gated here.
///
/// Value-parsed, NOT presence-parsed (fresh-eyes fix): every doc site
/// (`batch_scheduler`, `cli.rs`, the watchdog sweep) teaches `FOCR_BATCH_SPINE=0`
/// as "spine off" — the old `is_some()` parse ARMED the spine on exactly the
/// value users set to kill it, and made the sweep's `=0` control leg
/// meaningless. `0`/`off`/`false`/`no` (any case) now disable; any other
/// present value arms; unset stays off.
#[must_use]
pub fn batch_spine_enabled() -> bool {
    static FLAG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *FLAG.get_or_init(|| match std::env::var(BATCH_SPINE_ENV) {
        Ok(v) => !matches!(
            v.trim().to_ascii_lowercase().as_str(),
            "" | "0" | "off" | "false" | "no"
        ),
        Err(_) => false,
    })
}

/// The configured in-flight stream cap `B` ([`BATCH_SIZE_ENV`], read ONCE).
/// Always `>= 1`.
#[must_use]
pub fn batch_size_cap() -> usize {
    static CAP: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
    *CAP.get_or_init(|| {
        std::env::var(BATCH_SIZE_ENV)
            .ok()
            .and_then(|v| v.trim().parse::<usize>().ok())
            .filter(|&n| n > 0)
            .unwrap_or(DEFAULT_BATCH_SIZE)
    })
}

/// Batched (`M=B`) int8 GEMV: run ONE int8 GEMM of `B` stacked activation rows
/// against the shared `[n, k]` weight panel `qw`, returning `B` output rows each
/// length `n`. The `r`-th returned row is BYTE-FOR-BYTE identical to
/// [`gemv_i8`]`(rows[r], qw)`:
///
///  * each activation row is dynamically int8-quantized on its OWN
///    [`quantize_row_i8`] `a_scale` (exactly as the per-row GEMV);
///  * the i32 contraction `acc[r,o] = Σ_p xq[r,p]·w[o,p]` from one `M=B`
///    [`simd::igemm_s8s8`] equals the `m=1` contraction for row `r` (M-independent
///    — bd-1azu.2);
///  * the dequant `acc as f32 * a_scale[r] * scales[o]` uses the SAME operands in
///    the SAME left-associative order as [`gemv_i8`].
///
/// Output channels are fanned across the rayon pool in [`I8_GEMV_BLOCK`] chunks
/// (mirroring [`gemv_i8`]'s 64-row blocking); the row-block grouping never changes
/// a per-output reduction, so the result is invariant to it. Each chunk writes a
/// contiguous, disjoint slice of a channel-major `[n, b]` scratch (data-race-free
/// without `unsafe`), then we transpose into the per-stream rows the decode loop
/// slices for RoPE / attention.
fn gemv_i8_batched(rows: &[&[f32]], qw: &QInt8) -> Vec<Vec<f32>> {
    let b = rows.len();
    let k = qw.k;
    let n = qw.n;
    debug_assert_eq!(qw.w.len(), qw.expected_w_len());
    debug_assert_eq!(qw.scales.len(), n);
    if b == 0 {
        return Vec::new();
    }
    // Quantize each activation row ONCE (per-row symmetric int8), packed `[b, k]`.
    let mut xq = vec![0i8; b * k];
    let mut a_scales = vec![0.0f32; b];
    for (r, &row) in rows.iter().enumerate() {
        debug_assert_eq!(row.len(), k);
        let (q, a_scale) = quantize_row_i8(row);
        xq[r * k..(r + 1) * k].copy_from_slice(&q);
        a_scales[r] = a_scale;
    }
    // Channel-major output `[n, b]`: each `I8_GEMV_BLOCK`-channel chunk is a
    // contiguous, disjoint slice, so the rayon fan-out has no aliasing.
    let mut ycm = vec![0.0f32; n * b];
    ycm.par_chunks_mut(I8_GEMV_BLOCK * b)
        .enumerate()
        .for_each(|(blk, ys)| {
            let base = blk * I8_GEMV_BLOCK;
            let cnt = ys.len() / b;
            let mut acc = vec![0i32; b * cnt];
            igemm_i8_block(qw, &xq, b, base, cnt, &mut acc);
            for j in 0..cnt {
                let scale_o = qw.scales[base + j];
                for r in 0..b {
                    ys[j * b + r] = acc[r * cnt + j] as f32 * a_scales[r] * scale_o;
                }
            }
        });
    // Transpose channel-major `[n, b]` -> per-stream rows `[b][n]`.
    let mut out: Vec<Vec<f32>> = (0..b).map(|_| vec![0.0f32; n]).collect();
    for o in 0..n {
        let col = o * b;
        for r in 0..b {
            out[r][o] = ycm[col + r];
        }
    }
    out
}

/// One int8 incremental decode step over `B` in-flight page-streams — the Phase-6
/// batched twin of [`decode_step_with_cache_i8`]. Returns `B` hidden rows, each a
/// `[1, hidden]` [`Mat`]; the `s`-th row is BYTE-FOR-BYTE identical to
/// [`decode_step_with_cache_i8`] run for stream `s` alone (the lossless contract,
/// bd-1azu.3 — the per-projection invariant is proven by
/// `tests/batched_forward_parity.rs`, the kernel invariant by
/// `tests/batched_igemm_parity.rs`).
///
/// Per layer the q/k/v stack and `o_proj` run as ONE `M=B` [`gemv_i8_batched`]
/// each (the shared weight panel is read once and reused across all `B` rows),
/// while RoPE at each stream's true absolute `positions[s]`, the
/// [`BatchedRingCache`] ring write, [`rswa::decode_attention`], and the per-stream
/// MoE dispatch ([`decode_mlp_i8`]) stay a faithful per-stream loop over the
/// existing single-stream kernels. Each stream's K/V land in ITS OWN ring and the
/// score/softmax/value reduction is never cross-stream (bd-1waa-safe).
///
/// The MoE expert GEMMs deliberately STAY per-stream: the routed top-k experts
/// differ per stream, so there is no single shared weight panel to batch
/// losslessly. The dense layer-0 MLP and the per-layer shared expert DO share a
/// panel across streams — batching those is the (OFF-by-default) optimization
/// seam left for a later bead and intentionally not taken here (Doctrine #7:
/// lossless-but-unoptimized first).
///
/// `token_embeds[s]` is stream `s`'s `[1, hidden]` input embedding and
/// `positions[s]` its true absolute decode position; both slices must have length
/// `caches.num_streams()`. Holds NO model-cache lock (the caller owns `wc`); one
/// live forward (Doctrine #5).
///
/// # Errors
/// [`FocrError::Other`] on any shape/length mismatch, or as
/// [`decode_step_with_cache_i8`].
pub fn batched_decode_step_i8(
    wc: &DecoderWeightCacheI8,
    caches: &mut BatchedRingCache,
    token_embeds: &[Mat],
    positions: &[usize],
) -> FocrResult<Vec<Mat>> {
    // Advance EVERY stream in the cache, in cache order — the original full-batch
    // contract. Delegates to the subset-aware core with the identity stream map,
    // so this remains byte-for-byte the function its tests pin.
    let stream_ids: Vec<usize> = (0..caches.num_streams()).collect();
    batched_decode_step_i8_streams(wc, caches, &stream_ids, token_embeds, positions)
}

/// Subset-aware twin of [`batched_decode_step_i8`] (bd-1azu.11): advance ONLY the
/// `stream_ids` of the `caches` this step — the active set the continuous-batch
/// scheduler hands over, which shrinks as streams retire. `token_embeds[k]`,
/// `positions[k]`, and the returned row `k` all correspond to cache stream
/// `stream_ids[k]`. With `stream_ids == 0..num_streams()` this is exactly the
/// full-batch [`batched_decode_step_i8`].
///
/// LOSSLESS by construction, identically to [`batched_decode_step_i8`]: each
/// projection runs as ONE `M = stream_ids.len()` int8 GEMM whose per-row i32
/// contraction is `M`-independent (bd-1azu.2), so dropping retired rows from the
/// active set NEVER changes a surviving row — the throughput win of pruning
/// retired slots costs zero accuracy. RoPE at each stream's TRUE absolute
/// position, the R-SWA ring write into stream `stream_ids[k]`'s OWN ring,
/// `decode_attention`, and the MoE dispatch all stay a faithful per-stream loop.
///
/// # Errors
/// [`FocrError::Other`] if `token_embeds`/`positions` disagree with
/// `stream_ids.len()`, a `stream_ids` entry is out of range, the cache layer
/// count is wrong, or any `token_embeds[k]` is not `[1, hidden]`.
pub fn batched_decode_step_i8_streams(
    wc: &DecoderWeightCacheI8,
    caches: &mut BatchedRingCache,
    stream_ids: &[usize],
    token_embeds: &[Mat],
    positions: &[usize],
) -> FocrResult<Vec<Mat>> {
    let hidden = config::HIDDEN_SIZE;
    let qkv_dim = checked_shape_mul(
        "decoder::batched_decode_step_i8",
        config::NUM_ATTENTION_HEADS,
        config::HEAD_DIM,
        "num_heads*head_dim",
    )?;
    let eps = config::RMS_NORM_EPS;
    let b = stream_ids.len();
    if token_embeds.len() != b || positions.len() != b {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::batched_decode_step_i8: token_embeds {} / positions {} != active streams {b}",
            token_embeds.len(),
            positions.len()
        )));
    }
    if caches.num_layers() != config::NUM_HIDDEN_LAYERS {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::batched_decode_step_i8: {} cache layers != {} layers",
            caches.num_layers(),
            config::NUM_HIDDEN_LAYERS
        )));
    }
    let n_streams = caches.num_streams();
    for (k, &sid) in stream_ids.iter().enumerate() {
        if sid >= n_streams {
            return Err(FocrError::Other(anyhow::anyhow!(
                "decoder::batched_decode_step_i8: active slot {k} stream id {sid} >= cache streams {n_streams}"
            )));
        }
    }
    for (s, te) in token_embeds.iter().enumerate() {
        if te.rows != 1 || te.cols != hidden {
            return Err(FocrError::Other(anyhow::anyhow!(
                "decoder::batched_decode_step_i8: token_embeds[{s}] shape [{}, {}] != [1, {hidden}]",
                te.rows,
                te.cols
            )));
        }
    }
    let profiling = prof::enabled();
    // Per-stream running hidden `x[s]` (`[1, hidden]`), seeded from the embeds.
    let mut x: Vec<Mat> = token_embeds.to_vec();
    for layer in 0..config::NUM_HIDDEN_LAYERS {
        let cl = &wc.layers[layer];
        let t_attn = profiling.then(Instant::now);
        // 1. Per-stream pre-attention RMSNorm.
        let mut normed: Vec<Mat> = Vec::with_capacity(b);
        for s in 0..b {
            normed.push(nn::rms_norm(&x[s], Some(&cl.input_ln), eps)?);
        }
        let nrows: Vec<&[f32]> = normed.iter().map(|m| m.row(0)).collect();
        // 2. ONE `M=B` int8 GEMM per q/k/v projection (the fused stack when armed,
        //    else three) — byte-identical per stream to the matching branch of
        //    `decode_step_with_cache_i8`.
        let (mut q_mats, mut k_mats, v_rows): (Vec<Mat>, Vec<Mat>, Vec<Vec<f32>>) =
            if let Some(qkv) = &cl.qkv {
                let outs = gemv_i8_batched(&nrows, qkv);
                let mut qm = Vec::with_capacity(b);
                let mut km = Vec::with_capacity(b);
                let mut vr = Vec::with_capacity(b);
                for row in outs {
                    qm.push(Mat::from_vec(1, qkv_dim, row[0..qkv_dim].to_vec()));
                    km.push(Mat::from_vec(
                        1,
                        qkv_dim,
                        row[qkv_dim..2 * qkv_dim].to_vec(),
                    ));
                    vr.push(row[2 * qkv_dim..3 * qkv_dim].to_vec());
                }
                (qm, km, vr)
            } else {
                let q_out = gemv_i8_batched(&nrows, &cl.q_proj);
                let k_out = gemv_i8_batched(&nrows, &cl.k_proj);
                let v_out = gemv_i8_batched(&nrows, &cl.v_proj);
                let qm: Vec<Mat> = q_out
                    .into_iter()
                    .map(|r| Mat::from_vec(1, qkv_dim, r))
                    .collect();
                let km: Vec<Mat> = k_out
                    .into_iter()
                    .map(|r| Mat::from_vec(1, qkv_dim, r))
                    .collect();
                (qm, km, v_out)
            };
        // 3. Per-stream: RoPE at the stream's TRUE absolute position, ring write,
        //    R-SWA decode attention (never cross-stream).
        let mut context_rows: Vec<Mat> = Vec::with_capacity(b);
        for s in 0..b {
            let rope = RopeTable::build(&[positions[s]], config::HEAD_DIM, config::ROPE_THETA);
            apply_rope(&mut q_mats[s], &rope)?;
            apply_rope(&mut k_mats[s], &rope)?;
            caches.write_decode_step(stream_ids[s], layer, &k_mats[s].data, &v_rows[s])?;
            context_rows.push(rswa::decode_attention(
                caches.layer(stream_ids[s], layer),
                &q_mats[s].data,
            )?);
        }
        // 4. ONE `M=B` `o_proj` over the stacked per-stream contexts.
        let ctx_refs: Vec<&[f32]> = context_rows.iter().map(|m| m.row(0)).collect();
        let attn_rows = gemv_i8_batched(&ctx_refs, &cl.o_proj);
        if let Some(t) = t_attn {
            prof::add(&prof::ATTN_NS, t.elapsed().as_nanos() as u64);
        }
        // 5. Per-stream residual, post-attn RMSNorm, MoE/MLP dispatch, residual.
        for (s, attn) in attn_rows.into_iter().enumerate() {
            let attn_out = Mat::from_vec(1, hidden, attn);
            let h = add_residual(&x[s], &attn_out)?;
            let normed2 = nn::rms_norm(&h, Some(&cl.post_attn_ln), eps)?;
            let mlp_out = Mat::from_vec(
                1,
                hidden,
                decode_mlp_i8(&cl.mlp, &normed2, calib::enabled().then_some(layer))?,
            );
            x[s] = add_residual(&h, &mlp_out)?;
        }
    }
    Ok(x)
}

/// Batched (`M=B`) final RMSNorm + int8 `lm_head` over `B` decode hidden rows —
/// the Phase-6 twin of [`lm_head_cached_i8`]. Returns `B` logit rows, each
/// `[1, VOCAB_SIZE]`; row `s` is BYTE-FOR-BYTE identical to [`lm_head_cached_i8`]
/// for stream `s` (per-stream final RMSNorm, then ONE shared-panel `M=B`
/// [`gemv_i8_batched`] over the `lm_head`).
///
/// # Errors
/// [`FocrError::Other`] if any `hiddens[s]` is not a single `[1, hidden]` row.
pub fn batched_lm_head_i8(wc: &DecoderWeightCacheI8, hiddens: &[Mat]) -> FocrResult<Vec<Mat>> {
    let t = prof::enabled().then(Instant::now);
    let mut normed: Vec<Mat> = Vec::with_capacity(hiddens.len());
    for (s, h) in hiddens.iter().enumerate() {
        if h.rows != 1 {
            return Err(FocrError::Other(anyhow::anyhow!(
                "decoder::batched_lm_head_i8: hiddens[{s}] has {} rows, expected 1",
                h.rows
            )));
        }
        normed.push(nn::rms_norm(h, Some(&wc.final_norm), config::RMS_NORM_EPS)?);
    }
    let rows: Vec<&[f32]> = normed.iter().map(|m| m.row(0)).collect();
    let outs = gemv_i8_batched(&rows, &wc.lm_head);
    let logits = outs
        .into_iter()
        .map(|y| Mat::from_vec(1, config::VOCAB_SIZE, y))
        .collect();
    if let Some(t) = t {
        prof::add(&prof::LMHEAD_NS, t.elapsed().as_nanos() as u64);
    }
    Ok(logits)
}

// ── Batched K-token speculative VERIFY forward (Lever D/K, bd-1azu.30) ────────
//
// Speculative decode proposes `K` draft tokens after the current sequence; the
// verify half must compute the `K` next-token logits rows — one per draft
// position — in ONE forward, BIT-EXACT to running `K` sequential single-token
// decode steps. The win is QUERY-dim batching: the LINEAR projections that share
// a weight panel across the `K` draft queries (the fused q/k/v stack and
// `o_proj`) become ONE `M=K` int8 GEMM (read each weight row once, reuse across
// all `K` queries — bd-1azu.2's M-independent i32 contraction), while the
// attention stays a faithful per-position fold.
//
// LOSSLESS by construction, and NOT the rejected key-batch (docs/NEGATIVE_EVIDENCE.md):
//  * the only cross-token coupling in a decode forward is the KV cache, so
//    processing the `K` positions LAYER-major (all `K` through layer L, then
//    layer L+1) over a SHARED evolving ring reproduces the token-major sequential
//    KV evolution exactly — draft `i` at layer L writes its K/V then attends over
//    reference ++ ring ++ draft[0..=i] at L, precisely as the sequential step
//    would (draft[0..i] already written, draft[i+1..] not yet);
//  * that per-position causal attention is delegated to
//    [`rswa::verify_attention`], which replays the draft writes into a PRIVATE
//    clone of the layer's ring (the caller's caches are never mutated) — so each
//    context row is byte-for-byte the matching sequential [`rswa::attention`];
//  * every per-position op outside attention (RMSNorm, RoPE at the TRUE absolute
//    position `base_position + i`, the dense/MoE MLP, the residual adds, the final
//    RMSNorm + `lm_head`) is the SAME single-token kernel the sequential decode
//    runs, invoked once per draft position — row-independent, hence identical.
// The per-projection `M=K`-vs-`m=1` byte-identity is the bd-1azu.2 gate
// (`tests/batched_igemm_parity.rs` / `tests/batched_forward_parity.rs`); the
// per-position verify attention + no-mutation is the bd-1azu.30 gate
// (`tests/spec_verify_forward_parity.rs`).

/// `FOCR_SPEC_VERIFY`: presence kill-switch that ARMS the bd-1azu.30 batched
/// speculative verify forward ([`verify_forward`] / [`verify_forward_i8`]) at the
/// (future) decode driver. DEFAULT OFF — when unset the verify forwards are simply
/// unused and the live decode path is byte-for-byte today's. Consulted ONCE by the
/// speculative-decode driver (bd-1azu.35), NEVER inside the verify math, exactly
/// like [`BATCH_SPINE_ENV`] / [`QKV_FUSED_ENV`].
const SPEC_VERIFY_ENV: &str = "FOCR_SPEC_VERIFY";

/// Whether the batched speculative verify forward is armed (the [`SPEC_VERIFY_ENV`]
/// kill-switch, read ONCE into a process-wide bool — never touched per-token). The
/// verify kernels are pure and testable regardless of this flag; only the (later)
/// driver's decision to route through them is gated here.
#[must_use]
pub fn spec_verify_enabled() -> bool {
    static FLAG: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *FLAG.get_or_init(|| std::env::var_os(SPEC_VERIFY_ENV).is_some())
}

/// Validate the `K` draft `token_embeds` against `[1, hidden]` and the per-layer
/// `caches` count — the shared prologue of [`verify_forward`] / [`verify_forward_i8`].
#[allow(dead_code)] // bd-1azu.30 verify seam: used only by the (gated) verify forwards.
fn check_verify_inputs(
    context: &str,
    caches: &[RingCache],
    token_embeds: &[Mat],
) -> FocrResult<()> {
    let hidden = config::HIDDEN_SIZE;
    if caches.len() != config::NUM_HIDDEN_LAYERS {
        return Err(FocrError::Other(anyhow::anyhow!(
            "decoder::{context}: {} caches != {} layers",
            caches.len(),
            config::NUM_HIDDEN_LAYERS
        )));
    }
    for (i, te) in token_embeds.iter().enumerate() {
        if te.rows != 1 || te.cols != hidden {
            return Err(FocrError::Other(anyhow::anyhow!(
                "decoder::{context}: token_embeds[{i}] shape [{}, {}] != [1, {hidden}]",
                te.rows,
                te.cols
            )));
        }
    }
    Ok(())
}

/// Batched `K`-token speculative VERIFY forward (f32) — bd-1azu.30. Given the `K`
/// draft tokens' embeddings `token_embeds` (each `[1, hidden]`) proposed after the
/// sequence whose KV is in `caches`, with draft position `i` at TRUE absolute
/// position `base_position + i`, return the `K` next-token logits rows. Row `i` is
/// BYTE-FOR-BYTE the logits running `i+1` sequential [`decode_step_with_cache`]
/// calls (then [`lm_head_cached`]) would emit for draft position `i` — the
/// lossless verify contract (see the section note above).
///
/// `caches` is READ-ONLY: [`rswa::verify_attention`] folds each draft token into a
/// private clone of the per-layer ring, so the caller's caches are left untouched
/// (no checkpoint/rollback needed).
///
/// # Errors
/// [`FocrError::Other`] if `caches` is not `NUM_HIDDEN_LAYERS` long or any
/// `token_embeds[i]` is not `[1, hidden]`; propagates the per-layer kernel errors.
#[allow(dead_code)] // bd-1azu.30 verify seam: consumed by the speculative decode loop (later bead).
pub(crate) fn verify_forward(
    wc: &DecoderWeightCache,
    caches: &[RingCache],
    token_embeds: &[Mat],
    base_position: usize,
) -> FocrResult<Vec<Mat>> {
    let hidden = config::HIDDEN_SIZE;
    let qkv_dim = checked_shape_mul(
        "decoder::verify_forward",
        config::NUM_ATTENTION_HEADS,
        config::HEAD_DIM,
        "num_heads*head_dim",
    )?;
    let eps = config::RMS_NORM_EPS;
    check_verify_inputs("verify_forward", caches, token_embeds)?;
    let k = token_embeds.len();

    // Per-draft-position running hidden `x[i]` (`[1, hidden]`), seeded from embeds.
    let mut x: Vec<Mat> = token_embeds.to_vec();
    for layer in 0..config::NUM_HIDDEN_LAYERS {
        let cl = &wc.layers[layer];
        // 1. Per-position pre-attn RMSNorm + q/k/v projection (the bespoke f32
        //    `gemv`, identical to the sequential `decode_step_with_cache`) + RoPE
        //    at the draft position's TRUE absolute position.
        let mut q_mats: Vec<Mat> = Vec::with_capacity(k);
        let mut k_mats: Vec<Mat> = Vec::with_capacity(k);
        let mut v_rows: Vec<Vec<f32>> = Vec::with_capacity(k);
        for i in 0..k {
            let normed = nn::rms_norm(&x[i], Some(&cl.input_ln), eps)?;
            let nrow = normed.row(0);
            let mut q = Mat::from_vec(1, qkv_dim, gemv(nrow, &cl.q_proj, qkv_dim, hidden));
            let mut kk = Mat::from_vec(1, qkv_dim, gemv(nrow, &cl.k_proj, qkv_dim, hidden));
            let v = gemv(nrow, &cl.v_proj, qkv_dim, hidden);
            let rope = RopeTable::build(&[base_position + i], config::HEAD_DIM, config::ROPE_THETA);
            apply_rope(&mut q, &rope)?;
            apply_rope(&mut kk, &rope)?;
            q_mats.push(q);
            k_mats.push(kk);
            v_rows.push(v);
        }
        // 2. Query-batched verify attention: per draft position, fold its K/V into a
        //    private clone of THIS layer's ring then attend (causal among draft) —
        //    `caches[layer]` is not mutated.
        let q_refs: Vec<&[f32]> = q_mats.iter().map(|m| m.data.as_slice()).collect();
        let k_refs: Vec<&[f32]> = k_mats.iter().map(|m| m.data.as_slice()).collect();
        let v_refs: Vec<&[f32]> = v_rows.iter().map(|r| r.as_slice()).collect();
        let contexts = rswa::verify_attention(&caches[layer], &q_refs, &k_refs, &v_refs)?;
        // 3. Per-position `o_proj`, residual, post-attn RMSNorm, MLP, residual.
        for i in 0..k {
            let o = gemv(&contexts[i].data, &cl.o_proj, hidden, qkv_dim);
            let attn_out = Mat::from_vec(1, hidden, o);
            let h = add_residual(&x[i], &attn_out)?;
            let normed2 = nn::rms_norm(&h, Some(&cl.post_attn_ln), eps)?;
            let mlp_out = Mat::from_vec(
                1,
                hidden,
                decode_mlp(&cl.mlp, &normed2, calib::enabled().then_some(layer))?,
            );
            x[i] = add_residual(&h, &mlp_out)?;
        }
    }
    // Final RMSNorm + lm_head per draft position (each row-independent).
    x.iter().map(|h| lm_head_cached(wc, h)).collect()
}

/// Batched `K`-token speculative VERIFY forward (int8) — the int8 twin of
/// [`verify_forward`] (bd-1azu.30). The q/k/v stack and `o_proj` run as ONE `M=K`
/// [`gemv_i8_batched`] each (the shared weight panel read once, reused across all
/// `K` draft queries — byte-identical per row to the `m=1` [`gemv_i8`] the
/// sequential [`decode_step_with_cache_i8`] runs, bd-1azu.2); RoPE, the verify
/// attention, the MoE dispatch, and the `lm_head` stay a faithful per-position
/// loop over the existing single-token kernels. Row `i` is BYTE-FOR-BYTE the
/// logits `i+1` sequential [`decode_step_with_cache_i8`] + [`lm_head_cached_i8`]
/// calls would emit.
///
/// `caches` is READ-ONLY (see [`verify_forward`]).
///
/// # Errors
/// As [`verify_forward`].
#[allow(dead_code)] // bd-1azu.30 verify seam: consumed by the speculative decode loop (later bead).
pub(crate) fn verify_forward_i8(
    wc: &DecoderWeightCacheI8,
    caches: &[RingCache],
    token_embeds: &[Mat],
    base_position: usize,
) -> FocrResult<Vec<Mat>> {
    let hidden = config::HIDDEN_SIZE;
    let qkv_dim = checked_shape_mul(
        "decoder::verify_forward_i8",
        config::NUM_ATTENTION_HEADS,
        config::HEAD_DIM,
        "num_heads*head_dim",
    )?;
    let eps = config::RMS_NORM_EPS;
    check_verify_inputs("verify_forward_i8", caches, token_embeds)?;
    let k = token_embeds.len();

    let mut x: Vec<Mat> = token_embeds.to_vec();
    for layer in 0..config::NUM_HIDDEN_LAYERS {
        let cl = &wc.layers[layer];
        // 1. Per-position pre-attn RMSNorm.
        let mut normed: Vec<Mat> = Vec::with_capacity(k);
        for i in 0..k {
            normed.push(nn::rms_norm(&x[i], Some(&cl.input_ln), eps)?);
        }
        let nrows: Vec<&[f32]> = normed.iter().map(|m| m.row(0)).collect();
        // 2. ONE `M=K` int8 GEMM per q/k/v projection (the fused stack when armed,
        //    else three) — byte-identical per draft position to the matching branch
        //    of `decode_step_with_cache_i8`.
        let (mut q_mats, mut k_mats, v_rows): (Vec<Mat>, Vec<Mat>, Vec<Vec<f32>>) =
            if let Some(qkv) = &cl.qkv {
                let outs = gemv_i8_batched(&nrows, qkv);
                let mut qm = Vec::with_capacity(k);
                let mut km = Vec::with_capacity(k);
                let mut vr = Vec::with_capacity(k);
                for row in outs {
                    qm.push(Mat::from_vec(1, qkv_dim, row[0..qkv_dim].to_vec()));
                    km.push(Mat::from_vec(
                        1,
                        qkv_dim,
                        row[qkv_dim..2 * qkv_dim].to_vec(),
                    ));
                    vr.push(row[2 * qkv_dim..3 * qkv_dim].to_vec());
                }
                (qm, km, vr)
            } else {
                let q_out = gemv_i8_batched(&nrows, &cl.q_proj);
                let k_out = gemv_i8_batched(&nrows, &cl.k_proj);
                let v_out = gemv_i8_batched(&nrows, &cl.v_proj);
                let qm: Vec<Mat> = q_out
                    .into_iter()
                    .map(|r| Mat::from_vec(1, qkv_dim, r))
                    .collect();
                let km: Vec<Mat> = k_out
                    .into_iter()
                    .map(|r| Mat::from_vec(1, qkv_dim, r))
                    .collect();
                (qm, km, v_out)
            };
        // 3. RoPE each q/k at the draft position's TRUE absolute position.
        for i in 0..k {
            let rope = RopeTable::build(&[base_position + i], config::HEAD_DIM, config::ROPE_THETA);
            apply_rope(&mut q_mats[i], &rope)?;
            apply_rope(&mut k_mats[i], &rope)?;
        }
        // 4. Query-batched verify attention over a private clone of the ring.
        let q_refs: Vec<&[f32]> = q_mats.iter().map(|m| m.data.as_slice()).collect();
        let k_refs: Vec<&[f32]> = k_mats.iter().map(|m| m.data.as_slice()).collect();
        let v_refs: Vec<&[f32]> = v_rows.iter().map(|r| r.as_slice()).collect();
        let contexts = rswa::verify_attention(&caches[layer], &q_refs, &k_refs, &v_refs)?;
        // 5. ONE `M=K` `o_proj` over the stacked per-position contexts.
        let ctx_refs: Vec<&[f32]> = contexts.iter().map(|m| m.row(0)).collect();
        let attn_rows = gemv_i8_batched(&ctx_refs, &cl.o_proj);
        // 6. Per-position residual, post-attn RMSNorm, MoE dispatch, residual.
        for (i, attn) in attn_rows.into_iter().enumerate() {
            let attn_out = Mat::from_vec(1, hidden, attn);
            let h = add_residual(&x[i], &attn_out)?;
            let normed2 = nn::rms_norm(&h, Some(&cl.post_attn_ln), eps)?;
            let mlp_out = Mat::from_vec(
                1,
                hidden,
                decode_mlp_i8(&cl.mlp, &normed2, calib::enabled().then_some(layer))?,
            );
            x[i] = add_residual(&h, &mlp_out)?;
        }
    }
    x.iter().map(|h| lm_head_cached_i8(wc, h)).collect()
}

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

    const MOE_POLICY_CHILD_ENV: &str = "FOCR_MOE_POLICY_TEST_CASE";

    fn moe_policy_child_case() -> Option<(&'static str, bool)> {
        match std::env::var(MOE_POLICY_CHILD_ENV).ok().as_deref() {
            None => None,
            Some("unset") => Some(("unset", false)),
            Some("zero") => Some(("zero", false)),
            Some("one") => Some(("one", true)),
            Some(other) => {
                eprintln!("unknown MoE policy subprocess case {other:?}");
                None
            }
        }
    }

    fn int8_moe_policy_fixture(
        rollback: bool,
    ) -> (Mat, Vec<f32>, Vec<[QLinear; 3]>, [QLinear; 3], moe::Routing) {
        const TARGETS: [f32; moe::config::NUM_EXPERTS_PER_TOK] =
            [16_777_216.0, 1.0, -16_777_216.0, 1.0, 1.0, 1.0];
        let fixture: serde_json::Value =
            serde_json::from_str(include_str!("../../tests/fixtures/moe_torch_2_10_cpu.json"))
                .expect("valid pinned torch MoE oracle fixture");
        let unique = &fixture["cases"][0];
        let scores = unique["scores_f32_bits"]
            .as_array()
            .expect("fixture score bits")
            .iter()
            .map(|bits| f32::from_bits(bits.as_u64().expect("fixture f32 bits") as u32))
            .collect::<Vec<_>>();

        let hidden_size = moe::config::HIDDEN_SIZE;
        let mut hidden_data = vec![0.0f32; hidden_size];
        hidden_data[0] = 1.0;
        let hidden = Mat::from_vec(1, hidden_size, hidden_data);
        let mut gate = vec![0.0f32; moe::config::N_ROUTED_EXPERTS * hidden_size];
        for (expert, score) in scores.into_iter().enumerate() {
            gate[expert * hidden_size] = score;
        }
        let routing = moe::route_default(&hidden, &gate).expect("public route_default succeeds");
        let expected_field = if rollback {
            "torch_sorted_indices"
        } else {
            "torch_unsorted_indices"
        };
        let expected = unique[expected_field]
            .as_array()
            .expect("fixture route indices")
            .iter()
            .map(|item| item.as_u64().expect("fixture expert index") as usize)
            .collect::<Vec<_>>();
        assert_eq!(routing.indices[0].as_slice(), expected.as_slice());

        let mut unit_weights = vec![0i8; hidden_size];
        unit_weights[0] = 127;
        let unit = QInt8::new(unit_weights, vec![1.0 / 127.0], 1, hidden_size);
        let zero = QInt8::new(vec![0; hidden_size], vec![1.0], 1, hidden_size);
        let silu_one = 1.0f32 / (1.0 + (-1.0f32).exp());
        let mut expert_outputs = vec![0.0f32; moe::config::N_ROUTED_EXPERTS];
        for (slot, &expert) in routing.indices[0].iter().enumerate() {
            expert_outputs[expert] = TARGETS[slot] / routing.weights[0][slot];
        }
        let experts = expert_outputs
            .into_iter()
            .map(|expert_output| {
                let down = if expert_output == 0.0 {
                    QInt8::new(vec![0; hidden_size], vec![1.0; hidden_size], hidden_size, 1)
                } else {
                    QInt8::new(
                        vec![127; hidden_size],
                        vec![expert_output / (127.0 * silu_one); hidden_size],
                        hidden_size,
                        1,
                    )
                };
                [
                    QLinear::I8(unit.clone()),
                    QLinear::I8(unit.clone()),
                    QLinear::I8(down),
                ]
            })
            .collect::<Vec<_>>();
        let shared = [
            QLinear::I8(zero.clone()),
            QLinear::I8(zero),
            QLinear::I8(QInt8::new(
                vec![0; hidden_size],
                vec![1.0; hidden_size],
                hidden_size,
                1,
            )),
        ];
        (hidden, gate, experts, shared, routing)
    }

    fn fold_int8_moe_contributions(
        contributions: &[Vec<f32>; moe::config::NUM_EXPERTS_PER_TOK],
        indices: &[usize; moe::config::NUM_EXPERTS_PER_TOK],
        ascending_expert: bool,
    ) -> Vec<f32> {
        let mut slots: [usize; moe::config::NUM_EXPERTS_PER_TOK] = std::array::from_fn(|slot| slot);
        if ascending_expert {
            slots.sort_unstable_by_key(|&slot| indices[slot]);
        }
        let mut out = vec![0.0f32; moe::config::HIDDEN_SIZE];
        for slot in slots {
            for (dst, &value) in out.iter_mut().zip(contributions[slot].iter()) {
                *dst += value;
            }
        }
        out
    }

    #[test]
    fn moe_policy_subprocess_probe_int8() -> FocrResult<()> {
        let Some((case, rollback)) = moe_policy_child_case() else {
            return Ok(());
        };
        let (hidden, gate, experts, shared, routing) = int8_moe_policy_fixture(rollback);
        let indices = routing.indices[0];
        let weights = routing.weights[0];

        let mut prefill_contributions: [Vec<f32>; moe::config::NUM_EXPERTS_PER_TOK] =
            std::array::from_fn(|_| vec![0.0; moe::config::HIDDEN_SIZE]);
        let mut decode_contributions: [Vec<f32>; moe::config::NUM_EXPERTS_PER_TOK] =
            std::array::from_fn(|_| vec![0.0; moe::config::HIDDEN_SIZE]);
        for slot in 0..moe::config::NUM_EXPERTS_PER_TOK {
            let expert = indices[slot];
            let prefill = expert_mlp_q(
                &hidden,
                &experts[expert][0],
                &experts[expert][1],
                &experts[expert][2],
                None,
            )?;
            let decode = expert_gemv_i8_serial(
                hidden.row(0),
                &experts[expert][0],
                &experts[expert][1],
                &experts[expert][2],
                None,
            );
            for channel in 0..moe::config::HIDDEN_SIZE {
                prefill_contributions[slot][channel] = weights[slot] * prefill.data[channel];
                decode_contributions[slot][channel] = weights[slot] * decode[channel];
            }
        }
        let expected_prefill =
            fold_int8_moe_contributions(&prefill_contributions, &indices, rollback);
        let alternate_prefill =
            fold_int8_moe_contributions(&prefill_contributions, &indices, !rollback);
        let expected_decode =
            fold_int8_moe_contributions(&decode_contributions, &indices, rollback);
        let alternate_decode =
            fold_int8_moe_contributions(&decode_contributions, &indices, !rollback);
        assert_ne!(
            expected_prefill[0].to_bits(),
            alternate_prefill[0].to_bits()
        );
        assert_ne!(expected_decode[0].to_bits(), alternate_decode[0].to_bits());

        let mut batched_data = hidden.data.clone();
        batched_data.extend_from_slice(&hidden.data);
        let batched_hidden = Mat::from_vec(2, moe::config::HIDDEN_SIZE, batched_data);
        let prefill = moe_block_i8(&batched_hidden, &gate, &experts, &shared, None)?;
        assert_eq!(prefill.row(0), expected_prefill.as_slice());
        assert_eq!(prefill.row(1), expected_prefill.as_slice());

        let mlp = CachedMlpI8::Moe {
            gate,
            experts,
            shared,
        };
        let decode = decode_mlp_i8(&mlp, &hidden, None)?;
        assert_eq!(decode, expected_decode);
        eprintln!("FOCR_MOE_POLICY_PROBE_INT8={case}");
        Ok(())
    }

    fn quant_oc_fixture_builder() -> crate::quant::focrq::FocrqBuilder {
        let arch = crate::native_engine::model_arch::arch_by_id("got-ocr2")
            .expect("got-ocr2 test architecture must be registered");
        crate::quant::focrq::FocrqBuilder::new()
            .with_model_id(arch.id())
            .with_license_notice(arch.license_notice())
    }

    #[test]
    fn quant_oc_loaded_rejects_prequantized_output_row_mismatch() {
        let mut builder = quant_oc_fixture_builder();
        builder
            .add_quantized(
                "w",
                crate::quant::focrq::WriteDType::QInt8PerChan,
                vec![3, 2],
                vec![0; 6],
                vec![0; 3 * std::mem::size_of::<f32>()],
                0,
                0,
            )
            .expect("valid synthetic QInt8 record");
        let weights = Weights::from_bytes(builder.build()).expect("synthetic artifact loads");

        let q = quant_oc_loaded(&weights, "w", 3).expect("matching output rows load");
        assert_eq!((q.n, q.k), (3, 2), "input/K dimension must be preserved");

        let err = quant_oc_loaded(&weights, "w", 2).expect_err("row mismatch must fail");
        assert!(matches!(err, FocrError::FormatMismatch(_)));
        assert!(format!("{err}").contains("3 output rows; expected 2"));
    }

    #[test]
    fn quant_oc_loaded_rejects_high_precision_output_row_mismatch_and_wrong_rank() {
        let mut builder = quant_oc_fixture_builder();
        builder
            .add_tensor(
                "matrix",
                crate::quant::focrq::WriteDType::Bf16,
                vec![3, 2],
                vec![0; 3 * 2 * std::mem::size_of::<u16>()],
            )
            .expect("valid synthetic BF16 matrix");
        builder
            .add_tensor(
                "vector",
                crate::quant::focrq::WriteDType::Bf16,
                vec![6],
                vec![0; 6 * std::mem::size_of::<u16>()],
            )
            .expect("valid synthetic BF16 vector");
        let weights = Weights::from_bytes(builder.build()).expect("synthetic artifact loads");

        let q = quant_oc_loaded(&weights, "matrix", 3).expect("matching matrix loads");
        assert_eq!((q.n, q.k), (3, 2), "input/K dimension must be preserved");

        let err = quant_oc_loaded(&weights, "matrix", 2).expect_err("row mismatch must fail");
        assert!(matches!(err, FocrError::FormatMismatch(_)));
        assert!(format!("{err}").contains("3 output rows; expected 2"));

        let err = quant_oc_loaded(&weights, "vector", 1)
            .expect_err("a rank-1 vector is not an [out, in] weight");
        assert!(matches!(err, FocrError::FormatMismatch(_)));
        assert!(format!("{err}").contains("rank 1; expected 2"));
    }

    // ── Packed int4 expert consumption (bd-4l71) ─────────────────────────────

    /// Deterministic synthetic [`QInt4`]: packed nibbles covering the full
    /// `[-8, 7]` domain and mixed-sign non-power-of-two group scales.
    fn synthetic_qint4(n: usize, k: usize, group_size: usize, salt: u64) -> QInt4 {
        let packed: Vec<u8> = (0..n * k / 2)
            .map(|i| ((i as u64).wrapping_mul(2654435761).wrapping_add(salt * 977) & 0xFF) as u8)
            .collect();
        let groups = k / group_size;
        let scales: Vec<f32> = (0..n * groups)
            .map(|i| {
                let raw = ((i as u64).wrapping_mul(40503).wrapping_add(salt) % 4096) as f32;
                (raw / 4096.0 - 0.5) * 1.7e-2 + 1.0e-3
            })
            .collect();
        QInt4 {
            packed: packed.into(),
            scales: scales.into(),
            n,
            k,
            group_size,
            tier: 1,
        }
    }

    /// Independent decoder-level oracle for the packed-int4 GEMV: unpack the
    /// nibbles with the module-external [`simd::int4::unpack_to_i8`] oracle,
    /// contract each group in i32, dequant with the group scale in ascending
    /// group order, then apply the activation scale — the exact op order
    /// [`gemv_i4`]/[`gemv_i4_serial`] pin.
    fn reference_gemv_i4(xq: &[i8], a_scale: f32, q: &QInt4) -> Vec<f32> {
        let b = simd::int4::unpack_to_i8(&q.packed, q.n, q.k);
        let scales = q.scales.to_vec();
        let groups = q.k / q.group_size;
        let mut y = vec![0.0f32; q.n];
        for o in 0..q.n {
            let mut acc_f = 0.0f32;
            for g in 0..groups {
                let mut acc_i: i32 = 0;
                for kk in g * q.group_size..(g + 1) * q.group_size {
                    acc_i += i32::from(xq[kk]) * i32::from(b[o * q.k + kk]);
                }
                acc_f += scales[o * groups + g] * acc_i as f32;
            }
            y[o] = acc_f * a_scale;
        }
        y
    }

    /// bd-4l71 acceptance: the decoder's packed-int4 GEMV plumbing (parallel
    /// block-partitioned [`gemv_i4`] AND single-call [`gemv_i4_serial`]) is
    /// BIT-EXACT against the unpack-then-contract oracle across the REAL
    /// Unlimited expert dims (routed gate/up `[896,1280]` g32, routed down
    /// `[1280,896]` g16, dense layer-0 gate/up `[6848,1280]` g32 and down
    /// `[1280,6848]` g16) plus an odd-block shape straddling `I8_GEMV_BLOCK`.
    #[test]
    fn gemv_i4_bitexact_vs_unpack_oracle_over_real_expert_dims() {
        let cases: [(usize, usize, usize); 5] = [
            (896, 1280, 32),
            (1280, 896, 16),
            (6848, 1280, 32),
            (1280, 6848, 16),
            (130, 96, 16), // n crosses a 64-row block boundary with an odd tail
        ];
        for (case, &(n, k, group)) in cases.iter().enumerate() {
            let q = synthetic_qint4(n, k, group, case as u64 + 1);
            let x: Vec<f32> = (0..k)
                .map(|i| (i as f32 * 0.37).sin() * 2.5 - 0.4)
                .collect();
            let (xq, a_scale) = quantize_row_i8(&x);

            let want = reference_gemv_i4(&xq, a_scale, &q);
            let parallel = gemv_i4(&x, &q);
            let serial = gemv_i4_serial(&xq, a_scale, &q);

            let bits = |s: &[f32]| s.iter().map(|f| f.to_bits()).collect::<Vec<u32>>();
            assert_eq!(
                bits(&parallel),
                bits(&want),
                "gemv_i4 != oracle for (n={n},k={k},g={group})"
            );
            assert_eq!(
                bits(&serial),
                bits(&want),
                "gemv_i4_serial != oracle for (n={n},k={k},g={group})"
            );
        }
    }

    /// The prefill m>1 packed-int4 linear must reproduce, row for row and
    /// BIT-EXACTLY, the m=1 serial GEMV over that row alone (each row carries
    /// its own ties-to-even activation scale) — the same losslessness contract
    /// the int8 batched paths pin. Also holds through the [`QLinear`] dispatch.
    #[test]
    fn linear_int4_dynamic_rows_bitexact_vs_serial_gemv() {
        let (n, k, group) = (96usize, 64usize, 16usize);
        let q = synthetic_qint4(n, k, group, 7);
        let m = 3usize;
        let x = Mat::from_vec(
            m,
            k,
            (0..m * k)
                .map(|i| (i as f32 * 0.11).cos() * 1.9 + 0.2)
                .collect(),
        );
        let out = linear_int4_dynamic(&x, &q).expect("shapes agree");
        assert_eq!(out.shape(), (m, n));
        for r in 0..m {
            let (xq, a_scale) = quantize_row_i8(x.row(r));
            let want = gemv_i4_serial(&xq, a_scale, &q);
            assert_eq!(
                out.row(r).iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
                want.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
                "prefill row {r} != m=1 serial GEMV"
            );
        }
        // QLinear dispatch reaches the same kernel.
        let via_dispatch =
            linear_q_dynamic(&x, &QLinear::I4(q.clone())).expect("dispatch shapes agree");
        assert_eq!(via_dispatch, out);

        // Shape mismatch is rejected, not silently mis-contracted.
        let bad = Mat::zeros(1, k + group);
        assert!(linear_int4_dynamic(&bad, &q).is_err());
    }

    /// End-to-end tensor-level parity (bd-4l71 acceptance): a full SwiGLU
    /// expert over PACKED int4 weights (real routed-expert dims, g16 and g32
    /// mixed per tensor) matches the `dequantize_int4` → f32 reference path
    /// within the tolerance the int8 path's analogous test uses
    /// (`nn::linear_int8_dynamic_approximates_f32`-style closeness, scaled to
    /// the output magnitude) — and the decode (serial) and prefill (parallel /
    /// m=1 batched) drivers agree BIT-EXACTLY with each other, because the
    /// integer contraction contract is shared.
    #[test]
    fn expert_forward_int4_matches_dequantized_f32_reference() {
        use crate::quant::int4::{dequantize_int4, pack_int4_f32};

        let (hidden, inter) = (1280usize, 896usize);
        let wgen = |len: usize, salt: u64| -> Vec<f32> {
            (0..len)
                .map(|i| {
                    let raw = ((i as u64)
                        .wrapping_mul(6364136223846793005)
                        .wrapping_add(salt)
                        >> 33) as u32;
                    (raw as f32 / u32::MAX as f32 - 0.5) * 0.08
                })
                .collect()
        };
        let gate_f = wgen(inter * hidden, 1);
        let up_f = wgen(inter * hidden, 2);
        let down_f = wgen(hidden * inter, 3);
        // Mixed group sizes per tensor: the recipe allows g16 OR g32 per tensor.
        let gate_q = pack_int4_f32(&gate_f, inter, hidden, 16);
        let up_q = pack_int4_f32(&up_f, inter, hidden, 32);
        let down_q = pack_int4_f32(&down_f, hidden, inter, 16);
        let to_qint4 = |q: &crate::quant::int4::QuantizedInt4| QInt4 {
            packed: q.packed.clone().into(),
            scales: q.scales.clone().into(),
            n: q.n,
            k: q.k,
            group_size: q.group_size,
            tier: 1,
        };
        let gate = QLinear::I4(to_qint4(&gate_q));
        let up = QLinear::I4(to_qint4(&up_q));
        let down = QLinear::I4(to_qint4(&down_q));

        let x: Vec<f32> = (0..hidden)
            .map(|i| (i as f32 * 0.019).sin() * 1.3 - 0.1)
            .collect();

        // Decode (serial, cross-expert-parallel shape) and internally-parallel
        // paths agree bit-exactly — the same integer contraction contract.
        let decode = expert_gemv_i8_serial(&x, &gate, &up, &down, None);
        let parallel = expert_gemv_i8(&x, &gate, &up, &down, None);
        assert_eq!(
            decode.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
            parallel.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
            "serial vs parallel int4 expert paths must be bit-identical"
        );
        // The prefill (m=1) driver row equals the decode row bit-exactly too.
        let xm = Mat::from_vec(1, hidden, x.clone());
        let prefill = expert_mlp_q(&xm, &gate, &up, &down, None).expect("prefill expert runs");
        assert_eq!(
            prefill
                .row(0)
                .iter()
                .map(|f| f.to_bits())
                .collect::<Vec<u32>>(),
            decode.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
            "prefill m=1 vs decode int4 expert paths must be bit-identical"
        );

        // f32 reference over the EXACT dequantized int4 grid values: the only
        // divergence left is the dynamic int8 activation quantization the int8
        // path also performs — the near-lossless closeness the int8 test pins.
        let gate_d = dequantize_int4(&gate_q);
        let up_d = dequantize_int4(&up_q);
        let down_d = dequantize_int4(&down_q);
        let dot = |w: &[f32], row: usize, x: &[f32]| -> f32 {
            let k = x.len();
            w[row * k..(row + 1) * k]
                .iter()
                .zip(x)
                .map(|(&a, &b)| a * b)
                .sum()
        };
        let mut act = vec![0.0f32; inter];
        for o in 0..inter {
            let g = dot(&gate_d, o, &x);
            let u = dot(&up_d, o, &x);
            act[o] = silu(g) * u;
        }
        let reference: Vec<f32> = (0..hidden).map(|o| dot(&down_d, o, &act)).collect();

        let amax = reference.iter().fold(0.0f32, |m, &v| m.max(v.abs()));
        assert!(amax > 0.0, "reference expert output must be non-degenerate");
        for (o, (&got, &want)) in decode.iter().zip(reference.iter()).enumerate() {
            assert!(
                (got - want).abs() <= 0.02 * amax,
                "int4 expert channel {o}: {got} vs f32 reference {want} (amax {amax})"
            );
        }
    }

    /// [`quant_ffn_loaded`] reads a `QInt4PerGroup` artifact tensor VERBATIM
    /// into the packed [`QLinear::I4`] arm (nibbles + group scales untouched —
    /// no unpack at residency), rejects an output-row mismatch, and still
    /// resolves qint8/bf16 tensors through the historical int8 arm.
    #[test]
    fn quant_ffn_loaded_keeps_int4_packed_and_dispatches_int8_unchanged() {
        let (n, k, group) = (4usize, 32usize, 16usize);
        let q = synthetic_qint4(n, k, group, 11);
        let mut builder = quant_oc_fixture_builder();
        builder
            .add_quantized(
                "w4",
                crate::quant::focrq::WriteDType::QInt4PerGroup,
                vec![n, k],
                q.packed.to_vec(),
                q.scales
                    .to_vec()
                    .iter()
                    .flat_map(|s| s.to_le_bytes())
                    .collect(),
                group,
                1,
            )
            .expect("valid synthetic QInt4 record");
        builder
            .add_quantized(
                "w8",
                crate::quant::focrq::WriteDType::QInt8PerChan,
                vec![3, 2],
                vec![0; 6],
                vec![0; 3 * std::mem::size_of::<f32>()],
                0,
                0,
            )
            .expect("valid synthetic QInt8 record");
        let weights = Weights::from_bytes(builder.build()).expect("synthetic artifact loads");

        let loaded = quant_ffn_loaded(&weights, "w4", n).expect("int4 tensor loads");
        assert_eq!((loaded.n(), loaded.k()), (n, k));
        let QLinear::I4(loaded_q) = &loaded else {
            panic!("QInt4PerGroup record must take the packed int4 arm");
        };
        assert_eq!(
            loaded_q.packed, q.packed,
            "nibbles must stay packed verbatim"
        );
        assert_eq!(loaded_q.scales, q.scales);
        assert_eq!(loaded_q.group_size, group);

        // The loaded tensor computes bit-identically to the directly built one.
        let x: Vec<f32> = (0..k).map(|i| (i as f32) * 0.21 - 2.0).collect();
        assert_eq!(gemv_q(&x, &loaded), gemv_i4(&x, &q));

        let err = quant_ffn_loaded(&weights, "w4", n + 1).expect_err("row mismatch must fail");
        assert!(matches!(err, FocrError::FormatMismatch(_)));
        assert!(format!("{err}").contains("output rows; expected"));

        let loaded8 = quant_ffn_loaded(&weights, "w8", 3).expect("int8 tensor loads");
        assert!(
            matches!(loaded8, QLinear::I8(_)),
            "QInt8PerChan must keep the historical int8 arm"
        );
    }

    /// Decoder-level plumbing over a synthetic [`Weights`]: a dense
    /// [`CachedMlpI8`] built from artifact int4 tensors decodes BIT-EXACTLY to
    /// the same expert computed over directly constructed packed weights.
    #[test]
    fn decode_mlp_over_artifact_int4_weights_is_bitexact() {
        let (hidden, inter, group) = (64usize, 48usize, 16usize);
        let gate = synthetic_qint4(inter, hidden, group, 21);
        let up = synthetic_qint4(inter, hidden, 32, 22);
        let down = synthetic_qint4(hidden, inter, group, 23);
        let mut builder = quant_oc_fixture_builder();
        for (name, q) in [("gate", &gate), ("up", &up), ("down", &down)] {
            builder
                .add_quantized(
                    name,
                    crate::quant::focrq::WriteDType::QInt4PerGroup,
                    vec![q.n, q.k],
                    q.packed.to_vec(),
                    q.scales
                        .to_vec()
                        .iter()
                        .flat_map(|s| s.to_le_bytes())
                        .collect(),
                    q.group_size,
                    1,
                )
                .expect("valid synthetic QInt4 record");
        }
        let weights = Weights::from_bytes(builder.build()).expect("synthetic artifact loads");
        let mlp = CachedMlpI8::Dense {
            gate: quant_ffn_loaded(&weights, "gate", inter).expect("gate loads"),
            up: quant_ffn_loaded(&weights, "up", inter).expect("up loads"),
            down: quant_ffn_loaded(&weights, "down", hidden).expect("down loads"),
        };

        let x = Mat::from_vec(
            1,
            hidden,
            (0..hidden).map(|i| (i as f32 * 0.31).sin() * 1.1).collect(),
        );
        let decoded = decode_mlp_i8(&mlp, &x, None).expect("decode over artifact int4 runs");
        let want = expert_gemv_i8(
            x.row(0),
            &QLinear::I4(gate),
            &QLinear::I4(up),
            &QLinear::I4(down),
            None,
        );
        assert_eq!(
            decoded.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
            want.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
            "artifact-loaded int4 dense MLP must match direct packed compute"
        );
        // Prefill driver over the same cache agrees with the decode row.
        let prefill = prefill_mlp_i8(&mlp, &x, None).expect("prefill over artifact int4 runs");
        assert_eq!(
            prefill
                .row(0)
                .iter()
                .map(|f| f.to_bits())
                .collect::<Vec<u32>>(),
            decoded.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
        );
    }

    #[test]
    fn activation_quantization_uses_true_division_and_ties_to_even() {
        let ties = [127.0, 0.5, 2.5, -0.5, -2.5];
        let (q, scale) = quantize_row_i8(&ties);
        assert_eq!(scale, 1.0);
        assert_eq!(q, [127, 0, 2, 0, -2]);

        // This f32 pair sits on a boundary where true division yields
        // -125.49999 while reciprocal multiplication yields exactly -125.5.
        // The contract requires the former, so this keeps the division rule live.
        let boundary = [f32::from_bits(0x3cdd_67c9), f32::from_bits(0xbcda_ca56)];
        let (q, scale) = quantize_row_i8(&boundary);
        assert_eq!(scale.to_bits(), 0x395f_2615);
        assert_eq!(q, [127, -125]);
        let reciprocal = (boundary[1] * (1.0 / scale)).round_ties_even() as i8;
        assert_eq!(
            reciprocal, -126,
            "test vector must distinguish the two formulas"
        );
    }

    // ── Fused q/k/v GEMV bit-identity (FOCR_QKV_FUSED, bd-241s) ──────────────────

    /// bd-2mo.3: a [`WeightLayout::SmmlaPanels`] weight (the offline `--arch
    /// aarch64-smmla` packing kept by the loader on an SMMLA host) produces
    /// BYTE-IDENTICAL f32 output to the same weight in row-major layout,
    /// through the REAL decode entry points: [`gemv_i8`] (fresh quantize),
    /// [`gemv_i8_prequant`], and the batched prequant GEMM. `n` deliberately
    /// exceeds `I8_GEMV_BLOCK` so the even-base panel slicing at block
    /// boundaries is exercised, and includes an odd tail block (n = 130).
    #[test]
    fn smmla_panel_layout_is_byte_identical_through_the_gemv_paths() {
        let (n, k) = (130usize, 96usize);
        let mut w = vec![0i8; n * k];
        for (i, v) in w.iter_mut().enumerate() {
            *v = (((i as i64 * 37 + 5) % 255) - 127) as i8;
        }
        let scales: Vec<f32> = (0..n).map(|o| 1.0e-3 + o as f32 * 3.0e-5).collect();
        let row_major = QInt8::new(w.clone(), scales.clone(), n, k);
        let (panels, _, _) = crate::simd::pack::smmla_pack_panels(&w, 0, n, k, k);
        let packed = QInt8::new_smmla_panels(panels, scales, n, k);

        // gemv_i8: fresh activation quantize inside.
        let x: Vec<f32> = (0..k).map(|i| (i as f32) * 0.11 - 4.0).collect();
        let a = gemv_i8(&x, &row_major);
        let b = gemv_i8(&x, &packed);
        assert_eq!(a, b, "gemv_i8 must be layout-invariant (bit-exact)");

        // gemv_i8_prequant: caller-provided int8 row.
        let xq: Vec<i8> = (0..k)
            .map(|i| (((i * 91) % 255) as i32 - 127) as i8)
            .collect();
        let a = gemv_i8_prequant(&xq, 0.017, &row_major);
        let b = gemv_i8_prequant(&xq, 0.017, &packed);
        assert_eq!(
            a, b,
            "gemv_i8_prequant must be layout-invariant (bit-exact)"
        );
        println!(r#"{{"check":"smmla_panel_gemv_parity","n":{n},"k":{k},"result":"pass"}}"#);
    }

    /// bd-2mo.3: [`fuse_qkv`] over three panel-layout weights produces the
    /// panel pack of the fused matrix — i.e. fusing THEN packing equals
    /// packing THEN fusing — and the fused panel weight GEMVs bit-identically
    /// to the fused row-major weight.
    #[test]
    fn fuse_qkv_concatenates_smmla_panels_losslessly() {
        let (qkv_dim, hidden) = (64usize, 48usize);
        let mk = |salt: i64| -> (Vec<i8>, Vec<f32>) {
            let mut w = vec![0i8; qkv_dim * hidden];
            for (i, v) in w.iter_mut().enumerate() {
                *v = (((i as i64 * 31 + salt * 101) % 255) - 127) as i8;
            }
            let scales: Vec<f32> = (0..qkv_dim)
                .map(|o| 1.0e-3 + (o as f32 + salt as f32 * 0.5) * 1.0e-4)
                .collect();
            (w, scales)
        };
        let (qw, qs) = mk(1);
        let (kw, ks) = mk(2);
        let (vw, vs) = mk(3);
        let pack = |w: &[i8]| crate::simd::pack::smmla_pack_panels(w, 0, qkv_dim, hidden, hidden).0;

        let fused_rm = fuse_qkv(
            &QInt8::new(qw.clone(), qs.clone(), qkv_dim, hidden),
            &QInt8::new(kw.clone(), ks.clone(), qkv_dim, hidden),
            &QInt8::new(vw.clone(), vs.clone(), qkv_dim, hidden),
        );
        let fused_pk = fuse_qkv(
            &QInt8::new_smmla_panels(pack(&qw), qs, qkv_dim, hidden),
            &QInt8::new_smmla_panels(pack(&kw), ks, qkv_dim, hidden),
            &QInt8::new_smmla_panels(pack(&vw), vs, qkv_dim, hidden),
        );
        assert_eq!(fused_pk.layout, WeightLayout::SmmlaPanels);
        // fuse-then-pack == pack-then-fuse (byte equality of the panel stream)
        let repacked = crate::simd::pack::smmla_pack_panels(
            &fused_rm.w,
            0,
            fused_rm.n,
            fused_rm.k,
            fused_rm.k,
        )
        .0;
        assert_eq!(
            &fused_pk.w[..],
            &repacked[..],
            "panel concat must equal pack of the fused matrix"
        );
        // and the GEMV agrees bitwise
        let x: Vec<f32> = (0..hidden).map(|i| (i as f32) * 0.07 - 1.5).collect();
        assert_eq!(
            gemv_i8(&x, &fused_rm),
            gemv_i8(&x, &fused_pk),
            "fused panel GEMV must be bit-exact"
        );
        println!(r#"{{"check":"smmla_panel_fuse_qkv","result":"pass"}}"#);
    }

    /// The fused `[3*qkv_dim, hidden]` GEMV must reproduce, BYTE-FOR-BYTE, the q/k/v
    /// that three separate [`gemv_i8`] calls produce. `qkv_dim` is deliberately NOT
    /// a multiple of [`I8_GEMV_BLOCK`] (64) so the fused output's rayon row-chunks
    /// straddle the q|k|v seams — proving each output row is an independent dot
    /// product whose value is invariant to how rows are grouped into tasks.
    #[test]
    fn fused_qkv_gemv_is_byte_identical_to_three_calls() {
        let qkv_dim = 70usize; // crosses a 64-row block boundary, not a multiple of 64
        let hidden = 96usize;
        // Deterministic int8 weights (full [-127,127] range) + per-channel scales.
        let mk = |salt: i32| -> QInt8 {
            let mut w = vec![0i8; qkv_dim * hidden];
            for o in 0..qkv_dim {
                for i in 0..hidden {
                    let raw = ((o as i32 * 31 + i as i32 * 7 + salt * 101) % 255) - 127;
                    w[o * hidden + i] = raw as i8;
                }
            }
            let scales: Vec<f32> = (0..qkv_dim)
                .map(|o| 1.0e-3 + (o as f32 + salt as f32 * 0.5) * 1.0e-4)
                .collect();
            QInt8::new(w, scales, qkv_dim, hidden)
        };
        let q_proj = mk(0);
        let k_proj = mk(1);
        let v_proj = mk(2);
        // Deterministic activation row spanning negatives/positives (exercises the
        // dynamic per-row quantize: amax, clamp, round).
        let nrow: Vec<f32> = (0..hidden)
            .map(|i| (i as f32 * 0.37).sin() * 2.5 - 0.4)
            .collect();

        // Default path: three separate GEMVs (each re-quantizes `nrow`).
        let q_sep = gemv_i8(&nrow, &q_proj);
        let k_sep = gemv_i8(&nrow, &k_proj);
        let v_sep = gemv_i8(&nrow, &v_proj);

        // Fused path: one stacked weight, one GEMV, sliced back.
        let fused = fuse_qkv(&q_proj, &k_proj, &v_proj);
        assert_eq!(fused.n, 3 * qkv_dim);
        assert_eq!(fused.k, hidden);
        let y = gemv_i8(&nrow, &fused);

        // Compare RAW BIT PATTERNS (not approximate) — this is a bit-exact lever.
        let bits = |s: &[f32]| s.iter().map(|f| f.to_bits()).collect::<Vec<u32>>();
        assert_eq!(bits(&q_sep), bits(&y[0..qkv_dim]), "q mismatch");
        assert_eq!(bits(&k_sep), bits(&y[qkv_dim..2 * qkv_dim]), "k mismatch");
        assert_eq!(
            bits(&v_sep),
            bits(&y[2 * qkv_dim..3 * qkv_dim]),
            "v mismatch"
        );
    }

    // ── Fused norm->quant epilogue bit-identity (FOCR_FUSE_NORM_QUANT, bd-1azu.54) ──

    /// Lever 1: folding the per-row int8 quantize into the RMSNorm epilogue and
    /// reusing that ONE int8 row across q/k/v must reproduce, BYTE-FOR-BYTE, the
    /// default `nn::rms_norm` -> f32 row -> per-`gemv_i8` re-quantize path — both
    /// the dequantized projection outputs AND the intermediate `(xq, a_scale)`.
    /// `qkv_dim` crosses a 64-row [`I8_GEMV_BLOCK`] boundary so the block fan-out is
    /// exercised; the input row spans negatives/positives (dynamic amax/clamp/round).
    #[test]
    fn fused_norm_quant_is_byte_identical_to_norm_then_gemv() {
        let hidden = 96usize;
        let qkv_dim = 70usize; // crosses a 64-row I8_GEMV_BLOCK boundary
        let mk = |salt: i32| -> QInt8 {
            let mut w = vec![0i8; qkv_dim * hidden];
            for o in 0..qkv_dim {
                for i in 0..hidden {
                    let raw = ((o as i32 * 29 + i as i32 * 11 + salt * 97) % 255) - 127;
                    w[o * hidden + i] = raw as i8;
                }
            }
            let scales: Vec<f32> = (0..qkv_dim)
                .map(|o| 1.0e-3 + (o as f32 + salt as f32 * 0.5) * 1.0e-4)
                .collect();
            QInt8::new(w, scales, qkv_dim, hidden)
        };
        let q_proj = mk(0);
        let k_proj = mk(1);
        let v_proj = mk(2);
        let ln_w: Vec<f32> = (0..hidden).map(|i| 0.5 + (i as f32) * 0.03).collect();
        let x = Mat::from_vec(
            1,
            hidden,
            (0..hidden)
                .map(|i| (i as f32 * 0.41).cos() * 3.0 - 0.7)
                .collect(),
        );
        let eps = config::RMS_NORM_EPS;
        let bits = |s: &[f32]| s.iter().map(|f| f.to_bits()).collect::<Vec<u32>>();

        // Default path: norm -> f32 row -> three re-quantizing `gemv_i8`.
        let normed = nn::rms_norm(&x, Some(&ln_w), eps).unwrap();
        let nrow = normed.row(0);
        let q_def = gemv_i8(nrow, &q_proj);
        let k_def = gemv_i8(nrow, &k_proj);
        let v_def = gemv_i8(nrow, &v_proj);

        // Fused path: quantize ONCE in the norm epilogue, reuse for q/k/v.
        let (xq, a_scale) = rms_norm_quant_i8(&x, Some(&ln_w), eps).unwrap();
        let q_fused = gemv_i8_prequant(&xq, a_scale, &q_proj);
        let k_fused = gemv_i8_prequant(&xq, a_scale, &k_proj);
        let v_fused = gemv_i8_prequant(&xq, a_scale, &v_proj);

        assert_eq!(bits(&q_def), bits(&q_fused), "q mismatch");
        assert_eq!(bits(&k_def), bits(&k_fused), "k mismatch");
        assert_eq!(bits(&v_def), bits(&v_fused), "v mismatch");

        // The fused `(xq, a_scale)` themselves equal the separate norm-then-quantize.
        let (xq_ref, a_ref) = quantize_row_i8(nrow);
        assert_eq!(xq, xq_ref, "int8 activation bytes mismatch");
        assert_eq!(a_scale.to_bits(), a_ref.to_bits(), "a_scale bits mismatch");
    }

    // ── Fused ngram->lm_head epilogue bit-identity (FOCR_FUSE_NGRAM_LMHEAD, .54) ──

    /// Lever 3: folding the sliding-window no-repeat-ngram ban into the int8 lm_head
    /// dequant epilogue must reproduce, BYTE-FOR-BYTE, the default `gemv_i8` then
    /// the sampler's `masked_sliding_window_logits_if_needed` copy-mask. `vocab`
    /// crosses several [`I8_GEMV_BLOCK`] (64) boundaries. The sequence `[3,7,3]`
    /// bans the bigram completion `7` (current_prefix `[3]`, bigram `(3,7)` in
    /// window) — so the masked path is genuinely exercised — and the empty-ban case
    /// must be byte-identical to the plain head.
    #[test]
    fn fused_ngram_lmhead_is_byte_identical_to_separate_mask() {
        use super::super::sampler;
        let vocab = 130usize;
        let hidden = 48usize;
        let mut w = vec![0i8; vocab * hidden];
        for o in 0..vocab {
            for i in 0..hidden {
                let raw = ((o as i32 * 13 + i as i32 * 7 + 5) % 255) - 127;
                w[o * hidden + i] = raw as i8;
            }
        }
        let scales: Vec<f32> = (0..vocab).map(|o| 1.0e-3 + o as f32 * 1.0e-4).collect();
        let qw = QInt8::new(w, scales, vocab, hidden);
        let x: Vec<f32> = (0..hidden)
            .map(|i| (i as f32 * 0.31).sin() * 2.0 - 0.3)
            .collect();
        let bits = |s: &[f32]| s.iter().map(|f| f.to_bits()).collect::<Vec<u32>>();

        let ngram_size = 2usize;
        let window = 16usize;
        let seq: Vec<u32> = vec![3, 7, 3];

        // Default path: full head, then the sampler's copy-then-mask.
        let logits = gemv_i8(&x, &qw);
        let expected =
            sampler::masked_sliding_window_logits_if_needed(&logits, &seq, ngram_size, window, &[])
                .unwrap_or_else(|| logits.clone());

        // Fused path: ban set folded into the lm_head dequant epilogue.
        let banned =
            sampler::collect_sliding_window_ngram_bans(&seq, ngram_size, window, &[], vocab);
        let fused = gemv_i8_ngram_masked(&x, &qw, &banned);

        assert!(
            !banned.is_empty(),
            "test sequence should ban at least one token"
        );
        // token 7 is the bigram completion the sequence bans → its logit is -inf.
        assert_eq!(
            fused[7].to_bits(),
            f32::NEG_INFINITY.to_bits(),
            "banned token 7 should be masked to -inf"
        );
        assert_eq!(bits(&fused), bits(&expected), "masked logits row mismatch");

        // No-ban case: empty bans ⇒ byte-for-byte the plain head.
        let none = gemv_i8_ngram_masked(&x, &qw, &[]);
        assert_eq!(bits(&none), bits(&logits), "no-ban head must equal gemv_i8");
    }

    // ── Batched M=B projection GEMV bit-identity (bd-1azu.3) ─────────────────────

    /// [`gemv_i8_batched`] row `r` must reproduce, BYTE-FOR-BYTE, what the
    /// single-row [`gemv_i8`] produces for that activation row — the lossless claim
    /// the batched decode spine rests on. Exercises the REAL functions (not a
    /// replicated model) over decode-shaped panels: an `o_proj`-like square, a
    /// fused-qkv-like non-64-multiple width, and a wide panel whose `n` crosses
    /// many [`I8_GEMV_BLOCK`] (64) boundaries and is not a multiple of 64.
    #[test]
    fn batched_gemv_i8_is_byte_identical_to_per_row() {
        let mk = |n: usize, k: usize, salt: i32| -> QInt8 {
            let mut w = vec![0i8; n * k];
            for o in 0..n {
                for i in 0..k {
                    let raw = ((o as i32 * 17 + i as i32 * 5 + salt * 101) % 255) - 127;
                    w[o * k + i] = raw as i8;
                }
            }
            let scales: Vec<f32> = (0..n)
                .map(|o| 1.0e-3 + (o as f32 + salt as f32 * 0.5) * 1.0e-4)
                .collect();
            QInt8::new(w, scales, n, k)
        };
        let mkrow = |k: usize, seed: f32| -> Vec<f32> {
            (0..k)
                .map(|i| ((i as f32 + seed) * 0.37).sin() * 2.5 - 0.4)
                .collect()
        };
        let bits = |s: &[f32]| s.iter().map(|f| f.to_bits()).collect::<Vec<u32>>();
        for &(n, k) in &[(96usize, 96usize), (3 * 70, 96), (130, 70)] {
            let w = mk(n, k, 3);
            for b in 1..=3usize {
                let rows_owned: Vec<Vec<f32>> = (0..b).map(|s| mkrow(k, s as f32 * 1.7)).collect();
                let rows: Vec<&[f32]> = rows_owned.iter().map(|r| r.as_slice()).collect();
                let batched = gemv_i8_batched(&rows, &w);
                assert_eq!(batched.len(), b);
                for r in 0..b {
                    let single = gemv_i8(&rows_owned[r], &w);
                    assert_eq!(
                        bits(&single),
                        bits(&batched[r]),
                        "batched row {r} != per-row gemv_i8 (n={n} k={k} b={b})"
                    );
                }
            }
        }
    }

    /// The batched env kill-switch defaults OFF and the size cap defaults to a
    /// positive value when the environment is unset — the spine is additive,
    /// default-OFF (Doctrine #3). Guards the kill-switch assertion behind the env
    /// so a host that exports `FOCR_BATCH_SPINE` does not flake.
    #[test]
    fn batch_spine_defaults_off_with_positive_cap() {
        if std::env::var_os("FOCR_BATCH_SPINE").is_none() {
            assert!(!batch_spine_enabled());
        }
        assert!(batch_size_cap() >= 1);
    }

    // ── Chunked prefill kill-switch + tiling (FOCR_PREFILL_CHUNK, bd-1azu.9) ──────

    /// `FOCR_PREFILL_CHUNK` defaults to the monolithic schedule (`None`) when the
    /// environment is unset — the lever is additive, default-OFF (Doctrine #3).
    /// Guarded behind the env presence so a host that exports it does not flake.
    #[test]
    fn prefill_chunk_defaults_off() {
        if std::env::var_os("FOCR_PREFILL_CHUNK").is_none() {
            assert_eq!(prefill_chunk_size(), None);
        }
    }

    /// [`prefill_chunk_bounds`] tiles `[0, seq)` into CONTIGUOUS, gap-free,
    /// ascending `[c0, c1)` chunks of at most `chunk` tokens, covering exactly
    /// `[0, seq)` for every chunk size — including `1`, primes, sizes that do not
    /// divide `seq`, and sizes `>= seq` (a single chunk). This ascending coverage
    /// is what lets chunk `g` attend exactly the keys earlier chunks wrote.
    #[test]
    fn prefill_chunk_bounds_cover_contiguously_in_order() {
        for &seq in &[0usize, 1, 2, 5, 7, 16, 17] {
            for &chunk in &[1usize, 2, 3, 5, 7, 16, 64] {
                let bounds = prefill_chunk_bounds(seq, chunk);
                let mut cursor = 0usize;
                for &(c0, c1) in &bounds {
                    assert_eq!(c0, cursor, "seq={seq} chunk={chunk}: gap/overlap");
                    assert!(c1 > c0, "seq={seq} chunk={chunk}: empty/reversed chunk");
                    assert!(c1 - c0 <= chunk, "seq={seq} chunk={chunk}: chunk too wide");
                    assert!(c1 <= seq, "seq={seq} chunk={chunk}: chunk past seq");
                    cursor = c1;
                }
                assert_eq!(cursor, seq, "seq={seq} chunk={chunk}: must cover [0, seq)");
            }
        }
    }

    /// The assembled chunked attention is BYTE-FOR-BYTE the monolithic
    /// [`prefill_attention`] over the whole sequence, for every chunk size — the
    /// crux of the chunked-prefill lossless claim, exercised on the REAL
    /// [`chunk_prefill_attention`] kernel at a small synthetic head shape (a model
    /// is not needed: attention is shape-parametrized).
    #[test]
    fn chunked_attention_is_byte_identical_to_monolithic() {
        let (num_heads, head_dim, seq) = (3usize, 4usize, 11usize);
        let dim = num_heads * head_dim;
        // Deterministic synthetic q/k/v spanning negatives/positives.
        let mk = |salt: f32| -> Mat {
            let data: Vec<f32> = (0..seq * dim)
                .map(|i| ((i as f32 + salt) * 0.37).sin() * 1.7 - 0.2)
                .collect();
            Mat::from_vec(seq, dim, data)
        };
        let (q, k, v) = (mk(0.0), mk(11.0), mk(23.0));
        let monolithic = prefill_attention(&q, &k, &v, num_heads, head_dim).unwrap();
        let bits = |s: &[f32]| s.iter().map(|f| f.to_bits()).collect::<Vec<u32>>();
        // chunk 1, 2, a prime (3), and the full length (single chunk).
        for &chunk in &[1usize, 2, 3, seq] {
            let mut assembled = Mat::zeros(seq, dim);
            for (c0, c1) in prefill_chunk_bounds(seq, chunk) {
                let q_chunk = Mat::from_vec(c1 - c0, dim, q.data[c0 * dim..c1 * dim].to_vec());
                let kpre = Mat::from_vec(c1, dim, k.data[..c1 * dim].to_vec());
                let vpre = Mat::from_vec(c1, dim, v.data[..c1 * dim].to_vec());
                let ctx = chunk_prefill_attention(&q_chunk, &kpre, &vpre, num_heads, head_dim, c0)
                    .unwrap();
                assembled.data[c0 * dim..c1 * dim].copy_from_slice(&ctx.data);
            }
            assert_eq!(
                bits(&monolithic.data),
                bits(&assembled.data),
                "chunk={chunk}: chunked attention != monolithic prefill_attention"
            );
        }
    }

    // ── Vocab-tiled lm_head bit-identity (FOCR_LMHEAD_SHARD, bd-1azu.25) ──────────

    /// [`vocab_tile_ranges`] partitions `[0, n)` into CONTIGUOUS, gap-free,
    /// ascending tiles for every tile count — including counts that do not divide
    /// `n` and counts larger than `n` (which must clamp, never emit out-of-bounds
    /// or empty-then-restart ranges). The ascending coverage is what preserves the
    /// vocab column order under sharding.
    #[test]
    fn vocab_tile_ranges_cover_contiguously_in_order() {
        for &n in &[0usize, 1, 7, 64, 65, 129_280] {
            for &tiles in &[1usize, 2, 3, 7, 16, 1000, 200_000] {
                let ranges = vocab_tile_ranges(n, tiles);
                // Contiguous from 0, gap-free, non-overlapping, covering [0, n).
                let mut cursor = 0usize;
                for &(start, end) in &ranges {
                    assert_eq!(start, cursor, "n={n} tiles={tiles}: tile gap/overlap");
                    assert!(end >= start, "n={n} tiles={tiles}: reversed tile");
                    assert!(end <= n, "n={n} tiles={tiles}: tile end out of bounds");
                    cursor = end;
                }
                assert_eq!(cursor, n, "n={n} tiles={tiles}: ranges must cover [0, n)");
                // Tile count is clamped into `[1, n.max(1)]`.
                assert!(!ranges.is_empty(), "n={n} tiles={tiles}: at least one tile");
                assert!(
                    ranges.len() <= n.max(1),
                    "n={n} tiles={tiles}: more tiles than channels"
                );
            }
        }
    }

    /// [`gemv_i8_sharded`] must reproduce, BYTE-FOR-BYTE, what the monolithic
    /// [`gemv_i8`] produces — for several contiguous-tile counts INCLUDING ones
    /// that do not evenly divide the vocab and one larger than a 64-row block.
    /// `n` is a deliberately awkward small synthetic vocab (not a multiple of
    /// [`I8_GEMV_BLOCK`]) so tile and 64-row chunk boundaries straddle, proving
    /// each logit is an independent dot whose value is invariant to grouping.
    #[test]
    fn lmhead_shard_i8_is_byte_identical_to_monolithic() {
        let mk = |n: usize, k: usize, salt: i32| -> QInt8 {
            let mut w = vec![0i8; n * k];
            for o in 0..n {
                for i in 0..k {
                    let raw = ((o as i32 * 13 + i as i32 * 7 + salt * 101) % 255) - 127;
                    w[o * k + i] = raw as i8;
                }
            }
            let scales: Vec<f32> = (0..n)
                .map(|o| 1.0e-3 + (o as f32 + salt as f32 * 0.5) * 1.0e-4)
                .collect();
            QInt8::new(w, scales, n, k)
        };
        let bits = |s: &[f32]| s.iter().map(|f| f.to_bits()).collect::<Vec<u32>>();
        // Small synthetic vocab `n=150` (crosses two 64-blocks, not a multiple of 64).
        let (n, k) = (150usize, 96usize);
        let qw = mk(n, k, 3);
        let x: Vec<f32> = (0..k)
            .map(|i| (i as f32 * 0.41).sin() * 2.5 - 0.4)
            .collect();
        let monolithic = gemv_i8(&x, &qw);
        // 7 does not divide 150; 13 does not divide 150; 1 is the trivial tile;
        // 200 > n exercises the clamp; 150 is one-column-per-tile.
        for &tiles in &[1usize, 2, 7, 13, 64, 150, 200] {
            let sharded = gemv_i8_sharded(&x, &qw, tiles);
            assert_eq!(
                bits(&monolithic),
                bits(&sharded),
                "int8 lm_head: {tiles}-tile shard != monolithic gemv_i8 (n={n} k={k})"
            );
        }
    }

    /// [`gemv_sharded`] must reproduce, BYTE-FOR-BYTE, what the monolithic
    /// [`gemv`] produces — the f32-head twin of the int8 check, across non-dividing
    /// tile counts on a small synthetic vocab.
    #[test]
    fn lmhead_shard_f32_is_byte_identical_to_monolithic() {
        let (n, k) = (150usize, 96usize);
        let w: Vec<f32> = (0..n * k)
            .map(|idx| ((idx as f32 * 0.013).sin() * 1.7) - 0.3)
            .collect();
        let x: Vec<f32> = (0..k)
            .map(|i| (i as f32 * 0.29).cos() * 2.1 + 0.2)
            .collect();
        let bits = |s: &[f32]| s.iter().map(|f| f.to_bits()).collect::<Vec<u32>>();
        let monolithic = gemv(&x, &w, n, k);
        for &tiles in &[1usize, 2, 7, 13, 64, 150, 200] {
            let sharded = gemv_sharded(&x, &w, n, k, tiles);
            assert_eq!(
                bits(&monolithic),
                bits(&sharded),
                "f32 lm_head: {tiles}-tile shard != monolithic gemv (n={n} k={k})"
            );
        }
    }

    /// The vocab-shard env kill-switch defaults OFF and the tile count defaults to
    /// a positive value when the environment is unset — an additive, default-OFF
    /// lever (Doctrine #3). Guarded behind the env so a host that exports
    /// `FOCR_LMHEAD_SHARD` does not flake.
    #[test]
    fn lmhead_shard_defaults_off_with_positive_tiles() {
        if std::env::var_os("FOCR_LMHEAD_SHARD").is_none() {
            assert!(!lmhead_shard_enabled());
        }
        assert!(lmhead_shard_tiles() >= 1);
    }

    fn assert_err_contains<T>(res: FocrResult<T>, needle: &str) {
        let message = match res {
            Ok(_) => String::from("<ok>"),
            Err(err) => err.to_string(),
        };
        assert!(
            message.contains(needle),
            "error {message:?} did not contain {needle:?}"
        );
    }

    // ── R-SWA cache repack ([SPEC-091]) ─────────────────────────────────────────

    #[test]
    fn token_major_to_head_major_transposes_kv() -> FocrResult<()> {
        // seq=2, num_heads=2, head_dim=2 -> token-major row r = [h0d0,h0d1,h1d0,h1d1].
        // row0 = [0,1, 2,3]  row1 = [4,5, 6,7]
        let k = Mat::from_vec(2, 4, vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]);
        let v = Mat::from_vec(2, 4, vec![10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0]);
        let (kh, vh) = token_major_to_head_major(&k, &v, 2, 2, 2)?;
        // head-major [num_heads, seq, head_dim]: head0 = rows' first head across seq
        //   head0: s0=[0,1] s1=[4,5] ; head1: s0=[2,3] s1=[6,7]
        assert_eq!(kh, vec![0.0, 1.0, 4.0, 5.0, 2.0, 3.0, 6.0, 7.0]);
        assert_eq!(vh, vec![10.0, 11.0, 14.0, 15.0, 12.0, 13.0, 16.0, 17.0]);
        Ok(())
    }

    // ── Token embedding ([SPEC-070]) ────────────────────────────────────────

    #[test]
    fn embed_tokens_gathers_rows() -> FocrResult<()> {
        // vocab=3, hidden=2 table: row0=[0,1] row1=[2,3] row2=[4,5]
        let table = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0];
        let out = embed_tokens(&table, 3, 2, &[2, 0, 1, 2])?;
        assert_eq!(out.shape(), (4, 2));
        assert_eq!(out.row(0), &[4.0, 5.0]);
        assert_eq!(out.row(1), &[0.0, 1.0]);
        assert_eq!(out.row(2), &[2.0, 3.0]);
        assert_eq!(out.row(3), &[4.0, 5.0]);
        Ok(())
    }

    #[test]
    fn embed_tokens_rejects_oob_id() {
        let table = vec![0.0, 1.0, 2.0, 3.0];
        assert!(embed_tokens(&table, 2, 2, &[5]).is_err());
    }

    #[test]
    fn embed_tokens_rejects_bad_table_len() {
        let table = vec![0.0, 1.0, 2.0];
        assert!(embed_tokens(&table, 2, 2, &[0]).is_err());
    }

    #[test]
    fn embed_tokens_rejects_table_shape_product_overflow_without_panic() {
        assert_err_contains(embed_tokens(&[], usize::MAX, 2, &[]), "vocab*hidden");
    }

    #[test]
    fn embed_tokens_rejects_output_shape_product_overflow_without_panic() {
        assert_err_contains(embed_tokens(&[], 0, usize::MAX, &[0, 1]), "seq*hidden");
    }

    // ── RoPE ([SPEC-078]) ───────────────────────────────────────────────────

    #[test]
    fn rope_position_zero_is_identity() -> FocrResult<()> {
        // head_dim=4, single head; pos 0 => cos=1, sin=0 => x unchanged.
        let rope = RopeTable::build(&[0], 4, 10000.0);
        let mut x = Mat::from_vec(1, 4, vec![1.0, 2.0, 3.0, 4.0]);
        apply_rope(&mut x, &rope)?;
        for (got, want) in x.data.iter().zip([1.0f32, 2.0, 3.0, 4.0].iter()) {
            assert!((got - want).abs() < 1e-6, "{got} != {want}");
        }
        Ok(())
    }

    #[test]
    fn rope_matches_hand_computed_rotate_half() -> FocrResult<()> {
        // head_dim=2 => half=1, inv_freq[0] = 10000^0 = 1, so angle = pos.
        // rotate_half([a,b]) = [-b, a].
        // out[0] = a*cos - b*sin ; out[1] = b*cos + a*sin.
        let theta = 10000.0f32;
        let pos = 1usize;
        let rope = RopeTable::build(&[pos], 2, theta);
        let (a, b) = (3.0f32, 5.0f32);
        let mut x = Mat::from_vec(1, 2, vec![a, b]);
        apply_rope(&mut x, &rope)?;
        let (s, c) = (pos as f32).sin_cos();
        let want0 = a * c - b * s;
        let want1 = b * c + a * s;
        assert!((x.data[0] - want0).abs() < 1e-5, "{} != {want0}", x.data[0]);
        assert!((x.data[1] - want1).abs() < 1e-5, "{} != {want1}", x.data[1]);
        Ok(())
    }

    #[test]
    fn rope_preserves_norm_per_head() -> FocrResult<()> {
        // A rotation preserves the L2 norm of each head block.
        let rope = RopeTable::build(&[7, 13], 4, 10000.0);
        let mut x = Mat::from_vec(2, 8, (0..16).map(|i| (i as f32) * 0.25 - 2.0).collect());
        let head_norms = |data: &[f32]| -> Vec<f32> {
            let mut out = Vec::with_capacity(4);
            for t in 0..2 {
                for h in 0..2 {
                    let r = &data[t * 8 + h * 4..t * 8 + h * 4 + 4];
                    out.push(r.iter().map(|v| v * v).sum::<f32>());
                }
            }
            out
        };
        let before = head_norms(&x.data);
        apply_rope(&mut x, &rope)?;
        let after = head_norms(&x.data);
        for (b, a) in before.iter().zip(after.iter()) {
            assert!((b - a).abs() < 1e-4, "norm changed: {b} -> {a}");
        }
        Ok(())
    }

    #[test]
    fn rope_two_heads_share_phase() -> FocrResult<()> {
        // head_dim=2, two heads packed in one row -> each head rotated by the
        // SAME phase (RoPE is per-position, shared across heads).
        let rope = RopeTable::build(&[2], 2, 10000.0);
        let mut x = Mat::from_vec(1, 4, vec![1.0, 0.0, 0.0, 1.0]);
        apply_rope(&mut x, &rope)?;
        let (s, c) = (2.0f32).sin_cos();
        // head0 = [1,0] -> [c, s]; head1 = [0,1] -> [-s, c]
        assert!((x.data[0] - c).abs() < 1e-5);
        assert!((x.data[1] - s).abs() < 1e-5);
        assert!((x.data[2] - (-s)).abs() < 1e-5);
        assert!((x.data[3] - c).abs() < 1e-5);
        Ok(())
    }

    #[test]
    fn rope_rejects_bad_shape() {
        let rope = RopeTable::build(&[0, 1], 4, 10000.0);
        // cols (3) not a multiple of head_dim (4)
        let mut bad = Mat::zeros(2, 3);
        assert!(apply_rope(&mut bad, &rope).is_err());
        // wrong number of rows vs positions
        let mut wrong_rows = Mat::zeros(3, 4);
        assert!(apply_rope(&mut wrong_rows, &rope).is_err());
    }

    #[test]
    fn rope_rejects_malformed_backing_data_without_panic() {
        let rope = RopeTable::build(&[0], 4, 10000.0);
        let mut bad = Mat {
            rows: 1,
            cols: 4,
            data: vec![1.0, 2.0, 3.0],
        };
        assert_err_contains(apply_rope(&mut bad, &rope), "apply_rope x: data len 3");
    }

    #[test]
    #[should_panic(expected = "RopeTable: seq*head_dim overflow")]
    fn rope_table_rejects_shape_product_overflow_before_allocating() {
        let _ = RopeTable::build(&[0, 1], usize::MAX - 1, 10000.0);
    }

    #[test]
    fn rope_table_empty_sequence_does_not_allocate_by_head_dim() {
        let rope = RopeTable::build(&[], usize::MAX - 1, 10000.0);
        assert_eq!(rope.head_dim, usize::MAX - 1);
        assert!(rope.cos.is_empty());
        assert!(rope.sin.is_empty());
    }

    // ── Dense SwiGLU MLP ([SPEC-075]) ───────────────────────────────────────

    #[test]
    fn dense_mlp_matches_hand_computed() -> FocrResult<()> {
        // hidden=2, inter=2, single token x=[1,0].
        // gate_w=[[1,0],[0,1]] => gate = [1, 0]
        // up_w  =[[1,1],[1,1]] => up   = [1, 1]
        // silu(gate) = [silu(1), silu(0)] = [0.7310586, 0]
        // silu(gate)*up = [0.7310586, 0]
        // down_w=[[1,0],[0,1]] => out = [0.7310586, 0]
        let x = Mat::from_vec(1, 2, vec![1.0, 0.0]);
        let gate_w = vec![1.0, 0.0, 0.0, 1.0];
        let up_w = vec![1.0, 1.0, 1.0, 1.0];
        let down_w = vec![1.0, 0.0, 0.0, 1.0];
        let out = dense_mlp(&x, &gate_w, &up_w, &down_w, 2, 2)?;
        assert_eq!(out.shape(), (1, 2));
        let silu1 = 1.0f32 / (1.0 + (-1.0f32).exp());
        assert!((out.data[0] - silu1).abs() < 1e-5, "{}", out.data[0]);
        assert!(out.data[1].abs() < 1e-6, "{}", out.data[1]);
        Ok(())
    }

    #[test]
    fn dense_mlp_rejects_bad_hidden() {
        let x = Mat::from_vec(1, 3, vec![1.0, 0.0, 0.0]);
        assert!(dense_mlp(&x, &[1.0, 0.0], &[1.0, 0.0], &[1.0, 0.0], 2, 1).is_err());
    }

    // ── linear_no_bias (the F.linear [out,in] transpose) ────────────────────

    #[test]
    fn linear_no_bias_transposes_pytorch_layout() -> FocrResult<()> {
        // x=[1,2] (1x2 row); w=[out=3, in=2] row-major:
        //   w = [[1,0],[0,1],[1,1]]  => y = [x·w0, x·w1, x·w2] = [1, 2, 3]
        let x = Mat::from_vec(1, 2, vec![1.0, 2.0]);
        let w = vec![1.0, 0.0, 0.0, 1.0, 1.0, 1.0];
        let y = linear_no_bias(&x, &w, 2, 3)?;
        assert_eq!(y.shape(), (1, 3));
        assert_eq!(y.data, vec![1.0, 2.0, 3.0]);
        Ok(())
    }

    #[test]
    fn linear_no_bias_transposes_multirow_nonsquare_layout() -> FocrResult<()> {
        let x = Mat::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let w = vec![
            1.0, 0.0, 0.0, // row 0 selects x0
            0.0, 1.0, 0.0, // row 1 selects x1
            0.0, 0.0, 1.0, // row 2 selects x2
            1.0, 1.0, 1.0, // row 3 sums the input row
        ];
        let y = linear_no_bias(&x, &w, 3, 4)?;
        assert_eq!(y.shape(), (2, 4));
        assert_eq!(y.data, vec![1.0, 2.0, 3.0, 6.0, 4.0, 5.0, 6.0, 15.0]);
        Ok(())
    }

    #[test]
    fn linear_no_bias_rejects_weight_shape_product_overflow_without_panic() {
        let x = Mat::zeros(1, 2);
        assert_err_contains(linear_no_bias(&x, &[], 2, usize::MAX), "out*in");
    }

    #[test]
    fn linear_no_bias_rejects_malformed_input_backing_data_without_panic() {
        let x = Mat {
            rows: 1,
            cols: 2,
            data: vec![1.0],
        };
        assert_err_contains(
            linear_no_bias(&x, &[1.0, 0.0], 2, 1),
            "linear_no_bias x: data len 1",
        );
    }

    // ── Residual add ([SPEC-072]) ───────────────────────────────────────────

    #[test]
    fn add_residual_sums_elementwise() -> FocrResult<()> {
        let a = Mat::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]);
        let b = Mat::from_vec(2, 2, vec![10.0, 20.0, 30.0, 40.0]);
        let c = add_residual(&a, &b)?;
        assert_eq!(c.data, vec![11.0, 22.0, 33.0, 44.0]);
        Ok(())
    }

    #[test]
    fn add_residual_rejects_shape_mismatch() {
        let a = Mat::zeros(2, 2);
        let b = Mat::zeros(2, 3);
        assert!(add_residual(&a, &b).is_err());
    }

    #[test]
    fn add_residual_rejects_malformed_backing_data_without_panic() {
        let a = Mat {
            rows: 1,
            cols: 2,
            data: vec![1.0],
        };
        let b = Mat::zeros(1, 2);
        assert_err_contains(add_residual(&a, &b), "add_residual lhs: data len 1");
    }

    // ── Final norm + lm_head ([SPEC-071]/[SPEC-081]) ────────────────────────

    #[test]
    fn lm_head_proj_matches_matmul() -> FocrResult<()> {
        // hidden=2 vocab=3; head_w=[vocab,hidden] = [[1,0],[0,1],[1,1]]
        // hidden state [3,4] -> logits [3, 4, 7]
        let h = Mat::from_vec(1, 2, vec![3.0, 4.0]);
        let head_w = vec![1.0, 0.0, 0.0, 1.0, 1.0, 1.0];
        let logits = lm_head_proj(&h, &head_w, 3)?;
        assert_eq!(logits.shape(), (1, 3));
        assert_eq!(logits.data, vec![3.0, 4.0, 7.0]);
        Ok(())
    }

    #[test]
    fn norm_and_lm_head_composes_rmsnorm_then_head() -> FocrResult<()> {
        // hidden=2, vocab=2. x=[3,4], norm_w=[1,1], eps=0.
        // rmsnorm: mean(x^2)=12.5, rstd=1/sqrt(12.5); normed=[3*rstd,4*rstd]
        // head_w = identity [[1,0],[0,1]] => logits = normed
        let h = Mat::from_vec(1, 2, vec![3.0, 4.0]);
        let norm_w = vec![1.0, 1.0];
        let head_w = vec![1.0, 0.0, 0.0, 1.0];
        let logits = norm_and_lm_head(&h, &norm_w, &head_w, 2, 0.0)?;
        let rstd = 1.0f32 / 12.5f32.sqrt();
        assert!((logits.data[0] - 3.0 * rstd).abs() < 1e-5);
        assert!((logits.data[1] - 4.0 * rstd).abs() < 1e-5);
        Ok(())
    }

    #[test]
    fn lm_head_last_row_is_full_last_row() -> FocrResult<()> {
        // The "compute only the rows you read" structural win: the decode driver
        // feeds lm_head ONLY the last hidden row (mod.rs). RMSNorm and the lm_head
        // linear are per-row independent, so the last-row slice MUST be exactly
        // (bit-for-bit) the last row of projecting all rows — proving the prefill
        // optimization (skip the [seq-1, vocab] head GEMM) changes nothing observable.
        let (seq, hidden, vocab) = (5usize, 8usize, 7usize);
        let h: Vec<f32> = (0..seq * hidden)
            .map(|i| ((i as f32) * 0.37).sin())
            .collect();
        let full = Mat::from_vec(seq, hidden, h);
        let norm_w: Vec<f32> = (0..hidden).map(|i| 0.5 + (i as f32) * 0.1).collect();
        let head_w: Vec<f32> = (0..vocab * hidden)
            .map(|i| ((i as f32) * 0.13).cos())
            .collect();
        let eps = 1e-6;

        let full_logits = norm_and_lm_head(&full, &norm_w, &head_w, vocab, eps)?;
        let last = Mat::from_vec(1, hidden, full.row(seq - 1).to_vec());
        let last_logits = norm_and_lm_head(&last, &norm_w, &head_w, vocab, eps)?;

        assert_eq!((last_logits.rows, last_logits.cols), (1, vocab));
        let full_last = &full_logits.data[(seq - 1) * vocab..seq * vocab];
        assert_eq!(
            last_logits.data, full_last,
            "last-row lm_head must be bit-identical to full[last]"
        );
        Ok(())
    }

    // ── attn output proj / qkv ([SPEC-090]) ─────────────────────────────────

    #[test]
    fn attn_output_proj_projects_context() -> FocrResult<()> {
        // qkv_dim=2 hidden=2; context=[1,1]; o_proj=[[2,0],[0,3]] -> [2,3]
        let ctx = Mat::from_vec(1, 2, vec![1.0, 1.0]);
        let o = vec![2.0, 0.0, 0.0, 3.0];
        let out = attn_output_proj(&ctx, &o, 2, 2)?;
        assert_eq!(out.data, vec![2.0, 3.0]);
        Ok(())
    }

    #[test]
    fn qkv_with_rope_shapes_and_pos0_identity() -> FocrResult<()> {
        // hidden=2, single head head_dim=2 (qkv_dim=2), one token at pos 0.
        // q/k/v projections identity => q=k=v=normed; RoPE at pos0 = identity.
        let normed = Mat::from_vec(1, 2, vec![1.0, 2.0]);
        let ident = vec![1.0, 0.0, 0.0, 1.0];
        let lw = LayerWeights {
            input_ln: &[1.0, 1.0],
            post_attn_ln: &[1.0, 1.0],
            q_proj: &ident,
            k_proj: &ident,
            v_proj: &ident,
            o_proj: &ident,
            gate_w: &[],
            up_w: &[],
            down_w: &[],
        };
        let rope = RopeTable::build(&[0], 2, 10000.0);
        let (q, k, v) = qkv_with_rope(&normed, &lw, &rope, 2, 2)?;
        assert_eq!(q.shape(), (1, 2));
        assert_eq!(v.data, vec![1.0, 2.0]); // v never rope'd
        assert_eq!(q.data, vec![1.0, 2.0]); // pos0 rope = identity
        assert_eq!(k.data, vec![1.0, 2.0]);
        Ok(())
    }

    // ── Full layer driver ([SPEC-072]) ──────────────────────────────────────

    #[test]
    fn layer_forward_pre_norm_residual_identity_path() -> FocrResult<()> {
        // Build a layer where attention and MLP both return ZERO; then the
        // output must equal the input (pure residual passthrough), proving the
        // `x + sublayer(...)` wiring ([SPEC-072]).
        let hidden = 2usize;
        let qkv_dim = 2usize;
        let x = Mat::from_vec(2, hidden, vec![1.0, 2.0, 3.0, 4.0]);
        let ident = vec![1.0, 0.0, 0.0, 1.0];
        let lw = LayerWeights {
            input_ln: &[1.0, 1.0],
            post_attn_ln: &[1.0, 1.0],
            q_proj: &ident,
            k_proj: &ident,
            v_proj: &ident,
            o_proj: &ident,
            gate_w: &[],
            up_w: &[],
            down_w: &[],
        };
        let rope = RopeTable::build(&[0, 1], qkv_dim, 10000.0);
        let out = layer_forward(
            &x,
            &lw,
            &rope,
            hidden,
            qkv_dim,
            1e-6,
            // attention returns zeros
            |q, _k, _v| Ok(Mat::zeros(q.rows, q.cols)),
            // mlp returns zeros
            |n| Ok(Mat::zeros(n.rows, n.cols)),
        )?;
        assert_eq!(out.shape(), (2, hidden));
        for (got, want) in out.data.iter().zip(x.data.iter()) {
            assert!((got - want).abs() < 1e-6, "{got} != {want}");
        }
        Ok(())
    }

    #[test]
    fn layer_forward_adds_both_sublayers() -> FocrResult<()> {
        // attn closure returns a constant context that o_proj maps to a known
        // vector; mlp adds another known vector. Output = x + attn_out + mlp_out.
        let hidden = 2usize;
        let qkv_dim = 2usize;
        let x = Mat::from_vec(1, hidden, vec![10.0, 20.0]);
        let ident = vec![1.0, 0.0, 0.0, 1.0];
        let lw = LayerWeights {
            input_ln: &[1.0, 1.0],
            post_attn_ln: &[1.0, 1.0],
            q_proj: &ident,
            k_proj: &ident,
            v_proj: &ident,
            o_proj: &ident, // o_proj identity => attn_out == context
            gate_w: &[],
            up_w: &[],
            down_w: &[],
        };
        let rope = RopeTable::build(&[0], qkv_dim, 10000.0);
        let out = layer_forward(
            &x,
            &lw,
            &rope,
            hidden,
            qkv_dim,
            1e-6,
            |_q, _k, _v| Ok(Mat::from_vec(1, qkv_dim, vec![1.0, 1.0])),
            |_n| Ok(Mat::from_vec(1, hidden, vec![100.0, 100.0])),
        )?;
        // x + [1,1] + [100,100] = [111, 121]
        assert!((out.data[0] - 111.0).abs() < 1e-4, "{}", out.data[0]);
        assert!((out.data[1] - 121.0).abs() < 1e-4, "{}", out.data[1]);
        Ok(())
    }

    // ── Top-level entrypoints are wired to the Weights accessors ────────────

    #[test]
    fn top_level_entrypoints_error_cleanly_on_empty_weights() {
        // `forward`/`lm_head` now look tensors up by name and delegate to the
        // tested driver; an empty `Weights::default()` has none, so they must
        // surface a clean `FormatMismatch` (tensor not found), never panic.
        let w = Weights::default();
        let h = Mat::zeros(1, config::HIDDEN_SIZE);
        assert!(matches!(forward(&w, &h), Err(FocrError::FormatMismatch(_))));
        assert!(matches!(lm_head(&w, &h), Err(FocrError::FormatMismatch(_))));
    }

    // ── Prefill attention shape + causal correctness ────────────────────────

    #[test]
    fn prefill_attention_shapes_and_first_token_self_only() -> FocrResult<()> {
        // Single head, head_dim=2, seq=2. With a causal mask, row 0 attends only
        // to itself, so its context == v[0] regardless of v[1] / the scores.
        let (num_heads, head_dim, seq) = (1usize, 2usize, 2usize);
        let q = Mat::from_vec(seq, 2, vec![1.0, 0.0, 0.0, 1.0]);
        let k = Mat::from_vec(seq, 2, vec![1.0, 0.0, 0.0, 1.0]);
        let v = Mat::from_vec(seq, 2, vec![5.0, 6.0, 7.0, 8.0]);
        let out = prefill_attention(&q, &k, &v, num_heads, head_dim)?;
        assert_eq!(out.shape(), (seq, num_heads * head_dim));
        // Row 0 (causal): only key 0 visible -> context == v row 0.
        assert!((out.data[0] - 5.0).abs() < 1e-5, "{}", out.data[0]);
        assert!((out.data[1] - 6.0).abs() < 1e-5, "{}", out.data[1]);
        // Row 1 attends to both keys -> a convex blend of v[0] and v[1], so each
        // channel lies strictly between the two value rows.
        assert!(out.data[2] > 5.0 && out.data[2] < 7.0, "{}", out.data[2]);
        assert!(out.data[3] > 6.0 && out.data[3] < 8.0, "{}", out.data[3]);
        Ok(())
    }

    #[test]
    fn prefill_attention_rejects_bad_shape() {
        let q = Mat::zeros(2, 3); // 3 not a multiple of head_dim*heads=4
        let k = Mat::zeros(2, 4);
        let v = Mat::zeros(2, 4);
        assert!(prefill_attention(&q, &k, &v, 2, 2).is_err());
    }

    #[test]
    fn prefill_attention_rejects_malformed_backing_data_without_panic() {
        let q = Mat {
            rows: 2,
            cols: 4,
            data: vec![0.0; 7],
        };
        let k = Mat::zeros(2, 4);
        let v = Mat::zeros(2, 4);
        assert_err_contains(
            prefill_attention(&q, &k, &v, 2, 2),
            "prefill_attention q: data len 7",
        );
    }
}