[SpecDecode]Fix spec decode warmup device selection (#55245)
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com> Co-authored-by: GitHub Copilot <copilot@github.com>
Kunshang Ji committed
Sep 4, 2026 at 00:48 UTC
a8693df504404ea0288e480165143d8a216f4220
1 file changed
+1
-1
vllm/model_executor/warmup/spec_decode_rejection_warmup.py
+1
-1
@@ -44,7 +44,7 @@ def spec_decode_rejection_warmup(worker: Worker) -> None:
44
use_block_verification = rejection_method == "block"
45
use_synthetic = rejection_method == "synthetic"
46
47
- device = torch.device("cuda")
47
+ device = worker.device
48
num_reqs = 1
49
tokens_per_req = num_spec + 1
50
num_logits = num_reqs * tokens_per_req