nuvec_nn/rope_scaling.rs
1//! Rotary-embedding frequency scalings (mlx-lm 0.32.2's `rope_utils.py`): linear, Llama 3,
2//! YaRN, LongRoPE (`SuScaledRoPE`) and the proportional RoPE of Gemma 4's global layers.
3//! [`crate::Rope::scaled`] builds a rotary embedding from one.
4
5use nuvec::{Array, Static};
6
7/// How a rotary embedding's frequencies (or positions) are scaled, as mlx-lm's RoPE
8/// classes. Frequencies are computed in float32 with mlx-lm's formulas.
9#[derive(Debug, Clone, PartialEq, Static)]
10#[non_exhaustive]
11pub enum RopeScaling {
12 /// The plain `base^(-2i/dims)` frequencies (`default`).
13 None,
14 /// Positions divided by `factor` (`linear`).
15 Linear {
16 /// The division factor.
17 factor: f32,
18 },
19 /// Llama 3's frequency-dependent scaling (`Llama3RoPE`).
20 Llama3(Llama3),
21 /// YaRN (`YarnRoPE`): interpolated frequencies and an input scale (`mscale`).
22 Yarn(Yarn),
23 /// LongRoPE (`SuScaledRoPE`): per-pair long factors and an input scale.
24 LongRope(LongRope),
25 /// Gemma 4's proportional RoPE (`ProportionalRoPE`): the first `rotated_dims` of the
26 /// head's features rotate with exponents over the whole head, scaled by `factor`; the
27 /// rest keep their values (infinite wavelength).
28 Proportional {
29 /// Features that rotate (a multiple of 2, at most the head's).
30 rotated_dims: i32,
31 /// The frequencies' scale.
32 factor: f32,
33 },
34}
35
36/// [`RopeScaling::Llama3`]'s parameters (mlx-lm's defaults).
37#[derive(Debug, Clone, Copy, PartialEq, Static)]
38pub struct Llama3 {
39 /// Scale of low frequencies.
40 pub factor: f32,
41 /// Wavelength bound (in original contexts) of scaled low frequencies.
42 pub low_freq_factor: f32,
43 /// Wavelength bound of unscaled high frequencies.
44 pub high_freq_factor: f32,
45 /// The context length the model was trained with.
46 pub original_max_position_embeddings: f32,
47}
48
49impl Default for Llama3 {
50 fn default() -> Self {
51 Self {
52 factor: 1.0,
53 low_freq_factor: 1.0,
54 high_freq_factor: 4.0,
55 original_max_position_embeddings: 8192.0,
56 }
57 }
58}
59
60/// [`RopeScaling::Yarn`]'s parameters (mlx-lm's defaults; Hugging Face's `truncate` is
61/// ignored, as by mlx-lm).
62#[derive(Debug, Clone, Copy, PartialEq, Static)]
63pub struct Yarn {
64 /// The context's scale (`factor`).
65 pub factor: f32,
66 /// The context length the model was trained with.
67 pub original_max_position_embeddings: f32,
68 /// Rotations above which frequencies are kept (`beta_fast`).
69 pub beta_fast: f32,
70 /// Rotations below which frequencies are interpolated (`beta_slow`).
71 pub beta_slow: f32,
72 /// `mscale`.
73 pub mscale: f32,
74 /// `mscale_all_dim`.
75 pub mscale_all_dim: f32,
76}
77
78impl Default for Yarn {
79 fn default() -> Self {
80 Self {
81 factor: 1.0,
82 original_max_position_embeddings: 4096.0,
83 beta_fast: 32.0,
84 beta_slow: 1.0,
85 mscale: 1.0,
86 mscale_all_dim: 0.0,
87 }
88 }
89}
90
91/// [`RopeScaling::LongRope`]'s parameters (mlx-lm's defaults, long factors).
92#[derive(Debug, Clone, PartialEq, Static)]
93pub struct LongRope {
94 /// Per-pair factors (`long_factor`; one value applies to every pair).
95 pub long_factor: Vec<f32>,
96 /// The maximum context length.
97 pub max_position_embeddings: f32,
98 /// The context length the model was trained with.
99 pub original_max_position_embeddings: f32,
100 /// An explicit input scale (`long_mscale`).
101 pub long_mscale: Option<f32>,
102}
103
104impl Default for LongRope {
105 fn default() -> Self {
106 Self {
107 long_factor: vec![1.0],
108 max_position_embeddings: 131_072.0,
109 original_max_position_embeddings: 4096.0,
110 long_mscale: None,
111 }
112 }
113}
114
115/// `base^(i / dims)` for `i = 0, 2, …, upto − 2`, in float32 (mlx-lm's `base ** (arange(0,
116/// dims, 2) / dims)`).
117fn powers(base: f32, upto: i32, dims: i32) -> Vec<f32> {
118 (0..upto)
119 .step_by(2)
120 .map(|i| base.powf(i as f32 / dims as f32))
121 .collect()
122}
123
124impl RopeScaling {
125 /// The per-pair frequencies (`[dims / 2]`, float32) for `dims` rotated features and
126 /// `base`, or `None` when the plain `base` (and [`RopeScaling::position_scale`])
127 /// applies.
128 pub fn frequencies(&self, dims: i32, base: f32) -> Option<Array> {
129 let freqs: Vec<f32> = match self {
130 Self::None | Self::Linear { .. } => return None,
131 Self::Llama3(p) => {
132 let old = p.original_max_position_embeddings;
133 let (low_wavelen, high_wavelen) =
134 (old / p.low_freq_factor, old / p.high_freq_factor);
135 powers(base, dims, dims)
136 .into_iter()
137 .map(|freq| {
138 let wavelen = 2.0 * std::f32::consts::PI * freq;
139 let scaled = if wavelen > low_wavelen {
140 freq * p.factor
141 } else {
142 freq
143 };
144 if wavelen > high_wavelen && wavelen < low_wavelen {
145 let smooth = (old / wavelen - p.low_freq_factor)
146 / (p.high_freq_factor - p.low_freq_factor);
147 scaled / ((1.0 - smooth) / p.factor + smooth)
148 } else {
149 scaled
150 }
151 })
152 .collect()
153 }
154 Self::Yarn(p) => {
155 let correction = |rotations: f32| {
156 (dims as f32
157 * (p.original_max_position_embeddings
158 / (rotations * 2.0 * std::f32::consts::PI))
159 .ln())
160 / (2.0 * base.ln())
161 };
162 let low = correction(p.beta_fast).floor().max(0.0);
163 let mut high = correction(p.beta_slow).ceil().min((dims - 1) as f32);
164 if low == high {
165 high += 0.001;
166 }
167 powers(base, dims, dims)
168 .into_iter()
169 .enumerate()
170 .map(|(i, extra)| {
171 let inter = p.factor * extra;
172 let ramp = ((i as f32 - low) / (high - low)).clamp(0.0, 1.0);
173 let mask = 1.0 - ramp;
174 (inter * extra) / (inter * mask + extra * (1.0 - mask))
175 })
176 .collect()
177 }
178 Self::LongRope(p) => {
179 let f = powers(base, dims, dims);
180 let factor = |i: usize| match p.long_factor.as_slice() {
181 [one] => *one,
182 all => all.get(i).copied().unwrap_or(1.0),
183 };
184 f.into_iter()
185 .enumerate()
186 .map(|(i, x)| factor(i) * x)
187 .collect()
188 }
189 Self::Proportional {
190 rotated_dims,
191 factor,
192 } => {
193 let mut f: Vec<f32> = powers(base, *rotated_dims, dims)
194 .into_iter()
195 .map(|x| factor * x)
196 .collect();
197 f.resize((dims / 2) as usize, f32::INFINITY);
198 f
199 }
200 };
201 Some(Array::from(freqs))
202 }
203
204 /// The scale mlx-lm applies to the rotated features before the rotation (YaRN's
205 /// `mscale`, LongRoPE's; 1 otherwise).
206 pub fn mscale(&self) -> f32 {
207 let yarn = |scale: f32, m: f32| {
208 if scale <= 1.0 {
209 1.0
210 } else {
211 0.1 * m * scale.ln() + 1.0
212 }
213 };
214 match self {
215 Self::Yarn(p) => yarn(p.factor, p.mscale) / yarn(p.factor, p.mscale_all_dim),
216 Self::LongRope(p) => {
217 let factor = p.max_position_embeddings / p.original_max_position_embeddings;
218 p.long_mscale.unwrap_or(if factor <= 1.0 {
219 1.0
220 } else {
221 (1.0 + factor.ln() / p.original_max_position_embeddings.ln()).sqrt()
222 })
223 }
224 _ => 1.0,
225 }
226 }
227
228 /// The scale of positions (`1 / factor` for [`RopeScaling::Linear`]; 1 otherwise).
229 pub fn position_scale(&self) -> f32 {
230 match self {
231 Self::Linear { factor } => 1.0 / factor,
232 _ => 1.0,
233 }
234 }
235}