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 14, 05:30 UTC
[CPU] update torch 2.8 and fix missing fields in TorchSDPAMetadata (#25652)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Sep 25, 2025 at 21:29 UTC
eb32335e355f194c1b9d27a570176d4cf7fb9abf
7 files changed
+59
-53
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+2
-5
@@ -58,11 +58,8 @@ function cpu_tests() {
58
# pytest -x -v -s tests/kernels/attention/test_cache.py -m cpu_model
59
# pytest -x -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model
60
61
- # Note: disable Bart until supports V1
62
- pytest -x -v -s tests/models/language/generation -m cpu_model \
63
- --ignore=tests/models/language/generation/test_bart.py
64
- VLLM_CPU_SGL_KERNEL=1 pytest -x -v -s tests/models/language/generation -m cpu_model \
65
- --ignore=tests/models/language/generation/test_bart.py
61
+ pytest -x -v -s tests/models/language/generation -m cpu_model
62
+ VLLM_CPU_SGL_KERNEL=1 pytest -x -v -s tests/models/language/generation -m cpu_model
63
64
pytest -x -v -s tests/models/language/pooling -m cpu_model
65
pytest -x -v -s tests/models/multimodal/generation \
docker/Dockerfile.cpu
-3
@@ -114,9 +114,6 @@ WORKDIR /workspace/vllm
114
RUN --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
115
cp requirements/test.in requirements/cpu-test.in && \
116
sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
117
- sed -i 's/^torch==.*/torch==2.6.0/g' requirements/cpu-test.in && \
118
- sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \
119
- sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \
117
uv pip compile requirements/cpu-test.in -o requirements/cpu-test.txt --index-strategy unsafe-best-match --torch-backend cpu
118
119
RUN --mount=type=cache,target=/root/.cache/uv \
requirements/cpu-build.txt
+1
-3
@@ -1,12 +1,10 @@
1
-# Temporarily used for x86 CPU backend to avoid performance regression of torch>2.6.0+cpu,
2
-# see https://github.com/pytorch/pytorch/pull/151218
1
cmake>=3.26.1
2
ninja
3
packaging>=24.2
4
setuptools>=77.0.3,<80.0.0
5
setuptools-scm>=8
6
--extra-index-url https://download.pytorch.org/whl/cpu
9
-torch==2.6.0+cpu; platform_machine == "x86_64" # torch>2.6.0+cpu has performance regression on x86 platform, see https://github.com/pytorch/pytorch/pull/151218
7
+torch==2.8.0+cpu; platform_machine == "x86_64"
8
torch==2.8.0; platform_machine == "ppc64le" or platform_machine == "aarch64" or platform_system == "Darwin"
9
wheel
10
jinja2>=3.1.6
requirements/cpu.txt
+2
-2
@@ -8,7 +8,7 @@ numba == 0.61.2; python_version > '3.9' and platform_machine != "s390x"
8
packaging>=24.2
9
setuptools>=77.0.3,<80.0.0
10
--extra-index-url https://download.pytorch.org/whl/cpu
11
-torch==2.6.0+cpu; platform_machine == "x86_64" # torch>2.6.0+cpu has performance regression on x86 platform, see https://github.com/pytorch/pytorch/pull/151218
11
+torch==2.8.0+cpu; platform_machine == "x86_64"
12
torch==2.8.0; platform_system == "Darwin"
13
torch==2.8.0; platform_machine == "ppc64le" or platform_machine == "aarch64"
14
@@ -23,7 +23,7 @@ datasets # for benchmark scripts
23
24
# Intel Extension for PyTorch, only for x86_64 CPUs
25
intel-openmp==2024.2.1; platform_machine == "x86_64"
26
-intel_extension_for_pytorch==2.6.0; platform_machine == "x86_64" # torch>2.6.0+cpu has performance regression on x86 platform, see https://github.com/pytorch/pytorch/pull/151218
26
+intel_extension_for_pytorch==2.8.0; platform_machine == "x86_64"
27
triton==3.2.0; platform_machine == "x86_64" # Triton is required for torch 2.6+cpu, as it is imported in torch.compile.
28
29
# Use this to gather CPU info and optimize based on ARM Neoverse cores
vllm/v1/attention/backends/cpu_attn.py
+13
-1
@@ -85,6 +85,19 @@ class TorchSDPABackend(AttentionBackend):
85
86
@dataclass
87
class TorchSDPAMetadata(AttentionMetadata):
88
+ """Attention metadata for prefill and decode batched together."""
89
+ # Total number of prefill requests.
90
+ num_prefills: int
91
+ # Number of prefill tokens.
92
+ num_prefill_tokens: int
93
+ # Number of decode tokens. Note that it is equivalent to the number of
94
+ # decode requests.
95
+ num_decode_tokens: int
96
+ # (num_tokens,). The indices of the token slots that input tokens will be
97
+ # stored into. E.g., if `slot_mapping` is [35, 2, 17] and the block size
98
+ # is 16, the three tokens are stored in the 3rd slot in block 2, 2nd slot
99
+ # in block 0, and 1st slot in block 1, respectively.
100
+ slot_mapping: torch.Tensor
101
"""Metadata for PagedAttention."""
102
# (batch_size,). The length of sequences (entire tokens seen so far) per
103
# sequence.
@@ -420,7 +433,6 @@ class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
433
num_prompt_req], # prefill
434
query_start_loc=query_start_loc_cpu[:num_reqs +
435
1], # for logits index
423
- enable_kv_scales_calculation=False,
436
)
437
438
return attn_metadata
vllm/v1/sample/ops/topk_topp_sampler.py
+41
@@ -68,6 +68,8 @@ class TopKTopPSampler(nn.Module):
68
"native implementation of top-p & top-k sampling. For the "
69
"best performance, please install FlashInfer.")
70
self.forward = self.forward_native
71
+ elif current_platform.is_cpu():
72
+ self.forward = self.forward_cpu
73
else:
74
self.forward = self.forward_native
75
@@ -119,6 +121,45 @@ class TopKTopPSampler(nn.Module):
121
# because of slicing operation in logits_processor.
122
return flashinfer_sample(logits.contiguous(), k, p, generators), None
123
124
+ def forward_cpu(
125
+ self,
126
+ logits: torch.Tensor,
127
+ generators: dict[int, torch.Generator],
128
+ k: Optional[torch.Tensor],
129
+ p: Optional[torch.Tensor],
130
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
131
+ """
132
+ PyTorch-native implementation of top-k and top-p sampling for CPU.
133
+
134
+ The logits tensor may be updated in-place.
135
+ """
136
+ logits = self.apply_top_k_top_p(logits, k, p)
137
+ logits_to_return = None
138
+ if self.logprobs_mode == "processed_logits":
139
+ logits_to_return = logits
140
+ elif self.logprobs_mode == "processed_logprobs":
141
+ logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32)
142
+
143
+ # Note: this is a workaround for
144
+ # https://github.com/pytorch/pytorch/pull/151218
145
+ @torch.compile(dynamic=True)
146
+ def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor:
147
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
148
+ q = torch.empty_like(probs)
149
+ q.exponential_()
150
+ return probs.div(q).argmax(dim=-1).view(-1)
151
+
152
+ if len(generators) != logits.shape[0]:
153
+ return compiled_random_sample(logits), logits_to_return
154
+ else:
155
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
156
+ q = torch.empty_like(probs)
157
+ q.exponential_()
158
+ for i, generator in generators.items():
159
+ q[i].exponential_(generator=generator)
160
+
161
+ return probs.div_(q).argmax(dim=-1).view(-1), logits_to_return
162
+
163
164
def apply_top_k_top_p(
165
logits: torch.Tensor,
vllm/v1/worker/cpu_worker.py
-39
@@ -8,18 +8,13 @@ import torch
8
9
from vllm import envs
10
from vllm.config import VllmConfig
11
-from vllm.distributed.parallel_state import get_pp_group, get_tp_group
11
from vllm.logger import init_logger
12
from vllm.model_executor.utils import set_random_seed
13
from vllm.platforms import CpuArchEnum, current_platform
14
from vllm.platforms.cpu import CpuPlatform, LogicalCPUInfo
16
-from vllm.sequence import IntermediateTensors
17
-from vllm.v1.core.sched.output import SchedulerOutput
18
-from vllm.v1.outputs import ModelRunnerOutput
15
from vllm.v1.worker.cpu_model_runner import CPUModelRunner
16
from vllm.v1.worker.gpu_worker import (Worker,
17
init_worker_distributed_environment)
22
-from vllm.v1.worker.utils import is_residual_scattered_for_sp
18
19
logger = init_logger(__name__)
20
@@ -102,40 +97,6 @@ class CPUWorker(Worker):
97
set_random_seed(self.model_config.seed)
98
self.model_runner.warming_up_model()
99
105
- @torch.inference_mode()
106
- def execute_model(
107
- self,
108
- scheduler_output: "SchedulerOutput",
109
- ) -> Optional[ModelRunnerOutput]:
110
- intermediate_tensors = None
111
- num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
112
- num_input_tokens = self.model_runner._get_num_input_tokens(
113
- num_scheduled_tokens)
114
- all_gather_tensors = {
115
- "residual":
116
- not is_residual_scattered_for_sp(self.vllm_config,
117
- num_input_tokens)
118
- }
119
- if not get_pp_group().is_first_rank:
120
- intermediate_tensors = IntermediateTensors(
121
- get_pp_group().recv_tensor_dict(
122
- all_gather_group=get_tp_group(),
123
- all_gather_tensors=all_gather_tensors))
124
-
125
- output = self.model_runner.execute_model(scheduler_output,
126
- intermediate_tensors)
127
-
128
- if not get_pp_group().is_last_rank:
129
- assert isinstance(output, IntermediateTensors)
130
- get_pp_group().send_tensor_dict(
131
- output.tensors,
132
- all_gather_group=get_tp_group(),
133
- all_gather_tensors=all_gather_tensors)
134
- return None
135
-
136
- assert isinstance(output, ModelRunnerOutput)
137
- return output if self.is_driver_worker else None
138
-
100
def _get_autobind_cpu_ids(
101
self, cpu_selector: Callable[[list[LogicalCPUInfo]],
102
list[LogicalCPUInfo]]