biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[TPU] Support tensor parallelism in async llm engine (#6891)

Browse Source
This commit is contained in:
Earthwalker
2024-07-30 03:42:21 +08:00
committed by GitHub
parent 60d1c6e584
commit 7f8d612d24
2 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
vllm/engine/async_llm_engine.py
View File

@@ -407,8 +407,14 @@ class AsyncLLMEngine:
from vllm.executor.neuron_executor import NeuronExecutorAsync
executor_class = NeuronExecutorAsync
elif engine_config.device_config.device_type == "tpu":
from vllm.executor.tpu_executor import TPUExecutorAsync
executor_class = TPUExecutorAsync
if distributed_executor_backend == "ray":
initialize_ray_cluster(engine_config.parallel_config)
from vllm.executor.ray_tpu_executor import RayTPUExecutorAsync
executor_class = RayTPUExecutorAsync
else:
assert distributed_executor_backend is None
from vllm.executor.tpu_executor import TPUExecutorAsync
executor_class = TPUExecutorAsync
elif engine_config.device_config.device_type == "cpu":
from vllm.executor.cpu_executor import CPUExecutorAsync
executor_class = CPUExecutorAsync