[CPU] Ship triton-cpu wheel and fix several hardcoded pin_memory=True (#52092)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Aug 13, 2026 at 14:40 UTC
373592ef57d4a19b057237fb015b0bd1382daa03
8 files changed
+50
-41
.buildkite/hardware_tests/cpu.yaml
+7
-22
@@ -20,6 +20,9 @@ steps:
20
- tests/kernels/mamba/test_cpu_short_conv.py
21
- tests/kernels/mamba/test_causal_conv1d.py
22
- tests/kernels/mamba/test_mamba_ssm.py
23
+ - vllm/v1/sample/ops/topk_topp_triton.py
24
+ - vllm/v1/sample/ops/topk_topp_sampler.py
25
+ - tests/v1/sample/test_topk_topp_sampler.py
26
commands:
27
- |
28
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -32,7 +35,8 @@ steps:
35
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
36
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
37
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
35
- pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
38
+ pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py
39
+ pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
40
41
# Note: SDE can't be downloaded from CI host because of AWS WAF
42
# - label: CPU-Compatibility Tests
@@ -57,31 +61,12 @@ steps:
61
- vllm/
62
- tests/models/language/generation/
63
- tests/models/language/pooling/
64
+ - tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
65
commands:
66
- |
67
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 50m "
68
pytest -x -v -s tests/models/language/generation -m cpu_model
64
- pytest -x -v -s tests/models/language/pooling -m cpu_model"
65
-
66
-- label: CPU-ModelRunnerV2 Tests
67
- depends_on: []
68
- device: intel_cpu
69
- no_plugin: true
70
- soft_fail: true
71
- source_file_dependencies:
72
- - vllm/v1/worker/cpu/
73
- - vllm/v1/worker/gpu/
74
- - vllm/v1/sample/ops/topk_topp_triton.py
75
- - vllm/v1/sample/ops/topk_topp_sampler.py
76
- - tests/v1/sample/test_topk_topp_sampler.py
77
- - tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
78
- commands:
79
- - |
80
- bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
81
- uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
82
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
83
- # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
84
- pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp
69
+ pytest -x -v -s tests/models/language/pooling -m cpu_model
70
pytest -x -v -s tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py"
71
72
- label: CPU-Quantization Model Tests
docker/Dockerfile.cpu
+30
-15
@@ -180,11 +180,12 @@ RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
180
bash build_rust.sh && \
181
if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
182
183
-######################### BUILD IMAGE #########################
184
-FROM base-arch AS vllm-build
185
-
186
-ARG USE_SCCACHE
187
-ARG SCCACHE_ENDPOINT
183
+######################### SOURCE PREP IMAGE #########################
184
+# Shared prep (source, rust artifacts, build deps) for both vllm-build and
185
+# vllm-dev, so the two independent compiles (bdist_wheel vs. setup.py
186
+# develop) can run concurrently instead of vllm-dev waiting on a wheel
187
+# build it never uses.
188
+FROM base-arch AS vllm-src
189
190
ARG GIT_REPO_CHECK=0
191
# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
@@ -221,6 +222,12 @@ COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
222
223
RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
224
225
+######################### BUILD IMAGE #########################
226
+FROM vllm-src AS vllm-build
227
+
228
+ARG USE_SCCACHE
229
+ARG SCCACHE_ENDPOINT
230
+
231
RUN --mount=type=cache,target=/root/.cache/uv \
232
--mount=type=cache,target=/root/.cache/ccache \
233
--mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
@@ -254,7 +261,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
261
--mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
262
--mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
263
if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
257
- git clone --recurse-submodules "https://github.com/triton-lang/triton-cpu.git"; \
264
+ git clone --recurse-submodules --shallow-submodules --filter=blob:none "https://github.com/triton-lang/triton-cpu.git"; \
265
cd triton-cpu; \
266
git checkout "270e696d"; \
267
if [ "$USE_SCCACHE" = "1" ]; then \
@@ -286,7 +293,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
293
uv pip install -r requirements/test/cpu.txt --torch-backend cpu
294
295
######################### DEV IMAGE #########################
289
-FROM vllm-build AS vllm-dev
296
+FROM vllm-src AS vllm-dev
297
298
WORKDIR /vllm-workspace
299
@@ -312,6 +319,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
319
uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \
320
pre-commit install --hook-type pre-commit --hook-type commit-msg
321
322
+# Installed after requirements/test/cpu.txt (which pins triton==3.6.0) so the
323
+# triton-cpu wheel isn't overwritten by that pin.
324
+RUN --mount=type=cache,target=/root/.cache/uv \
325
+ --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
326
+ if ls dist/*.whl >/dev/null 2>&1; then \
327
+ uv pip install dist/*.whl; \
328
+ fi
329
+
330
ENTRYPOINT ["bash"]
331
332
######################### TEST IMAGE #########################
@@ -323,6 +338,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
338
--mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
339
uv pip install dist/*.whl
340
341
+RUN --mount=type=cache,target=/root/.cache/uv \
342
+ --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
343
+ if ls dist/*.whl >/dev/null 2>&1; then \
344
+ uv pip install dist/*.whl; \
345
+ fi
346
+
347
ADD ./tests/ ./tests/
348
ADD ./examples/ ./examples/
349
ADD ./benchmarks/ ./benchmarks/
@@ -343,11 +364,6 @@ ENV HF_HUB_DOWNLOAD_TIMEOUT 60
364
######################### RELEASE IMAGE #########################
365
FROM base-arch AS vllm-openai
366
346
-# Re-declared here because this stage is `FROM base-arch` (not `vllm-build`),
347
-# so the RUN below that gates the triton-cpu wheel install on $VLLM_CPU_X86
348
-# would otherwise see an empty value and try to install it on non-x86 targets.
349
-ARG VLLM_CPU_X86=0
350
-
367
WORKDIR /vllm-workspace
368
369
RUN --mount=type=cache,target=/root/.cache/uv \
@@ -356,10 +372,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
372
uv pip install "$(realpath dist/*.whl)[audio]"
373
374
RUN --mount=type=cache,target=/root/.cache/uv \
359
- --mount=type=cache,target=/root/.cache/ccache \
375
--mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
361
- if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
362
- uv pip install "$(realpath dist/*.whl)"; \
376
+ if ls dist/*.whl >/dev/null 2>&1; then \
377
+ uv pip install dist/*.whl; \
378
fi
379
380
# Add labels to document build configuration
vllm/model_executor/models/granite_speech.py
+3
-1
@@ -61,6 +61,7 @@ from vllm.sequence import IntermediateTensors
61
from vllm.tokenizers import cached_tokenizer_from_config
62
from vllm.transformers_utils.processor import cached_processor_from_config
63
from vllm.utils.tensor_schema import TensorSchema, TensorShape
64
+from vllm.utils.torch_utils import PIN_MEMORY
65
66
from .blip2 import Blip2QFormerModel
67
from .interfaces import (
@@ -732,7 +733,8 @@ class GraniteSpeechForConditionalGeneration(
733
target_device = self.encoder.input_linear.weight.device
734
if target_device == input_features_mask.device:
735
return input_features_mask
735
- return input_features_mask.pin_memory().to(target_device, non_blocking=True)
736
+ masked = input_features_mask.pin_memory() if PIN_MEMORY else input_features_mask
737
+ return masked.to(target_device, non_blocking=True)
738
739
def _pad_and_stack_input_features(
740
self,
vllm/model_executor/models/qwen3_vl.py
+2
-1
@@ -112,6 +112,7 @@ from vllm.tokenizers.registry import cached_tokenizer_from_config
112
from vllm.triton_utils import HAS_TRITON, tl, triton
113
from vllm.utils.collection_utils import is_list_of
114
from vllm.utils.math_utils import round_up
115
+from vllm.utils.torch_utils import PIN_MEMORY
116
from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers
117
118
from ...utils.torch_utils import async_tensor_h2d
@@ -709,7 +710,7 @@ class Qwen3_VisionTransformer(nn.Module):
710
pinned = torch.empty(
711
(num_pos, pos_ids[0].shape[1]),
712
dtype=pos_ids[0].dtype,
712
- pin_memory=True,
713
+ pin_memory=PIN_MEMORY,
714
)
715
pos_ids = torch.cat(pos_ids, dim=0, out=pinned).to(
716
self.device, non_blocking=True
vllm/v1/worker/cpu/shm.py
+1
@@ -63,6 +63,7 @@ torch.cuda.Event = _EventPlaceholder
63
torch.cuda.Stream = _StreamPlaceholder
64
torch.cuda.set_stream = noop
65
torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()
66
+torch.cuda.stream = lambda *args, **kwargs: _StreamPlaceholder()
67
torch.accelerator.synchronize = noop
68
torch.accelerator.empty_cache = empty_cache_noop
69
torch.Tensor.pin_memory = fake_pin_memory
vllm/v1/worker/cpu_worker.py
+3
@@ -171,6 +171,9 @@ class CPUWorker(Worker):
171
self.local_rank,
172
current_platform.dist_backend,
173
)
174
+ if self.use_v2_model_runner:
175
+ logger.info_once("Using V2 Model Runner")
176
+
177
# Set random seed.
178
set_random_seed(self.model_config.seed)
179
vllm/v1/worker/gpu/mm/encoder_runner.py
+2
-1
@@ -17,6 +17,7 @@ from vllm.multimodal.utils import (
17
group_and_batch_mm_kwargs,
18
set_mm_embedding_modality,
19
)
20
+from vllm.utils.torch_utils import PIN_MEMORY
21
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
22
from vllm.v1.worker.utils import (
23
EncoderTimingStats,
@@ -115,7 +116,7 @@ class EncoderRunner:
116
) -> list[torch.Tensor]:
117
encoder_outputs: list[torch.Tensor] = []
118
for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs(
118
- mm_kwargs, device=self.device, pin_memory=True
119
+ mm_kwargs, device=self.device, pin_memory=PIN_MEMORY
120
):
121
batch_outputs = self.model.embed_multimodal(**mm_kwargs_batch)
122
sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items)
vllm/v1/worker/gpu/structured_outputs.py
+2
-1
@@ -5,6 +5,7 @@ import torch
5
6
from vllm.triton_utils import tl, triton
7
from vllm.utils.math_utils import cdiv
8
+from vllm.utils.torch_utils import PIN_MEMORY
9
from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
10
from vllm.v1.worker.gpu.input_batch import InputBatch
11
@@ -63,7 +64,7 @@ class StructuredOutputsWorker:
64
# Asynchronously copy the mapping to GPU.
65
with torch.cuda.stream(self.copy_stream):
66
logits_indices = torch.tensor(
66
- mapping, dtype=torch.int32, device="cpu", pin_memory=True
67
+ mapping, dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY
68
)
69
logits_indices = self.logits_indices[: len(mapping)].copy_(
70
logits_indices, non_blocking=True