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

spark_model/layers/ops/
activations.rs
