Skip to content

Commit 2eeccce

Browse files
lsy323shreyankg
authored andcommitted
[V1][TPU] Add assertion on multi-step-scheduler (vllm-project#14707)
Signed-off-by: Siyuan Liu <lsiyuan@google.com>
1 parent 9e58ebb commit 2eeccce

File tree

1 file changed

+11
-5
lines changed

1 file changed

+11
-5
lines changed

vllm/platforms/tpu.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -91,13 +91,19 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
9191
parallel_config = vllm_config.parallel_config
9292
scheduler_config = vllm_config.scheduler_config
9393
if parallel_config.worker_cls == "auto":
94-
if envs.VLLM_USE_V1:
95-
parallel_config.worker_cls = \
96-
"vllm.v1.worker.tpu_worker.TPUWorker"
97-
else:
98-
if scheduler_config.is_multi_step:
94+
if scheduler_config.is_multi_step:
95+
if envs.VLLM_USE_V1:
96+
raise NotImplementedError(
97+
"Multi-step scheduling is not supported (and not "
98+
"needed) on vLLM V1. Please launch without "
99+
"--num-scheduler-steps.")
100+
else:
99101
parallel_config.worker_cls = \
100102
"vllm.worker.multi_step_tpu_worker.MultiStepTPUWorker"
103+
else:
104+
if envs.VLLM_USE_V1:
105+
parallel_config.worker_cls = \
106+
"vllm.v1.worker.tpu_worker.TPUWorker"
101107
else:
102108
parallel_config.worker_cls = \
103109
"vllm.worker.tpu_worker.TPUWorker"

0 commit comments

Comments
 (0)