Skip to content

GitLab

Unverified Commit 73f635a7 authored by XiongfeiWei's avatar XiongfeiWei Committed by GitHub
Browse files

[Bug] Add TPU backend option (#32438) 


Signed-off-by: default avatarXiongfei Wei <isaacwxf23@gmail.com>
parent 35bf5d08
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 0 deletions
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
View file @ 73f635a7
...@@ -32,6 +32,7 @@ class UnquantizedMoeBackend(Enum): ...@@ -32,6 +32,7 @@ class UnquantizedMoeBackend(Enum):
TRITON = "TRITON" TRITON = "TRITON"
CPU = "CPU" CPU = "CPU"
XPU = "XPU" XPU = "XPU"
TPU = "TPU"
# NOTE(zyongye): Unsupported backend means backend # NOTE(zyongye): Unsupported backend means backend
...@@ -40,6 +41,7 @@ class UnquantizedMoeBackend(Enum): ...@@ -40,6 +41,7 @@ class UnquantizedMoeBackend(Enum):
UNSUPPORTED_BACKEND = [ UNSUPPORTED_BACKEND = [
UnquantizedMoeBackend.CPU, UnquantizedMoeBackend.CPU,
UnquantizedMoeBackend.XPU, UnquantizedMoeBackend.XPU,
UnquantizedMoeBackend.TPU,
] ]
...@@ -91,6 +93,8 @@ def select_unquantized_moe_backend( ...@@ -91,6 +93,8 @@ def select_unquantized_moe_backend(
backend = UnquantizedMoeBackend.XPU backend = UnquantizedMoeBackend.XPU
if current_platform.is_cpu(): if current_platform.is_cpu():
backend = UnquantizedMoeBackend.CPU backend = UnquantizedMoeBackend.CPU
if current_platform.is_tpu():
backend = UnquantizedMoeBackend.TPU
logger.info_once(_make_log_backend(backend), scope="local") logger.info_once(_make_log_backend(backend), scope="local")
return backend return backend
......
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