[CPU] Ship triton-cpu wheel and fix several hardcoded pin_memory=True (#52092)

Signed-off-by: jiang1.li <jiang1.li@intel.com>

.buildkite/hardware_tests/cpu.yaml

@@ -20,6 +20,9 @@ steps:
   - tests/kernels/mamba/test_cpu_short_conv.py
   - tests/kernels/mamba/test_causal_conv1d.py
   - tests/kernels/mamba/test_mamba_ssm.py
+  - vllm/v1/sample/ops/topk_topp_triton.py
+  - vllm/v1/sample/ops/topk_topp_sampler.py
+  - tests/v1/sample/test_topk_topp_sampler.py
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -32,7 +35,8 @@ steps:
       pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
       pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
       pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
-      pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
+      pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py
+      pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
 
 # Note: SDE can't be downloaded from CI host because of AWS WAF
 # - label: CPU-Compatibility Tests
@@ -57,31 +61,12 @@ steps:
   - vllm/
   - tests/models/language/generation/
   - tests/models/language/pooling/
+  - tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 50m "
       pytest -x -v -s tests/models/language/generation -m cpu_model
-      pytest -x -v -s tests/models/language/pooling -m cpu_model"
-
-- label: CPU-ModelRunnerV2 Tests
-  depends_on: []
-  device: intel_cpu
-  no_plugin: true
-  soft_fail: true
-  source_file_dependencies:
-  - vllm/v1/worker/cpu/
-  - vllm/v1/worker/gpu/
-  - vllm/v1/sample/ops/topk_topp_triton.py
-  - vllm/v1/sample/ops/topk_topp_sampler.py
-  - tests/v1/sample/test_topk_topp_sampler.py
-  - tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
-  commands:
-    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
-      uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
-      VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
-      # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
-      pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp
+      pytest -x -v -s tests/models/language/pooling -m cpu_model
       pytest -x -v -s tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py"
 
 - label: CPU-Quantization Model Tests

docker/Dockerfile.cpu

@@ -180,11 +180,12 @@ RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
     bash build_rust.sh && \
     if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
 
-######################### BUILD IMAGE #########################
-FROM base-arch AS vllm-build
-
-ARG USE_SCCACHE
-ARG SCCACHE_ENDPOINT
+######################### SOURCE PREP IMAGE #########################
+# Shared prep (source, rust artifacts, build deps) for both vllm-build and
+# vllm-dev, so the two independent compiles (bdist_wheel vs. setup.py
+# develop) can run concurrently instead of vllm-dev waiting on a wheel
+# build it never uses.
+FROM base-arch AS vllm-src
 
 ARG GIT_REPO_CHECK=0
 # 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/
 
 RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
 
+######################### BUILD IMAGE #########################
+FROM vllm-src AS vllm-build
+
+ARG USE_SCCACHE
+ARG SCCACHE_ENDPOINT
+
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
@@ -254,7 +261,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
     --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
     if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
-        git clone --recurse-submodules "https://github.com/triton-lang/triton-cpu.git"; \
+        git clone --recurse-submodules --shallow-submodules --filter=blob:none "https://github.com/triton-lang/triton-cpu.git"; \
         cd triton-cpu; \
         git checkout "270e696d"; \
         if [ "$USE_SCCACHE" = "1" ]; then \
@@ -286,7 +293,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -r requirements/test/cpu.txt --torch-backend cpu
 
 ######################### DEV IMAGE #########################
-FROM vllm-build AS vllm-dev
+FROM vllm-src AS vllm-dev
 
 WORKDIR /vllm-workspace
 
@@ -312,6 +319,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \
     pre-commit install --hook-type pre-commit --hook-type commit-msg
 
+# Installed after requirements/test/cpu.txt (which pins triton==3.6.0) so the
+# triton-cpu wheel isn't overwritten by that pin.
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
+    if ls dist/*.whl >/dev/null 2>&1; then \
+        uv pip install dist/*.whl; \
+    fi
+
 ENTRYPOINT ["bash"]
 
 ######################### TEST IMAGE #########################
@@ -323,6 +338,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
     uv pip install dist/*.whl
 
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
+    if ls dist/*.whl >/dev/null 2>&1; then \
+        uv pip install dist/*.whl; \
+    fi
+
 ADD ./tests/ ./tests/
 ADD ./examples/ ./examples/
 ADD ./benchmarks/ ./benchmarks/
@@ -343,11 +364,6 @@ ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 ######################### RELEASE IMAGE #########################
 FROM base-arch AS vllm-openai
 
-# Re-declared here because this stage is `FROM base-arch` (not `vllm-build`),
-# so the RUN below that gates the triton-cpu wheel install on $VLLM_CPU_X86
-# would otherwise see an empty value and try to install it on non-x86 targets.
-ARG VLLM_CPU_X86=0
-
 WORKDIR /vllm-workspace
 
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -356,10 +372,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install "$(realpath dist/*.whl)[audio]"
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
-    if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
-        uv pip install "$(realpath dist/*.whl)"; \
+    if ls dist/*.whl >/dev/null 2>&1; then \
+        uv pip install dist/*.whl; \
     fi
 
 # Add labels to document build configuration

vllm/model_executor/models/granite_speech.py

@@ -61,6 +61,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.transformers_utils.processor import cached_processor_from_config
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
+from vllm.utils.torch_utils import PIN_MEMORY
 
 from .blip2 import Blip2QFormerModel
 from .interfaces import (
@@ -732,7 +733,8 @@ class GraniteSpeechForConditionalGeneration(
         target_device = self.encoder.input_linear.weight.device
         if target_device == input_features_mask.device:
             return input_features_mask
-        return input_features_mask.pin_memory().to(target_device, non_blocking=True)
+        masked = input_features_mask.pin_memory() if PIN_MEMORY else input_features_mask
+        return masked.to(target_device, non_blocking=True)
 
     def _pad_and_stack_input_features(
         self,

vllm/model_executor/models/qwen3_vl.py

@@ -112,6 +112,7 @@ from vllm.tokenizers.registry import cached_tokenizer_from_config
 from vllm.triton_utils import HAS_TRITON, tl, triton
 from vllm.utils.collection_utils import is_list_of
 from vllm.utils.math_utils import round_up
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers
 
 from ...utils.torch_utils import async_tensor_h2d
@@ -709,7 +710,7 @@ class Qwen3_VisionTransformer(nn.Module):
         pinned = torch.empty(
             (num_pos, pos_ids[0].shape[1]),
             dtype=pos_ids[0].dtype,
-            pin_memory=True,
+            pin_memory=PIN_MEMORY,
         )
         pos_ids = torch.cat(pos_ids, dim=0, out=pinned).to(
             self.device, non_blocking=True

vllm/v1/worker/cpu/shm.py

@@ -63,6 +63,7 @@ torch.cuda.Event = _EventPlaceholder
 torch.cuda.Stream = _StreamPlaceholder
 torch.cuda.set_stream = noop
 torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()
+torch.cuda.stream = lambda *args, **kwargs: _StreamPlaceholder()
 torch.accelerator.synchronize = noop
 torch.accelerator.empty_cache = empty_cache_noop
 torch.Tensor.pin_memory = fake_pin_memory

vllm/v1/worker/cpu_worker.py

@@ -171,6 +171,9 @@ class CPUWorker(Worker):
             self.local_rank,
             current_platform.dist_backend,
         )
+        if self.use_v2_model_runner:
+            logger.info_once("Using V2 Model Runner")
+
         # Set random seed.
         set_random_seed(self.model_config.seed)
 

vllm/v1/worker/gpu/mm/encoder_runner.py

@@ -17,6 +17,7 @@ from vllm.multimodal.utils import (
     group_and_batch_mm_kwargs,
     set_mm_embedding_modality,
 )
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.utils import (
     EncoderTimingStats,
@@ -115,7 +116,7 @@ class EncoderRunner:
     ) -> list[torch.Tensor]:
         encoder_outputs: list[torch.Tensor] = []
         for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs(
-            mm_kwargs, device=self.device, pin_memory=True
+            mm_kwargs, device=self.device, pin_memory=PIN_MEMORY
         ):
             batch_outputs = self.model.embed_multimodal(**mm_kwargs_batch)
             sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items)

vllm/v1/worker/gpu/structured_outputs.py

@@ -5,6 +5,7 @@ import torch
 
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
 from vllm.v1.worker.gpu.input_batch import InputBatch
 
@@ -63,7 +64,7 @@ class StructuredOutputsWorker:
         # Asynchronously copy the mapping to GPU.
         with torch.cuda.stream(self.copy_stream):
             logits_indices = torch.tensor(
-                mapping, dtype=torch.int32, device="cpu", pin_memory=True
+                mapping, dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY
             )
             logits_indices = self.logits_indices[: len(mapping)].copy_(
                 logits_indices, non_blocking=True