From 22f9066285861cc7cdb49d5caad995582ae3cd36 Mon Sep 17 00:00:00 2001 From: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com> Date: Mon, 9 Dec 2024 17:22:28 -0500 Subject: [PATCH] Setting the value for the scpecilative decoding worker class on rocm platform (#313) Signed-off-by: Gregory Shtrasberg --- vllm/platforms/rocm.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 00efa056f7ef0..d2f7cd40e25b2 100644 --- 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"