spark_model/layers/ops/
activations.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/// Fused SiLU activation: output = SiLU(gate) * up.
17///
18/// Kernel: `silu_mul_separate(gate, up, output, n)`
19/// Grid: (ceil(n/256), 1, 1)  Block: (256, 1, 1)
20pub fn silu_mul(
21    gpu: &dyn GpuBackend,
22    kernel: KernelHandle,
23    gate: DevicePtr,
24    up: DevicePtr,
25    output: DevicePtr,
26    num_elements: u32,
27    stream: u64,
28) -> Result<()> {
29    KernelLaunch::new(gpu, kernel)
30        .grid([div_ceil(num_elements, 256), 1, 1])
31        .block([256, 1, 1])
32        .arg_ptr(gate)
33        .arg_ptr(up)
34        .arg_ptr(output)
35        .arg_u32(num_elements)
36        .launch(stream)
37}
38
39/// [`silu_mul`] over ROW-STRIDED operands — the consumer of the fused
40/// dense-FFN gate+up projection (#927).
41///
42/// `gate` and `up` are two column halves of ONE `[rows, in_stride]` BF16
43/// matrix (`up = gate.offset(cols * 2)`, `in_stride = 2 * cols`), which is
44/// what a single cuBLASLt call at `N = 2 * intermediate` produces; `output` is
45/// the contiguous `[rows, cols]` the down projection reads. Passing
46/// `in_stride == out_stride == cols` reproduces [`silu_mul`] exactly.
47///
48/// Kernel: `silu_mul_strided(gate, up, output, rows, cols, in_stride,
49/// out_stride)`, grid `(ceil(cols/256), rows, 1)` block `(256,1,1)` — a 2-D
50/// grid so no thread divides to recover its row. Rule:
51/// `layers/dense_ffn_gateup_fused.rs`.
52#[allow(clippy::too_many_arguments)]
53pub fn silu_mul_strided(
54    gpu: &dyn GpuBackend,
55    kernel: KernelHandle,
56    gate: DevicePtr,
57    up: DevicePtr,
58    output: DevicePtr,
59    rows: u32,
60    cols: u32,
61    in_stride: u32,
62    out_stride: u32,
63    stream: u64,
64) -> Result<()> {
65    KernelLaunch::new(gpu, kernel)
66        .grid([div_ceil(cols, 256), rows, 1])
67        .block([256, 1, 1])
68        .arg_ptr(gate)
69        .arg_ptr(up)
70        .arg_ptr(output)
71        .arg_u32(rows)
72        .arg_u32(cols)
73        .arg_u32(in_stride)
74        .arg_u32(out_stride)
75        .launch(stream)
76}
77
78/// Fused SiLU·mul + per-token-group(128) FP8-E4M3 quantization — replaces the
79/// `silu_mul` → `per_token_group_quant_fp8` pair on the W8A8 prefill down-path
80/// without materializing the BF16 intermediate. Bit-identical to the pair
81/// (product rounds through BF16 before the group max; same reduction order,
82/// scale floor, and SATFINITE encode).
83///
84/// `out_bf16` is nullable (`DevicePtr::NULL`): pass the post-SiLU BF16 buffer
85/// only when a downstream consumer needs it (expert down_proj LoRA fold).
86///
87/// Kernel: `silu_mul_quant_fp8(gate, up, out_fp8, a_scale, out_bf16, M, K)`
88/// Grid: (M, 1, 1)  Block: (128, 1, 1). Caller must ensure `k % 128 == 0`
89/// and `k / 128 <= 16` (SILU_QUANT_MAX_GROUPS) — fall back to the unfused
90/// pair otherwise.
91#[allow(clippy::too_many_arguments)]
92pub fn silu_mul_quant_fp8(
93    gpu: &dyn GpuBackend,
94    kernel: KernelHandle,
95    gate: DevicePtr,
96    up: DevicePtr,
97    out_fp8: DevicePtr,
98    a_scale: DevicePtr,
99    out_bf16: DevicePtr,
100    m: u32,
101    k: u32,
102    stream: u64,
103) -> Result<()> {
104    KernelLaunch::new(gpu, kernel)
105        .grid([m, 1, 1])
106        .block([128, 1, 1])
107        .arg_ptr(gate)
108        .arg_ptr(up)
109        .arg_ptr(out_fp8)
110        .arg_ptr(a_scale)
111        .arg_ptr(out_bf16)
112        .arg_u32(m)
113        .arg_u32(k)
114        .launch(stream)
115}
116
117/// L2 normalization (in-place): `data[i] = data[i] / sqrt(sum(data^2) + eps)`.
118///
119/// Applied per head: data is [num_heads, head_dim], each head normalized independently.
120/// Required for Gated Delta Net Q/K normalization (use_qk_l2norm_in_kernel=True).
121///
122/// Kernel: `l2_norm_bf16(data, head_dim, eps)`
123/// Grid: (num_heads, 1, 1)  Block: (min(head_dim, 1024), 1, 1)
124pub fn l2_norm(
125    gpu: &dyn GpuBackend,
126    kernel: KernelHandle,
127    data: DevicePtr,
128    num_heads: u32,
129    head_dim: u32,
130    eps: f32,
131    num_tokens: u32,
132    stride: u32,
133    stream: u64,
134) -> Result<()> {
135    KernelLaunch::new(gpu, kernel)
136        .grid([num_heads, num_tokens, 1])
137        .block([head_dim.min(1024), 1, 1])
138        .arg_ptr(data)
139        .arg_u32(head_dim)
140        .arg_f32(eps)
141        .arg_u32(stride)
142        .launch(stream)
143}
144
145/// Element-wise sigmoid gate: `output[i] = input[i] * sigmoid(gate[i])`.
146///
147/// Used for gated attention in Qwen3: attn_output = attn_output * sigmoid(q_gate).
148///
149/// Kernel: `sigmoid_gate_mul(input, gate, output, n)`
150/// Grid: (ceil(n/256), 1, 1)  Block: (256, 1, 1)
151pub fn sigmoid_gate_mul(
152    gpu: &dyn GpuBackend,
153    kernel: KernelHandle,
154    input: DevicePtr,
155    gate: DevicePtr,
156    output: DevicePtr,
157    num_elements: u32,
158    stream: u64,
159) -> Result<()> {
160    KernelLaunch::new(gpu, kernel)
161        .grid([div_ceil(num_elements, 256), 1, 1])
162        .block([256, 1, 1])
163        .arg_ptr(input)
164        .arg_ptr(gate)
165        .arg_ptr(output)
166        .arg_u32(num_elements)
167        .launch(stream)
168}
169
170/// Per-head sigmoid gate multiply with broadcast over head_dim.
171///
172/// Step 3.7 attention gate: `g_proj` produces one BF16 scalar per head.
173/// This kernel applies `output[t,h,d] = input[t,h,d] * sigmoid(gate[t,h])`
174/// where the sigmoid gate is broadcast across all `hd` dimensions of each head.
175///
176/// Kernel: `sigmoid_gate_mul_head_broadcast(input, gate, output, nq, hd, total)`
177/// Grid: (ceil(total/256), 1, 1)  Block: (256, 1, 1)
178pub fn sigmoid_gate_mul_head_broadcast(
179    gpu: &dyn GpuBackend,
180    kernel: KernelHandle,
181    input: DevicePtr,
182    gate: DevicePtr,
183    output: DevicePtr,
184    nq: u32,
185    hd: u32,
186    num_tokens: u32,
187    stream: u64,
188) -> Result<()> {
189    let total = num_tokens * nq * hd;
190    KernelLaunch::new(gpu, kernel)
191        .grid([div_ceil(total, 256), 1, 1])
192        .block([256, 1, 1])
193        .arg_ptr(input)
194        .arg_ptr(gate)
195        .arg_ptr(output)
196        .arg_u32(nq)
197        .arg_u32(hd)
198        .arg_u32(total)
199        .launch(stream)
200}
201
202/// Per-head softplus gate multiply with broadcast over `head_dim`.
203#[allow(clippy::too_many_arguments)]
204pub fn softplus_gate_mul_head_broadcast(
205    gpu: &dyn GpuBackend,
206    kernel: KernelHandle,
207    input: DevicePtr,
208    gate: DevicePtr,
209    output: DevicePtr,
210    nq: u32,
211    hd: u32,
212    num_tokens: u32,
213    stream: u64,
214) -> Result<()> {
215    let total = num_tokens * nq * hd;
216    KernelLaunch::new(gpu, kernel)
217        .grid([div_ceil(total, 256), 1, 1])
218        .block([256, 1, 1])
219        .arg_ptr(input)
220        .arg_ptr(gate)
221        .arg_ptr(output)
222        .arg_u32(nq)
223        .arg_u32(hd)
224        .arg_u32(total)
225        .launch(stream)
226}
227
228/// BF16 residual add: `residual[i] += src[i]` (in-place).
229///
230/// Kernel: `bf16_residual_add(residual, src, n)`
231/// Grid: (ceil(n/256), 1, 1)  Block: (256, 1, 1)
232pub fn residual_add(
233    gpu: &dyn GpuBackend,
234    kernel: KernelHandle,
235    residual: DevicePtr,
236    src: DevicePtr,
237    num_elements: u32,
238    stream: u64,
239) -> Result<()> {
240    KernelLaunch::new(gpu, kernel)
241        .grid([div_ceil(num_elements, 256), 1, 1])
242        .block([256, 1, 1])
243        .arg_ptr(residual)
244        .arg_ptr(src)
245        .arg_u32(num_elements)
246        .launch(stream)
247}
248
249/// BF16 scaled accumulate: `output[i] += scale * src[i]`.
250///
251/// Kernel: `bf16_scaled_add(output, src, scale, n)`
252/// Grid: (ceil(n/256), 1, 1)  Block: (256, 1, 1)
253pub fn scaled_add(
254    gpu: &dyn GpuBackend,
255    kernel: KernelHandle,
256    output: DevicePtr,
257    src: DevicePtr,
258    scale: f32,
259    num_elements: u32,
260    stream: u64,
261) -> Result<()> {
262    KernelLaunch::new(gpu, kernel)
263        .grid([div_ceil(num_elements, 256), 1, 1])
264        .block([256, 1, 1])
265        .arg_ptr(output)
266        .arg_ptr(src)
267        .arg_f32(scale)
268        .arg_u32(num_elements)
269        .launch(stream)
270}
271
272/// Sigmoid-gated blend: output = output + sigmoid_gate * src.
273///
274/// Kernel: `bf16_sigmoid_blend(output, src, sigmoid_gate, n)`
275/// Grid: (ceil(n/256), 1, 1)  Block: (256, 1, 1)
276pub fn sigmoid_blend(
277    gpu: &dyn GpuBackend,
278    kernel: KernelHandle,
279    output: DevicePtr,
280    src: DevicePtr,
281    sigmoid_gate: f32,
282    num_elements: u32,
283    stream: u64,
284) -> Result<()> {
285    KernelLaunch::new(gpu, kernel)
286        .grid([div_ceil(num_elements, 256), 1, 1])
287        .block([256, 1, 1])
288        .arg_ptr(output)
289        .arg_ptr(src)
290        .arg_f32(sigmoid_gate)
291        .arg_u32(num_elements)
292        .launch(stream)
293}
294
295// ── SSM Preprocessing ─────────────────────────────────────────────
spark_model/layers/ops/activations.rs

spark_model/layers/ops/
activations.rs
