metrale_model_layers/layers/ops/
embeddings.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Rotary position embedding launchers: plain, strided, proportional,
4//! interleaved multimodal (MRoPE) and table-driven (YaRN).
5//!
6//! Owner: model-layers ops.
7//! Invariants: none beyond the types.
8
9#![allow(unused_imports)]
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
14
15use crate::layers::moe;
16use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
17
18use super::*;
19
20/// 2026-09-25: [`rope`] for rows that are not packed: Q rows `q_row_stride` and
21/// K rows `k_row_stride` elements apart, all `num_tokens` rows in one launch.
22/// `rope_forward_strided` differs from `rope_forward` only in the row address,
23/// so the packed strides give the same result as [`rope`].
24#[allow(clippy::too_many_arguments)]
25pub fn rope_strided(
26    gpu: &dyn GpuBackend,
27    kernel: KernelHandle,
28    q: DevicePtr,
29    k: DevicePtr,
30    positions: DevicePtr,
31    num_tokens: u32,
32    num_q_heads: u32,
33    num_kv_heads: u32,
34    head_dim: u32,
35    rotary_dim: u32,
36    theta: f32,
37    q_row_stride: u32,
38    k_row_stride: u32,
39    stream: u64,
40) -> Result<()> {
41    assert!(
42        rotary_dim > 0,
43        "rope_strided: rotary_dim=0, nq={num_q_heads} nkv={num_kv_heads} hd={head_dim}"
44    );
45    let half_rot = (rotary_dim / 2).max(1);
46    let pos_per_block = (128 / half_rot).max(1);
47    let seq_blocks = div_ceil(num_tokens, pos_per_block);
48    KernelLaunch::new(gpu, kernel)
49        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
50        .block([128, 1, 1])
51        .arg_ptr(q)
52        .arg_ptr(k)
53        .arg_ptr(positions)
54        .arg_u32(num_tokens)
55        .arg_u32(num_q_heads)
56        .arg_u32(num_kv_heads)
57        .arg_u32(head_dim)
58        .arg_u32(rotary_dim)
59        .arg_f32(theta)
60        .arg_u32(q_row_stride)
61        .arg_u32(k_row_stride)
62        .launch(stream)
63}
64
65/// 2026-09-25: Apply rotary position embeddings to Q and K in place, rows packed
66/// (`num_q_heads * head_dim` and `num_kv_heads * head_dim` elements apart).
67/// `positions` is a device `u32[seq_len]`.
68pub fn rope(
69    gpu: &dyn GpuBackend,
70    kernel: KernelHandle,
71    q: DevicePtr,
72    k: DevicePtr,
73    positions: DevicePtr,
74    seq_len: u32,
75    num_q_heads: u32,
76    num_kv_heads: u32,
77    head_dim: u32,
78    rotary_dim: u32,
79    theta: f32,
80    stream: u64,
81) -> Result<()> {
82    assert!(
83        rotary_dim > 0,
84        "rope: rotary_dim=0, nq={num_q_heads} nkv={num_kv_heads} hd={head_dim}"
85    );
86    let half_rot = (rotary_dim / 2).max(1);
87    let pos_per_block = (128 / half_rot).max(1);
88    let seq_blocks = div_ceil(seq_len, pos_per_block);
89    KernelLaunch::new(gpu, kernel)
90        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
91        .block([128, 1, 1])
92        .arg_ptr(q)
93        .arg_ptr(k)
94        .arg_ptr(positions)
95        .arg_u32(seq_len)
96        .arg_u32(num_q_heads)
97        .arg_u32(num_kv_heads)
98        .arg_u32(head_dim)
99        .arg_u32(rotary_dim)
100        .arg_f32(theta)
101        .launch(stream)
102}
103
104/// 2026-09-25: Proportional RoPE (Gemma-4 full-attention layers).
105///
106/// Rotation pairs are (i, i + head_dim/2) for i in [0, rope_angles), and the
107/// frequency exponent's denominator is `head_dim`.
108#[allow(clippy::too_many_arguments)]
109pub fn rope_proportional(
110    gpu: &dyn GpuBackend,
111    kernel: KernelHandle,
112    q: DevicePtr,
113    k: DevicePtr,
114    positions: DevicePtr,
115    seq_len: u32,
116    num_q_heads: u32,
117    num_kv_heads: u32,
118    head_dim: u32,
119    rope_angles: u32,
120    theta: f32,
121    stream: u64,
122) -> Result<()> {
123    assert!(rope_angles > 0, "rope_proportional: rope_angles=0");
124    let pairs_per_pos = rope_angles.max(1);
125    let pos_per_block = (128 / pairs_per_pos).max(1);
126    let seq_blocks = div_ceil(seq_len, pos_per_block);
127    KernelLaunch::new(gpu, kernel)
128        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
129        .block([128, 1, 1])
130        .arg_ptr(q)
131        .arg_ptr(k)
132        .arg_ptr(positions)
133        .arg_u32(seq_len)
134        .arg_u32(num_q_heads)
135        .arg_u32(num_kv_heads)
136        .arg_u32(head_dim)
137        .arg_u32(rope_angles)
138        .arg_f32(theta)
139        .launch(stream)
140}
141
142/// 2026-09-25: Interleaved multimodal RoPE (MRoPE) with three position streams
143/// (`pos_t`, `pos_h`, `pos_w`); rotary pair `i` takes its position from stream
144/// `i % 3`. With the same pointer for all three, the math is that of `rope`.
145#[allow(clippy::too_many_arguments)]
146pub fn rope_mrope_interleaved(
147    gpu: &dyn GpuBackend,
148    kernel: KernelHandle,
149    q: DevicePtr,
150    k: DevicePtr,
151    pos_t: DevicePtr,
152    pos_h: DevicePtr,
153    pos_w: DevicePtr,
154    seq_len: u32,
155    num_q_heads: u32,
156    num_kv_heads: u32,
157    head_dim: u32,
158    rotary_dim: u32,
159    theta: f32,
160    stream: u64,
161) -> Result<()> {
162    assert!(rotary_dim > 0, "rope_mrope_interleaved: rotary_dim=0");
163    let half_rot = (rotary_dim / 2).max(1);
164    let pos_per_block = (128 / half_rot).max(1);
165    let seq_blocks = div_ceil(seq_len, pos_per_block);
166    KernelLaunch::new(gpu, kernel)
167        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
168        .block([128, 1, 1])
169        .arg_ptr(q)
170        .arg_ptr(k)
171        .arg_ptr(pos_t)
172        .arg_ptr(pos_h)
173        .arg_ptr(pos_w)
174        .arg_u32(seq_len)
175        .arg_u32(num_q_heads)
176        .arg_u32(num_kv_heads)
177        .arg_u32(head_dim)
178        .arg_u32(rotary_dim)
179        .arg_f32(theta)
180        .launch(stream)
181}
182
183/// 2026-09-25: [`rope_mrope_interleaved`] for K only, for when Q was already
184/// rotated by a fused prefill kernel.
185#[allow(clippy::too_many_arguments)]
186pub fn rope_mrope_interleaved_k_only(
187    gpu: &dyn GpuBackend,
188    kernel: KernelHandle,
189    k: DevicePtr,
190    pos_t: DevicePtr,
191    pos_h: DevicePtr,
192    pos_w: DevicePtr,
193    seq_len: u32,
194    num_kv_heads: u32,
195    head_dim: u32,
196    rotary_dim: u32,
197    theta: f32,
198    stream: u64,
199) -> Result<()> {
200    assert!(
201        rotary_dim > 0,
202        "rope_mrope_interleaved_k_only: rotary_dim=0"
203    );
204    let half_rot = (rotary_dim / 2).max(1);
205    let pos_per_block = (128 / half_rot).max(1);
206    let seq_blocks = div_ceil(seq_len, pos_per_block);
207    KernelLaunch::new(gpu, kernel)
208        .grid([num_kv_heads, seq_blocks, 1])
209        .block([128, 1, 1])
210        .arg_ptr(k)
211        .arg_ptr(pos_t)
212        .arg_ptr(pos_h)
213        .arg_ptr(pos_w)
214        .arg_u32(seq_len)
215        .arg_u32(num_kv_heads)
216        .arg_u32(head_dim)
217        .arg_u32(rotary_dim)
218        .arg_f32(theta)
219        .launch(stream)
220}
221
222/// 2026-09-25: RoPE with frequencies from a precomputed `inv_freq` table
223/// (`[rotary_dim/2]` FP32); `rope_forward_yarn` ignores `theta`.
224#[allow(clippy::too_many_arguments)]
225pub fn rope_yarn(
226    gpu: &dyn GpuBackend,
227    kernel: KernelHandle,
228    q: DevicePtr,
229    k: DevicePtr,
230    positions: DevicePtr,
231    seq_len: u32,
232    num_q_heads: u32,
233    num_kv_heads: u32,
234    head_dim: u32,
235    rotary_dim: u32,
236    inv_freq: DevicePtr,
237    theta: f32,
238    stream: u64,
239) -> Result<()> {
240    assert!(
241        rotary_dim > 0,
242        "rope: rotary_dim=0, nq={num_q_heads} nkv={num_kv_heads} hd={head_dim}"
243    );
244    let half_rot = (rotary_dim / 2).max(1);
245    let pos_per_block = (128 / half_rot).max(1);
246    let seq_blocks = div_ceil(seq_len, pos_per_block);
247    KernelLaunch::new(gpu, kernel)
248        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
249        .block([128, 1, 1])
250        .arg_ptr(q)
251        .arg_ptr(k)
252        .arg_ptr(positions)
253        .arg_u32(seq_len)
254        .arg_u32(num_q_heads)
255        .arg_u32(num_kv_heads)
256        .arg_u32(head_dim)
257        .arg_u32(rotary_dim)
258        .arg_ptr(inv_freq)
259        .arg_f32(theta)
260        .launch(stream)
261}
262
263/// 2026-09-25: RoPE from an `inv_freq` table with cosine and sine scaled by
264/// `attention_factor`.
265#[allow(clippy::too_many_arguments)]
266pub fn rope_yarn_scaled(
267    gpu: &dyn GpuBackend,
268    kernel: KernelHandle,
269    q: DevicePtr,
270    k: DevicePtr,
271    positions: DevicePtr,
272    seq_len: u32,
273    num_q_heads: u32,
274    num_kv_heads: u32,
275    head_dim: u32,
276    rotary_dim: u32,
277    inv_freq: DevicePtr,
278    attention_factor: f32,
279    stream: u64,
280) -> Result<()> {
281    assert!(rotary_dim > 0, "rope_yarn_scaled: rotary_dim=0");
282    let half_rot = (rotary_dim / 2).max(1);
283    let pos_per_block = (128 / half_rot).max(1);
284    let seq_blocks = div_ceil(seq_len, pos_per_block);
285    KernelLaunch::new(gpu, kernel)
286        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
287        .block([128, 1, 1])
288        .arg_ptr(q)
289        .arg_ptr(k)
290        .arg_ptr(positions)
291        .arg_u32(seq_len)
292        .arg_u32(num_q_heads)
293        .arg_u32(num_kv_heads)
294        .arg_u32(head_dim)
295        .arg_u32(rotary_dim)
296        .arg_ptr(inv_freq)
297        .arg_f32(attention_factor)
298        .launch(stream)
299}
metrale_model_layers/layers/ops/embeddings.rs

metrale_model_layers/layers/ops/
embeddings.rs
