metrale_model_layers/layers/ops/
ssm_gdn_a2.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the token-sequential GDN prefill kernels: the
4//! plain, split and split4 forms, and the persistent forms that keep H in
5//! shared memory.
6//!
7//! Owner: model-layers ops (GDN).
8//! Invariants: none beyond the types.
9
10#![allow(unused_imports)]
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use crate::layers::moe;
17use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
18
19use super::*;
20
21/// 2026-09-25: Token-sequential GDN prefill (`gated_delta_rule_prefill`): one
22/// block per (v-head, sequence) walks all `seq_len` tokens. Q/K/V and
23/// gate/beta are read at their own per-token strides, so the packed conv
24/// output needs no repacking.
25#[allow(clippy::too_many_arguments)]
26pub fn gdn_prefill(
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    batch_size: u32,
37    seq_len: u32,
38    num_k_heads: u32,
39    num_v_heads: u32,
40    k_dim: u32,
41    v_dim: u32,
42    qk_stride: u32,
43    v_stride: u32,
44    gb_stride: u32,
45    stream: u64,
46) -> Result<()> {
47    KernelLaunch::new(gpu, kernel)
48        .grid([num_v_heads, batch_size, 1])
49        .block([128, 1, 1])
50        .shared_mem(4 * k_dim * 4)
51        .arg_ptr(h_state)
52        .arg_ptr(query)
53        .arg_ptr(key)
54        .arg_ptr(value)
55        .arg_ptr(gate)
56        .arg_ptr(beta)
57        .arg_ptr(output)
58        .arg_u32(batch_size)
59        .arg_u32(seq_len)
60        .arg_u32(num_k_heads)
61        .arg_u32(num_v_heads)
62        .arg_u32(k_dim)
63        .arg_u32(v_dim)
64        .arg_u32(qk_stride)
65        .arg_u32(v_stride)
66        .arg_u32(gb_stride)
67        .launch(stream)
68}
69
70/// 2026-09-25: [`gdn_prefill`] with v_dim split across 2 blocks per v-head
71/// (`gated_delta_rule_prefill_split`).
72#[allow(clippy::too_many_arguments)]
73pub fn gdn_prefill_split(
74    gpu: &dyn GpuBackend,
75    kernel: KernelHandle,
76    h_state: DevicePtr,
77    query: DevicePtr,
78    key: DevicePtr,
79    value: DevicePtr,
80    gate: DevicePtr,
81    beta: DevicePtr,
82    output: DevicePtr,
83    batch_size: u32,
84    seq_len: u32,
85    num_k_heads: u32,
86    num_v_heads: u32,
87    k_dim: u32,
88    v_dim: u32,
89    qk_stride: u32,
90    v_stride: u32,
91    gb_stride: u32,
92    stream: u64,
93) -> Result<()> {
94    KernelLaunch::new(gpu, kernel)
95        .grid([num_v_heads * 2, batch_size, 1])
96        .block([64, 1, 1])
97        .shared_mem(4 * k_dim * 4)
98        .arg_ptr(h_state)
99        .arg_ptr(query)
100        .arg_ptr(key)
101        .arg_ptr(value)
102        .arg_ptr(gate)
103        .arg_ptr(beta)
104        .arg_ptr(output)
105        .arg_u32(batch_size)
106        .arg_u32(seq_len)
107        .arg_u32(num_k_heads)
108        .arg_u32(num_v_heads)
109        .arg_u32(k_dim)
110        .arg_u32(v_dim)
111        .arg_u32(qk_stride)
112        .arg_u32(v_stride)
113        .arg_u32(gb_stride)
114        .launch(stream)
115}
116
117/// 2026-09-25: [`gdn_prefill`] with v_dim split across 4 blocks of 32 threads
118/// per v-head (`gated_delta_rule_prefill_split4`).
119#[allow(clippy::too_many_arguments)]
120pub fn gdn_prefill_split4(
121    gpu: &dyn GpuBackend,
122    kernel: KernelHandle,
123    h_state: DevicePtr,
124    query: DevicePtr,
125    key: DevicePtr,
126    value: DevicePtr,
127    gate: DevicePtr,
128    beta: DevicePtr,
129    output: DevicePtr,
130    batch_size: u32,
131    seq_len: u32,
132    num_k_heads: u32,
133    num_v_heads: u32,
134    k_dim: u32,
135    v_dim: u32,
136    qk_stride: u32,
137    v_stride: u32,
138    gb_stride: u32,
139    stream: u64,
140) -> Result<()> {
141    KernelLaunch::new(gpu, kernel)
142        .grid([num_v_heads * 4, batch_size, 1])
143        .block([32, 1, 1])
144        .shared_mem(4 * k_dim * 4) // 2026-09-25: double-buffered k and q, 4 * k_dim floats.
145        .arg_ptr(h_state)
146        .arg_ptr(query)
147        .arg_ptr(key)
148        .arg_ptr(value)
149        .arg_ptr(gate)
150        .arg_ptr(beta)
151        .arg_ptr(output)
152        .arg_u32(batch_size)
153        .arg_u32(seq_len)
154        .arg_u32(num_k_heads)
155        .arg_u32(num_v_heads)
156        .arg_u32(k_dim)
157        .arg_u32(v_dim)
158        .arg_u32(qk_stride)
159        .arg_u32(v_stride)
160        .arg_u32(gb_stride)
161        .launch(stream)
162}
163
164/// 2026-09-25: Persistent GDN prefill (`gated_delta_rule_prefill_persistent`):
165/// each block loads its v-head's H into shared memory once, runs every token
166/// against it, and writes it back to global memory once at the end.
167#[allow(clippy::too_many_arguments)]
168pub fn gdn_prefill_persistent(
169    gpu: &dyn GpuBackend,
170    kernel: KernelHandle,
171    h_state: DevicePtr,
172    query: DevicePtr,
173    key: DevicePtr,
174    value: DevicePtr,
175    gate: DevicePtr,
176    beta: DevicePtr,
177    output: DevicePtr,
178    batch_size: u32,
179    seq_len: u32,
180    num_k_heads: u32,
181    num_v_heads: u32,
182    k_dim: u32,
183    v_dim: u32,
184    qk_stride: u32,
185    v_stride: u32,
186    gb_stride: u32,
187    stream: u64,
188) -> Result<()> {
189    let smem = k_dim * v_dim * 4 + 4 * k_dim * 4; // 2026-09-25: H, then double-buffered k and q.
190    KernelLaunch::new(gpu, kernel)
191        .grid([num_v_heads, batch_size, 1])
192        .block([128, 1, 1])
193        .shared_mem(smem)
194        .arg_ptr(h_state)
195        .arg_ptr(query)
196        .arg_ptr(key)
197        .arg_ptr(value)
198        .arg_ptr(gate)
199        .arg_ptr(beta)
200        .arg_ptr(output)
201        .arg_u32(batch_size)
202        .arg_u32(seq_len)
203        .arg_u32(num_k_heads)
204        .arg_u32(num_v_heads)
205        .arg_u32(k_dim)
206        .arg_u32(v_dim)
207        .arg_u32(qk_stride)
208        .arg_u32(v_stride)
209        .arg_u32(gb_stride)
210        .launch(stream)
211}
212
213/// 2026-09-25: Persistent GDN prefill whose shared-memory size the caller
214/// computes, for kernels with a larger layout (the WY4-persistent and WY64
215/// prefill kernels).
216#[allow(clippy::too_many_arguments)]
217pub fn gdn_prefill_persistent_smem(
218    gpu: &dyn GpuBackend,
219    kernel: KernelHandle,
220    h_state: DevicePtr,
221    query: DevicePtr,
222    key: DevicePtr,
223    value: DevicePtr,
224    gate: DevicePtr,
225    beta: DevicePtr,
226    output: DevicePtr,
227    batch_size: u32,
228    seq_len: u32,
229    num_k_heads: u32,
230    num_v_heads: u32,
231    k_dim: u32,
232    v_dim: u32,
233    qk_stride: u32,
234    v_stride: u32,
235    gb_stride: u32,
236    smem: u32,
237    stream: u64,
238) -> Result<()> {
239    KernelLaunch::new(gpu, kernel)
240        .grid([num_v_heads, batch_size, 1])
241        .block([128, 1, 1])
242        .shared_mem(smem)
243        .arg_ptr(h_state)
244        .arg_ptr(query)
245        .arg_ptr(key)
246        .arg_ptr(value)
247        .arg_ptr(gate)
248        .arg_ptr(beta)
249        .arg_ptr(output)
250        .arg_u32(batch_size)
251        .arg_u32(seq_len)
252        .arg_u32(num_k_heads)
253        .arg_u32(num_v_heads)
254        .arg_u32(k_dim)
255        .arg_u32(v_dim)
256        .arg_u32(qk_stride)
257        .arg_u32(v_stride)
258        .arg_u32(gb_stride)
259        .launch(stream)
260}
metrale_model_layers/layers/ops/ssm_gdn_a2.rs

metrale_model_layers/layers/ops/
ssm_gdn_a2.rs
