A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 17:30 UTC
[CPU] Support for Whisper (#30062)
Signed-off-by: Aditya Tewari <aditya.tewari@arm.com>
Aditya Tewari committed
Dec 10, 2025 at 12:58 UTC
cebda2a4afa9ec9c6656c0aa5e96d0003e9b185d
5 files changed
+49
-24
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+5
@@ -36,6 +36,11 @@ function cpu_tests() {
36
set -e
37
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"
38
39
+ # Run model tests
40
+ docker exec cpu-test bash -c "
41
+ set -e
42
+ pytest -x -v -s tests/models/multimodal/generation/test_whisper.py -m cpu_model"
43
+
44
# Run kernel tests
45
docker exec cpu-test bash -c "
46
set -e
csrc/cpu/cpu_attn.cpp
-1
@@ -117,7 +117,6 @@ torch::Tensor get_scheduler_metadata(
117
input.casual = casual;
118
input.isa = isa;
119
input.enable_kv_split = enable_kv_split;
120
- TORCH_CHECK(casual, "Only supports casual mask for now.");
120
121
VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
122
CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] {
tests/models/multimodal/generation/test_whisper.py
+19
-2
@@ -92,13 +92,14 @@ def run_test(
92
*,
93
tensor_parallel_size: int,
94
distributed_executor_backend: str | None = None,
95
+ dtype: str = "half",
96
) -> None:
97
prompt_list = PROMPTS * 10
98
expected_list = EXPECTED[model] * 10
99
100
with vllm_runner(
101
model,
101
- dtype="half",
102
+ dtype=dtype,
103
max_model_len=448,
104
tensor_parallel_size=tensor_parallel_size,
105
distributed_executor_backend=distributed_executor_backend,
@@ -120,12 +121,28 @@ def run_test(
121
122
@pytest.mark.core_model
123
@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
124
+@pytest.mark.parametrize("dtype", ["half"])
125
@create_new_process_for_each_test()
124
-def test_models(vllm_runner, model) -> None:
126
+def test_models(vllm_runner, model, dtype) -> None:
127
run_test(
128
vllm_runner,
129
model,
130
tensor_parallel_size=1,
131
+ dtype=dtype,
132
+ )
133
+
134
+
135
+@pytest.mark.cpu_model
136
+@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
137
+@pytest.mark.parametrize("dtype", ["half"])
138
+def test_models_cpu(vllm_runner, model, dtype) -> None:
139
+ # @create_new_process_for_each_test() does not work for some runners
140
+ # TODO: to fix cpu privilege issues in run-cpu-test-arm.sh
141
+ run_test(
142
+ vllm_runner,
143
+ model,
144
+ tensor_parallel_size=1,
145
+ dtype=dtype,
146
)
147
148
vllm/v1/attention/backends/cpu_attn.py
+19
-19
@@ -21,7 +21,7 @@ from vllm.v1.attention.backends.utils import (
21
CommonAttentionMetadata,
22
split_decodes_and_prefills,
23
)
24
-from vllm.v1.kv_cache_interface import AttentionSpec
24
+from vllm.v1.kv_cache_interface import AttentionSpec, CrossAttentionSpec
25
26
logger = init_logger(__name__)
27
@@ -50,11 +50,13 @@ class CPUAttentionBackend(AttentionBackend):
50
51
@classmethod
52
def supports_attn_type(cls, attn_type: str) -> bool:
53
- """CPU attention supports decoder and encoder-only attention."""
53
+ """CPU attention supports decoder,
54
+ encoder-only and encoder-decoder attention."""
55
return attn_type in (
56
AttentionType.DECODER,
57
AttentionType.ENCODER,
58
AttentionType.ENCODER_ONLY,
59
+ AttentionType.ENCODER_DECODER,
60
)
61
62
@staticmethod
@@ -136,6 +138,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
138
self.window_size = -1
139
self.block_size = vllm_config.cache_config.block_size
140
self.isa = _get_attn_isa(self.dtype, self.block_size)
141
+ self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec)
142
143
def build(
144
self,
@@ -151,7 +154,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
154
seq_lens = common_attn_metadata.seq_lens
155
block_table_tensor = common_attn_metadata.block_table_tensor
156
slot_mapping = common_attn_metadata.slot_mapping
154
- causal = common_attn_metadata.causal
157
+ causal = False if self.is_cross_attention else common_attn_metadata.causal
158
159
sdpa_start_loc = query_start_loc
160
num_decode_tokens = 0
@@ -171,22 +174,19 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
174
query_start_loc = query_start_loc[: num_decodes + 1]
175
block_table_tensor = block_table_tensor[:num_decodes]
176
174
- sheduler_metadata = None
175
- if causal:
176
- # for decode batch, use the custom kernel
177
- sheduler_metadata = ops.cpu_attn_get_scheduler_metadata(
178
- num_reqs=num_reqs,
179
- num_heads=self.num_heads,
180
- num_kv_heads=self.num_kv_heads,
181
- head_dim=self.head_dim,
182
- seq_lens=seq_lens,
183
- dtype=self.dtype,
184
- query_start_loc=query_start_loc,
185
- causal=causal,
186
- sliding_window_size=self.window_size,
187
- isa=self.isa,
188
- enable_kv_split=True,
189
- )
177
+ sheduler_metadata = ops.cpu_attn_get_scheduler_metadata(
178
+ num_reqs=num_reqs,
179
+ num_heads=self.num_heads,
180
+ num_kv_heads=self.num_kv_heads,
181
+ head_dim=self.head_dim,
182
+ seq_lens=seq_lens,
183
+ dtype=self.dtype,
184
+ query_start_loc=query_start_loc,
185
+ causal=causal,
186
+ sliding_window_size=self.window_size,
187
+ isa=self.isa,
188
+ enable_kv_split=True,
189
+ )
190
191
attn_metadata = CPUAttentionMetadata(
192
isa=self.isa,
vllm/v1/worker/utils.py
+6
-2
@@ -313,8 +313,12 @@ def bind_kv_cache(
313
# TODO - analyze where runner_kv_caches is used and the right
314
# way to ensure it properly reflects multiple attention layers
315
# in the same decoder block.
316
- if current_platform.is_cuda_alike() or current_platform.is_xpu():
317
- # We know that the GPU runner is not impacted by this
316
+ if (
317
+ current_platform.is_cuda_alike()
318
+ or current_platform.is_xpu()
319
+ or current_platform.is_cpu()
320
+ ):
321
+ # We know that the GPU / CPU runner is not impacted by this
322
# case. Some test code depends on runner_kv_caches, but
323
# not in a way that's impacted by ignoring this.
324
pass