A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[V1][TPU] Add assertion on multi-step-scheduler (#14707)
Signed-off-by: Siyuan Liu <lsiyuan@google.com>
Siyuan Liu committed
Mar 12, 2025 at 21:37 UTC
1bc3b739c4421a5d63b527fc2b5335e90e450204
1 file changed
+11
-5
vllm/platforms/tpu.py
+11
-5
@@ -91,13 +91,19 @@ class TpuPlatform(Platform):
91
parallel_config = vllm_config.parallel_config
92
scheduler_config = vllm_config.scheduler_config
93
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:
101
parallel_config.worker_cls = \
102
"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"
107
else:
108
parallel_config.worker_cls = \
109
"vllm.worker.tpu_worker.TPUWorker"