spark_model/layers/ops/
ssm_preproc.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/// Deinterleave QKVZ projection output from GQA-grouped to sequential layout.
17///
18/// Input: interleaved [num_groups × (kd + kd + vpg*vd + vpg*vd)]
19/// Output: sequential [Q_total | K_total | V_total | Z_total]
20///
21/// Kernel: `deinterleave_qkvz(interleaved, output, num_groups, head_k_dim,
22///          vheads_per_group, head_v_dim)`
23/// Grid: (ceil(total/256), 1, 1)  Block: (256, 1, 1)
24pub fn deinterleave_qkvz(
25    gpu: &dyn GpuBackend,
26    kernel: KernelHandle,
27    interleaved: DevicePtr,
28    output: DevicePtr,
29    num_tokens: u32,
30    num_groups: u32,
31    head_k_dim: u32,
32    vheads_per_group: u32,
33    head_v_dim: u32,
34    stream: u64,
35) -> Result<()> {
36    let group_dim = 2 * head_k_dim + 2 * vheads_per_group * head_v_dim;
37    let total = num_groups * group_dim;
38    KernelLaunch::new(gpu, kernel)
39        .grid([num_tokens, div_ceil(total, 256), 1])
40        .block([256, 1, 1])
41        .arg_ptr(interleaved)
42        .arg_ptr(output)
43        .arg_u32(num_groups)
44        .arg_u32(head_k_dim)
45        .arg_u32(vheads_per_group)
46        .arg_u32(head_v_dim)
47        .launch(stream)
48}
49
50/// Deinterleave Q/Gate from per-head interleaved to contiguous layout (in-place).
51///
52/// Input layout:  [Q_h0(hd), G_h0(hd), Q_h1(hd), G_h1(hd), ...]
53/// Output layout: [Q_h0(hd), Q_h1(hd), ..., G_h0(hd), G_h1(hd), ...]
54///
55/// Kernel: `deinterleave_qg(data, num_heads, head_dim)`
56/// Grid: (1, 1, 1)  Block: (256, 1, 1)
57/// Dynamic shared memory: num_heads * head_dim * 2 * 2 bytes
58pub fn deinterleave_qg(
59    gpu: &dyn GpuBackend,
60    kernel: KernelHandle,
61    data: DevicePtr,
62    num_tokens: u32,
63    num_heads: u32,
64    head_dim: u32,
65    stride: u32,
66    stream: u64,
67) -> Result<()> {
68    let shared_bytes = num_heads * head_dim * 2 * 2; // BF16 = 2 bytes each
69    KernelLaunch::new(gpu, kernel)
70        .grid([num_tokens, 1, 1])
71        .block([256, 1, 1])
72        .shared_mem(shared_bytes)
73        .arg_ptr(data)
74        .arg_u32(num_heads)
75        .arg_u32(head_dim)
76        .arg_u32(stride)
77        .launch(stream)
78}
79
80/// Deinterleave Q/Gate with split output — Q to separate contiguous buffer.
81///
82/// Same as [`deinterleave_qg`] but writes Q to `q_out` (contiguous [N, q_dim])
83/// instead of in-place. Gate is still written back to `data` in-place.
84/// Eliminates the per-token D2D copy loop.
85pub fn deinterleave_qg_split(
86    gpu: &dyn GpuBackend,
87    kernel: KernelHandle,
88    data: DevicePtr,
89    q_out: DevicePtr,
90    num_tokens: u32,
91    num_heads: u32,
92    head_dim: u32,
93    stride: u32,
94    stream: u64,
95) -> Result<()> {
96    let shared_bytes = num_heads * head_dim * 2 * 2;
97    KernelLaunch::new(gpu, kernel)
98        .grid([num_tokens, 1, 1])
99        .block([256, 1, 1])
100        .shared_mem(shared_bytes)
101        .arg_ptr(data)
102        .arg_ptr(q_out)
103        .arg_u32(num_heads)
104        .arg_u32(head_dim)
105        .arg_u32(stride)
106        .launch(stream)
107}
108
109/// Fused deinterleave Q/Gate + per-head Q RMS norm.
110///
111/// Combines [`deinterleave_qg_split`] + Q RMS norm into a single kernel,
112/// eliminating one global memory round-trip for Q data.
113/// Gate is deinterleaved to `data[q_total..]`, Q is deinterleaved → normalized → `q_out`.
114pub fn deinterleave_qg_split_qnorm(
115    gpu: &dyn GpuBackend,
116    kernel: KernelHandle,
117    data: DevicePtr,
118    q_out: DevicePtr,
119    q_norm_weight: DevicePtr,
120    num_tokens: u32,
121    num_heads: u32,
122    head_dim: u32,
123    stride: u32,
124    eps: f32,
125    stream: u64,
126) -> Result<()> {
127    let shared_bytes = num_heads * head_dim * 2 * 2;
128    KernelLaunch::new(gpu, kernel)
129        .grid([num_tokens, 1, 1])
130        .block([256, 1, 1])
131        .shared_mem(shared_bytes)
132        .arg_ptr(data)
133        .arg_ptr(q_out)
134        .arg_ptr(q_norm_weight)
135        .arg_u32(num_heads)
136        .arg_u32(head_dim)
137        .arg_u32(stride)
138        .arg_f32(eps)
139        .launch(stream)
140}
141
142/// Fused deinterleave Q/Gate + per-head Q RMS norm + MRoPE.
143///
144/// Holo/Qwen3.6 MRoPE prefill fast path. Gate is deinterleaved to
145/// `data[q_total..]`; Q is deinterleaved, normalized, MRoPE-rotated, then
146/// written to `q_out`.
147#[allow(clippy::too_many_arguments)]
148pub fn deinterleave_qg_split_qnorm_mrope(
149    gpu: &dyn GpuBackend,
150    kernel: KernelHandle,
151    data: DevicePtr,
152    q_out: DevicePtr,
153    q_norm_weight: DevicePtr,
154    pos_t: DevicePtr,
155    pos_h: DevicePtr,
156    pos_w: DevicePtr,
157    num_tokens: u32,
158    num_heads: u32,
159    head_dim: u32,
160    stride: u32,
161    rotary_dim: u32,
162    eps: f32,
163    theta: f32,
164    stream: u64,
165) -> Result<()> {
166    let raw_shared = num_heads * head_dim * 2 * 2;
167    let norm_shared = num_heads * head_dim * 2;
168    KernelLaunch::new(gpu, kernel)
169        .grid([num_tokens, 1, 1])
170        .block([256, 1, 1])
171        .shared_mem(raw_shared + norm_shared)
172        .arg_ptr(data)
173        .arg_ptr(q_out)
174        .arg_ptr(q_norm_weight)
175        .arg_ptr(pos_t)
176        .arg_ptr(pos_h)
177        .arg_ptr(pos_w)
178        .arg_u32(num_heads)
179        .arg_u32(head_dim)
180        .arg_u32(stride)
181        .arg_u32(rotary_dim)
182        .arg_f32(eps)
183        .arg_f32(theta)
184        .launch(stream)
185}
186
187/// Batched sigmoid gate multiply across multiple tokens.
188///
189/// Replaces per-token [`sigmoid_gate_mul`] launches with a single kernel.
190/// `gate` is strided (gate_stride elements between tokens in gate buffer).
191pub fn sigmoid_gate_mul_batched(
192    gpu: &dyn GpuBackend,
193    kernel: KernelHandle,
194    input: DevicePtr,
195    gate: DevicePtr,
196    output: DevicePtr,
197    dim: u32,
198    gate_stride: u32,
199    num_tokens: u32,
200    stream: u64,
201) -> Result<()> {
202    let total = num_tokens * dim;
203    KernelLaunch::new(gpu, kernel)
204        .grid([div_ceil(total, 256), 1, 1])
205        .block([256, 1, 1])
206        .arg_ptr(input)
207        .arg_ptr(gate)
208        .arg_ptr(output)
209        .arg_u32(dim)
210        .arg_u32(gate_stride)
211        .arg_u32(total)
212        .launch(stream)
213}
214
215/// Compute GDN gates from interleaved BA projection + learned A_log/dt_bias.
216///
217/// Outputs FP32 gate (decay) and beta (write gate) for each value head.
218///
219/// Kernel: `compute_gdn_gates(ba_interleaved, A_log, dt_bias, gate_out,
220///          beta_out, num_v_heads, num_groups, vheads_per_group)`
221/// Grid: (1, 1, 1)  Block: (num_v_heads, 1, 1)
222pub fn compute_gdn_gates(
223    gpu: &dyn GpuBackend,
224    kernel: KernelHandle,
225    ba_interleaved: DevicePtr,
226    a_log: DevicePtr,
227    dt_bias: DevicePtr,
228    gate_out: DevicePtr,
229    beta_out: DevicePtr,
230    num_tokens: u32,
231    num_v_heads: u32,
232    num_groups: u32,
233    vheads_per_group: u32,
234    ba_stride: u32,
235    stream: u64,
236) -> Result<()> {
237    KernelLaunch::new(gpu, kernel)
238        .grid([num_tokens, 1, 1])
239        .block([num_v_heads, 1, 1])
240        .arg_ptr(ba_interleaved)
241        .arg_ptr(a_log)
242        .arg_ptr(dt_bias)
243        .arg_ptr(gate_out)
244        .arg_ptr(beta_out)
245        .arg_u32(num_v_heads)
246        .arg_u32(num_groups)
247        .arg_u32(vheads_per_group)
248        .arg_u32(ba_stride)
249        .launch(stream)
250}
251
252/// Fused BA projection + GDN gates: dense GEMV + gate/beta transforms.
253///
254/// Combines `dense_gemv(input, ba_weight, ba_out, N, K)` and
255/// `compute_gdn_gates(ba_out, a_log, dt_bias, gate, beta)` into a single
256/// kernel, eliminating the intermediate ba_out buffer and one graph node.
257///
258/// Kernel: `dense_gemv_ba_gates(A, B, A_log, dt_bias, gate, beta, N, K, vpg)`
259/// Grid: (ceil(N/4), 1, 1)  Block: (256, 1, 1)
260#[allow(clippy::too_many_arguments)]
261pub fn dense_gemv_ba_gates(
262    gpu: &dyn GpuBackend,
263    kernel: KernelHandle,
264    input: DevicePtr,
265    ba_weight: &DenseWeight,
266    a_log: DevicePtr,
267    dt_bias: DevicePtr,
268    gate_out: DevicePtr,
269    beta_out: DevicePtr,
270    n: u32,
271    k: u32,
272    vheads_per_group: u32,
273    stream: u64,
274) -> Result<()> {
275    KernelLaunch::new(gpu, kernel)
276        .grid([div_ceil(n, 4), 1, 1])
277        .block([256, 1, 1])
278        .arg_ptr(input)
279        .arg_ptr(ba_weight.weight)
280        .arg_ptr(a_log)
281        .arg_ptr(dt_bias)
282        .arg_ptr(gate_out)
283        .arg_ptr(beta_out)
284        .arg_u32(n)
285        .arg_u32(k)
286        .arg_u32(vheads_per_group)
287        .launch(stream)
288}
289
290/// Fused BA GEMM + GDN gates for prefill (token-parallel).
291///
292/// Replaces `dense_gemm(normed, ba_weight) + compute_gdn_gates` in the prefill path.
293/// Uses vectorized uint4 loads and warp-shuffle reduction per token, adding a token
294/// dimension via blockIdx.y. Skips the intermediate ba_out buffer entirely.
295///
296/// Output layout (shared gate_out buffer):
297///   gate_out[token * gate_stride + vh]      = gate (alpha→exp transform)
298///   gate_out[token * gate_stride + nv + vh] = beta (sigmoid)
299///
300/// Kernel: `dense_gemm_ba_gates_prefill(A, B, A_log, dt_bias, gate_out, M, N, K,
301///          K_stride, gate_stride, nv, vpg)`
302/// Grid: (ceil(N/4), M_tokens, 1)  Block: (256, 1, 1)
303#[allow(clippy::too_many_arguments)]
304pub fn dense_gemm_ba_gates_prefill(
305    gpu: &dyn GpuBackend,
306    kernel: KernelHandle,
307    input: DevicePtr,        // [M, K_stride] activations (BF16)
308    ba_weight: &DenseWeight, // [N, K] BA weight (BF16, row-major)
309    a_log: DevicePtr,
310    dt_bias: DevicePtr,
311    gate_out: DevicePtr, // [M, gate_stride] FP32 unified gate+beta buffer
312    m: u32,              // num_tokens
313    n: u32,              // ba_size (64)
314    k: u32,              // hidden_size (2048)
315    k_stride: u32,       // BF16 elements between tokens in input (= k)
316    gate_stride: u32,    // FP32 elements between tokens in gate_out (= 2*nv)
317    nv: u32,             // num_v_heads (32)
318    vheads_per_group: u32,
319    stream: u64,
320) -> Result<()> {
321    KernelLaunch::new(gpu, kernel)
322        .grid([div_ceil(n, 4), m, 1])
323        .block([256, 1, 1])
324        .arg_ptr(input)
325        .arg_ptr(ba_weight.weight)
326        .arg_ptr(a_log)
327        .arg_ptr(dt_bias)
328        .arg_ptr(gate_out)
329        .arg_u32(m)
330        .arg_u32(n)
331        .arg_u32(k)
332        .arg_u32(k_stride)
333        .arg_u32(gate_stride)
334        .arg_u32(nv)
335        .arg_u32(vheads_per_group)
336        .launch(stream)
337}
338
339// ── Sampling ─────────────────────────────────────────────────────
spark_model/layers/ops/ssm_preproc.rs

spark_model/layers/ops/
ssm_preproc.rs
