add LFM2-1.2B, LFM2.5-1.2B-Instruct

src/models/common/mod.rs

@@ -411,6 +411,145 @@ impl QKVCatAttention {
     }
 }
 
+pub struct QKNormAttention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    q_norm: RmsNorm,
+    k_norm: RmsNorm,
+    num_attention_heads: usize,
+    num_key_value_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    scaling: f64,
+    kv_cache: Option<(Tensor, Tensor)>,
+}
+
+impl QKNormAttention {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        num_attention_heads: usize,
+        head_dim: Option<usize>,
+        num_key_value_heads: Option<usize>,
+        attention_bias: bool,
+        rms_norm_eps: f64,
+        q_proj_pp_name: Option<&str>,
+        k_proj_pp_name: Option<&str>,
+        v_proj_pp_name: Option<&str>,
+        o_proj_pp_name: Option<&str>,
+        q_norm_pp_name: Option<&str>,
+        k_norm_pp_name: Option<&str>,
+    ) -> Result<Self> {
+        let head_dim = head_dim.unwrap_or(hidden_size / num_attention_heads);
+        let num_key_value_heads = num_key_value_heads.unwrap_or(num_attention_heads);
+        let num_kv_groups = num_attention_heads / num_key_value_heads;
+        let scaling = 1f64 / f64::sqrt(head_dim as f64);
+        let q_proj_pp_name = q_proj_pp_name.unwrap_or("q_proj");
+        let k_proj_pp_name = k_proj_pp_name.unwrap_or("k_proj");
+        let v_proj_pp_name = v_proj_pp_name.unwrap_or("v_proj");
+        let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj");
+        let q_norm_pp_name = q_norm_pp_name.unwrap_or("q_norm");
+        let k_norm_pp_name = k_norm_pp_name.unwrap_or("k_norm");
+        let q_proj = linear_b(
+            hidden_size,
+            num_attention_heads * head_dim,
+            attention_bias,
+            vb.pp(q_proj_pp_name),
+        )?;
+        let k_proj = linear_b(
+            hidden_size,
+            num_key_value_heads * head_dim,
+            attention_bias,
+            vb.pp(k_proj_pp_name),
+        )?;
+        let v_proj = linear_b(
+            hidden_size,
+            num_key_value_heads * head_dim,
+            attention_bias,
+            vb.pp(v_proj_pp_name),
+        )?;
+        let o_proj = linear_b(
+            num_attention_heads * head_dim,
+            hidden_size,
+            attention_bias,
+            vb.pp(o_proj_pp_name),
+        )?;
+        let q_norm = rms_norm(head_dim, rms_norm_eps, vb.pp(q_norm_pp_name))?;
+        let k_norm = rms_norm(head_dim, rms_norm_eps, vb.pp(k_norm_pp_name))?;
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            q_norm,
+            k_norm,
+            num_attention_heads,
+            num_key_value_heads,
+            num_kv_groups,
+            head_dim,
+            scaling,
+            kv_cache: None,
+        })
+    }
+
+    pub fn forward(
+        &mut self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (b_sz, q_len, _) = xs.dims3()?;
+        let query_states = self.q_proj.forward(xs)?.reshape((
+            b_sz,
+            q_len,
+            self.num_attention_heads,
+            self.head_dim,
+        ))?;
+        let query_states = self.q_norm.forward(&query_states)?.transpose(1, 2)?;
+        let key_states = self.k_proj.forward(xs)?.reshape((
+            b_sz,
+            q_len,
+            self.num_key_value_heads,
+            self.head_dim,
+        ))?;
+        let key_states = self.k_norm.forward(&key_states)?.transpose(1, 2)?;
+        let value_states = self.v_proj.forward(xs)?;
+        let value_states = value_states
+            .reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let (query_states, key_states) =
+            apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?;
+        let (key_states, value_states) = match &self.kv_cache {
+            None => (key_states, value_states),
+            Some((prev_k, prev_v)) => {
+                let key_states = Tensor::cat(&[prev_k, &key_states], 2)?;
+                let value_states = Tensor::cat(&[prev_v, &value_states], 2)?;
+                (key_states, value_states)
+            }
+        };
+        self.kv_cache = Some((key_states.clone(), value_states.clone()));
+        let attn_output = eager_attention_forward(
+            &query_states,
+            &key_states,
+            &value_states,
+            Some(self.num_kv_groups),
+            attention_mask,
+            self.scaling,
+        )?;
+        let attn_output =
+            attn_output.reshape((b_sz, q_len, self.num_attention_heads * self.head_dim))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        self.kv_cache = None
+    }
+}
+
 pub struct NaiveAttnTwoLinearMLPBlock {
     self_attn: NaiveAttention,
     mlp: TwoLinearMLP,