Medusa: keep draft proposals on GPU

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

Author: dongbo910220

vllm/v1/spec_decode/medusa.py

@@ -38,16 +38,24 @@ def propose(
         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
-        # TODO(woosuk): OPTIMIZATION: Return GPU tensor without GPU-CPU
-        # synchronization.
-        draft_tokens = [logit.argmax(dim=-1).tolist() for logit in logits]
-        return [list(row) for row in zip(*draft_tokens)]
+        # Compute argmax for each Medusa head and stack into a single tensor
+        # Shape: [batch_size, num_heads]
+        draft_tokens = torch.stack([logit.argmax(dim=-1) for logit in logits], dim=1)
+
+        # Sanity check to catch any unexpected shape mismatch early
+        batch_size = target_hidden_states.shape[0]
+        num_heads = len(logits)
+        assert draft_tokens.shape == (
+            batch_size,
+            num_heads,
+        ), f"Expected shape ({batch_size}, {num_heads}), got {draft_tokens.shape}"
+
+        return draft_tokens
 
     def load_model(self, target_model: nn.Module) -> None:
         from vllm.compilation.backends import set_model_tag