Setting the value for the scpecilative decoding worker class on rocm …

…platform (#313) Signed-off-by: Gregory Shtrasberg <[email protected]>
ROCm · Dec 9, 2024 · 22f9066 · 22f9066
1 parent fb82bf1
commit 22f9066
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
@@ -150,6 +150,8 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
             elif vllm_config.speculative_config:
                 parallel_config.worker_cls = \
                     "vllm.spec_decode.spec_decode_worker.create_spec_worker"
+                parallel_config.sd_worker_cls = \
+                    "vllm.worker.worker.Worker"
             else:
                 parallel_config.worker_cls = "vllm.worker.worker.Worker"