metrale_model_layers/layers/ops/
activations.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launch wrappers for element-wise kernels: activation-gated products (SiLU or
4//! GeLU), sigmoid and softplus gates, per-head L2 norm, residual and scaled adds.
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: `output = act(gate) * up` over `num_elements` BF16 values, where `kernel`
21/// fixes the activation: callers pass `moe_silu_mul` (SiLU) or, for a GeLU FFN, `gelu_mul`.
22/// The `moe_silu_mul.cu` copies in `gb10/deepseek-v4-flash` (also compiled for
23/// `longcat-flash-lite`) and `gb10/step3p7-flash` first clamp `gate` to at most 10 and `up`
24/// to [-10, 10].
25pub fn silu_mul(
26    gpu: &dyn GpuBackend,
27    kernel: KernelHandle,
28    gate: DevicePtr,
29    up: DevicePtr,
30    output: DevicePtr,
31    num_elements: u32,
32    stream: u64,
33) -> Result<()> {
34    KernelLaunch::new(gpu, kernel)
35        .grid([div_ceil(num_elements, 256), 1, 1])
36        .block([256, 1, 1])
37        .arg_ptr(gate)
38        .arg_ptr(up)
39        .arg_ptr(output)
40        .arg_u32(num_elements)
41        .launch(stream)
42}
43
44/// 2026-09-25: `SiLU(gate) * up` over row-strided BF16 operands (kernel
45/// `silu_mul_strided`): row `r` reads `gate` and `up` at `r * in_stride + c` for
46/// `c < cols` and writes `output` at `r * out_stride + c`. The fused dense-FFN gate+up
47/// projection (`layers/dense_ffn_gateup_fused.rs`) passes the two column halves of one
48/// `[rows, 2 * cols]` matrix (`up = gate.offset(cols * 2)`) and a contiguous output.
49/// One grid row per matrix row.
50#[allow(clippy::too_many_arguments)]
51pub fn silu_mul_strided(
52    gpu: &dyn GpuBackend,
53    kernel: KernelHandle,
54    gate: DevicePtr,
55    up: DevicePtr,
56    output: DevicePtr,
57    rows: u32,
58    cols: u32,
59    in_stride: u32,
60    out_stride: u32,
61    stream: u64,
62) -> Result<()> {
63    KernelLaunch::new(gpu, kernel)
64        .grid([div_ceil(cols, 256), rows, 1])
65        .block([256, 1, 1])
66        .arg_ptr(gate)
67        .arg_ptr(up)
68        .arg_ptr(output)
69        .arg_u32(rows)
70        .arg_u32(cols)
71        .arg_u32(in_stride)
72        .arg_u32(out_stride)
73        .launch(stream)
74}
75
76/// 2026-09-25: `SiLU(gate) * up` quantized to FP8 E4M3 with one scale per 128-element
77/// group of each row, in one kernel (`silu_mul_quant_fp8`); the FP8 MoE prefill uses it
78/// before the down projection (`moe/forward_prefill_fp8.rs`).
79///
80/// `out_bf16` may be `DevicePtr::NULL`; when it is not, the kernel also writes the
81/// post-SiLU BF16 values there. The caller must ensure `k % 128 == 0` and
82/// `k / 128 <= 16` (the kernel's `SILU_QUANT_MAX_GROUPS`).
83#[allow(clippy::too_many_arguments)]
84pub fn silu_mul_quant_fp8(
85    gpu: &dyn GpuBackend,
86    kernel: KernelHandle,
87    gate: DevicePtr,
88    up: DevicePtr,
89    out_fp8: DevicePtr,
90    a_scale: DevicePtr,
91    out_bf16: DevicePtr,
92    m: u32,
93    k: u32,
94    stream: u64,
95) -> Result<()> {
96    KernelLaunch::new(gpu, kernel)
97        .grid([m, 1, 1])
98        .block([128, 1, 1])
99        .arg_ptr(gate)
100        .arg_ptr(up)
101        .arg_ptr(out_fp8)
102        .arg_ptr(a_scale)
103        .arg_ptr(out_bf16)
104        .arg_u32(m)
105        .arg_u32(k)
106        .launch(stream)
107}
108
109/// 2026-09-25: In-place L2 normalization of each head, `x / sqrt(sum(x^2) + eps)` (kernel
110/// `l2_norm_bf16`). `data` is `[num_tokens, stride]` BF16 with `num_heads * head_dim`
111/// used per token; one block per (head, token).
112pub fn l2_norm(
113    gpu: &dyn GpuBackend,
114    kernel: KernelHandle,
115    data: DevicePtr,
116    num_heads: u32,
117    head_dim: u32,
118    eps: f32,
119    num_tokens: u32,
120    stride: u32,
121    stream: u64,
122) -> Result<()> {
123    KernelLaunch::new(gpu, kernel)
124        .grid([num_heads, num_tokens, 1])
125        .block([head_dim.min(1024), 1, 1])
126        .arg_ptr(data)
127        .arg_u32(head_dim)
128        .arg_f32(eps)
129        .arg_u32(stride)
130        .launch(stream)
131}
132
133/// 2026-09-25: `output[i] = input[i] * sigmoid(gate[i])` (kernel `sigmoid_gate_mul`).
134pub fn sigmoid_gate_mul(
135    gpu: &dyn GpuBackend,
136    kernel: KernelHandle,
137    input: DevicePtr,
138    gate: DevicePtr,
139    output: DevicePtr,
140    num_elements: u32,
141    stream: u64,
142) -> Result<()> {
143    KernelLaunch::new(gpu, kernel)
144        .grid([div_ceil(num_elements, 256), 1, 1])
145        .block([256, 1, 1])
146        .arg_ptr(input)
147        .arg_ptr(gate)
148        .arg_ptr(output)
149        .arg_u32(num_elements)
150        .launch(stream)
151}
152
153/// 2026-09-25: `output[t, h, d] = input[t, h, d] * sigmoid(gate[t, h])`: one BF16 gate
154/// value per head, `[num_tokens, nq]`, broadcast over the head's `hd` elements (kernel
155/// `sigmoid_gate_mul_head_broadcast`).
156pub fn sigmoid_gate_mul_head_broadcast(
157    gpu: &dyn GpuBackend,
158    kernel: KernelHandle,
159    input: DevicePtr,
160    gate: DevicePtr,
161    output: DevicePtr,
162    nq: u32,
163    hd: u32,
164    num_tokens: u32,
165    stream: u64,
166) -> Result<()> {
167    let total = num_tokens * nq * hd;
168    KernelLaunch::new(gpu, kernel)
169        .grid([div_ceil(total, 256), 1, 1])
170        .block([256, 1, 1])
171        .arg_ptr(input)
172        .arg_ptr(gate)
173        .arg_ptr(output)
174        .arg_u32(nq)
175        .arg_u32(hd)
176        .arg_u32(total)
177        .launch(stream)
178}
179
180/// 2026-09-25: Like [`sigmoid_gate_mul_head_broadcast`] with softplus in place of sigmoid
181/// (kernel `softplus_gate_mul_head_broadcast`).
182#[allow(clippy::too_many_arguments)]
183pub fn softplus_gate_mul_head_broadcast(
184    gpu: &dyn GpuBackend,
185    kernel: KernelHandle,
186    input: DevicePtr,
187    gate: DevicePtr,
188    output: DevicePtr,
189    nq: u32,
190    hd: u32,
191    num_tokens: u32,
192    stream: u64,
193) -> Result<()> {
194    let total = num_tokens * nq * hd;
195    KernelLaunch::new(gpu, kernel)
196        .grid([div_ceil(total, 256), 1, 1])
197        .block([256, 1, 1])
198        .arg_ptr(input)
199        .arg_ptr(gate)
200        .arg_ptr(output)
201        .arg_u32(nq)
202        .arg_u32(hd)
203        .arg_u32(total)
204        .launch(stream)
205}
206
207/// 2026-09-25: In-place BF16 `residual[i] += src[i]` (kernel `bf16_residual_add`).
208pub fn residual_add(
209    gpu: &dyn GpuBackend,
210    kernel: KernelHandle,
211    residual: DevicePtr,
212    src: DevicePtr,
213    num_elements: u32,
214    stream: u64,
215) -> Result<()> {
216    KernelLaunch::new(gpu, kernel)
217        .grid([div_ceil(num_elements, 256), 1, 1])
218        .block([256, 1, 1])
219        .arg_ptr(residual)
220        .arg_ptr(src)
221        .arg_u32(num_elements)
222        .launch(stream)
223}
224
225/// 2026-09-25: In-place BF16 `output[i] += scale * src[i]` (kernel `bf16_scaled_add`).
226pub fn scaled_add(
227    gpu: &dyn GpuBackend,
228    kernel: KernelHandle,
229    output: DevicePtr,
230    src: DevicePtr,
231    scale: f32,
232    num_elements: u32,
233    stream: u64,
234) -> Result<()> {
235    KernelLaunch::new(gpu, kernel)
236        .grid([div_ceil(num_elements, 256), 1, 1])
237        .block([256, 1, 1])
238        .arg_ptr(output)
239        .arg_ptr(src)
240        .arg_f32(scale)
241        .arg_u32(num_elements)
242        .launch(stream)
243}
244
245/// 2026-09-25: In-place BF16 `output[i] += sigmoid_gate * src[i]`, where `sigmoid_gate`
246/// is a host value the caller has already passed through the sigmoid. The arguments are
247/// those of kernel `bf16_sigmoid_blend`; nothing in the tree calls this function.
248pub fn sigmoid_blend(
249    gpu: &dyn GpuBackend,
250    kernel: KernelHandle,
251    output: DevicePtr,
252    src: DevicePtr,
253    sigmoid_gate: f32,
254    num_elements: u32,
255    stream: u64,
256) -> Result<()> {
257    KernelLaunch::new(gpu, kernel)
258        .grid([div_ceil(num_elements, 256), 1, 1])
259        .block([256, 1, 1])
260        .arg_ptr(output)
261        .arg_ptr(src)
262        .arg_f32(sigmoid_gate)
263        .arg_u32(num_elements)
264        .launch(stream)
265}
metrale_model_layers/layers/ops/activations.rs

metrale_model_layers/layers/ops/
activations.rs
