fuel-transformers 0.19.1

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Gemma 4 text decoder ported to the lazy-graph API.
//!
//! Phase D specialized port. Gemma 4 is a **multimodal** model
//! (text + vision + audio). v1 ports the **text-only** decoder;
//! the vision tower (ViT-style encoder) and audio tower (audio
//! mel-spec encoder) plus the multimodal embedding integration
//! are deferred to follow-up sessions.
//!
//! The text decoder extends [`crate::models::lazy_gemma3`] with several
//! architectural wrinkles:
//!
//!   1. **Per-layer-type head_dim and num_kv_heads.** Sliding
//!      layers use `head_dim` + `num_key_value_heads`; full
//!      ("global") layers use `global_head_dim` +
//!      `num_global_key_value_heads`. This is configured by a
//!      per-layer `layer_types: Vec<Gemma4LayerType>`.
//!   2. **Partial rotary on global layers.** Only the first
//!      `partial_rotary_factor * head_dim` features per head
//!      are rotated; the rest pass through unchanged. Sliding
//!      layers use full rotary.
//!   3. **Per-layer-type RoPE base.** Sliding layers use
//!      `rope_local_base_freq` (~10_000); global layers use
//!      `rope_theta` (~1_000_000).
//!   4. **V normalization** (pure RmsNorm without a learned
//!      weight, just rsqrt of mean-of-squares). Applied
//!      *after* the V projection and reshape, before attention.
//!   5. **Per-head Q/K RmsNorm with `(gain + 1)` offset**
//!      (Gemma family convention; same as Gemma3).
//!   6. **4-norm block structure** (Gemma3-shape):
//!      input + post-attn + pre-FFN + post-FFN.
//!   7. **`query_pre_attn_scalar`** — Gemma4 honors the eager
//!      reference's choice to compute scale as
//!      `1 / sqrt(head_dim)` (the divisor matches the actual
//!      head_dim per layer type); the config-stored
//!      `query_pre_attn_scalar` field is unused by the eager
//!      forward, so v1 ignores it.
//!   8. **Final-logit soft-cap** when `final_logit_softcapping`
//!      is set (no per-layer attention-score cap in Gemma4 unlike
//!      Gemma3 — that field is absent).
//!
//! # Scope (v1)
//!
//! Forward-only, single sequence (`batch == 1`), no KV cache,
//! F32. Multimodal injection (image and audio embeddings
//! interleaved with text token embeddings via mask) is deferred.

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

pub use crate::models::lazy_gemma::GemmaActivation as Gemma4Activation;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Gemma4LayerType {
    SlidingAttention,
    FullAttention,
}

#[derive(Debug, Clone, PartialEq)]
pub struct Gemma4TextConfig {
    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,
    /// For full-attention layers (default = `num_key_value_heads`).
    pub num_global_key_value_heads: Option<usize>,
    /// Head dim for sliding layers.
    pub head_dim: usize,
    /// Head dim for full-attention layers.
    pub global_head_dim: usize,
    pub rms_norm_eps: f64,
    /// RoPE base for full-attention layers (default 1_000_000).
    pub rope_theta: f64,
    /// RoPE base for sliding layers (default 10_000).
    pub rope_local_base_freq: f64,
    /// Fraction of `global_head_dim` rotated on full-attn layers.
    pub partial_rotary_factor: f64,
    pub sliding_window: usize,
    pub layer_types: Vec<Gemma4LayerType>,
    pub attention_bias: bool,
    pub hidden_activation: Gemma4Activation,
    pub final_logit_softcapping: Option<f64>,
    pub tie_word_embeddings: bool,
}

impl Gemma4TextConfig {
    pub fn num_global_kv(&self) -> usize {
        self.num_global_key_value_heads
            .unwrap_or(self.num_key_value_heads)
    }
    pub fn layer_type(&self, layer_idx: usize) -> Gemma4LayerType {
        self.layer_types
            .get(layer_idx)
            .copied()
            .unwrap_or(Gemma4LayerType::SlidingAttention)
    }
    pub fn global_rope_dim(&self) -> usize {
        let r = (self.partial_rotary_factor * self.global_head_dim as f64) as usize;
        r & !1
    }
}

#[derive(Debug, Clone)]
pub struct Gemma4LayerWeights {
    pub attn_q: WeightStorage,
    pub attn_q_bias: Option<Arc<[f32]>>,
    pub attn_k: WeightStorage,
    pub attn_k_bias: Option<Arc<[f32]>>,
    pub attn_v: WeightStorage,
    pub attn_v_bias: Option<Arc<[f32]>>,
    pub attn_o: WeightStorage,
    pub attn_o_bias: Option<Arc<[f32]>>,
    /// Per-head Q RmsNorm gain on the layer-effective head_dim.
    pub q_norm_gain: Arc<[f32]>,
    pub k_norm_gain: Arc<[f32]>,
    pub input_norm_gain: Arc<[f32]>,
    pub post_attn_norm_gain: Arc<[f32]>,
    pub pre_ffn_norm_gain: Arc<[f32]>,
    pub post_ffn_norm_gain: Arc<[f32]>,
    pub ffn_gate: WeightStorage,
    pub ffn_up: WeightStorage,
    pub ffn_down: WeightStorage,
}

#[derive(Debug, Clone)]
pub struct Gemma4TextWeights {
    pub token_embedding: Arc<[f32]>,
    pub layers: Vec<Gemma4LayerWeights>,
    pub final_norm_gain: Arc<[f32]>,
}

#[derive(Debug, Clone)]
pub struct Gemma4TextModel {
    pub config: Gemma4TextConfig,
    pub weights: Gemma4TextWeights,
}

impl Gemma4TextModel {
    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 offset RmsNorm
    /// and return per-token hidden states `(1, seq, hidden_size)`.
    /// Gemma4-text-specific: alternating
    /// `SlidingAttention` / `FullAttention` layer types each
    /// with their own head_dim, num_kv_heads, RoPE table, and
    /// mask — all honored by the shared backbone.
    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 (typically the
    /// concatenation of vision/audio-projected embeddings + text
    /// token embeddings).
    ///
    /// `scaled_embeds` shape: `(1, seq, hidden_size)`. The caller
    /// must apply Gemma's `sqrt(hidden_size)` scaling before
    /// invoking — matching the convention used by lazy_paligemma /
    /// lazy_llava / lazy_voxtral so the multimodal composition
    /// layer owns the scaling decision.
    ///
    /// Returns logits `(1, seq, vocab_size)` (with optional logit
    /// softcap applied if the config enables it).
    pub fn forward_embeds(&self, scaled_embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        let h_norm = self.decode_from_scaled_embeds(scaled_embeds, start_pos)?;
        self.apply_lm_head(&h_norm)
    }

    /// Hidden-state variant of [`Self::forward_embeds`]. Returns the
    /// post-final-RmsNorm states `(1, seq, hidden_size)`. Used by
    /// retrieval / embedding consumers that project hidden states
    /// rather than predict tokens.
    pub fn forward_hidden_embeds(
        &self,
        scaled_embeds: &Tensor,
        start_pos: usize,
    ) -> Result<Tensor> {
        self.decode_from_scaled_embeds(scaled_embeds, start_pos)
    }

    /// Build per-token embeddings without running the decoder. Used by
    /// the multimodal composition to obtain text-side embeddings that
    /// will be spliced with vision/audio features before
    /// [`Self::forward_embeds`] is called.
    ///
    /// Returns shape `(1, seq, hidden_size)`. The caller is responsible
    /// for the `sqrt(hidden_size)` scaling (matches PaliGemma/LLaVA
    /// convention so the composition owns the splice ordering).
    ///
    /// `anchor` selects the graph the const-emitted embedding table
    /// lives on, so the result can be concatenated with image/audio
    /// embeddings on the same graph.
    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 lm_head = WeightStorage::F32(self.weights.token_embedding.clone());
        let logits = lm_head.apply_linear(h_norm, cfg.hidden_size, cfg.vocab_size)?;
        match cfg.final_logit_softcapping {
            None => Ok(logits),
            Some(sc) => Ok(logits.mul_scalar(1.0 / sc).tanh().mul_scalar(sc)),
        }
    }

    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);

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

    fn decode_from_scaled_embeds(
        &self,
        scaled_embeds: &Tensor,
        start_pos: usize,
    ) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let dims = scaled_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!(
                "Gemma4TextModel::forward_embeds: expected scaled_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(
                "Gemma4TextModel::forward_embeds: seq must be > 0".into(),
            )
            .bt());
        }
        if cfg.num_hidden_layers == 0 {
            return Err(fuel_ir::error::Error::Msg(
                "Gemma4TextModel::forward_embeds: num_hidden_layers must be > 0".into(),
            )
            .bt());
        }
        if cfg.layer_types.len() != cfg.num_hidden_layers {
            return Err(fuel_ir::error::Error::Msg(format!(
                "Gemma4TextModel::forward_embeds: layer_types length ({}) \
                 must match num_hidden_layers ({})",
                cfg.layer_types.len(),
                cfg.num_hidden_layers,
            ))
            .bt());
        }
        let mut h = scaled_embeds.clone();

        let (rope_cos_l, rope_sin_l) =
            h.rope_tables_const(cfg.rope_local_base_freq, start_pos, seq, cfg.head_dim);

        let rope_dim_global = cfg.global_rope_dim();
        let (rope_cos_g, rope_sin_g) =
            h.rope_tables_const(cfg.rope_theta, start_pos, seq, rope_dim_global);

        let full_mask = self.build_mask(&h, seq, None);
        let sliding_mask = self.build_mask(&h, seq, Some(cfg.sliding_window));

        for (idx, layer) in weights.layers.iter().enumerate() {
            let kind = cfg.layer_type(idx);
            let (head_dim, num_kv, rope_cos, rope_sin, rope_dim, mask, is_global) = match kind {
                Gemma4LayerType::SlidingAttention => (
                    cfg.head_dim,
                    cfg.num_key_value_heads,
                    &rope_cos_l,
                    &rope_sin_l,
                    cfg.head_dim,
                    &sliding_mask,
                    false,
                ),
                Gemma4LayerType::FullAttention => (
                    cfg.global_head_dim,
                    cfg.num_global_kv(),
                    &rope_cos_g,
                    &rope_sin_g,
                    rope_dim_global,
                    &full_mask,
                    true,
                ),
            };
            h = self.apply_layer(
                &h, layer, head_dim, num_kv, rope_cos, rope_sin, rope_dim, mask, is_global,
            )?;
        }
        h.rms_norm_affine_with_offset(&weights.final_norm_gain, 1.0, cfg.rms_norm_eps)
    }

    fn build_mask(&self, anchor: &Tensor, seq: usize, sliding: Option<usize>) -> Tensor {
        let window = sliding.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_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",
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn apply_layer(
        &self,
        x: &Tensor,
        layer: &Gemma4LayerWeights,
        head_dim: usize,
        num_kv: usize,
        rope_cos: &Tensor,
        rope_sin: &Tensor,
        rope_dim: usize,
        mask: &Tensor,
        _is_global: bool,
    ) -> 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_heads = cfg.num_attention_heads;
        let q_dim = n_heads * head_dim;
        let kv_dim = num_kv * head_dim;

        // Pre-attention norm.
        let residual = x.clone();
        let x_norm =
            x.rms_norm_affine_with_offset(&layer.input_norm_gain, 1.0, cfg.rms_norm_eps)?;

        // Q / K / V projections.
        let q = layer
            .attn_q
            .apply_linear(&x_norm, cfg.hidden_size, q_dim)?
            .add_optional_trailing_bias(layer.attn_q_bias.as_ref())?;
        let k = layer
            .attn_k
            .apply_linear(&x_norm, cfg.hidden_size, kv_dim)?
            .add_optional_trailing_bias(layer.attn_k_bias.as_ref())?;
        let v = layer
            .attn_v
            .apply_linear(&x_norm, cfg.hidden_size, kv_dim)?
            .add_optional_trailing_bias(layer.attn_v_bias.as_ref())?;

        let _ = (batch, seq);
        let q = q.split_heads(n_heads, head_dim)?;
        let k = k.split_heads(num_kv, head_dim)?;
        let v = v.split_heads(num_kv, head_dim)?;

        // Per-head Q/K RmsNorm with `(gain + 1)` offset.
        let q = q.rms_norm_affine_with_offset(&layer.q_norm_gain, 1.0, cfg.rms_norm_eps)?;
        let k = k.rms_norm_affine_with_offset(&layer.k_norm_gain, 1.0, cfg.rms_norm_eps)?;
        // V normalization: pure RmsNorm without a learned weight.
        let v = v_rms_norm(&v, cfg.rms_norm_eps)?;

        // RoPE. Global uses partial rotary; sliding uses full.
        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 expand.
        let n_rep = n_heads / num_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 / (head_dim as f64).sqrt();
        let scores = q_r.matmul(&k_t)?;
        let scores_scaled = scores.mul_scalar(scale);
        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_o
            .apply_linear(&merged, q_dim, cfg.hidden_size)?
            .add_optional_trailing_bias(layer.attn_o_bias.as_ref())?;
        let attn_out_norm = attn_out.rms_norm_affine_with_offset(
            &layer.post_attn_norm_gain,
            1.0,
            cfg.rms_norm_eps,
        )?;
        let h1 = residual.add(&attn_out_norm)?;

        // FFN sublayer.
        let residual2 = h1.clone();
        let h1_norm =
            h1.rms_norm_affine_with_offset(&layer.pre_ffn_norm_gain, 1.0, cfg.rms_norm_eps)?;
        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 activated = match cfg.hidden_activation {
            Gemma4Activation::Gelu => gate.gelu(),
            Gemma4Activation::GeluPytorchTanh => gate.gelu_tanh(),
        };
        let ffn_in = activated.mul(&up)?;
        let ffn_out =
            layer
                .ffn_down
                .apply_linear(&ffn_in, cfg.intermediate_size, cfg.hidden_size)?;
        let ffn_out_norm = ffn_out.rms_norm_affine_with_offset(
            &layer.post_ffn_norm_gain,
            1.0,
            cfg.rms_norm_eps,
        )?;
        residual2.add(&ffn_out_norm)
    }
}

/// V RmsNorm: pure rsqrt(mean of squares + eps), no learned weight.
fn v_rms_norm(v: &Tensor, eps: f64) -> Result<Tensor> {
    v.rms_norm_last_dim(eps)
}

impl Gemma4TextWeights {
    /// Load Gemma 4 text-decoder weights from a `MmapedSafetensors`
    /// file using the standard HuggingFace naming. Gemma 4 ties
    /// `lm_head.weight` to `model.embed_tokens.weight` — there is no
    /// separate output projection field on `Gemma4TextWeights` (the
    /// `apply_lm_head` path reuses the embedding table directly).
    ///
    /// Tensor names mirrored from `fuel_transformers::models::llm::gemma4::text`:
    ///   - `model.embed_tokens.weight` → [`Gemma4TextWeights::token_embedding`]
    ///   - `model.layers.{i}.self_attn.{q,k,v,o}_proj.weight`
    ///   - `model.layers.{i}.self_attn.{q,k,v,o}_proj.bias`
    ///     (loaded only when `attention_bias == true`)
    ///   - `model.layers.{i}.self_attn.q_norm.weight` → `q_norm_gain`
    ///   - `model.layers.{i}.self_attn.k_norm.weight` → `k_norm_gain`
    ///   - `model.layers.{i}.input_layernorm.weight` → `input_norm_gain`
    ///   - `model.layers.{i}.post_attention_layernorm.weight` → `post_attn_norm_gain`
    ///   - `model.layers.{i}.pre_feedforward_layernorm.weight` → `pre_ffn_norm_gain`
    ///   - `model.layers.{i}.post_feedforward_layernorm.weight` → `post_ffn_norm_gain`
    ///   - `model.layers.{i}.mlp.{gate,up,down}_proj.weight` → `ffn_{gate,up,down}`
    ///   - `model.norm.weight` → `final_norm_gain`
    ///
    /// Per-layer Q/K/V/O dims are layer-type dependent: sliding layers
    /// use `cfg.head_dim` + `cfg.num_key_value_heads`; full ("global")
    /// layers use `cfg.global_head_dim` + `cfg.num_global_kv()`. The
    /// q_norm/k_norm weights are sized to the layer-effective head_dim.
    pub fn load_from_mmapped(
        st: &fuel_loaders::safetensors::MmapedSafetensors,
        cfg: &Gemma4TextConfig,
    ) -> Result<Self> {
        use fuel_tensor::{lazy::load_tensor_as_f32, lazy::load_transposed_matrix_preserve_dtype};

        let h = cfg.hidden_size;
        let i_dim = cfg.intermediate_size;
        let n_heads = cfg.num_attention_heads;

        if cfg.layer_types.len() != cfg.num_hidden_layers {
            fuel_ir::bail!(
                "Gemma4TextWeights::load_from_mmapped: layer_types length ({}) \
                 must match num_hidden_layers ({})",
                cfg.layer_types.len(),
                cfg.num_hidden_layers,
            );
        }

        let token_embedding = load_tensor_as_f32(st, "model.embed_tokens.weight")?;
        if token_embedding.len() != cfg.vocab_size * h {
            fuel_ir::bail!(
                "model.embed_tokens.weight: {} elts, expected {} ({}×{})",
                token_embedding.len(),
                cfg.vocab_size * h,
                cfg.vocab_size,
                h,
            );
        }

        let mut layers: Vec<Gemma4LayerWeights> = Vec::with_capacity(cfg.num_hidden_layers);
        for li in 0..cfg.num_hidden_layers {
            let kind = cfg.layer_type(li);
            let (head_dim, num_kv) = match kind {
                Gemma4LayerType::SlidingAttention => (cfg.head_dim, cfg.num_key_value_heads),
                Gemma4LayerType::FullAttention => (cfg.global_head_dim, cfg.num_global_kv()),
            };
            let q_dim = n_heads * head_dim;
            let kv_dim = num_kv * head_dim;

            let p = format!("model.layers.{li}");
            let attn_q = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.self_attn.q_proj.weight"),
                q_dim,
                h,
            )?;
            let attn_k = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.self_attn.k_proj.weight"),
                kv_dim,
                h,
            )?;
            let attn_v = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.self_attn.v_proj.weight"),
                kv_dim,
                h,
            )?;
            let attn_o = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.self_attn.o_proj.weight"),
                h,
                q_dim,
            )?;
            let attn_q_bias = if cfg.attention_bias {
                load_tensor_as_f32(st, &format!("{p}.self_attn.q_proj.bias"))
                    .ok()
                    .map(Arc::from)
            } else {
                None
            };
            let attn_k_bias = if cfg.attention_bias {
                load_tensor_as_f32(st, &format!("{p}.self_attn.k_proj.bias"))
                    .ok()
                    .map(Arc::from)
            } else {
                None
            };
            let attn_v_bias = if cfg.attention_bias {
                load_tensor_as_f32(st, &format!("{p}.self_attn.v_proj.bias"))
                    .ok()
                    .map(Arc::from)
            } else {
                None
            };
            let attn_o_bias = if cfg.attention_bias {
                load_tensor_as_f32(st, &format!("{p}.self_attn.o_proj.bias"))
                    .ok()
                    .map(Arc::from)
            } else {
                None
            };
            let q_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.self_attn.q_norm.weight"),
            )?);
            let k_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.self_attn.k_norm.weight"),
            )?);
            let input_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.input_layernorm.weight"),
            )?);
            let post_attn_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.post_attention_layernorm.weight"),
            )?);
            let pre_ffn_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.pre_feedforward_layernorm.weight"),
            )?);
            let post_ffn_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.post_feedforward_layernorm.weight"),
            )?);
            let ffn_gate = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.mlp.gate_proj.weight"),
                i_dim,
                h,
            )?;
            let ffn_up = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.mlp.up_proj.weight"),
                i_dim,
                h,
            )?;
            let ffn_down = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.mlp.down_proj.weight"),
                h,
                i_dim,
            )?;
            layers.push(Gemma4LayerWeights {
                attn_q,
                attn_q_bias,
                attn_k,
                attn_k_bias,
                attn_v,
                attn_v_bias,
                attn_o,
                attn_o_bias,
                q_norm_gain,
                k_norm_gain,
                input_norm_gain,
                post_attn_norm_gain,
                pre_ffn_norm_gain,
                post_ffn_norm_gain,
                ffn_gate,
                ffn_up,
                ffn_down,
            });
        }

        let final_norm_gain = Arc::from(load_tensor_as_f32(st, "model.norm.weight")?);

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

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

    fn vec_of(n: usize, next: &mut dyn FnMut() -> f32) -> Arc<[f32]> {
        Arc::from((0..n).map(|_| next()).collect::<Vec<_>>())
    }

    fn tiny_weights(cfg: &Gemma4TextConfig) -> Gemma4TextWeights {
        let mut s: u32 = 42949;
        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 h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        let mut nb: Box<dyn FnMut() -> f32> = Box::new(next);
        let token_embedding = vec_of(cfg.vocab_size * h, &mut *nb);

        let layers: Vec<Gemma4LayerWeights> = (0..cfg.num_hidden_layers)
            .map(|idx| {
                let kind = cfg.layer_type(idx);
                let (head_dim, num_kv) = match kind {
                    Gemma4LayerType::SlidingAttention => (cfg.head_dim, cfg.num_key_value_heads),
                    Gemma4LayerType::FullAttention => (cfg.global_head_dim, cfg.num_global_kv()),
                };
                let q_dim = cfg.num_attention_heads * head_dim;
                let kv_dim = num_kv * head_dim;
                Gemma4LayerWeights {
                    attn_q: WeightStorage::F32(vec_of(h * q_dim, &mut *nb)),
                    attn_q_bias: if cfg.attention_bias {
                        Some(vec_of(q_dim, &mut *nb))
                    } else {
                        None
                    },
                    attn_k: WeightStorage::F32(vec_of(h * kv_dim, &mut *nb)),
                    attn_k_bias: if cfg.attention_bias {
                        Some(vec_of(kv_dim, &mut *nb))
                    } else {
                        None
                    },
                    attn_v: WeightStorage::F32(vec_of(h * kv_dim, &mut *nb)),
                    attn_v_bias: if cfg.attention_bias {
                        Some(vec_of(kv_dim, &mut *nb))
                    } else {
                        None
                    },
                    attn_o: WeightStorage::F32(vec_of(q_dim * h, &mut *nb)),
                    attn_o_bias: if cfg.attention_bias {
                        Some(vec_of(h, &mut *nb))
                    } else {
                        None
                    },
                    q_norm_gain: Arc::from(vec![0.05_f32; head_dim]),
                    k_norm_gain: Arc::from(vec![0.05_f32; head_dim]),
                    input_norm_gain: Arc::from(vec![0.05_f32; h]),
                    post_attn_norm_gain: Arc::from(vec![0.05_f32; h]),
                    pre_ffn_norm_gain: Arc::from(vec![0.05_f32; h]),
                    post_ffn_norm_gain: Arc::from(vec![0.05_f32; h]),
                    ffn_gate: WeightStorage::F32(vec_of(h * i, &mut *nb)),
                    ffn_up: WeightStorage::F32(vec_of(h * i, &mut *nb)),
                    ffn_down: WeightStorage::F32(vec_of(i * h, &mut *nb)),
                }
            })
            .collect();
        let final_norm_gain = Arc::from(vec![0.05_f32; h]);
        Gemma4TextWeights {
            token_embedding,
            layers,
            final_norm_gain,
        }
    }

    fn tiny_config() -> Gemma4TextConfig {
        // 4 layers: pattern [sliding, sliding, sliding, full].
        Gemma4TextConfig {
            vocab_size: 32,
            hidden_size: 24,
            intermediate_size: 32,
            num_hidden_layers: 4,
            num_attention_heads: 4,
            num_key_value_heads: 2,
            num_global_key_value_heads: Some(2),
            head_dim: 4,
            global_head_dim: 8,
            rms_norm_eps: 1e-6,
            rope_theta: 1_000_000.0,
            rope_local_base_freq: 10_000.0,
            partial_rotary_factor: 0.5, // global_rope_dim = 4
            sliding_window: 3,
            layer_types: vec![
                Gemma4LayerType::SlidingAttention,
                Gemma4LayerType::SlidingAttention,
                Gemma4LayerType::SlidingAttention,
                Gemma4LayerType::FullAttention,
            ],
            attention_bias: false,
            hidden_activation: Gemma4Activation::GeluPytorchTanh,
            final_logit_softcapping: Some(30.0),
            tie_word_embeddings: true,
        }
    }

    #[test]
    fn forward_shape_and_finite() {
        let cfg = tiny_config();
        let model = Gemma4TextModel {
            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}");
        }
    }

    /// Per-layer head_dim selection: the global layer's head_dim
    /// is 8 while sliding layers use 4. The model must still
    /// produce coherent output. Swap the final layer from full
    /// to sliding (different head_dim, different num_kv,
    /// different RoPE table) — output must differ.
    #[test]
    fn per_layer_head_dim_alters_output() {
        let cfg_a = tiny_config();
        let cfg_b = Gemma4TextConfig {
            layer_types: vec![
                Gemma4LayerType::SlidingAttention,
                Gemma4LayerType::SlidingAttention,
                Gemma4LayerType::SlidingAttention,
                Gemma4LayerType::SlidingAttention,
            ],
            ..cfg_a.clone()
        };
        let weights_a = tiny_weights(&cfg_a);
        let weights_b = tiny_weights(&cfg_b);
        let m_a = Gemma4TextModel {
            config: cfg_a,
            weights: weights_a,
        };
        let m_b = Gemma4TextModel {
            config: cfg_b,
            weights: weights_b,
        };
        let toks = [1_u32, 2, 3];
        let a = m_a.forward(&toks, 0).unwrap().realize_f32();
        let b = m_b.forward(&toks, 0).unwrap().realize_f32();
        let mut max_diff = 0.0_f32;
        for (x, y) in a.iter().zip(b.iter()) {
            max_diff = max_diff.max((x - y).abs());
        }
        assert!(
            max_diff > 1e-6,
            "layer-type change must alter output, max_diff = {max_diff}"
        );
    }

    /// Final logit soft-cap must measurably change output.
    #[test]
    fn final_softcap_changes_output() {
        let cfg_no = Gemma4TextConfig {
            final_logit_softcapping: None,
            ..tiny_config()
        };
        let cfg_yes = Gemma4TextConfig {
            final_logit_softcapping: Some(5.0),
            ..tiny_config()
        };
        let weights = tiny_weights(&cfg_no);
        let m_no = Gemma4TextModel {
            config: cfg_no,
            weights: weights.clone(),
        };
        let m_yes = Gemma4TextModel {
            config: cfg_yes,
            weights,
        };
        let toks = [3_u32, 7, 11];
        let a = m_no.forward(&toks, 0).unwrap().realize_f32();
        let b = m_yes.forward(&toks, 0).unwrap().realize_f32();
        let mut max_diff = 0.0_f32;
        for (x, y) in a.iter().zip(b.iter()) {
            max_diff = max_diff.max((x - y).abs());
        }
        assert!(
            max_diff > 1e-6,
            "final logit soft-cap must alter output, max_diff = {max_diff}"
        );
    }

    /// V normalization is wired: a different magnitude of input
    /// to the V projection should still produce sane output. We
    /// just check that the v_rms_norm helper actually runs on a
    /// sample tensor (smoke test).
    #[test]
    fn v_rms_norm_smoke() {
        let dev = Device::cpu();
        let x = Tensor::from_f32(
            Arc::from((0..16).map(|i| (i as f32 + 1.0) * 0.1).collect::<Vec<_>>()),
            Shape::from_dims(&[1, 2, 2, 4]),
            &dev,
        )
        .unwrap();
        let normed = v_rms_norm(&x, 1e-6).unwrap().realize_f32();
        // mean-squared per last-dim group should be ~1 after RMS norm.
        for chunk in normed.chunks(4) {
            let mean_sq: f32 = chunk.iter().map(|v| v * v).sum::<f32>() / 4.0;
            assert!(
                (mean_sq - 1.0).abs() < 1e-3,
                "v_rms_norm did not produce unit-RMS chunk: mean_sq = {mean_sq}"
            );
        }
    }

    #[test]
    fn forward_hidden_shape_and_finite() {
        let cfg = tiny_config();
        let model = Gemma4TextModel {
            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() {
        // forward_embeds(scaled embed_tokens(tokens), start_pos) must
        // produce the same logits as forward(tokens, start_pos) when the
        // caller does the sqrt scaling externally — this is the contract
        // PaliGemma / LLaVA / Voxtral all rely on.
        let cfg = tiny_config();
        let model = Gemma4TextModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        let tokens: Vec<u32> = vec![1, 2, 3];

        // Reference: forward(tokens).
        let logits_ref = model.forward(&tokens, 0).unwrap().realize_f32();

        // Path 2: embed_tokens_anchored → scale → forward_embeds.
        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 scaled = embeds.mul_scalar((cfg.hidden_size as f64).sqrt());
        let logits_via_embeds = model.forward_embeds(&scaled, 0).unwrap().realize_f32();

        assert_eq!(logits_ref.len(), logits_via_embeds.len());
        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,
            "forward vs forward_embeds (post-scale) must agree (max diff {max_diff})"
        );
    }

    #[test]
    fn forward_embeds_rejects_bad_shape() {
        let cfg = tiny_config();
        let model = Gemma4TextModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg),
        };
        // Wrong hidden size.
        let bad_embeds = 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_embeds, 0).is_err());
        // Wrong rank.
        let rank2 = Tensor::from_f32(
            vec![0.0_f32; cfg.hidden_size],
            Shape::from_dims(&[1, cfg.hidden_size]),
            &Device::cpu(),
        )
        .unwrap();
        assert!(model.forward_embeds(&rank2, 0).is_err());
    }

    #[test]
    fn forward_hidden_embeds_returns_pre_lm_head_states() {
        // forward_hidden_embeds should produce the same hidden states as
        // forward_hidden(tokens) when fed the post-scale embeddings.
        let cfg = tiny_config();
        let model = Gemma4TextModel {
            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 scaled = embeds.mul_scalar((cfg.hidden_size as f64).sqrt());
        let h_via_embeds = model
            .forward_hidden_embeds(&scaled, 0)
            .unwrap()
            .realize_f32();

        assert_eq!(h_ref.len(), h_via_embeds.len());
        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,
            "forward_hidden vs forward_hidden_embeds (post-scale) must agree (max diff {max_diff})"
        );
    }
}