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
[Bugfix][CPU] Fallback oneDNN linear to torch linear to fix half gemm support on legecy platforms (#27526)
Signed-off-by: jiang1.li <jiang1.li@intel.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Li, Jiang committed
Oct 28, 2025 at 14:25 UTC
d34f5fe939eafa1eb4c6ecbe84020f09ebe09284
2 files changed
+22
-10
docker/Dockerfile.cpu
+1
-1
@@ -79,7 +79,7 @@ RUN echo 'ulimit -c 0' >> ~/.bashrc
79
######################### BUILD IMAGE #########################
80
FROM base AS vllm-build
81
82
-ARG max_jobs=2
82
+ARG max_jobs=32
83
ENV MAX_JOBS=${max_jobs}
84
85
ARG GIT_REPO_CHECK=0
vllm/model_executor/layers/utils.py
+21
-9
@@ -8,9 +8,12 @@ import torch
8
9
from vllm import _custom_ops as ops
10
from vllm import envs
11
+from vllm.logger import init_logger
12
from vllm.platforms import CpuArchEnum, current_platform
13
from vllm.utils.torch_utils import direct_register_custom_op
14
15
+logger = init_logger(__name__)
16
+
17
18
def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
19
# Shuffle weight along the last dimension so that
@@ -178,19 +181,28 @@ def dispatch_cpu_unquantized_gemm(
181
)
182
if remove_weight:
183
layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
184
+ return
185
elif (
186
ops._supports_onednn
187
and current_platform.get_cpu_architecture() != CpuArchEnum.POWERPC
188
):
185
- origin_weight = layer.weight
186
- if remove_weight:
187
- layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
188
- handler = ops.create_onednn_mm(origin_weight.t(), 32)
189
- layer.cpu_linear = lambda x, weight, bias: ops.onednn_mm(handler, x, bias)
190
- else:
191
- layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear(
192
- x, weight, bias
193
- )
189
+ try:
190
+ origin_weight = layer.weight
191
+ handler = ops.create_onednn_mm(origin_weight.t(), 32)
192
+ layer.cpu_linear = lambda x, weight, bias: ops.onednn_mm(handler, x, bias)
193
+ if remove_weight:
194
+ layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
195
+ return
196
+ except RuntimeError as e:
197
+ logger.warning_once(
198
+ "Failed to create oneDNN linear, fallback to torch linear."
199
+ f" Exception: {e}"
200
+ )
201
+
202
+ # fallback case
203
+ layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear(
204
+ x, weight, bias
205
+ )
206
207
208
def cpu_unquantized_gemm(