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 6, 17:00 UTC
In-Tree AMD Zen CPU Backend via zentorch [1/N] (#35970)
Signed-off-by: Lalithnarayan C <Lalithnarayan.C@amd.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Chinmay-Kulkarni-AMD <Chinmay.Kulkarni@amd.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Lalithnarayan C committed
Mar 16, 2026 at 05:05 UTC
7acaea634c53c6786c04c97e39f9c169f5fbddf9
9 files changed
+261
-2
docker/Dockerfile.cpu
+17
@@ -9,6 +9,7 @@
9
#
10
# Build targets:
11
# vllm-openai (default): used for serving deployment
12
+# vllm-openai-zen: vLLM from source + zentorch from PyPI via vllm[zen]
13
# vllm-test: used for CI tests
14
# vllm-dev: used for development
15
#
@@ -222,3 +223,19 @@ LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
223
LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
224
225
ENTRYPOINT ["vllm", "serve"]
226
+
227
+
228
+######################### ZEN CPU PYPI IMAGE #########################
229
+FROM vllm-openai AS vllm-openai-zen
230
+
231
+ARG TARGETARCH
232
+
233
+RUN if [ "$TARGETARCH" != "amd64" ]; then \
234
+ echo "ERROR: vllm-openai-amd only supports --platform=linux/amd64"; \
235
+ exit 1; \
236
+ fi
237
+
238
+RUN --mount=type=cache,target=/root/.cache/uv \
239
+ uv pip install "vllm[zen]"
240
+
241
+ENTRYPOINT ["vllm", "serve"]
setup.py
+2
@@ -966,6 +966,8 @@ setup(
966
ext_modules=ext_modules,
967
install_requires=get_requirements(),
968
extras_require={
969
+ # AMD Zen CPU optimizations via zentorch
970
+ "zen": ["zentorch"],
971
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
972
"tensorizer": ["tensorizer==2.10.1"],
973
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
tests/model_executor/test_cpu_unquantized_gemm_dispatch.py
new
+68
@@ -0,0 +1,68 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for CPU unquantized GEMM dispatch behavior."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.layers import utils
9
+from vllm.platforms import current_platform
10
+
11
+
12
+@pytest.fixture(scope="module")
13
+def _mock_zentorch_linear_unary():
14
+ """Register a mock zentorch_linear_unary op when zentorch is not installed.
15
+
16
+ Allows the dispatch tests to run in CI without a real zentorch build.
17
+ Skips registration when zentorch is already available.
18
+ """
19
+ if hasattr(torch.ops.zentorch, "zentorch_linear_unary"):
20
+ yield
21
+ return
22
+
23
+ lib_def = torch.library.Library("zentorch", "DEF")
24
+ lib_def.define(
25
+ "zentorch_linear_unary("
26
+ "Tensor input, "
27
+ "Tensor weight, "
28
+ "Tensor? bias, "
29
+ "bool is_weight_prepacked=False"
30
+ ") -> Tensor"
31
+ )
32
+
33
+ lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
34
+ lib_impl.impl(
35
+ "zentorch_linear_unary",
36
+ lambda input, weight, bias, is_weight_prepacked=False: (
37
+ torch.nn.functional.linear(input, weight, bias)
38
+ ),
39
+ )
40
+
41
+ yield
42
+
43
+ lib_impl._destroy()
44
+ lib_def._destroy()
45
+
46
+
47
+@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
48
+def test_dispatch_cpu_unquantized_gemm_uses_zentorch_on_zen(monkeypatch):
49
+ monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
50
+
51
+ layer = torch.nn.Linear(16, 8, bias=True)
52
+ x = torch.randn(4, 16)
53
+ expected = torch.nn.functional.linear(x, layer.weight, layer.bias)
54
+
55
+ utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
56
+ output = layer.cpu_linear(x, layer.weight, layer.bias)
57
+
58
+ torch.testing.assert_close(output, expected)
59
+
60
+
61
+@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
62
+def test_dispatch_cpu_unquantized_gemm_zen_remove_weight(monkeypatch):
63
+ monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
64
+
65
+ layer = torch.nn.Linear(16, 8, bias=True)
66
+ utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
67
+
68
+ assert layer.weight.numel() == 0
tests/test_zen_cpu_platform_detection.py
new
+37
@@ -0,0 +1,37 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from unittest.mock import mock_open, patch
4
+
5
+from vllm.platforms import _is_amd_zen_cpu
6
+
7
+
8
+def test_is_amd_zen_cpu_detects_amd_with_avx512():
9
+ cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2 avx512f avx512bw"
10
+ with (
11
+ patch("os.path.exists", return_value=True),
12
+ patch("builtins.open", mock_open(read_data=cpuinfo)),
13
+ ):
14
+ assert _is_amd_zen_cpu()
15
+
16
+
17
+def test_is_amd_zen_cpu_returns_false_for_amd_without_avx512():
18
+ cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2"
19
+ with (
20
+ patch("os.path.exists", return_value=True),
21
+ patch("builtins.open", mock_open(read_data=cpuinfo)),
22
+ ):
23
+ assert not _is_amd_zen_cpu()
24
+
25
+
26
+def test_is_amd_zen_cpu_returns_false_for_intel_with_avx512():
27
+ cpuinfo = "vendor_id: GenuineIntel\nflags: avx avx2 avx512f"
28
+ with (
29
+ patch("os.path.exists", return_value=True),
30
+ patch("builtins.open", mock_open(read_data=cpuinfo)),
31
+ ):
32
+ assert not _is_amd_zen_cpu()
33
+
34
+
35
+def test_is_amd_zen_cpu_returns_false_when_cpuinfo_missing():
36
+ with patch("os.path.exists", return_value=False):
37
+ assert not _is_amd_zen_cpu()
vllm/envs.py
+7
@@ -51,6 +51,7 @@ if TYPE_CHECKING:
51
VLLM_CPU_OMP_THREADS_BIND: str = "auto"
52
VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None
53
VLLM_CPU_SGL_KERNEL: bool = False
54
+ VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True
55
VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
56
VLLM_XLA_CHECK_RECOMPILATION: bool = False
57
VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"
@@ -709,6 +710,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
710
else None,
711
# (CPU backend only) whether to use SGL kernels, optimized for small batch.
712
"VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
713
+ # (Zen CPU backend) eagerly prepack weights into ZenDNN blocked layout
714
+ # at model load time. Eliminates per-inference layout conversion overhead.
715
+ "VLLM_ZENTORCH_WEIGHT_PREPACK": lambda: bool(
716
+ int(os.getenv("VLLM_ZENTORCH_WEIGHT_PREPACK", "1"))
717
+ ),
718
# If the env var is set, Ray Compiled Graph uses the specified
719
# channel type to communicate between workers belonging to
720
# different pipeline-parallel stages.
@@ -1768,6 +1774,7 @@ def compile_factors() -> dict[str, object]:
1774
"VLLM_V1_OUTPUT_PROC_CHUNK_SIZE",
1775
"VLLM_CPU_KVCACHE_SPACE",
1776
"VLLM_CPU_MOE_PREPACK",
1777
+ "VLLM_ZENTORCH_WEIGHT_PREPACK",
1778
"VLLM_TEST_FORCE_LOAD_FORMAT",
1779
"VLLM_ENABLE_CUDA_COMPATIBILITY",
1780
"VLLM_CUDA_COMPATIBILITY_PATH",
vllm/model_executor/layers/utils.py
+24
@@ -231,6 +231,30 @@ def dispatch_cpu_unquantized_gemm(
231
N, K = layer.weight.size()
232
dtype = layer.weight.dtype
233
234
+ # Zen CPU path: zentorch_linear_unary with optional eager weight prepacking.
235
+ if current_platform.is_zen_cpu() and hasattr(
236
+ torch.ops.zentorch, "zentorch_linear_unary"
237
+ ):
238
+ zen_weight = layer.weight.detach()
239
+ is_prepacked = False
240
+
241
+ if envs.VLLM_ZENTORCH_WEIGHT_PREPACK and hasattr(
242
+ torch.ops.zentorch, "zentorch_weight_prepack_for_linear"
243
+ ):
244
+ zen_weight = torch.ops.zentorch.zentorch_weight_prepack_for_linear(
245
+ zen_weight
246
+ )
247
+ is_prepacked = True
248
+
249
+ layer.cpu_linear = lambda x, weight, bias, _p=is_prepacked: (
250
+ torch.ops.zentorch.zentorch_linear_unary(
251
+ x, zen_weight, bias, is_weight_prepacked=_p
252
+ )
253
+ )
254
+ if remove_weight:
255
+ layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
256
+ return
257
+
258
if envs.VLLM_CPU_SGL_KERNEL and check_cpu_sgl_kernel(N, K, dtype):
259
packed_weight = torch.ops._C.convert_weight_packed(layer.weight)
260
if getattr(layer, "bias", None) is not None:
vllm/platforms/__init__.py
+36
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import logging
4
+import os
5
import traceback
6
from itertools import chain
7
from typing import TYPE_CHECKING
@@ -150,6 +151,15 @@ def xpu_platform_plugin() -> str | None:
151
return "vllm.platforms.xpu.XPUPlatform" if is_xpu else None
152
153
154
+def _is_amd_zen_cpu() -> bool:
155
+ """Detect AMD CPU with AVX-512 via /proc/cpuinfo."""
156
+ if not os.path.exists("/proc/cpuinfo"):
157
+ return False
158
+ with open("/proc/cpuinfo") as f:
159
+ cpuinfo = f.read()
160
+ return "AuthenticAMD" in cpuinfo and "avx512" in cpuinfo
161
+
162
+
163
def cpu_platform_plugin() -> str | None:
164
is_cpu = False
165
logger.debug("Checking if CPU platform is available.")
@@ -171,7 +181,24 @@ def cpu_platform_plugin() -> str | None:
181
except Exception as e:
182
logger.debug("CPU platform is not available because: %s", str(e))
183
174
- return "vllm.platforms.cpu.CpuPlatform" if is_cpu else None
184
+ if not is_cpu:
185
+ return None
186
+
187
+ if _is_amd_zen_cpu():
188
+ try:
189
+ import zentorch # noqa: F401
190
+
191
+ logger.debug(
192
+ "AMD Zen CPU detected with zentorch installed, using ZenCpuPlatform."
193
+ )
194
+ return "vllm.platforms.zen_cpu.ZenCpuPlatform"
195
+ except ImportError:
196
+ logger.debug(
197
+ "AMD Zen CPU detected but zentorch not installed, "
198
+ "falling back to CpuPlatform."
199
+ )
200
+
201
+ return "vllm.platforms.cpu.CpuPlatform"
202
203
204
builtin_platform_plugins = {
@@ -269,4 +296,11 @@ def __setattr__(name: str, value):
296
raise AttributeError(f"No attribute named '{name}' exists in {__name__}.")
297
298
272
-__all__ = ["Platform", "PlatformEnum", "current_platform", "CpuArchEnum", "_init_trace"]
299
+__all__ = [
300
+ "Platform",
301
+ "PlatformEnum",
302
+ "current_platform",
303
+ "CpuArchEnum",
304
+ "_init_trace",
305
+ "_is_amd_zen_cpu",
306
+]
vllm/platforms/interface.py
+3
@@ -167,6 +167,9 @@ class Platform:
167
def is_cpu(self) -> bool:
168
return self._enum == PlatformEnum.CPU
169
170
+ def is_zen_cpu(self) -> bool:
171
+ return False
172
+
173
def is_out_of_tree(self) -> bool:
174
return self._enum == PlatformEnum.OOT
175
vllm/platforms/zen_cpu.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import TYPE_CHECKING
5
+
6
+from vllm.logger import init_logger
7
+from vllm.platforms.cpu import CpuPlatform
8
+from vllm.utils.torch_utils import is_torch_equal_or_newer
9
+
10
+logger = init_logger(__name__)
11
+
12
+if TYPE_CHECKING:
13
+ from vllm.config import VllmConfig
14
+
15
+
16
+class ZenCpuPlatform(CpuPlatform):
17
+ """CPU platform with AMD Zen (ZenDNN/zentorch) optimizations.
18
+
19
+ Model-load time (dispatch_cpu_unquantized_gemm in layers/utils.py):
20
+ - Routes linear ops to zentorch_linear_unary.
21
+ - When VLLM_ZENTORCH_WEIGHT_PREPACK=1 (default), eagerly prepacks
22
+ weights via zentorch_weight_prepack_for_linear.
23
+ """
24
+
25
+ device_name: str = "cpu"
26
+ device_type: str = "cpu"
27
+
28
+ def is_zen_cpu(self) -> bool:
29
+ # is_cpu() also returns True for this platform (inherited from CpuPlatform).
30
+ return True
31
+
32
+ @classmethod
33
+ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
34
+ super().check_and_update_config(vllm_config)
35
+ cls._apply_pytorch_backports()
36
+
37
+ @classmethod
38
+ def _apply_pytorch_backports(cls):
39
+ """Backport PyTorch mainline fixes missing in 2.10.
40
+
41
+ PyTorch 2.10 has a bug in FxGraphCachePickler.dumps that doesn't
42
+ catch ValueError, causing torch.compile cache misses. Remove this
43
+ once we drop PyTorch 2.10 support. PT mainline already has this fix.
44
+ """
45
+ if not is_torch_equal_or_newer("2.10.0") or is_torch_equal_or_newer("2.11.0"):
46
+ return
47
+
48
+ cls._patch_fxgraphcache_pickle()
49
+
50
+ @classmethod
51
+ def _patch_fxgraphcache_pickle(cls):
52
+ """Backport mainline ValueError fix to FxGraphCachePickler.dumps()."""
53
+ from torch._inductor.codecache import BypassFxGraphCache, FxGraphCachePickler
54
+
55
+ original_dumps = FxGraphCachePickler.dumps
56
+ if hasattr(original_dumps, "_zen_patched"):
57
+ return
58
+
59
+ def patched_dumps(self, obj):
60
+ try:
61
+ return original_dumps(self, obj)
62
+ except ValueError as e:
63
+ raise BypassFxGraphCache("Failed to pickle cache key") from e
64
+
65
+ patched_dumps._zen_patched = True # type: ignore[attr-defined]
66
+ FxGraphCachePickler.dumps = patched_dumps
67
+ logger.info("[zen_cpu] Patched FxGraphCachePickler.dumps (ValueError fix)")