Multi step scheduling

vllm/executor/hpu_executor.py

@@ -54,9 +54,16 @@ def _create_worker(self,
                        local_rank: int = 0,
                        rank: int = 0,
                        distributed_init_method: Optional[str] = None):
+        if self.scheduler_config.is_multi_step:
+            module_name = "vllm.worker.multi_step_hpu_worker"
+            class_name = "MultiStepHPUWorker"
+        else:
+            module_name = "vllm.worker.hpu_worker"
+            class_name = "HPUWorker"
+
         wrapper = WorkerWrapperBase(
-            worker_module_name="vllm.worker.hpu_worker",
-            worker_class_name="HPUWorker",
+            worker_module_name=module_name,
+            worker_class_name=class_name,
         )
         wrapper.init_worker(**self._get_worker_kwargs(local_rank, rank,
                                                       distributed_init_method))

vllm/executor/ray_hpu_executor.py

@@ -87,8 +87,8 @@ def _get_worker_module_and_class(
                                            Type[WorkerBase]]]]:  # noqa: F821
         worker_class_fn = None
         if self.scheduler_config.is_multi_step:
-            raise NotImplementedError(
-                "Multi-step execution is not implemented for HPU")
+            worker_module_name = "vllm.worker.multi_step_hpu_worker"
+            worker_class_name = "MultiStepHPUWorker"
         elif self.speculative_config:
             raise NotImplementedError(
                 "Speculative decoding is not implemented for HPU")