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/// Strided gated RMS norm for MULTI-SEQ DECODE: all `(head, sequence)` pairs
253/// in ONE launch instead of one launch per sequence.
254///
255/// WHY (#927). The H100 nsys trace of a batch-16 decode step
256/// (2026-09-11 round 7, `Qwen/Qwen3.8-27B-FP8`, step 43.595 ms) showed
257/// `gated_rms_norm_f32_input` firing **768** times — 48 SSM layers x 16
258/// sequences — for 1.612 ms, i.e. 2.1 us each. At that size it is pure
259/// launch/tail overhead, not work, and it was the ONLY per-layer kernel in the
260/// step still scaling with the row count: `gated_delta_rule_decode_f32_strided`
261/// and `causal_conv1d_update_l2norm_f32_strided` next to it are already at 48.
262/// This entry point makes it 48 too, recovering most of 3.70% of the step.
263///
264/// BIT-IDENTICAL to `gated_rms_norm` at the same addresses: one block per
265/// `(sequence, head)` row either way, same reduction over the same elements in
266/// the same order. Only the base address differs, so no cross-row interaction
267/// is introduced — the same argument `rms_norm_strided` makes.
268///
269/// Strides are in ELEMENTS of each buffer's own type: `input_seq_stride` in
270/// f32, `gate_seq_stride`/`output_seq_stride` in BF16.
271///
272/// Grid: (heads_per_seq, num_seqs, 1)   Block: (min(hidden_size, 1024), 1, 1)
273#[allow(clippy::too_many_arguments)]
274pub fn gated_rms_norm_strided(
275    gpu: &dyn GpuBackend,
276    kernel: KernelHandle,
277    input: DevicePtr,
278    gate: DevicePtr,
279    weight: &DenseWeight,
280    output: DevicePtr,
281    heads_per_seq: u32,
282    num_seqs: u32,
283    hidden_size: u32,
284    gate_stride: u32,
285    eps: f32,
286    group_size: u32,
287    input_seq_stride: u32,
288    gate_seq_stride: u32,
289    output_seq_stride: u32,
290    stream: u64,
291) -> Result<()> {
292    KernelLaunch::new(gpu, kernel)
293        .grid([heads_per_seq, num_seqs, 1])
294        .block([hidden_size.min(1024), 1, 1])
295        .arg_ptr(input)
296        .arg_ptr(gate)
297        .arg_ptr(weight.weight)
298        .arg_ptr(output)
299        .arg_u32(hidden_size)
300        .arg_f32(eps)
301        .arg_u32(gate_stride)
302        .arg_u32(group_size)
303        .arg_u32(input_seq_stride)
304        .arg_u32(gate_seq_stride)
305        .arg_u32(output_seq_stride)
306        .launch(stream)
307}
308
309/// Batched gated RMS norm for prefill: all (head, actual_token) pairs in one launch.
310///
311/// Grid: (heads_per_token, num_actual_tokens, 1)
312/// Block: (min(head_dim, 1024), 1, 1)
313#[allow(clippy::too_many_arguments)]
314pub fn gated_rms_norm_prefill(
315    gpu: &dyn GpuBackend,
316    kernel: KernelHandle,
317    input: DevicePtr,
318    gate: DevicePtr,
319    weight: &DenseWeight,
320    output: DevicePtr,
321    heads_per_token: u32,
322    head_dim: u32,
323    eps: f32,
324    num_actual_tokens: u32,
325    input_token_stride: u32,
326    gate_token_stride: u32,
327    stream: u64,
328) -> Result<()> {
329    KernelLaunch::new(gpu, kernel)
330        .grid([heads_per_token, num_actual_tokens, 1])
331        .block([head_dim.min(1024), 1, 1])
332        .arg_ptr(input)
333        .arg_ptr(gate)
334        .arg_ptr(weight.weight)
335        .arg_ptr(output)
336        .arg_u32(head_dim)
337        .arg_f32(eps)
338        .arg_u32(input_token_stride)
339        .arg_u32(gate_token_stride)
340        .launch(stream)
341}
342
343// ── GEMM ───────────────────────────────────────────────────────────
spark_model/layers/ops/norm.rs

spark_model/layers/ops/
norm.rs
