fuel-transformers 0.20.0

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Mistral decoder ported to the lazy-graph API.
//!
//! Phase D LLM port. The Mistral architecture is "LLaMA with sliding-
//! window attention" — every other piece (GQA, RmsNorm, SwiGLU FFN,
//! RoPE, bias-free Q/K/V/O projections) is shared with LLaMA. We
//! reuse [`fuel_tensor::lazy::LayerWeights`] directly since the per-layer
//! parameter shape is identical; only the attention mask differs.
//!
//! # Sliding-window mask
//!
//! Eager Mistral builds a `[seq, seq]` (or `[seq, total_seq]` for
//! decode with KV cache) mask where `mask[i, j] = -inf` if either:
//!   - `j > i` (causal — can't attend to future)
//!   - `j + sliding_window <= i` (sliding-window — can't attend past
//!     `sliding_window` tokens back)
//!
//! With `sliding_window = None`, this reduces to a strict lower-
//! triangular causal mask (LLaMA semantics). The lazy port mirrors
//! this — every other piece of `apply_layer` matches the LLaMA path
//! (`fuel_model_llama::LlamaModel::apply_layer`).
//!
//! # KV cache
//!
//! v1 of this port does NOT use a KV cache — every `forward` call
//! recomputes attention over the full input. This matches the LLaMA
//! lazy port's v1 contract and keeps the prefill / decode patterns
//! uniform. Adding [`crate::models::lazy_kv_cache::LazyKvCache::append_rotating`]
//! (Phase C) on top is orthogonal plumbing for a later session;
//! `MistralModel` will gain a `forward_with_cache` then, mirroring
//! the LLaMA pair.
//!
//! # Scope
//!
//! - Forward-only. Autograd through Mistral is not exercised here.
//! - Single sequence (`batch == 1`). Multi-batch is a future addition
//!   when InferenceContext starts batching real workloads.
//! - F32 activations; weights via [`fuel_tensor::lazy::WeightStorage`]
//!   (F32 / BF16 / Q4_0 all work through `apply_linear`).
//!
//! # Weight names (HuggingFace safetensors)
//!
//! Mirrors the eager `fuel_transformers::models::llm::mistral` field
//! ordering for safetensors loader interop:
//!   - `model.embed_tokens.weight` → [`MistralWeights::token_embedding`]
//!   - `model.layers.{i}.self_attn.{q,k,v,o}_proj.weight` → `attn_{q,k,v,o}`
//!   - `model.layers.{i}.mlp.{gate,up,down}_proj.weight` → `ffn_{gate,up,down}`
//!   - `model.layers.{i}.input_layernorm.weight` → `attn_norm_gain`
//!   - `model.layers.{i}.post_attention_layernorm.weight` → `ffn_norm_gain`
//!   - `model.norm.weight` → `final_norm_gain`
//!   - `lm_head.weight` → `output`
//!
//! All Q/K/V/O biases are `None` (Mistral uses `linear_no_bias`).

use fuel_ir::Shape;
use fuel_ir::error::Result;
use fuel_tensor::Device;
use fuel_tensor::{
    lazy::LayerWeights, lazy::Tensor, lazy::WeightStorage, lazy::load_tensor_as_f32,
    lazy::load_transposed_matrix_preserve_dtype,
};
use std::sync::Arc;

/// Mistral / Mixtral model configuration.
///
/// Mirrors HuggingFace `config.json` naming where practical so a
/// downloaded config can be deserialized directly (when the
/// safetensors loader is wired). Fields:
///   - `sliding_window`: window size for the sliding-window mask.
///     `Some(4096)` for Mistral-7B-v0.1; `None` reverts to a strict
///     causal mask (i.e. dense attention, equivalent to LLaMA).
#[derive(Debug, Clone, PartialEq)]
pub struct MistralConfig {
    pub vocab_size: usize,
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub num_hidden_layers: usize,
    pub num_attention_heads: usize,
    pub num_key_value_heads: usize,
    pub head_dim: usize,
    pub rms_norm_eps: f64,
    pub rope_theta: f64,
    pub max_position_embeddings: usize,
    pub sliding_window: Option<usize>,
}

impl MistralConfig {
    /// Preset for `mistralai/Mistral-7B-v0.1`. Same layout the eager
    /// `Config::config_7b_v0_1` returns.
    pub fn mistral_7b_v0_1() -> Self {
        Self {
            vocab_size: 32000,
            hidden_size: 4096,
            intermediate_size: 14336,
            num_hidden_layers: 32,
            num_attention_heads: 32,
            num_key_value_heads: 8,
            head_dim: 128,
            rms_norm_eps: 1e-5,
            rope_theta: 10_000.0,
            max_position_embeddings: 32_768,
            sliding_window: Some(4096),
        }
    }
}

// ROADMAP item 8 (II): config-from-path as a capability of the config TYPE.
// Mirrors #57's Qwen2 template — a `serde` raw carrying the HF field names plus
// constant defaults, then `resolve` fills the sibling-derived values
// (num_key_value_heads, head_dim) through the shared `hf_config` rules.
#[derive(Debug, Clone, serde::Deserialize)]
struct MistralConfigRaw {
    vocab_size: usize,
    hidden_size: usize,
    intermediate_size: usize,
    num_hidden_layers: usize,
    num_attention_heads: usize,
    #[serde(default)]
    num_key_value_heads: Option<usize>,
    #[serde(default)]
    head_dim: Option<usize>,
    #[serde(default = "default_mistral_rms_norm_eps")]
    rms_norm_eps: f64,
    #[serde(default = "default_mistral_rope_theta")]
    rope_theta: f64,
    max_position_embeddings: usize,
    #[serde(default)]
    sliding_window: Option<usize>,
}

fn default_mistral_rms_norm_eps() -> f64 {
    1e-5
}
fn default_mistral_rope_theta() -> f64 {
    10_000.0
}

impl MistralConfigRaw {
    fn from_json_str(json: &str) -> Result<Self> {
        serde_json::from_str(json)
            .map_err(|e| fuel_ir::error::Error::Msg(format!("parsing Mistral config.json: {e}")))
    }

    fn resolve(self) -> Result<MistralConfig> {
        Ok(MistralConfig {
            vocab_size: self.vocab_size,
            hidden_size: self.hidden_size,
            intermediate_size: self.intermediate_size,
            num_hidden_layers: self.num_hidden_layers,
            num_attention_heads: self.num_attention_heads,
            num_key_value_heads: fuel_loaders::hf_config::num_key_value_heads(
                self.num_key_value_heads,
                self.num_attention_heads,
            )?,
            head_dim: fuel_loaders::hf_config::head_dim(
                self.head_dim,
                self.hidden_size,
                self.num_attention_heads,
            )?,
            rms_norm_eps: self.rms_norm_eps,
            rope_theta: self.rope_theta,
            max_position_embeddings: self.max_position_embeddings,
            // Mistral keeps sliding_window as Option (None = no window); pass as-is.
            sliding_window: self.sliding_window,
        })
    }
}

impl MistralConfig {
    /// Parse a HuggingFace `config.json` string into a [`MistralConfig`].
    ///
    /// ROADMAP item 8 (II): reads the artifact rather than returning a preset —
    /// see the born-red `mistral_config_from_hf_json_parses_the_artifact_not_a_preset`.
    pub fn from_hf_json_str(json: &str) -> Result<Self> {
        MistralConfigRaw::from_json_str(json)?.resolve()
    }
}

/// Mistral weights. Per-layer K/V/Q/O/FFN tensors reuse
/// [`fuel_tensor::lazy::LayerWeights`] (the LLaMA shape works unchanged
/// because Mistral's per-layer parameter layout is identical to
/// bias-free LLaMA).
#[derive(Debug, Clone)]
pub struct MistralWeights {
    /// `[vocab_size, hidden_size]` token embedding table.
    pub token_embedding: Arc<[f32]>,
    /// Per-layer weights.
    pub layers: Vec<LayerWeights>,
    /// `[hidden_size]` RmsNorm gain before the lm_head.
    pub final_norm_gain: Arc<[f32]>,
    /// `[hidden_size, vocab_size]` lm_head projection.
    pub output: WeightStorage,
}

/// Mistral LM, lazy-graph form.
#[derive(Debug, Clone)]
pub struct MistralModel {
    pub config: MistralConfig,
    pub weights: MistralWeights,
}

impl MistralModel {
    /// Run a forward pass on `tokens` and return the final logits
    /// `[1, seq, vocab_size]` as a [`Tensor`]. Call `.realize_f32()`
    /// on the result to materialize.
    ///
    /// `start_pos` offsets the RoPE frequencies. Pass `0` for the
    /// first forward call of a sequence; for v1 (no KV cache) callers
    /// always re-feed the full prefix so `start_pos` is typically `0`.
    pub fn forward(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let seq = tokens.len();
        let _batch = 1;
        assert!(seq > 0, "MistralModel::forward: tokens must be non-empty");

        // Embedding lookup.
        let h = Tensor::embed_tokens(
            weights.token_embedding.clone(),
            cfg.vocab_size,
            cfg.hidden_size,
            tokens,
            &Device::cpu(),
        )?;

        self.forward_embeds(&h, start_pos)
    }

    /// Forward from pre-computed input embeddings of shape
    /// `(batch, seq, hidden_size)`. Used by multimodal models
    /// (Pixtral, etc.) that interleave image embeddings with
    /// text embeddings before running the Mistral decoder.
    pub fn forward_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let dims = embeds.shape();
        let dims = dims.dims();
        embeds
            .shape()
            .dims3()
            .map_err(|e| e.context("MistralModel::forward_embeds: embeds"))?;
        let seq = dims[1];
        // GAP-281: a typed decline, not a panic. The relation is a CONFIG
        // property a real checkpoint can violate, so it is an error to return,
        // not an invariant to assert. Matches the conformant form used by the
        // other 21 models (see `lazy_qwen3.rs`).
        if cfg.num_attention_heads * cfg.head_dim != cfg.hidden_size {
            return Err(fuel_ir::error::Error::Msg(
                "MistralConfig: num_attention_heads * head_dim must equal hidden_size".into(),
            )
            .bt());
        }
        // GAP-281: a typed decline, not a panic. Copied from the conformant
        // form in `lazy_mixtral.rs`; the message keeps its values via format!.
        if !cfg
            .num_attention_heads
            .is_multiple_of(cfg.num_key_value_heads)
        {
            return Err(fuel_ir::error::Error::Msg(format!(
                "MistralConfig: num_attention_heads ({}) must be a multiple of num_key_value_heads ({})",
                cfg.num_attention_heads, cfg.num_key_value_heads,
            ))
            .bt());
        }

        let mut h = embeds.clone();

        let (rope_cos, rope_sin) =
            h.rope_tables_const(cfg.rope_theta, start_pos, seq, cfg.head_dim);
        let mask = self.build_sliding_window_mask(embeds, seq);

        for layer in &weights.layers {
            h = self.apply_layer(&h, layer, &rope_cos, &rope_sin, &mask)?;
        }

        let h_norm = h.rms_norm_affine(
            std::sync::Arc::clone(&weights.final_norm_gain),
            cfg.rms_norm_eps,
        )?;
        weights
            .output
            .apply_linear(&h_norm, cfg.hidden_size, cfg.vocab_size)
    }

    /// Run the decoder forward up to the final RmsNorm and
    /// return per-token hidden states `(1, seq, hidden_size)`.
    /// Skips the `lm_head` projection — useful for embedding
    /// adapters (NV-Embed v2's perceiver-style latent-attention
    /// head, custom poolers, etc.) that swap the causal LM head
    /// for a downstream projector.
    pub fn forward_hidden(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let seq = tokens.len();
        let _batch = 1;
        assert!(
            seq > 0,
            "MistralModel::forward_hidden: tokens must be non-empty"
        );

        let h = Tensor::embed_tokens(
            weights.token_embedding.clone(),
            cfg.vocab_size,
            cfg.hidden_size,
            tokens,
            &Device::cpu(),
        )?;
        self.forward_hidden_embeds(&h, start_pos)
    }

    /// Like [`Self::forward_embeds`] but skips the `lm_head`
    /// projection and returns the post-RmsNorm hidden states.
    /// Used by multimodal+embedding compositions that mix
    /// pre-computed embeddings with a custom pooling head.
    pub fn forward_hidden_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        let _weights = &self.weights;
        let dims = embeds.shape();
        let dims = dims.dims();
        embeds
            .shape()
            .dims3()
            .map_err(|e| e.context("MistralModel::forward_hidden_embeds: embeds"))?;
        let seq = dims[1];
        let mask = self.build_sliding_window_mask(embeds, seq);
        self.forward_hidden_embeds_with_mask_impl(embeds, &mask, start_pos)
    }

    /// Like [`Self::forward_hidden_embeds`] but takes a
    /// caller-supplied additive attention mask of shape
    /// `(1, 1, seq, seq)` instead of building the standard
    /// sliding-window causal mask internally. Pass `0` for
    /// keep and `-inf` (or a large negative) for mask.
    ///
    /// Used by NV-Embed v2 and other Mistral-shape embedding
    /// models that run with bidirectional attention (mask is
    /// just the padding mask, no causal triangle).
    pub fn forward_hidden_embeds_with_mask(
        &self,
        embeds: &Tensor,
        attention_mask: &Tensor,
        start_pos: usize,
    ) -> Result<Tensor> {
        self.forward_hidden_embeds_with_mask_impl(embeds, attention_mask, start_pos)
    }

    fn forward_hidden_embeds_with_mask_impl(
        &self,
        embeds: &Tensor,
        mask: &Tensor,
        start_pos: usize,
    ) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let dims = embeds.shape();
        let dims = dims.dims();
        embeds
            .shape()
            .dims3()
            .map_err(|e| e.context("MistralModel::forward_hidden_embeds_with_mask_impl: embeds"))?;
        let seq = dims[1];
        let mut h = embeds.clone();
        let (rope_cos, rope_sin) =
            h.rope_tables_const(cfg.rope_theta, start_pos, seq, cfg.head_dim);

        for layer in &weights.layers {
            h = self.apply_layer(&h, layer, &rope_cos, &rope_sin, mask)?;
        }
        h.rms_norm_affine(
            std::sync::Arc::clone(&weights.final_norm_gain),
            cfg.rms_norm_eps,
        )
    }

    /// Build the sliding-window causal mask for a single forward
    /// (v1, no cache). Shape `[1, 1, seq, seq]`, broadcast-ready for
    /// `scores`.
    ///
    /// `mask[i, j] = -inf` if `j > i` (causal) OR
    /// `j + sliding_window <= i` (sliding-window). With
    /// `sliding_window == None` this reduces to a strict lower-tri
    /// causal mask.
    fn build_sliding_window_mask(&self, anchor: &Tensor, seq: usize) -> Tensor {
        let cfg = &self.config;
        let window = cfg.sliding_window.unwrap_or(seq + 1);
        let mut mask_data = vec![0.0_f32; seq * seq];
        for i in 0..seq {
            for j in 0..seq {
                if j > i || j + window <= i {
                    mask_data[i * seq + j] = f32::NEG_INFINITY;
                }
            }
        }
        anchor.const_f32_like(mask_data, Shape::from_dims(&[1, 1, seq, seq]))
        .expect(
                "build_sliding_window_mask: buffer is vec![_; seq*seq] and the shape's elem_count is seq*seq -- \n             both derived from `seq` in this function; the loop writes in place",
        )
    }

    /// Single transformer layer. Mirrors
    /// `fuel_model_llama::LlamaModel::apply_layer` except the attention
    /// mask is injected by the caller — usually a sliding-window
    /// causal mask from `build_sliding_window_mask` but can be any
    /// additive `(1, 1, seq, seq)` shape (e.g. bidirectional
    /// padding mask for NV-Embed v2).
    fn apply_layer(
        &self,
        x: &Tensor,
        layer: &LayerWeights,
        rope_cos: &Tensor,
        rope_sin: &Tensor,
        mask: &Tensor,
    ) -> Result<Tensor> {
        let cfg = &self.config;
        let x_shape = x.shape();
        let dims = x_shape.dims();
        let batch = dims[0];
        let seq = dims[1];
        let kv_dim = cfg.num_key_value_heads * cfg.head_dim;

        // Pre-attention RmsNorm (affine).
        let x_norm = x.rms_norm_affine(
            std::sync::Arc::clone(&layer.attn_norm_gain),
            cfg.rms_norm_eps,
        )?;

        // Q / K / V projections (bias-free for Mistral).
        let q = layer
            .attn_q
            .apply_linear(&x_norm, cfg.hidden_size, cfg.hidden_size)?;
        let k = layer
            .attn_k
            .apply_linear(&x_norm, cfg.hidden_size, kv_dim)?;
        let v = layer
            .attn_v
            .apply_linear(&x_norm, cfg.hidden_size, kv_dim)?;

        // Reshape to per-head and transpose to [batch, heads, seq, head_dim].
        let _ = (batch, seq);
        let q = q.split_heads(cfg.num_attention_heads, cfg.head_dim)?;
        let k = k.split_heads(cfg.num_key_value_heads, cfg.head_dim)?;
        let v = v.split_heads(cfg.num_key_value_heads, cfg.head_dim)?;

        // Apply RoPE on Q and K only.
        let q_r = q.rope_with_tables(rope_cos, rope_sin)?;
        let k_r = k.rope_with_tables(rope_cos, rope_sin)?;

        let n_rep = cfg.num_attention_heads / cfg.num_key_value_heads;
        let k_full = k_r.repeat_interleave(1_usize, n_rep)?;
        let v_full = v.repeat_interleave(1_usize, n_rep)?;

        // Scaled dot-product attention with the caller-supplied mask.
        let k_t = k_full.transpose()?;
        let scale = 1.0_f64 / (cfg.head_dim as f64).sqrt();
        let scores = q_r.matmul(&k_t)?;
        let scores_scaled = scores.mul_scalar(scale);
        let _ = seq; // silence unused after refactor; mask already sized for seq.
        let scores_masked = scores_scaled.broadcast_add(mask)?;
        let attn = scores_masked.softmax_last_dim()?;
        let attn_v = attn.matmul(&v_full)?;

        // Merge heads + output projection.
        let merged = attn_v.merge_heads()?;
        let attn_out = layer
            .attn_o
            .apply_linear(&merged, cfg.hidden_size, cfg.hidden_size)?;

        // First residual.
        let h1 = x.add(&attn_out)?;

        // Pre-FFN RmsNorm (affine).
        let h1_norm = h1.rms_norm_affine(
            std::sync::Arc::clone(&layer.ffn_norm_gain),
            cfg.rms_norm_eps,
        )?;

        // SwiGLU FFN: `down(silu(gate) * up)`.
        let gate = layer
            .ffn_gate
            .apply_linear(&h1_norm, cfg.hidden_size, cfg.intermediate_size)?;
        let up = layer
            .ffn_up
            .apply_linear(&h1_norm, cfg.hidden_size, cfg.intermediate_size)?;
        let swiglu = gate.silu().mul(&up)?;
        let ffn_out =
            layer
                .ffn_down
                .apply_linear(&swiglu, cfg.intermediate_size, cfg.hidden_size)?;

        // Second residual.
        h1.add(&ffn_out)
    }
}

// ---- Safetensors loader ----------------------------------------------------

/// Load Mistral weights from a `MmapedSafetensors` file using a
/// configurable prefix (e.g. `""` for a plain Mistral checkpoint
/// or `"language_model."` for a multimodal checkpoint that embeds
/// the LM under a sub-prefix like Pixtral does). Mistral has no
/// Q/K/V biases (`linear_no_bias`) so the loader returns `None`
/// for all three. Output projection falls back to tied embeddings
/// when `lm_head.weight` is absent.
pub fn load_mistral_weights_with_prefix(
    st: &fuel_loaders::safetensors::MmapedSafetensors,
    cfg: &MistralConfig,
    prefix: &str,
) -> Result<MistralWeights> {
    let h = cfg.hidden_size;
    let kv = cfg.num_key_value_heads * cfg.head_dim;
    let token_embedding = load_tensor_as_f32(st, &format!("{prefix}model.embed_tokens.weight"))?;
    if token_embedding.len() != cfg.vocab_size * h {
        fuel_ir::bail!(
            "{prefix}model.embed_tokens.weight: {} elts, expected {}",
            token_embedding.len(),
            cfg.vocab_size * h,
        );
    }

    let mut layers: Vec<LayerWeights> = Vec::with_capacity(cfg.num_hidden_layers);
    for i in 0..cfg.num_hidden_layers {
        let p = format!("{prefix}model.layers.{i}");
        let attn_q = load_transposed_matrix_preserve_dtype(
            st,
            &format!("{p}.self_attn.q_proj.weight"),
            h,
            h,
        )?;
        let attn_k = load_transposed_matrix_preserve_dtype(
            st,
            &format!("{p}.self_attn.k_proj.weight"),
            kv,
            h,
        )?;
        let attn_v = load_transposed_matrix_preserve_dtype(
            st,
            &format!("{p}.self_attn.v_proj.weight"),
            kv,
            h,
        )?;
        let attn_o = load_transposed_matrix_preserve_dtype(
            st,
            &format!("{p}.self_attn.o_proj.weight"),
            h,
            h,
        )?;
        let ffn_gate = load_transposed_matrix_preserve_dtype(
            st,
            &format!("{p}.mlp.gate_proj.weight"),
            cfg.intermediate_size,
            h,
        )?;
        let ffn_up = load_transposed_matrix_preserve_dtype(
            st,
            &format!("{p}.mlp.up_proj.weight"),
            cfg.intermediate_size,
            h,
        )?;
        let ffn_down = load_transposed_matrix_preserve_dtype(
            st,
            &format!("{p}.mlp.down_proj.weight"),
            h,
            cfg.intermediate_size,
        )?;
        let attn_norm_gain = load_tensor_as_f32(st, &format!("{p}.input_layernorm.weight"))?;
        let ffn_norm_gain =
            load_tensor_as_f32(st, &format!("{p}.post_attention_layernorm.weight"))?;
        layers.push(LayerWeights {
            attn_q,
            attn_q_bias: None,
            attn_k,
            attn_k_bias: None,
            attn_v,
            attn_v_bias: None,
            attn_o,
            ffn_gate,
            ffn_up,
            ffn_down,
            attn_norm_gain: Arc::from(attn_norm_gain),
            ffn_norm_gain: Arc::from(ffn_norm_gain),
        });
    }

    let final_norm_gain = load_tensor_as_f32(st, &format!("{prefix}model.norm.weight"))?;
    // Mistral typically does NOT tie embeddings (Mistral-7B has a
    // distinct lm_head); fall back to tied for safety only.
    let output: WeightStorage = match load_transposed_matrix_preserve_dtype(
        st,
        &format!("{prefix}lm_head.weight"),
        cfg.vocab_size,
        h,
    ) {
        Ok(w) => w,
        Err(_) => {
            let mut transposed = vec![0.0_f32; h * cfg.vocab_size];
            for i in 0..cfg.vocab_size {
                for j in 0..h {
                    transposed[j * cfg.vocab_size + i] = token_embedding[i * h + j];
                }
            }
            WeightStorage::F32(Arc::from(transposed))
        }
    };

    Ok(MistralWeights {
        token_embedding: Arc::from(token_embedding),
        layers,
        final_norm_gain: Arc::from(final_norm_gain),
        output,
    })
}

impl MistralWeights {
    /// Load Mistral weights from a `MmapedSafetensors` file using
    /// the standard HuggingFace top-level naming (`model.*` +
    /// `lm_head.*`). See [`load_mistral_weights_with_prefix`] for
    /// the multimodal-embedded form.
    pub fn load_from_mmapped(
        st: &fuel_loaders::safetensors::MmapedSafetensors,
        cfg: &MistralConfig,
    ) -> Result<Self> {
        load_mistral_weights_with_prefix(st, cfg, "")
    }
}

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

    /// LCG-seeded weights so the shape test is deterministic; values
    /// kept small so the forward stays in normal-float range.
    fn tiny_weights(cfg: &MistralConfig) -> MistralWeights {
        let mut s: u32 = 12345;
        let next = || -> f32 {
            s = s.wrapping_mul(1103515245).wrapping_add(12345);
            ((s >> 16) as u16 as f32 / 65535.0 - 0.5) * 0.05
        };
        let vec_of = |n: usize, next: &mut dyn FnMut() -> f32| -> Arc<[f32]> {
            Arc::from((0..n).map(|_| next()).collect::<Vec<_>>())
        };
        let h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        let kv = cfg.num_key_value_heads * cfg.head_dim;
        let mut next_box: Box<dyn FnMut() -> f32> = Box::new(next);
        let token_embedding = vec_of(cfg.vocab_size * h, &mut *next_box);
        let layers: Vec<LayerWeights> = (0..cfg.num_hidden_layers)
            .map(|_| LayerWeights {
                attn_q: WeightStorage::F32(vec_of(h * h, &mut *next_box)),
                attn_q_bias: None,
                attn_k: WeightStorage::F32(vec_of(h * kv, &mut *next_box)),
                attn_k_bias: None,
                attn_v: WeightStorage::F32(vec_of(h * kv, &mut *next_box)),
                attn_v_bias: None,
                attn_o: WeightStorage::F32(vec_of(h * h, &mut *next_box)),
                ffn_gate: WeightStorage::F32(vec_of(h * i, &mut *next_box)),
                ffn_up: WeightStorage::F32(vec_of(h * i, &mut *next_box)),
                ffn_down: WeightStorage::F32(vec_of(i * h, &mut *next_box)),
                attn_norm_gain: Arc::from(vec![1.0_f32; h]),
                ffn_norm_gain: Arc::from(vec![1.0_f32; h]),
            })
            .collect();
        let final_norm_gain = Arc::from(vec![1.0_f32; h]);
        let output = WeightStorage::F32(vec_of(h * cfg.vocab_size, &mut *next_box));
        MistralWeights {
            token_embedding,
            layers,
            final_norm_gain,
            output,
        }
    }

    /// GAP-326: the three embeds entry points check rank with `dims3()`, so a
    /// rank-2 input is a typed rank error naming the function that declined
    /// it. All three used to be `assert_eq!` panics.
    /// A one-layer Mistral small enough to build in a test.
    fn rank_test_model() -> MistralModel {
        let cfg = MistralConfig {
            vocab_size: 32,
            hidden_size: 16,
            intermediate_size: 32,
            num_hidden_layers: 1,
            num_attention_heads: 4,
            num_key_value_heads: 2,
            head_dim: 4,
            rms_norm_eps: 1e-5,
            rope_theta: 10_000.0,
            max_position_embeddings: 32,
            sliding_window: None,
        };
        MistralModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        }
    }

    /// An all-zero f32 CPU tensor of the given dims.
    fn zeros(dims: &[usize]) -> Tensor {
        let count = dims.iter().product();
        Tensor::from_f32(
            Arc::from(vec![0.0_f32; count]),
            Shape::from_dims(dims),
            &Device::cpu(),
        )
        .unwrap()
    }

    /// GAP-326: the three embeds entry points check rank with `dims3()`, so a
    /// rank-2 input is a typed rank error naming the function that declined
    /// it. All three used to be `assert_eq!` panics.
    #[test]
    fn embeds_of_the_wrong_rank_are_a_typed_error_not_a_panic() {
        let model = rank_test_model();
        let cfg = model.config.clone();
        let tensor = zeros;
        // Positive control: rank-3 embeds still build.
        model
            .forward_hidden_embeds(&tensor(&[1, 3, cfg.hidden_size]), 0)
            .expect("rank-3 embeds must build");
        let rank2 = tensor(&[3, cfg.hidden_size]);
        let mask = tensor(&[1, 1, 3, 3]);
        let results = [
            (
                "MistralModel::forward_embeds: embeds",
                model.forward_embeds(&rank2, 0),
            ),
            (
                "MistralModel::forward_hidden_embeds: embeds",
                model.forward_hidden_embeds(&rank2, 0),
            ),
            (
                "MistralModel::forward_hidden_embeds_with_mask_impl: embeds",
                model.forward_hidden_embeds_with_mask(&rank2, &mask, 0),
            ),
        ];
        for (entry, result) in results {
            let text = result.expect_err(entry).to_string();
            assert!(text.contains(entry), "{entry}: {text}");
            assert!(
                text.contains("unexpected rank, expected: 3, got: 2"),
                "{entry}: {text}"
            );
        }
    }

    /// Smoke: a tiny 2-layer Mistral-shape forward produces logits of
    /// the expected shape and contains no NaNs.
    /// GAP-281: the `num_attention_heads * head_dim == hidden_size` relation must
    /// DECLINE with a typed error, not panic. "Never panic on production paths."
    ///
    /// Born-red: against the production `assert_eq!` this test FAILS BY PANIC
    /// rather than by assertion, which is the point -- a panic is not a decline.
    #[test]
    fn head_dim_mismatch_declines_rather_than_panicking() {
        let valid = MistralConfig {
            vocab_size: 32,
            hidden_size: 16,
            intermediate_size: 32,
            num_hidden_layers: 1,
            num_attention_heads: 4,
            num_key_value_heads: 2,
            head_dim: 4, // 4 * 4 == 16 == hidden_size
            rms_norm_eps: 1e-5,
            rope_theta: 10_000.0,
            max_position_embeddings: 32,
            sliding_window: None,
        };
        let tokens: Vec<u32> = vec![1, 2, 3];

        // POSITIVE CONTROL first: the conforming config must SUCCEED, or "it
        // returned Err" below is satisfied by a fixture that is simply broken.
        let good = MistralModel {
            config: valid.clone(),
            weights: tiny_weights(&valid),
        };
        good.forward(&tokens, 0)
            .expect("control: the conforming config must still run");

        let mut bad = valid.clone();
        bad.head_dim = 8; // 4 * 8 = 32 != 16
        let model = MistralModel {
            config: bad,
            weights: tiny_weights(&valid),
        };
        let err = model
            .forward(&tokens, 0)
            .expect_err("a decoupled head_dim must DECLINE, not panic and not succeed");
        let msg = format!("{err}");
        assert!(
            msg.contains("head_dim") && msg.contains("hidden_size"),
            "the decline must name the relation it rejected, got: {msg}"
        );
    }

    #[test]
    fn forward_shape_and_finite_2_layer() {
        let cfg = MistralConfig {
            vocab_size: 32,
            hidden_size: 16,
            intermediate_size: 32,
            num_hidden_layers: 2,
            num_attention_heads: 4,
            num_key_value_heads: 2,
            head_dim: 4,
            rms_norm_eps: 1e-5,
            rope_theta: 10_000.0,
            max_position_embeddings: 64,
            sliding_window: Some(4),
        };
        let model = MistralModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let tokens: Vec<u32> = vec![1, 2, 3, 4, 5, 6, 7, 8];
        let logits = model.forward(&tokens, 0).unwrap();
        assert_eq!(logits.shape().dims(), &[1, tokens.len(), cfg.vocab_size]);
        let out = logits.realize_f32();
        assert_eq!(out.len(), tokens.len() * cfg.vocab_size);
        for (i, &v) in out.iter().enumerate() {
            assert!(v.is_finite(), "logits[{i}] = {v} not finite");
        }
    }

    /// With `sliding_window = None`, the mask reduces to strict
    /// lower-triangular causal — the first row should attend only to
    /// itself, the last row should attend to every prior + itself.
    /// We verify by comparing the model output against a parallel
    /// run with `sliding_window = Some(seq + 1)` (large enough to
    /// not actually mask anything past the causal cutoff). They must
    /// be element-wise identical.
    #[test]
    fn sliding_window_none_matches_large_window() {
        let mut cfg = MistralConfig {
            vocab_size: 16,
            hidden_size: 8,
            intermediate_size: 16,
            num_hidden_layers: 1,
            num_attention_heads: 2,
            num_key_value_heads: 1,
            head_dim: 4,
            rms_norm_eps: 1e-5,
            rope_theta: 10_000.0,
            max_position_embeddings: 32,
            sliding_window: None,
        };
        let weights = tiny_weights(&cfg);
        let tokens: Vec<u32> = vec![3, 1, 4, 1, 5, 9];
        let out_no_window = MistralModel {
            config: cfg.clone(),
            weights: weights.clone(),
        }
        .forward(&tokens, 0)
        .unwrap()
        .realize_f32();
        cfg.sliding_window = Some(tokens.len() + 1);
        let out_large_window = MistralModel {
            config: cfg,
            weights,
        }
        .forward(&tokens, 0)
        .unwrap()
        .realize_f32();
        assert_eq!(out_no_window, out_large_window);
    }

    /// `sliding_window = Some(2)` must NOT match a strict-causal run
    /// — for seq_len > 2 the window actually restricts attention.
    #[test]
    fn sliding_window_2_differs_from_strict_causal() {
        let cfg_small = MistralConfig {
            vocab_size: 16,
            hidden_size: 8,
            intermediate_size: 16,
            num_hidden_layers: 1,
            num_attention_heads: 2,
            num_key_value_heads: 1,
            head_dim: 4,
            rms_norm_eps: 1e-5,
            rope_theta: 10_000.0,
            max_position_embeddings: 32,
            sliding_window: Some(2),
        };
        let mut cfg_full = cfg_small.clone();
        cfg_full.sliding_window = None;
        let weights = tiny_weights(&cfg_small);
        let tokens: Vec<u32> = vec![1, 2, 3, 4, 5];
        let out_small = MistralModel {
            config: cfg_small,
            weights: weights.clone(),
        }
        .forward(&tokens, 0)
        .unwrap()
        .realize_f32();
        let out_full = MistralModel {
            config: cfg_full,
            weights,
        }
        .forward(&tokens, 0)
        .unwrap()
        .realize_f32();
        // The two must differ — sliding_window=2 masks history past 2 tokens
        // while full causal attends to everything below.
        let any_diff = out_small
            .iter()
            .zip(out_full.iter())
            .any(|(&a, &b)| (a - b).abs() > 1e-5);
        assert!(any_diff, "sliding_window=2 should differ from full causal");
    }

    /// `forward_hidden` returns post-RmsNorm hidden states with
    /// shape `(1, seq, hidden)`, not logits. Skipping the
    /// `lm_head` projection unlocks embedding-adapter
    /// composition (NV-Embed v2 perceiver, etc.).
    #[test]
    fn forward_hidden_skips_lm_head() {
        let cfg = MistralConfig {
            vocab_size: 32,
            hidden_size: 16,
            intermediate_size: 32,
            num_hidden_layers: 2,
            num_attention_heads: 4,
            num_key_value_heads: 2,
            head_dim: 4,
            rms_norm_eps: 1e-5,
            rope_theta: 10_000.0,
            max_position_embeddings: 32,
            sliding_window: None,
        };
        let model = MistralModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let tokens: Vec<u32> = vec![1, 2, 3, 4];
        let hidden = model.forward_hidden(&tokens, 0).unwrap();
        assert_eq!(hidden.shape().dims(), &[1, tokens.len(), cfg.hidden_size]);
        for &v in &hidden.realize_f32() {
            assert!(v.is_finite(), "non-finite hidden: {v}");
        }
    }

    // ROADMAP item 8 (II) born-red: config-from-path reads the artifact and routes
    // GQA + head_dim through the shared `hf_config` rules (mirrors #57's Qwen2 rows).
    #[test]
    fn mistral_config_from_hf_json_parses_the_artifact_not_a_preset() {
        // Mistral-7B-v0.1 shape WITHOUT head_dim (derived) — proves it reads the
        // file, not a hardcoded preset.
        let json = r#"{
            "vocab_size": 32000, "hidden_size": 4096, "intermediate_size": 14336,
            "num_hidden_layers": 32, "num_attention_heads": 32, "num_key_value_heads": 8,
            "rms_norm_eps": 1e-5, "rope_theta": 10000.0,
            "max_position_embeddings": 32768, "sliding_window": 4096
        }"#;
        let cfg = MistralConfig::from_hf_json_str(json).expect("parse mistral config.json");
        assert_eq!(cfg.vocab_size, 32000);
        assert_eq!(cfg.hidden_size, 4096);
        assert_eq!(cfg.intermediate_size, 14336);
        assert_eq!(cfg.num_hidden_layers, 32);
        assert_eq!(cfg.num_attention_heads, 32);
        assert_eq!(cfg.num_key_value_heads, 8);
        assert_eq!(cfg.head_dim, 128); // derived: 4096 / 32
        assert_eq!(cfg.rope_theta, 10000.0);
        assert_eq!(cfg.rms_norm_eps, 1e-5);
        assert_eq!(cfg.max_position_embeddings, 32768);
        assert_eq!(cfg.sliding_window, Some(4096));
    }

    #[test]
    fn mistral_config_from_hf_json_reads_a_second_distinct_config() {
        // Different values, explicit head_dim, no sliding_window — a preset return
        // would ignore all of these.
        let json = r#"{
            "vocab_size": 32768, "hidden_size": 5120, "intermediate_size": 14336,
            "num_hidden_layers": 40, "num_attention_heads": 32, "num_key_value_heads": 8,
            "head_dim": 128, "rms_norm_eps": 1e-5, "rope_theta": 1000000.0,
            "max_position_embeddings": 32768
        }"#;
        let cfg = MistralConfig::from_hf_json_str(json).expect("parse second config");
        assert_eq!(cfg.hidden_size, 5120);
        assert_eq!(cfg.num_hidden_layers, 40);
        assert_eq!(cfg.rope_theta, 1000000.0);
        assert_eq!(cfg.head_dim, 128); // explicit, taken as-is
        assert_eq!(cfg.sliding_window, None); // absent -> None (no window)
    }

    #[test]
    fn mistral_config_gqa_defaults_to_num_heads_when_absent() {
        let json = r#"{
            "vocab_size": 32000, "hidden_size": 4096, "intermediate_size": 14336,
            "num_hidden_layers": 32, "num_attention_heads": 32,
            "max_position_embeddings": 32768
        }"#;
        let cfg = MistralConfig::from_hf_json_str(json).expect("parse");
        // absent num_key_value_heads -> num_attention_heads (MHA), via the shared rule.
        assert_eq!(cfg.num_key_value_heads, 32);
    }

    #[test]
    fn mistral_config_preserves_true_mqa() {
        let json = r#"{
            "vocab_size": 32000, "hidden_size": 4096, "intermediate_size": 14336,
            "num_hidden_layers": 32, "num_attention_heads": 32, "num_key_value_heads": 1,
            "max_position_embeddings": 32768
        }"#;
        let cfg = MistralConfig::from_hf_json_str(json).expect("parse");
        assert_eq!(cfg.num_key_value_heads, 1); // true MQA survives, not collapsed to 32
    }

    /// GAP-281 (GQA relation): `num_attention_heads % num_key_value_heads == 0` must
    /// DECLINE with a typed error, not panic. "Never panic on production paths."
    ///
    /// The fixture violates ONLY the GQA axis: head_dim is left conforming so the
    /// head_dim guard (converted separately) cannot answer first and make this vacuous.
    /// 4 heads over 3 kv-heads gives 4 % 3 == 1.
    ///
    /// Born-red: against the production `assert_eq!` this FAILS BY PANIC.
    #[test]
    fn gqa_head_multiple_mismatch_declines_rather_than_panicking() {
        let valid = MistralConfig {
            vocab_size: 32,
            hidden_size: 16,
            intermediate_size: 32,
            num_hidden_layers: 1,
            num_attention_heads: 4,
            num_key_value_heads: 2,
            head_dim: 4,
            rms_norm_eps: 1e-5,
            rope_theta: 10_000.0,
            max_position_embeddings: 32,
            sliding_window: None,
        };
        let tokens: Vec<u32> = vec![1, 2, 3];
        // POSITIVE CONTROL: the conforming config must SUCCEED.
        let good = MistralModel {
            config: valid.clone(),
            weights: tiny_weights(&valid),
        };
        good.forward(&tokens, 0)
            .expect("control: the conforming config must still run");

        let mut bad = valid.clone();
        bad.num_key_value_heads = 3; // 4 % 3 == 1, and head_dim stays conforming
        let model = MistralModel {
            config: bad,
            weights: tiny_weights(&valid),
        };
        let err = model
            .forward(&tokens, 0)
            .expect_err("a non-multiple kv-head count must DECLINE, not panic");
        let msg = format!("{err}");
        assert!(
            msg.contains("num_key_value_heads"),
            "the decline must name the relation it rejected, got: {msg}"
        );
    }
}