feat(cpu): add CPU support for Mamba ShortConv (#35059)
Signed-off-by: Rahul Vishwakarma <Rahul.Vishwakarma2@ibm.com>
Rahul Vishwakarma committed
Jul 7, 2026 at 08:17 UTC
9dd24658966a7224388650e653d2935daacdb71f
4 files changed
+277
-4
.buildkite/hardware_tests/cpu.yaml
+2
@@ -17,12 +17,14 @@ steps:
17
- tests/kernels/test_awq_int4_to_int8.py
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
commands:
22
- |
23
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
24
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
25
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
26
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
27
+ pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
28
pytest -x -v -s tests/kernels/test_onednn.py
29
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
30
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+2
-2
@@ -39,7 +39,8 @@ function cpu_tests() {
39
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
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"
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"
44
45
# skip tests requiring model downloads if HF_TOKEN is not set
46
# due to rate-limits
@@ -63,7 +64,6 @@ function cpu_tests() {
64
set -e
65
pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"
66
66
-
67
# basic online serving
68
docker exec cpu-test bash -c '
69
set -e
tests/kernels/mamba/test_cpu_short_conv.py
new
+180
@@ -0,0 +1,180 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+from unittest.mock import MagicMock, patch
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.config import CompilationConfig, VllmConfig
11
+from vllm.forward_context import set_forward_context
12
+from vllm.model_executor.layers.mamba.short_conv import ShortConv
13
+from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm
14
+from vllm.platforms import current_platform
15
+from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata
16
+
17
+if not current_platform.is_cpu():
18
+ pytest.skip("skipping CPU-only tests", allow_module_level=True)
19
+
20
+
21
+@pytest.fixture(autouse=True)
22
+def mock_dist():
23
+ with (
24
+ patch(
25
+ "vllm.model_executor.layers.linear.get_tensor_model_parallel_rank",
26
+ return_value=0,
27
+ ),
28
+ patch(
29
+ "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
30
+ return_value=1,
31
+ ),
32
+ patch(
33
+ "vllm.distributed.parallel_state.model_parallel_is_initialized",
34
+ return_value=True,
35
+ ),
36
+ patch(
37
+ "vllm.distributed.parallel_state.get_tp_group",
38
+ return_value=MagicMock(rank_in_group=0),
39
+ ),
40
+ ):
41
+ yield
42
+
43
+
44
+@pytest.fixture
45
+def vllm_config():
46
+ # ShortConv only needs compilation_config from the current vLLM config, so a
47
+ # minimal config (model_config=None) avoids mocking ModelConfig and the
48
+ # associated VllmConfig validation churn.
49
+ return VllmConfig(compilation_config=CompilationConfig())
50
+
51
+
52
+def test_short_conv_forward_native_prefill(vllm_config):
53
+ prefix = "test_layer"
54
+ config = SimpleNamespace(conv_L_cache=4, conv_bias=True)
55
+ dim = 16
56
+
57
+ from vllm.config import set_current_vllm_config
58
+
59
+ with set_current_vllm_config(vllm_config):
60
+ layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)
61
+
62
+ layer.to("cpu")
63
+ dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
64
+ dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)
65
+
66
+ # Mock AttentionMetadata
67
+ num_prefills = 1
68
+ num_prefill_tokens = 5
69
+ query_start_loc_p = torch.tensor([0, 5], dtype=torch.int32)
70
+ state_indices_tensor_p = torch.tensor([0], dtype=torch.int32)
71
+
72
+ # ShortConvAttentionMetadata
73
+ attn_metadata = ShortConvAttentionMetadata(
74
+ num_prefills=num_prefills,
75
+ num_prefill_tokens=num_prefill_tokens,
76
+ num_decodes=0,
77
+ num_decode_tokens=0,
78
+ num_reqs=1,
79
+ query_start_loc_p=query_start_loc_p,
80
+ has_initial_states_p=torch.tensor([False]),
81
+ state_indices_tensor_p=state_indices_tensor_p,
82
+ state_indices_tensor_d=torch.empty((0, 1), dtype=torch.int32),
83
+ num_accepted_tokens=None,
84
+ query_start_loc_d=None,
85
+ block_idx_last_scheduled_token=None,
86
+ block_idx_first_scheduled_token_p=None,
87
+ block_idx_last_computed_token=None,
88
+ block_idx_last_scheduled_token_prev_step=None,
89
+ num_computed_tokens_p=None,
90
+ seq_lens=torch.tensor([5]),
91
+ )
92
+
93
+ # Mock KV cache
94
+ # conv_state shape (num_blocks, L_cache - 1, dim)
95
+ conv_state = torch.zeros((1, config.conv_L_cache - 1, dim))
96
+ layer.kv_cache = (conv_state,)
97
+
98
+ hidden_states = torch.randn((num_prefill_tokens, dim))
99
+ output = torch.zeros_like(hidden_states)
100
+
101
+ attn_metadata_dict = {prefix: attn_metadata}
102
+ with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config):
103
+ layer.forward_native(hidden_states, output)
104
+
105
+ # Check if KV cache was updated
106
+ assert not torch.allclose(conv_state, torch.zeros_like(conv_state))
107
+
108
+
109
+def test_short_conv_forward_native_decode(vllm_config):
110
+ prefix = "test_layer_decode"
111
+ config = SimpleNamespace(conv_L_cache=4, conv_bias=True)
112
+ dim = 16
113
+
114
+ from vllm.config import set_current_vllm_config
115
+
116
+ with set_current_vllm_config(vllm_config):
117
+ layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)
118
+
119
+ layer.to("cpu")
120
+ dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
121
+ dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)
122
+
123
+ # Mock AttentionMetadata for 2 decode requests
124
+ num_decodes = 2
125
+ state_indices_tensor_d = torch.tensor([0, 1], dtype=torch.int32)
126
+
127
+ attn_metadata = ShortConvAttentionMetadata(
128
+ num_prefills=0,
129
+ num_prefill_tokens=0,
130
+ num_decodes=num_decodes,
131
+ num_decode_tokens=num_decodes,
132
+ num_reqs=num_decodes,
133
+ query_start_loc_p=None,
134
+ has_initial_states_p=None,
135
+ state_indices_tensor_p=torch.empty((0,), dtype=torch.int32),
136
+ state_indices_tensor_d=state_indices_tensor_d,
137
+ num_accepted_tokens=None,
138
+ query_start_loc_d=torch.tensor([0, 1, 2], dtype=torch.int32),
139
+ block_idx_last_scheduled_token=None,
140
+ block_idx_first_scheduled_token_p=None,
141
+ block_idx_last_computed_token=None,
142
+ block_idx_last_scheduled_token_prev_step=None,
143
+ num_computed_tokens_p=None,
144
+ seq_lens=torch.tensor([1, 1]),
145
+ )
146
+
147
+ # Mock KV cache (2 blocks for 2 requests)
148
+ conv_state = torch.randn((2, config.conv_L_cache - 1, dim))
149
+ layer.kv_cache = (conv_state,)
150
+
151
+ hidden_states = torch.randn((num_decodes, dim))
152
+ output = torch.zeros_like(hidden_states)
153
+
154
+ old_conv_state = conv_state.clone()
155
+
156
+ attn_metadata_dict = {prefix: attn_metadata}
157
+ with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config):
158
+ layer.forward_native(hidden_states, output)
159
+
160
+ # Check if KV cache was updated
161
+ assert not torch.allclose(conv_state, old_conv_state)
162
+
163
+
164
+def test_dispatch_cpu_unquantized_gemm_conv_layer():
165
+ # Convolution layers have >2D weights; dispatch should skip them gracefully.
166
+ # Shape/dtype are AMX-pack safe (bf16, width==4, dim % block_size == 0) so
167
+ # the AMX prepack branch does not raise on AMX-capable CPUs.
168
+ class MockConvLayer(torch.nn.Module):
169
+ def __init__(self):
170
+ super().__init__()
171
+ self.weight = torch.nn.Parameter(
172
+ torch.randn(32, 1, 4, dtype=torch.bfloat16)
173
+ )
174
+ self.bias = torch.nn.Parameter(torch.randn(32, dtype=torch.bfloat16))
175
+
176
+ layer = MockConvLayer()
177
+ # The ndim != 2 guard returns early without raising.
178
+ dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
179
+ # No cpu_linear set — conv layers are handled elsewhere.
180
+ assert not hasattr(layer, "cpu_linear")
vllm/model_executor/layers/mamba/short_conv.py
+93
-2
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
23
causal_conv1d_fn,
24
causal_conv1d_update,
25
)
26
+from vllm.platforms import current_platform
27
from vllm.utils.torch_utils import direct_register_custom_op
28
from vllm.v1.attention.backend import AttentionMetadata
29
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
@@ -90,7 +91,94 @@ class ShortConv(MambaBase, CustomOp):
91
hidden_states: torch.Tensor,
92
output: torch.Tensor,
93
):
93
- return
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,
98
+ causal_conv1d_update_torch,
99
+ )
100
+
101
+ forward_context = get_forward_context()
102
+ attn_metadata_raw = forward_context.attn_metadata
103
+ attn_metadata: AttentionMetadata | None = None
104
+ if attn_metadata_raw is not None:
105
+ assert isinstance(attn_metadata_raw, dict)
106
+ attn_metadata = attn_metadata_raw[self.prefix]
107
+ assert isinstance(attn_metadata, ShortConvAttentionMetadata)
108
+
109
+ BCx, _ = self.in_proj(hidden_states)
110
+ B, C, x = BCx.chunk(3, dim=-1)
111
+
112
+ # (dim, kernel_size) — same reshape as forward_cuda
113
+ conv_weights = self.conv.weight.view(
114
+ self.conv.weight.size(0), self.conv.weight.size(2)
115
+ )
116
+
117
+ if attn_metadata is None:
118
+ # Profile run — output value doesn't matter
119
+ Bx = (B * x).contiguous()
120
+ output_tensor, _ = self.out_proj(C * Bx)
121
+ output[: hidden_states.shape[0]] = output_tensor
122
+ return
123
+
124
+ conv_state = (
125
+ self.kv_cache[0]
126
+ if is_conv_state_dim_first()
127
+ else self.kv_cache[0].transpose(-1, -2)
128
+ ) # (num_blocks, dim, state_len)
129
+
130
+ num_prefills = attn_metadata.num_prefills
131
+ num_decodes = attn_metadata.num_decode_tokens
132
+ num_prefill_tokens = attn_metadata.num_prefill_tokens
133
+ has_prefill = num_prefills > 0
134
+ has_decode = num_decodes > 0
135
+ num_actual_tokens = num_decodes + num_prefill_tokens
136
+
137
+ B_d, B_p = torch.split(
138
+ B[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0
139
+ )
140
+ C_d, C_p = torch.split(
141
+ C[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0
142
+ )
143
+ x_d, x_p = torch.split(
144
+ x[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0
145
+ )
146
+
147
+ conv_output_list = []
148
+
149
+ if has_prefill:
150
+ assert attn_metadata.state_indices_tensor_p is not None
151
+ Bx_p = (B_p * x_p).transpose(0, 1) # (dim, num_prefill_tokens)
152
+ out_p = causal_conv1d_torch(
153
+ Bx_p,
154
+ conv_weights,
155
+ self.conv.bias,
156
+ conv_state,
157
+ attn_metadata.query_start_loc_p,
158
+ attn_metadata.state_indices_tensor_p.flatten(),
159
+ attn_metadata.has_initial_states_p,
160
+ activation=None,
161
+ ).transpose(0, 1)[:num_prefill_tokens] # (num_prefill_tokens, dim)
162
+ conv_output_list.append(C_p * out_p)
163
+
164
+ if has_decode:
165
+ assert attn_metadata.state_indices_tensor_d is not None
166
+ 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
178
+ conv_output_list.insert(0, C_d * out_d)
179
+
180
+ hidden_states_out = torch.vstack(conv_output_list)
181
+ output[:num_actual_tokens], _ = self.out_proj(hidden_states_out)
182
183
def forward(
184
self,
@@ -235,7 +323,10 @@ def short_conv(
323
) -> None:
324
forward_context: ForwardContext = get_forward_context()
325
self = forward_context.no_compile_layers[layer_name]
238
- self.forward_cuda(hidden_states=hidden_states, output=output)
326
+ if not current_platform.is_cpu():
327
+ self.forward_cuda(hidden_states=hidden_states, output=output)
328
+ else:
329
+ self.forward_native(hidden_states=hidden_states, output=output)
330
331
332
def short_conv_fake(