mlx_native/ops/
rope.rs

1//! Rotary Position Embedding (RoPE) GPU dispatch.
2//!
3//! Applies rotation in pairs of elements using cos/sin of
4//! `position * theta^(-2i/d)`.  Gemma 4 uses theta=10000 for sliding
5//! attention layers and theta=1000000 for global attention layers.
6
7use metal::MTLSize;
8
9use crate::buffer::MlxBuffer;
10use crate::dtypes::DType;
11use crate::encoder::CommandEncoder;
12use crate::error::{MlxError, Result};
13use crate::kernel_registry::KernelRegistry;
14
15/// MSL source for the RoPE kernels (embedded at compile time).
16pub static ROPE_SHADER_SOURCE: &str = include_str!("../shaders/rope.metal");
17
18/// Register RoPE shader sources with the given kernel registry.
19pub fn register(registry: &mut KernelRegistry) {
20    registry.register_source("rope_f32", ROPE_SHADER_SOURCE);
21    registry.register_source("rope_f16", ROPE_SHADER_SOURCE);
22    registry.register_source("rope_bf16", ROPE_SHADER_SOURCE);
23    registry.register_source("rope_neox_bf16", ROPE_SHADER_SOURCE);
24    registry.register_source("rope_neox_f32", ROPE_SHADER_SOURCE);
25}
26
27/// Dispatch a RoPE operation on the GPU.
28///
29/// # Arguments
30///
31/// * `encoder`      - Command encoder to record the dispatch into.
32/// * `registry`     - Kernel registry (must have RoPE sources registered).
33/// * `device`       - Metal device for pipeline compilation.
34/// * `input`        - Input buffer of shape `[seq_len, head_dim]` (f32 or f16).
35/// * `output`       - Output buffer (same dtype and shape as input).
36/// * `params_buf`   - Params buffer containing `[theta, head_dim, 0, 0]` as f32.
37/// * `positions_buf` - Positions buffer containing `[pos_0, pos_1, ...]` as u32.
38/// * `seq_len`      - Number of sequence positions.
39/// * `head_dim`     - Dimension of each head (must be even).
40///
41/// # Errors
42///
43/// Returns `MlxError::InvalidArgument` if:
44/// - Input dtype is not f32 or f16.
45/// - head_dim is not even.
46/// - Input and output element counts do not match.
47pub fn dispatch_rope(
48    encoder: &mut CommandEncoder,
49    registry: &mut KernelRegistry,
50    device: &metal::DeviceRef,
51    input: &MlxBuffer,
52    output: &MlxBuffer,
53    params_buf: &MlxBuffer,
54    positions_buf: &MlxBuffer,
55    seq_len: u32,
56    head_dim: u32,
57) -> Result<()> {
58    if head_dim % 2 != 0 {
59        return Err(MlxError::InvalidArgument(format!(
60            "RoPE head_dim must be even, got {}",
61            head_dim
62        )));
63    }
64    if head_dim == 0 || seq_len == 0 {
65        return Err(MlxError::InvalidArgument(
66            "RoPE head_dim and seq_len must be > 0".into(),
67        ));
68    }
69
70    let expected_elements = (seq_len as usize) * (head_dim as usize);
71    if input.element_count() != expected_elements {
72        return Err(MlxError::InvalidArgument(format!(
73            "RoPE input element count {} != seq_len({}) * head_dim({})",
74            input.element_count(),
75            seq_len,
76            head_dim
77        )));
78    }
79    if output.element_count() != expected_elements {
80        return Err(MlxError::InvalidArgument(format!(
81            "RoPE output element count {} != seq_len({}) * head_dim({})",
82            output.element_count(),
83            seq_len,
84            head_dim
85        )));
86    }
87    // hf2q ADR-030 iter-113 — defense-in-depth dtype-coherence check.
88    if input.dtype() != output.dtype() {
89        return Err(MlxError::InvalidArgument(format!(
90            "RoPE dtype mismatch: input={} != output={}",
91            input.dtype(), output.dtype(),
92        )));
93    }
94
95    let kernel_name = match input.dtype() {
96        DType::F32 => "rope_f32",
97        DType::F16 => "rope_f16",
98        DType::BF16 => "rope_bf16",
99        _ => {
100            return Err(MlxError::InvalidArgument(format!(
101                "RoPE unsupported dtype: {}",
102                input.dtype()
103            )));
104        }
105    };
106
107    let pipeline = registry.get_pipeline(kernel_name, device)?;
108    let half_dim = head_dim / 2;
109
110    // Grid: (half_dim, seq_len) — one thread per pair per position
111    // Threadgroup: use a reasonable size for the pair dimension
112    let tg_x = std::cmp::min(64, half_dim as u64);
113    let tg_y = std::cmp::min(4, seq_len as u64);
114
115    encoder.encode(
116        pipeline,
117        &[
118            (0, input),
119            (1, output),
120            (2, params_buf),
121            (3, positions_buf),
122        ],
123        MTLSize::new(half_dim as u64, seq_len as u64, 1),
124        MTLSize::new(tg_x, tg_y, 1),
125    );
126
127    Ok(())
128}
129
130/// GPU params for the neox RoPE kernel's auxiliary params buffer.
131///
132/// Must match the uint array in `rope_neox_bf16` buffer(4).
133#[repr(C)]
134#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
135struct GpuRopeNeoxParams {
136    n_heads: u32,
137    _pad: u32,
138}
139
140/// Dispatch a Neox/split-convention RoPE operation on the GPU (bf16 only).
141///
142/// The Neox convention pairs `(d[i], d[i + half_rope_dim])` instead of
143/// `(d[2i], d[2i+1])`.  Supports partial rotary where only the first
144/// `rope_dim` dimensions are rotated.
145///
146/// # Arguments
147///
148/// * `encoder`       - Command encoder to record the dispatch into.
149/// * `registry`      - Kernel registry (must have rope_neox_bf16 registered).
150/// * `device`        - Metal device for pipeline compilation.
151/// * `input`         - Input buffer of shape `[seq_len * n_heads, head_dim]` (bf16).
152/// * `output`        - Output buffer (same shape and dtype as input).
153/// * `params_buf`    - Params buffer containing `[theta, head_dim, rope_dim, 0]` as f32.
154/// * `positions_buf` - Positions buffer containing `[pos_0, pos_1, ...]` as u32 (length = seq_len).
155/// * `seq_len`       - Number of sequence positions.
156/// * `n_heads`       - Number of attention heads.
157/// * `head_dim`      - Dimension of each head.
158/// * `rope_dim`      - Number of dimensions to rotate (must be even, <= head_dim).
159///
160/// # Errors
161///
162/// Returns `MlxError::InvalidArgument` if parameters are invalid.
163#[allow(clippy::too_many_arguments)]
164pub fn dispatch_rope_neox_bf16(
165    encoder: &mut CommandEncoder,
166    registry: &mut KernelRegistry,
167    device: &metal::DeviceRef,
168    input: &MlxBuffer,
169    output: &MlxBuffer,
170    params_buf: &MlxBuffer,
171    positions_buf: &MlxBuffer,
172    seq_len: u32,
173    n_heads: u32,
174    head_dim: u32,
175    rope_dim: u32,
176) -> Result<()> {
177    use super::encode_helpers::{as_bytes, encode_with_args, KernelArg};
178
179    if rope_dim % 2 != 0 {
180        return Err(MlxError::InvalidArgument(format!(
181            "RoPE neox rope_dim must be even, got {}",
182            rope_dim
183        )));
184    }
185    if rope_dim > head_dim {
186        return Err(MlxError::InvalidArgument(format!(
187            "RoPE neox rope_dim ({}) must be <= head_dim ({})",
188            rope_dim, head_dim
189        )));
190    }
191    if head_dim == 0 || seq_len == 0 || n_heads == 0 {
192        return Err(MlxError::InvalidArgument(
193            "RoPE neox head_dim, seq_len, and n_heads must be > 0".into(),
194        ));
195    }
196
197    let n_rows = (seq_len as usize) * (n_heads as usize);
198    let expected_elements = n_rows * (head_dim as usize);
199    if input.element_count() != expected_elements {
200        return Err(MlxError::InvalidArgument(format!(
201            "RoPE neox input element count {} != seq_len({}) * n_heads({}) * head_dim({})",
202            input.element_count(),
203            seq_len,
204            n_heads,
205            head_dim
206        )));
207    }
208    if output.element_count() != expected_elements {
209        return Err(MlxError::InvalidArgument(format!(
210            "RoPE neox output element count {} != seq_len({}) * n_heads({}) * head_dim({})",
211            output.element_count(),
212            seq_len,
213            n_heads,
214            head_dim
215        )));
216    }
217
218    let pipeline = registry.get_pipeline("rope_neox_bf16", device)?;
219    let half_rope = rope_dim / 2;
220
221    let gpu_rope_params = GpuRopeNeoxParams {
222        n_heads,
223        _pad: 0,
224    };
225
226    // Grid: (half_rope, n_rows) — one thread per pair per row
227    let tg_x = std::cmp::min(64, half_rope as u64);
228    let tg_y = std::cmp::min(4, n_rows as u64);
229
230    encode_with_args(
231        encoder,
232        pipeline,
233        &[
234            (0, KernelArg::Buffer(input)),
235            (1, KernelArg::Buffer(output)),
236            (2, KernelArg::Buffer(params_buf)),
237            (3, KernelArg::Buffer(positions_buf)),
238            (4, KernelArg::Bytes(as_bytes(&gpu_rope_params))),
239        ],
240        MTLSize::new(half_rope as u64, n_rows as u64, 1),
241        MTLSize::new(tg_x, tg_y, 1),
242    );
243
244    Ok(())
245}
246
247/// GPU params for the neox f32 RoPE kernel with freq_factors support.
248///
249/// Must match the uint array in `rope_neox_f32` buffer(4).
250#[repr(C)]
251#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
252struct GpuRopeNeoxF32Params {
253    n_heads: u32,
254    has_freq_factors: u32,
255}
256
257/// Dispatch a Neox/split-convention RoPE operation on the GPU (f32) with
258/// optional freq_factors support.
259///
260/// The Neox convention pairs `(d[i], d[i + half_rope_dim])` instead of
261/// `(d[2i], d[2i+1])`.  Supports partial rotary where only the first
262/// `rope_dim` dimensions are rotated.
263///
264/// When `freq_factors` is `Some`, each pair's base frequency is divided by
265/// `freq_factors[pair_idx]`.  Gemma 4's global attention layers use this to
266/// mask out rotation for certain dimensions (freq_factor=1e30 -> identity).
267///
268/// # Arguments
269///
270/// * `encoder`       - Command encoder to record the dispatch into.
271/// * `registry`      - Kernel registry (must have rope_neox_f32 registered).
272/// * `device`        - Metal device for pipeline compilation.
273/// * `input`         - Input buffer of shape `[seq_len * n_heads, head_dim]` (f32).
274/// * `output`        - Output buffer (same shape and dtype as input).
275/// * `params_buf`    - Params buffer containing `[theta, head_dim, rope_dim, 0]` as f32.
276/// * `positions_buf` - Positions buffer containing `[pos_0, pos_1, ...]` as u32 (length = seq_len).
277/// * `freq_factors`  - Optional freq_factors buffer of shape `[rope_dim/2]` (f32).
278///                     Pass `None` for standard RoPE (equivalent to all-ones).
279/// * `seq_len`       - Number of sequence positions.
280/// * `n_heads`       - Number of attention heads.
281/// * `head_dim`      - Dimension of each head.
282/// * `rope_dim`      - Number of dimensions to rotate (must be even, <= head_dim).
283///
284/// # Errors
285///
286/// Returns `MlxError::InvalidArgument` if parameters are invalid.
287#[allow(clippy::too_many_arguments)]
288pub fn dispatch_rope_neox_f32(
289    encoder: &mut CommandEncoder,
290    registry: &mut KernelRegistry,
291    device: &metal::DeviceRef,
292    input: &MlxBuffer,
293    output: &MlxBuffer,
294    params_buf: &MlxBuffer,
295    positions_buf: &MlxBuffer,
296    freq_factors: Option<&MlxBuffer>,
297    seq_len: u32,
298    n_heads: u32,
299    head_dim: u32,
300    rope_dim: u32,
301) -> Result<()> {
302    use super::encode_helpers::{as_bytes, encode_with_args, KernelArg};
303
304    if rope_dim % 2 != 0 {
305        return Err(MlxError::InvalidArgument(format!(
306            "RoPE neox f32 rope_dim must be even, got {}",
307            rope_dim
308        )));
309    }
310    if rope_dim > head_dim {
311        return Err(MlxError::InvalidArgument(format!(
312            "RoPE neox f32 rope_dim ({}) must be <= head_dim ({})",
313            rope_dim, head_dim
314        )));
315    }
316    if head_dim == 0 || seq_len == 0 || n_heads == 0 {
317        return Err(MlxError::InvalidArgument(
318            "RoPE neox f32 head_dim, seq_len, and n_heads must be > 0".into(),
319        ));
320    }
321
322    let n_rows = (seq_len as usize) * (n_heads as usize);
323    let expected_elements = n_rows * (head_dim as usize);
324    if input.element_count() != expected_elements {
325        return Err(MlxError::InvalidArgument(format!(
326            "RoPE neox f32 input element count {} != seq_len({}) * n_heads({}) * head_dim({})",
327            input.element_count(),
328            seq_len,
329            n_heads,
330            head_dim
331        )));
332    }
333    if output.element_count() != expected_elements {
334        return Err(MlxError::InvalidArgument(format!(
335            "RoPE neox f32 output element count {} != seq_len({}) * n_heads({}) * head_dim({})",
336            output.element_count(),
337            seq_len,
338            n_heads,
339            head_dim
340        )));
341    }
342
343    let pipeline = registry.get_pipeline("rope_neox_f32", device)?;
344    let half_rope = rope_dim / 2;
345
346    let has_ff = freq_factors.is_some();
347    let gpu_rope_params = GpuRopeNeoxF32Params {
348        n_heads,
349        has_freq_factors: u32::from(has_ff),
350    };
351
352    // When no freq_factors buffer is provided, bind the input buffer as a
353    // harmless dummy at buffer(5) — Metal requires all declared buffers to
354    // be bound. The shader checks has_freq_factors before reading.
355    let ff_buf = freq_factors.unwrap_or(input);
356
357    // Grid: (half_rope, n_rows) — one thread per pair per row
358    let tg_x = std::cmp::min(64, half_rope as u64);
359    let tg_y = std::cmp::min(4, n_rows as u64);
360
361    encode_with_args(
362        encoder,
363        pipeline,
364        &[
365            (0, KernelArg::Buffer(input)),
366            (1, KernelArg::Buffer(output)),
367            (2, KernelArg::Buffer(params_buf)),
368            (3, KernelArg::Buffer(positions_buf)),
369            (4, KernelArg::Bytes(as_bytes(&gpu_rope_params))),
370            (5, KernelArg::Buffer(ff_buf)),
371        ],
372        MTLSize::new(half_rope as u64, n_rows as u64, 1),
373        MTLSize::new(tg_x, tg_y, 1),
374    );
375
376    Ok(())
377}
mlx_native/ops/rope.rs

mlx_native/ops/
rope.rs
