| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | // |
| 4 | // CPU at::Tensor wrappers for Mamba decode-step kernels defined in |
| 5 | // mamba_kernels.hpp. |
| 6 | |
| 7 | #include "cpu/mamba_kernels.hpp" |
| 8 | |
| 9 | #include <ATen/ATen.h> |
| 10 | #include <torch/library.h> |
| 11 | #include <c10/util/Optional.h> |
| 12 | |
| 13 | #include "cpu_types.hpp" |
| 14 | |
| 15 | // --------------------------------------------------------------------------- |
| 16 | // causal_conv1d_update |
| 17 | // --------------------------------------------------------------------------- |
| 18 | at::Tensor causal_conv1d_update_cpu_impl( |
| 19 | at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight, |
| 20 | const c10::optional<at::Tensor>& bias, |
| 21 | const c10::optional<std::string>& activation, |
| 22 | const c10::optional<at::Tensor>& conv_state_indices, |
| 23 | const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) { |
| 24 | bool do_silu = false; |
| 25 | if (activation.has_value()) { |
| 26 | const std::string& act = activation.value(); |
| 27 | do_silu = (act == "silu" || act == "swish"); |
| 28 | } |
| 29 | |
| 30 | at::ScalarType dtype = x.scalar_type(); |
| 31 | |
| 32 | // Input x: contiguous in native dtype. |
| 33 | at::Tensor x_c = x.is_contiguous() ? x : x.contiguous(); |
| 34 | |
| 35 | // conv_state: NEVER copy the full paged tensor just for layout reasons. |
| 36 | // If the dtype matches we work directly on conv_state (contiguous or not) |
| 37 | // by extracting strides and passing them to the kernel. |
| 38 | // Only a dtype-conversion copy is made when types differ (rare for BF16). |
| 39 | bool state_type_ok = (conv_state.scalar_type() == dtype); |
| 40 | at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype); |
| 41 | // state_c and conv_state may be non-contiguous — that is intentional. |
| 42 | |
| 43 | // Weight: coerce to same dtype if needed (should match in practice) |
| 44 | at::Tensor w_c = |
| 45 | (weight.scalar_type() != dtype) |
| 46 | ? weight.to(dtype).contiguous() |
| 47 | : (weight.is_contiguous() ? weight : weight.contiguous()); |
| 48 | |
| 49 | // Bias stays float32 (small scalar, used only for fp32 accumulation) |
| 50 | at::Tensor bias_f32; |
| 51 | if (bias.has_value() && bias.value().defined()) |
| 52 | bias_f32 = bias.value().to(at::kFloat).contiguous(); |
| 53 | |
| 54 | int64_t batch = x_c.size(0); |
| 55 | int64_t dim = x_c.size(1); |
| 56 | int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1; |
| 57 | int64_t width = w_c.size(1); |
| 58 | int64_t state_len = state_c.size(2); |
| 59 | |
| 60 | // Extract strides — works for contiguous AND non-contiguous (transposed) |
| 61 | // state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in |
| 62 | // contiguous) stride(1): between conv channels (dim stride) stride(2): |
| 63 | // between state elements (=1 when contiguous, =dim when transposed) |
| 64 | int64_t stride_s_slot = state_c.stride(0); |
| 65 | int64_t stride_s_dim = state_c.stride(1); |
| 66 | int64_t stride_s_state = state_c.stride(2); |
| 67 | |
| 68 | at::Tensor out = x_c.clone(); // native dtype, no float32 alloc |
| 69 | |
| 70 | const int32_t* cache_idx_ptr = nullptr; |
| 71 | at::Tensor cache_idx_int; |
| 72 | if (conv_state_indices.has_value()) { |
| 73 | cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous(); |
| 74 | cache_idx_ptr = cache_idx_int.data_ptr<int32_t>(); |
| 75 | } |
| 76 | |
| 77 | VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] { |
| 78 | mamba_cpu::causal_conv1d_update_kernel<scalar_t>( |
| 79 | x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot, |
| 80 | stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(), |
| 81 | bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr, |
| 82 | out.data_ptr<scalar_t>(), cache_idx_ptr, |
| 83 | static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len, |
| 84 | do_silu); |
| 85 | }); |
| 86 | |
| 87 | // Write back only when a type-conversion copy was made. |
| 88 | // Layout-only non-contiguity is handled via strides above — no copy needed. |
| 89 | if (!state_type_ok) conv_state.copy_(state_c); |
| 90 | |
| 91 | return out; |
| 92 | } |
| 93 | |
| 94 | // --------------------------------------------------------------------------- |
| 95 | // selective_state_update |
| 96 | // --------------------------------------------------------------------------- |
| 97 | void selective_state_update_cpu_impl( |
| 98 | at::Tensor& state, // (nstates, nheads, dim, dstate) |
| 99 | const at::Tensor& x, // (N, nheads, dim) |
| 100 | const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B, |
| 101 | const at::Tensor& C, const c10::optional<at::Tensor>& D, |
| 102 | const c10::optional<at::Tensor>& z, |
| 103 | const c10::optional<at::Tensor>& dt_bias, bool dt_softplus, |
| 104 | const c10::optional<at::Tensor>& state_batch_indices, |
| 105 | const c10::optional<at::Tensor>& dst_state_batch_indices, |
| 106 | int64_t null_block_id, at::Tensor& out, |
| 107 | const c10::optional<at::Tensor>& num_accepted_tokens, |
| 108 | const c10::optional<at::Tensor>& cu_seqlens) { |
| 109 | at::ScalarType state_type = state.scalar_type(); |
| 110 | at::ScalarType input_type = x.scalar_type(); |
| 111 | |
| 112 | // x, B, C must be contiguous and match input_type |
| 113 | auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor { |
| 114 | at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t; |
| 115 | return r.is_contiguous() ? r : r.contiguous(); |
| 116 | }; |
| 117 | at::Tensor x_in = ensure_input(x); |
| 118 | at::Tensor B_in = ensure_input(B); |
| 119 | at::Tensor C_in = ensure_input(C); |
| 120 | at::Tensor z_in; |
| 121 | if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value()); |
| 122 | |
| 123 | // A, D, dt_bias are float32 model parameters that arrive here as expanded |
| 124 | // tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0). |
| 125 | // We need just the scalar value per head as a (nheads,) 1-D array so that |
| 126 | // A_ptr[h] in the kernel correctly reads head h's value. |
| 127 | // |
| 128 | // Strategy: peel trailing expanded (stride=0) dims via .select(), which is |
| 129 | // a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0) |
| 130 | // → .select(2,0) → (nheads, head_dim) strides (1,0) |
| 131 | // → .select(1,0) → (nheads,) stride (1,) ← contiguous, free. |
| 132 | // No allocation, no type conversion (A is already float32). |
| 133 | auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor { |
| 134 | at::Tensor r = t; |
| 135 | // Peel trailing dimensions that are broadcast (stride=0 or size=1) |
| 136 | while (r.dim() > 1) r = r.select(r.dim() - 1, 0); |
| 137 | if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat); |
| 138 | return r.is_contiguous() ? r : r.contiguous(); |
| 139 | }; |
| 140 | |
| 141 | at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32 |
| 142 | at::Tensor D_f32, dt_bias_f32; |
| 143 | if (D.has_value() && D.value().defined()) |
| 144 | D_f32 = to_per_head_1d_f32(D.value()); |
| 145 | if (dt_bias.has_value() && dt_bias.value().defined()) |
| 146 | dt_bias_f32 = to_per_head_1d_f32(dt_bias.value()); |
| 147 | |
| 148 | // dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE |
| 149 | // the type conversion so we convert head_dim x fewer elements. |
| 150 | at::Tensor dt_f32; |
| 151 | { |
| 152 | // If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2, |
| 153 | // take a zero-copy view of index 0 along that dim first. |
| 154 | at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads) |
| 155 | at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2; |
| 156 | dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous(); |
| 157 | } |
| 158 | |
| 159 | int64_t nheads = state.size(1); |
| 160 | int64_t dim = state.size(2); |
| 161 | int64_t dstate = state.size(3); |
| 162 | int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined()) |
| 163 | ? cu_seqlens.value().size(0) - 1 |
| 164 | : x_in.size(0); |
| 165 | int64_t ngroups = B_in.size(1); |
| 166 | |
| 167 | // Strides |
| 168 | int64_t stride_state_n = state.stride(0); |
| 169 | int64_t stride_state_h = state.stride(1); |
| 170 | int64_t stride_state_d = state.stride(2); |
| 171 | int64_t stride_x_n = x_in.stride(0); |
| 172 | int64_t stride_x_h = x_in.stride(1); |
| 173 | int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads) |
| 174 | int64_t stride_BC_n = B_in.stride(0); |
| 175 | int64_t stride_BC_g = B_in.stride(1); |
| 176 | int64_t stride_out_n = out.stride(0); |
| 177 | int64_t stride_out_h = out.stride(1); |
| 178 | |
| 179 | // Optional index pointers |
| 180 | auto get_int32_ptr = |
| 181 | [](const c10::optional<at::Tensor>& opt) -> const int32_t* { |
| 182 | return (opt.has_value() && opt.value().defined()) |
| 183 | ? opt.value().data_ptr<int32_t>() |
| 184 | : nullptr; |
| 185 | }; |
| 186 | const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices); |
| 187 | const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices); |
| 188 | const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens); |
| 189 | const int32_t* csl_ptr = get_int32_ptr(cu_seqlens); |
| 190 | |
| 191 | // Dispatch on (state_t, input_t, out_t): write directly into `out` |
| 192 | // without any intermediate float32 buffer. |
| 193 | VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] { |
| 194 | using state_t = scalar_t; |
| 195 | VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] { |
| 196 | using input_t = scalar_t; |
| 197 | VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] { |
| 198 | using out_t = scalar_t; |
| 199 | mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>( |
| 200 | state.data_ptr<state_t>(), stride_state_n, stride_state_h, |
| 201 | stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h, |
| 202 | dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(), |
| 203 | B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n, |
| 204 | stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr, |
| 205 | z_in.defined() ? z_in.data_ptr<input_t>() : nullptr, |
| 206 | dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr, |
| 207 | out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr, |
| 208 | dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N, |
| 209 | nheads, ngroups, dim, dstate, dt_softplus); |
| 210 | }); |
| 211 | }); |
| 212 | }); |
| 213 | } |
| 214 | |
| 215 | // --------------------------------------------------------------------------- |
| 216 | // mamba_chunk_scan_fwd_cpu |
| 217 | // --------------------------------------------------------------------------- |
| 218 | void mamba_chunk_scan_fwd_cpu_impl( |
| 219 | at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller |
| 220 | at::Tensor& |
| 221 | final_states, // [batch, nheads, headdim, dstate] float32 contiguous |
| 222 | const at::Tensor& x, // [seqlen, nheads, headdim] |
| 223 | const at::Tensor& |
| 224 | dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp) |
| 225 | const at::Tensor& A, // [nheads] float32 |
| 226 | const at::Tensor& B, // [seqlen, ngroups, dstate] |
| 227 | const at::Tensor& C, // [seqlen, ngroups, dstate] |
| 228 | const c10::optional<at::Tensor>& D, // [nheads] float32 (optional) |
| 229 | const c10::optional<at::Tensor>& z, // [seqlen, nheads, headdim] (optional) |
| 230 | const at::Tensor& cu_seqlens // [batch+1] int32 |
| 231 | ) { |
| 232 | const at::ScalarType input_type = x.scalar_type(); |
| 233 | |
| 234 | auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor { |
| 235 | at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t; |
| 236 | return r.is_contiguous() ? r : r.contiguous(); |
| 237 | }; |
| 238 | at::Tensor x_in = ensure_contig(x); |
| 239 | at::Tensor B_in = ensure_contig(B); |
| 240 | at::Tensor C_in = ensure_contig(C); |
| 241 | at::Tensor z_in; |
| 242 | if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value()); |
| 243 | |
| 244 | // A and D are float32 model parameters, potentially broadcast-expanded. |
| 245 | // Strip trailing broadcast dims to get a contiguous (nheads,) array. |
| 246 | auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor { |
| 247 | at::Tensor r = t; |
| 248 | while (r.dim() > 1) r = r.select(r.dim() - 1, 0); |
| 249 | if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat); |
| 250 | return r.is_contiguous() ? r : r.contiguous(); |
| 251 | }; |
| 252 | at::Tensor A_f32 = to_per_head_f32(A); |
| 253 | at::Tensor D_f32; |
| 254 | if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value()); |
| 255 | |
| 256 | // dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in |
| 257 | // Python. |
| 258 | at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous(); |
| 259 | if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat); |
| 260 | |
| 261 | at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous(); |
| 262 | |
| 263 | const int64_t batch = final_states.size(0); |
| 264 | const int64_t nheads = final_states.size(1); |
| 265 | const int64_t headdim = final_states.size(2); |
| 266 | const int64_t dstate = final_states.size(3); |
| 267 | const int64_t ngroups = B_in.size(1); |
| 268 | |
| 269 | TORCH_CHECK(final_states.is_contiguous(), |
| 270 | "mamba_chunk_scan_fwd_cpu: final_states must be contiguous"); |
| 271 | TORCH_CHECK(out.is_contiguous(), |
| 272 | "mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via " |
| 273 | "raw data_ptr)"); |
| 274 | |
| 275 | VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] { |
| 276 | mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>( |
| 277 | final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(), |
| 278 | dt_c.data_ptr<float>(), A_f32.data_ptr<float>(), |
| 279 | B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(), |
| 280 | D_f32.defined() ? D_f32.data_ptr<float>() : nullptr, |
| 281 | z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr, |
| 282 | out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads, |
| 283 | ngroups, headdim, dstate); |
| 284 | }); |
| 285 | } |