spark_model/layers/ops/
embeddings.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `ops.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
10
11use crate::layers::moe;
12use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
13
14use super::*;
15
16/// RoPE: apply rotary position embeddings to Q and K in-place.
17///
18/// Kernel: `rope_forward(Q, K, positions, seq_len, num_q_heads,
19///          num_kv_heads, head_dim, rotary_dim, theta)`
20/// Grid: (num_q_heads + num_kv_heads, ceil(seq_len/4), 1)
21/// Block: (128, 1, 1)
22///
23/// `positions` must be a device pointer to a `u32[seq_len]` array.
24/// Strided RoPE: rotates ALL `num_tokens` rows in ONE launch.
25///
26/// `rope` above derives each row's address from a PACKED layout
27/// (`num_*_heads * head_dim` between tokens). The multi-seq decode buffer is not
28/// packed — Q and K live inside one interleaved `[Q|K|V|gate]` block whose rows
29/// sit `per_seq_qkv` apart — so that path was calling `rope` once per sequence
30/// with `seq_len = 1`: 258 launches/step at 4.6 us = 1.18 ms across the 16
31/// attention layers.
32///
33/// Bit-identical to n packed launches: same math, same ordering, only the row
34/// address differs. Passing the packed strides reproduces `rope` exactly.
35#[allow(clippy::too_many_arguments)]
36pub fn rope_strided(
37    gpu: &dyn GpuBackend,
38    kernel: KernelHandle,
39    q: DevicePtr,
40    k: DevicePtr,
41    positions: DevicePtr,
42    num_tokens: u32,
43    num_q_heads: u32,
44    num_kv_heads: u32,
45    head_dim: u32,
46    rotary_dim: u32,
47    theta: f32,
48    q_row_stride: u32,
49    k_row_stride: u32,
50    stream: u64,
51) -> Result<()> {
52    assert!(
53        rotary_dim > 0,
54        "rope_strided: rotary_dim=0, nq={num_q_heads} nkv={num_kv_heads} hd={head_dim}"
55    );
56    let half_rot = (rotary_dim / 2).max(1);
57    let pos_per_block = (128 / half_rot).max(1);
58    let seq_blocks = div_ceil(num_tokens, pos_per_block);
59    KernelLaunch::new(gpu, kernel)
60        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
61        .block([128, 1, 1])
62        .arg_ptr(q)
63        .arg_ptr(k)
64        .arg_ptr(positions)
65        .arg_u32(num_tokens)
66        .arg_u32(num_q_heads)
67        .arg_u32(num_kv_heads)
68        .arg_u32(head_dim)
69        .arg_u32(rotary_dim)
70        .arg_f32(theta)
71        .arg_u32(q_row_stride)
72        .arg_u32(k_row_stride)
73        .launch(stream)
74}
75
76pub fn rope(
77    gpu: &dyn GpuBackend,
78    kernel: KernelHandle,
79    q: DevicePtr,
80    k: DevicePtr,
81    positions: DevicePtr,
82    seq_len: u32,
83    num_q_heads: u32,
84    num_kv_heads: u32,
85    head_dim: u32,
86    rotary_dim: u32,
87    theta: f32,
88    stream: u64,
89) -> Result<()> {
90    assert!(
91        rotary_dim > 0,
92        "rope: rotary_dim=0, nq={num_q_heads} nkv={num_kv_heads} hd={head_dim}"
93    );
94    let half_rot = (rotary_dim / 2).max(1);
95    let pos_per_block = (128 / half_rot).max(1);
96    let seq_blocks = div_ceil(seq_len, pos_per_block);
97    KernelLaunch::new(gpu, kernel)
98        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
99        .block([128, 1, 1])
100        .arg_ptr(q)
101        .arg_ptr(k)
102        .arg_ptr(positions)
103        .arg_u32(seq_len)
104        .arg_u32(num_q_heads)
105        .arg_u32(num_kv_heads)
106        .arg_u32(head_dim)
107        .arg_u32(rotary_dim)
108        .arg_f32(theta)
109        .launch(stream)
110}
111
112/// Proportional RoPE (Gemma-4 full-attention layers).
113///
114/// Rotation pairs are (i, i + head_dim/2) for i in [0, rope_angles).
115/// Frequency denominator is `head_dim` (not rotary_dim).
116#[allow(clippy::too_many_arguments)]
117pub fn rope_proportional(
118    gpu: &dyn GpuBackend,
119    kernel: KernelHandle,
120    q: DevicePtr,
121    k: DevicePtr,
122    positions: DevicePtr,
123    seq_len: u32,
124    num_q_heads: u32,
125    num_kv_heads: u32,
126    head_dim: u32,
127    rope_angles: u32,
128    theta: f32,
129    stream: u64,
130) -> Result<()> {
131    assert!(rope_angles > 0, "rope_proportional: rope_angles=0");
132    let pairs_per_pos = rope_angles.max(1);
133    let pos_per_block = (128 / pairs_per_pos).max(1);
134    let seq_blocks = div_ceil(seq_len, pos_per_block);
135    KernelLaunch::new(gpu, kernel)
136        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
137        .block([128, 1, 1])
138        .arg_ptr(q)
139        .arg_ptr(k)
140        .arg_ptr(positions)
141        .arg_u32(seq_len)
142        .arg_u32(num_q_heads)
143        .arg_u32(num_kv_heads)
144        .arg_u32(head_dim)
145        .arg_u32(rope_angles)
146        .arg_f32(theta)
147        .launch(stream)
148}
149
150/// MRoPE (interleaved multi-modal rotary) for Qwen3.6.
151///
152/// Applies rotary embedding using three separate position-ID streams
153/// (`pos_t`, `pos_h`, `pos_w`). Each rotary pair is owned by one of the
154/// three sections, chosen by `pair_idx % 3` (round-robin). For text-only
155/// serving pass the same pointer for all three streams — the result is
156/// bit-identical to scalar RoPE.
157///
158/// Kernel: `rope_forward_mrope_interleaved`
159#[allow(clippy::too_many_arguments)]
160pub fn rope_mrope_interleaved(
161    gpu: &dyn GpuBackend,
162    kernel: KernelHandle,
163    q: DevicePtr,
164    k: DevicePtr,
165    pos_t: DevicePtr,
166    pos_h: DevicePtr,
167    pos_w: DevicePtr,
168    seq_len: u32,
169    num_q_heads: u32,
170    num_kv_heads: u32,
171    head_dim: u32,
172    rotary_dim: u32,
173    theta: f32,
174    stream: u64,
175) -> Result<()> {
176    assert!(rotary_dim > 0, "rope_mrope_interleaved: rotary_dim=0");
177    let half_rot = (rotary_dim / 2).max(1);
178    let pos_per_block = (128 / half_rot).max(1);
179    let seq_blocks = div_ceil(seq_len, pos_per_block);
180    KernelLaunch::new(gpu, kernel)
181        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
182        .block([128, 1, 1])
183        .arg_ptr(q)
184        .arg_ptr(k)
185        .arg_ptr(pos_t)
186        .arg_ptr(pos_h)
187        .arg_ptr(pos_w)
188        .arg_u32(seq_len)
189        .arg_u32(num_q_heads)
190        .arg_u32(num_kv_heads)
191        .arg_u32(head_dim)
192        .arg_u32(rotary_dim)
193        .arg_f32(theta)
194        .launch(stream)
195}
196
197/// MRoPE for K only. Used when Q was already rotated by a fused Q prefill
198/// kernel.
199#[allow(clippy::too_many_arguments)]
200pub fn rope_mrope_interleaved_k_only(
201    gpu: &dyn GpuBackend,
202    kernel: KernelHandle,
203    k: DevicePtr,
204    pos_t: DevicePtr,
205    pos_h: DevicePtr,
206    pos_w: DevicePtr,
207    seq_len: u32,
208    num_kv_heads: u32,
209    head_dim: u32,
210    rotary_dim: u32,
211    theta: f32,
212    stream: u64,
213) -> Result<()> {
214    assert!(
215        rotary_dim > 0,
216        "rope_mrope_interleaved_k_only: rotary_dim=0"
217    );
218    let half_rot = (rotary_dim / 2).max(1);
219    let pos_per_block = (128 / half_rot).max(1);
220    let seq_blocks = div_ceil(seq_len, pos_per_block);
221    KernelLaunch::new(gpu, kernel)
222        .grid([num_kv_heads, seq_blocks, 1])
223        .block([128, 1, 1])
224        .arg_ptr(k)
225        .arg_ptr(pos_t)
226        .arg_ptr(pos_h)
227        .arg_ptr(pos_w)
228        .arg_u32(seq_len)
229        .arg_u32(num_kv_heads)
230        .arg_u32(head_dim)
231        .arg_u32(rotary_dim)
232        .arg_f32(theta)
233        .launch(stream)
234}
235
236/// RoPE with precomputed YaRN inv_freq table (Mistral Small 4).
237/// The kernel reads frequencies from the table instead of computing from theta.
238#[allow(clippy::too_many_arguments)]
239pub fn rope_yarn(
240    gpu: &dyn GpuBackend,
241    kernel: KernelHandle,
242    q: DevicePtr,
243    k: DevicePtr,
244    positions: DevicePtr,
245    seq_len: u32,
246    num_q_heads: u32,
247    num_kv_heads: u32,
248    head_dim: u32,
249    rotary_dim: u32,
250    inv_freq: DevicePtr,
251    theta: f32,
252    stream: u64,
253) -> Result<()> {
254    assert!(
255        rotary_dim > 0,
256        "rope: rotary_dim=0, nq={num_q_heads} nkv={num_kv_heads} hd={head_dim}"
257    );
258    let half_rot = (rotary_dim / 2).max(1);
259    let pos_per_block = (128 / half_rot).max(1);
260    let seq_blocks = div_ceil(seq_len, pos_per_block);
261    KernelLaunch::new(gpu, kernel)
262        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
263        .block([128, 1, 1])
264        .arg_ptr(q)
265        .arg_ptr(k)
266        .arg_ptr(positions)
267        .arg_u32(seq_len)
268        .arg_u32(num_q_heads)
269        .arg_u32(num_kv_heads)
270        .arg_u32(head_dim)
271        .arg_u32(rotary_dim)
272        .arg_ptr(inv_freq)
273        .arg_f32(theta)
274        .launch(stream)
275}
276
277/// YaRN RoPE for standard attention with explicit cosine/sine amplitude.
278#[allow(clippy::too_many_arguments)]
279pub fn rope_yarn_scaled(
280    gpu: &dyn GpuBackend,
281    kernel: KernelHandle,
282    q: DevicePtr,
283    k: DevicePtr,
284    positions: DevicePtr,
285    seq_len: u32,
286    num_q_heads: u32,
287    num_kv_heads: u32,
288    head_dim: u32,
289    rotary_dim: u32,
290    inv_freq: DevicePtr,
291    attention_factor: f32,
292    stream: u64,
293) -> Result<()> {
294    assert!(rotary_dim > 0, "rope_yarn_scaled: rotary_dim=0");
295    let half_rot = (rotary_dim / 2).max(1);
296    let pos_per_block = (128 / half_rot).max(1);
297    let seq_blocks = div_ceil(seq_len, pos_per_block);
298    KernelLaunch::new(gpu, kernel)
299        .grid([num_q_heads + num_kv_heads, seq_blocks, 1])
300        .block([128, 1, 1])
301        .arg_ptr(q)
302        .arg_ptr(k)
303        .arg_ptr(positions)
304        .arg_u32(seq_len)
305        .arg_u32(num_q_heads)
306        .arg_u32(num_kv_heads)
307        .arg_u32(head_dim)
308        .arg_u32(rotary_dim)
309        .arg_ptr(inv_freq)
310        .arg_f32(attention_factor)
311        .launch(stream)
312}
spark_model/layers/ops/embeddings.rs

spark_model/layers/ops/
embeddings.rs
