[ROCm][bugfix] scpecilative decoding worker class (vllm-project#11035)

Signed-off-by: Gregory Shtrasberg <[email protected]>
ROCm · Dec 9, 2024 · b63ba84 · b63ba84
1 parent 9c6459e
commit b63ba84
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
@@ -93,6 +93,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"