[V1][Spec Decode] Optimize Medusa proposer to avoid GPU-CPU sync (#29723)

Signed-off-by: dongbo910220 <1275604947@qq.com>

vllm/v1/spec_decode/medusa.py

@@ -38,16 +38,16 @@ class MedusaProposer:
         self,
         target_hidden_states: torch.Tensor,
         sampling_metadata: SamplingMetadata,
-    ) -> list[list[int]]:
+    ) -> torch.Tensor:
         # Generate blocks and compute logits
         blocks = self.model(target_hidden_states)
         logits = self.model.compute_logits(blocks)
 
-        # Get draft tokens and transpose the result
+        # Compute argmax for each Medusa head and stack into a single tensor
-        # TODO(woosuk): OPTIMIZATION: Return GPU tensor without GPU-CPU
+        # Shape: [batch_size, num_heads]
-        # synchronization.
+        draft_tokens = torch.stack([logit.argmax(dim=-1) for logit in logits], dim=1)
-        draft_tokens = [logit.argmax(dim=-1).tolist() for logit in logits]
+
-        return [list(row) for row in zip(*draft_tokens)]
+        return draft_tokens
 
     def load_model(self, target_model: nn.Module) -> None:
         from vllm.compilation.backends import set_model_tag