spark_model/layers/ops/
ssm_gdn_batched.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Q12 Phase 2b: batched GDN prefill ops, hoisted from `ssm_gdn_a.rs`
4//! to keep that file under the 500-LoC file-size cap.
5
6#![allow(unused_imports)]
7
8use anyhow::Result;
9use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
10use spark_runtime::kernel_args::KernelLaunch;
11
12// ───────────────────────────────────────────────────────────────────────
13// Q12 Phase 2b: batched GDN prefill ops.
14//
15// Same-chunk-len batched variants for the four GDN kernels patched in
16// commits 37c44cc / 81d76fa / 5a93095 / 43f7c25. Caller is responsible
17// for:
18//   1. Uploading a device array `h_state_ptrs: [float*; batch_size]`
19//      containing the per-stream `SsmLayerState::h_state` pointers.
20//   2. Stacking QKV / gate / beta / output for all batched streams
21//      contiguously as `[batch_size * seq_len, conv_dim]` etc. (each
22//      stream's input lands at `b * seq_len * stride`).
23//   3. Ensuring all batched streams share the same `seq_len` (the
24//      scheduler `can_batch_prefill_only` gate enforces this).
25//
26// Validation status: kernels unvalidated against hardware. Wiring
27// these ops into `Qwen3SsmLayer::prefill_batched` is the next step;
28// they're committed first so the kernel patches and the Rust bindings
29// stay in lockstep.
30// ───────────────────────────────────────────────────────────────────────
31
32/// Batched WY32 persistent GDN prefill (uses gated_delta_rule_prefill_wy64_batched).
33#[allow(clippy::too_many_arguments)]
34pub fn gdn_prefill_persistent_smem_batched(
35    gpu: &dyn GpuBackend,
36    kernel: KernelHandle,
37    h_state_ptrs: DevicePtr,
38    query: DevicePtr,
39    key: DevicePtr,
40    value: DevicePtr,
41    gate: DevicePtr,
42    beta: DevicePtr,
43    output: DevicePtr,
44    batch_size: u32,
45    seq_len: u32,
46    num_k_heads: u32,
47    num_v_heads: u32,
48    k_dim: u32,
49    v_dim: u32,
50    qk_stride: u32,
51    v_stride: u32,
52    gb_stride: u32,
53    smem: u32,
54    stream: u64,
55) -> Result<()> {
56    KernelLaunch::new(gpu, kernel)
57        .grid([num_v_heads, batch_size, 1])
58        .block([128, 1, 1])
59        .shared_mem(smem)
60        .arg_ptr(h_state_ptrs)
61        .arg_ptr(query)
62        .arg_ptr(key)
63        .arg_ptr(value)
64        .arg_ptr(gate)
65        .arg_ptr(beta)
66        .arg_ptr(output)
67        .arg_u32(batch_size)
68        .arg_u32(seq_len)
69        .arg_u32(num_k_heads)
70        .arg_u32(num_v_heads)
71        .arg_u32(k_dim)
72        .arg_u32(v_dim)
73        .arg_u32(qk_stride)
74        .arg_u32(v_stride)
75        .arg_u32(gb_stride)
76        .launch(stream)
77}
78
79/// Batched persistent GDN prefill (uses gated_delta_rule_prefill_persistent_batched
80/// or gated_delta_rule_prefill_persistent_wy4_batched).
81#[allow(clippy::too_many_arguments)]
82pub fn gdn_prefill_persistent_batched(
83    gpu: &dyn GpuBackend,
84    kernel: KernelHandle,
85    h_state_ptrs: DevicePtr,
86    query: DevicePtr,
87    key: DevicePtr,
88    value: DevicePtr,
89    gate: DevicePtr,
90    beta: DevicePtr,
91    output: DevicePtr,
92    batch_size: u32,
93    seq_len: u32,
94    num_k_heads: u32,
95    num_v_heads: u32,
96    k_dim: u32,
97    v_dim: u32,
98    qk_stride: u32,
99    v_stride: u32,
100    gb_stride: u32,
101    stream: u64,
102) -> Result<()> {
103    let smem = k_dim * v_dim * 4 + 4 * k_dim * 4;
104    KernelLaunch::new(gpu, kernel)
105        .grid([num_v_heads, batch_size, 1])
106        .block([128, 1, 1])
107        .shared_mem(smem)
108        .arg_ptr(h_state_ptrs)
109        .arg_ptr(query)
110        .arg_ptr(key)
111        .arg_ptr(value)
112        .arg_ptr(gate)
113        .arg_ptr(beta)
114        .arg_ptr(output)
115        .arg_u32(batch_size)
116        .arg_u32(seq_len)
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        .launch(stream)
125}
126
127/// Batched split4 GDN prefill (uses gated_delta_rule_prefill_split4_batched).
128#[allow(clippy::too_many_arguments)]
129pub fn gdn_prefill_split4_batched(
130    gpu: &dyn GpuBackend,
131    kernel: KernelHandle,
132    h_state_ptrs: DevicePtr,
133    query: DevicePtr,
134    key: DevicePtr,
135    value: DevicePtr,
136    gate: DevicePtr,
137    beta: DevicePtr,
138    output: DevicePtr,
139    batch_size: u32,
140    seq_len: u32,
141    num_k_heads: u32,
142    num_v_heads: u32,
143    k_dim: u32,
144    v_dim: u32,
145    qk_stride: u32,
146    v_stride: u32,
147    gb_stride: u32,
148    stream: u64,
149) -> Result<()> {
150    KernelLaunch::new(gpu, kernel)
151        .grid([num_v_heads * 4, batch_size, 1])
152        .block([32, 1, 1])
153        .shared_mem(4 * k_dim * 4)
154        .arg_ptr(h_state_ptrs)
155        .arg_ptr(query)
156        .arg_ptr(key)
157        .arg_ptr(value)
158        .arg_ptr(gate)
159        .arg_ptr(beta)
160        .arg_ptr(output)
161        .arg_u32(batch_size)
162        .arg_u32(seq_len)
163        .arg_u32(num_k_heads)
164        .arg_u32(num_v_heads)
165        .arg_u32(k_dim)
166        .arg_u32(v_dim)
167        .arg_u32(qk_stride)
168        .arg_u32(v_stride)
169        .arg_u32(gb_stride)
170        .launch(stream)
171}
spark_model/layers/ops/ssm_gdn_batched.rs

spark_model/layers/ops/
ssm_gdn_batched.rs
