biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[core] Multi Step Scheduling (#7000)

Browse Source 
Co-authored-by: afeldman-nm <156691304+afeldman-nm@users.noreply.github.com>
This commit is contained in:
William Lin
2024-08-19 13:52:13 -07:00
committed by GitHub
parent dad961ef5c
commit 47b65a5508
13 changed files with 1004 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
vllm/executor/gpu_executor.py
View File

@@ -69,13 +69,19 @@ class GPUExecutor(ExecutorBase):
distributed_init_method: Optional[str] = None) -> Dict:
worker_kwargs = self._get_worker_kwargs(local_rank, rank,
distributed_init_method)
if self.speculative_config is None:
worker_kwargs.update(worker_module_name="vllm.worker.worker",
worker_class_name="Worker")
else:
if self.scheduler_config.is_multi_step:
worker_kwargs.update(
worker_module_name="vllm.worker.multi_step_worker",
worker_class_name="MultiStepWorker")
elif self.speculative_config:
worker_kwargs.update(
worker_module_name="vllm.spec_decode.spec_decode_worker",
worker_class_name="create_spec_worker")
else:
worker_kwargs.update(worker_module_name="vllm.worker.worker",
worker_class_name="Worker")
return worker_kwargs
def _create_worker(self,

3
vllm/executor/ray_gpu_executor.py
View File

@@ -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"