| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | // |
| 4 | // Fused CPU vector kernels for Mamba decode-step hotspots: |
| 5 | // - causal_conv1d_update (depthwise 1-D conv state roll + compute) |
| 6 | // - selective_state_update (SSM recurrence, single-step) |
| 7 | |
| 8 | #pragma once |
| 9 | |
| 10 | #include "cpu_types.hpp" |
| 11 | #include <cmath> |
| 12 | #include <cstring> |
| 13 | #include <cstdint> |
| 14 | #include <algorithm> |
| 15 | |
| 16 | namespace mamba_cpu { |
| 17 | |
| 18 | // --------------------------------------------------------------------------- |
| 19 | // causal_conv1d_update — templated for native BF16/FP32 |
| 20 | // |
| 21 | // state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor. |
| 22 | // Explicit strides are passed so the kernel writes directly into the |
| 23 | // correct memory locations without making a contiguous copy of the full |
| 24 | // paged tensor (which was the source of the 34-41% direct_copy_kernel). |
| 25 | // |
| 26 | // stride_s_slot = state.stride(0) — between cache slots |
| 27 | // stride_s_dim = state.stride(1) — between conv_dim channels |
| 28 | // stride_s_state = state.stride(2) — between state elements |
| 29 | // |
| 30 | // When stride_s_state == 1 (contiguous), the memmove fast path is used. |
| 31 | // --------------------------------------------------------------------------- |
| 32 | template <typename scalar_t> |
| 33 | inline void causal_conv1d_update_kernel( |
| 34 | const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr, |
| 35 | int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state, |
| 36 | const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr, |
| 37 | scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs, |
| 38 | int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen, |
| 39 | int64_t width, int64_t state_len, bool do_silu) { |
| 40 | #pragma omp parallel for |
| 41 | for (int64_t b = 0; b < batch; ++b) { |
| 42 | int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b; |
| 43 | if (cache_idx == pad_slot_id) continue; |
| 44 | |
| 45 | for (int64_t t = 0; t < seqlen; ++t) { |
| 46 | const scalar_t* x_b = x_ptr + (b * dim * seqlen + t); |
| 47 | scalar_t* out_b = out_ptr + (b * dim * seqlen + t); |
| 48 | // Base of this slot in the (possibly non-contiguous) paged state |
| 49 | scalar_t* s_base = state_ptr + cache_idx * stride_s_slot; |
| 50 | |
| 51 | for (int64_t d = 0; d < dim; ++d) { |
| 52 | float x_val = static_cast<float>(x_b[d * seqlen]); |
| 53 | scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state |
| 54 | const scalar_t* w = weight_ptr + d * width; |
| 55 | |
| 56 | // Accumulate in float32 for precision |
| 57 | float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f; |
| 58 | for (int64_t k = 0; k < state_len; ++k) { |
| 59 | acc += static_cast<float>(w[k]) * |
| 60 | static_cast<float>(sd[k * stride_s_state]); |
| 61 | } |
| 62 | acc += static_cast<float>(w[state_len]) * x_val; |
| 63 | |
| 64 | // Shift state left and append new input. |
| 65 | // Use memmove when contiguous (stride==1); element loop otherwise. |
| 66 | if (stride_s_state == 1) { |
| 67 | if (state_len > 1) |
| 68 | std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t)); |
| 69 | if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val); |
| 70 | } else { |
| 71 | for (int64_t k = 0; k < state_len - 1; ++k) |
| 72 | sd[k * stride_s_state] = sd[(k + 1) * stride_s_state]; |
| 73 | if (state_len > 0) |
| 74 | sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val); |
| 75 | } |
| 76 | |
| 77 | if (do_silu) { |
| 78 | float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc)) |
| 79 | : std::exp(acc) / (1.0f + std::exp(acc)); |
| 80 | acc *= sigmoid; |
| 81 | } |
| 82 | out_b[d * seqlen] = static_cast<scalar_t>(acc); |
| 83 | } |
| 84 | } |
| 85 | } |
| 86 | } |
| 87 | |
| 88 | // --------------------------------------------------------------------------- |
| 89 | // selective_state_update |
| 90 | // |
| 91 | // Template parameters: |
| 92 | // state_t - dtype of ssm_state cache (typically BFloat16) |
| 93 | // input_t - dtype of x, B, C (typically BFloat16) |
| 94 | // out_t - dtype of output tensor (typically BFloat16) |
| 95 | // Write directly — no float32 intermediate buffer needed. |
| 96 | // |
| 97 | // A, D, dt_bias are accepted as const float* (they are always float32 |
| 98 | // model parameters in Mamba2). This eliminates the per-call float32→BF16 |
| 99 | // conversion and the .contiguous() materialisation of the broadcast-expand. |
| 100 | // |
| 101 | // dt is accepted as a (N, nheads) scalar-per-head tensor, not as the |
| 102 | // (N, nheads, head_dim) expansion, so no .contiguous() copy is needed. |
| 103 | // --------------------------------------------------------------------------- |
| 104 | template <typename state_t, typename input_t, typename out_t = float> |
| 105 | inline void selective_state_update_kernel( |
| 106 | state_t* __restrict__ state_ptr, int64_t stride_state_n, |
| 107 | int64_t stride_state_h, int64_t stride_state_d, |
| 108 | const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h, |
| 109 | // dt: (N, nheads) — scalar per head, NOT expanded to head_dim |
| 110 | const float* __restrict__ dt_ptr, int64_t stride_dt_n, |
| 111 | // A: (nheads,) float32 — scalar per head |
| 112 | const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr, |
| 113 | const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g, |
| 114 | // D: (nheads,) float32 — scalar per head (nullptr if not used) |
| 115 | const float* __restrict__ D_ptr, |
| 116 | // z: same shape as x (optional) |
| 117 | const input_t* __restrict__ z_ptr, |
| 118 | // dt_bias: (nheads,) float32 — scalar per head (nullptr if not used) |
| 119 | const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr, |
| 120 | int64_t stride_out_n, int64_t stride_out_h, |
| 121 | const int32_t* __restrict__ state_batch_indices, |
| 122 | const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id, |
| 123 | const int32_t* __restrict__ num_accepted_tokens, |
| 124 | const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads, |
| 125 | int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) { |
| 126 | using state_vec_t = vec_op::vec_t<state_t>; |
| 127 | using input_vec_t = vec_op::vec_t<input_t>; |
| 128 | constexpr int VEC_ELEM_NUM = 8; |
| 129 | |
| 130 | int64_t nheads_per_group = nheads / ngroups; |
| 131 | |
| 132 | for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) { |
| 133 | int64_t bos, seq_len; |
| 134 | if (cu_seqlens != nullptr) { |
| 135 | bos = cu_seqlens[seq_idx]; |
| 136 | seq_len = cu_seqlens[seq_idx + 1] - bos; |
| 137 | } else { |
| 138 | bos = seq_idx; |
| 139 | seq_len = 1; |
| 140 | } |
| 141 | |
| 142 | int64_t state_read_idx = (state_batch_indices != nullptr) |
| 143 | ? state_batch_indices[seq_idx] |
| 144 | : seq_idx; |
| 145 | if (state_read_idx == null_block_id) continue; |
| 146 | |
| 147 | int64_t state_write_idx = (num_accepted_tokens == nullptr) |
| 148 | ? ((dst_state_batch_indices != nullptr) |
| 149 | ? dst_state_batch_indices[seq_idx] |
| 150 | : state_read_idx) |
| 151 | : -1; |
| 152 | |
| 153 | state_t* s = state_ptr + state_read_idx * stride_state_n; |
| 154 | |
| 155 | for (int64_t t = 0; t < seq_len; ++t) { |
| 156 | int64_t token_idx = bos + t; |
| 157 | const input_t* x_tok = x_ptr + token_idx * stride_x_n; |
| 158 | // dt: (N, nheads) — one float per head per token |
| 159 | const float* dt_tok = dt_ptr + token_idx * stride_dt_n; |
| 160 | const input_t* B_tok = B_ptr + token_idx * stride_BC_n; |
| 161 | const input_t* C_tok = C_ptr + token_idx * stride_BC_n; |
| 162 | out_t* out_tok = out_ptr + token_idx * stride_out_n; |
| 163 | |
| 164 | #pragma omp parallel for |
| 165 | for (int64_t h = 0; h < nheads; ++h) { |
| 166 | int64_t g = h / nheads_per_group; |
| 167 | const input_t* x_h = x_tok + h * stride_x_h; |
| 168 | const input_t* B_g = B_tok + g * stride_BC_g; |
| 169 | const input_t* C_g = C_tok + g * stride_BC_g; |
| 170 | out_t* out_h = out_tok + h * stride_out_h; |
| 171 | state_t* s_h = s + h * stride_state_h; |
| 172 | |
| 173 | // Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing |
| 174 | float dt_val = dt_tok[h]; |
| 175 | if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h]; |
| 176 | if (dt_softplus) { |
| 177 | dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val; |
| 178 | } |
| 179 | const float A_val = A_ptr[h]; // scalar: same for all dim, dstate |
| 180 | const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f; |
| 181 | |
| 182 | const input_t* z_h = |
| 183 | (z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h |
| 184 | : nullptr; |
| 185 | |
| 186 | vec_op::FP32Vec8 dt_vec(dt_val); |
| 187 | // dA = exp(A * dt): A and dt are SCALARS per head, so compute once |
| 188 | // and broadcast. This saves 7 redundant std::exp() calls that |
| 189 | // FP32Vec8::exp() would otherwise make on the broadcast vector. |
| 190 | const float dA_scalar = std::exp(A_val * dt_val); |
| 191 | vec_op::FP32Vec8 dA(dA_scalar); // broadcast |
| 192 | |
| 193 | for (int64_t d = 0; d < dim; ++d) { |
| 194 | float x_val = static_cast<float>(x_h[d]); |
| 195 | |
| 196 | vec_op::FP32Vec8 out_vec(0.0f); |
| 197 | state_t* s_hd = s_h + d * stride_state_d; |
| 198 | const input_t* B_g_base = B_g; |
| 199 | const input_t* C_g_base = C_g; |
| 200 | |
| 201 | vec_op::FP32Vec8 x_vec(x_val); |
| 202 | // dBx = B * x * dt — same dA for all dstate (A is scalar) |
| 203 | // s_new = s * dA + B * x * dt |
| 204 | |
| 205 | int64_t n = 0; |
| 206 | for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) { |
| 207 | vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n))); |
| 208 | vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n))); |
| 209 | vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n))); |
| 210 | |
| 211 | vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec; |
| 212 | vec_op::FP32Vec8 s_new = s_v * dA + dBx; |
| 213 | |
| 214 | state_vec_t(s_new).save(s_hd + n); |
| 215 | out_vec = out_vec + s_new * C_v; |
| 216 | } |
| 217 | |
| 218 | float out_val = out_vec.reduce_sum(); |
| 219 | for (; n < dstate; ++n) { |
| 220 | // Reuse dA_scalar computed once per head — no exp() re-call |
| 221 | float dBx = static_cast<float>(B_g[n]) * x_val * dt_val; |
| 222 | float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx; |
| 223 | s_hd[n] = static_cast<state_t>(s_new); |
| 224 | out_val += s_new * static_cast<float>(C_g[n]); |
| 225 | } |
| 226 | |
| 227 | if (D_ptr != nullptr) out_val += x_val * D_val; |
| 228 | if (z_h != nullptr) { |
| 229 | float z_val = static_cast<float>(z_h[d]); |
| 230 | float sigmoid = (z_val >= 0) |
| 231 | ? 1.0f / (1.0f + std::exp(-z_val)) |
| 232 | : std::exp(z_val) / (1.0f + std::exp(z_val)); |
| 233 | out_val *= z_val * sigmoid; |
| 234 | } |
| 235 | out_h[d] = static_cast<out_t>(out_val); |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | if (num_accepted_tokens != nullptr && |
| 240 | dst_state_batch_indices != nullptr) { |
| 241 | int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t]; |
| 242 | if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) { |
| 243 | state_t* dst_s = state_ptr + token_dst_idx * stride_state_n; |
| 244 | std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t)); |
| 245 | } |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | if (num_accepted_tokens == nullptr && state_write_idx != null_block_id && |
| 250 | state_write_idx != state_read_idx) { |
| 251 | state_t* dst_s = state_ptr + state_write_idx * stride_state_n; |
| 252 | std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t)); |
| 253 | } |
| 254 | } |
| 255 | } |
| 256 | |
| 257 | // --------------------------------------------------------------------------- |
| 258 | // mamba_chunk_scan_fwd |
| 259 | // |
| 260 | // Prefill SSM recurrence for Mamba2 / SSD models. |
| 261 | // |
| 262 | // Key difference from selective_state_update_kernel (decode path): |
| 263 | // - #pragma omp parallel for collapse(2) is OUTSIDE the time loop. |
| 264 | // Each thread owns a (batch, head) slice and runs the entire token |
| 265 | // sequence without any per-token OpenMP synchronisation overhead. |
| 266 | // For seqlen=256, this eliminates 256 thread-barrier launches per batch. |
| 267 | // |
| 268 | // `dt` arrives already processed (float32, after bias + softplus + clamp) |
| 269 | // to keep this kernel simple. Preprocessing is done in the Python wrapper. |
| 270 | // |
| 271 | // `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output |
| 272 | // tensor, pre-initialised by the caller (zero or from initial_states). |
| 273 | // Each (b, h) slice is private to exactly one thread via collapse(2), so |
| 274 | // there are no write conflicts. |
| 275 | // |
| 276 | // D is treated as a scalar per head ([nheads] float32). |
| 277 | // --------------------------------------------------------------------------- |
| 278 | template <typename input_t> |
| 279 | inline void mamba_chunk_scan_fwd_kernel( |
| 280 | float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32 |
| 281 | const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim] |
| 282 | const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed) |
| 283 | const float* __restrict__ A_ptr, // [nheads] f32 |
| 284 | const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate] |
| 285 | const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate] |
| 286 | const float* __restrict__ D_ptr, // [nheads] f32 (nullable) |
| 287 | const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable) |
| 288 | input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim] |
| 289 | const int32_t* __restrict__ cu_seqlens, // [batch+1] int32 |
| 290 | int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim, |
| 291 | int64_t dstate) { |
| 292 | using input_vec_t = vec_op::vec_t<input_t>; |
| 293 | constexpr int VEC_ELEM_NUM = 8; |
| 294 | |
| 295 | const int64_t nheads_per_group = nheads / ngroups; |
| 296 | // states layout: [batch, nheads, headdim, dstate] contiguous (caller |
| 297 | // guarantee) |
| 298 | const int64_t stride_s_b = nheads * headdim * dstate; |
| 299 | const int64_t stride_s_h = headdim * dstate; |
| 300 | // stride_s_d = dstate, stride_s_n = 1 |
| 301 | |
| 302 | #pragma omp parallel for collapse(2) schedule(static) |
| 303 | for (int64_t b = 0; b < batch; ++b) { |
| 304 | for (int64_t h = 0; h < nheads; ++h) { |
| 305 | const int64_t seq_start = cu_seqlens[b]; |
| 306 | const int64_t seq_end = cu_seqlens[b + 1]; |
| 307 | const int64_t g = h / nheads_per_group; |
| 308 | |
| 309 | const float A_val = A_ptr[h]; |
| 310 | const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f; |
| 311 | |
| 312 | // Working state slice: states[b, h, :, :] — float32, headdim * dstate. |
| 313 | // Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB). |
| 314 | float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h; |
| 315 | |
| 316 | for (int64_t t = seq_start; t < seq_end; ++t) { |
| 317 | const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim; |
| 318 | const float* dt_h = dt_ptr + t * nheads + h; |
| 319 | const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate; |
| 320 | const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate; |
| 321 | const input_t* z_h = (z_ptr != nullptr) |
| 322 | ? z_ptr + t * nheads * headdim + h * headdim |
| 323 | : nullptr; |
| 324 | input_t* out_h = out_ptr + t * nheads * headdim + h * headdim; |
| 325 | |
| 326 | const float dt_val = *dt_h; |
| 327 | const float dA_val = std::exp(A_val * dt_val); |
| 328 | const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar |
| 329 | const vec_op::FP32Vec8 dt_vec(dt_val); |
| 330 | |
| 331 | for (int64_t d = 0; d < headdim; ++d) { |
| 332 | const float x_val = static_cast<float>(x_h[d]); |
| 333 | float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32 |
| 334 | |
| 335 | // Vectorised SSM update + readout over dstate: |
| 336 | // s_new = s * dA + x * dt * B |
| 337 | // y += s_new * C |
| 338 | int64_t n = 0; |
| 339 | vec_op::FP32Vec8 y_vec(0.0f); |
| 340 | const vec_op::FP32Vec8 x_vec(x_val); |
| 341 | |
| 342 | for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) { |
| 343 | const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n))); |
| 344 | const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n))); |
| 345 | const vec_op::FP32Vec8 s_v(s_bhd + n); |
| 346 | |
| 347 | const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v; |
| 348 | s_new.save(s_bhd + n); |
| 349 | y_vec = y_vec + s_new * C_v; |
| 350 | } |
| 351 | |
| 352 | float y_val = y_vec.reduce_sum(); |
| 353 | |
| 354 | // Scalar tail for remaining dstate elements |
| 355 | for (; n < dstate; ++n) { |
| 356 | const float B_n = static_cast<float>(B_g[n]); |
| 357 | const float C_n = static_cast<float>(C_g[n]); |
| 358 | const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n; |
| 359 | s_bhd[n] = s_new; |
| 360 | y_val += s_new * C_n; |
| 361 | } |
| 362 | |
| 363 | // D skip connection (scalar per head) |
| 364 | if (D_ptr != nullptr) y_val += x_val * D_val; |
| 365 | |
| 366 | // z gating: out = y * z * sigmoid(z) (SiLU) |
| 367 | if (z_h != nullptr) { |
| 368 | const float z_val = static_cast<float>(z_h[d]); |
| 369 | const float sigmoid = |
| 370 | (z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val)) |
| 371 | : std::exp(z_val) / (1.0f + std::exp(z_val)); |
| 372 | y_val *= z_val * sigmoid; |
| 373 | } |
| 374 | |
| 375 | out_h[d] = static_cast<input_t>(y_val); |
| 376 | } |
| 377 | } |
| 378 | } |
| 379 | } |
| 380 | } |
| 381 | |
| 382 | } // namespace mamba_cpu |