llama-rs 0.16.1

//! LLaMA model architecture implementation
//!
//! This module implements the LLaMA transformer architecture, supporting:
//! - LLaMA 1, 2, 3 variants
//! - Gemma2 (final logit softcapping)
//! - Grouped Query Attention (GQA)
//! - RoPE position embeddings
//! - Quantized weights

use std::sync::Arc;

use crate::backend::Backend;
use crate::tensor::{DType, Tensor};

use super::config::ModelConfig;
use super::deltanet::RecurrentConfig;
use super::error::{ModelError, ModelResult};
use super::layers::{Linear, NormLayer, TransformerLayer};
use super::deltanet::DeltaNetConfig;
use super::{Architecture, InferenceContext, Model};

/// LLaMA model implementation
pub struct LlamaModel {
    /// Model configuration
    config: ModelConfig,
    /// Token embedding matrix [vocab_size, hidden_size]
    token_embedding: Tensor,
    /// Transformer layers
    layers: Vec<TransformerLayer>,
    /// Final normalization (RMSNorm or LayerNorm depending on architecture)
    norm: NormLayer,
    /// Output projection (may share weights with embedding)
    output: Linear,
    /// Model architecture variant
    architecture: Architecture,
    /// Per-layer recurrent flag (true = delta-net or Mamba, false = full attention)
    recurrent_mask: Vec<bool>,
    /// Recurrent config (None if no recurrent layers)
    recurrent_config: Option<RecurrentConfig>,
    /// Per-layer token embedding [n_epl * n_layers, vocab_size] (Gemma 4 PLIE)
    per_layer_token_embd: Option<Tensor>,
    /// Shared projection: hidden_size -> n_epl * n_layers (Gemma 4 PLIE)
    per_layer_model_proj: Option<Linear>,
    /// RMS norm for per-layer projection (applied per n_epl chunk)
    per_layer_proj_norm: Option<super::layers::RMSNorm>,
    /// Per-layer embedding dimension (256 for Gemma 4, 0 if PLIE inactive)
    n_epl: usize,
}

impl LlamaModel {
    /// Create a new LLaMA model from loaded weights
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        config: ModelConfig,
        token_embedding: Tensor,
        layers: Vec<TransformerLayer>,
        norm: NormLayer,
        output: Linear,
        architecture: Architecture,
        per_layer_token_embd: Option<Tensor>,
        per_layer_model_proj: Option<Linear>,
        per_layer_proj_norm: Option<super::layers::RMSNorm>,
        n_epl: usize,
    ) -> ModelResult<Self> {
        if layers.len() != config.num_layers {
            return Err(ModelError::ConfigError(format!(
                "Expected {} layers, got {}",
                config.num_layers,
                layers.len()
            )));
        }

        let recurrent_mask: Vec<bool> = layers.iter().map(|l| l.is_recurrent()).collect();
        let has_recurrent = recurrent_mask.iter().any(|&r| r);

        let recurrent_config = if has_recurrent && config.has_ssm() {
            let is_mamba =
                matches!(architecture, Architecture::Mamba | Architecture::Mamba2);
            Some(if is_mamba {
                RecurrentConfig::Mamba(super::mamba::MambaConfig {
                    d_inner: config.ssm_d_inner,
                    d_state: config.ssm_d_state,
                    dt_rank: config.ssm_dt_rank,
                    conv_kernel: config.ssm_conv_kernel.max(1),
                })
            } else {
                let d_inner = config.ssm_d_inner;
                let d_state = config.ssm_d_state;
                let num_v_heads = config.ssm_dt_rank;
                let num_k_heads = config.ssm_n_group.max(1);
                let head_v_dim = d_inner / num_v_heads.max(1);
                let head_k_dim = d_state;
                let conv_kernel = config.ssm_conv_kernel;
                let q_dim = num_k_heads * head_k_dim;
                let k_dim = num_k_heads * head_k_dim;
                let qkv_dim = q_dim + k_dim + d_inner;
                RecurrentConfig::DeltaNet(DeltaNetConfig {
                    d_inner,
                    d_state,
                    num_v_heads,
                    num_k_heads,
                    head_v_dim,
                    head_k_dim,
                    conv_kernel,
                    qkv_dim,
                })
            })
        } else {
            None
        };

        Ok(Self {
            config,
            token_embedding,
            layers,
            norm,
            output,
            architecture,
            recurrent_mask,
            recurrent_config,
            per_layer_token_embd,
            per_layer_model_proj,
            per_layer_proj_norm,
            n_epl,
        })
    }

    /// Create an InferenceContext appropriate for this model (with recurrent state if needed).
    pub fn create_context(&self, backend: Arc<dyn Backend>) -> InferenceContext {
        if let Some(ref rc) = self.recurrent_config {
            InferenceContext::new_with_recurrent(
                &self.config,
                backend,
                &self.recurrent_mask,
                rc,
            )
        } else {
            InferenceContext::new(&self.config, backend)
        }
    }

    /// Get model configuration
    pub fn config(&self) -> &ModelConfig {
        &self.config
    }

    /// Get transformer layers
    pub fn layers(&self) -> &[TransformerLayer] {
        &self.layers
    }

    /// Decompose the model into its parts for GPU hybrid inference.
    /// The layers are moved out to avoid duplicating the weights.
    #[allow(clippy::type_complexity)]
    pub fn into_parts(
        self,
    ) -> (
        ModelConfig,
        Tensor,
        Vec<TransformerLayer>,
        NormLayer,
        Linear,
        Architecture,
        Vec<bool>,
        Option<RecurrentConfig>,
        Option<Tensor>,
        Option<Linear>,
        Option<super::layers::RMSNorm>,
        usize,
    ) {
        (
            self.config,
            self.token_embedding,
            self.layers,
            self.norm,
            self.output,
            self.architecture,
            self.recurrent_mask,
            self.recurrent_config,
            self.per_layer_token_embd,
            self.per_layer_model_proj,
            self.per_layer_proj_norm,
            self.n_epl,
        )
    }

    /// Get final normalization layer
    pub fn norm(&self) -> &NormLayer {
        &self.norm
    }

    /// Get output projection layer  
    pub fn output(&self) -> &Linear {
        &self.output
    }

    /// Get token embedding tensor
    pub fn token_embedding(&self) -> &Tensor {
        &self.token_embedding
    }

    /// Dequantize the embedding table once and return it as a `Cow`.
    ///
    /// For F32 embeddings this borrows the existing data (zero-copy).
    /// For quantized embeddings this dequantizes once into an owned `Vec`.
    fn dequantize_embeddings<'a>(
        &'a self,
        backend: &dyn Backend,
    ) -> ModelResult<std::borrow::Cow<'a, [f32]>> {
        if self.token_embedding.dtype() == DType::F32 {
            Ok(std::borrow::Cow::Borrowed(self.token_embedding.as_f32()?))
        } else {
            let numel = self.token_embedding.numel();
            let mut dequant = Tensor::zeros(vec![numel], DType::F32);
            backend.dequantize(&self.token_embedding, &mut dequant)?;
            Ok(std::borrow::Cow::Owned(dequant.as_f32()?.to_vec()))
        }
    }

    /// Dequantize the per-layer token embedding table once and return as `Cow`.
    fn dequantize_per_layer_embeddings<'a>(
        &'a self,
        backend: &dyn Backend,
    ) -> ModelResult<Option<std::borrow::Cow<'a, [f32]>>> {
        match self.per_layer_token_embd {
            None => Ok(None),
            Some(ref embd) => {
                if embd.dtype() == DType::F32 {
                    Ok(Some(std::borrow::Cow::Borrowed(embd.as_f32()?)))
                } else {
                    let numel = embd.numel();
                    let mut dequant = Tensor::zeros(vec![numel], DType::F32);
                    backend.dequantize(embd, &mut dequant)?;
                    Ok(Some(std::borrow::Cow::Owned(dequant.as_f32()?.to_vec())))
                }
            }
        }
    }

    /// Compute PLIE (per-layer input embedding) signals for all layers and tokens.
    ///
    /// Returns a flat Vec<f32> of shape [num_tokens, n_epl * n_layers] that the
    /// caller slices per-token and per-layer during the layer loop.
    /// Returns None if PLIE is not active.
    fn compute_plie_inputs(
        &self,
        tokens: &[u32],
        hiddens: &[Tensor],
        backend: &dyn Backend,
    ) -> ModelResult<Option<Vec<f32>>> {
        let n_epl = self.n_epl;
        if n_epl == 0 {
            return Ok(None);
        }
        let (per_layer_embd_data, per_layer_model_proj, per_layer_proj_norm) =
            match (
                &self.per_layer_token_embd,
                &self.per_layer_model_proj,
                &self.per_layer_proj_norm,
            ) {
                (Some(_), Some(proj), Some(norm)) => {
                    let embd_data = self.dequantize_per_layer_embeddings(backend)?
                        .expect("PLIE embedding must exist");
                    (embd_data, proj, norm)
                }
                _ => return Ok(None),
            };

        let n_layers = self.config.num_layers;
        let hidden_size = self.config.hidden_size;
        let total_pl_dim = n_epl * n_layers;
        let num_tokens = tokens.len();

        let sqrt_n_epl = (n_epl as f32).sqrt();
        let inv_sqrt_hidden = 1.0 / (hidden_size as f32).sqrt();
        let inv_sqrt_2: f32 = 1.0 / 2.0f32.sqrt();
        let norm_eps = per_layer_proj_norm.eps;
        let norm_weight = per_layer_proj_norm.weight.as_f32()?;

        // Result: [num_tokens * total_pl_dim]
        let mut result = vec![0.0f32; num_tokens * total_pl_dim];

        for (tok_idx, &token) in tokens.iter().enumerate() {
            let token_id = token as usize;

            // 1. Look up per-layer embedding for this token
            // per_layer_token_embd is [total_pl_dim, vocab_size] in GGUF convention,
            // but data layout is [vocab_size][total_pl_dim] (row-major by token)
            let embd_start = token_id * total_pl_dim;
            let embd_end = embd_start + total_pl_dim;
            if embd_end > per_layer_embd_data.len() {
                return Err(super::error::ModelError::InvalidMetadata {
                    key: "per_layer_token_embd".into(),
                    message: format!(
                        "PLIE embedding out of bounds: token={}, need {}, have {}",
                        token, embd_end, per_layer_embd_data.len()
                    ),
                });
            }
            let mut pl_embed = per_layer_embd_data[embd_start..embd_end].to_vec();

            // Scale by sqrt(n_epl)
            for v in pl_embed.iter_mut() {
                *v *= sqrt_n_epl;
            }

            // 2. Project main embedding through per_layer_model_proj
            let hidden_tensor = &hiddens[tok_idx];
            let mut pl_proj_tensor = Tensor::zeros(vec![total_pl_dim], DType::F32);
            per_layer_model_proj.forward(hidden_tensor, &mut pl_proj_tensor, backend)?;

            let pl_proj = pl_proj_tensor.as_f32_mut()?;

            // Scale by 1/sqrt(hidden_size)
            for v in pl_proj.iter_mut() {
                *v *= inv_sqrt_hidden;
            }

            // 3. RMS norm each n_epl chunk independently
            for layer in 0..n_layers {
                let start = layer * n_epl;
                let end = start + n_epl;
                let chunk = &mut pl_proj[start..end];

                // Compute RMS
                let ss: f32 = chunk.iter().map(|x| x * x).sum::<f32>() / n_epl as f32;
                let rms_inv = 1.0 / (ss + norm_eps).sqrt();

                // Apply: chunk[i] = chunk[i] / rms * norm_weight[i]
                for (i, v) in chunk.iter_mut().enumerate() {
                    *v = *v * rms_inv * norm_weight[i];
                }
            }

            // 4. Combine: (pl_proj + pl_embed) * inv_sqrt_2
            let out_start = tok_idx * total_pl_dim;
            let out_slice = &mut result[out_start..out_start + total_pl_dim];
            for i in 0..total_pl_dim {
                out_slice[i] = (pl_proj[i] + pl_embed[i]) * inv_sqrt_2;
            }
        }

        Ok(Some(result))
    }

    /// Get token embedding for given token IDs (public for testing)
    pub fn embed_tokens(&self, tokens: &[u32], backend: &dyn Backend) -> ModelResult<Tensor> {
        let hidden_size = self.config.hidden_size;
        let vocab_size = self.config.vocab_size;
        let seq_len = tokens.len();

        let embedding_data = self.dequantize_embeddings(backend)?;

        let mut output = vec![0.0f32; seq_len * hidden_size];

        // GGUF stores embeddings with shape listed as [hidden_size, vocab_size]
        // but in GGML convention, this means the data is laid out as [vocab_size][hidden_size]
        // i.e., each row is a token's embedding vector
        // So embedding for token t starts at t * hidden_size
        for (i, &token) in tokens.iter().enumerate() {
            let token_idx = token as usize;
            if token_idx >= vocab_size {
                return Err(ModelError::InvalidMetadata {
                    key: "token".into(),
                    message: format!("Token ID {} exceeds vocab size {}", token, vocab_size),
                });
            }

            let src_start = token_idx * hidden_size;
            let src_end = src_start + hidden_size;

            if src_end > embedding_data.len() {
                return Err(ModelError::InvalidMetadata {
                    key: "embedding".into(),
                    message: format!(
                        "Embedding index out of bounds: token_idx={}, src_end={}, embedding_len={}",
                        token_idx,
                        src_end,
                        embedding_data.len()
                    ),
                });
            }

            let dst_start = i * hidden_size;
            output[dst_start..dst_start + hidden_size]
                .copy_from_slice(&embedding_data[src_start..src_end]);
        }

        if seq_len == 1 {
            Tensor::from_f32(&output, vec![hidden_size])
        } else {
            Tensor::from_f32(&output, vec![seq_len, hidden_size])
        }
        .map_err(|e| e.into())
    }

    /// Compute logits from hidden state
    fn compute_logits(&self, hidden: &Tensor, backend: &dyn Backend) -> ModelResult<Tensor> {
        // Apply final normalization
        let mut normed = Tensor::zeros(hidden.shape().to_vec(), DType::F32);
        self.norm.forward(hidden, &mut normed, backend)?;

        // Project to vocabulary
        let mut logits = Tensor::zeros(vec![self.config.vocab_size], DType::F32);
        self.output.forward(&normed, &mut logits, backend)?;

        // Final logit softcapping (Gemma2): logits = cap * tanh(logits / cap)
        if self.config.final_logit_softcap > 0.0 {
            let cap = self.config.final_logit_softcap;
            let data = logits.as_f32_mut()?;
            for v in data.iter_mut() {
                *v = cap * (*v / cap).tanh();
            }
        }

        Ok(logits)
    }
}

impl Model for LlamaModel {
    fn create_context(&self, backend: Arc<dyn Backend>) -> InferenceContext {
        self.create_context(backend)
    }

    /// Forward pass. Supports LLaMA 1/2/3, Gemma2 (final logit softcapping when `final_logit_softcap` > 0).
    fn forward(&self, tokens: &[u32], ctx: &mut InferenceContext) -> ModelResult<Tensor> {
        let backend = ctx.backend.as_ref();
        let num_tokens = tokens.len();

        // Check context length
        let new_pos = ctx.position + num_tokens;
        if new_pos > self.config.max_seq_len {
            return Err(ModelError::ContextLengthExceeded {
                current: new_pos,
                max: self.config.max_seq_len,
            });
        }

        let embedding_data = self.dequantize_embeddings(backend)?;
        let hidden_size = self.config.hidden_size;
        let vocab_size = self.config.vocab_size;

        // Build per-token hidden states from embeddings
        let mut hiddens: Vec<Tensor> = Vec::with_capacity(num_tokens);
        for &token in tokens {
            let token_idx = token as usize;
            if token_idx >= vocab_size {
                return Err(ModelError::InvalidMetadata {
                    key: "token".into(),
                    message: format!("Token ID {} exceeds vocab size {}", token, vocab_size),
                });
            }
            let src = token_idx * hidden_size;
            hiddens.push(Tensor::from_f32(
                &embedding_data[src..src + hidden_size],
                vec![hidden_size],
            )?);
        }

        if std::env::var("LLAMA_DEBUG").is_ok() && ctx.position == 0 {
            let h = hiddens.last().unwrap().as_f32().unwrap();
            let n = h.len().min(8);
            eprintln!("[DBG] tokens: {:?}", tokens);
            eprintln!("[DBG] embed[0..{}]: {:?}", n, &h[..n]);
        }

        // Gemma scales token embeddings by sqrt(hidden_size)
        if self.architecture.is_gemma() {
            let scale = (hidden_size as f32).sqrt();
            for hidden in &mut hiddens {
                let data = hidden.as_f32_mut()?;
                for v in data.iter_mut() {
                    *v *= scale;
                }
            }
        }

        // Gemma 4 PLIE: pre-compute per-layer input embeddings for all tokens.
        // This must happen AFTER Gemma embedding scaling but BEFORE the layer loop,
        // because the projection uses the scaled main embedding as input.
        let plie_inputs = self.compute_plie_inputs(tokens, &hiddens, backend)?;
        let n_epl = self.n_epl;
        let n_layers = self.config.num_layers;
        let total_pl_dim = n_epl * n_layers;

        for (layer_idx, layer) in self.layers.iter().enumerate() {
            for (token_offset, hidden) in hiddens.iter_mut().enumerate() {
                let current_pos = ctx.position + token_offset;
                let recurrent_state = ctx
                    .recurrent_state
                    .as_mut()
                    .and_then(|rs| rs.states[layer_idx].as_mut());

                // Extract per-layer input slice for this token + layer
                let pl_slice = plie_inputs.as_ref().map(|all| {
                    let token_base = token_offset * total_pl_dim;
                    let layer_base = token_base + layer_idx * n_epl;
                    &all[layer_base..layer_base + n_epl]
                });

                let kv_slot = ctx.kv_cache.kv_source_layer[layer_idx];
                let skip_kv = kv_slot != layer_idx;
                *hidden = layer.forward(
                    hidden,
                    &mut ctx.kv_cache.k_cache[kv_slot],
                    &mut ctx.kv_cache.v_cache[kv_slot],
                    current_pos,
                    self.config.rope_config.freq_base,
                    self.config.rope_config.freq_scale,
                    backend,
                    recurrent_state,
                    pl_slice,
                    skip_kv,
                )?;
            }

            if std::env::var("LLAMA_DEBUG").is_ok() && ctx.position == 0 {
                let show = std::env::var("LLAMA_DEBUG_FULL").is_ok()
                    || layer_idx < 4 || layer_idx == self.layers.len() - 1;
                if show {
                    let h = hiddens.last().unwrap().as_f32().unwrap();
                    let rms: f32 = (h.iter().map(|x| x * x).sum::<f32>() / h.len() as f32).sqrt();
                    eprintln!("[DBG] layer {:>2}: rms={:.6} first4=[{:.6},{:.6},{:.6},{:.6}]",
                        layer_idx, rms, h[0], h[1], h[2], h[3]);
                }
            }
        }

        ctx.position = new_pos;
        ctx.kv_cache.seq_len = new_pos;

        // Compute logits from the last token's hidden state
        self.compute_logits(hiddens.last().unwrap(), backend)
    }

    fn config(&self) -> &ModelConfig {
        &self.config
    }

    fn architecture(&self) -> Architecture {
        self.architecture
    }
}

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

    #[test]
    fn test_llama_config() {
        let config = ModelConfig::llama_7b();
        assert_eq!(config.vocab_size, 32000);
        assert_eq!(config.hidden_size, 4096);
        assert_eq!(config.num_layers, 32);
        assert_eq!(config.num_heads, 32);
    }
}