metrale_model_layers/layers/ops/
ssm_gdn_snap.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the exact-verify `_snap` kernels: the fused-norm
4//! GDN decodes with an inline per-token h-state snapshot, and the FP32-output
5//! fused verify conv.
6//!
7//! The kernels are optional (`try_kernel`; only the qwen3.6-27b/nvfp4 tree
8//! compiles them). Where a handle is 0, the caller runs the parent kernel and
9//! copies the snapshots with `copy_d2d_async`.
10//!
11//! Owner: model-layers ops (GDN).
12//! Invariants: none beyond the types.
13
14use anyhow::Result;
15use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
16use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
17
18use crate::weight_map::DenseWeight;
19
20/// 2026-09-25: [`super::gdn_decode_f32_norm`] plus an inline h-state snapshot.
21///
22/// `h_inter` receives the updated H after the state-norm clamp, the same
23/// values left in `h_state`; NULL skips the snapshot. Same grid, block and
24/// argument order as the parent, with `h_inter` appended.
25#[allow(clippy::too_many_arguments)]
26pub fn gdn_decode_f32_norm_snap(
27    gpu: &dyn GpuBackend,
28    kernel: KernelHandle,
29    h_state: DevicePtr,
30    query: DevicePtr,
31    key: DevicePtr,
32    value: DevicePtr,
33    gate: DevicePtr,
34    beta: DevicePtr,
35    z_gate: DevicePtr,
36    norm_weight: DevicePtr,
37    output: DevicePtr,
38    h_inter: DevicePtr,
39    batch_size: u32,
40    num_k_heads: u32,
41    num_v_heads: u32,
42    k_dim: u32,
43    v_dim: u32,
44    eps: f32,
45    stream: u64,
46) -> Result<()> {
47    KernelLaunch::new(gpu, kernel)
48        .grid([num_v_heads, batch_size, 1])
49        .block([128, 1, 1])
50        .arg_ptr(h_state)
51        .arg_ptr(query)
52        .arg_ptr(key)
53        .arg_ptr(value)
54        .arg_ptr(gate)
55        .arg_ptr(beta)
56        .arg_ptr(z_gate)
57        .arg_ptr(norm_weight)
58        .arg_ptr(output)
59        .arg_u32(batch_size)
60        .arg_u32(num_k_heads)
61        .arg_u32(num_v_heads)
62        .arg_u32(k_dim)
63        .arg_u32(v_dim)
64        .arg_f32(eps)
65        .arg_ptr(h_inter)
66        .launch(stream)
67}
68
69/// 2026-09-25: [`super::gdn_decode_f32_strided_norm`] plus an inline h-state
70/// snapshot, for the batched-verify arm over `batch_size` sequences.
71///
72/// `h_inter` is the snapshot base for this token position, and sequences'
73/// snapshots are `h_inter_seq_stride` FP32 elements apart. The caller measures
74/// that stride from the pool's pointers; it does not follow from the head
75/// dims. NULL skips the snapshot.
76#[allow(clippy::too_many_arguments)]
77pub fn gdn_decode_f32_strided_norm_snap(
78    gpu: &dyn GpuBackend,
79    kernel: KernelHandle,
80    h_state: DevicePtr,
81    query: DevicePtr,
82    key: DevicePtr,
83    value: DevicePtr,
84    gate: DevicePtr,
85    beta: DevicePtr,
86    z_gate: DevicePtr,
87    norm_weight: DevicePtr,
88    output: DevicePtr,
89    h_inter: DevicePtr,
90    h_inter_seq_stride: u64,
91    batch_size: u32,
92    num_k_heads: u32,
93    num_v_heads: u32,
94    k_dim: u32,
95    v_dim: u32,
96    qk_stride: u32,
97    v_stride: u32,
98    gb_stride: u32,
99    z_stride: u32,
100    out_stride: u32,
101    eps: f32,
102    stream: u64,
103) -> Result<()> {
104    KernelLaunch::new(gpu, kernel)
105        .grid([num_v_heads, batch_size, 1])
106        .block([128, 1, 1])
107        .arg_ptr(h_state)
108        .arg_ptr(query)
109        .arg_ptr(key)
110        .arg_ptr(value)
111        .arg_ptr(gate)
112        .arg_ptr(beta)
113        .arg_ptr(z_gate)
114        .arg_ptr(norm_weight)
115        .arg_ptr(output)
116        .arg_u32(batch_size)
117        .arg_u32(num_k_heads)
118        .arg_u32(num_v_heads)
119        .arg_u32(k_dim)
120        .arg_u32(v_dim)
121        .arg_u32(qk_stride)
122        .arg_u32(v_stride)
123        .arg_u32(gb_stride)
124        .arg_u32(z_stride)
125        .arg_u32(out_stride)
126        .arg_f32(eps)
127        .arg_ptr(h_inter)
128        .arg_u64(h_inter_seq_stride)
129        .launch(stream)
130}
131
132/// 2026-09-25: FP32-output twin of [`super::gdn_verify_fused_conv_kn`]: one
133/// launch runs conv1d + SiLU + L2 norm for all K verify positions, writes FP32
134/// conv rows, and writes each per-token conv-state snapshot inline.
135/// `output_stride` is in FP32 elements.
136#[allow(clippy::too_many_arguments)]
137pub fn gdn_verify_fused_conv_kn_f32(
138    gpu: &dyn GpuBackend,
139    kernel: KernelHandle,
140    conv_state: DevicePtr,
141    new_input: DevicePtr,
142    weight: &DenseWeight,
143    output: DevicePtr,
144    conv_state_inter: DevicePtr,
145    num_tokens: u32,
146    dim: u32,
147    d_conv: u32,
148    qk_channels: u32,
149    head_dim: u32,
150    input_stride: u32,
151    output_stride: u32,
152    inter_stride: u32,
153    l2_eps: f32,
154    stream: u64,
155) -> Result<()> {
156    KernelLaunch::new(gpu, kernel)
157        .grid([div_ceil(dim, 256), 1, 1])
158        .block([256, 1, 1])
159        .arg_ptr(conv_state)
160        .arg_ptr(new_input)
161        .arg_ptr(weight.weight)
162        .arg_ptr(output)
163        .arg_ptr(conv_state_inter)
164        .arg_u32(num_tokens)
165        .arg_u32(dim)
166        .arg_u32(d_conv)
167        .arg_u32(qk_channels)
168        .arg_u32(head_dim)
169        .arg_u32(input_stride)
170        .arg_u32(output_stride)
171        .arg_u32(inter_stride)
172        .arg_f32(l2_eps)
173        .launch(stream)
174}
metrale_model_layers/layers/ops/ssm_gdn_snap.rs

metrale_model_layers/layers/ops/
ssm_gdn_snap.rs
