| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn as nn |
| 6 | |
| 7 | from vllm.config import VllmConfig |
| 8 | from vllm.forward_context import set_forward_context |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.model_executor.model_loader import get_model |
| 11 | from vllm.model_executor.models.interfaces import is_mixture_of_experts |
| 12 | from vllm.v1.sample.metadata import SamplingMetadata |
| 13 | |
| 14 | # Initialize logger |
| 15 | logger = init_logger(__name__) |
| 16 | |
| 17 | |
| 18 | class MedusaProposer: |
| 19 | """Medusa proposer class for generating token sequences.""" |
| 20 | |
| 21 | def __init__( |
| 22 | self, |
| 23 | vllm_config: VllmConfig, |
| 24 | device: torch.device, |
| 25 | ): |
| 26 | # Save config parameters |
| 27 | self.vllm_config = vllm_config |
| 28 | assert vllm_config.speculative_config is not None, ( |
| 29 | "Speculative config must be set" |
| 30 | ) |
| 31 | self.spec_config = vllm_config.speculative_config |
| 32 | self.device = device |
| 33 | self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 34 | self.hidden_size = self.spec_config.draft_model_config.get_hidden_size() |
| 35 | self.dtype = vllm_config.model_config.dtype |
| 36 | self.num_speculative_tokens = self.spec_config.num_speculative_tokens |
| 37 | |
| 38 | def propose( |
| 39 | self, |
| 40 | num_speculative_tokens: int, |
| 41 | target_hidden_states: torch.Tensor, |
| 42 | sampling_metadata: SamplingMetadata, |
| 43 | slot_mappings: dict[str, torch.Tensor] |
| 44 | | list[dict[str, torch.Tensor]] |
| 45 | | None = None, # unused |
| 46 | ) -> torch.Tensor: |
| 47 | assert num_speculative_tokens == self.num_speculative_tokens |
| 48 | # Generate blocks and compute logits |
| 49 | blocks = self.model(target_hidden_states) |
| 50 | logits = self.model.compute_logits(blocks) |
| 51 | |
| 52 | # Compute argmax for each Medusa head and stack into a single tensor |
| 53 | # Shape: [batch_size, num_heads] |
| 54 | draft_tokens = torch.stack([logit.argmax(dim=-1) for logit in logits], dim=1) |
| 55 | |
| 56 | return draft_tokens |
| 57 | |
| 58 | def load_model(self, target_model: nn.Module) -> None: |
| 59 | from vllm.compilation.backends import set_model_tag |
| 60 | |
| 61 | with set_model_tag("medusa_head"): |
| 62 | self.model = get_model( |
| 63 | vllm_config=self.vllm_config, |
| 64 | model_config=self.spec_config.draft_model_config, |
| 65 | ) |
| 66 | assert not ( |
| 67 | is_mixture_of_experts(self.model) |
| 68 | and self.vllm_config.parallel_config.enable_eplb |
| 69 | ), "EPLB for Medusa is not supported" |
| 70 | |
| 71 | @torch.inference_mode() |
| 72 | def dummy_run(self, num_tokens: int) -> None: |
| 73 | hidden_states = torch.zeros( |
| 74 | (self.max_num_tokens, self.hidden_size), |
| 75 | dtype=self.dtype, |
| 76 | device=self.device, |
| 77 | ) |
| 78 | with set_forward_context(None, self.vllm_config, num_tokens=num_tokens): |
| 79 | self.model(hidden_states) |