spark_model/layers/ops/
ssm_gdn_b.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `ops.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
10
11use crate::layers::moe;
12use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
13
14use super::*;
15
16/// Fused 3-token GDN decode (K=3 speculative verification).
17///
18/// Processes exactly 3 tokens through GDN in a single kernel launch.
19/// Saves 2 intermediate H states (H_1, H_2) for rollback on draft rejection.
20/// 4 passes vs 6 for 3× sequential decode.
21///
22/// Kernel: `gated_delta_rule_chunk3(h_state, query, key, value, gate, beta,
23///          output, h_inter0, h_inter1, batch_size, num_k_heads,
24///          num_v_heads, k_dim, v_dim, qk_stride, v_stride, gb_stride)`
25/// Grid: (num_v_heads, batch, 1)  Block: (128, 1, 1)
26#[allow(clippy::too_many_arguments)]
27pub fn gdn_decode_chunk3(
28    gpu: &dyn GpuBackend,
29    kernel: KernelHandle,
30    h_state: DevicePtr,
31    query: DevicePtr,
32    key: DevicePtr,
33    value: DevicePtr,
34    gate: DevicePtr,
35    beta: DevicePtr,
36    output: DevicePtr,
37    h_state_inter0: DevicePtr,
38    h_state_inter1: 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    stream: u64,
48) -> Result<()> {
49    KernelLaunch::new(gpu, kernel)
50        .grid([num_v_heads, batch_size, 1])
51        .block([128, 1, 1])
52        .arg_ptr(h_state)
53        .arg_ptr(query)
54        .arg_ptr(key)
55        .arg_ptr(value)
56        .arg_ptr(gate)
57        .arg_ptr(beta)
58        .arg_ptr(output)
59        .arg_ptr(h_state_inter0)
60        .arg_ptr(h_state_inter1)
61        .arg_u32(batch_size)
62        .arg_u32(num_k_heads)
63        .arg_u32(num_v_heads)
64        .arg_u32(k_dim)
65        .arg_u32(v_dim)
66        .arg_u32(qk_stride)
67        .arg_u32(v_stride)
68        .arg_u32(gb_stride)
69        .launch(stream)
70}
71
72/// WY-chunkwise 2-token GDN decode (2-pass algorithm).
73///
74/// Drop-in replacement for `gdn_decode_chunk2`. Computes both H^T @ k_t
75/// dot products in a single pass over H, then applies WY algebraic correction.
76/// 2 passes vs 3, reducing memory traffic by 33%.
77///
78/// Grid: (num_v_heads, batch, 1)  Block: (128, 1, 1)
79#[allow(clippy::too_many_arguments)]
80pub fn gdn_decode_wy2(
81    gpu: &dyn GpuBackend,
82    kernel: KernelHandle,
83    h_state: DevicePtr,
84    query: DevicePtr,
85    key: DevicePtr,
86    value: DevicePtr,
87    gate: DevicePtr,
88    beta: DevicePtr,
89    output: DevicePtr,
90    h_state_intermediate: DevicePtr,
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    // false = contiguous state bases indexed by (b*num_v_heads+vh);
100    // true  = device pointer TABLES, one entry per sequence. See
101    // `gdn_decode_wy4` for the full rationale — contiguous is only correct at
102    // batch_size==1 because the intermediate's pool stride is
103    // num_intermediates x h_state's.
104    state_is_table: bool,
105    stream: u64,
106) -> Result<()> {
107    // HARD guard, not debug_assert: this compiles out in release, which is
108    // exactly where the corruption would be silent.
109    anyhow::ensure!(
110        state_is_table || batch_size == 1,
111        "gdn_decode_wy2: contiguous state addressing is only valid at \
112         batch_size==1 (got {batch_size}) — the intermediate's pool stride is \
113         num_intermediates x h_state's, so sequence 1's Hi0 would land on \
114         sequence 0's Hi1. Stage pointer tables and pass state_is_table=true."
115    );
116    KernelLaunch::new(gpu, kernel)
117        .grid([num_v_heads, batch_size, 1])
118        .block([128, 1, 1])
119        .arg_ptr(h_state)
120        .arg_ptr(query)
121        .arg_ptr(key)
122        .arg_ptr(value)
123        .arg_ptr(gate)
124        .arg_ptr(beta)
125        .arg_ptr(output)
126        .arg_ptr(h_state_intermediate)
127        .arg_u32(batch_size)
128        .arg_u32(num_k_heads)
129        .arg_u32(num_v_heads)
130        .arg_u32(k_dim)
131        .arg_u32(v_dim)
132        .arg_u32(qk_stride)
133        .arg_u32(v_stride)
134        .arg_u32(gb_stride)
135        .arg_u32(u32::from(state_is_table))
136        .launch(stream)
137}
138
139/// WY-chunkwise 3-token GDN decode (2-pass algorithm).
140///
141/// Drop-in replacement for `gdn_decode_chunk3`. All 3 H^T @ k_t dot products
142/// computed in a single pass. 2 passes vs 4, reducing memory traffic by 50%.
143///
144/// Grid: (num_v_heads, batch, 1)  Block: (128, 1, 1)
145#[allow(clippy::too_many_arguments)]
146pub fn gdn_decode_wy3(
147    gpu: &dyn GpuBackend,
148    kernel: KernelHandle,
149    h_state: DevicePtr,
150    query: DevicePtr,
151    key: DevicePtr,
152    value: DevicePtr,
153    gate: DevicePtr,
154    beta: DevicePtr,
155    output: DevicePtr,
156    h_state_inter0: DevicePtr,
157    h_state_inter1: DevicePtr,
158    batch_size: u32,
159    num_k_heads: u32,
160    num_v_heads: u32,
161    k_dim: u32,
162    v_dim: u32,
163    qk_stride: u32,
164    v_stride: u32,
165    gb_stride: u32,
166    // false = contiguous state bases indexed by (b*num_v_heads+vh);
167    // true  = device pointer TABLES, one entry per sequence. See
168    // `gdn_decode_wy4` for the full rationale — contiguous is only correct at
169    // batch_size==1 because the intermediates' pool stride is
170    // num_intermediates x h_state's.
171    state_is_table: bool,
172    stream: u64,
173) -> Result<()> {
174    // HARD guard, not debug_assert: this compiles out in release, which is
175    // exactly where the corruption would be silent.
176    anyhow::ensure!(
177        state_is_table || batch_size == 1,
178        "gdn_decode_wy3: contiguous state addressing is only valid at \
179         batch_size==1 (got {batch_size}) — the intermediates' pool stride is \
180         num_intermediates x h_state's, so sequence 1's Hi0 would land on \
181         sequence 0's Hi1. Stage pointer tables and pass state_is_table=true."
182    );
183    KernelLaunch::new(gpu, kernel)
184        .grid([num_v_heads, batch_size, 1])
185        .block([128, 1, 1])
186        .arg_ptr(h_state)
187        .arg_ptr(query)
188        .arg_ptr(key)
189        .arg_ptr(value)
190        .arg_ptr(gate)
191        .arg_ptr(beta)
192        .arg_ptr(output)
193        .arg_ptr(h_state_inter0)
194        .arg_ptr(h_state_inter1)
195        .arg_u32(batch_size)
196        .arg_u32(num_k_heads)
197        .arg_u32(num_v_heads)
198        .arg_u32(k_dim)
199        .arg_u32(v_dim)
200        .arg_u32(qk_stride)
201        .arg_u32(v_stride)
202        .arg_u32(gb_stride)
203        .arg_u32(u32::from(state_is_table))
204        .launch(stream)
205}
spark_model/layers/ops/ssm_gdn_b.rs

spark_model/layers/ops/
ssm_gdn_b.rs
