spark_model/layers/ops/
ssm_gdn_snap.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Launch wrappers for the exact-verify `_snap` kernel twins (issue #435
4//! route (a)): the fused-norm GDN decode kernels with an inline per-token
5//! h-state rollback snapshot, and the FP32-output fused verify conv.
6//!
7//! All three are OPTIONAL kernels (`try_kernel`, model-shadow staged): the
8//! exact-verify arm falls back to the parent kernel + `copy_d2d_async`
9//! snapshots when a handle is 0 — same bits, more launches.
10
11use anyhow::Result;
12use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
14
15use crate::weight_map::DenseWeight;
16
17/// [`super::gdn_decode_f32_norm`] + inline h-state snapshot.
18///
19/// `h_inter` receives the post-update (post-state-norm-clamp) H — the same
20/// bits left in `h_state` — or is skipped when NULL (the final verify
21/// position, whose snapshot index has no reader). Same grid/block and
22/// argument order as the parent, with `h_inter` appended.
23#[allow(clippy::too_many_arguments)]
24pub fn gdn_decode_f32_norm_snap(
25    gpu: &dyn GpuBackend,
26    kernel: KernelHandle,
27    h_state: DevicePtr,
28    query: DevicePtr,
29    key: DevicePtr,
30    value: DevicePtr,
31    gate: DevicePtr,
32    beta: DevicePtr,
33    z_gate: DevicePtr,
34    norm_weight: DevicePtr,
35    output: DevicePtr,
36    h_inter: DevicePtr,
37    batch_size: u32,
38    num_k_heads: u32,
39    num_v_heads: u32,
40    k_dim: u32,
41    v_dim: u32,
42    eps: f32,
43    stream: u64,
44) -> Result<()> {
45    KernelLaunch::new(gpu, kernel)
46        .grid([num_v_heads, batch_size, 1])
47        .block([128, 1, 1])
48        .arg_ptr(h_state)
49        .arg_ptr(query)
50        .arg_ptr(key)
51        .arg_ptr(value)
52        .arg_ptr(gate)
53        .arg_ptr(beta)
54        .arg_ptr(z_gate)
55        .arg_ptr(norm_weight)
56        .arg_ptr(output)
57        .arg_u32(batch_size)
58        .arg_u32(num_k_heads)
59        .arg_u32(num_v_heads)
60        .arg_u32(k_dim)
61        .arg_u32(v_dim)
62        .arg_f32(eps)
63        .arg_ptr(h_inter)
64        .launch(stream)
65}
66
67/// [`super::gdn_decode_f32_strided_norm`] + inline h-state snapshot, for the
68/// batched-verify arm at `batch_size = n` sequences.
69///
70/// `h_inter` is the snapshot base for THIS token position; sequences are
71/// `h_inter_seq_stride` FP32 elements apart (the ssm-pool per-slot
72/// intermediate stride — passed, not inferred, because pool slots are
73/// `num_intermediates` snapshots wide while H itself is dense). NULL skips.
74#[allow(clippy::too_many_arguments)]
75pub fn gdn_decode_f32_strided_norm_snap(
76    gpu: &dyn GpuBackend,
77    kernel: KernelHandle,
78    h_state: DevicePtr,
79    query: DevicePtr,
80    key: DevicePtr,
81    value: DevicePtr,
82    gate: DevicePtr,
83    beta: DevicePtr,
84    z_gate: DevicePtr,
85    norm_weight: DevicePtr,
86    output: DevicePtr,
87    h_inter: DevicePtr,
88    h_inter_seq_stride: u64,
89    batch_size: u32,
90    num_k_heads: u32,
91    num_v_heads: u32,
92    k_dim: u32,
93    v_dim: u32,
94    qk_stride: u32,
95    v_stride: u32,
96    gb_stride: u32,
97    z_stride: u32,
98    out_stride: u32,
99    eps: f32,
100    stream: u64,
101) -> Result<()> {
102    KernelLaunch::new(gpu, kernel)
103        .grid([num_v_heads, batch_size, 1])
104        .block([128, 1, 1])
105        .arg_ptr(h_state)
106        .arg_ptr(query)
107        .arg_ptr(key)
108        .arg_ptr(value)
109        .arg_ptr(gate)
110        .arg_ptr(beta)
111        .arg_ptr(z_gate)
112        .arg_ptr(norm_weight)
113        .arg_ptr(output)
114        .arg_u32(batch_size)
115        .arg_u32(num_k_heads)
116        .arg_u32(num_v_heads)
117        .arg_u32(k_dim)
118        .arg_u32(v_dim)
119        .arg_u32(qk_stride)
120        .arg_u32(v_stride)
121        .arg_u32(gb_stride)
122        .arg_u32(z_stride)
123        .arg_u32(out_stride)
124        .arg_f32(eps)
125        .arg_ptr(h_inter)
126        .arg_u64(h_inter_seq_stride)
127        .launch(stream)
128}
129
130/// FP32-output twin of [`super::gdn_verify_fused_conv_kn`]: one launch for
131/// all K verify positions of conv1d+SiLU+L2norm, FP32 conv rows (what the
132/// sequential-decode-exact GDN chain reads), every per-token conv-state
133/// rollback snapshot written inline. `output_stride` is in FP32 elements.
134///
135/// Kernel: `gdn_verify_fused_conv_kn_f32(conv_state, new_input, weight,
136///          output, conv_state_inter, num_tokens, dim, d_conv, qk_channels,
137///          head_dim, input_stride, output_stride, inter_stride, l2_eps)`
138/// Grid: (ceil(dim/256), 1, 1)  Block: (256, 1, 1)
139#[allow(clippy::too_many_arguments)]
140pub fn gdn_verify_fused_conv_kn_f32(
141    gpu: &dyn GpuBackend,
142    kernel: KernelHandle,
143    conv_state: DevicePtr,
144    new_input: DevicePtr,
145    weight: &DenseWeight,
146    output: DevicePtr,
147    conv_state_inter: DevicePtr,
148    num_tokens: u32,
149    dim: u32,
150    d_conv: u32,
151    qk_channels: u32,
152    head_dim: u32,
153    input_stride: u32,
154    output_stride: u32,
155    inter_stride: u32,
156    l2_eps: f32,
157    stream: u64,
158) -> Result<()> {
159    KernelLaunch::new(gpu, kernel)
160        .grid([div_ceil(dim, 256), 1, 1])
161        .block([256, 1, 1])
162        .arg_ptr(conv_state)
163        .arg_ptr(new_input)
164        .arg_ptr(weight.weight)
165        .arg_ptr(output)
166        .arg_ptr(conv_state_inter)
167        .arg_u32(num_tokens)
168        .arg_u32(dim)
169        .arg_u32(d_conv)
170        .arg_u32(qk_channels)
171        .arg_u32(head_dim)
172        .arg_u32(input_stride)
173        .arg_u32(output_stride)
174        .arg_u32(inter_stride)
175        .arg_f32(l2_eps)
176        .launch(stream)
177}
spark_model/layers/ops/ssm_gdn_snap.rs

spark_model/layers/ops/
ssm_gdn_snap.rs
