fix

Signed-off-by: Woosuk Kwon <[email protected]>

Author: WoosukKwon

vllm/v1/worker/gpu/spec_decode/eagle.py

@@ -139,7 +139,7 @@ def generate_draft(
         num_tokens_across_dp: torch.Tensor | None,
     ) -> None:
         pos = self.input_buffers.positions[:num_reqs]
-        query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
+        query_start_loc = self.input_buffers.query_start_loc.gpu[: num_reqs + 1]
         for step in range(1, self.num_speculative_steps):
             # Run the eagle model.
             last_hidden_states, hidden_states = self.run_model(