fuel-transformers 0.19.1

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Phi (Phi-1.5 / Phi-2) decoder ported to the lazy-graph API.
//!
//! Phase D specialized port. Microsoft's older Phi series (before
//! Phi-3) has its own architectural signature distinct from the
//! Llama family:
//!
//!   1. **Parallel attention + MLP** — same residual block as
//!      Falcon: `out = residual + attn(LN(x)) + mlp(LN(x))`.
//!      A single `input_layernorm` feeds both paths. No
//!      separate `post_attention_layernorm`.
//!   2. **Partial rotary** — only the first
//!      `partial_rotary_factor * head_dim` features of each
//!      head are rotated; the tail passes through unchanged.
//!      Same shape as the StableLM / Persimmon partial rotary
//!      handled by [`Tensor::rope_partial`].
//!   3. **LayerNorm (with bias) on hidden_size** — not RmsNorm.
//!   4. **Sequential MLP** — `fc2(act(fc1(x)))`. No SwiGLU gate.
//!      Activation is config-driven (the reference Phi-2 uses
//!      "new-GELU"; we expose both Gelu and GeluPytorchTanh
//!      variants and treat them as the same approximation
//!      family — the underlying `Tensor::gelu_tanh()` is the
//!      tanh-approximation, matching `Activation::NewGelu`).
//!   5. **All linear layers carry biases** — Q/K/V/dense and
//!      fc1/fc2 and lm_head all include bias terms.
//!
//! The `qk_layernorm` flag is exposed but v1 only supports
//! `qk_layernorm == false` (which is the Phi-2 default — the
//! flag was added for some Phi-2 fine-tunes; the reference
//! checkpoint ships with it off).
//!
//! # Scope (v1)
//!
//! Forward-only, single sequence (`batch == 1`), no KV cache
//! (recomputes each call), F32 activations.

use fuel_ir::Shape;
use fuel_ir::error::Result;
use fuel_tensor::Device;
use fuel_tensor::{lazy::Tensor, lazy::WeightStorage};
use serde::Deserialize;
use std::sync::Arc;

/// Which GELU variant the MLP's activation uses.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PhiActivation {
    /// Standard ERF GELU `0.5 * x * (1 + erf(x / sqrt(2)))`.
    Gelu,
    /// PyTorch's `approximate="tanh"` variant — same family as
    /// `NewGelu` in HF transformers. Default for Phi-2.
    GeluPytorchTanh,
}

#[derive(Debug, Clone, PartialEq)]
pub struct PhiConfig {
    pub vocab_size: usize,
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub num_hidden_layers: usize,
    pub num_attention_heads: usize,
    /// Defaults to `num_attention_heads` if `None` (MHA, no GQA).
    pub num_key_value_heads: Option<usize>,
    pub head_dim: usize,
    pub layer_norm_eps: f64,
    pub rope_theta: f64,
    pub max_position_embeddings: usize,
    /// Fraction of `head_dim` that gets rotated. Phi-2 uses 0.4.
    pub partial_rotary_factor: f64,
    /// v1 only supports `false`. `true` would require a per-head
    /// LayerNorm applied post-reshape — left to follow-up.
    pub qk_layernorm: bool,
    pub hidden_activation: PhiActivation,
}

fn default_layer_norm_eps() -> f64 {
    1e-5
}
fn default_rope_theta() -> f64 {
    10_000.0
}
fn default_partial_rotary_factor() -> f64 {
    0.4
}

/// A Phi `config.json` exactly as HuggingFace ships it.
///
/// Every field here is either present in the file or has a constant default,
/// so this struct is pure `serde` and carries no logic. Anything that has to
/// look at a *sibling* field — which `#[serde(default = "...")]` cannot do —
/// lives in [`PhiConfigRaw::resolve`] instead. That split is the whole reason
/// this type exists; see [`fuel_loaders::hf_config`] for the measurement behind it.
///
/// `head_dim` and `num_key_value_heads` are `Option` on purpose: absent and
/// present-but-equal-to-the-default are different facts, and collapsing them
/// here would discard the distinction before `resolve` can act on it.
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct PhiConfigRaw {
    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_layer_norm_eps")]
    layer_norm_eps: f64,
    #[serde(default = "default_rope_theta")]
    rope_theta: f64,
    max_position_embeddings: usize,
    #[serde(default = "default_partial_rotary_factor")]
    partial_rotary_factor: f64,
    #[serde(default)]
    qk_layernorm: bool,
}

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

    /// Apply the cross-field defaults and the v1 validation.
    fn resolve(self) -> Result<PhiConfig> {
        if self.qk_layernorm {
            return Err(fuel_ir::error::Error::Msg(
                "Phi: qk_layernorm=true is not yet supported by the standalone port (v1)".into(),
            )
            .bt());
        }
        Ok(PhiConfig {
            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,
            // Deliberately NOT resolved: `PhiConfig` keeps the `Option`, and
            // the MHA fallback happens at use site rather than at parse time.
            num_key_value_heads: self.num_key_value_heads,
            head_dim: fuel_loaders::hf_config::head_dim(
                self.head_dim,
                self.hidden_size,
                self.num_attention_heads,
            )?,
            layer_norm_eps: self.layer_norm_eps,
            rope_theta: self.rope_theta,
            max_position_embeddings: self.max_position_embeddings,
            partial_rotary_factor: self.partial_rotary_factor,
            qk_layernorm: self.qk_layernorm,
            hidden_activation: PhiActivation::GeluPytorchTanh,
        })
    }
}

impl PhiConfig {
    /// ⚠️ GAP-282: returns `Result` because the shared rule now REJECTS a
    /// kv-head count that does not divide `num_attention_heads` -- a config
    /// property that truncates silently when GQA groups query heads. The
    /// accessor computes a value that can be invalid, so it reports rather
    /// than asserting.
    pub fn num_kv_heads(&self) -> Result<usize> {
        // Route through the shared rule rather than hand-rolling the fallback
        // (item-8-II membership gate). Verified identical to the previous
        // `unwrap_or(num_attention_heads)` on every input, incl. Some(1) → 1.
        fuel_loaders::hf_config::num_key_value_heads(
            self.num_key_value_heads,
            self.num_attention_heads,
        )
    }
    pub fn rope_dim(&self) -> usize {
        let r = (self.partial_rotary_factor * self.head_dim as f64) as usize;
        // RoPE expects an even rope_dim (pairs of features).
        r & !1
    }
}

#[derive(Debug, Clone)]
pub struct PhiLayerWeights {
    pub input_ln_gain: Arc<[f32]>,
    pub input_ln_bias: Arc<[f32]>,
    pub attn_q: WeightStorage,
    pub attn_q_bias: Arc<[f32]>,
    pub attn_k: WeightStorage,
    pub attn_k_bias: Arc<[f32]>,
    pub attn_v: WeightStorage,
    pub attn_v_bias: Arc<[f32]>,
    /// Phi calls the output projection `dense`.
    pub attn_dense: WeightStorage,
    pub attn_dense_bias: Arc<[f32]>,
    pub fc1: WeightStorage,
    pub fc1_bias: Arc<[f32]>,
    pub fc2: WeightStorage,
    pub fc2_bias: Arc<[f32]>,
}

#[derive(Debug, Clone)]
pub struct PhiWeights {
    pub token_embedding: Arc<[f32]>,
    pub layers: Vec<PhiLayerWeights>,
    pub final_ln_gain: Arc<[f32]>,
    pub final_ln_bias: Arc<[f32]>,
    pub lm_head: WeightStorage,
    /// Optional `lm_head` bias. HF Phi-2 ships with this bias present;
    /// some Phi forks (Dolphin-Phi, fine-tunes) drop it. Loaders MUST
    /// produce `None` rather than zero-padded all-zeros when the bias
    /// is absent on disk — this preserves the no-bias semantics
    /// instead of doing a redundant broadcast-add of zeros and avoids
    /// a panic on checkpoints that genuinely lack the parameter.
    pub lm_head_bias: Option<Arc<[f32]>>,
}

#[derive(Debug, Clone)]
pub struct PhiModel {
    pub config: PhiConfig,
    pub weights: PhiWeights,
}

impl PhiModel {
    pub fn forward(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        let h_norm = self.run_backbone(tokens, start_pos)?;
        self.apply_lm_head(&h_norm)
    }

    /// Run the decoder forward up to the final LayerNorm and
    /// return per-token hidden states `(1, seq, hidden_size)`.
    /// Skips the `lm_head` projection AND its bias.
    pub fn forward_hidden(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        self.run_backbone(tokens, start_pos)
    }

    /// Multimodal entry point. Skips the token-embedding step and runs
    /// the decoder over pre-embedded inputs (e.g. concatenated vision +
    /// text embeddings for a Phi-based multimodal model like
    /// Moondream).
    ///
    /// `embeds` shape: `(1, seq, hidden_size)`. Phi does NOT scale
    /// embeddings — `embeds` is passed raw.
    pub fn forward_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        let h_norm = self.run_backbone_embeds(embeds, start_pos)?;
        self.apply_lm_head(&h_norm)
    }

    /// Hidden-state variant of [`Self::forward_embeds`]. Returns the
    /// post-final-LayerNorm states `(1, seq, hidden_size)`.
    pub fn forward_hidden_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        self.run_backbone_embeds(embeds, start_pos)
    }

    /// Build per-token embeddings without running the decoder. Used by
    /// multimodal compositions. Returns shape `(1, seq, hidden_size)`.
    pub fn embed_tokens_anchored(&self, anchor: &Tensor, tokens: &[u32]) -> Result<Tensor> {
        let cfg = &self.config;
        anchor.embed_tokens_anchored(
            self.weights.token_embedding.clone(),
            cfg.vocab_size,
            cfg.hidden_size,
            tokens,
        )
    }

    fn apply_lm_head(&self, h_norm: &Tensor) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        match &weights.lm_head_bias {
            Some(b) => weights.lm_head.apply_linear_with_bias(
                h_norm,
                cfg.hidden_size,
                cfg.vocab_size,
                Arc::clone(b),
            ),
            None => Ok(weights
                .lm_head
                .apply_linear(h_norm, cfg.hidden_size, cfg.vocab_size)?),
        }
    }

    fn run_backbone(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let seq = tokens.len();
        assert!(seq > 0, "PhiModel: tokens must be non-empty");

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

    fn run_backbone_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();
        if dims.len() != 3 || dims[0] != 1 || dims[2] != cfg.hidden_size {
            return Err(fuel_ir::error::Error::Msg(format!(
                "PhiModel::forward_embeds: expected embeds shape \
                 (1, seq, hidden_size={}), got {:?}",
                cfg.hidden_size, dims,
            ))
            .bt());
        }
        let seq = dims[1];
        if seq == 0 {
            return Err(fuel_ir::error::Error::Msg(
                "PhiModel::forward_embeds: seq must be > 0".into(),
            )
            .bt());
        }
        if cfg.num_attention_heads * cfg.head_dim != cfg.hidden_size {
            return Err(fuel_ir::error::Error::Msg(
                "PhiConfig: num_attention_heads * head_dim must equal hidden_size".into(),
            )
            .bt());
        }
        if cfg.qk_layernorm {
            return Err(fuel_ir::error::Error::Msg(
                "PhiModel v1: qk_layernorm not yet supported (reference Phi-2 sets it false)"
                    .into(),
            )
            .bt());
        }
        let rope_dim = cfg.rope_dim();
        if rope_dim == 0 || rope_dim > cfg.head_dim {
            return Err(fuel_ir::error::Error::Msg(format!(
                "PhiConfig: rope_dim ({}) out of [1, head_dim ({})]",
                rope_dim, cfg.head_dim,
            ))
            .bt());
        }
        let mut h = embeds.clone();

        let (rope_cos, rope_sin) = h.rope_tables_const(cfg.rope_theta, start_pos, seq, rope_dim);

        for layer in &weights.layers {
            h = self.apply_layer(&h, layer, &rope_cos, &rope_sin)?;
        }
        h.layer_norm_affine(
            std::sync::Arc::clone(&weights.final_ln_gain),
            std::sync::Arc::clone(&weights.final_ln_bias),
            cfg.layer_norm_eps,
        )
    }

    fn apply_layer(
        &self,
        x: &Tensor,
        layer: &PhiLayerWeights,
        rope_cos: &Tensor,
        rope_sin: &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 n_kv = cfg.num_kv_heads()?;
        let kv_dim = n_kv * cfg.head_dim;
        let rope_dim = cfg.rope_dim();

        // Single LN feeds BOTH attention and MLP paths (Phi parallel block).
        let x_norm = x.layer_norm_affine(
            std::sync::Arc::clone(&layer.input_ln_gain),
            std::sync::Arc::clone(&layer.input_ln_bias),
            cfg.layer_norm_eps,
        )?;

        // ---- Attention path -------------------------------------------------
        let q = layer.attn_q.apply_linear_with_bias(
            &x_norm,
            cfg.hidden_size,
            cfg.hidden_size,
            std::sync::Arc::clone(&layer.attn_q_bias),
        )?;
        let k = layer.attn_k.apply_linear_with_bias(
            &x_norm,
            cfg.hidden_size,
            kv_dim,
            std::sync::Arc::clone(&layer.attn_k_bias),
        )?;
        let v = layer.attn_v.apply_linear_with_bias(
            &x_norm,
            cfg.hidden_size,
            kv_dim,
            std::sync::Arc::clone(&layer.attn_v_bias),
        )?;

        let _ = (batch, seq);
        let q = q.split_heads(cfg.num_attention_heads, cfg.head_dim)?;
        let k = k.split_heads(n_kv, cfg.head_dim)?;
        let v = v.split_heads(n_kv, cfg.head_dim)?;

        // Partial rotary on the first `rope_dim` features.
        let q_r = q.rope_partial(rope_cos, rope_sin, rope_dim)?;
        let k_r = k.rope_partial(rope_cos, rope_sin, rope_dim)?;

        // GQA replication (Phi-2 uses MHA so n_rep = 1 usually).
        let n_rep = cfg.num_attention_heads / n_kv;
        let k_full = k_r.repeat_interleave(1_usize, n_rep)?;
        let v_full = v.repeat_interleave(1_usize, n_rep)?;

        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);
        // Strict causal mask.
        let mut mask_data = vec![0.0_f32; seq * seq];
        for i in 0..seq {
            for j in (i + 1)..seq {
                mask_data[i * seq + j] = f32::NEG_INFINITY;
            }
        }
        let mask = x.const_f32_like(mask_data, Shape::from_dims(&[1, 1, seq, seq]))?;
        let scores_masked = scores_scaled.broadcast_add(&mask)?;
        let attn = scores_masked.softmax_last_dim()?;
        let attn_v = attn.matmul(&v_full)?;

        let merged = attn_v.merge_heads()?;
        let attn_out = layer.attn_dense.apply_linear_with_bias(
            &merged,
            cfg.hidden_size,
            cfg.hidden_size,
            std::sync::Arc::clone(&layer.attn_dense_bias),
        )?;

        // ---- MLP path (uses the SAME x_norm) --------------------------------
        let fc1_out = layer.fc1.apply_linear_with_bias(
            &x_norm,
            cfg.hidden_size,
            cfg.intermediate_size,
            std::sync::Arc::clone(&layer.fc1_bias),
        )?;
        let activated = match cfg.hidden_activation {
            PhiActivation::Gelu => fc1_out.gelu(),
            PhiActivation::GeluPytorchTanh => fc1_out.gelu_tanh(),
        };
        let mlp_out = layer.fc2.apply_linear_with_bias(
            &activated,
            cfg.intermediate_size,
            cfg.hidden_size,
            std::sync::Arc::clone(&layer.fc2_bias),
        )?;

        // Parallel combine: residual + attn + mlp.
        x.add(&attn_out)?.add(&mlp_out)
    }
}

// ---- HuggingFace config + safetensors loading ----------------------------

impl PhiConfig {
    /// Parse a Phi-2 `config.json` from HuggingFace.
    ///
    /// HF native names mapped to standalone field names:
    ///   `hidden_size` → `hidden_size`,
    ///   `intermediate_size` → `intermediate_size`,
    ///   `num_hidden_layers` → `num_hidden_layers`,
    ///   `num_attention_heads` → `num_attention_heads`,
    ///   `num_key_value_heads` → `num_key_value_heads` (Option — falls
    ///   back to `num_attention_heads` for MHA configs),
    ///   `head_dim` → `head_dim` (defaults to `hidden_size / num_attention_heads`),
    ///   `layer_norm_eps` → `layer_norm_eps` (defaults to 1e-5),
    ///   `rope_theta` → `rope_theta` (defaults to 10000.0),
    ///   `max_position_embeddings` → `max_position_embeddings`,
    ///   `partial_rotary_factor` → `partial_rotary_factor` (defaults to 0.4 — Phi-2's),
    ///   `qk_layernorm` → `qk_layernorm` (defaults to false; v1 only
    ///   supports false — set to true bails at load time).
    ///
    /// Two stages, because they answer different questions. `PhiConfigRaw`
    /// is the wire shape and is pure `serde`; `PhiConfigRaw::resolve`
    /// applies the defaults `serde` cannot express — those that read another
    /// field — and the one validation this port enforces.
    pub fn from_hf_json_str(json: &str) -> Result<Self> {
        PhiConfigRaw::from_json_str(json)?.resolve()
    }
}

impl PhiWeights {
    /// Load Phi-2 weights from memory-mapped safetensors. Loads weights
    /// as `WeightStorage::F32` after on-load transposition from HF's
    /// `[out, in]` layout to fuel's `[in, out]`.
    ///
    /// `lm_head.bias` is loaded as `None` when absent on disk — matches
    /// the `Option<Arc<[f32]>>` field type and avoids the previous
    /// panicking design (closed in commit b723ddf8).
    pub fn load_from_mmapped(
        st: &fuel_loaders::safetensors::MmapedSafetensors,
        cfg: &PhiConfig,
    ) -> Result<Self> {
        let h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        let n_kv = cfg.num_kv_heads()?;
        let kv_dim = n_kv * cfg.head_dim;

        let token_embedding = Arc::from(fuel_tensor::lazy::load_tensor_as_f32(
            st,
            "model.embed_tokens.weight",
        )?);

        let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
        for li in 0..cfg.num_hidden_layers {
            let load_lin = |suffix: &str, out_f: usize, in_f: usize| -> Result<WeightStorage> {
                let name = format!("model.layers.{li}.{suffix}");
                let v = fuel_tensor::lazy::load_transposed_matrix(st, &name, out_f, in_f)?;
                Ok(WeightStorage::F32(Arc::from(v)))
            };
            let load_vec = |suffix: &str| -> Result<Arc<[f32]>> {
                let name = format!("model.layers.{li}.{suffix}");
                Ok(Arc::from(fuel_tensor::lazy::load_tensor_as_f32(st, &name)?))
            };

            layers.push(PhiLayerWeights {
                input_ln_gain: load_vec("input_layernorm.weight")?,
                input_ln_bias: load_vec("input_layernorm.bias")?,
                attn_q: load_lin("self_attn.q_proj.weight", h, h)?,
                attn_q_bias: load_vec("self_attn.q_proj.bias")?,
                attn_k: load_lin("self_attn.k_proj.weight", kv_dim, h)?,
                attn_k_bias: load_vec("self_attn.k_proj.bias")?,
                attn_v: load_lin("self_attn.v_proj.weight", kv_dim, h)?,
                attn_v_bias: load_vec("self_attn.v_proj.bias")?,
                attn_dense: load_lin("self_attn.dense.weight", h, h)?,
                attn_dense_bias: load_vec("self_attn.dense.bias")?,
                fc1: load_lin("mlp.fc1.weight", i, h)?,
                fc1_bias: load_vec("mlp.fc1.bias")?,
                fc2: load_lin("mlp.fc2.weight", h, i)?,
                fc2_bias: load_vec("mlp.fc2.bias")?,
            });
        }

        let final_ln_gain = Arc::from(fuel_tensor::lazy::load_tensor_as_f32(
            st,
            "model.final_layernorm.weight",
        )?);
        let final_ln_bias = Arc::from(fuel_tensor::lazy::load_tensor_as_f32(
            st,
            "model.final_layernorm.bias",
        )?);
        let lm_head = WeightStorage::F32(Arc::from(fuel_tensor::lazy::load_transposed_matrix(
            st,
            "lm_head.weight",
            cfg.vocab_size,
            h,
        )?));
        let lm_head_bias = fuel_tensor::lazy::load_tensor_as_f32(st, "lm_head.bias")
            .ok()
            .map(Arc::from);

        Ok(PhiWeights {
            token_embedding,
            layers,
            final_ln_gain,
            final_ln_bias,
            lm_head,
            lm_head_bias,
        })
    }
}

impl PhiModel {
    /// Download a Phi-2 checkpoint from the HuggingFace Hub and build
    /// a ready-to-forward `PhiModel`.
    pub fn from_hub(repo_id: &str) -> Result<Self> {
        let api = hf_hub::api::sync::Api::new()
            .map_err(|e| fuel_ir::error::Error::Msg(format!("hf-hub api init: {e}")))?;
        let repo = api.model(repo_id.to_string());

        let config_path = repo
            .get("config.json")
            .map_err(|e| fuel_ir::error::Error::Msg(format!("hf-hub config.json: {e}")))?;
        let config_str = std::fs::read_to_string(&config_path)?;
        let config = PhiConfig::from_hf_json_str(&config_str)?;

        let weight_paths: Vec<std::path::PathBuf> =
            match repo.get("model.safetensors.index.json") {
                Ok(index_path) => {
                    let index_str = std::fs::read_to_string(&index_path)?;
                    let index: serde_json::Value = serde_json::from_str(&index_str)
                        .map_err(|e| fuel_ir::error::Error::Msg(format!("parsing index: {e}")))?;
                    let weight_map = index
                        .get("weight_map")
                        .and_then(|x| x.as_object())
                        .ok_or_else(|| {
                            fuel_ir::error::Error::Msg("index: missing weight_map".into())
                        })?;
                    let mut unique = std::collections::HashSet::new();
                    for v in weight_map.values() {
                        if let Some(s) = v.as_str() {
                            unique.insert(s.to_string());
                        }
                    }
                    let mut paths = Vec::new();
                    for name in &unique {
                        paths.push(repo.get(name).map_err(|e| {
                            fuel_ir::error::Error::Msg(format!("hf-hub {name}: {e}"))
                        })?);
                    }
                    paths
                }
                Err(_) => vec![repo.get("model.safetensors").map_err(|e| {
                    fuel_ir::error::Error::Msg(format!("hf-hub model.safetensors: {e}"))
                })?],
            };

        let st = unsafe { fuel_loaders::safetensors::MmapedSafetensors::multi(&weight_paths) }?;
        let weights = PhiWeights::load_from_mmapped(&st, &config)?;

        Ok(PhiModel { config, weights })
    }
}

#[cfg(test)]
impl PhiConfig {
    /// The hand-rolled parser this type used before the `serde` split, kept
    /// verbatim and test-only as the differential's oracle.
    ///
    /// It is NOT dead code awaiting deletion: a differential that ran once at
    /// refactor time proves the two agreed at one commit, and a differential
    /// that keeps running proves they still do. Deleting this turns the first
    /// into the only claim anyone can make.
    pub(crate) fn from_hf_json_str_legacy(json: &str) -> Result<Self> {
        let v: serde_json::Value = serde_json::from_str(json)
            .map_err(|e| fuel_ir::error::Error::Msg(format!("parsing Phi config.json: {e}")))?;
        let get_usize = |key: &str| -> Result<usize> {
            v.get(key)
                .and_then(|x| x.as_u64())
                .map(|x| x as usize)
                .ok_or_else(|| {
                    fuel_ir::error::Error::Msg(format!("Phi config.json: missing/invalid {key:?}"))
                })
        };
        let get_f64 = |key: &str| -> Option<f64> { v.get(key).and_then(|x| x.as_f64()) };
        let get_bool = |key: &str| -> Option<bool> { v.get(key).and_then(|x| x.as_bool()) };

        let vocab_size = get_usize("vocab_size")?;
        let hidden_size = get_usize("hidden_size")?;
        let intermediate_size = get_usize("intermediate_size")?;
        let num_hidden_layers = get_usize("num_hidden_layers")?;
        let num_attention_heads = get_usize("num_attention_heads")?;
        let num_key_value_heads = v
            .get("num_key_value_heads")
            .and_then(|x| x.as_u64())
            .map(|x| x as usize);
        let head_dim = v
            .get("head_dim")
            .and_then(|x| x.as_u64())
            .map(|x| x as usize)
            .unwrap_or(hidden_size / num_attention_heads);
        let layer_norm_eps = get_f64("layer_norm_eps").unwrap_or(1e-5);
        let rope_theta = get_f64("rope_theta").unwrap_or(10_000.0);
        let max_position_embeddings = get_usize("max_position_embeddings")?;
        let partial_rotary_factor = get_f64("partial_rotary_factor").unwrap_or(0.4);
        let qk_layernorm = get_bool("qk_layernorm").unwrap_or(false);
        if qk_layernorm {
            return Err(fuel_ir::error::Error::Msg(
                "Phi: qk_layernorm=true is not yet supported by the standalone port (v1)".into(),
            )
            .bt());
        }

        Ok(PhiConfig {
            vocab_size,
            hidden_size,
            intermediate_size,
            num_hidden_layers,
            num_attention_heads,
            num_key_value_heads,
            head_dim,
            layer_norm_eps,
            rope_theta,
            max_position_embeddings,
            partial_rotary_factor,
            qk_layernorm,
            hidden_activation: PhiActivation::GeluPytorchTanh,
        })
    }
}

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

    fn tiny_weights(cfg: &PhiConfig) -> PhiWeights {
        let mut s: u32 = 31337;
        let next = move || -> 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 n_kv = cfg.num_kv_heads().unwrap();
        let kv = n_kv * cfg.head_dim;
        let mut nb: Box<dyn FnMut() -> f32> = Box::new(next);
        let token_embedding = vec_of(cfg.vocab_size * h, &mut *nb);

        let layers: Vec<PhiLayerWeights> = (0..cfg.num_hidden_layers)
            .map(|_| PhiLayerWeights {
                input_ln_gain: Arc::from(vec![1.0_f32; h]),
                input_ln_bias: Arc::from(vec![0.0_f32; h]),
                attn_q: WeightStorage::F32(vec_of(h * h, &mut *nb)),
                attn_q_bias: vec_of(h, &mut *nb),
                attn_k: WeightStorage::F32(vec_of(h * kv, &mut *nb)),
                attn_k_bias: vec_of(kv, &mut *nb),
                attn_v: WeightStorage::F32(vec_of(h * kv, &mut *nb)),
                attn_v_bias: vec_of(kv, &mut *nb),
                attn_dense: WeightStorage::F32(vec_of(h * h, &mut *nb)),
                attn_dense_bias: vec_of(h, &mut *nb),
                fc1: WeightStorage::F32(vec_of(h * i, &mut *nb)),
                fc1_bias: vec_of(i, &mut *nb),
                fc2: WeightStorage::F32(vec_of(i * h, &mut *nb)),
                fc2_bias: vec_of(h, &mut *nb),
            })
            .collect();
        let final_ln_gain = Arc::from(vec![1.0_f32; h]);
        let final_ln_bias = Arc::from(vec![0.0_f32; h]);
        let lm_head = WeightStorage::F32(vec_of(h * cfg.vocab_size, &mut *nb));
        let lm_head_bias = Some(vec_of(cfg.vocab_size, &mut *nb));
        PhiWeights {
            token_embedding,
            layers,
            final_ln_gain,
            final_ln_bias,
            lm_head,
            lm_head_bias,
        }
    }

    /// Every `config.json` fixture exercised in this module, as ONE corpus.
    ///
    /// Stated as a population rather than scattered through the tests: the
    /// differential below is only as good as this list, and a differential
    /// over three fixtures and one over three hundred are the same green tick.
    /// **Corpus: 4 fixtures, covering the sole `PhiConfig` parser in this file
    /// — the target is fully covered, which is a claim worth making rather
    /// than an absence worth assuming.**
    const DIFFERENTIAL_CORPUS: &[(&str, &str)] = &[
        (
            "phi-2 canonical (explicit head_dim, explicit kv_heads)",
            r#"{
                "vocab_size": 51200, "hidden_size": 2560, "intermediate_size": 10240,
                "num_hidden_layers": 32, "num_attention_heads": 32,
                "num_key_value_heads": 32, "head_dim": 80,
                "layer_norm_eps": 1e-5, "rope_theta": 10000.0,
                "max_position_embeddings": 2048, "partial_rotary_factor": 0.4
            }"#,
        ),
        (
            "minimal (every optional absent - exercises all defaults)",
            r#"{
                "vocab_size": 51200, "hidden_size": 1024, "intermediate_size": 4096,
                "num_hidden_layers": 16, "num_attention_heads": 16,
                "max_position_embeddings": 2048
            }"#,
        ),
        (
            "qk_layernorm = true (both paths must REJECT)",
            r#"{
                "vocab_size": 51200, "hidden_size": 1024, "intermediate_size": 4096,
                "num_hidden_layers": 16, "num_attention_heads": 16,
                "max_position_embeddings": 2048, "qk_layernorm": true
            }"#,
        ),
        (
            "explicit head_dim that DISAGREES with hidden_size/num_attention_heads",
            r#"{
                "vocab_size": 51200, "hidden_size": 2560, "intermediate_size": 10240,
                "num_hidden_layers": 32, "num_attention_heads": 32, "head_dim": 96,
                "max_position_embeddings": 2048
            }"#,
        ),
    ];

    // item-8-II membership guard for the phi FORK fix: num_kv_heads() now routes
    // through hf_config::num_key_value_heads instead of a hand-rolled unwrap_or.
    // Verified BEHAVIOURALLY (GAP-209) so the repoint is checked by a test rather
    // than by the two expressions looking alike; the load-bearing row is MQA.
    #[test]
    fn phi_num_kv_heads_defaults_to_num_heads_when_absent() {
        let json = r#"{
            "vocab_size": 51200, "hidden_size": 1024, "intermediate_size": 4096,
            "num_hidden_layers": 16, "num_attention_heads": 16, "max_position_embeddings": 2048
        }"#;
        let cfg = PhiConfig::from_hf_json_str(json).unwrap();
        assert_eq!(cfg.num_kv_heads().unwrap(), 16); // absent → num_attention_heads (MHA)
    }

    #[test]
    fn phi_num_kv_heads_preserves_true_mqa() {
        let json = r#"{
            "vocab_size": 51200, "hidden_size": 1024, "intermediate_size": 4096,
            "num_hidden_layers": 16, "num_attention_heads": 16, "num_key_value_heads": 1,
            "max_position_embeddings": 2048
        }"#;
        let cfg = PhiConfig::from_hf_json_str(json).unwrap();
        assert_eq!(cfg.num_kv_heads().unwrap(), 1); // Some(1) → 1: TRUE MQA survives
    }

    #[test]
    fn serde_path_agrees_with_the_legacy_parser_on_every_fixture() {
        assert_eq!(
            DIFFERENTIAL_CORPUS.len(),
            4,
            "corpus shrank - a differential silently covering less is the failure              mode this assert exists for",
        );
        for (name, json) in DIFFERENTIAL_CORPUS {
            let new = PhiConfig::from_hf_json_str(json);
            let old = PhiConfig::from_hf_json_str_legacy(json);
            match (new, old) {
                // VALUES, not shapes: PhiConfig's PartialEq compares every
                // field, so two configs with the same field SET and different
                // numbers fail here.
                (Ok(a), Ok(b)) => assert_eq!(a, b, "differential mismatch on {name}"),
                (Err(_), Err(_)) => {}
                (Ok(_), Err(e)) => panic!("{name}: serde path accepted, legacy rejected: {e}"),
                (Err(e), Ok(_)) => panic!("{name}: serde path rejected, legacy accepted: {e}"),
            }
        }
    }

    #[test]
    fn explicit_head_dim_survives_both_paths() {
        // The property the corpus alone cannot check: every in-tree fixture
        // that specifies head_dim specifies the QUOTIENT (phi-2 ships 80
        // against 2560/32), so a resolver that ignored the explicit value
        // would pass the differential. 96 != 2560/32 = 80.
        let json = DIFFERENTIAL_CORPUS[3].1;
        let cfg = PhiConfig::from_hf_json_str(json).unwrap();
        assert_eq!(
            cfg.head_dim, 96,
            "explicit head_dim must not be overwritten"
        );
        assert_ne!(cfg.head_dim, 2560 / 32);
        let legacy = PhiConfig::from_hf_json_str_legacy(json).unwrap();
        assert_eq!(legacy.head_dim, 96, "legacy must agree - it is the oracle");
    }

    #[test]
    fn from_hf_json_str_parses_canonical_phi2_fields() {
        // Excerpt from microsoft/phi-2/config.json (excluding the
        // wrapper fields like `model_type`, `transformers_version`,
        // etc. — only the architecture-relevant subset is needed).
        let json = r#"{
            "vocab_size": 51200,
            "hidden_size": 2560,
            "intermediate_size": 10240,
            "num_hidden_layers": 32,
            "num_attention_heads": 32,
            "num_key_value_heads": 32,
            "head_dim": 80,
            "layer_norm_eps": 1e-5,
            "rope_theta": 10000.0,
            "max_position_embeddings": 2048,
            "partial_rotary_factor": 0.4
        }"#;
        let cfg = PhiConfig::from_hf_json_str(json).unwrap();
        assert_eq!(cfg.vocab_size, 51200);
        assert_eq!(cfg.hidden_size, 2560);
        assert_eq!(cfg.intermediate_size, 10240);
        assert_eq!(cfg.num_hidden_layers, 32);
        assert_eq!(cfg.num_attention_heads, 32);
        assert_eq!(cfg.num_key_value_heads, Some(32));
        assert_eq!(cfg.head_dim, 80);
        assert!((cfg.layer_norm_eps - 1e-5).abs() < 1e-12);
        assert!((cfg.rope_theta - 10000.0).abs() < 1e-9);
        assert_eq!(cfg.max_position_embeddings, 2048);
        assert!((cfg.partial_rotary_factor - 0.4).abs() < 1e-9);
        assert!(!cfg.qk_layernorm);
        // rope_dim = (0.4 * 80) & !1 = 32 — even
        assert_eq!(cfg.rope_dim(), 32);
    }

    #[test]
    fn from_hf_json_str_applies_optional_defaults() {
        // Minimal Phi-2-like config (no rope_theta, no head_dim,
        // no partial_rotary_factor, no GQA).
        let json = r#"{
            "vocab_size": 51200,
            "hidden_size": 1024,
            "intermediate_size": 4096,
            "num_hidden_layers": 16,
            "num_attention_heads": 16,
            "max_position_embeddings": 2048
        }"#;
        let cfg = PhiConfig::from_hf_json_str(json).unwrap();
        assert_eq!(cfg.num_key_value_heads, None); // MHA — no GQA
        assert_eq!(cfg.head_dim, 64); // 1024 / 16
        assert!((cfg.layer_norm_eps - 1e-5).abs() < 1e-12);
        assert!((cfg.rope_theta - 10000.0).abs() < 1e-9);
        assert!((cfg.partial_rotary_factor - 0.4).abs() < 1e-9);
        assert!(!cfg.qk_layernorm);
    }

    #[test]
    fn from_hf_json_str_bails_on_qk_layernorm_true() {
        // v1 only supports qk_layernorm = false; setting it true must
        // fail explicitly (not silently produce wrong results).
        let json = r#"{
            "vocab_size": 51200,
            "hidden_size": 1024,
            "intermediate_size": 4096,
            "num_hidden_layers": 16,
            "num_attention_heads": 16,
            "max_position_embeddings": 2048,
            "qk_layernorm": true
        }"#;
        let result = PhiConfig::from_hf_json_str(json);
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(
            err.contains("qk_layernorm"),
            "expected error to mention qk_layernorm, got: {err}",
        );
    }

    fn tiny_config() -> PhiConfig {
        PhiConfig {
            vocab_size: 32,
            hidden_size: 16,
            intermediate_size: 32,
            num_hidden_layers: 2,
            num_attention_heads: 4,
            num_key_value_heads: None,
            head_dim: 4,
            layer_norm_eps: 1e-5,
            rope_theta: 10_000.0,
            max_position_embeddings: 64,
            partial_rotary_factor: 0.5, // rope_dim = 2
            qk_layernorm: false,
            hidden_activation: PhiActivation::GeluPytorchTanh,
        }
    }

    #[test]
    fn forward_shape_and_finite() {
        let cfg = tiny_config();
        let model = PhiModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let tokens: Vec<u32> = vec![1, 2, 3, 4, 5];
        let logits = model.forward(&tokens, 0).unwrap();
        assert_eq!(logits.shape().dims(), &[1, tokens.len(), cfg.vocab_size]);
        for &v in &logits.realize_f32() {
            assert!(v.is_finite(), "got non-finite logit {v}");
        }
    }

    /// No-lm-head-bias checkpoints (some Phi fine-tunes drop the
    /// bias) must run without panic via the `None` branch.
    /// Regression guard for the previous panicking design where
    /// `lm_head_bias` was a non-Optional `Arc<[f32]>`.
    #[test]
    fn forward_without_lm_head_bias() {
        let cfg = tiny_config();
        let mut weights = tiny_weights(&cfg);
        weights.lm_head_bias = None;
        let model = PhiModel {
            config: cfg.clone(),
            weights,
        };
        let tokens: Vec<u32> = vec![1, 2, 3];
        let logits = model.forward(&tokens, 0).unwrap();
        assert_eq!(logits.shape().dims(), &[1, tokens.len(), cfg.vocab_size]);
        for &v in &logits.realize_f32() {
            assert!(v.is_finite(), "got non-finite logit {v}");
        }
    }

    /// With lm_head_bias = Some(zeros), output must equal the
    /// no-bias forward. Verifies the Option branch produces
    /// arithmetically equivalent results.
    #[test]
    fn lm_head_bias_zero_matches_none() {
        let cfg = tiny_config();
        let base = tiny_weights(&cfg);

        // Path A: bias = None.
        let mut weights_none = base.clone();
        weights_none.lm_head_bias = None;
        let m_none = PhiModel {
            config: cfg.clone(),
            weights: weights_none,
        };
        let a = m_none.forward(&[1, 2, 3], 0).unwrap().realize_f32();

        // Path B: bias = Some(zeros).
        let mut weights_zero = base.clone();
        weights_zero.lm_head_bias = Some(Arc::from(vec![0.0_f32; cfg.vocab_size]));
        let m_zero = PhiModel {
            config: cfg.clone(),
            weights: weights_zero,
        };
        let b = m_zero.forward(&[1, 2, 3], 0).unwrap().realize_f32();

        assert_eq!(a.len(), b.len());
        for (i, (av, bv)) in a.iter().zip(b.iter()).enumerate() {
            assert!(
                (av - bv).abs() < 1e-5,
                "logit[{i}]: none={av} vs zero-bias={bv}"
            );
        }
    }

    /// rope_dim equal to head_dim must run (full rotary degenerate case).
    #[test]
    fn full_rotary_factor() {
        let mut cfg = tiny_config();
        cfg.partial_rotary_factor = 1.0;
        assert_eq!(cfg.rope_dim(), cfg.head_dim);
        let model = PhiModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let logits = model.forward(&[1, 2, 3], 0).unwrap().realize_f32();
        assert_eq!(logits.len(), 3 * cfg.vocab_size);
    }

    /// Parallel block must NOT collapse to sequential: differentiating
    /// attention and MLP at the residual sum is the structural test.
    /// Zero out fc1/fc2 (kills MLP path) and confirm output equals
    /// `residual + attn`. Then zero out Q/K/V/dense (kills attn) and
    /// confirm output equals `residual + mlp`. Both must be possible
    /// independently.
    #[test]
    fn parallel_paths_are_additive() {
        let cfg = PhiConfig {
            num_hidden_layers: 1,
            ..tiny_config()
        };
        let base = tiny_weights(&cfg);

        let mut no_mlp = base.clone();
        let zero = |n: usize| WeightStorage::F32(Arc::from(vec![0.0_f32; n]));
        let zero_b = |n: usize| Arc::from(vec![0.0_f32; n]) as Arc<[f32]>;
        no_mlp.layers[0].fc1 = zero(cfg.hidden_size * cfg.intermediate_size);
        no_mlp.layers[0].fc1_bias = zero_b(cfg.intermediate_size);
        no_mlp.layers[0].fc2 = zero(cfg.intermediate_size * cfg.hidden_size);
        no_mlp.layers[0].fc2_bias = zero_b(cfg.hidden_size);

        let mut no_attn = base.clone();
        no_attn.layers[0].attn_q = zero(cfg.hidden_size * cfg.hidden_size);
        no_attn.layers[0].attn_q_bias = zero_b(cfg.hidden_size);
        no_attn.layers[0].attn_k = zero(cfg.hidden_size * cfg.hidden_size);
        no_attn.layers[0].attn_k_bias = zero_b(cfg.hidden_size);
        no_attn.layers[0].attn_v = zero(cfg.hidden_size * cfg.hidden_size);
        no_attn.layers[0].attn_v_bias = zero_b(cfg.hidden_size);
        no_attn.layers[0].attn_dense = zero(cfg.hidden_size * cfg.hidden_size);
        no_attn.layers[0].attn_dense_bias = zero_b(cfg.hidden_size);

        // Both should produce finite, non-trivial outputs different
        // from each other (different paths active).
        let m_no_mlp = PhiModel {
            config: cfg.clone(),
            weights: no_mlp,
        };
        let m_no_attn = PhiModel {
            config: cfg.clone(),
            weights: no_attn,
        };
        let toks = [3_u32, 7, 2];
        let a = m_no_mlp.forward(&toks, 0).unwrap().realize_f32();
        let b = m_no_attn.forward(&toks, 0).unwrap().realize_f32();
        assert!(a.iter().all(|v| v.is_finite()));
        assert!(b.iter().all(|v| v.is_finite()));
        let mut max_diff = 0.0_f32;
        for (av, bv) in a.iter().zip(b.iter()) {
            max_diff = max_diff.max((av - bv).abs());
        }
        assert!(
            max_diff > 1e-6,
            "no-mlp vs no-attn must produce different outputs: {max_diff}"
        );
    }

    #[test]
    fn num_kv_heads_default_is_num_attention_heads() {
        let cfg = tiny_config();
        assert_eq!(cfg.num_kv_heads().unwrap(), cfg.num_attention_heads);
    }

    /// `forward_hidden` returns post-LayerNorm hidden states
    /// `(1, seq, hidden_size)` without the lm_head matmul or
    /// its bias.
    #[test]
    fn forward_hidden_shape_and_finite() {
        let cfg = tiny_config();
        let model = PhiModel {
            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}");
        }
    }

    #[test]
    fn forward_embeds_matches_forward_after_token_lookup() {
        let cfg = tiny_config();
        let model = PhiModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let tokens: Vec<u32> = vec![1, 2, 3];
        let logits_ref = model.forward(&tokens, 0).unwrap().realize_f32();
        let anchor =
            Tensor::from_f32(vec![0.0_f32], Shape::from_dims(&[1]), &Device::cpu()).unwrap();
        let embeds = model.embed_tokens_anchored(&anchor, &tokens).unwrap();
        let logits_via_embeds = model.forward_embeds(&embeds, 0).unwrap().realize_f32();
        let max_diff = logits_ref
            .iter()
            .zip(logits_via_embeds.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f32, f32::max);
        assert!(
            max_diff < 1e-5,
            "Phi forward vs forward_embeds must agree (max diff {max_diff})"
        );
    }

    #[test]
    fn forward_embeds_rejects_bad_shape() {
        let cfg = tiny_config();
        let model = PhiModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let bad = Tensor::from_f32(
            vec![0.0_f32; 3 * (cfg.hidden_size + 1)],
            Shape::from_dims(&[1, 3, cfg.hidden_size + 1]),
            &Device::cpu(),
        )
        .unwrap();
        assert!(model.forward_embeds(&bad, 0).is_err());
    }

    #[test]
    fn forward_hidden_embeds_matches_forward_hidden() {
        let cfg = tiny_config();
        let model = PhiModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let tokens: Vec<u32> = vec![5, 7];
        let h_ref = model.forward_hidden(&tokens, 0).unwrap().realize_f32();
        let anchor =
            Tensor::from_f32(vec![0.0_f32], Shape::from_dims(&[1]), &Device::cpu()).unwrap();
        let embeds = model.embed_tokens_anchored(&anchor, &tokens).unwrap();
        let h_via_embeds = model
            .forward_hidden_embeds(&embeds, 0)
            .unwrap()
            .realize_f32();
        let max_diff = h_ref
            .iter()
            .zip(h_via_embeds.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f32, f32::max);
        assert!(
            max_diff < 1e-5,
            "Phi forward_hidden vs forward_hidden_embeds must agree (max diff {max_diff})"
        );
    }
}