metrale_model_layers/layers/ops/
ssm_mamba_gdn_verify.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the fused GDN verify kernels: conv1d + L2 norm over
4//! K=2 or K=n draft positions (one sequence or a batch), and the K=2 gated RMS norm.
5//!
6//! Owner: model-layers ops.
7//! Invariants: none beyond the types.
8
9use super::*;
10
11/// 2026-09-25: K=2 MTP-verify conv1d + SiLU + L2 norm: both draft positions in
12/// one launch, with the position-0 conv-state snapshot written inline.
13///
14/// On return `conv_state` holds the committed (post position-1) window and
15/// `conv_state_inter` the position-0 rollback snapshot. The
16/// `gdn_verify_fused_microtest` example checks both states, and the gated-norm
17/// output computed from this kernel's output, against the per-token kernels at
18/// cos >= 0.99999.
19///
20/// Kernel: `gdn_verify_fused_conv_k2(conv_state, new_input, weight, output,
21///          conv_state_inter, dim, d_conv, qk_channels, head_dim,
22///          input_stride, output_stride, l2_eps)`
23/// Grid: (ceil(dim/256), 1, 1)  Block: (256, 1, 1)
24#[allow(clippy::too_many_arguments)]
25pub fn gdn_verify_fused_conv_k2(
26    gpu: &dyn GpuBackend,
27    kernel: KernelHandle,
28    conv_state: DevicePtr,
29    new_input: DevicePtr,
30    weight: &DenseWeight,
31    output: DevicePtr,
32    conv_state_inter: DevicePtr,
33    d_inner: u32,
34    d_conv: u32,
35    qk_channels: u32,
36    head_dim: u32,
37    input_stride: u32,
38    output_stride: u32,
39    l2_eps: f32,
40    stream: u64,
41) -> Result<()> {
42    KernelLaunch::new(gpu, kernel)
43        .grid([div_ceil(d_inner, 256), 1, 1])
44        .block([256, 1, 1])
45        .arg_ptr(conv_state)
46        .arg_ptr(new_input)
47        .arg_ptr(weight.weight)
48        .arg_ptr(output)
49        .arg_ptr(conv_state_inter)
50        .arg_u32(d_inner)
51        .arg_u32(d_conv)
52        .arg_u32(qk_channels)
53        .arg_u32(head_dim)
54        .arg_u32(input_stride)
55        .arg_u32(output_stride)
56        .arg_f32(l2_eps)
57        .launch(stream)
58}
59
60/// 2026-09-25: [`gdn_verify_fused_conv_kn`] for `n_seq` sequences in one launch,
61/// one grid row (`blockIdx.y`) per sequence.
62///
63/// The kernel offsets `conv_state`, `new_input`, `output` and
64/// `conv_state_inter` by `seq` times their `*_seq_stride`, then runs the same
65/// per-sequence body as `gdn_verify_fused_conv_kn`, so each sequence gets the
66/// result a separate call would give it.
67///
68/// Grid: (ceil(d_inner/256), n_seq, 1)  Block: (256, 1, 1)
69#[allow(clippy::too_many_arguments)]
70pub fn gdn_verify_fused_conv_kn_batched(
71    gpu: &dyn GpuBackend,
72    kernel: KernelHandle,
73    conv_state: DevicePtr,
74    new_input: DevicePtr,
75    weight: &DenseWeight,
76    output: DevicePtr,
77    conv_state_inter: DevicePtr,
78    num_tokens: u32,
79    d_inner: u32,
80    d_conv: u32,
81    qk_channels: u32,
82    head_dim: u32,
83    input_stride: u32,
84    output_stride: u32,
85    inter_stride: u32,
86    l2_eps: f32,
87    n_seq: u32,
88    conv_state_seq_stride: u32,
89    input_seq_stride: u32,
90    output_seq_stride: u32,
91    inter_seq_stride: u32,
92    stream: u64,
93) -> Result<()> {
94    KernelLaunch::new(gpu, kernel)
95        .grid([div_ceil(d_inner, 256), n_seq, 1])
96        .block([256, 1, 1])
97        .arg_ptr(conv_state)
98        .arg_ptr(new_input)
99        .arg_ptr(weight.weight)
100        .arg_ptr(output)
101        .arg_ptr(conv_state_inter)
102        .arg_u32(num_tokens)
103        .arg_u32(d_inner)
104        .arg_u32(d_conv)
105        .arg_u32(qk_channels)
106        .arg_u32(head_dim)
107        .arg_u32(input_stride)
108        .arg_u32(output_stride)
109        .arg_u32(inter_stride)
110        .arg_f32(l2_eps)
111        .arg_u32(conv_state_seq_stride)
112        .arg_u32(input_seq_stride)
113        .arg_u32(output_seq_stride)
114        .arg_u32(inter_seq_stride)
115        .launch(stream)
116}
117
118/// 2026-09-25: K-position verify conv1d + SiLU + L2 norm: all `num_tokens`
119/// draft positions in one launch. Position `t`'s conv-state snapshot is
120/// written to `conv_state_inter + t * inter_stride`.
121///
122/// On return `conv_state` holds the window after the last position, which is
123/// also snapshot `num_tokens - 1`.
124///
125/// Kernel: `gdn_verify_fused_conv_kn(conv_state, new_input, weight, output,
126///          conv_state_inter, num_tokens, dim, d_conv, qk_channels, head_dim,
127///          input_stride, output_stride, inter_stride, l2_eps)`
128/// Grid: (ceil(dim/256), 1, 1)  Block: (256, 1, 1)
129#[allow(clippy::too_many_arguments)]
130pub fn gdn_verify_fused_conv_kn(
131    gpu: &dyn GpuBackend,
132    kernel: KernelHandle,
133    conv_state: DevicePtr,
134    new_input: DevicePtr,
135    weight: &DenseWeight,
136    output: DevicePtr,
137    conv_state_inter: DevicePtr,
138    num_tokens: u32,
139    d_inner: u32,
140    d_conv: u32,
141    qk_channels: u32,
142    head_dim: u32,
143    input_stride: u32,
144    output_stride: u32,
145    inter_stride: u32,
146    l2_eps: f32,
147    stream: u64,
148) -> Result<()> {
149    KernelLaunch::new(gpu, kernel)
150        .grid([div_ceil(d_inner, 256), 1, 1])
151        .block([256, 1, 1])
152        .arg_ptr(conv_state)
153        .arg_ptr(new_input)
154        .arg_ptr(weight.weight)
155        .arg_ptr(output)
156        .arg_ptr(conv_state_inter)
157        .arg_u32(num_tokens)
158        .arg_u32(d_inner)
159        .arg_u32(d_conv)
160        .arg_u32(qk_channels)
161        .arg_u32(head_dim)
162        .arg_u32(input_stride)
163        .arg_u32(output_stride)
164        .arg_u32(inter_stride)
165        .arg_f32(l2_eps)
166        .launch(stream)
167}
168
169/// 2026-09-25: K=2 MTP-verify gated RMS norm: both draft positions in one
170/// launch. The Z gate is read from the deinterleaved [Q|K|V|Z] buffer at
171/// `z_offset` within each position. The `gdn_verify_fused_microtest` example
172/// checks it against the per-token `gated_rms_norm` at cos >= 0.99999.
173///
174/// Kernel: `gdn_verify_fused_norm_k2(gdn_out, deint, weight, output,
175///          hidden_size, eps, deint_stride, z_offset, out_stride)`
176/// Grid: (num_v_heads, 2, 1)  Block: (hidden_size, 1, 1)
177#[allow(clippy::too_many_arguments)]
178pub fn gdn_verify_fused_norm_k2(
179    gpu: &dyn GpuBackend,
180    kernel: KernelHandle,
181    gdn_out: DevicePtr,
182    deint: DevicePtr,
183    weight: &DenseWeight,
184    output: DevicePtr,
185    num_v_heads: u32,
186    hidden_size: u32,
187    eps: f32,
188    deint_stride: u32,
189    z_offset: u32,
190    out_stride: u32,
191    stream: u64,
192) -> Result<()> {
193    KernelLaunch::new(gpu, kernel)
194        .grid([num_v_heads, 2, 1])
195        .block([hidden_size, 1, 1])
196        .arg_ptr(gdn_out)
197        .arg_ptr(deint)
198        .arg_ptr(weight.weight)
199        .arg_ptr(output)
200        .arg_u32(hidden_size)
201        .arg_f32(eps)
202        .arg_u32(deint_stride)
203        .arg_u32(z_offset)
204        .arg_u32(out_stride)
205        .launch(stream)
206}
metrale_model_layers/layers/ops/ssm_mamba_gdn_verify.rs

metrale_model_layers/layers/ops/
ssm_mamba_gdn_verify.rs
