metrale_model_layers/layers/ops/
ssm_ssd.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Mamba-2 SSD chunked-scan launchers (cumsum, CB bmm, fused scan)
4//! and their tiling and shared-memory constants.
5//!
6//! Owner: model-layers ops.
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 super::*;
16
17/// 2026-09-25: SSD chunk length. Must equal `SSD_L` in `mamba2_ssd_chunk.cu`.
18pub const SSD_L: u32 = 64;
19/// 2026-09-25: head_dim rows per SSD scan block. Must equal `SSD_PT` in
20/// `mamba2_ssd_chunk.cu`.
21pub const SSD_PT: u32 = 64;
22
23/// 2026-09-25: K1: per-chunk dt (softplus and clamp) and the inclusive cumsum
24/// of the log-decay.
25#[allow(clippy::too_many_arguments)]
26pub fn mamba2_ssd_cumsum(
27    gpu: &dyn GpuBackend,
28    kernel: KernelHandle,
29    dt_raw: DevicePtr,
30    a_log: DevicePtr,
31    dt_bias: DevicePtr,
32    dt_out: DevicePtr,
33    da_cs: DevicePtr,
34    seq_len: u32,
35    num_heads: u32,
36    nchunks: u32,
37    batch_size: u32,
38    dt_stride: u32,
39    dt_min: f32,
40    dt_max: f32,
41    stream: u64,
42) -> Result<()> {
43    KernelLaunch::new(gpu, kernel)
44        .grid([nchunks, num_heads, batch_size])
45        .block([SSD_L, 1, 1])
46        .arg_ptr(dt_raw)
47        .arg_ptr(a_log)
48        .arg_ptr(dt_bias)
49        .arg_ptr(dt_out)
50        .arg_ptr(da_cs)
51        .arg_u32(seq_len)
52        .arg_u32(num_heads)
53        .arg_u32(nchunks)
54        .arg_u32(dt_stride)
55        .arg_f32(dt_min)
56        .arg_f32(dt_max)
57        .launch(stream)
58}
59
60/// 2026-09-25: K2: `CB[c][g][t][s] = C_t . B_s` (raw, FP32).
61#[allow(clippy::too_many_arguments)]
62pub fn mamba2_ssd_bmm(
63    gpu: &dyn GpuBackend,
64    kernel: KernelHandle,
65    b_proj: DevicePtr,
66    c_proj: DevicePtr,
67    cb: DevicePtr,
68    seq_len: u32,
69    nchunks: u32,
70    n_groups: u32,
71    state_size: u32,
72    batch_size: u32,
73    bc_stride: u32,
74    stream: u64,
75) -> Result<()> {
76    // 2026-09-25: smem: sC[L][N] + sB[L][N] BF16.
77    let smem = 2 * SSD_L * state_size * 2;
78    KernelLaunch::new(gpu, kernel)
79        .grid([nchunks, n_groups, batch_size])
80        .block([128, 1, 1])
81        .shared_mem(smem)
82        .arg_ptr(b_proj)
83        .arg_ptr(c_proj)
84        .arg_ptr(cb)
85        .arg_u32(seq_len)
86        .arg_u32(nchunks)
87        .arg_u32(n_groups)
88        .arg_u32(state_size)
89        .arg_u32(bc_stride)
90        .launch(stream)
91}
92
93/// 2026-09-25: Largest dynamic shared memory, in bytes, one block may opt into
94/// on the GB10 target (sm_121).
95pub const MAX_DYNAMIC_SMEM: u32 = 101_376;
96
97/// 2026-09-25: Dynamic shared memory [`mamba2_ssd_scan`] requests for a given
98/// SSM state size. It grows linearly in `state_size`, so a large SSM state can
99/// exceed [`MAX_DYNAMIC_SMEM`]. The launch calls this function, so the two
100/// cannot drift.
101pub fn ssd_scan_smem(state_size: u32) -> u32 {
102    SSD_PT * (state_size + 1) * 4
103        + 2 * SSD_L * state_size * 2
104        + 2 * SSD_L * state_size * 2
105        + 2 * SSD_L * SSD_PT * 2
106        + 2 * SSD_L * 4
107        + 2 * SSD_L * 4
108}
109
110/// 2026-09-25: Whether the SSD chunked scan fits in shared memory for this
111/// `state_size`.
112///
113/// Callers check this before selecting the SSD path. Over the limit the
114/// launch's `cuFuncSetAttribute(MAX_DYNAMIC_SHARED)` fails and the launch
115/// returns that error; nothing falls back. `state_size = 96` needs 91,392 B
116/// and fits; `state_size = 128` needs 115,968 B and does not.
117pub fn ssd_scan_fits(state_size: u32) -> bool {
118    ssd_scan_smem(state_size) <= MAX_DYNAMIC_SMEM
119}
120
121/// 2026-09-25: K3: chunk_state, state_passing and chunk_scan fused in one
122/// kernel. The running state h0 stays in shared memory, so no per-chunk state
123/// tensor is written to global memory.
124#[allow(clippy::too_many_arguments)]
125pub fn mamba2_ssd_scan(
126    gpu: &dyn GpuBackend,
127    kernel: KernelHandle,
128    h_state: DevicePtr,
129    x: DevicePtr,
130    b_proj: DevicePtr,
131    c_proj: DevicePtr,
132    d_param: DevicePtr,
133    dt_f32: DevicePtr,
134    da_cs: DevicePtr,
135    cb: DevicePtr,
136    output: DevicePtr,
137    seq_len: u32,
138    num_heads: u32,
139    head_dim: u32,
140    state_size: u32,
141    n_groups: u32,
142    nchunks: u32,
143    batch_size: u32,
144    x_stride: u32,
145    bc_stride: u32,
146    y_stride: u32,
147    stream: u64,
148) -> Result<()> {
149    // 2026-09-25: sH[PT][N+1] FP32 | double-buffered streaming tiles: sB[2][L][N] |
150    // sCM[2][L][N] | sX[2][L][PT] BF16 | sdA[2][L] + sdt[2][L] FP32.
151    let smem = ssd_scan_smem(state_size);
152    KernelLaunch::new(gpu, kernel)
153        .grid([num_heads, head_dim / SSD_PT, batch_size])
154        .block([512, 1, 1]) // 2026-09-25: 16 warps, each owning two n-tiles (see kernel).
155        .shared_mem(smem)
156        .arg_ptr(h_state)
157        .arg_ptr(x)
158        .arg_ptr(b_proj)
159        .arg_ptr(c_proj)
160        .arg_ptr(d_param)
161        .arg_ptr(dt_f32)
162        .arg_ptr(da_cs)
163        .arg_ptr(cb)
164        .arg_ptr(output)
165        .arg_u32(seq_len)
166        .arg_u32(num_heads)
167        .arg_u32(head_dim)
168        .arg_u32(state_size)
169        .arg_u32(n_groups)
170        .arg_u32(nchunks)
171        .arg_u32(x_stride)
172        .arg_u32(bc_stride)
173        .arg_u32(y_stride)
174        .launch(stream)
175}
metrale_model_layers/layers/ops/ssm_ssd.rs

metrale_model_layers/layers/ops/
ssm_ssd.rs
