[CPU] Use pre-built triton (#58140)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Sep 24, 2026 at 16:12 UTC
b44895cf932be890235b700b49b0753261cf5b97
3 files changed
+19
-77
docker/Dockerfile.cpu
+5
-77
@@ -107,61 +107,6 @@ ENV UV_LINK_MODE="copy"
107
108
ENV TARGETARCH=${TARGETARCH}
109
110
-######################### TRITON-CPU BUILD IMAGE #########################
111
-# Branches off base-common directly (before requirements/cpu.txt is installed
112
-# below) so a torch/etc. version bump doesn't invalidate this stage's cache --
113
-# it builds triton-cpu in an isolated `uv build` environment and never touches
114
-# these packages.
115
-FROM base-common AS vllm-triton-cpu-build
116
-
117
-# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
118
-# Re-declared here because this stage is `FROM base-common` (not `vllm-build`),
119
-# so it does not inherit the ARG/ENV defined there. Without it, the guard
120
-# below would see an empty value and build triton-cpu on non-x86 targets
121
-# (e.g. arm64).
122
-ARG VLLM_CPU_X86=0
123
-
124
-ARG USE_SCCACHE
125
-ARG SCCACHE_ENDPOINT
126
-ARG SCCACHE_LOCAL_ONLY=0
127
-
128
-WORKDIR /vllm-workspace
129
-
130
-RUN mkdir dist
131
-
132
-RUN --mount=type=cache,target=/root/.cache/uv \
133
- --mount=type=cache,target=/root/.cache/ccache \
134
- --mount=type=cache,target=/root/.cache/sccache-local,sharing=locked \
135
- --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
136
- --mount=type=cache,target=/root/.triton \
137
- --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
138
- set -ex; \
139
- if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
140
- git clone --recurse-submodules --shallow-submodules --filter=blob:none "https://github.com/triton-lang/triton-cpu.git"; \
141
- cd triton-cpu; \
142
- git checkout "270e696d"; \
143
- if [ ! -f third_party/sleef/CMakeLists.txt ]; then \
144
- echo "sleef submodule empty after shallow fetch; retrying unshallow"; \
145
- git submodule foreach --recursive 'git fetch --unshallow origin || true'; \
146
- git submodule update --init --recursive --force; \
147
- fi; \
148
- test -f third_party/sleef/CMakeLists.txt || { echo "FATAL: sleef submodule unavailable"; exit 1; }; \
149
- if [ "$USE_SCCACHE" = "1" ]; then \
150
- if [ "$SCCACHE_LOCAL_ONLY" = "1" ]; then \
151
- unset SCCACHE_BUCKET SCCACHE_REGION; \
152
- export SCCACHE_DIR=/root/.cache/sccache-local SCCACHE_CACHE_SIZE=20G; \
153
- elif [ -n "${SCCACHE_ENDPOINT}" ]; then \
154
- export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; \
155
- fi; \
156
- export TRITON_APPEND_CMAKE_ARGS="-DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache"; \
157
- sccache --show-stats; \
158
- fi; \
159
- uv build --wheel --out-dir=../dist; \
160
- if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi; \
161
- cd ..; \
162
- rm -rf triton-cpu; \
163
- fi
164
-
110
######################### COMMON BASE IMAGE + PYTHON DEPS #########################
111
FROM base-common AS base-common-pydeps
112
@@ -385,9 +330,12 @@ WORKDIR /vllm-workspace
330
COPY requirements/test/cpu.txt requirements/test/cpu.txt
331
332
# cpu.txt is compiled for x86_64, so platform markers are resolved away. Drop
388
-# packages unavailable on aarch64 (decord, terratorch) for arm builds.
333
+# packages unavailable on aarch64 (decord, terratorch) for arm builds. Also
334
+# drop the generic PyPI triton pulled in transitively (via xgrammar) on x86
335
+# builds, so it doesn't clobber the triton-cpu wheel from requirements/cpu.txt.
336
RUN case "$(uname -m)" in \
337
aarch64|arm64) sed -i '/^decord==/d; /^terratorch==/d' requirements/test/cpu.txt ;; \
338
+ x86_64) sed -i '/^triton==/d' requirements/test/cpu.txt ;; \
339
esac
340
341
# pytrec-eval-terrier bundles pre-C23 trec_eval sources. GCC 15 defaults to C23;
@@ -440,14 +388,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
388
fi && \
389
pre-commit install --hook-type pre-commit --hook-type commit-msg
390
443
-# Installed after requirements/test/cpu.txt (which pins triton==3.6.0) so the
444
-# triton-cpu wheel isn't overwritten by that pin.
445
-RUN --mount=type=cache,target=/root/.cache/uv \
446
- --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
447
- if ls dist/*.whl >/dev/null 2>&1; then \
448
- uv pip install dist/*.whl; \
449
- fi
450
-
391
ENTRYPOINT ["bash"]
392
393
######################### TEST IMAGE #########################
@@ -459,12 +399,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
399
--mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
400
uv pip install dist/*.whl
401
462
-RUN --mount=type=cache,target=/root/.cache/uv \
463
- --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
464
- if ls dist/*.whl >/dev/null 2>&1; then \
465
- uv pip install dist/*.whl; \
466
- fi
467
-
402
ADD ./tests/ ./tests/
403
ADD ./examples/ ./examples/
404
ADD ./benchmarks/ ./benchmarks/
@@ -494,13 +428,7 @@ WORKDIR /vllm-workspace
428
RUN --mount=type=cache,target=/root/.cache/uv \
429
--mount=type=cache,target=/root/.cache/ccache \
430
--mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
497
- uv pip install "$(realpath dist/*.whl)[audio]"
498
-
499
-RUN --mount=type=cache,target=/root/.cache/uv \
500
- --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
501
- if ls dist/*.whl >/dev/null 2>&1; then \
502
- uv pip install dist/*.whl; \
503
- fi
431
+ uv pip install "$(realpath dist/*.whl)[audio,bench]"
432
433
# Add labels to document build configuration
434
LABEL org.opencontainers.image.title="vLLM CPU"
requirements/cpu.txt
+5
@@ -21,5 +21,10 @@ torchcodec >= 0.14; platform_machine != "s390x" and platform_machine != "riscv64
21
# Intel Extension for PyTorch, only for x86_64 CPUs
22
intel-openmp==2024.2.1; platform_machine == "x86_64"
23
24
+# triton-cpu, only for x86_64 CPUs
25
+triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version == "3.10"
26
+triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version == "3.11"
27
+triton @ https://github.com/triton-lang/triton-cpu/releases/download/v3.8.0/triton-3.8.0+cpu-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl ; platform_machine == "x86_64" and python_version >= "3.12"
28
+
29
# Use this to gather CPU info and optimize based on ARM Neoverse cores
30
py-cpuinfo; platform_machine == "aarch64"
vllm/platforms/cpu.py
+9
@@ -199,6 +199,15 @@ class CpuPlatform(Platform):
199
if model_config is not None:
200
model_config.disable_cascade_attn = True
201
202
+ # Import lazily: vllm.triton_utils imports vllm.platforms.current_platform,
203
+ # which is still being resolved while this platform class is loading.
204
+ from vllm.triton_utils import HAS_TRITON
205
+
206
+ if cls.get_cpu_architecture() == CpuArchEnum.X86 and not HAS_TRITON:
207
+ logger.warning_once(
208
+ "Triton is not installed. triton-cpu is expected on x86 CPUs."
209
+ )
210
+
211
cache_config = vllm_config.cache_config
212
213
is_deepseek_v4 = (