[core] Multi Step Scheduling (#7000)

Co-authored-by: afeldman-nm <156691304+afeldman-nm@users.noreply.github.com>
2024-08-19 13:52:13 -07:00
parent dad961ef5c
commit 47b65a5508
13 changed files with 1004 additions and 34 deletions
--- a/vllm/executor/ray_gpu_executor.py
+++ b/vllm/executor/ray_gpu_executor.py
@@ -94,6 +94,9 @@ class RayGPUExecutor(DistributedGPUExecutor):
        if self.speculative_config is not None:
            worker_module_name = "vllm.spec_decode.spec_decode_worker"
            worker_class_name = "create_spec_worker"
+        elif self.scheduler_config.is_multi_step:
+            worker_module_name = "vllm.worker.multi_step_worker"
+            worker_class_name = "MultiStepWorker"
        else:
            worker_module_name = "vllm.worker.worker"
            worker_class_name = "Worker"