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 3, 01:45 UTC
[CPU] Experimentally enable Triton and MRV2 (#43225)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
May 22, 2026 at 16:48 UTC
65b7a812a2dabd212d78c7b5b8a320b4efb9750d
11 files changed
+169
-13
.buildkite/hardware_tests/cpu.yaml
+14
@@ -54,6 +54,20 @@ steps:
54
pytest -x -v -s tests/models/language/generation -m cpu_model
55
pytest -x -v -s tests/models/language/pooling -m cpu_model"
56
57
+- label: CPU-ModelRunnerV2 Tests
58
+ depends_on: []
59
+ device: intel_cpu
60
+ no_plugin: true
61
+ soft_fail: true
62
+ source_file_dependencies:
63
+ - vllm/v1/worker/cpu/
64
+ - vllm/v1/worker/gpu/
65
+ commands:
66
+ - |
67
+ bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
68
+ uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
69
+ VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model"
70
+
71
- label: CPU-Quantization Model Tests
72
depends_on: []
73
device: intel_cpu
docker/Dockerfile.cpu
+5
-6
@@ -27,11 +27,14 @@ WORKDIR /workspace
27
ARG PYTHON_VERSION=3.12
28
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
29
30
+ARG max_jobs=32
31
+ENV MAX_JOBS=${max_jobs}
32
+
33
# Install minimal dependencies and uv
34
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
35
--mount=type=cache,target=/var/lib/apt,sharing=locked \
36
apt-get update -y \
34
- && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates \
37
+ && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates zlib1g-dev \
38
gcc-12 g++-12 libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \
39
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
40
&& curl -LsSf https://astral.sh/uv/install.sh | sh
@@ -123,9 +126,6 @@ RUN --mount=type=cache,target=/root/.cargo/registry \
126
######################### BUILD IMAGE #########################
127
FROM base AS vllm-build
128
126
-ARG max_jobs=32
127
-ENV MAX_JOBS=${max_jobs}
128
-
129
ARG GIT_REPO_CHECK=0
130
# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
131
ARG VLLM_CPU_X86=0
@@ -257,8 +257,7 @@ WORKDIR /vllm-workspace
257
RUN --mount=type=cache,target=/root/.cache/uv \
258
--mount=type=cache,target=/root/.cache/ccache \
259
--mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
260
- uv pip install dist/*.whl && \
261
- uv pip install "vllm[audio]"
260
+ uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]"
261
262
# Add labels to document build configuration
263
LABEL org.opencontainers.image.title="vLLM CPU"
setup.py
+5
@@ -1195,6 +1195,11 @@ setup(
1195
"opentelemetry-exporter-otlp>=1.26.0",
1196
"opentelemetry-semantic-conventions-ai>=0.4.1",
1197
],
1198
+ "triton-cpu": [
1199
+ "triton @ "
1200
+ "git+https://github.com/triton-lang/triton-cpu.git@270e696d ; "
1201
+ "platform_machine == 'x86_64'",
1202
+ ], # Remove after stable release
1203
},
1204
cmdclass=cmdclass,
1205
package_data=package_data,
vllm/platforms/cpu.py
+20
-3
@@ -14,6 +14,7 @@ from vllm.logger import init_logger
14
from vllm.utils.cpu_resource_utils import (
15
DEVICE_CONTROL_ENV_VAR,
16
get_memory_node_info,
17
+ get_visible_memory_node,
18
)
19
from vllm.utils.mem_constants import GiB_bytes
20
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -135,9 +136,13 @@ class CpuPlatform(Platform):
136
scheduler_config.async_scheduling = False
137
138
parallel_config = vllm_config.parallel_config
138
- # OMP requires the MP executor to function correctly, UniProc is not
139
- # supported as it is not possible to set the OMP environment correctly
140
- if parallel_config.distributed_executor_backend == "uni":
139
+ if (
140
+ os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING", "1") == "1"
141
+ and parallel_config.distributed_executor_backend == "uni"
142
+ ):
143
+ # OMP requires the MP executor to function correctly, UniProc
144
+ # is not supported as it is not possible to set the OMP
145
+ # environment correctly
146
parallel_config.distributed_executor_backend = "mp"
147
148
if parallel_config.worker_cls == "auto":
@@ -481,3 +486,15 @@ class CpuPlatform(Platform):
486
slot_mapping,
487
isa,
488
)
489
+
490
+ @classmethod
491
+ def get_current_memory_usage(
492
+ cls, device: torch.types.Device | None = None
493
+ ) -> float:
494
+ allowed_mem_node_list = get_visible_memory_node()
495
+ mem_status_list = [get_memory_node_info(i) for i in allowed_mem_node_list]
496
+ memory_usage = 0
497
+ for s in mem_status_list:
498
+ memory_usage += s.total_memory - s.available_memory
499
+
500
+ return memory_usage
vllm/triton_utils/importing.py
+12
@@ -3,6 +3,7 @@
3
4
import os
5
import types
6
+from importlib.metadata import version
7
from importlib.util import find_spec
8
9
from vllm.logger import init_logger
@@ -48,6 +49,17 @@ if HAS_TRITON:
49
len(active_drivers),
50
)
51
HAS_TRITON = False
52
+
53
+ # Check Triton CPU
54
+ if "cpu" in version("vllm"):
55
+ if "cpu" in backends:
56
+ HAS_TRITON = True
57
+ else:
58
+ logger.warning(
59
+ "Triton is installed, but doesn't include CPU backend. "
60
+ "Disabling Triton."
61
+ )
62
+ HAS_TRITON = False
63
except ImportError:
64
# This can occur if Triton is partially installed or triton.backends
65
# is missing.
vllm/utils/platform_utils.py
+3
-1
@@ -50,8 +50,10 @@ def is_pin_memory_available() -> bool:
50
def is_uva_available() -> bool:
51
"""Check if Unified Virtual Addressing (UVA) is available."""
52
# UVA requires pinned memory.
53
+ from vllm.platforms import current_platform
54
+
55
# TODO: Add more requirements for UVA if needed.
54
- return is_pin_memory_available()
56
+ return is_pin_memory_available() or current_platform.is_cpu()
57
58
59
@cache
vllm/v1/worker/cpu/__init__.py
vllm/v1/worker/cpu/buffer_utils.py
new
+16
@@ -0,0 +1,16 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Sequence
4
+
5
+import torch
6
+
7
+from vllm.utils.platform_utils import is_uva_available
8
+
9
+
10
+class UvaBuffer:
11
+ def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
12
+ if not is_uva_available():
13
+ raise RuntimeError("UVA is not available")
14
+ self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
15
+ self.np = self.cpu.numpy()
16
+ self.uva = self.cpu
vllm/v1/worker/cpu/model_runner.py
new
+16
@@ -0,0 +1,16 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.logger import init_logger
5
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
6
+
7
+logger = init_logger(__name__)
8
+
9
+
10
+class CPUModelRunner(GPUModelRunner):
11
+ # TBD: Whether need to move this to Worker?
12
+ def warming_up_model(self) -> None:
13
+ logger.info("Warming up model for the compilation...")
14
+ # Only generate graph for the generic shape
15
+ self.profile_run()
16
+ logger.info("Warming up done.")
vllm/v1/worker/cpu/shm.py
new
+62
@@ -0,0 +1,62 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# isort: skip_file
5
+# ruff: noqa: E402
6
+# mypy: disable-error-code="misc, assignment"
7
+
8
+from typing import Any
9
+
10
+# Patch torch APIs
11
+import torch
12
+
13
+
14
+def noop(*args: Any, **kwargs: Any) -> None:
15
+ pass
16
+
17
+
18
+class _EventPlaceholder:
19
+ def __init__(self, *args, **kwargs) -> None:
20
+ self.record = noop
21
+ self.synchronize = noop
22
+
23
+
24
+class _StreamPlaceholder:
25
+ def __init__(self, *args, **kwargs) -> None:
26
+ self.wait_stream = noop
27
+
28
+ def __enter__(self, *args, **kwargs):
29
+ return self
30
+
31
+ def __exit__(self, exc_type, exc_val, exc_tb):
32
+ pass
33
+
34
+
35
+torch.Event = _EventPlaceholder
36
+torch.cuda.Event = _EventPlaceholder
37
+torch.cuda.Stream = _StreamPlaceholder
38
+torch.cuda.set_stream = noop
39
+torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()
40
+torch.accelerator.synchronize = noop
41
+torch.accelerator.empty_cache = noop
42
+
43
+# Patch vLLM torch utils
44
+import vllm.utils.torch_utils as torch_utils
45
+
46
+
47
+def async_tensor_h2d(
48
+ data: list,
49
+ dtype: torch.dtype,
50
+ device: str | torch.device,
51
+ pin_memory: bool = False,
52
+) -> torch.Tensor:
53
+ return torch.tensor(data, dtype=dtype, device="cpu")
54
+
55
+
56
+torch_utils.async_tensor_h2d = async_tensor_h2d
57
+
58
+# Patch model runner APIs
59
+import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils
60
+import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils
61
+
62
+gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer
vllm/v1/worker/cpu_worker.py
+16
-3
@@ -1,5 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Must be imported firstly
5
+import vllm.v1.worker.cpu.shm # noqa # isort: skip
6
+
7
import math
8
import os
9
import sys
@@ -101,6 +105,8 @@ class CPUWorker(Worker):
105
)
106
107
def init_device(self):
108
+ self.device = torch.device("cpu")
109
+
110
# Check whether critical libraries are loaded
111
def check_preloaded_libs(name: str):
112
ld_preload_list = os.environ.get("LD_PRELOAD", "")
@@ -141,9 +147,16 @@ class CPUWorker(Worker):
147
set_random_seed(self.model_config.seed)
148
149
# Construct the model runner
144
- self.model_runner: CPUModelRunner = CPUModelRunner(
145
- self.vllm_config, torch.device("cpu")
146
- )
150
+ if self.use_v2_model_runner:
151
+ from vllm.v1.worker.cpu.model_runner import (
152
+ CPUModelRunner as CPUModelRunnerV2,
153
+ )
154
+
155
+ self.model_runner: CPUModelRunner = CPUModelRunnerV2( # type: ignore
156
+ self.vllm_config, self.device
157
+ )
158
+ else:
159
+ self.model_runner = CPUModelRunner(self.vllm_config, torch.device("cpu"))
160
161
def sleep(self, level: int = 1) -> None:
162
logger.warning("sleep mode is not supported on CPU, ignore it.")