zsfm_burn/
rope.rs

1use anyhow::Result;
2use burn::tensor::{Device, Tensor, TensorData};
3
4/// Build Llama-style RoPE cos/sin tables for positions `[0, max_seq)`.
5///
6/// Returns `(cos, sin)` each shaped `[max_seq, head_dim]` (interleaved
7/// `cat([freqs, freqs])` layout). Shared by Chronos, Sundial, TimesFM, Toto.
8pub fn rope_tables(
9    head_dim: usize,
10    max_seq: usize,
11    theta: f64,
12    device: &Device,
13) -> Result<(Tensor<2>, Tensor<2>)> {
14    let half = head_dim / 2;
15    let inv_freq: Vec<f32> = (0..half)
16        .map(|i| 1.0 / theta.powf(2.0 * i as f64 / head_dim as f64) as f32)
17        .collect();
18    build_rope_tables(&inv_freq, head_dim, max_seq, device)
19}
20
21/// Build cos/sin tables from explicit `inv_freq` (`[half]`).
22pub fn build_rope_tables(
23    inv_freq: &[f32],
24    head_dim: usize,
25    max_seq: usize,
26    device: &Device,
27) -> Result<(Tensor<2>, Tensor<2>)> {
28    let half = head_dim / 2;
29    anyhow::ensure!(inv_freq.len() == half, "inv_freq len must be head_dim/2");
30    let inv = Tensor::<1>::from_data(TensorData::new(inv_freq.to_vec(), [half]), device);
31    let positions: Vec<f32> = (0..max_seq).map(|p| p as f32).collect();
32    let positions = Tensor::<1>::from_data(TensorData::new(positions, [max_seq]), device);
33    let pos_col = positions.unsqueeze_dim::<2>(1);
34    let inv_row = inv.unsqueeze_dim::<2>(0);
35    let freqs = pos_col * inv_row;
36    let emb = Tensor::cat(vec![freqs.clone(), freqs], 1);
37    Ok((emb.clone().cos(), emb.sin()))
38}
39
40/// Apply interleaved (paired) rotary embeddings in place on host data.
41///
42/// Rotates pairs `(2i, 2i+1)` of the first `rope_dim` dims of each
43/// `[n_heads, seq_len, head_dim]` row using flat `[max_pos * half_rope]`
44/// tables. Shared by Moirai-2, TabICL, TabPFN, Lag-Llama (rank arg covers
45/// their dim ordering differences).
46#[allow(clippy::too_many_arguments)]
47pub fn apply_interleaved_rope(
48    x_data: &mut [f32],
49    time_ids: &[usize],
50    n_heads: usize,
51    seq_len: usize,
52    head_dim: usize,
53    rope_dim: usize,
54    cos_table: &[f32],
55    sin_table: &[f32],
56) -> anyhow::Result<()> {
57    let half_rope = rope_dim / 2;
58    anyhow::ensure!(rope_dim % 2 == 0, "rope_dim must be even");
59    anyhow::ensure!(
60        x_data.len() == n_heads * seq_len * head_dim,
61        "x_data len mismatch"
62    );
63    for h in 0..n_heads {
64        for (t, &pos) in time_ids.iter().enumerate() {
65            let base = h * seq_len * head_dim + t * head_dim;
66            let trig_base = pos * half_rope;
67            anyhow::ensure!(
68                trig_base + half_rope <= cos_table.len(),
69                "rope table too small for pos {pos}"
70            );
71            for i in 0..half_rope {
72                let cos_v = cos_table[trig_base + i];
73                let sin_v = sin_table[trig_base + i];
74                let ie = base + 2 * i;
75                let io = base + 2 * i + 1;
76                let ve = x_data[ie];
77                let vo = x_data[io];
78                x_data[ie] = cos_v * ve - sin_v * vo;
79                x_data[io] = sin_v * ve + cos_v * vo;
80            }
81        }
82    }
83    Ok(())
84}
85
86/// Llama-style rotary embedding apply on `[b, heads, seq, head_dim]`
87/// with prebuilt `[max_seq, head_dim]` cos/sin tables (positions `start..`).
88///
89/// `rotate_half(x) = cat([-x[half:], x[:half]])`; out = `x*cos + rot*sin`.
90/// Shared by Chronos, Sundial, Toto (all use this exact form).
91pub fn apply_llama_rope(
92    x: Tensor<4>,
93    cos_table: Tensor<2>,
94    sin_table: Tensor<2>,
95    start_pos: usize,
96) -> Tensor<4> {
97    let dims = x.dims();
98    let seq = dims[2];
99    let hd = dims[3];
100    let half = hd / 2;
101    let cos = cos_table
102        .narrow(0, start_pos, seq)
103        .unsqueeze_dim::<3>(0)
104        .unsqueeze_dim::<4>(0);
105    let sin = sin_table
106        .narrow(0, start_pos, seq)
107        .unsqueeze_dim::<3>(0)
108        .unsqueeze_dim::<4>(0);
109    let x1 = x.clone().narrow(3, 0, half);
110    let x2 = x.clone().narrow(3, half, half);
111    let rotated = Tensor::cat(vec![x2.mul_scalar(-1.0), x1], 3);
112    x * cos + rotated * sin
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    fn dev() -> Device {
120        Device::flex()
121    }
122
123    #[test]
124    fn rope_tables_shape() {
125        let device = dev();
126        let (cos, sin) = rope_tables(8, 4, 10_000.0, &device).unwrap();
127        assert_eq!(cos.dims(), [4, 8]);
128        assert_eq!(sin.dims(), [4, 8]);
129    }
130
131    #[test]
132    fn interleaved_rope_identity_at_pos0() {
133        // At pos 0, cos=1 sin=0 → input unchanged.
134        let mut x = vec![1.0f32, 2.0, 3.0, 4.0];
135        let cos = vec![1.0f32, 1.0];
136        let sin = vec![0.0f32, 0.0];
137        apply_interleaved_rope(&mut x, &[0], 1, 1, 4, 4, &cos, &sin).unwrap();
138        assert_eq!(x, vec![1.0, 2.0, 3.0, 4.0]);
139    }
140}
zsfm_burn/rope.rs

zsfm_burn/
rope.rs
