[SpecDecode]Fix spec decode warmup device selection (#55245)

Signed-off-by: Kunshang Ji <jikunshang95@gmail.com> Co-authored-by: GitHub Copilot <copilot@github.com>
2026-09-04 00:48:45 -07:00
parent a5c9179e73
commit a8693df504
1 changed files with 1 additions and 1 deletions
@@ -44,7 +44,7 @@ def spec_decode_rejection_warmup(worker: Worker) -> None:
    use_block_verification = rejection_method == "block"
    use_synthetic = rejection_method == "synthetic"

-    device = torch.device("cuda")
+    device = worker.device
    num_reqs = 1
    tokens_per_req = num_spec + 1
    num_logits = num_reqs * tokens_per_req