Skip to content

GitLab

Unverified Commit 1bc3b739 authored by Siyuan Liu's avatar Siyuan Liu Committed by GitHub
Browse files

[V1][TPU] Add assertion on multi-step-scheduler (#14707) 


Signed-off-by: default avatarSiyuan Liu <lsiyuan@google.com>
parent 1bd32bc8
Hide whitespace changes
Inline Side-by-side
Showing with 11 additions and 5 deletions
vllm/platforms/tpu.py
View file @ 1bc3b739
...@@ -91,13 +91,19 @@ class TpuPlatform(Platform): ...@@ -91,13 +91,19 @@ class TpuPlatform(Platform):
parallel_config = vllm_config.parallel_config parallel_config = vllm_config.parallel_config
scheduler_config = vllm_config.scheduler_config scheduler_config = vllm_config.scheduler_config
if parallel_config.worker_cls == "auto": if parallel_config.worker_cls == "auto":
if envs.VLLM_USE_V1: if scheduler_config.is_multi_step:
parallel_config.worker_cls = \ if envs.VLLM_USE_V1:
"vllm.v1.worker.tpu_worker.TPUWorker" raise NotImplementedError(
else: "Multi-step scheduling is not supported (and not "
if scheduler_config.is_multi_step: "needed) on vLLM V1. Please launch without "
"--num-scheduler-steps.")
else:
parallel_config.worker_cls = \ parallel_config.worker_cls = \
"vllm.worker.multi_step_tpu_worker.MultiStepTPUWorker" "vllm.worker.multi_step_tpu_worker.MultiStepTPUWorker"
else:
if envs.VLLM_USE_V1:
parallel_config.worker_cls = \
"vllm.v1.worker.tpu_worker.TPUWorker"
else: else:
parallel_config.worker_cls = \ parallel_config.worker_cls = \
"vllm.worker.tpu_worker.TPUWorker" "vllm.worker.tpu_worker.TPUWorker"
......
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