@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5c9f6557d7

[Hardware][CPU] Enable granite-4 model on cpu (#47641)

Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com> Signed-off-by: Akash Kaothalkar <akashkaothalkar@dhcp-9-123-5-76.bl1-in.ibm.com> Signed-off-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan> Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Akash Kaothalkar <akashkaothalkar@dhcp-9-123-5-76.bl1-in.ibm.com> Co-authored-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan> Co-authored-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com> Co-authored-by: Akash kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

Akash kaothalkar committed Jul 20, 2026 at 11:45 UTC 5c9f6557d7845b52917ef013011c83618439b17f
21 files changed +1283 -62
.buildkite/hardware_tests/cpu.yaml
+5 -1
@@ -18,6 +18,8 @@ steps:
18 - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
19 - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
20 - tests/kernels/mamba/test_cpu_short_conv.py
21 + - tests/kernels/mamba/test_causal_conv1d.py
22 + - tests/kernels/mamba/test_mamba_ssm.py
23 commands:
24 - |
25 bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -28,7 +30,9 @@ steps:
30 pytest -x -v -s tests/kernels/test_onednn.py
31 pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
32 pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
31 - pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
33 + pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
34 + pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
35 + pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
36
37 # Note: SDE can't be downloaded from CI host because of AWS WAF
38 # - label: CPU-Compatibility Tests
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+4 -1
@@ -40,7 +40,9 @@ function cpu_tests() {
40 pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
41 pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
42 pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
43 - pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
43 + pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
44 + pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
45 + pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
46
47 # skip tests requiring model downloads if HF_TOKEN is not set
48 # due to rate-limits
@@ -97,3 +99,4 @@ function cpu_tests() {
99 # All of CPU tests are expected to be finished less than 40 mins.
100 export -f cpu_tests
101 timeout 2h bash -c cpu_tests
102 +
cmake/cpu_extension.cmake
+3
@@ -430,6 +430,7 @@ set(VLLM_EXT_SRC
430 "csrc/cpu/layernorm.cpp"
431 "csrc/cpu/mla_decode.cpp"
432 "csrc/cpu/pos_encoding.cpp"
433 + "csrc/cpu/mamba_cpu.cpp"
434 "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
435 "csrc/cpu/cpu_attn.cpp"
436 "csrc/cpu/torch_bindings.cpp")
@@ -489,6 +490,7 @@ if (ENABLE_X86_ISA)
490 "csrc/cpu/spec_decode_utils.cpp"
491 "csrc/cpu/cpu_attn.cpp"
492 "csrc/cpu/dnnl_kernels.cpp"
493 + "csrc/cpu/mamba_cpu.cpp"
494 "csrc/cpu/torch_bindings.cpp"
495 # TODO: Remove these files
496 "csrc/cpu/activation.cpp"
@@ -502,6 +504,7 @@ if (ENABLE_X86_ISA)
504 "csrc/cpu/utils.cpp"
505 "csrc/cpu/spec_decode_utils.cpp"
506 "csrc/cpu/cpu_attn.cpp"
507 + "csrc/cpu/mamba_cpu.cpp"
508 "csrc/cpu/dnnl_kernels.cpp"
509 "csrc/cpu/torch_bindings.cpp"
510 # TODO: Remove these files
csrc/cpu/cpu_types_vsx.hpp
+8 -7
@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
336 reg.val[1] = fp16_to_fp32_bits(raw_lo);
337 }
338 float reduce_sum() const {
339 - AliasReg ar;
340 - ar.reg = reg;
341 - float result = 0;
342 - unroll_loop<int, VEC_ELEM_NUM>(
343 - [&result, &ar](int i) { result += ar.values[i]; });
344 -
345 - return result;
339 + // VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
340 + // Step 1: pairwise sum of the two 4-wide halves
341 + __vector float s = vec_add(reg.val[0], reg.val[1]);
342 + // Step 2: rotate by 8 bytes (2 floats) and add
343 + s = vec_add(s, vec_sld(s, s, 8));
344 + // Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total
345 + s = vec_add(s, vec_sld(s, s, 4));
346 + return vec_extract(s, 0);
347 }
348 FP32Vec8 exp() const {
349 f32x4x2_t out;
csrc/cpu/mamba_cpu.cpp new
+285
@@ -0,0 +1,285 @@
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 +}
csrc/cpu/mamba_kernels.hpp new
+382
@@ -0,0 +1,382 @@
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
csrc/cpu/torch_bindings.cpp
+50
@@ -213,6 +213,32 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
213 torch::Tensor slot_mapping,
214 const int64_t block_size);
215
216 +at::Tensor causal_conv1d_update_cpu_impl(
217 + at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
218 + const c10::optional<at::Tensor>& bias,
219 + const c10::optional<std::string>& activation,
220 + const c10::optional<at::Tensor>& conv_state_indices,
221 + const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
222 +
223 +void selective_state_update_cpu_impl(
224 + at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
225 + const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
226 + const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
227 + const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
228 + const c10::optional<at::Tensor>& state_batch_indices,
229 + const c10::optional<at::Tensor>& dst_state_batch_indices,
230 + int64_t null_block_id, at::Tensor& out,
231 + const c10::optional<at::Tensor>& num_accepted_tokens,
232 + const c10::optional<at::Tensor>& cu_seqlens);
233 +
234 +void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
235 + const at::Tensor& x, const at::Tensor& dt,
236 + const at::Tensor& A, const at::Tensor& B,
237 + const at::Tensor& C,
238 + const c10::optional<at::Tensor>& D,
239 + const c10::optional<at::Tensor>& z,
240 + const at::Tensor& cu_seqlens);
241 +
242 void init_cpu_memory_env(std::vector<int64_t> node_ids);
243
244 namespace cpu_utils {
@@ -595,6 +621,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
621 "block_size) -> ()",
622 &compute_slot_mapping_kernel_impl);
623
624 + // Mamba CPU kernels
625 + ops.def(
626 + "causal_conv1d_update_cpu_vec("
627 + "Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
628 + "Tensor? bias, str? activation, Tensor? conv_state_indices, "
629 + "Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
630 + &causal_conv1d_update_cpu_impl);
631 +
632 + ops.def(
633 + "selective_state_update_cpu("
634 + "Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
635 + "Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
636 + "Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
637 + "SymInt null_block_id, Tensor(a13!) out, "
638 + "Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
639 + &selective_state_update_cpu_impl);
640 +
641 + ops.def(
642 + "mamba_chunk_scan_fwd_cpu("
643 + "Tensor(a0!) out, Tensor(a1!) final_states, "
644 + "Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
645 + "Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
646 + &mamba_chunk_scan_fwd_cpu_impl);
647 +
648 ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
649
650 // Speculative decoding kernels
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+1 -1
@@ -425,7 +425,7 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No
425 match the single-call result.
426 """
427 from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
428 - causal_conv1d_torch,
428 + causal_conv1d_fn_cpu as causal_conv1d_torch,
429 )
430
431 x, weight, bias = _conv_inputs(total_tokens)
tests/kernels/mamba/test_causal_conv1d.py
+8 -3
@@ -18,8 +18,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
18 DEVICE = current_platform.device_type
19
20 pytestmark = pytest.mark.skipif(
21 - not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
22 - reason="causal_conv1d Triton kernels require CUDA-alike or XPU",
21 + not (
22 + current_platform.is_cuda_alike()
23 + or current_platform.is_xpu()
24 + or current_platform.is_cpu()
25 + ),
26 + reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU",
27 )
28
29
@@ -284,7 +288,8 @@ def test_causal_conv1d_varlen(
288 batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
289 ):
290 device = DEVICE
287 - torch.accelerator.empty_cache()
291 + if not current_platform.is_cpu():
292 + torch.accelerator.empty_cache()
293 rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
294 if itype == torch.bfloat16:
295 rtol, atol = 1e-2, 5e-2
tests/kernels/mamba/test_mamba_ssm.py
+48 -2
@@ -20,8 +20,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
20 DEVICE = current_platform.device_type
21
22 pytestmark = pytest.mark.skipif(
23 - not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
24 - reason="mamba_ssm kernels require CUDA-alike or XPU",
23 + not (
24 + current_platform.is_cuda_alike()
25 + or current_platform.is_xpu()
26 + or current_platform.is_cpu()
27 + ),
28 + reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU",
29 )
30
31 # selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op,
@@ -342,6 +346,13 @@ def test_selective_scan(
346 @pytest.mark.parametrize("has_z", [False, True])
347 @pytest.mark.parametrize("dstate", [16, 64])
348 @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
349 +@pytest.mark.skipif(
350 + current_platform.is_cpu(),
351 + reason=(
352 + "CPU kernel for selective_state_update only supports "
353 + "Mamba 2 (scalar A/dt), not Mamba 1."
354 + ),
355 +)
356 def test_selective_state_update(dim, dstate, has_z, itype):
357 device = DEVICE
358 rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
@@ -436,6 +447,13 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r
447 @pytest.mark.parametrize("dstate", [16, 64])
448 @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
449 @pytest.mark.parametrize("max_seq_len", [1, 2, 4])
450 +@pytest.mark.skipif(
451 + current_platform.is_cpu(),
452 + reason=(
453 + "CPU kernel for selective_state_update only supports "
454 + "Mamba 2 (scalar A/dt), not Mamba 1."
455 + ),
456 +)
457 def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
458 device = DEVICE
459 rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
@@ -697,6 +715,13 @@ def test_selective_scan_varlen(
715 @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
716 # tests correctness in case subset of the sequences are padded
717 @pytest.mark.parametrize("with_padding", [True, False])
718 +@pytest.mark.skipif(
719 + current_platform.is_cpu(),
720 + reason=(
721 + "CPU kernel for selective_state_update only supports "
722 + "Mamba 2 (scalar A/dt), not Mamba 1."
723 + ),
724 +)
725 def test_selective_state_update_with_batch_indices(
726 with_padding, dim, dstate, has_z, itype
727 ):
@@ -789,6 +814,13 @@ def test_selective_state_update_with_batch_indices(
814 @pytest.mark.parametrize("ngroups", [1, 4])
815 @pytest.mark.parametrize("dstate", [16, 64])
816 @pytest.mark.parametrize("dim", [2048, 4096])
817 +@pytest.mark.skipif(
818 + current_platform.is_cpu(),
819 + reason=(
820 + "CPU kernel for selective_state_update only supports "
821 + "Mamba 2 (scalar A/dt), not Mamba 1."
822 + ),
823 +)
824 def test_selective_state_update_with_heads_with_batch_indices(
825 dim, dstate, ngroups, has_z, tie_hdim, itype
826 ):
@@ -862,6 +894,13 @@ def test_selective_state_update_with_heads_with_batch_indices(
894 @pytest.mark.parametrize("dstate", [16, 64])
895 @pytest.mark.parametrize("dim", [2048, 4096])
896 @pytest.mark.parametrize("max_seq_len", [2, 4])
897 +@pytest.mark.skipif(
898 + current_platform.is_cpu(),
899 + reason=(
900 + "CPU kernel for selective_state_update only supports "
901 + "Mamba 2 (scalar A/dt), not Mamba 1."
902 + ),
903 +)
904 def test_selective_state_update_with_num_accepted_tokens(
905 dim, dstate, has_z, itype, max_seq_len
906 ):
@@ -988,6 +1027,13 @@ def test_selective_state_update_with_num_accepted_tokens(
1027 @pytest.mark.parametrize("dstate", [16, 64])
1028 @pytest.mark.parametrize("dim", [2048, 4096])
1029 @pytest.mark.parametrize("max_seq_len", [2, 4])
1030 +@pytest.mark.skipif(
1031 + current_platform.is_cpu(),
1032 + reason=(
1033 + "CPU kernel for selective_state_update only supports "
1034 + "Mamba 2 (scalar A/dt), not Mamba 1."
1035 + ),
1036 +)
1037 def test_selective_state_update_varlen_with_num_accepted(
1038 dim, dstate, has_z, itype, max_seq_len
1039 ):
vllm/_custom_ops.py
+87
@@ -2070,6 +2070,93 @@ def selective_scan_fwd(
2070 )
2071
2072
2073 +def causal_conv1d_update_cpu_vec(
2074 + x: torch.Tensor,
2075 + conv_state: torch.Tensor,
2076 + weight: torch.Tensor,
2077 + bias: torch.Tensor | None = None,
2078 + activation: str | None = None,
2079 + conv_state_indices: torch.Tensor | None = None,
2080 + query_start_loc: torch.Tensor | None = None,
2081 + pad_slot_id: int = 0,
2082 +) -> torch.Tensor:
2083 + return torch.ops._C.causal_conv1d_update_cpu_vec(
2084 + x,
2085 + conv_state,
2086 + weight,
2087 + bias,
2088 + activation,
2089 + conv_state_indices,
2090 + query_start_loc,
2091 + pad_slot_id,
2092 + )
2093 +
2094 +
2095 +def selective_state_update_cpu(
2096 + state: torch.Tensor,
2097 + x: torch.Tensor,
2098 + dt: torch.Tensor,
2099 + A: torch.Tensor,
2100 + B: torch.Tensor,
2101 + C: torch.Tensor,
2102 + D: torch.Tensor | None,
2103 + z: torch.Tensor | None,
2104 + dt_bias: torch.Tensor | None,
2105 + dt_softplus: bool,
2106 + state_batch_indices: torch.Tensor | None,
2107 + dst_state_batch_indices: torch.Tensor | None,
2108 + null_block_id: int,
2109 + out: torch.Tensor,
2110 + num_accepted_tokens: torch.Tensor | None,
2111 + cu_seqlens: torch.Tensor | None,
2112 +):
2113 + torch.ops._C.selective_state_update_cpu(
2114 + state,
2115 + x,
2116 + dt,
2117 + A,
2118 + B,
2119 + C,
2120 + D,
2121 + z,
2122 + dt_bias,
2123 + dt_softplus,
2124 + state_batch_indices,
2125 + dst_state_batch_indices,
2126 + null_block_id,
2127 + out,
2128 + num_accepted_tokens,
2129 + cu_seqlens,
2130 + )
2131 +
2132 +
2133 +def mamba_chunk_scan_fwd_cpu(
2134 + out: torch.Tensor,
2135 + final_states: torch.Tensor,
2136 + x: torch.Tensor,
2137 + dt: torch.Tensor,
2138 + A: torch.Tensor,
2139 + B: torch.Tensor,
2140 + C: torch.Tensor,
2141 + D: torch.Tensor | None,
2142 + z: torch.Tensor | None,
2143 + cu_seqlens: torch.Tensor,
2144 +) -> None:
2145 + """Prefill SSM scan kernel. out and final_states are written in-place."""
2146 + torch.ops._C.mamba_chunk_scan_fwd_cpu(
2147 + out,
2148 + final_states,
2149 + x,
2150 + dt,
2151 + A,
2152 + B,
2153 + C,
2154 + D,
2155 + z,
2156 + cu_seqlens,
2157 + )
2158 +
2159 +
2160 # ROCm skinny gemms
2161 def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor:
2162 return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)
vllm/config/mamba.py
+1
@@ -27,6 +27,7 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta):
27
28 TRITON = "triton"
29 FLASHINFER = "flashinfer"
30 + CPU = "cpu"
31
32
33 @config
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+12
@@ -1237,3 +1237,15 @@ def causal_conv1d_update(
1237 if unsqueeze:
1238 out = out.squeeze(-1)
1239 return out.to(original_x_dtype)
1240 +
1241 +
1242 +from vllm.platforms import current_platform # noqa: E402
1243 +
1244 +if current_platform.is_cpu():
1245 + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
1246 + causal_conv1d_fn_cpu,
1247 + causal_conv1d_update_cpu,
1248 + )
1249 +
1250 + causal_conv1d_fn = causal_conv1d_fn_cpu # type: ignore
1251 + causal_conv1d_update = causal_conv1d_update_cpu # type: ignore
vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py
+68 -8
@@ -6,18 +6,31 @@ from __future__ import annotations
6 import torch
7 import torch.nn.functional as F
8
9 +from vllm._custom_ops import causal_conv1d_update_cpu_vec
10 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
11
10 -# for prefill
11 -def causal_conv1d_torch(
12 +
13 +def causal_conv1d_fn_cpu(
14 x: torch.Tensor,
15 weight: torch.Tensor,
16 bias: torch.Tensor | None,
17 conv_states: torch.Tensor,
18 query_start_loc: torch.Tensor,
17 - cache_indices: torch.Tensor,
18 - has_initial_state: torch.Tensor,
19 + cache_indices: torch.Tensor | None = None,
20 + has_initial_state: torch.Tensor | None = None,
21 activation: str | None = "silu",
22 + pad_slot_id: int = PAD_SLOT_ID,
23 + **kwargs,
24 ) -> torch.Tensor:
25 + """CPU implementation for causal_conv1d_fwd."""
26 + if isinstance(activation, bool) and activation:
27 + activation = "silu"
28 + elif isinstance(activation, bool):
29 + activation = None
30 +
31 + original_x_dtype = x.dtype
32 + x = x.to(conv_states.dtype)
33 +
34 out = torch.empty_like(x)
35 state_len = weight.shape[1] - 1
36 assert activation in {None, "silu", "swish"}
@@ -27,11 +40,21 @@ def causal_conv1d_torch(
40 for idx in range(query_start_loc.shape[0] - 1)
41 ]
42 weight = weight.unsqueeze(1)
43 +
44 for seq_idx, (bos, eos) in enumerate(seq_begin_end_idx):
31 - slot = int(cache_indices[seq_idx].item())
45 + if bos == eos:
46 + continue
47 +
48 + slot = (
49 + int(cache_indices[seq_idx].item()) if cache_indices is not None else seq_idx
50 + )
51 +
52 + if slot == pad_slot_id:
53 + continue
54
55 seq_x = x[:, bos:eos].unsqueeze(0)
34 - if bool(has_initial_state[seq_idx].item()):
56 +
57 + if has_initial_state is not None and bool(has_initial_state[seq_idx].item()):
58 initial_state = conv_states[slot, :, :state_len].unsqueeze(0)
59 else:
60 initial_state = torch.zeros(
@@ -51,16 +74,48 @@ def causal_conv1d_torch(
74 groups=weight.shape[0],
75 )
76 seq_out = seq_out[..., -seq_x.shape[-1] :].to(dtype=x.dtype)
77 +
78 if activation in ("silu", "swish"):
79 seq_out = F.silu(seq_out)
80
81 out[:, bos:eos] = seq_out.squeeze(0)
82 conv_states[slot, :, :state_len].copy_(conv_input[..., -state_len:].squeeze(0))
83
60 - return out
84 + return out.to(original_x_dtype)
85 +
86 +
87 +def causal_conv1d_update_cpu(
88 + x: torch.Tensor,
89 + conv_state: torch.Tensor,
90 + weight: torch.Tensor,
91 + bias: torch.Tensor | None = None,
92 + activation: bool | str | None = None,
93 + conv_state_indices: torch.Tensor | None = None,
94 + query_start_loc: torch.Tensor | None = None,
95 + pad_slot_id: int | None = None,
96 + **kwargs,
97 +) -> torch.Tensor:
98 + """CPU implementation for causal_conv1d_update."""
99 + if isinstance(activation, bool):
100 + activation = "silu" if activation else None
101 +
102 + if pad_slot_id is None:
103 + pad_slot_id = kwargs.get("null_block_id", NULL_BLOCK_ID)
104 + if pad_slot_id is None:
105 + pad_slot_id = NULL_BLOCK_ID
106 +
107 + return causal_conv1d_update_cpu_vec(
108 + x,
109 + conv_state,
110 + weight,
111 + bias,
112 + activation,
113 + conv_state_indices,
114 + query_start_loc,
115 + pad_slot_id,
116 + )
117
118
63 -# for decode
119 def causal_conv1d_update_torch(
120 x: torch.Tensor,
121 conv_state: torch.Tensor,
@@ -68,6 +123,11 @@ def causal_conv1d_update_torch(
123 bias: torch.Tensor | None = None,
124 activation: str | None = None,
125 ) -> torch.Tensor:
126 + """
127 + Pure PyTorch fallback for causal_conv1d_update.
128 + Currently used as a fallback for Arm (aarch64) to leverage
129 + oneDNN/ACL F.conv1d kernels for batched decoding.
130 + """
131 assert activation in {None, "silu", "swish"}
132
133 _, dim, seq_len = x.shape
vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
+45 -23
@@ -10,9 +10,13 @@ import vllm._custom_ops as ops
10 from vllm.forward_context import ForwardContext, get_forward_context
11 from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
12 from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
13 - causal_conv1d_torch,
13 + causal_conv1d_fn_cpu as causal_conv1d_torch,
14 +)
15 +from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
16 + causal_conv1d_update_cpu,
17 causal_conv1d_update_torch,
18 )
19 +from vllm.platforms import CpuArchEnum, current_platform
20 from vllm.utils.torch_utils import (
21 LayerNameType,
22 _resolve_layer_name,
@@ -140,21 +144,30 @@ def _cpu_gdn_attention_nonspec(
144 conv_states=conv_state,
145 weight=layer.conv1d.weight,
146 bias=layer.conv1d.bias,
143 - silu_activation=layer.activation == "silu",
147 + silu_activation=(layer.activation == "silu"),
148 conv_state_indices=decode_state_indices,
149 is_vnni=True,
150 )
151 else:
148 - decode_conv_state = conv_state[decode_state_indices].contiguous()
149 - decode_mixed_qkv = causal_conv1d_update_torch(
150 - # [B, dim] -> [B, dim, 1]
151 - x=decode_mixed_qkv.unsqueeze(-1),
152 - conv_state=decode_conv_state,
153 - weight=conv_weights,
154 - bias=layer.conv1d.bias,
155 - activation=layer.activation,
156 - ).squeeze(-1)
157 - conv_state[decode_state_indices] = decode_conv_state
152 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
153 + decode_conv_state = conv_state[decode_state_indices].contiguous()
154 + decode_mixed_qkv = causal_conv1d_update_torch(
155 + x=decode_mixed_qkv.unsqueeze(-1),
156 + conv_state=decode_conv_state,
157 + weight=conv_weights,
158 + bias=layer.conv1d.bias,
159 + activation=layer.activation,
160 + ).squeeze(-1)
161 + conv_state[decode_state_indices] = decode_conv_state
162 + else:
163 + decode_mixed_qkv = causal_conv1d_update_cpu(
164 + x=decode_mixed_qkv,
165 + conv_state=conv_state,
166 + weight=conv_weights,
167 + bias=layer.conv1d.bias,
168 + activation=layer.activation,
169 + conv_state_indices=decode_state_indices,
170 + )
171
172 query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
173
@@ -495,17 +508,26 @@ def _spec_aware_nonspec(
508 decode_a = a[:num_decode_tokens]
509 decode_state_indices = state_indices_tensor[:num_decodes]
510 # Only the first ``width-1`` columns hold the real conv state.
498 - decode_conv_state = conv_buf[decode_state_indices][
499 - :, :, : width - 1
500 - ].contiguous()
501 - decode_mixed_qkv = causal_conv1d_update_torch(
502 - x=decode_mixed_qkv.unsqueeze(-1),
503 - conv_state=decode_conv_state,
504 - weight=conv_weights,
505 - bias=layer.conv1d.bias,
506 - activation=layer.activation,
507 - ).squeeze(-1)
508 - conv_buf[decode_state_indices, :, : width - 1] = decode_conv_state
511 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
512 + conv_state_view = conv_buf[:, :, : width - 1]
513 + decode_conv_state = conv_state_view[decode_state_indices].contiguous()
514 + decode_mixed_qkv = causal_conv1d_update_torch(
515 + x=decode_mixed_qkv.unsqueeze(-1),
516 + conv_state=decode_conv_state,
517 + weight=conv_weights,
518 + bias=layer.conv1d.bias,
519 + activation=layer.activation,
520 + ).squeeze(-1)
521 + conv_state_view[decode_state_indices] = decode_conv_state
522 + else:
523 + decode_mixed_qkv = causal_conv1d_update_cpu(
524 + x=decode_mixed_qkv,
525 + conv_state=conv_buf[:, :, : width - 1],
526 + weight=conv_weights,
527 + bias=layer.conv1d.bias,
528 + activation=layer.activation,
529 + conv_state_indices=decode_state_indices,
530 + )
531
532 query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
533 # rearrange_mixed_qkv can return views whose last dim is not
vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py new
+144
@@ -0,0 +1,144 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +import vllm._custom_ops as ops
7 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
8 +
9 +
10 +def _mamba_chunk_scan_combined_fwd_cpu(
11 + x,
12 + dt,
13 + A,
14 + B,
15 + C,
16 + chunk_size,
17 + out,
18 + D=None,
19 + z=None,
20 + dt_bias=None,
21 + initial_states=None,
22 + return_intermediate_states=False,
23 + seq_idx=None,
24 + cu_seqlens=None,
25 + cu_chunk_seqlens=None,
26 + last_chunk_indices=None,
27 + dt_softplus=False,
28 + dt_limit=(0.0, float("inf")),
29 + state_dtype=None,
30 + **kwargs,
31 +):
32 + seqlen, nheads, headdim = x.shape
33 + _, ngroups, dstate = B.shape
34 +
35 + assert cu_seqlens is not None
36 + batch = cu_seqlens.size(0) - 1
37 +
38 + dt_f = dt.float()
39 + if dt_bias is not None:
40 + dt_f = dt_f + dt_bias.float().unsqueeze(0)
41 + if dt_softplus:
42 + dt_f = torch.nn.functional.softplus(dt_f)
43 + if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"):
44 + dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1])
45 +
46 + all_states = torch.zeros(
47 + batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device
48 + )
49 + if initial_states is not None:
50 + all_states.copy_(initial_states.float())
51 +
52 + assert out.is_contiguous(), (
53 + "_mamba_chunk_scan_combined_fwd_cpu: `out` must be "
54 + "pre-allocated as a contiguous tensor"
55 + )
56 +
57 + D_1d = None
58 + if D is not None:
59 + d = D.float()
60 + while d.dim() > 1 and d.stride(-1) == 0:
61 + d = d.squeeze(-1)
62 + D_1d = d.contiguous()
63 +
64 + ops.mamba_chunk_scan_fwd_cpu(
65 + out,
66 + all_states,
67 + x,
68 + dt_f,
69 + A,
70 + B,
71 + C,
72 + D_1d,
73 + z,
74 + cu_seqlens.to(torch.int32),
75 + )
76 +
77 + out_dtype = state_dtype if state_dtype is not None else x.dtype
78 + all_states = all_states.to(out_dtype)
79 +
80 + return all_states
81 +
82 +
83 +def selective_state_update(
84 + state,
85 + x,
86 + dt,
87 + A,
88 + B,
89 + C,
90 + D=None,
91 + z=None,
92 + dt_bias=None,
93 + dt_softplus=False,
94 + state_batch_indices=None,
95 + dst_state_batch_indices=None,
96 + null_block_id=NULL_BLOCK_ID,
97 + out=None,
98 + num_accepted_tokens=None,
99 + cu_seqlens=None,
100 + is_blackwell=False,
101 + enable_stochastic_rounding=False,
102 + cache_philox_rounds=0,
103 +):
104 + """CPU implementation for selective_state_update."""
105 + # Ensure out tensor exists
106 + if out is None:
107 + out = torch.empty_like(x if x.dim() == 2 else x)
108 +
109 + _state = state.unsqueeze(1) if state.dim() == 3 else state
110 + _x = x.unsqueeze(1) if x.dim() == 2 else x
111 + _dt = dt.unsqueeze(1) if dt.dim() == 2 else dt
112 + _A = A.unsqueeze(0) if A.dim() == 2 else A
113 + _B = B.unsqueeze(1) if B.dim() == 2 else B
114 + _C = C.unsqueeze(1) if C.dim() == 2 else C
115 + _D = D.unsqueeze(0) if (D is not None and D.dim() == 1) else D
116 + _z = z.unsqueeze(1) if (z is not None and z.dim() == 2) else z
117 + _dt_bias = (
118 + dt_bias.unsqueeze(0)
119 + if (dt_bias is not None and dt_bias.dim() == 1)
120 + else dt_bias
121 + )
122 + _out = out.unsqueeze(1) if out.dim() == 2 else out
123 +
124 + _sbi = state_batch_indices
125 + _dsbi = dst_state_batch_indices
126 + ops.selective_state_update_cpu(
127 + _state,
128 + _x,
129 + _dt,
130 + _A,
131 + _B,
132 + _C,
133 + _D,
134 + _z,
135 + _dt_bias,
136 + dt_softplus,
137 + _sbi,
138 + _dsbi,
139 + null_block_id,
140 + _out,
141 + num_accepted_tokens,
142 + cu_seqlens,
143 + )
144 + return _out.squeeze(1) if out.dim() == 2 else _out
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+10
@@ -845,3 +845,13 @@ def selective_scan_fn(
845 return delta # output written inplace to delta
846 else:
847 return z # output written inplace to z
848 +
849 +
850 +from vllm.platforms import current_platform # noqa: E402
851 +
852 +if current_platform.is_cpu():
853 + from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import (
854 + selective_state_update as selective_state_update_cpu,
855 + )
856 +
857 + selective_state_update = selective_state_update_cpu # type: ignore
vllm/model_executor/layers/mamba/ops/ssd_combined.py
+8
@@ -225,3 +225,11 @@ def mamba_chunk_scan_combined_varlen(
225 )
226
227 return varlen_states
228 +
229 +
230 +from vllm.platforms import current_platform # noqa: E402
231 +
232 +if current_platform.is_cpu():
233 + import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_mamba_ssm
234 +
235 + _mamba_chunk_scan_combined_fwd = cpu_mamba_ssm._mamba_chunk_scan_combined_fwd_cpu # type: ignore
vllm/model_executor/layers/mamba/ops/ssu_dispatch.py
+83 -2
@@ -4,8 +4,9 @@
4 Dispatch module for Mamba selective state update (SSU) backends.
5
6 Provides a unified `selective_state_update` function that dispatches to
7 -either the Triton or FlashInfer backend based on the configured
8 -`MambaBackendEnum`. Follows SGLang's dispatch pattern adapted for vLLM.
7 +the Triton, FlashInfer, or CPU backend based on the configured
8 +`MambaBackendEnum`. On CPU-only platforms (PowerPC, x86 without CUDA)
9 +the backend defaults to 'cpu'.
10 """
11
12 from abc import ABC, abstractmethod
@@ -182,9 +183,75 @@ class FlashInferSSUBackend(MambaSSUBackend):
183 )
184
185
186 +class CPUSSUBackend(MambaSSUBackend):
187 + """CPU SSU backend using the compiled C++ VSX/scalar kernel.
188 +
189 + On CPU-only platforms (PowerPC, x86 without CUDA) this dispatches to
190 + the vectorized C++ kernel registered as ``torch.ops._C.selective_state_update_cpu``.
191 + That kernel uses vec_op SIMD intrinsics (VSX on ppc64le, AVX2 on x86,
192 + scalar fallback elsewhere) and is parallelised with OpenMP across heads.
193 +
194 + Falls back to the pure-PyTorch implementation only if the C++ op is
195 + unavailable (e.g. a CPU-less build).
196 + """
197 +
198 + def __init__(self, mamba_config: MambaConfig):
199 + super().__init__(mamba_config)
200 + from vllm import _custom_ops as ops
201 +
202 + self._cpp_kernel = ops.selective_state_update_cpu
203 + logger.info("CPUSSUBackend: using compiled C++ selective_state_update kernel.")
204 +
205 + @property
206 + def name(self) -> str:
207 + return "cpu"
208 +
209 + def __call__(
210 + self,
211 + state: torch.Tensor,
212 + x: torch.Tensor,
213 + dt: torch.Tensor,
214 + A: torch.Tensor,
215 + B: torch.Tensor,
216 + C: torch.Tensor,
217 + D: torch.Tensor,
218 + dt_bias: torch.Tensor,
219 + z: torch.Tensor | None = None,
220 + dt_softplus: bool = False,
221 + state_batch_indices: torch.Tensor | None = None,
222 + dst_state_batch_indices: torch.Tensor | None = None,
223 + null_block_id: int = NULL_BLOCK_ID,
224 + out: torch.Tensor | None = None,
225 + num_accepted_tokens: torch.Tensor | None = None,
226 + cu_seqlens: torch.Tensor | None = None,
227 + is_blackwell: bool = False,
228 + ) -> None:
229 + # C++ kernel: state shape expected as (nstates, nheads, dim, dstate)
230 + # The kernel writes in-place into `out` and updates `state`.
231 + self._cpp_kernel(
232 + state,
233 + x,
234 + dt,
235 + A,
236 + B,
237 + C,
238 + D,
239 + z,
240 + dt_bias,
241 + dt_softplus,
242 + state_batch_indices,
243 + dst_state_batch_indices,
244 + null_block_id,
245 + out,
246 + num_accepted_tokens,
247 + cu_seqlens,
248 + )
249 +
250 +
251 _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
252 MambaBackendEnum.TRITON: TritonSSUBackend,
253 MambaBackendEnum.FLASHINFER: FlashInferSSUBackend,
254 + MambaBackendEnum.CPU: CPUSSUBackend,
255 }
256
257 _mamba_ssu_backend: MambaSSUBackend | None = None
@@ -210,6 +277,20 @@ def initialize_mamba_ssu_backend(
277 global _mamba_ssu_backend
278
279 backend = mamba_config.backend
280 +
281 + # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is
282 + # unstable or unavailable. Silently fall back to the CPU
283 + # backend unless the user explicitly chose something other than "triton".
284 + if backend == MambaBackendEnum.TRITON:
285 + from vllm.platforms import current_platform
286 +
287 + if current_platform.is_cpu():
288 + logger.info(
289 + "CPU platform detected: overriding Mamba SSU backend "
290 + "from 'triton' to 'cpu'."
291 + )
292 + backend = MambaBackendEnum.CPU
293 +
294 if backend not in _BACKEND_REGISTRY:
295 raise ValueError(
296 f"Unknown Mamba SSU backend: {backend}. "
vllm/model_executor/layers/mamba/short_conv.py
+25 -12
@@ -94,9 +94,13 @@ class ShortConv(MambaBase, CustomOp):
94 # Reference torch causal conv1d; runs on all CPU platforms. AMX kernels
95 # for causal conv can be plugged in here later.
96 from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
97 - causal_conv1d_torch,
97 + causal_conv1d_fn_cpu as causal_conv1d_torch,
98 + )
99 + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
100 + causal_conv1d_update_cpu,
101 causal_conv1d_update_torch,
102 )
103 + from vllm.platforms import CpuArchEnum, current_platform
104
105 forward_context = get_forward_context()
106 attn_metadata_raw = forward_context.attn_metadata
@@ -164,17 +168,26 @@ class ShortConv(MambaBase, CustomOp):
168 if has_decode:
169 assert attn_metadata.state_indices_tensor_d is not None
170 state_indices_d = attn_metadata.state_indices_tensor_d.flatten()
167 - Bx_d = (B_d * x_d).unsqueeze(-1) # (num_decodes, dim, 1)
168 - # Advanced indexing returns a copy; update in-place then scatter back
169 - gathered = conv_state[state_indices_d] # (num_decodes, dim, state_len)
170 - out_d = causal_conv1d_update_torch(
171 - Bx_d,
172 - gathered,
173 - conv_weights,
174 - self.conv.bias,
175 - activation=None,
176 - ).squeeze(-1) # (num_decodes, dim)
177 - conv_state[state_indices_d] = gathered
171 + Bx_d = B_d * x_d # (num_decodes, dim)
172 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
173 + conv_state_view = conv_state[state_indices_d].contiguous()
174 + out_d = causal_conv1d_update_torch(
175 + Bx_d.unsqueeze(-1),
176 + conv_state_view,
177 + conv_weights,
178 + self.conv.bias,
179 + activation=None,
180 + ).squeeze(-1)
181 + conv_state[state_indices_d] = conv_state_view
182 + else:
183 + out_d = causal_conv1d_update_cpu(
184 + Bx_d,
185 + conv_state,
186 + conv_weights,
187 + self.conv.bias,
188 + activation=None,
189 + conv_state_indices=state_indices_d,
190 + )
191 conv_output_list.insert(0, C_d * out_d)
192
193 hidden_states_out = torch.vstack(conv_output_list)
vllm/model_executor/layers/utils.py
+6 -2
@@ -234,10 +234,14 @@ def dispatch_cpu_unquantized_gemm(
234 layer.cpu_linear = torch.nn.functional.linear
235 return
236
237 + # Skip CPU GEMM dispatch for non-2D weights (e.g. MoE 3D expert weights).
238 + # These layers are handled by their own specialized methods.
239 if layer.weight.ndim != 2:
240 # this is not a linear layer
239 - # For now it should be a causal_conv1d op
240 - if torch.cpu._is_amx_tile_supported():
241 + # For now it should be a causal_conv1d op or MoE 3D expert weights
242 + if torch.cpu._is_amx_tile_supported() and hasattr(
243 + ops, "causal_conv1d_weight_pack"
244 + ):
245 # prepack conv weight
246 unpacked = (
247 layer.weight.view(