zsfm_sundial/infer/
rope.rs

1use anyhow::Result;
2use candle_core::{DType, Device, Tensor, D};
3
4pub struct RopeCache {
5    cos_t: Tensor, // [max_seq, head_dim]
6    sin_t: Tensor, // [max_seq, head_dim]
7}
8
9impl RopeCache {
10    pub fn new(head_dim: usize, max_seq: usize, theta: f64, device: &Device) -> Result<Self> {
11        let half = head_dim / 2;
12        let inv_freq: Vec<f64> = (0..half)
13            .map(|i| 1.0 / theta.powf(2.0 * i as f64 / head_dim as f64))
14            .collect();
15        let mut cos = vec![0.0f32; max_seq * head_dim];
16        let mut sin = vec![0.0f32; max_seq * head_dim];
17        for p in 0..max_seq {
18            for i in 0..half {
19                let angle = (p as f64 * inv_freq[i]) as f32;
20                let (s, c) = angle.sin_cos();
21                // cos[i] == cos[half+i], sin[i] == sin[half+i] — both halves are symmetric
22                cos[p * head_dim + i] = c;
23                cos[p * head_dim + half + i] = c;
24                sin[p * head_dim + i] = s;
25                sin[p * head_dim + half + i] = s;
26            }
27        }
28        let cos_t = Tensor::from_vec(cos, (max_seq, head_dim), device)?;
29        let sin_t = Tensor::from_vec(sin, (max_seq, head_dim), device)?;
30        Ok(Self { cos_t, sin_t })
31    }
32
33    /// Apply RoPE to x with shape [b, n_heads, seq, head_dim].
34    ///
35    /// Llama rotation: x * cos + rotate_half(x) * sin
36    /// rotate_half(x) = cat([-x[..., half:], x[..., :half]], dim=-1)
37    pub fn apply(&self, x: &Tensor, start_pos: usize) -> Result<Tensor> {
38        let dims = x.dims();
39        let seq = dims[2];
40        let hd = dims[3];
41        let half = hd / 2;
42        let dtype = x.dtype();
43
44        // Zero-copy narrow: produces a view into the pre-built Tensor
45        // cos_t and sin_t already satisfy cos[i] == cos[half+i], so use directly —
46        // no need to narrow to half then cat back to full width.
47        let cos_t = self.cos_t.narrow(0, start_pos, seq)?.unsqueeze(0)?.unsqueeze(0)?;
48        let sin_t = self.sin_t.narrow(0, start_pos, seq)?.unsqueeze(0)?.unsqueeze(0)?;
49
50        let x32 = x.to_dtype(DType::F32)?;
51        let x1 = x32.narrow(D::Minus1, 0, half)?;
52        let x2 = x32.narrow(D::Minus1, half, half)?;
53
54        // rotate_half: cat([-x2, x1])
55        let rotated = Tensor::cat(&[&x2.neg()?, &x1], D::Minus1)?;
56
57        let out = (x32.broadcast_mul(&cos_t)? + rotated.broadcast_mul(&sin_t)?)?;
58        Ok(out.to_dtype(dtype)?)
59    }
60}
zsfm_sundial/infer/rope.rs

zsfm_sundial/infer/
rope.rs
