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 15, 09:45 UTC
[Model] vLLM v1 supports Medusa (#17956)
Signed-off-by: lisiqi23 <lisiqi23@xiaomi.com> Signed-off-by: skylee-01 <497627264@qq.com> Co-authored-by: lisiqi23 <lisiqi23@xiaomi.com>
Sky Lee committed
May 16, 2025 at 12:05 UTC
f4937a51c138978928f38da6a2d3b30c53286240
4 files changed
+108
-2
vllm/engine/arg_utils.py
+4
-1
@@ -1324,19 +1324,22 @@ class EngineArgs:
1324
# Only Ngram speculative decoding so far.
1325
is_ngram_enabled = False
1326
is_eagle_enabled = False
1327
+ is_medusa_enabled = False
1328
if self.speculative_config is not None:
1329
# This is supported but experimental (handled below).
1330
speculative_method = self.speculative_config.get("method")
1331
if speculative_method:
1332
if speculative_method in ("ngram", "[ngram]"):
1333
is_ngram_enabled = True
1334
+ elif speculative_method == "medusa":
1335
+ is_medusa_enabled = True
1336
elif speculative_method in ("eagle", "eagle3"):
1337
is_eagle_enabled = True
1338
else:
1339
speculative_model = self.speculative_config.get("model")
1340
if speculative_model in ("ngram", "[ngram]"):
1341
is_ngram_enabled = True
1339
- if not (is_ngram_enabled or is_eagle_enabled):
1342
+ if not (is_ngram_enabled or is_eagle_enabled or is_medusa_enabled):
1343
# Other speculative decoding methods are not supported yet.
1344
_raise_or_fallback(feature_name="Speculative Decoding",
1345
recommend_to_remove=False)
vllm/model_executor/models/medusa.py
+4
-1
@@ -51,7 +51,10 @@ class Medusa(nn.Module):
51
needs to have truncated_vocab_size (=k) as an attribute."""
52
53
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
54
- config = vllm_config.model_config.hf_config
54
+ if hasattr(vllm_config, 'draft_model_config'):
55
+ config = vllm_config.draft_model_config.hf_config
56
+ else:
57
+ config = vllm_config.model_config.hf_config
58
super().__init__()
59
self.config = config
60
self.blocks = nn.ModuleList([
vllm/v1/spec_decode/medusa.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import torch
4
+import torch.nn as nn
5
+
6
+from vllm.config import VllmConfig, set_current_vllm_config
7
+from vllm.forward_context import set_forward_context
8
+from vllm.logger import init_logger
9
+from vllm.model_executor.model_loader import get_model_loader
10
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype
11
+from vllm.model_executor.models.medusa import Medusa
12
+from vllm.v1.sample.metadata import SamplingMetadata
13
+
14
+# Initialize logger
15
+logger = init_logger(__name__)
16
+
17
+
18
+class MedusaProposer:
19
+ """
20
+ Medusa proposer class for generating token sequences
21
+ """
22
+
23
+ def __init__(
24
+ self,
25
+ vllm_config: VllmConfig,
26
+ device: torch.device,
27
+ ):
28
+ # Save config parameters
29
+ self.vllm_config = vllm_config
30
+ self.device = device
31
+ self.max_num_tokens = (
32
+ vllm_config.scheduler_config.max_num_batched_tokens)
33
+ self.hidden_size = vllm_config.speculative_config.\
34
+ draft_model_config.get_hidden_size(
35
+ )
36
+ self.dtype = vllm_config.model_config.dtype
37
+
38
+ def propose(
39
+ self,
40
+ target_hidden_states: torch.Tensor,
41
+ sampling_metadata: SamplingMetadata,
42
+ ) -> torch.Tensor:
43
+ # Generate blocks and compute logits
44
+ blocks = self.model(target_hidden_states)
45
+ logits = self.model.compute_logits(blocks, None)
46
+
47
+ # Get draft tokens and transpose the result
48
+ draft_tokens = [logit.argmax(dim=-1).tolist() for logit in logits]
49
+ return [list(row) for row in zip(*draft_tokens)]
50
+
51
+ def load_model(self, target_model: nn.Module) -> None:
52
+ # Get model loader and config
53
+ loader = get_model_loader(self.vllm_config.load_config)
54
+ draft_config = self.vllm_config.speculative_config.draft_model_config
55
+
56
+ # Load model with proper dtype and config
57
+ with set_default_torch_dtype(draft_config.dtype), \
58
+ set_current_vllm_config(self.vllm_config):
59
+ self.model = Medusa(
60
+ vllm_config=self.vllm_config.speculative_config).to(
61
+ self.device)
62
+
63
+ # Load model weights
64
+ weights = loader.get_all_weights(draft_config, self.model)
65
+ self.model.load_weights(weights)
66
+
67
+ @torch.inference_mode()
68
+ def dummy_run(self, num_tokens: int) -> None:
69
+ hidden_states = torch.zeros((self.max_num_tokens, self.hidden_size),
70
+ dtype=self.dtype,
71
+ device=self.device)
72
+ with set_forward_context(None, self.vllm_config,
73
+ num_tokens=num_tokens):
74
+ self.model(hidden_states)
vllm/v1/worker/gpu_model_runner.py
+26
@@ -47,6 +47,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
47
from vllm.v1.sample.rejection_sampler import RejectionSampler
48
from vllm.v1.sample.sampler import Sampler
49
from vllm.v1.spec_decode.eagle import EagleProposer
50
+from vllm.v1.spec_decode.medusa import MedusaProposer
51
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
52
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
53
from vllm.v1.spec_decode.utils import is_spec_decode_supported
@@ -156,6 +157,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
157
self.device) # type: ignore
158
if self.speculative_config.method == "eagle3":
159
self.use_aux_hidden_state_outputs = True
160
+ elif self.speculative_config.method == "medusa":
161
+ self.drafter = MedusaProposer(
162
+ vllm_config=self.vllm_config,
163
+ device=self.device) # type: ignore
164
else:
165
raise ValueError("Unknown speculative decoding method: "
166
f"{self.speculative_config.method}")
@@ -1254,6 +1259,27 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1259
assert isinstance(self.drafter, NgramProposer)
1260
spec_token_ids = self.generate_draft_token_ids(
1261
valid_sampled_token_ids, sampling_metadata)
1262
+ elif self.speculative_config.method == "medusa":
1263
+ assert isinstance(self.drafter, MedusaProposer)
1264
+ if max_gen_len == 1:
1265
+ hidden_states = sample_hidden_states
1266
+ else:
1267
+ indices = []
1268
+ offset = 0
1269
+ for num_draft, tokens in zip(
1270
+ spec_decode_metadata.num_draft_tokens,
1271
+ valid_sampled_token_ids):
1272
+ indices.append(offset + len(tokens) - 1)
1273
+ offset += num_draft + 1
1274
+
1275
+ indices = torch.tensor(indices,
1276
+ device=sample_hidden_states.device)
1277
+ hidden_states = sample_hidden_states[indices]
1278
+
1279
+ spec_token_ids = self.drafter.propose(
1280
+ target_hidden_states=hidden_states,
1281
+ sampling_metadata=sampling_metadata,
1282
+ )
1283
elif self.speculative_config.use_eagle():
1284
assert isinstance(self.drafter, EagleProposer)
1285
# TODO(woosuk): Refactor the loop.