A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 17:30 UTC
[V1][Spec Decode] Optimize Medusa proposer to avoid GPU-CPU sync (#29723)
Signed-off-by: dongbo910220 <1275604947@qq.com>
dongbo910220 committed
Dec 10, 2025 at 08:15 UTC
03b5f940fdcff25024ce5d37c357d770344a8f20
1 file changed
+6
-6
vllm/v1/spec_decode/medusa.py
+6
-6
@@ -38,16 +38,16 @@ class MedusaProposer:
38
self,
39
target_hidden_states: torch.Tensor,
40
sampling_metadata: SamplingMetadata,
41
- ) -> list[list[int]]:
41
+ ) -> torch.Tensor:
42
# Generate blocks and compute logits
43
blocks = self.model(target_hidden_states)
44
logits = self.model.compute_logits(blocks)
45
46
- # Get draft tokens and transpose the result
47
- # TODO(woosuk): OPTIMIZATION: Return GPU tensor without GPU-CPU
48
- # synchronization.
49
- draft_tokens = [logit.argmax(dim=-1).tolist() for logit in logits]
50
- return [list(row) for row in zip(*draft_tokens)]
46
+ # Compute argmax for each Medusa head and stack into a single tensor
47
+ # Shape: [batch_size, num_heads]
48
+ draft_tokens = torch.stack([logit.argmax(dim=-1) for logit in logits], dim=1)
49
+
50
+ return draft_tokens
51
52
def load_model(self, target_model: nn.Module) -> None:
53
from vllm.compilation.backends import set_model_tag