spark_model/layers/ops/
norm.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// ── Normalization ──────────────────────────────────────────────────
17
18/// RMS normalization: output = rms_norm(input) * weight.
19///
20/// Kernel: `rms_norm(input, weight, output, hidden_size, eps)`
21/// Grid: (num_tokens, 1, 1)  Block: (min(hidden_size, 1024), 1, 1)
22/// Strided RMS norm: `num_groups` groups of `rows_per_group` rows in ONE launch,
23/// groups `row_stride` ELEMENTS apart, rows packed at `hidden_size` inside a group.
24///
25/// `rms_norm` above assumes one packed [num_tokens, hidden_size] block. The
26/// multi-seq q/k head-norms are packed only WITHIN a sequence — each sequence's
27/// heads sit inside its own interleaved [Q|K|V|gate] block — so that path was
28/// launching the packed kernel once per sequence (516 launches/step, 0.76 ms).
29/// Bit-identical: one block per row either way, same math, same reduction.
30#[allow(clippy::too_many_arguments)]
31pub fn rms_norm_strided(
32    gpu: &dyn GpuBackend,
33    kernel: KernelHandle,
34    input: DevicePtr,
35    weight: &DenseWeight,
36    output: DevicePtr,
37    rows_per_group: u32,
38    num_groups: u32,
39    hidden_size: u32,
40    eps: f32,
41    row_stride: u32,
42    stream: u64,
43) -> Result<()> {
44    KernelLaunch::new(gpu, kernel)
45        .grid([rows_per_group, num_groups, 1])
46        .block([hidden_size.min(1024), 1, 1])
47        .arg_ptr(input)
48        .arg_ptr(weight.weight)
49        .arg_ptr(output)
50        .arg_u32(hidden_size)
51        .arg_f32(eps)
52        .arg_u32(row_stride)
53        .launch(stream)
54}
55
56pub fn rms_norm(
57    gpu: &dyn GpuBackend,
58    kernel: KernelHandle,
59    input: DevicePtr,
60    weight: &DenseWeight,
61    output: DevicePtr,
62    num_tokens: u32,
63    hidden_size: u32,
64    eps: f32,
65    stream: u64,
66) -> Result<()> {
67    KernelLaunch::new(gpu, kernel)
68        .grid([num_tokens, 1, 1])
69        .block([hidden_size.min(1024), 1, 1])
70        .arg_ptr(input)
71        .arg_ptr(weight.weight)
72        .arg_ptr(output)
73        .arg_u32(hidden_size)
74        .arg_f32(eps)
75        .launch(stream)
76}
77
78/// Warp-per-row RMS norm for SHORT rows — one warp per row instead of one
79/// block, so the grid shrinks 8x and the reduction needs no shared memory or
80/// barrier. Profitable exactly for the Qwen3 per-head `q_norm`/`k_norm` during
81/// prefill (`num_rows = heads * seq`, `hidden_size = head_dim`), where the
82/// block-per-row kernel measured ~43x above its bandwidth floor.
83pub fn rms_norm_warp_row(
84    gpu: &dyn GpuBackend,
85    kernel: KernelHandle,
86    input: DevicePtr,
87    weight: &DenseWeight,
88    output: DevicePtr,
89    num_rows: u32,
90    hidden_size: u32,
91    eps: f32,
92    stream: u64,
93) -> Result<()> {
94    const ROWS_PER_BLOCK: u32 = 8;
95    KernelLaunch::new(gpu, kernel)
96        .grid([num_rows.div_ceil(ROWS_PER_BLOCK), 1, 1])
97        .block([32 * ROWS_PER_BLOCK, 1, 1])
98        .arg_ptr(input)
99        .arg_ptr(weight.weight)
100        .arg_ptr(output)
101        .arg_u32(num_rows)
102        .arg_u32(hidden_size)
103        .arg_f32(eps)
104        .launch(stream)
105}
106
107/// Gate for [`rms_norm_warp_row`]: short even rows, many of them.
108/// Disable with `ATLAS_RMS_NORM_WARP_ROW=0`.
109pub fn rms_norm_short_row_eligible(num_rows: u32, hidden_size: u32) -> bool {
110    use std::sync::OnceLock;
111    static ON: OnceLock<bool> = OnceLock::new();
112    let on = *ON.get_or_init(|| std::env::var("ATLAS_RMS_NORM_WARP_ROW").as_deref() != Ok("0"));
113    on && hidden_size <= 256 && hidden_size.is_multiple_of(2) && num_rows >= 1024
114}
115
116/// Fused RMS norm + residual save: normed = rms_norm(input), residual = input.
117///
118/// Eliminates a separate D2D copy by writing the raw input to the residual
119/// buffer in the same pass as the normalized output write.
120///
121/// Kernel: `rms_norm_residual(input, weight, output, residual, hidden_size, eps)`
122/// Grid: (num_tokens, 1, 1)  Block: (min(hidden_size, 1024), 1, 1)
123pub fn rms_norm_residual(
124    gpu: &dyn GpuBackend,
125    kernel: KernelHandle,
126    input: DevicePtr,
127    weight: &DenseWeight,
128    output: DevicePtr,
129    residual: DevicePtr,
130    num_tokens: u32,
131    hidden_size: u32,
132    eps: f32,
133    stream: u64,
134) -> Result<()> {
135    KernelLaunch::new(gpu, kernel)
136        .grid([num_tokens, 1, 1])
137        .block([hidden_size.min(1024), 1, 1])
138        .arg_ptr(input)
139        .arg_ptr(weight.weight)
140        .arg_ptr(output)
141        .arg_ptr(residual)
142        .arg_u32(hidden_size)
143        .arg_f32(eps)
144        .launch(stream)
145}
146
147/// Fused residual add + RMS norm + residual save.
148///
149/// `hidden[i] += src[i]; normed = rms_norm(hidden) * (1+weight); residual = hidden`.
150/// Eliminates one kernel launch per fusion site (48 per decode step).
151///
152/// Kernel: `residual_add_rms_norm(hidden, src, weight, output, residual, hidden_size, eps)`
153/// Grid: (num_tokens, 1, 1)  Block: (min(hidden_size, 1024), 1, 1)
154#[allow(clippy::too_many_arguments)]
155pub fn residual_add_rms_norm(
156    gpu: &dyn GpuBackend,
157    kernel: KernelHandle,
158    hidden: DevicePtr,
159    src: DevicePtr,
160    weight: &DenseWeight,
161    output: DevicePtr,
162    residual: DevicePtr,
163    num_tokens: u32,
164    hidden_size: u32,
165    eps: f32,
166    stream: u64,
167) -> Result<()> {
168    KernelLaunch::new(gpu, kernel)
169        .grid([num_tokens, 1, 1])
170        .block([hidden_size.min(1024), 1, 1])
171        .arg_ptr(hidden)
172        .arg_ptr(src)
173        .arg_ptr(weight.weight)
174        .arg_ptr(output)
175        .arg_ptr(residual)
176        .arg_u32(hidden_size)
177        .arg_f32(eps)
178        .launch(stream)
179}
180
181/// Dual-output fused residual add + RMS norm (ATLAS_FP32_ROUTING).
182///
183/// Same as `residual_add_rms_norm` (bf16 hidden/residual/output unchanged) but
184/// ALSO writes the normed output in FP32 to `output_f32` for the MoE router GEMM,
185/// removing the norm's bf16-store rounding from the routing-critical path.
186///
187/// Kernel: `residual_add_rms_norm_gatef32(hidden, src, weight, output,
188///          output_f32, residual, hidden_size, eps)`
189/// Grid: (num_tokens, 1, 1)  Block: (min(hidden_size, 1024), 1, 1)
190#[allow(clippy::too_many_arguments)]
191pub fn residual_add_rms_norm_gatef32(
192    gpu: &dyn GpuBackend,
193    kernel: KernelHandle,
194    hidden: DevicePtr,
195    src: DevicePtr,
196    weight: &DenseWeight,
197    output: DevicePtr,
198    output_f32: DevicePtr,
199    residual: DevicePtr,
200    num_tokens: u32,
201    hidden_size: u32,
202    eps: f32,
203    stream: u64,
204) -> Result<()> {
205    KernelLaunch::new(gpu, kernel)
206        .grid([num_tokens, 1, 1])
207        .block([hidden_size.min(1024), 1, 1])
208        .arg_ptr(hidden)
209        .arg_ptr(src)
210        .arg_ptr(weight.weight)
211        .arg_ptr(output)
212        .arg_ptr(output_f32)
213        .arg_ptr(residual)
214        .arg_u32(hidden_size)
215        .arg_f32(eps)
216        .launch(stream)
217}
218
219/// Gated RMS norm (norm_before_gate=False, per-group):
220///   output = rms_norm_per_group(input * silu(gate), weight, group_size)
221///
222/// Kernel: `gated_rms_norm(input, gate, weight, output, hidden_size, eps, gate_stride, group_size)`
223/// Grid: (num_tokens, 1, 1)  Block: (min(hidden_size, 1024), 1, 1)
224pub fn gated_rms_norm(
225    gpu: &dyn GpuBackend,
226    kernel: KernelHandle,
227    input: DevicePtr,
228    gate: DevicePtr,
229    weight: &DenseWeight,
230    output: DevicePtr,
231    num_tokens: u32,
232    hidden_size: u32,
233    gate_stride: u32,
234    eps: f32,
235    group_size: u32,
236    stream: u64,
237) -> Result<()> {
238    KernelLaunch::new(gpu, kernel)
239        .grid([num_tokens, 1, 1])
240        .block([hidden_size.min(1024), 1, 1])
241        .arg_ptr(input)
242        .arg_ptr(gate)
243        .arg_ptr(weight.weight)
244        .arg_ptr(output)
245        .arg_u32(hidden_size)
246        .arg_f32(eps)
247        .arg_u32(gate_stride)
248        .arg_u32(group_size)
249        .launch(stream)
250}
251
252/// Batched gated RMS norm for prefill: all (head, actual_token) pairs in one launch.
253///
254/// Grid: (heads_per_token, num_actual_tokens, 1)
255/// Block: (min(head_dim, 1024), 1, 1)
256#[allow(clippy::too_many_arguments)]
257pub fn gated_rms_norm_prefill(
258    gpu: &dyn GpuBackend,
259    kernel: KernelHandle,
260    input: DevicePtr,
261    gate: DevicePtr,
262    weight: &DenseWeight,
263    output: DevicePtr,
264    heads_per_token: u32,
265    head_dim: u32,
266    eps: f32,
267    num_actual_tokens: u32,
268    input_token_stride: u32,
269    gate_token_stride: u32,
270    stream: u64,
271) -> Result<()> {
272    KernelLaunch::new(gpu, kernel)
273        .grid([heads_per_token, num_actual_tokens, 1])
274        .block([head_dim.min(1024), 1, 1])
275        .arg_ptr(input)
276        .arg_ptr(gate)
277        .arg_ptr(weight.weight)
278        .arg_ptr(output)
279        .arg_u32(head_dim)
280        .arg_f32(eps)
281        .arg_u32(input_token_stride)
282        .arg_u32(gate_token_stride)
283        .launch(stream)
284}
285
286// ── GEMM ───────────────────────────────────────────────────────────
spark_model/layers/ops/norm.rs

spark_model/layers/ops/
norm.rs
