Skip to content

GitLab

Unverified Commit 8c1d4acb authored by Li, Jiang's avatar Li, Jiang Committed by GitHub
Browse files

[CPU] Disable oneDNN linear on non-x86 platforms (#25166) 


Signed-off-by: default avatarjiang1.li <jiang1.li@intel.com>
parent 486c5599
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
vllm/model_executor/layers/utils.py
View file @ 8c1d4acb
...@@ -7,7 +7,7 @@ import torch ...@@ -7,7 +7,7 @@ import torch
from vllm import _custom_ops as ops from vllm import _custom_ops as ops
from vllm import envs from vllm import envs
from vllm.platforms import current_platform from vllm.platforms import CpuArchEnum, current_platform
from vllm.utils import direct_register_custom_op from vllm.utils import direct_register_custom_op
...@@ -167,7 +167,8 @@ def dispatch_cpu_unquantized_gemm( ...@@ -167,7 +167,8 @@ def dispatch_cpu_unquantized_gemm(
if remove_weight: if remove_weight:
layer.weight = torch.nn.Parameter(torch.empty(0), layer.weight = torch.nn.Parameter(torch.empty(0),
requires_grad=False) requires_grad=False)
elif ops._supports_onednn: elif (ops._supports_onednn
and current_platform.get_cpu_architecture() == CpuArchEnum.X86):
origin_weight = layer.weight origin_weight = layer.weight
if remove_weight: if remove_weight:
layer.weight = torch.nn.Parameter(torch.empty(0), layer.weight = torch.nn.Parameter(torch.empty(0),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment