metrale_model_layers/layers/ops/
ssm_gdn_carry.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Launchers for the carried-state GDN verify
4//! (`kernels/gb10/common/gated_delta_rule_carry.cu`): the K = 2..4 verify
5//! kernels, the conv twin, and their standalone folds.
6//!
7//! Owner: model-layers ops (GDN).
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
12use metrale_gpu_runtime::kernel_args::KernelLaunch;
13
14/// 2026-09-26: Carried-state verify (`gdn_carry_wy{2,3,4}` or its `_lazy` form, chosen
15/// by the caller's `kernel`). It applies each sequence's pending rows
16/// (`pend[slot_tab[b]]`) to its copy of H and writes H back (the lazy form only when
17/// this verify's rows might not fit behind them in the stash,
18/// `pending + K > GDN_CARRY_CAP`), writes the output, and stashes the rows it verified
19/// at `carry_base + slot * seq_floats`. Engaged word b becomes 2 after a write-back,
20/// else 1.
21#[allow(clippy::too_many_arguments)]
22pub fn gdn_carry_wy(
23    gpu: &dyn GpuBackend,
24    kernel: KernelHandle,
25    h_table: DevicePtr,
26    query: DevicePtr,
27    key: DevicePtr,
28    value: DevicePtr,
29    gate: DevicePtr,
30    beta: DevicePtr,
31    output: DevicePtr,
32    carry_base: DevicePtr,
33    slot_tab: DevicePtr,
34    pend: DevicePtr,
35    seq_floats: u32,
36    batch_size: u32,
37    num_k_heads: u32,
38    num_v_heads: u32,
39    qk_stride: u32,
40    v_stride: u32,
41    gb_stride: u32,
42    k_dim: u32,
43    engaged_flag: DevicePtr,
44    stream: u64,
45) -> Result<()> {
46    anyhow::ensure!(
47        !h_table.is_null()
48            && !carry_base.is_null()
49            && !slot_tab.is_null()
50            && !pend.is_null()
51            && !engaged_flag.is_null(),
52        "gdn_carry_wy: null table/stash/slot/pend/flag"
53    );
54    KernelLaunch::new(gpu, kernel)
55        .grid([num_v_heads, batch_size, 1])
56        .block([128, 1, 1])
57        .arg_ptr(h_table)
58        .arg_ptr(query)
59        .arg_ptr(key)
60        .arg_ptr(value)
61        .arg_ptr(gate)
62        .arg_ptr(beta)
63        .arg_ptr(output)
64        .arg_ptr(carry_base)
65        .arg_ptr(slot_tab)
66        .arg_ptr(pend)
67        .arg_u32(seq_floats)
68        .arg_u32(batch_size)
69        .arg_u32(num_k_heads)
70        .arg_u32(num_v_heads)
71        .arg_u32(qk_stride)
72        .arg_u32(v_stride)
73        .arg_u32(gb_stride)
74        .arg_u32(k_dim)
75        .arg_ptr(engaged_flag)
76        .launch(stream)
77}
78
79/// 2026-09-26: Standalone fold (`gdn_carry_flush`) over `layers` GDN layers:
80/// layer `l` reads its h pointers at entry `l * table_layer_entries` of
81/// `h_table`, its stash at `carry_base + l * carry_layer_floats` floats and
82/// its counts at `pend + l * pend_layer_entries`. A count of 0 leaves H alone;
83/// `pend` is not cleared.
84#[allow(clippy::too_many_arguments)]
85pub fn gdn_carry_flush(
86    gpu: &dyn GpuBackend,
87    kernel: KernelHandle,
88    h_table: DevicePtr,
89    table_layer_entries: u64,
90    carry_base: DevicePtr,
91    carry_layer_floats: u64,
92    slot_tab: DevicePtr,
93    pend: DevicePtr,
94    pend_layer_entries: u32,
95    seq_floats: u32,
96    batch_size: u32,
97    num_v_heads: u32,
98    layers: u32,
99    stream: u64,
100) -> Result<()> {
101    KernelLaunch::new(gpu, kernel)
102        .grid([num_v_heads, batch_size, layers])
103        .block([128, 1, 1])
104        .arg_ptr(h_table)
105        .arg_u64(table_layer_entries)
106        .arg_ptr(carry_base)
107        .arg_u64(carry_layer_floats)
108        .arg_ptr(slot_tab)
109        .arg_ptr(pend)
110        .arg_u32(pend_layer_entries)
111        .arg_u32(seq_floats)
112        .arg_u32(batch_size)
113        .arg_u32(num_v_heads)
114        .launch(stream)
115}
116
117/// 2026-09-26: Carried-state conv verify (`gdn_carry_conv`): the twin of
118/// `gdn_verify_fused_conv_kn_batched` that first shifts each sequence's pending input rows
119/// into its window (writing it back when there were any), writes no snapshot and no final
120/// window, and stashes the position inputs at `conv_stash + slot * stash_seq_elems`. `lazy`
121/// must match the verify kernel of the same run (`gdn_carry_wy{K}_lazy`): the window is then
122/// written back only when the stash is full, else on every launch.
123#[allow(clippy::too_many_arguments)]
124pub fn gdn_carry_conv(
125    gpu: &dyn GpuBackend,
126    kernel: KernelHandle,
127    conv_state: DevicePtr,
128    new_input: DevicePtr,
129    weight: &crate::weight_map::DenseWeight,
130    output: DevicePtr,
131    conv_stash: DevicePtr,
132    slot_tab: DevicePtr,
133    pend: DevicePtr,
134    stash_seq_elems: u32,
135    num_tokens: u32,
136    dim: u32,
137    d_conv: u32,
138    qk_channels: u32,
139    head_dim: u32,
140    input_stride: u32,
141    output_stride: u32,
142    l2_eps: f32,
143    n_seq: u32,
144    conv_state_seq_stride: u32,
145    input_seq_stride: u32,
146    output_seq_stride: u32,
147    lazy: bool,
148    stream: u64,
149) -> Result<()> {
150    KernelLaunch::new(gpu, kernel)
151        .grid([dim.div_ceil(256), n_seq, 1])
152        .block([256, 1, 1])
153        .arg_ptr(conv_state)
154        .arg_ptr(new_input)
155        .arg_ptr(weight.weight)
156        .arg_ptr(output)
157        .arg_ptr(conv_stash)
158        .arg_ptr(slot_tab)
159        .arg_ptr(pend)
160        .arg_u32(stash_seq_elems)
161        .arg_u32(num_tokens)
162        .arg_u32(dim)
163        .arg_u32(d_conv)
164        .arg_u32(qk_channels)
165        .arg_u32(head_dim)
166        .arg_u32(input_stride)
167        .arg_u32(output_stride)
168        .arg_f32(l2_eps)
169        .arg_u32(conv_state_seq_stride)
170        .arg_u32(input_seq_stride)
171        .arg_u32(output_seq_stride)
172        .arg_u32(u32::from(lazy))
173        .launch(stream)
174}
175
176/// 2026-10-01: Carried-state exact verify (`gdn_exact_carry{2,3,4}` or its `_lazy` form, chosen
177/// by the caller's `kernel`; `gdn_exact_carry.cu` in the model directory): the stash protocol of
178/// [`gdn_carry_wy`] around the strided FP32 decode's per-token chain. Rows are sequence-major,
179/// `b * K + t`: FP32 q and k at `qk_stride`, v at `v_stride`, gate and beta at `gb_stride`, the
180/// FP32 output at `out_stride`, all in floats.
181#[allow(clippy::too_many_arguments)]
182pub fn gdn_exact_carry(
183    gpu: &dyn GpuBackend,
184    kernel: KernelHandle,
185    h_table: DevicePtr,
186    query: DevicePtr,
187    key: DevicePtr,
188    value: DevicePtr,
189    gate: DevicePtr,
190    beta: DevicePtr,
191    output: DevicePtr,
192    carry_base: DevicePtr,
193    slot_tab: DevicePtr,
194    pend: DevicePtr,
195    seq_floats: u32,
196    batch_size: u32,
197    num_k_heads: u32,
198    num_v_heads: u32,
199    k_dim: u32,
200    strides: [u32; 4],
201    engaged_flag: DevicePtr,
202    stream: u64,
203) -> Result<()> {
204    anyhow::ensure!(
205        !h_table.is_null()
206            && !carry_base.is_null()
207            && !slot_tab.is_null()
208            && !pend.is_null()
209            && !engaged_flag.is_null(),
210        "gdn_exact_carry: null table/stash/slot/pend/flag"
211    );
212    let [qk_stride, v_stride, gb_stride, out_stride] = strides;
213    KernelLaunch::new(gpu, kernel)
214        .grid([num_v_heads, batch_size, 1])
215        .block([128, 1, 1])
216        .arg_ptr(h_table)
217        .arg_ptr(query)
218        .arg_ptr(key)
219        .arg_ptr(value)
220        .arg_ptr(gate)
221        .arg_ptr(beta)
222        .arg_ptr(output)
223        .arg_ptr(carry_base)
224        .arg_ptr(slot_tab)
225        .arg_ptr(pend)
226        .arg_u32(seq_floats)
227        .arg_u32(batch_size)
228        .arg_u32(num_k_heads)
229        .arg_u32(num_v_heads)
230        .arg_u32(k_dim)
231        .arg_u32(qk_stride)
232        .arg_u32(v_stride)
233        .arg_u32(gb_stride)
234        .arg_u32(out_stride)
235        .arg_ptr(engaged_flag)
236        .launch(stream)
237}
238
239/// 2026-10-01: Carried-state FP32 conv verify (`gdn_carry_conv_f32`): the twin of
240/// `causal_conv1d_update_l2norm_f32_strided` for all `num_tokens` positions of `n_seq`
241/// sequences, with [`gdn_carry_conv`]'s pending-row fold and stash. Sequence b's window is at
242/// `conv_state + b * dim * d_conv` floats; row `b * num_tokens + t` reads `new_input` at
243/// `input_stride` and writes the FP32 `output` at `output_stride`. `lazy` must match the
244/// verify kernel of the same run.
245#[allow(clippy::too_many_arguments)]
246pub fn gdn_carry_conv_f32(
247    gpu: &dyn GpuBackend,
248    kernel: KernelHandle,
249    conv_state: DevicePtr,
250    new_input: DevicePtr,
251    weight: &crate::weight_map::DenseWeight,
252    output: DevicePtr,
253    conv_stash: DevicePtr,
254    slot_tab: DevicePtr,
255    pend: DevicePtr,
256    stash_seq_elems: u32,
257    num_tokens: u32,
258    dim: u32,
259    d_conv: u32,
260    qk_channels: u32,
261    head_dim: u32,
262    l2_eps: f32,
263    input_stride: u32,
264    output_stride: u32,
265    n_seq: u32,
266    lazy: bool,
267    stream: u64,
268) -> Result<()> {
269    KernelLaunch::new(gpu, kernel)
270        .grid([dim.div_ceil(256), n_seq, 1])
271        .block([256, 1, 1])
272        .arg_ptr(conv_state)
273        .arg_ptr(new_input)
274        .arg_ptr(weight.weight)
275        .arg_ptr(output)
276        .arg_ptr(conv_stash)
277        .arg_ptr(slot_tab)
278        .arg_ptr(pend)
279        .arg_u32(stash_seq_elems)
280        .arg_u32(num_tokens)
281        .arg_u32(dim)
282        .arg_u32(d_conv)
283        .arg_u32(qk_channels)
284        .arg_u32(head_dim)
285        .arg_f32(l2_eps)
286        .arg_u32(input_stride)
287        .arg_u32(output_stride)
288        .arg_u32(u32::from(lazy))
289        .launch(stream)
290}
291
292/// 2026-10-01: The single-sequence exact verify (`gdn_exact_chain{2,3,4}`, chosen by the caller's
293/// `kernel`): K rows of one sequence's strided-decode chain in one launch. `h_state` is updated to
294/// the state after the last row, `h_inter[t]` receives the state after row t (t < K - 1; the other
295/// entries must still be valid pointers and are not written). Rows are `strides =
296/// [qk, v, gate|beta, out]` floats apart.
297#[allow(clippy::too_many_arguments)]
298pub fn gdn_exact_chain(
299    gpu: &dyn GpuBackend,
300    kernel: KernelHandle,
301    h_state: DevicePtr,
302    query: DevicePtr,
303    key: DevicePtr,
304    value: DevicePtr,
305    gate: DevicePtr,
306    beta: DevicePtr,
307    output: DevicePtr,
308    h_inter: [DevicePtr; 3],
309    num_k_heads: u32,
310    num_v_heads: u32,
311    k_dim: u32,
312    strides: [u32; 4],
313    stream: u64,
314) -> Result<()> {
315    let [qk_stride, v_stride, gb_stride, out_stride] = strides;
316    KernelLaunch::new(gpu, kernel)
317        .grid([num_v_heads, 1, 1])
318        .block([128, 1, 1])
319        .arg_ptr(h_state)
320        .arg_ptr(query)
321        .arg_ptr(key)
322        .arg_ptr(value)
323        .arg_ptr(gate)
324        .arg_ptr(beta)
325        .arg_ptr(output)
326        .arg_ptr(h_inter[0])
327        .arg_ptr(h_inter[1])
328        .arg_ptr(h_inter[2])
329        .arg_u32(num_k_heads)
330        .arg_u32(num_v_heads)
331        .arg_u32(k_dim)
332        .arg_u32(qk_stride)
333        .arg_u32(v_stride)
334        .arg_u32(gb_stride)
335        .arg_u32(out_stride)
336        .launch(stream)
337}
338
339/// 2026-10-01: The single-sequence exact verify's FP32 conv (`gdn_conv_chain_f32`): `num_tokens`
340/// positions of `causal_conv1d_update_l2norm_f32_strided` for one sequence in one launch, the
341/// window after position t (t < num_tokens - 1) at `conv_inter[t]`, the final window at
342/// `conv_state`. The unused `conv_inter` entries must be valid pointers; they are not written.
343#[allow(clippy::too_many_arguments)]
344pub fn gdn_conv_chain_f32(
345    gpu: &dyn GpuBackend,
346    kernel: KernelHandle,
347    conv_state: DevicePtr,
348    new_input: DevicePtr,
349    weight: &crate::weight_map::DenseWeight,
350    output: DevicePtr,
351    conv_inter: [DevicePtr; 3],
352    num_tokens: u32,
353    dim: u32,
354    d_conv: u32,
355    qk_channels: u32,
356    head_dim: u32,
357    l2_eps: f32,
358    input_stride: u32,
359    output_stride: u32,
360    stream: u64,
361) -> Result<()> {
362    KernelLaunch::new(gpu, kernel)
363        .grid([dim.div_ceil(256), 1, 1])
364        .block([256, 1, 1])
365        .arg_ptr(conv_state)
366        .arg_ptr(new_input)
367        .arg_ptr(weight.weight)
368        .arg_ptr(output)
369        .arg_ptr(conv_inter[0])
370        .arg_ptr(conv_inter[1])
371        .arg_ptr(conv_inter[2])
372        .arg_u32(num_tokens)
373        .arg_u32(dim)
374        .arg_u32(d_conv)
375        .arg_u32(qk_channels)
376        .arg_u32(head_dim)
377        .arg_f32(l2_eps)
378        .arg_u32(input_stride)
379        .arg_u32(output_stride)
380        .launch(stream)
381}
382
383/// 2026-09-26: Standalone conv fold (`gdn_carry_conv_flush`) over `layers` GDN layers, laid
384/// out as in [`gdn_carry_flush`] with conv-state pointers and a BF16 stash.
385#[allow(clippy::too_many_arguments)]
386pub fn gdn_carry_conv_flush(
387    gpu: &dyn GpuBackend,
388    kernel: KernelHandle,
389    state_table: DevicePtr,
390    table_layer_entries: u64,
391    conv_stash: DevicePtr,
392    stash_layer_elems: u64,
393    slot_tab: DevicePtr,
394    pend: DevicePtr,
395    pend_layer_entries: u32,
396    stash_seq_elems: u32,
397    batch_size: u32,
398    dim: u32,
399    d_conv: u32,
400    layers: u32,
401    stream: u64,
402) -> Result<()> {
403    KernelLaunch::new(gpu, kernel)
404        .grid([dim.div_ceil(256), batch_size, layers])
405        .block([256, 1, 1])
406        .arg_ptr(state_table)
407        .arg_u64(table_layer_entries)
408        .arg_ptr(conv_stash)
409        .arg_u64(stash_layer_elems)
410        .arg_ptr(slot_tab)
411        .arg_ptr(pend)
412        .arg_u32(pend_layer_entries)
413        .arg_u32(stash_seq_elems)
414        .arg_u32(batch_size)
415        .arg_u32(dim)
416        .arg_u32(d_conv)
417        .launch(stream)
418}
419
420/// 2026-09-26: Rows the carry stash holds per (layer, slot): `CARRY_CAP` in
421/// gated_delta_rule_carry.cu. A verify whose K rows might not fit behind the pending ones
422/// (`pending + K > GDN_CARRY_CAP`) writes the state back and restarts the stash.
423pub const GDN_CARRY_CAP: usize = 8;
424
425/// 2026-09-26: Launch width (sequences) from which a run takes the lazy verify kernels,
426/// which write the state back only when the stash is full. Measured on dgx3 (k=1,
427/// --mtp-gate force): from 16 the lazy form is 2.5 % faster end to end at C=16; a threshold
428/// of 8 instead of 16 gives C=8 214.4 -> 216.2 tok/s (three reps each) with C=4 unchanged;
429/// at 2 and 4 sequences the eager form is ahead.
430pub const GDN_CARRY_LAZY_MIN_SEQS: usize = 8;
431
432/// 2026-09-26: Most pending rows an eager verify kernel folds; a slot holding more (only
433/// after lazy verifies) is folded by `gdn_carry_flush` first.
434pub const GDN_CARRY_EAGER_MAX_PENDING: usize = 4;
435
436/// 2026-09-26: Per-(layer, slot) conv stash length in BF16 elements: `GDN_CARRY_CAP` input
437/// rows.
438pub const fn gdn_carry_conv_seq_elems(conv_dim: usize) -> usize {
439    GDN_CARRY_CAP * conv_dim
440}
441
442/// 2026-09-26: Per-(layer, slot) stash width in floats:
443/// `vn[CAP][nv][vd] | g[CAP][nv] | sk[CAP][nv][kd]`, the layout of the kernel's
444/// `CARRY_VN`/`CARRY_G`/`CARRY_SK` macros.
445pub const fn gdn_carry_seq_floats(nv: usize, kd: usize, vd: usize) -> usize {
446    GDN_CARRY_CAP * (nv * vd + nv + nv * kd)
447}
metrale_model_layers/layers/ops/ssm_gdn_carry.rs

metrale_model_layers/layers/ops/
ssm_gdn_carry.rs
