combs_models/
rope.rs

1//! Rotary positional embeddings (standard Llama RoPE, half-split /
2//! `rotate_half` convention as used by HuggingFace `transformers`).
3
4use burn::tensor::{Tensor, TensorData, backend::Backend, Device};
5
6/// Precomputed RoPE cosine/sine tables.
7///
8/// Frequencies: `inv_freq[i] = theta^(-2i / head_dim)` for `i in 0..head_dim/2`.
9/// Tables are `[max_position, head_dim]` with the half-dim frequencies
10/// duplicated (`cat([cos, cos])`) to match the half-split convention
11/// (`config.rope_interleaved = false`, the HF Llama default).
12pub struct RotaryEmbedding<B: Backend> {
13    cos: Tensor<B, 2>,
14    sin: Tensor<B, 2>,
15}
16
17/// Computes the `[max_position, head_dim]` cos/sin tables on the host.
18/// Exposed for unit tests.
19pub fn build_tables(
20    head_dim: usize,
21    theta: f64,
22    max_position: usize,
23) -> (Vec<f32>, Vec<f32>) {
24    let half = head_dim / 2;
25    let inv_freq: Vec<f64> = (0..half)
26        .map(|i| theta.powf(-2.0 * i as f64 / head_dim as f64))
27        .collect();
28    let mut cos = Vec::with_capacity(max_position * head_dim);
29    let mut sin = Vec::with_capacity(max_position * head_dim);
30    for pos in 0..max_position {
31        // Half-split layout: the half-dim frequencies appear twice so the
32        // table can be applied elementwise against [x1, x2].
33        for _ in 0..2 {
34            for f in &inv_freq {
35                let angle = pos as f64 * f;
36                cos.push(angle.cos() as f32);
37                sin.push(angle.sin() as f32);
38            }
39        }
40    }
41    (cos, sin)
42}
43
44impl<B: Backend> RotaryEmbedding<B> {
45    /// Builds tables on `device`.
46    pub fn new(head_dim: usize, theta: f64, max_position: usize, device: &Device<B>) -> Self {
47        let (cos, sin) = build_tables(head_dim, theta, max_position);
48        RotaryEmbedding {
49            cos: Tensor::from_data(
50                TensorData::new(cos, [max_position, head_dim]),
51                device,
52            ),
53            sin: Tensor::from_data(
54                TensorData::new(sin, [max_position, head_dim]),
55                device,
56            ),
57        }
58    }
59
60    /// Applies RoPE to a `[batch, heads, seq, head_dim]` tensor whose first
61    /// sequence position is at absolute position `pos`.
62    pub fn apply(&self, x: Tensor<B, 4>, pos: usize) -> Tensor<B, 4> {
63        let [batch, heads, seq, dim] = x.dims();
64        let half = dim / 2;
65        let cos = self
66            .cos
67            .clone()
68            .narrow(0, pos, seq)
69            .reshape([1, 1, seq, dim]);
70        let sin = self
71            .sin
72            .clone()
73            .narrow(0, pos, seq)
74            .reshape([1, 1, seq, dim]);
75
76        // rotate_half(x) = cat([-x2, x1]) along the head_dim axis.
77        let x1 = x.clone().narrow(3, 0, half);
78        let x2 = x.clone().narrow(3, half, half);
79        let rotated = Tensor::cat(vec![x2.neg(), x1], 3);
80
81        let out = x * cos + rotated * sin;
82        debug_assert_eq!(out.dims(), [batch, heads, seq, dim]);
83        out
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn tables_match_known_math() {
93        let head_dim = 8;
94        let theta = 10000.0f64;
95        let max_pos = 4;
96        let (cos, sin) = build_tables(head_dim, theta, max_pos);
97        // inv_freq[i] = theta^(-2i/d): i=0 -> 1, i=1 -> 10000^-0.25 = 0.1,
98        // i=2 -> 0.01, i=3 -> 0.001.
99        let inv = [1.0f64, 0.1, 0.01, 0.001];
100        for pos in 0..max_pos {
101            for i in 0..4 {
102                let angle = pos as f64 * inv[i];
103                // half-split layout: index i and i + half share the frequency.
104                for idx in [pos * head_dim + i, pos * head_dim + i + 4] {
105                    assert!(
106                        (cos[idx] as f64 - angle.cos()).abs() < 1e-5,
107                        "cos mismatch at pos {pos} i {i}"
108                    );
109                    assert!(
110                        (sin[idx] as f64 - angle.sin()).abs() < 1e-5,
111                        "sin mismatch at pos {pos} i {i}"
112                    );
113                }
114            }
115        }
116        // Position 0 must be the identity rotation.
117        for i in 0..head_dim {
118            assert!((cos[i] - 1.0).abs() < 1e-6);
119            assert!(sin[i].abs() < 1e-6);
120        }
121    }
122}
combs_models/rope.rs

combs_models/
rope.rs
