@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4555143ea7
Fork 0

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 15, 07:15 UTC

[CPU] V1 support for the CPU backend (#16441)

Li, Jiang committed Jun 4, 2025 at 09:43 UTC 4555143ea7fdd2b2f0106e40889bfbab49879237
15 files changed +465 -40
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+5 -8
@@ -6,6 +6,7 @@ set -ex
6
7 # allow to bind to different cores
8 CORE_RANGE=${CORE_RANGE:-48-95}
9 +OMP_CORE_RANGE=${OMP_CORE_RANGE:-48-95}
10 NUMA_NODE=${NUMA_NODE:-1}
11
12 export CMAKE_BUILD_PARALLEL_LEVEL=32
@@ -23,10 +24,8 @@ numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --tag cpu-test-"$NUMA_NODE
24 numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" --tag cpu-test-"$NUMA_NODE"-avx2 --target vllm-test -f docker/Dockerfile.cpu .
25
26 # Run the image, setting --shm-size=4g for tensor parallel.
26 -docker run -itd --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --cpuset-cpus="$CORE_RANGE" \
27 - --cpuset-mems="$NUMA_NODE" --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"
28 -docker run -itd --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --cpuset-cpus="$CORE_RANGE" \
29 - --cpuset-mems="$NUMA_NODE" --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --shm-size=4g --name cpu-test-"$NUMA_NODE"-avx2 cpu-test-"$NUMA_NODE"-avx2
27 +docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --env VLLM_CPU_OMP_THREADS_BIND="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"
28 +docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --env VLLM_CPU_OMP_THREADS_BIND="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE"-avx2 cpu-test-"$NUMA_NODE"-avx2
29
30 function cpu_tests() {
31 set -e
@@ -56,7 +55,7 @@ function cpu_tests() {
55 # Run AWQ test
56 docker exec cpu-test-"$NUMA_NODE" bash -c "
57 set -e
59 - pytest -s -v \
58 + VLLM_USE_V1=0 pytest -s -v \
59 tests/quantization/test_ipex_quant.py"
60
61 # Run chunked-prefill and prefix-cache test
@@ -68,8 +67,6 @@ function cpu_tests() {
67 # online serving
68 docker exec cpu-test-"$NUMA_NODE" bash -c "
69 set -e
71 - export VLLM_CPU_KVCACHE_SPACE=10
72 - export VLLM_CPU_OMP_THREADS_BIND=$1
70 python3 -m vllm.entrypoints.openai.api_server --model facebook/opt-125m --dtype half &
71 timeout 600 bash -c 'until curl localhost:8000/v1/models; do sleep 1; done' || exit 1
72 python3 benchmarks/benchmark_serving.py \
@@ -89,4 +86,4 @@ function cpu_tests() {
86
87 # All of CPU tests are expected to be finished less than 40 mins.
88 export -f cpu_tests
92 -timeout 40m bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
89 +timeout 1h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
docs/usage/v1_guide.md
+2
@@ -40,6 +40,8 @@ This living user guide outlines a few known **important changes and limitations*
40 | **NVIDIA** | <nobr>🚀 Natively Supported</nobr> |
41 | **AMD** | <nobr>🚧 WIP</nobr> |
42 | **TPU** | <nobr>🚧 WIP</nobr> |
43 +| **CPU** | <nobr>🚧 WIP</nobr> |
44 +
45 #### Feature / Model
46
47 | Feature / Model | Status |
requirements/cpu.txt
+3
@@ -1,6 +1,9 @@
1 # Common dependencies
2 -r common.txt
3
4 +numba == 0.60.0; python_version == '3.9' # v0.61 doesn't support Python 3.9. Required for N-gram speculative decoding
5 +numba == 0.61.2; python_version > '3.9'
6 +
7 # Dependencies for CPUs
8 packaging>=24.2
9 setuptools>=77.0.3,<80.0.0
tests/kernels/attention/test_attention_selector.py
+4 -1
@@ -85,7 +85,10 @@ def test_env(
85 CpuPlatform()):
86 backend = get_attn_backend(16, torch.float16, torch.float16,
87 block_size, False)
88 - assert backend.get_name() == "TORCH_SDPA"
88 + if use_v1:
89 + assert backend.get_name() == "TORCH_SDPA_VLLM_V1"
90 + else:
91 + assert backend.get_name() == "TORCH_SDPA"
92
93 elif device == "hip":
94 with patch("vllm.attention.selector.current_platform",
tests/models/language/generation/test_common.py
-1
@@ -87,7 +87,6 @@ AITER_MODEL_LIST = [
87 pytest.param("bigcode/starcoder2-3b"), # starcoder2
88 pytest.param(
89 "TitanML/tiny-mixtral", # mixtral
90 - marks=[pytest.mark.cpu_model],
90 )
91 ])
92 @pytest.mark.parametrize("max_tokens", [32])
vllm/attention/backends/cpu_mla.py
+3 -3
@@ -178,7 +178,7 @@ class CPUMLAMetadataBuilder(AttentionMetadataBuilder[CPUMLAMetadata]):
178 seq_lens_tensor=seq_lens_tensor,
179 max_query_len=max_query_len,
180 max_kv_len=max_kv_len,
181 - query_start_loc=query_start_loc,
181 + prefill_query_start_loc=query_start_loc,
182 kv_start_loc=kv_start_loc,
183 max_decode_seq_len=input_data.max_decode_seq_len,
184 num_prefills=input_data.num_prefills,
@@ -264,8 +264,8 @@ class CPUMLAImpl(MLACommonImpl[CPUMLAMetadata]):
264 key=k,
265 value=v_padded,
266 out=output,
267 - seqlen_q=prefill_metadata.query_start_loc,
268 - seqlen_k=prefill_metadata.query_start_loc,
267 + seqlen_q=prefill_metadata.prefill_query_start_loc,
268 + seqlen_k=prefill_metadata.prefill_query_start_loc,
269 max_seqlen_q=prefill_metadata.max_query_len,
270 max_seqlen_k=prefill_metadata.max_query_len,
271 pdropout=0.0,
vllm/attention/backends/torch_sdpa.py
+12 -4
@@ -87,10 +87,13 @@ class TorchSDPAMetadata(AttentionMetadata, PagedAttentionMetadata):
87 # For chunked prefill only
88 max_query_len: Optional[int] = None
89 max_kv_len: Optional[int] = None
90 - query_start_loc: Optional[torch.Tensor] = None
90 + prefill_query_start_loc: Optional[torch.Tensor] = None
91 kv_start_loc: Optional[torch.Tensor] = None
92 prefill_block_tables: Optional[torch.Tensor] = None
93
94 + # For V1 logits index only
95 + query_start_loc: Optional[torch.Tensor] = None
96 +
97 # Begin encoder attn & enc/dec cross-attn fields...
98 # Encoder sequence lengths representation
99 encoder_seq_lens: Optional[List[int]] = None
@@ -375,7 +378,7 @@ class TorchSDPAMetadataBuilder(AttentionMetadataBuilder[TorchSDPAMetadata]):
378 seq_lens_tensor=seq_lens_tensor,
379 max_query_len=max_query_len,
380 max_kv_len=max_kv_len,
378 - query_start_loc=query_start_loc,
381 + prefill_query_start_loc=query_start_loc,
382 kv_start_loc=kv_start_loc,
383 max_decode_seq_len=input_data.max_decode_seq_len,
384 num_prefills=input_data.num_prefills,
@@ -470,6 +473,11 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
473 Returns:
474 shape = [num_tokens, num_heads * head_size]
475 """
476 +
477 + # For warming-up
478 + if attn_metadata is None:
479 + return query
480 +
481 attn_type = self.attn_type
482 if (attn_type == AttentionType.ENCODER
483 and (not attn_metadata.is_all_encoder_attn_metadata_set)):
@@ -537,8 +545,8 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
545
546 output = torch.empty_like(query)
547 if prefill_meta := attn_metadata.prefill_metadata:
540 - assert attn_metadata.seq_lens is not None
548 if not prefill_meta.prefill_metadata.chunked_prefill: # type: ignore
549 + assert attn_metadata.seq_lens is not None
550 self._run_sdpa_forward(output,
551 query,
552 key,
@@ -555,7 +563,7 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
563 query[:prefill_meta.num_prefill_tokens, :, :],
564 key_cache,
565 value_cache,
558 - prefill_meta.query_start_loc,
566 + prefill_meta.prefill_query_start_loc,
567 prefill_meta.kv_start_loc,
568 prefill_meta.max_query_len,
569 prefill_meta.max_kv_len,
vllm/compilation/wrapper.py
+6 -1
@@ -41,11 +41,16 @@ class TorchCompileWrapperWithCustomDispatcher:
41 # compiling the forward method
42
43 backend = vllm_config.compilation_config.init_backend(vllm_config)
44 + options = None
45 + if isinstance(backend, str) and backend == "inductor":
46 + options = get_current_vllm_config(
47 + ).compilation_config.inductor_compile_config
48
49 compiled_callable = torch.compile(
50 self.forward,
51 fullgraph=envs.VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE,
48 - backend=backend)
52 + backend=backend,
53 + options=options)
54
55 self.compiled_callable = compiled_callable
56 self.original_code_object = self.__class__.forward.__code__
vllm/engine/arg_utils.py
+3 -1
@@ -1399,6 +1399,7 @@ class EngineArgs:
1399 "FLASHINFER",
1400 "FLASHINFER_VLLM_V1",
1401 "ROCM_AITER_MLA",
1402 + "TORCH_SDPA_VLLM_V1",
1403 ]
1404 if (envs.is_set("VLLM_ATTENTION_BACKEND")
1405 and envs.VLLM_ATTENTION_BACKEND not in V1_BACKENDS):
@@ -1431,7 +1432,8 @@ class EngineArgs:
1432
1433 # Non-[CUDA, TPU] may be supported on V1, but off by default for now.
1434 v0_hardware = not any(
1434 - (current_platform.is_cuda(), current_platform.is_tpu()))
1435 + (current_platform.is_cuda(), current_platform.is_tpu(),
1436 + current_platform.is_cpu()))
1437 if v0_hardware and _warn_or_fallback( # noqa: SIM103
1438 current_platform.device_name):
1439 return False
vllm/platforms/cpu.py
+57 -10
@@ -57,7 +57,10 @@ class CpuPlatform(Platform):
57 logger.info("Using CPU MLA backend.")
58 return "vllm.attention.backends.cpu_mla.CPUMLABackend"
59 logger.info("Using Torch SDPA backend.")
60 - return "vllm.attention.backends.torch_sdpa.TorchSDPABackend"
60 + if use_v1:
61 + return "vllm.v1.attention.backends.cpu_attn.TorchSDPABackend"
62 + else:
63 + return "vllm.attention.backends.torch_sdpa.TorchSDPABackend"
64
65 @classmethod
66 def get_device_total_memory(cls, device_id: int = 0) -> int:
@@ -81,6 +84,8 @@ class CpuPlatform(Platform):
84 if not model_config.enforce_eager:
85 model_config.enforce_eager = True
86
87 + model_config.disable_cascade_attn = True
88 +
89 cache_config = vllm_config.cache_config
90
91 ipex_available = find_spec("intel_extension_for_pytorch") is not None
@@ -128,7 +133,8 @@ class CpuPlatform(Platform):
133 f" {kv_cache_space}, expect a positive integer value.")
134
135 parallel_config = vllm_config.parallel_config
131 - if (parallel_config.distributed_executor_backend is not None
136 + if (parallel_config.world_size > 1
137 + and parallel_config.distributed_executor_backend is not None
138 and parallel_config.distributed_executor_backend != "mp"):
139 logger.warning(("%s is not supported on CPU, fallback to mp "
140 "distributed executor backend."),
@@ -141,7 +147,38 @@ class CpuPlatform(Platform):
147 parallel_config.sd_worker_cls = \
148 "vllm.worker.cpu_worker.CPUWorker"
149 else:
144 - parallel_config.worker_cls = "vllm.worker.cpu_worker.CPUWorker"
150 + if envs.VLLM_USE_V1:
151 + parallel_config.worker_cls = \
152 + "vllm.v1.worker.cpu_worker.CPUWorker"
153 + else:
154 + parallel_config.worker_cls = \
155 + "vllm.worker.cpu_worker.CPUWorker"
156 +
157 + # Note: workaround for v1 gpu_model_runner
158 + from vllm.config import CompilationLevel
159 + vllm_config.compilation_config.cudagraph_capture_sizes = []
160 +
161 + compilation_config = vllm_config.compilation_config
162 + if (envs.VLLM_USE_V1 and vllm_config.compilation_config.level
163 + == CompilationLevel.PIECEWISE):
164 + compilation_config.level = CompilationLevel.DYNAMO_ONCE
165 + compilation_config.backend = "eager"
166 + compilation_config.custom_ops += ["none"]
167 + compilation_config.inductor_compile_config.update({
168 + "dce":
169 + True,
170 + "size_asserts":
171 + False,
172 + "nan_asserts":
173 + False,
174 + "memory_planning":
175 + True,
176 + "epilogue_fusion":
177 + True,
178 + })
179 +
180 + if vllm_config.lora_config is not None:
181 + compilation_config.level = CompilationLevel.NO_COMPILATION
182
183 assert vllm_config.device_config.device_type == "cpu"
184
@@ -149,6 +186,12 @@ class CpuPlatform(Platform):
186 # Environment variables for CPU executor
187 #
188
189 + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
190 +
191 + # Note: to avoid the error 'nthreads cannot be larger than environment
192 + # variable "NUMEXPR_MAX_THREADS" (64)'.
193 + os.environ["NUMEXPR_MAX_THREADS"] = str(len(os.sched_getaffinity(0)))
194 +
195 # Set default threads num for OpenMP parallel
196 os.environ["OMP_NUM_THREADS"] = str(torch.get_num_threads())
197
@@ -171,13 +214,6 @@ class CpuPlatform(Platform):
214 # To hint IPEX uses shared memory based AllReduce
215 os.environ["LOCAL_WORLD_SIZE"] = str(
216 vllm_config.parallel_config.tensor_parallel_size)
174 - if sys.platform == "darwin" and \
175 - envs.VLLM_WORKER_MULTIPROC_METHOD == "fork":
176 - if os.environ.get('VLLM_WORKER_MULTIPROC_METHOD', None) is None:
177 - logger.warning(
178 - "Default to spawn method on MacOS. If this is not desired,"
179 - " set VLLM_WORKER_MULTIPROC_METHOD to fork explicitly.")
180 - os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'
217
218 if vllm_config.model_config and vllm_config.model_config.use_mla:
219 logger.info(
@@ -204,3 +240,14 @@ class CpuPlatform(Platform):
240 Get device specific communicator class for distributed communication.
241 """
242 return "vllm.distributed.device_communicators.cpu_communicator.CpuCommunicator" # noqa
243 +
244 + @classmethod
245 + def supports_structured_output(cls) -> bool:
246 + return True
247 +
248 + @classmethod
249 + def supports_v1(cls, model_config) -> bool:
250 + """Returns whether the current platform can support v1 for the supplied
251 + model configuration.
252 + """
253 + return True
vllm/v1/attention/backends/cpu_attn.py new
+163
@@ -0,0 +1,163 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import numpy as np
3 +import torch
4 +
5 +from vllm.attention.backends.abstract import AttentionMetadata
6 +from vllm.attention.backends.torch_sdpa import (TorchSDPABackendImpl,
7 + TorchSDPAMetadata)
8 +from vllm.attention.backends.utils import CommonAttentionState
9 +from vllm.attention.ops.ipex_attn import PagedAttention
10 +from vllm.v1.attention.backends.utils import CommonAttentionMetadata
11 +from vllm.v1.core.sched.output import SchedulerOutput
12 +from vllm.v1.kv_cache_interface import AttentionSpec
13 +from vllm.v1.worker.block_table import BlockTable
14 +from vllm.v1.worker.cpu_model_runner import CPUModelRunner
15 +from vllm.v1.worker.gpu_input_batch import InputBatch
16 +
17 +
18 +class TorchSDPABackend:
19 + accept_output_buffer: bool = False
20 +
21 + @staticmethod
22 + def get_name() -> str:
23 + return "TORCH_SDPA_VLLM_V1"
24 +
25 + @staticmethod
26 + def get_impl_cls() -> type["TorchSDPABackendImpl"]:
27 + return TorchSDPABackendImpl
28 +
29 + @staticmethod
30 + def get_metadata_cls() -> type["AttentionMetadata"]:
31 + return TorchSDPAMetadata
32 +
33 + @staticmethod
34 + def get_state_cls() -> type["CommonAttentionState"]:
35 + return CommonAttentionState
36 +
37 + @staticmethod
38 + def get_builder_cls() -> type["TorchSDPAMetadataBuilderV1"]:
39 + return TorchSDPAMetadataBuilderV1
40 +
41 + @staticmethod
42 + def get_kv_cache_shape(
43 + num_blocks: int,
44 + block_size: int,
45 + num_kv_heads: int,
46 + head_size: int,
47 + ) -> tuple[int, ...]:
48 + return PagedAttention.get_kv_cache_shape(num_blocks, block_size,
49 + num_kv_heads, head_size)
50 +
51 + @staticmethod
52 + def use_cascade_attention(*args, **kwargs) -> bool:
53 + return False
54 +
55 +
56 +class TorchSDPAMetadataBuilderV1:
57 +
58 + def __init__(self, runner: CPUModelRunner, kv_cache_spec: AttentionSpec,
59 + block_table: BlockTable) -> None:
60 + self.runner = runner
61 + self.block_table = block_table
62 +
63 + # For reorder
64 + self.reorder_prompt_req_index_list = np.empty(self.runner.max_num_reqs,
65 + dtype=np.int64)
66 + self.reorder_decode_req_index_list = np.empty(self.runner.max_num_reqs,
67 + dtype=np.int64)
68 + self.num_prompt_req: int = 0
69 +
70 + self.seq_start_loc_cpu = torch.zeros(
71 + runner.max_num_reqs + 1,
72 + dtype=torch.int32,
73 + device="cpu",
74 + )
75 + self.seq_start_loc_np = self.seq_start_loc_cpu.numpy()
76 +
77 + def reorder_batch(self, input_batch: InputBatch,
78 + scheduler_output: SchedulerOutput) -> bool:
79 + prompt_list_idx = 0
80 + decode_list_idx = 0
81 + for req_index in range(input_batch.num_reqs):
82 + if input_batch.num_computed_tokens_cpu[
83 + req_index] < input_batch.num_prompt_tokens[req_index]:
84 + # prompt stage
85 + self.reorder_prompt_req_index_list[prompt_list_idx] = req_index
86 + prompt_list_idx += 1
87 + else:
88 + # decode stage
89 + self.reorder_decode_req_index_list[decode_list_idx] = req_index
90 + decode_list_idx += 1
91 + assert decode_list_idx + prompt_list_idx == input_batch.num_reqs
92 +
93 + # Update prompt requests number
94 + self.num_prompt_req = prompt_list_idx
95 +
96 + reorder_req_num = 0
97 + for req_index in range(decode_list_idx):
98 + if self.reorder_decode_req_index_list[req_index] < prompt_list_idx:
99 + reorder_req_num += 1
100 + else:
101 + break
102 +
103 + if reorder_req_num == 0:
104 + return False
105 +
106 + reorder_prompt_list = (
107 + self.reorder_prompt_req_index_list[:prompt_list_idx]
108 + [-reorder_req_num:])
109 + reorder_decode_list = (
110 + self.reorder_decode_req_index_list[:decode_list_idx]
111 + [:reorder_req_num])
112 + assert reorder_decode_list.size == reorder_prompt_list.size
113 +
114 + for idx in range(reorder_req_num):
115 + prompt_req_index = reorder_prompt_list[idx].item()
116 + decode_req_index = reorder_decode_list[idx].item()
117 + input_batch.swap_states(prompt_req_index, decode_req_index)
118 +
119 + return True
120 +
121 + def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
122 + common_prefix_len: int,
123 + common_attn_metadata: CommonAttentionMetadata):
124 + runner = self.runner
125 + block_table = self.block_table
126 + seq_lens_np = runner.seq_lens_np[:num_reqs]
127 + num_prompt_req = self.num_prompt_req
128 + max_prefill_seq_len = seq_lens_np[:num_prompt_req].max().item(
129 + ) if num_prompt_req > 0 else 0
130 + max_decode_seq_len = seq_lens_np[num_prompt_req:num_reqs].max().item(
131 + ) if num_prompt_req < num_reqs else 0
132 + self.seq_start_loc_np[0] = 0
133 + np.cumsum(seq_lens_np, out=self.seq_start_loc_np[1:num_reqs + 1])
134 + num_prefill_tokens = runner.query_start_loc_np[num_prompt_req].item()
135 + num_decode_tokens = runner.query_start_loc_np[num_reqs].item(
136 + ) - num_prefill_tokens
137 + slot_mapping = block_table.slot_mapping_cpu[:num_actual_tokens].long()
138 + block_table_tensor = block_table.get_device_tensor()
139 + attn_metadata = TorchSDPAMetadata(
140 + num_prefills=num_prompt_req,
141 + num_prefill_tokens=num_prefill_tokens,
142 + num_decode_tokens=num_decode_tokens,
143 + slot_mapping=slot_mapping,
144 + seq_lens_tensor=runner.
145 + seq_lens_cpu[num_prompt_req:num_reqs], # decode
146 + max_decode_seq_len=max_decode_seq_len, # decode
147 + block_tables=block_table_tensor[num_prompt_req:num_reqs], # decode
148 + chunked_prefill=True,
149 + max_query_len=max_query_len,
150 + max_kv_len=max_prefill_seq_len,
151 + prefill_query_start_loc=runner.
152 + query_start_loc_cpu[:num_prompt_req + 1], # prefill
153 + kv_start_loc=self.seq_start_loc_cpu[:num_prompt_req +
154 + 1], # prefill
155 + prefill_block_tables=block_table_tensor[:
156 + num_prompt_req], # prefill
157 + query_start_loc=runner.query_start_loc_cpu[:num_reqs +
158 + 1], # for logits index
159 + multi_modal_placeholder_index_maps=None,
160 + enable_kv_scales_calculation=False,
161 + )
162 +
163 + return attn_metadata
vllm/v1/worker/cpu_model_runner.py new
+86
@@ -0,0 +1,86 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from contextlib import contextmanager
3 +from typing import Any
4 +
5 +import torch
6 +
7 +from vllm.config import VllmConfig
8 +from vllm.logger import init_logger
9 +from vllm.model_executor.model_loader import get_model
10 +from vllm.v1.worker.gpu_model_runner import GPUModelRunner
11 +
12 +logger = init_logger(__name__)
13 +
14 +
15 +class CPUModelRunner(GPUModelRunner):
16 +
17 + def __init__(self, vllm_config: VllmConfig, device: torch.device):
18 + super().__init__(vllm_config, device)
19 +
20 + assert device == torch.device("cpu")
21 + assert self.speculative_config is None, "spec decode is not supported."
22 +
23 + self.use_cuda_graph = False
24 + self.cascade_attn_enabled = False
25 +
26 + self._postprocess_tenosrs()
27 +
28 + def _postprocess_tenosrs(self) -> None:
29 + # Note: replace device tensors with cpu tensors
30 + def replace_tensor(obj: Any, cpu_attr_name: str,
31 + device_attr_name) -> None:
32 + cpu_tensor = getattr(obj, cpu_attr_name, None)
33 + device_tensor = getattr(obj, device_attr_name, None)
34 + if cpu_tensor is not None and device_tensor is not None:
35 + assert isinstance(cpu_tensor, torch.Tensor)
36 + assert isinstance(device_tensor, torch.Tensor)
37 + setattr(obj, device_attr_name, cpu_tensor)
38 +
39 + for k, v in vars(self).items():
40 + if k.endswith("_cpu") and isinstance(v, torch.Tensor):
41 + replace_tensor(self, k, k[:-4])
42 +
43 + for k, v in vars(self.input_batch).items():
44 + if k.endswith("_cpu_tensor") and isinstance(v, torch.Tensor):
45 + replace_tensor(self.input_batch, k, k[:-11])
46 +
47 + for k, v in vars(self.input_batch.block_table).items():
48 + if k.endswith("_cpu") and isinstance(v, torch.Tensor):
49 + replace_tensor(self.input_batch.block_table, k, k[:-4])
50 +
51 + def load_model(self) -> None:
52 + logger.info("Starting to load model %s...", self.model_config.model)
53 + self.model = get_model(vllm_config=self.vllm_config)
54 +
55 + if self.lora_config:
56 + self.model = self.load_lora_model(self.model, self.model_config,
57 + self.scheduler_config,
58 + self.lora_config, self.device)
59 +
60 + def warming_up_model(self) -> None:
61 + logger.info("Warming up model for the compilation...")
62 + # Only generate graph for the generic shape
63 + self._dummy_run(max(16, self.max_num_reqs))
64 + logger.info("Warming up done.")
65 +
66 + def _init_device_properties(self) -> None:
67 + pass
68 +
69 + def _sync_device(self) -> None:
70 + pass
71 +
72 +
73 +@contextmanager
74 +def _set_global_compilation_settings():
75 + import torch._inductor.config
76 +
77 + # Note: The CPPGEMM backend requires freezing parameters.
78 + freezing_value = torch._inductor.config.freezing
79 + torch._inductor.config.freezing = True
80 + # Note: workaround for "ValueError: fast mode: can't pickle cyclic objects
81 + # including object type dict"
82 + force_disable_caches = torch._inductor.config.force_disable_caches
83 + torch._inductor.config.force_disable_caches = True
84 + yield
85 + torch._inductor.config.freezing = freezing_value
86 + torch._inductor.config.force_disable_caches = force_disable_caches
vllm/v1/worker/cpu_worker.py new
+101
@@ -0,0 +1,101 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import os
3 +from typing import Optional
4 +
5 +import torch
6 +
7 +from vllm import envs
8 +from vllm.config import VllmConfig
9 +from vllm.distributed.parallel_state import get_pp_group, get_tp_group
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.utils import set_random_seed
12 +from vllm.sequence import IntermediateTensors
13 +from vllm.v1.core.sched.output import SchedulerOutput
14 +from vllm.v1.outputs import ModelRunnerOutput
15 +from vllm.v1.worker.cpu_model_runner import CPUModelRunner
16 +from vllm.v1.worker.gpu_worker import (Worker,
17 + init_worker_distributed_environment)
18 +
19 +logger = init_logger(__name__)
20 +
21 +
22 +class CPUWorker(Worker):
23 +
24 + def __init__(self,
25 + vllm_config: VllmConfig,
26 + local_rank: int,
27 + rank: int,
28 + distributed_init_method: str,
29 + is_driver_worker: bool = False):
30 + super().__init__(vllm_config,
31 + local_rank,
32 + rank,
33 + distributed_init_method,
34 + is_driver_worker=is_driver_worker)
35 +
36 + self.parallel_config.disable_custom_all_reduce = True
37 +
38 + def init_device(self):
39 + # Setup OpenMP threads affinity.
40 + omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND
41 + if omp_cpuids == "all":
42 + self.local_omp_cpuid = "all"
43 + else:
44 + self.local_omp_cpuid = omp_cpuids.split("|")[self.rank]
45 + ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)
46 + if ret:
47 + logger.info(ret)
48 +
49 + # Note: unique identifier for creating allreduce shared memory
50 + os.environ["VLLM_DIST_IDENT"] = self.distributed_init_method.split(
51 + ":")[-1]
52 + # Initialize the distributed environment.
53 + init_worker_distributed_environment(self.vllm_config, self.rank,
54 + self.distributed_init_method,
55 + self.local_rank, "gloo")
56 + # Set random seed.
57 + set_random_seed(self.model_config.seed)
58 +
59 + # Construct the model runner
60 + self.model_runner: CPUModelRunner = CPUModelRunner(
61 + self.vllm_config, torch.device("cpu"))
62 +
63 + def sleep(self, level: int = 1) -> None:
64 + logger.warning("sleep mode is not supported on CPU, ignore it.")
65 + pass
66 +
67 + def wake_up(self, tags: Optional[list[str]] = None) -> None:
68 + logger.warning("sleep mode is not supported on CPU, ignore it.")
69 + pass
70 +
71 + def determine_available_memory(self) -> int:
72 + return self.cache_config.cpu_kvcache_space_bytes # type: ignore
73 +
74 + def compile_or_warm_up_model(self) -> None:
75 + # Reset the seed to ensure that the random state is not affected by
76 + # the model initialization and profiling.
77 + set_random_seed(self.model_config.seed)
78 + self.model_runner.warming_up_model()
79 +
80 + @torch.inference_mode()
81 + def execute_model(
82 + self,
83 + scheduler_output: "SchedulerOutput",
84 + ) -> Optional[ModelRunnerOutput]:
85 + intermediate_tensors = None
86 + if not get_pp_group().is_first_rank:
87 + intermediate_tensors = IntermediateTensors(
88 + get_pp_group().recv_tensor_dict(
89 + all_gather_group=get_tp_group()))
90 +
91 + output = self.model_runner.execute_model(scheduler_output,
92 + intermediate_tensors)
93 +
94 + if not get_pp_group().is_last_rank:
95 + assert isinstance(output, IntermediateTensors)
96 + get_pp_group().send_tensor_dict(output.tensors,
97 + all_gather_group=get_tp_group())
98 + return None
99 +
100 + assert isinstance(output, ModelRunnerOutput)
101 + return output if self.is_driver_worker else None
vllm/v1/worker/gpu_model_runner.py
+18 -10
@@ -5,7 +5,7 @@ import copy
5 import gc
6 import time
7 import weakref
8 -from typing import TYPE_CHECKING, Optional, Union
8 +from typing import TYPE_CHECKING, Any, Optional, Union
9
10 import numpy as np
11 import torch
@@ -38,7 +38,6 @@ from vllm.sequence import IntermediateTensors
38 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
39 GiB_bytes, LazyLoader, async_tensor_h2d, cdiv,
40 check_use_alibi, is_pin_memory_available)
41 -from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
41 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
42 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
43 from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,
@@ -203,8 +202,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
202 self.vllm_config.compilation_config.cudagraph_capture_sizes))
203
204 # Cache the device properties.
206 - self.device_properties = torch.cuda.get_device_properties(self.device)
207 - self.num_sms = self.device_properties.multi_processor_count
205 + self._init_device_properties()
206
207 # Persistent buffers for CUDA graphs.
208 self.input_ids = torch.zeros(self.max_num_tokens,
@@ -315,6 +313,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
313 self.input_batch, scheduler_output)
314 return batch_reordered
315
316 + # Note: used for model runner override.
317 + def _init_device_properties(self) -> None:
318 + """Initialize attributes from torch.cuda.get_device_properties
319 + """
320 + self.device_properties = torch.cuda.get_device_properties(self.device)
321 + self.num_sms = self.device_properties.multi_processor_count
322 +
323 + # Note: used for model runner override.
324 + def _sync_device(self) -> None:
325 + torch.cuda.synchronize()
326 +
327 def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
328 """Update the cached states and the persistent batch with the scheduler
329 output.
@@ -538,8 +547,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
547 def _prepare_inputs(
548 self,
549 scheduler_output: "SchedulerOutput",
541 - ) -> tuple[dict[str, FlashAttentionMetadata], torch.Tensor,
542 - Optional[SpecDecodeMetadata]]:
550 + ) -> tuple[dict[str, Any], torch.Tensor, Optional[SpecDecodeMetadata]]:
551 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
552 assert total_num_scheduled_tokens > 0
553 num_reqs = self.input_batch.num_reqs
@@ -652,7 +660,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
660 common_attn_metadata = CommonAttentionMetadata(
661 query_start_loc=query_start_loc, seq_lens=seq_lens)
662
655 - attn_metadata: dict[str, FlashAttentionMetadata] = {}
663 + attn_metadata: dict[str, Any] = {}
664 # Prepare the attention metadata for each KV cache group and make layers
665 # in the same group share the same metadata.
666 for kv_cache_group_id, kv_cache_group_spec in enumerate(
@@ -1710,7 +1718,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1718
1719 # Must synchronize the non-blocking GPU->CPU transfers.
1720 if prompt_logprobs_dict:
1713 - torch.cuda.synchronize()
1721 + self._sync_device()
1722
1723 return prompt_logprobs_dict
1724
@@ -1740,7 +1748,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1748 dtype=np.int32)
1749
1750 if skip_attn:
1743 - attn_metadata: Optional[dict[str, FlashAttentionMetadata]] = None
1751 + attn_metadata: Optional[dict[str, Any]] = None
1752 else:
1753 query_start_loc = self.query_start_loc[:num_reqs + 1]
1754 seq_lens = self.seq_lens[:num_reqs]
@@ -1964,7 +1972,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1972 sampler_output = self._dummy_sampler_run(hidden_states)
1973 else:
1974 sampler_output = None
1967 - torch.cuda.synchronize()
1975 + self._sync_device()
1976 del hidden_states, sampler_output
1977 self.encoder_cache.clear()
1978 gc.collect()
vllm/v1/worker/gpu_worker.py
+2 -1
@@ -342,13 +342,14 @@ def init_worker_distributed_environment(
342 rank: int,
343 distributed_init_method: Optional[str] = None,
344 local_rank: int = -1,
345 + backend: str = "nccl",
346 ) -> None:
347 """Initialize the distributed environment."""
348 parallel_config = vllm_config.parallel_config
349 set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
350
351 init_distributed_environment(parallel_config.world_size, rank,
351 - distributed_init_method, local_rank)
352 + distributed_init_method, local_rank, backend)
353
354 ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
355 parallel_config.pipeline_parallel_size)