add glm-asr-nano

src/position_embed/rope.rs

@@ -115,6 +115,49 @@ pub fn apply_rotary_pos_emb(
     Ok((q_embed, k_embed))
 }
 
+pub fn glm_asr_apply_rotary_pos_emb(
+    q: &Tensor,
+    k: &Tensor,
+    cos: &Tensor,
+    sin: &Tensor,
+    tof32: bool,
+) -> Result<(Tensor, Tensor)> {
+    // sin/cos: to (bs, 1, seq_len, head_dim/2)
+    // q/k: (bs, n_head, seq_len, head_dim)
+    let mut cos = cos.clone();
+    let mut sin = sin.clone();
+    if cos.rank() == 2 {
+        // (seq_len, head_dim/2) -> (1, 1, seq_len, head_dim/2)
+        cos = cos.unsqueeze(0)?.unsqueeze(0)?;
+        sin = sin.unsqueeze(0)?.unsqueeze(0)?;
+    }
+    if cos.rank() == 3 {
+        // (bs, seq_len, head_dim/2) -> (bs, 1, seq_len, head_dim/2)
+        cos = cos.unsqueeze(1)?;
+        sin = sin.unsqueeze(1)?;
+    }
+    let orig_dtype = q.dtype();
+    let q = if tof32 { &q.to_dtype(DType::F32)? } else { q };
+    let k = if tof32 { &k.to_dtype(DType::F32)? } else { k };
+    let cos = cos.to_dtype(q.dtype())?;
+    let sin = sin.to_dtype(q.dtype())?;
+    let rotary_dim = cos.dim(D::Minus1)?;
+    let q_rot = q.narrow(D::Minus1, 0, rotary_dim)?;
+    let q_pass = q.narrow(D::Minus1, rotary_dim, rotary_dim)?;
+    let k_rot = k.narrow(D::Minus1, 0, rotary_dim)?;
+    let k_pass = k.narrow(D::Minus1, rotary_dim, rotary_dim)?;
+
+    let q_embed = q_rot
+        .broadcast_mul(&cos)?
+        .add(&rotate_half(&q_rot)?.broadcast_mul(&sin)?)?;
+    let k_embed = k_rot
+        .broadcast_mul(&cos)?
+        .add(&rotate_half(&k_rot)?.broadcast_mul(&sin)?)?;
+    let q_embed = Tensor::cat(&[q_embed, q_pass], D::Minus1)?.to_dtype(orig_dtype)?;
+    let k_embed = Tensor::cat(&[k_embed, k_pass], D::Minus1)?.to_dtype(orig_dtype)?;
+    Ok((q_embed, k_embed))
+}
+
 #[derive(Debug, Clone)]
 pub struct Qwen2_5VLTextRotaryEmbedding {
     inv_freq: Vec<f32>,