metrale_model_layers/layers/ops/
ssm_gdn_batched.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the batched GDN prefill kernels, which run
4//! several streams' prefill chunks in one launch.
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;
14
15// 2026-09-25: The caller's contract for every launcher in this file:
16//   1. `h_state_ptrs` is a device array of `batch_size` per-stream h-state
17//      pointers;
18//   2. Q/K/V, gate/beta and the output are stacked per stream: stream `b`
19//      starts at `b * seq_len * stride`;
20//   3. every stream has the same `seq_len`, since the kernels take one.
21
22/// 2026-09-25: Batched persistent prefill whose shared memory the caller
23/// computes (`gated_delta_rule_prefill_wy64_batched` or
24/// `gated_delta_rule_prefill_persistent_wy4_batched`).
25#[allow(clippy::too_many_arguments)]
26pub fn gdn_prefill_persistent_smem_batched(
27    gpu: &dyn GpuBackend,
28    kernel: KernelHandle,
29    h_state_ptrs: 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    smem: u32,
46    stream: u64,
47) -> Result<()> {
48    KernelLaunch::new(gpu, kernel)
49        .grid([num_v_heads, batch_size, 1])
50        .block([128, 1, 1])
51        .shared_mem(smem)
52        .arg_ptr(h_state_ptrs)
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_u32(batch_size)
60        .arg_u32(seq_len)
61        .arg_u32(num_k_heads)
62        .arg_u32(num_v_heads)
63        .arg_u32(k_dim)
64        .arg_u32(v_dim)
65        .arg_u32(qk_stride)
66        .arg_u32(v_stride)
67        .arg_u32(gb_stride)
68        .launch(stream)
69}
70
71/// 2026-09-25: Batched persistent prefill with H plus double-buffered k and q
72/// in shared memory (`gated_delta_rule_prefill_persistent_batched`).
73#[allow(clippy::too_many_arguments)]
74pub fn gdn_prefill_persistent_batched(
75    gpu: &dyn GpuBackend,
76    kernel: KernelHandle,
77    h_state_ptrs: DevicePtr,
78    query: DevicePtr,
79    key: DevicePtr,
80    value: DevicePtr,
81    gate: DevicePtr,
82    beta: DevicePtr,
83    output: DevicePtr,
84    batch_size: u32,
85    seq_len: 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    stream: u64,
94) -> Result<()> {
95    let smem = k_dim * v_dim * 4 + 4 * k_dim * 4;
96    KernelLaunch::new(gpu, kernel)
97        .grid([num_v_heads, batch_size, 1])
98        .block([128, 1, 1])
99        .shared_mem(smem)
100        .arg_ptr(h_state_ptrs)
101        .arg_ptr(query)
102        .arg_ptr(key)
103        .arg_ptr(value)
104        .arg_ptr(gate)
105        .arg_ptr(beta)
106        .arg_ptr(output)
107        .arg_u32(batch_size)
108        .arg_u32(seq_len)
109        .arg_u32(num_k_heads)
110        .arg_u32(num_v_heads)
111        .arg_u32(k_dim)
112        .arg_u32(v_dim)
113        .arg_u32(qk_stride)
114        .arg_u32(v_stride)
115        .arg_u32(gb_stride)
116        .launch(stream)
117}
118
119/// 2026-09-25: Batched split4 prefill
120/// (`gated_delta_rule_prefill_split4_batched`).
121#[allow(clippy::too_many_arguments)]
122pub fn gdn_prefill_split4_batched(
123    gpu: &dyn GpuBackend,
124    kernel: KernelHandle,
125    h_state_ptrs: DevicePtr,
126    query: DevicePtr,
127    key: DevicePtr,
128    value: DevicePtr,
129    gate: DevicePtr,
130    beta: DevicePtr,
131    output: DevicePtr,
132    batch_size: u32,
133    seq_len: u32,
134    num_k_heads: u32,
135    num_v_heads: u32,
136    k_dim: u32,
137    v_dim: u32,
138    qk_stride: u32,
139    v_stride: u32,
140    gb_stride: u32,
141    stream: u64,
142) -> Result<()> {
143    KernelLaunch::new(gpu, kernel)
144        .grid([num_v_heads * 4, batch_size, 1])
145        .block([32, 1, 1])
146        .shared_mem(4 * k_dim * 4)
147        .arg_ptr(h_state_ptrs)
148        .arg_ptr(query)
149        .arg_ptr(key)
150        .arg_ptr(value)
151        .arg_ptr(gate)
152        .arg_ptr(beta)
153        .arg_ptr(output)
154        .arg_u32(batch_size)
155        .arg_u32(seq_len)
156        .arg_u32(num_k_heads)
157        .arg_u32(num_v_heads)
158        .arg_u32(k_dim)
159        .arg_u32(v_dim)
160        .arg_u32(qk_stride)
161        .arg_u32(v_stride)
162        .arg_u32(gb_stride)
163        .launch(stream)
164}
metrale_model_layers/layers/ops/ssm_gdn_batched.rs

metrale_model_layers/layers/ops/
ssm_gdn_batched.rs
