needle_core/
rope.rs

1//! Rotary Position Embeddings (RoPE).
2//! Applied to Q and K before the attention dot product.
3//! Uses base=10000, split-half rotation: first half × cos − second half × sin,
4//! second half × cos + first half × sin. Matches Python apply_rope exactly.
5
6use crate::math;
7use alloc::vec;
8use alloc::vec::Vec;
9
10/// Precomputed cosine/sine tables for RoPE.
11/// Shape: [max_len, head_dim/2]
12pub struct RopeCache {
13    pub cos: Vec<f32>,
14    pub sin: Vec<f32>,
15    pub max_len: usize,
16    pub half_dim: usize,
17}
18
19impl RopeCache {
20    pub fn new(max_len: usize, head_dim: usize, base: f32) -> Self {
21        let half_dim = head_dim / 2;
22        let mut cos = vec![0.0f32; max_len * half_dim];
23        let mut sin = vec![0.0f32; max_len * half_dim];
24
25        for pos in 0..max_len {
26            for i in 0..half_dim {
27                let theta = (pos as f32) / math::powf(base, 2.0 * i as f32 / head_dim as f32);
28                cos[pos * half_dim + i] = math::cos(theta);
29                sin[pos * half_dim + i] = math::sin(theta);
30            }
31        }
32
33        Self {
34            cos,
35            sin,
36            max_len,
37            half_dim,
38        }
39    }
40
41    /// Apply RoPE to a query/key tensor in-place.
42    /// `x`: shape [seq_len, num_heads, head_dim]
43    /// `offset`: position offset (for KV cache / decoder incremental decoding)
44    pub fn apply(
45        &self,
46        x: &mut [f32],
47        seq_len: usize,
48        num_heads: usize,
49        head_dim: usize,
50        offset: usize,
51    ) {
52        debug_assert_eq!(x.len(), seq_len * num_heads * head_dim);
53        debug_assert_eq!(head_dim, self.half_dim * 2);
54
55        for t in 0..seq_len {
56            let pos = t + offset;
57            debug_assert!(
58                pos < self.max_len,
59                "RoPE position {} exceeds max_len {}",
60                pos,
61                self.max_len
62            );
63            let cos_row = &self.cos[pos * self.half_dim..(pos + 1) * self.half_dim];
64            let sin_row = &self.sin[pos * self.half_dim..(pos + 1) * self.half_dim];
65
66            for h in 0..num_heads {
67                let base = (t * num_heads + h) * head_dim;
68                let (left, right) = x[base..base + head_dim].split_at_mut(self.half_dim);
69                for i in 0..self.half_dim {
70                    let x0 = left[i];
71                    let x1 = right[i];
72                    left[i] = x0 * cos_row[i] - x1 * sin_row[i];
73                    right[i] = x1 * cos_row[i] + x0 * sin_row[i];
74                }
75            }
76        }
77    }
78}
needle_core/rope.rs

needle_core/
rope.rs
