metrale_model_layers/layers/ops/
ssm_gdn_wyn.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the WY-chunkwise verify decodes at K = 4 and
4//! K >= 5 (contiguous and pointer-table forms), and the 2-token conv1d update.
5//!
6//! Owner: model-layers ops (GDN).
7//! Invariants: none beyond the types.
8
9// SPDX-License-Identifier: MIT OR Apache-2.0
10
11#![allow(unused_imports)]
12
13use anyhow::Result;
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
15use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
16
17use crate::layers::moe;
18use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
19
20use super::*;
21
22/// 2026-09-25: WY-chunkwise 4-token verify decode (`gated_delta_rule_wy4`):
23/// the [`gdn_decode_wy2`] scheme for four tokens, with three intermediate
24/// states.
25#[allow(clippy::too_many_arguments)]
26pub fn gdn_decode_wy4(
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    output: DevicePtr,
36    h_state_inter0: DevicePtr,
37    h_state_inter1: DevicePtr,
38    h_state_inter2: DevicePtr,
39    batch_size: u32,
40    num_k_heads: u32,
41    num_v_heads: u32,
42    k_dim: u32,
43    v_dim: u32,
44    qk_stride: u32,
45    v_stride: u32,
46    gb_stride: u32,
47    // 2026-09-25: false: the state arguments are contiguous bases indexed by
48    // `(b * num_v_heads + vh)`, which assumes every intermediate shares
49    // h_state's per-sequence stride. The SSM pool does not lay them out that
50    // way, so the contiguous form is correct only at batch_size == 1.
51    // true: device pointer tables, one entry per sequence; any batched verify
52    // passes true.
53    state_is_table: bool,
54    stream: u64,
55) -> Result<()> {
56    // 2026-09-25: `ensure!`, because a `debug_assert!` compiles out of release
57    // builds, where the misaddressing would be silent.
58    anyhow::ensure!(
59        state_is_table || batch_size == 1,
60        "gdn_decode_wy4: contiguous state addressing is only valid at \
61         batch_size==1 (got {batch_size}) — the intermediates' pool stride is \
62         num_intermediates x h_state's, so sequence 1's Hi0 would land on \
63         sequence 0's Hi1. Stage pointer tables and pass state_is_table=true."
64    );
65    KernelLaunch::new(gpu, kernel)
66        .grid([num_v_heads, batch_size, 1])
67        .block([128, 1, 1])
68        .arg_ptr(h_state)
69        .arg_ptr(query)
70        .arg_ptr(key)
71        .arg_ptr(value)
72        .arg_ptr(gate)
73        .arg_ptr(beta)
74        .arg_ptr(output)
75        .arg_ptr(h_state_inter0)
76        .arg_ptr(h_state_inter1)
77        .arg_ptr(h_state_inter2)
78        .arg_u32(batch_size)
79        .arg_u32(num_k_heads)
80        .arg_u32(num_v_heads)
81        .arg_u32(k_dim)
82        .arg_u32(v_dim)
83        .arg_u32(qk_stride)
84        .arg_u32(v_stride)
85        .arg_u32(gb_stride)
86        .arg_u32(u32::from(state_is_table))
87        .launch(stream)
88}
89
90/// 2026-09-25: WY-chunkwise verify decode for K tokens, contiguous form: the
91/// launch for `gated_delta_rule_wy5` .. `wy16` (and their `_f16` twins) from
92/// `gated_delta_rule_wyn.cu`, and for `gated_delta_rule_wy17`. K is
93/// compile-time in the kernel, so the `kernel` handle selects it. The kernel
94/// writes Hi_0..Hi_{K-2} and the final H.
95///
96/// Hi_t is at `h_state_inter_base + t * inter_stride_floats`, plus the
97/// per-(b, vh) offset; the stride is in the kernel's h element size.
98#[allow(clippy::too_many_arguments)]
99pub fn gdn_decode_wyn(
100    gpu: &dyn GpuBackend,
101    kernel: KernelHandle,
102    h_state: DevicePtr,
103    query: DevicePtr,
104    key: DevicePtr,
105    value: DevicePtr,
106    gate: DevicePtr,
107    beta: DevicePtr,
108    output: DevicePtr,
109    h_state_inter_base: DevicePtr,
110    inter_stride_floats: u32,
111    batch_size: u32,
112    num_k_heads: u32,
113    num_v_heads: u32,
114    k_dim: u32,
115    v_dim: u32,
116    qk_stride: u32,
117    v_stride: u32,
118    gb_stride: u32,
119    stream: u64,
120) -> Result<()> {
121    // 2026-09-25: The contiguous form offsets both h_state and the
122    // intermediates by `(b * num_v_heads + vh) * hv`, which is right only at
123    // batch_size == 1. Batched verify goes through [`gdn_decode_wyn_table`].
124    anyhow::ensure!(
125        batch_size == 1,
126        "gdn_decode_wyn: contiguous state addressing is only valid at batch_size==1 \
127         (got {batch_size}); port the wy4 `state_is_table` pointer-table form first"
128    );
129    KernelLaunch::new(gpu, kernel)
130        .grid([num_v_heads, batch_size, 1])
131        .block([128, 1, 1])
132        .arg_ptr(h_state)
133        .arg_ptr(query)
134        .arg_ptr(key)
135        .arg_ptr(value)
136        .arg_ptr(gate)
137        .arg_ptr(beta)
138        .arg_ptr(output)
139        .arg_ptr(h_state_inter_base)
140        .arg_u32(inter_stride_floats)
141        .arg_u32(batch_size)
142        .arg_u32(num_k_heads)
143        .arg_u32(num_v_heads)
144        .arg_u32(k_dim)
145        .arg_u32(v_dim)
146        .arg_u32(qk_stride)
147        .arg_u32(v_stride)
148        .arg_u32(gb_stride)
149        .launch(stream)
150}
151
152/// 2026-09-25: Cross-sequence batched-verify launch of the wyN pointer-table
153/// twins (`gated_delta_rule_wy{5..16}_table` and their `_f16_table` forms).
154/// `state_is_table` is compiled into the symbol, so the argument list is the
155/// same as [`gdn_decode_wyn`], but the state arguments are tables:
156/// - `h_tables`: the layer's staged table slice; slab 0 holds one
157///   per-sequence H base pointer per entry (`VERIFY_WY_TABLE_SEQS` entries).
158/// - `hi_tables`: slab 1 (Hi0); the following Hi slabs are
159///   `VERIFY_WY_TABLE_SEQS` entries apart.
160/// - `slab_entry_stride` must be `VERIFY_WY_TABLE_SEQS as u32`: the kernel
161///   reads the contiguous form's stride argument as the number of pointer
162///   entries between Hi slabs.
163///
164/// The tables are staged by `upload_verify_wy_tables` (model-engine
165/// `verify_e2.rs`).
166#[allow(clippy::too_many_arguments)]
167#[allow(dead_code)]
168pub fn gdn_decode_wyn_table(
169    gpu: &dyn GpuBackend,
170    kernel: KernelHandle,
171    h_tables: DevicePtr,
172    query: DevicePtr,
173    key: DevicePtr,
174    value: DevicePtr,
175    gate: DevicePtr,
176    beta: DevicePtr,
177    output: DevicePtr,
178    hi_tables: DevicePtr,
179    slab_entry_stride: u32,
180    batch_size: u32,
181    num_k_heads: u32,
182    num_v_heads: u32,
183    k_dim: u32,
184    v_dim: u32,
185    qk_stride: u32,
186    v_stride: u32,
187    gb_stride: u32,
188    stream: u64,
189) -> Result<()> {
190    anyhow::ensure!(
191        !h_tables.is_null() && !hi_tables.is_null(),
192        "gdn_decode_wyn_table: staged pointer tables are required (NULL table \
193         means upload_verify_wy_tables declined — run the per-sequence loop)"
194    );
195    KernelLaunch::new(gpu, kernel)
196        .grid([num_v_heads, batch_size, 1])
197        .block([128, 1, 1])
198        .arg_ptr(h_tables)
199        .arg_ptr(query)
200        .arg_ptr(key)
201        .arg_ptr(value)
202        .arg_ptr(gate)
203        .arg_ptr(beta)
204        .arg_ptr(output)
205        .arg_ptr(hi_tables)
206        .arg_u32(slab_entry_stride)
207        .arg_u32(batch_size)
208        .arg_u32(num_k_heads)
209        .arg_u32(num_v_heads)
210        .arg_u32(k_dim)
211        .arg_u32(v_dim)
212        .arg_u32(qk_stride)
213        .arg_u32(v_stride)
214        .arg_u32(gb_stride)
215        .launch(stream)
216}
217
218/// 2026-09-25: Conv1d update + SiLU for two tokens per sequence in one launch
219/// (`causal_conv1d_update_chunk2`), one thread per channel. The conv state
220/// after the first token goes to `conv_state_intermediate`, for rollback.
221#[allow(clippy::too_many_arguments)]
222pub fn conv1d_update_chunk2(
223    gpu: &dyn GpuBackend,
224    kernel: KernelHandle,
225    conv_state: DevicePtr,
226    input: DevicePtr,
227    weight: &DenseWeight,
228    output: DevicePtr,
229    conv_state_intermediate: DevicePtr,
230    d_inner: u32,
231    d_conv: u32,
232    batch_size: u32,
233    stream: u64,
234) -> Result<()> {
235    let bias_ptr = DevicePtr::NULL;
236    KernelLaunch::new(gpu, kernel)
237        .grid([div_ceil(d_inner, 256), batch_size, 1])
238        .block([256, 1, 1])
239        .arg_ptr(conv_state)
240        .arg_ptr(input)
241        .arg_ptr(weight.weight)
242        .arg_ptr(bias_ptr)
243        .arg_ptr(output)
244        .arg_ptr(conv_state_intermediate)
245        .arg_u32(batch_size)
246        .arg_u32(d_inner)
247        .arg_u32(d_conv)
248        .launch(stream)
249}
metrale_model_layers/layers/ops/ssm_gdn_wyn.rs

metrale_model_layers/layers/ops/
ssm_gdn_wyn.rs
