metrale_model_layers/layers/ops/
ssm_gdn_b.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the 3-token GDN verify decode and the
4//! WY-chunkwise 2- and 3-token verify decodes.
5//!
6//! Owner: model-layers ops (GDN).
7//! Invariants: none beyond the types.
8
9#![allow(unused_imports)]
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
14
15use crate::layers::moe;
16use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
17
18use super::*;
19
20/// 2026-09-25: Three tokens through the GDN recurrence in one launch
21/// (`gated_delta_rule_chunk3`). The states after the first and second tokens
22/// go to `h_state_inter0` and `h_state_inter1`, for rollback when drafts are
23/// rejected.
24#[allow(clippy::too_many_arguments)]
25pub fn gdn_decode_chunk3(
26    gpu: &dyn GpuBackend,
27    kernel: KernelHandle,
28    h_state: DevicePtr,
29    query: DevicePtr,
30    key: DevicePtr,
31    value: DevicePtr,
32    gate: DevicePtr,
33    beta: DevicePtr,
34    output: DevicePtr,
35    h_state_inter0: DevicePtr,
36    h_state_inter1: DevicePtr,
37    batch_size: 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        .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(output)
57        .arg_ptr(h_state_inter0)
58        .arg_ptr(h_state_inter1)
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_u32(qk_stride)
65        .arg_u32(v_stride)
66        .arg_u32(gb_stride)
67        .launch(stream)
68}
69
70/// 2026-09-25: WY-chunkwise 2-token verify decode (`gated_delta_rule_wy2`).
71/// One pass over H computes both tokens' `H^T k` products, the WY correction
72/// is applied, and a second pass updates H and writes the intermediate state.
73#[allow(clippy::too_many_arguments)]
74pub fn gdn_decode_wy2(
75    gpu: &dyn GpuBackend,
76    kernel: KernelHandle,
77    h_state: DevicePtr,
78    query: DevicePtr,
79    key: DevicePtr,
80    value: DevicePtr,
81    gate: DevicePtr,
82    beta: DevicePtr,
83    output: DevicePtr,
84    h_state_intermediate: DevicePtr,
85    batch_size: u32,
86    num_k_heads: u32,
87    num_v_heads: u32,
88    k_dim: u32,
89    v_dim: u32,
90    qk_stride: u32,
91    v_stride: u32,
92    gb_stride: u32,
93    // 2026-09-25: false: the state arguments are contiguous bases indexed by
94    // `(b * num_v_heads + vh)`, which assumes every intermediate shares
95    // h_state's per-sequence stride. The SSM pool does not lay them out that
96    // way, so the contiguous form is correct only at batch_size == 1.
97    // true: device pointer tables, one entry per sequence.
98    state_is_table: bool,
99    stream: u64,
100) -> Result<()> {
101    // 2026-09-25: `ensure!`, because a `debug_assert!` compiles out of release
102    // builds, where the misaddressing would be silent.
103    anyhow::ensure!(
104        state_is_table || batch_size == 1,
105        "gdn_decode_wy2: contiguous state addressing is only valid at \
106         batch_size==1 (got {batch_size}) — the intermediate's pool stride is \
107         num_intermediates x h_state's, so sequence 1's Hi0 would land on \
108         sequence 0's Hi1. Stage pointer tables and pass state_is_table=true."
109    );
110    KernelLaunch::new(gpu, kernel)
111        .grid([num_v_heads, batch_size, 1])
112        .block([128, 1, 1])
113        .arg_ptr(h_state)
114        .arg_ptr(query)
115        .arg_ptr(key)
116        .arg_ptr(value)
117        .arg_ptr(gate)
118        .arg_ptr(beta)
119        .arg_ptr(output)
120        .arg_ptr(h_state_intermediate)
121        .arg_u32(batch_size)
122        .arg_u32(num_k_heads)
123        .arg_u32(num_v_heads)
124        .arg_u32(k_dim)
125        .arg_u32(v_dim)
126        .arg_u32(qk_stride)
127        .arg_u32(v_stride)
128        .arg_u32(gb_stride)
129        .arg_u32(u32::from(state_is_table))
130        .launch(stream)
131}
132
133/// 2026-09-25: WY-chunkwise 3-token verify decode (`gated_delta_rule_wy3`):
134/// the [`gdn_decode_wy2`] scheme for three tokens, with two intermediate
135/// states.
136#[allow(clippy::too_many_arguments)]
137pub fn gdn_decode_wy3(
138    gpu: &dyn GpuBackend,
139    kernel: KernelHandle,
140    h_state: DevicePtr,
141    query: DevicePtr,
142    key: DevicePtr,
143    value: DevicePtr,
144    gate: DevicePtr,
145    beta: DevicePtr,
146    output: DevicePtr,
147    h_state_inter0: DevicePtr,
148    h_state_inter1: DevicePtr,
149    batch_size: u32,
150    num_k_heads: u32,
151    num_v_heads: u32,
152    k_dim: u32,
153    v_dim: u32,
154    qk_stride: u32,
155    v_stride: u32,
156    gb_stride: u32,
157    // 2026-09-25: false: the state arguments are contiguous bases indexed by
158    // `(b * num_v_heads + vh)`, which assumes every intermediate shares
159    // h_state's per-sequence stride. The SSM pool does not lay them out that
160    // way, so the contiguous form is correct only at batch_size == 1.
161    // true: device pointer tables, one entry per sequence.
162    state_is_table: bool,
163    stream: u64,
164) -> Result<()> {
165    // 2026-09-25: `ensure!`, because a `debug_assert!` compiles out of release
166    // builds, where the misaddressing would be silent.
167    anyhow::ensure!(
168        state_is_table || batch_size == 1,
169        "gdn_decode_wy3: contiguous state addressing is only valid at \
170         batch_size==1 (got {batch_size}) — the intermediates' pool stride is \
171         num_intermediates x h_state's, so sequence 1's Hi0 would land on \
172         sequence 0's Hi1. Stage pointer tables and pass state_is_table=true."
173    );
174    KernelLaunch::new(gpu, kernel)
175        .grid([num_v_heads, batch_size, 1])
176        .block([128, 1, 1])
177        .arg_ptr(h_state)
178        .arg_ptr(query)
179        .arg_ptr(key)
180        .arg_ptr(value)
181        .arg_ptr(gate)
182        .arg_ptr(beta)
183        .arg_ptr(output)
184        .arg_ptr(h_state_inter0)
185        .arg_ptr(h_state_inter1)
186        .arg_u32(batch_size)
187        .arg_u32(num_k_heads)
188        .arg_u32(num_v_heads)
189        .arg_u32(k_dim)
190        .arg_u32(v_dim)
191        .arg_u32(qk_stride)
192        .arg_u32(v_stride)
193        .arg_u32(gb_stride)
194        .arg_u32(u32::from(state_is_table))
195        .launch(stream)
196}
metrale_model_layers/layers/ops/ssm_gdn_b.rs

metrale_model_layers/layers/ops/
ssm_gdn_b.rs
