spark_model/layers/ops/
ssm_ssd.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Mamba-2 SSD chunked-scan launchers (cumsum / CB bmm / fused scan) and their
4//! tiling constants. Split from `ssm_mamba.rs` (500-LoC cap).
5
6#![allow(unused_imports)]
7
8use anyhow::Result;
9use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
10use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
11
12use super::*;
13
14/// SSD chunk length (must match `SSD_L` in the kernel).
15pub const SSD_L: u32 = 64;
16/// head_dim rows per SSD scan block (must match `SSD_PT` in the kernel).
17pub const SSD_PT: u32 = 64;
18
19/// K1: per-chunk dt (softplus+clamp) and inclusive cumsum of the log-decay.
20#[allow(clippy::too_many_arguments)]
21pub fn mamba2_ssd_cumsum(
22    gpu: &dyn GpuBackend,
23    kernel: KernelHandle,
24    dt_raw: DevicePtr,
25    a_log: DevicePtr,
26    dt_bias: DevicePtr,
27    dt_out: DevicePtr,
28    da_cs: DevicePtr,
29    seq_len: u32,
30    num_heads: u32,
31    nchunks: u32,
32    batch_size: u32,
33    dt_stride: u32,
34    dt_min: f32,
35    dt_max: f32,
36    stream: u64,
37) -> Result<()> {
38    KernelLaunch::new(gpu, kernel)
39        .grid([nchunks, num_heads, batch_size])
40        .block([SSD_L, 1, 1])
41        .arg_ptr(dt_raw)
42        .arg_ptr(a_log)
43        .arg_ptr(dt_bias)
44        .arg_ptr(dt_out)
45        .arg_ptr(da_cs)
46        .arg_u32(seq_len)
47        .arg_u32(num_heads)
48        .arg_u32(nchunks)
49        .arg_u32(dt_stride)
50        .arg_f32(dt_min)
51        .arg_f32(dt_max)
52        .launch(stream)
53}
54
55/// K2: `CB[c][g][t][s] = C_t . B_s` (raw, fp32).
56#[allow(clippy::too_many_arguments)]
57pub fn mamba2_ssd_bmm(
58    gpu: &dyn GpuBackend,
59    kernel: KernelHandle,
60    b_proj: DevicePtr,
61    c_proj: DevicePtr,
62    cb: DevicePtr,
63    seq_len: u32,
64    nchunks: u32,
65    n_groups: u32,
66    state_size: u32,
67    batch_size: u32,
68    bc_stride: u32,
69    stream: u64,
70) -> Result<()> {
71    // smem: sC[L][N] + sB[L][N] bf16
72    let smem = 2 * SSD_L * state_size * 2;
73    KernelLaunch::new(gpu, kernel)
74        .grid([nchunks, n_groups, batch_size])
75        .block([128, 1, 1])
76        .shared_mem(smem)
77        .arg_ptr(b_proj)
78        .arg_ptr(c_proj)
79        .arg_ptr(cb)
80        .arg_u32(seq_len)
81        .arg_u32(nchunks)
82        .arg_u32(n_groups)
83        .arg_u32(state_size)
84        .arg_u32(bc_stride)
85        .launch(stream)
86}
87
88/// Largest dynamic shared-memory block a kernel may opt into on the GB10 target.
89/// Measured on the device (sm_121):
90/// `CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN = 101376`
91/// (static default is only 49152; per-SM total is 102400).
92pub const MAX_DYNAMIC_SMEM: u32 = 101_376;
93
94/// Dynamic shared memory [`mamba2_ssd_scan`] requests for a given SSM state size.
95/// Grows ~linearly in `state_size`, so a checkpoint with a large SSM state can
96/// exceed [`MAX_DYNAMIC_SMEM`]. Kept next to the launch so the two cannot drift.
97pub fn ssd_scan_smem(state_size: u32) -> u32 {
98    SSD_PT * (state_size + 1) * 4
99        + 2 * SSD_L * state_size * 2
100        + 2 * SSD_L * state_size * 2
101        + 2 * SSD_L * SSD_PT * 2
102        + 2 * SSD_L * 4
103        + 2 * SSD_L * 4
104}
105
106/// Whether the SSD chunked scan physically fits for this `state_size`.
107///
108/// Callers MUST check this before selecting the SSD path: exceeding the limit
109/// makes `cuFuncSetAttribute(MAX_DYNAMIC_SHARED)` fail with
110/// `CUDA_ERROR_INVALID_VALUE`, which aborts the layer instead of degrading.
111/// Measured: Puzzle-75B `state_size=96` -> 91392 (fits); Nemotron Nano-30B
112/// `state_size=128` -> 115968 (does NOT fit, and made the model unservable).
113pub fn ssd_scan_fits(state_size: u32) -> bool {
114    ssd_scan_smem(state_size) <= MAX_DYNAMIC_SMEM
115}
116
117/// K3: fused chunk_state + state_passing + chunk_scan (h0 stays in shared memory,
118/// so the per-chunk `states` tensor vLLM round-trips through DRAM never exists).
119#[allow(clippy::too_many_arguments)]
120pub fn mamba2_ssd_scan(
121    gpu: &dyn GpuBackend,
122    kernel: KernelHandle,
123    h_state: DevicePtr,
124    x: DevicePtr,
125    b_proj: DevicePtr,
126    c_proj: DevicePtr,
127    d_param: DevicePtr,
128    dt_f32: DevicePtr,
129    da_cs: DevicePtr,
130    cb: DevicePtr,
131    output: DevicePtr,
132    seq_len: u32,
133    num_heads: u32,
134    head_dim: u32,
135    state_size: u32,
136    n_groups: u32,
137    nchunks: u32,
138    batch_size: u32,
139    x_stride: u32,
140    bc_stride: u32,
141    y_stride: u32,
142    stream: u64,
143) -> Result<()> {
144    // sH[PT][N+1] f32 | double-buffered streaming tiles: sB[2][L][N] |
145    // sCM[2][L][N] | sX[2][L][PT] bf16 | sdA[2][L] + sdt[2][L] f32.
146    // (sHb and sXt were dropped -- derived on the fly; see the kernel.)
147    let smem = ssd_scan_smem(state_size);
148    KernelLaunch::new(gpu, kernel)
149        .grid([num_heads, head_dim / SSD_PT, batch_size])
150        .block([512, 1, 1]) // 16 warps, 2 warp-tasks each (see kernel)
151        .shared_mem(smem)
152        .arg_ptr(h_state)
153        .arg_ptr(x)
154        .arg_ptr(b_proj)
155        .arg_ptr(c_proj)
156        .arg_ptr(d_param)
157        .arg_ptr(dt_f32)
158        .arg_ptr(da_cs)
159        .arg_ptr(cb)
160        .arg_ptr(output)
161        .arg_u32(seq_len)
162        .arg_u32(num_heads)
163        .arg_u32(head_dim)
164        .arg_u32(state_size)
165        .arg_u32(n_groups)
166        .arg_u32(nchunks)
167        .arg_u32(x_stride)
168        .arg_u32(bc_stride)
169        .arg_u32(y_stride)
170        .launch(stream)
171}
spark_model/layers/ops/ssm_ssd.rs

spark_model/layers/ops/
ssm_ssd.rs
