iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Spec Decode] Remove Medusa and MLP speculator support (#61701)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Woosuk Kwon
2026-10-14 17:29:23 -07:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 273a04964e
commit 8530aed2b7
27 changed files with 27 additions and 885 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test-amd.yaml
-2
View File
@@ -2669,7 +2669,6 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/model_executor/models/mlp_speculator.py
- tests/v1/spec_decode/test_acceptance_length.py
- vllm/platforms/rocm.py
commands:
@@ -5978,7 +5977,6 @@ steps:
- vllm/model_executor/models/gemma4_mtp.py
- vllm/model_executor/models/llama.py
- vllm/model_executor/models/llama_eagle3.py
- vllm/model_executor/models/medusa.py
- vllm/model_executor/models/registry.py
- tests/evals/gsm8k/
- tests/utils.py
.buildkite/test_areas/misc.yaml
-2
View File
@@ -841,7 +841,6 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/model_executor/models/mlp_speculator.py
- tests/v1/spec_decode/test_acceptance_length.py
commands:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
@@ -857,6 +856,5 @@ steps:
working_dir: /vllm-workspace/tests
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/model_executor/models/mlp_speculator.py
- tests/v1/spec_decode/test_acceptance_length.py
- vllm/platforms/rocm.py
.buildkite/test_areas/spec_decode.yaml
-1
View File
@@ -594,7 +594,6 @@ steps:
- vllm/model_executor/models/gemma4_mtp.py
- vllm/model_executor/models/llama.py
- vllm/model_executor/models/llama_eagle3.py
- vllm/model_executor/models/medusa.py
- vllm/model_executor/models/registry.py
- tests/evals/gsm8k/
- tests/utils.py
.github/mergify.yml
+1 -2
View File
@@ -494,8 +494,7 @@ pull_request_rules:
- files~=^tests/v1/e2e/spec_decode/
- files~=^examples/features/speculative_decoding/
- files~=^vllm/model_executor/models/.*eagle.*\.py
- files=vllm/model_executor/models/mlp_speculator.py
- files~=^vllm/transformers_utils/configs/(eagle|medusa|mlp_speculator)\.py
- files~=^vllm/transformers_utils/configs/eagle\.py
actions:
label:
add:
.github/workflows/issue_autolabel.yml
-1
View File
@@ -279,7 +279,6 @@ jobs:
"speculative-decoding": {
keywords: [
{ term: "eagle", searchIn: "title" },
{ term: "medusa", searchIn: "title" },
],
substrings: [
{ term: "speculative decoding", searchIn: "title" },
docs/features/speculative_decoding/README.md
+1 -3
View File
@@ -6,14 +6,13 @@ To train your own draft models for optimized speculative decoding, see [vllm-pro
## vLLM Speculation Methods
vLLM supports a variety of methods of speculative decoding. Model-based methods such as EAGLE, MTP, draft models, PARD and MLP provide the best latency reduction, while simpler methods such as n-gram and suffix decoding provide modest speedups without increasing workload during peak traffic.
vLLM supports a variety of methods of speculative decoding. Model-based methods such as EAGLE, MTP, draft models and PARD provide the best latency reduction, while simpler methods such as n-gram and suffix decoding provide modest speedups without increasing workload during peak traffic.
- [EAGLE](eagle.md)
- [LiLiCorr](lilicorr.md)
- [Multi-Token Prediction (MTP)](mtp.md)
- [Draft Model](draft_model.md)
- [Parallel Draft Model (PARD)](parallel_draft_model.md)
- [Multi-Layer Perceptron](mlp.md)
- [N-Gram](n_gram.md)
- [Suffix Decoding](suffix.md)
- [Hidden State Extraction](extract_hidden_states.md)
@@ -33,7 +32,6 @@ depend on your model family, traffic pattern, hardware, and sampling settings.
| MTP | High gain | Medium to high gain | Best when the target model has native MTP support. |
| Draft model | High gain | Medium gain | Needs a separate draft model. |
| Parallel Draft Model | High gain | Medium to high gain | Low draft model latency. |
| MLP speculator | Medium to high gain | Medium gain | Good when compatible MLP speculators are available. |
| N-gram | Low to medium gain | Medium gain | Lightweight and easy to enable. |
| Suffix decoding | Low to medium gain | Medium gain | No extra draft model; dynamic speculation depth. |
| Custom Proposer | Varies | Varies | Bring your own proposer class (experimental). |
docs/features/speculative_decoding/mlp.md
-48
View File
@@ -1,48 +0,0 @@
# MLP Draft Models
The following code configures vLLM to use speculative decoding where proposals are generated by draft models that condition draft predictions on both context vectors and sampled tokens. For more information, see [The Hitchhiker's Guide to Speculative Decoding](https://pytorch.org/blog/hitchhikers-guide-speculative-decoding/) and [IBM Research's Technical Report](https://arxiv.org/abs/2404.19124).
## MLP Drafter Example
```python
from vllm import LLM, SamplingParams
prompts = ["The future of AI is"]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
llm = LLM(
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
tensor_parallel_size=1,
speculative_config={
"model": "ibm-ai-platform/llama3-8b-accelerator",
"draft_tensor_parallel_size": 1,
"method": "mlp_speculator",
},
)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
```
!!! warning "Known issue"
`ibm-ai-platform/llama3-70b-accelerator` can fail with:
`AttributeError: 'MLPSpeculatorConfig' object has no attribute 'num_attention_heads'`.
Track status in [#34106](https://github.com/vllm-project/vllm/issues/34106)
and [#34163](https://github.com/vllm-project/vllm/pull/34163).
## Pre-Trained MLP Drafter Models
A variety of speculative models of this type are available on HF hub:
- [llama-13b-accelerator](https://huggingface.co/ibm-ai-platform/llama-13b-accelerator)
- [llama3-8b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-8b-accelerator)
- [codellama-34b-accelerator](https://huggingface.co/ibm-ai-platform/codellama-34b-accelerator)
- [llama2-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama2-70b-accelerator)
- [llama3-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-70b-accelerator)
- [granite-3b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-3b-code-instruct-accelerator)
- [granite-8b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-8b-code-instruct-accelerator)
- [granite-7b-instruct-accelerator](https://huggingface.co/ibm-granite/granite-7b-instruct-accelerator)
- [granite-20b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-20b-code-instruct-accelerator)
docs/mkdocs/gen_files/generate_examples.py
-1
View File
@@ -38,7 +38,6 @@ def title(text: str) -> str:
"openai": "OpenAI",
"lmcache": "LMCache",
"multilora": "MultiLoRA",
"mlpspeculator": "MLPSpeculator",
r"fp\d+": lambda x: x.group(0).upper(), # e.g. fp16, fp32
r"int\d+": lambda x: x.group(0).upper(), # e.g. int8, int16
}
tests/config/draft_model_arch_groundtruth.json
-17
View File
@@ -1,21 +1,4 @@
{
"abhigoyal/vllm-medusa-llama-68m-random": {
"architectures": [
"MedusaModel"
],
"model_type": "medusa",
"text_model_type": "medusa",
"hidden_size": 768,
"total_num_hidden_layers": 1,
"total_num_attention_heads": 0,
"head_size": "Error: integer division or modulo by zero",
"vocab_size": 32000,
"total_num_kv_heads": 0,
"num_experts": 0,
"is_deepseek_mla": false,
"is_multimodal_model": false,
"dtype": "torch.float32"
},
"luccafong/deepseek_mtp_draft_random": {
"architectures": [
"DeepSeekMTPModel"
tests/config/test_model_arch_config.py
+6 -21
View File
@@ -55,7 +55,6 @@ BASE_MODELS_TO_TEST = [
# (target_model, draft_model, trust_remote_code)
SPECULATIVE_MODELS = [
("JackFram/llama-68m", "abhigoyal/vllm-medusa-llama-68m-random", False),
("luccafong/deepseek_mtp_main_random", "luccafong/deepseek_mtp_draft_random", True),
("eagle618/deepseek-v3-random", "eagle618/eagle-deepseek-v3-random", True),
("meta-llama/Meta-Llama-3-8B-Instruct", "yuhuili/EAGLE-LLaMA3-Instruct-8B", True),
@@ -136,9 +135,7 @@ def test_qwen35_decision_config_selects_d3(tmp_path, has_decision_config, archit
assert not hasattr(config, "decision_config")
def _assert_model_arch_config(
model_config, expected: dict, check_head_size: bool = True
):
def _assert_model_arch_config(model_config, expected: dict):
"""Assert model_arch_config matches expected values."""
model_arch_config = model_config.model_arch_config
assert model_arch_config.architectures == expected["architectures"]
@@ -164,14 +161,10 @@ def _assert_model_arch_config(
config_format="hf",
)
assert str(torch_dtype) == expected["dtype"]
if check_head_size:
assert model_arch_config.head_size == expected["head_size"]
assert model_arch_config.head_size == expected["head_size"]
def _assert_model_config_methods(
model_config, expected: dict, check_head_size: bool = True
):
def _assert_model_config_methods(model_config, expected: dict):
"""Assert model_config methods return expected values."""
assert model_config.architectures == expected["architectures"]
assert model_config.get_vocab_size() == expected["vocab_size"]
@@ -182,9 +175,7 @@ def _assert_model_config_methods(
model_config.get_total_num_hidden_layers()
== expected["total_num_hidden_layers"]
)
if check_head_size:
assert model_config.get_head_size() == expected["head_size"]
assert model_config.get_head_size() == expected["head_size"]
def test_head_size_falls_back_when_head_dim_is_zero():
@@ -519,14 +510,8 @@ def test_draft_model_arch_config(
)
model_config = speculative_config.draft_model_config
# For medusa models, head_size may cause division by zero before
# model_arch_config was introduced, so we conditionally check it
check_head_size = isinstance(expected["head_size"], int)
_assert_model_arch_config(model_config, expected, check_head_size=check_head_size)
_assert_model_config_methods(
model_config, expected, check_head_size=check_head_size
)
_assert_model_arch_config(model_config, expected)
_assert_model_config_methods(model_config, expected)
def test_deepseek_v4_convertor_splits_vision_architecture():
tests/models/registry.py
-10
View File
@@ -1456,16 +1456,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
_SPECULATIVE_DECODING_EXAMPLE_MODELS = {
# [Medusa]
"MedusaModel": _HfExamplesInfo(
"JackFram/llama-68m", speculative_model="abhigoyal/vllm-medusa-llama-68m-random"
),
# Temporarily disabled.
# TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
# "MLPSpeculatorPreTrainedModel": _HfExamplesInfo(
# "JackFram/llama-160m",
# speculative_model="ibm-ai-platform/llama-160m-accelerator"
# ),
# [DFlash]
"DFlashDraftModel": _HfExamplesInfo(
"Qwen/Qwen3-4B",
tests/models/test_registry.py
-3
View File
@@ -113,9 +113,6 @@ def test_registry_model_property(model_arch, is_mm, score_type):
@pytest.mark.parametrize(
"model_arch,is_pp",
[
# TODO(woosuk): Re-enable this once the MLP Speculator is supported
# in V1.
# ("MLPSpeculatorPreTrainedModel", False),
("DeepseekV2ForCausalLM", True),
("Qwen2VLForConditionalGeneration", True),
],
tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py
-50
View File
@@ -1,50 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts
from vllm import SamplingParams
from vllm.config import CompilationConfig
from ...utils import compute_acceptance_rate
def test_medusa_acceptance_rate(
sampling_config: SamplingParams,
vllm_runner,
):
"""Verify a trained Medusa checkpoint achieves nonzero acceptance rate.
Uses the canonical FasterDecoding vicuna-7b checkpoint to confirm the
speculation path actually accepts tokens — unlike test_medusa_correctness,
which uses a random head and only validates output correctness.
"""
target_model = "lmsys/vicuna-7b-v1.3"
medusa_model = "FasterDecoding/medusa-vicuna-7b-v1.3"
prompts = _build_gsm8k_prompts(num_questions=10, num_shots=1)[0]
with vllm_runner(
target_model,
block_size=None,
trust_remote_code=False,
speculative_config={
"method": "medusa",
"model": medusa_model,
"num_speculative_tokens": 3,
},
max_model_len=1024,
enforce_eager=True,
enable_chunked_prefill=None,
disable_log_stats=False,
compilation_config=CompilationConfig(),
) as spec_runner:
spec_runner.llm.generate(prompts, sampling_config)
metrics = spec_runner.llm.get_metrics()
acceptance_rate = compute_acceptance_rate(metrics)
min_acceptance_rate = 0.198
print(f"Medusa acceptance rate: {acceptance_rate:.4f} (min {min_acceptance_rate})")
# Regression guard at 90% of the measured baseline.
assert acceptance_rate >= min_acceptance_rate, (
f"Medusa acceptance rate {acceptance_rate:.4f} below min {min_acceptance_rate}"
)
tests/v1/worker/test_gpu_warmup_blocks.py
-2
View File
@@ -334,8 +334,6 @@ def test_reserved_block_count_matches_real_kv_cache_manager():
("dflash", None, NUM_SPEC_STEPS + 1),
("ngram", None, 0),
("ngram_gpu", None, 0),
("medusa", None, 0),
("mlp_speculator", None, 0),
("suffix", None, 0),
("extract_hidden_states", None, 0),
],
vllm/config/speculative.py
+9 -52
View File
@@ -72,8 +72,6 @@ EagleModelTypes = Literal[
]
SpeculativeMethod = Literal[
"ngram",
"medusa",
"mlp_speculator",
"draft_model",
"suffix",
"custom_class",
@@ -1272,20 +1270,12 @@ class SpeculativeConfig:
self.prompt_lookup_min = 0
if self.model is not None:
# Old-format Medusa checkpoints (e.g. FasterDecoding/medusa-*)
# lack a model_type key in config.json, so AutoConfig cannot
# detect them. When the method is explicitly "medusa", inject
# model_type so MedusaConfig.from_pretrained is used instead.
draft_hf_overrides: HfOverrides
if self.method == "medusa":
draft_hf_overrides = {"model_type": "medusa"}
else:
# Compose any callable hf_overrides set on the target so the
# draft config receives the same transform (e.g. the test
# shrink). Dict overrides stay target-only.
draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides(
self.target_model_config.hf_overrides
)
# Compose any callable hf_overrides set on the target so the
# draft config receives the same transform (e.g. the test
# shrink). Dict overrides stay target-only.
draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides(
self.target_model_config.hf_overrides
)
self.draft_model_config = ModelConfig(
model=self.model,
runner="draft",
@@ -1312,17 +1302,6 @@ class SpeculativeConfig:
config_format=self.target_model_config.config_format,
)
# Old-format Medusa checkpoints (e.g. FasterDecoding/medusa-*)
# omit vocab_size in config.json, so MedusaConfig falls back to
# its default (32001). Align with the target model's vocab size
# to avoid shape mismatches when loading LM-head weights.
if self.method == "medusa":
target_vocab = self.target_model_config.hf_config.vocab_size
draft_hf = self.draft_model_config.hf_config
if draft_hf.vocab_size != target_vocab:
draft_hf.vocab_size = target_vocab
draft_hf.truncated_vocab_size = target_vocab
# Automatically detect the method
if self.method in ("eagle", "eagle3", "dflash", "dspark"):
pass
@@ -1353,10 +1332,6 @@ class SpeculativeConfig:
)
):
self.method = "dspark"
elif self.draft_model_config.hf_config.model_type == "medusa":
self.method = "medusa"
elif self.draft_model_config.hf_config.model_type == "mlp_speculator":
self.method = "mlp_speculator"
elif self.draft_model_config.hf_config.model_type in get_args(
MTPModelTypes
):
@@ -1468,13 +1443,6 @@ class SpeculativeConfig:
if self.method in ("dflash", "dspark"):
self.parallel_drafting = True
if self.num_speculative_tokens is not None and hasattr(
self.draft_model_config.hf_config, "num_lookahead_tokens"
):
self.draft_model_config.hf_config.num_lookahead_tokens = (
self.num_speculative_tokens
)
n_predict = getattr(
self.draft_model_config.hf_config, "n_predict", None
)
@@ -1550,7 +1518,6 @@ class SpeculativeConfig:
SpeculativeConfig._verify_and_get_draft_tp(
self.target_parallel_config,
self.draft_tensor_parallel_size,
self.draft_model_config.hf_config,
)
)
self.draft_model_config.max_model_len = (
@@ -1712,7 +1679,6 @@ class SpeculativeConfig:
def _verify_and_get_draft_tp(
target_parallel_config: ParallelConfig,
speculative_draft_tensor_parallel_size: int | None,
draft_hf_config: PreTrainedConfig,
) -> int:
"""Verifies and adjusts the tensor parallel size for a draft model
specified using speculative_draft_tensor_parallel_size.
@@ -1720,18 +1686,9 @@ class SpeculativeConfig:
# If speculative_draft_tensor_parallel_size is unset then set it
# appropriately else verify that it is set correctly.
if speculative_draft_tensor_parallel_size is None:
if draft_hf_config.model_type == "mlp_speculator":
speculative_draft_tensor_parallel_size = 1
if target_parallel_config.tensor_parallel_size > 1:
logger.warning(
"%s cannot currently be run with tp>1; "
"setting speculative_draft_tensor_parallel_size=1",
draft_hf_config.model_type,
)
else:
speculative_draft_tensor_parallel_size = (
target_parallel_config.tensor_parallel_size
)
speculative_draft_tensor_parallel_size = (
target_parallel_config.tensor_parallel_size
)
elif speculative_draft_tensor_parallel_size not in (
1,
target_parallel_config.tensor_parallel_size,
vllm/config/vllm.py
+1 -6
View File
@@ -2938,12 +2938,7 @@ class VllmConfig:
unsupported.append("pipeline parallelism with external_launcher")
if speculative_config is not None:
if speculative_config.method in (
"suffix",
"medusa",
"mlp_speculator",
"custom_class",
):
if speculative_config.method in ("suffix", "custom_class"):
unsupported.append(f"speculative method '{speculative_config.method}'")
# V2 EagleSpeculator does not support parallel_drafting (for P-Eagle).
vllm/model_executor/models/interfaces_base.py
+8 -12
View File
@@ -32,19 +32,13 @@ else:
logger = init_logger(__name__)
# The type of hidden states
# Currently, T = torch.Tensor for all models except for Medusa
# which has T = list[torch.Tensor]
T = TypeVar("T", default=torch.Tensor)
T_co = TypeVar("T_co", default=torch.Tensor, covariant=True)
# NOTE: Unlike those in `interfaces.py`, we don't define `ClassVar` tags
# for the base interfaces to avoid breaking OOT registration for existing models
# that don't inherit from the base interface classes
@runtime_checkable
class VllmModel(Protocol[T_co]):
class VllmModel(Protocol):
"""The interface required for all models in vLLM."""
def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: ...
@@ -53,7 +47,9 @@ class VllmModel(Protocol[T_co]):
"""Apply token embeddings to `input_ids`."""
...
def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> T_co: ...
def forward(
self, input_ids: torch.Tensor, positions: torch.Tensor
) -> torch.Tensor: ...
def _check_vllm_model_init(model: type[object] | object) -> bool:
@@ -111,13 +107,13 @@ def is_vllm_model(
@runtime_checkable
class VllmModelForTextGeneration(VllmModel[T], Protocol[T]):
class VllmModelForTextGeneration(VllmModel, Protocol):
"""The interface required for all generative models in vLLM."""
def compute_logits(
self,
hidden_states: T,
) -> T | None:
hidden_states: torch.Tensor,
) -> torch.Tensor | None:
"""Return `None` if TP rank > 0."""
...
@@ -145,7 +141,7 @@ def is_text_generation_model(
@runtime_checkable
class VllmModelForPooling(VllmModel[T_co], Protocol[T_co]):
class VllmModelForPooling(VllmModel, Protocol):
"""The interface required for all pooling models in vLLM."""
is_pooling_model: ClassVar[Literal[True]] = True
vllm/model_executor/models/medusa.py
-202
View File
@@ -1,202 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from .utils import maybe_prefix
class ResidualBlock(nn.Module):
def __init__(self, config: VllmConfig, hidden_size: int, num_layers: int) -> None:
super().__init__()
self.layers = nn.ModuleList(
[
nn.Linear(
hidden_size,
hidden_size,
bias=getattr(config, "medusa_fc_bias", False),
)
for _ in range(num_layers)
]
)
self.act = nn.SiLU()
def forward(self, x: torch.Tensor) -> torch.Tensor:
for layer in self.layers:
x = x + self.act(layer(x))
return x
class Medusa(nn.Module):
"""This class implements the Medusa draft model from the paper: https://arxiv.org/abs/2401.10774
Reference implementation: https://github.com/FasterDecoding/Medusa
Differences from reference implementation:
1. Currently this only supports generating proposals from top-1 tokens.
2. We have an optional token_map which reduces draft vocab to most
frequently used tokens to give some additional speed-up by reducing
sampling overhead. This is disabled unless the checkpoint file has
explicit token_map tensor and config has an optional attribute
truncated_vocab_size < vocab_size. To use this technique, one has to find
the top-k most frequent tokens in target dataset and add that as a tensor
in the draft checkpoint (using key token_map). Also, the draft config
needs to have truncated_vocab_size (=k) as an attribute."""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
speculative_config = vllm_config.speculative_config
assert speculative_config is not None
config = speculative_config.draft_model_config.hf_config
super().__init__()
self.config = config
self.blocks = nn.ModuleList(
[
ResidualBlock(
config=config,
hidden_size=self.config.hidden_size,
num_layers=self.config.num_hidden_layers,
)
for _ in range(self.config.num_heads)
]
)
self.orig_vocab_size = config.vocab_size
self.truncated_vocab_size = config.truncated_vocab_size
if getattr(config, "original_lm_head", False):
self.lm_head = ParallelLMHead(
self.truncated_vocab_size,
config.hidden_size,
prefix=maybe_prefix(prefix, "lm_head"),
)
self.lm_heads = [self.lm_head for _ in range(self.config.num_heads)]
else:
self.lm_heads = nn.ModuleList(
[
ParallelLMHead(
config.vocab_size,
config.hidden_size,
prefix=maybe_prefix(prefix, f"lm_heads.{i}"),
)
for i in range(self.config.num_heads)
]
)
logit_scale = getattr(config, "logit_scale", 1.0)
self.logits_processor = LogitsProcessor(
config.vocab_size, self.truncated_vocab_size, logit_scale
)
# Token map is a idx to token mapping to reduce the vocab size for
# the draft model. Using smaller vocab size for draft, containing
# only most frequent tokens reduces the speculation overhead. This
# doesn't affect the acceptance rate much and thus gives more speed
# -up. By default, this is disabled and is only used if the EAGLE
# checkpoint file has token_map tensor.
self.token_map = None
def forward(self, hidden_states: torch.Tensor) -> list[torch.Tensor]:
return [block(hidden_states) for block in self.blocks]
def compute_logits(
self,
hidden_states: list[torch.Tensor],
) -> list[torch.Tensor]:
logits_lst: list[torch.Tensor] = []
for hs, lm_head in zip(hidden_states, self.lm_heads):
_logits = self.logits_processor(lm_head, hs)
if _logits is None:
# _logits should only be None on rank > 0, in which case
# it should remain true for every lm_head
assert len(logits_lst) == 0
continue
if self.token_map is None:
logits_lst.append(_logits)
else:
logits_lst.append(
-torch.inf
* torch.ones(
size=(*_logits.shape[:-1], self.orig_vocab_size),
device=_logits.device,
dtype=_logits.dtype,
)
)
logits_lst[-1][..., self.token_map] = _logits
return logits_lst
@staticmethod
def _remap_old_checkpoint_key(name: str) -> str:
"""Map old FasterDecoding checkpoint keys to vLLM parameter names.
Old format uses bare numeric prefixes:
'{head}.{layer}.linear.weight' -> 'blocks.{head}.layers.{layer}.weight'
'{head}.{layer}.linear.bias' -> 'blocks.{head}.layers.{layer}.bias'
'{head}.{N}.weight' -> 'lm_heads.{head}.weight'
"""
parts = name.split(".")
if len(parts) >= 3 and parts[0].isdigit() and parts[1].isdigit():
head, layer = parts[0], parts[1]
rest = parts[2:]
if "linear" in rest:
param = ".".join(rest[rest.index("linear") + 1 :])
return f"blocks.{head}.layers.{layer}.{param}"
else:
return f"lm_heads.{head}.{'.'.join(rest)}"
return name
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
weights_map = {}
for name, loaded_weight in weights:
name = name.replace("medusa_heads.", "")
name = self._remap_old_checkpoint_key(name)
if name == "token_map":
if self.truncated_vocab_size < self.orig_vocab_size:
self.token_map = nn.Parameter(loaded_weight, requires_grad=False)
elif name in params_dict:
weights_map[name] = loaded_weight
elif (
getattr(self.config, "original_lm_head", False)
and name == "lm_heads.0.weight"
):
weights_map["lm_head.weight"] = loaded_weight
for name, loaded_weight in weights_map.items():
if (
"lm_head" in name
and self.token_map is not None
and loaded_weight.shape[0] > self.token_map.shape[0]
):
loaded_weight = loaded_weight[self.token_map]
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
if self.token_map is not None:
self.token_map.to(device=self.lm_heads[0].weight.device)
assert (self.truncated_vocab_size == self.orig_vocab_size) or (
self.token_map is not None
)
return loaded_params
vllm/model_executor/models/mlp_speculator.py
-182
View File
@@ -1,182 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
from collections.abc import Iterable
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from .utils import maybe_prefix
class MLPSpeculatorLayerNorm(nn.Module):
"""A L2 normalization implementation
...
Args:
----
normalized_shape : int
Dimensionality of input data (size of final tensor axis)
eps : float
Safety term to prevent division by zero. Make sure the chosen value
fits in the range of your encoding scheme
(i.e. fp16 requires eps >= 6e-8).
elementwise_scale_and_shift : bool
Include a learned scaling and shift term after normalization.
"""
def __init__(
self,
normalized_shape,
eps=1e-06,
elementwise_scale_and_shift=True,
):
super().__init__()
self.elementwise_scale_and_shift = elementwise_scale_and_shift
if self.elementwise_scale_and_shift:
self.weight = nn.Parameter(torch.empty(normalized_shape))
self.bias = nn.Parameter(torch.empty(normalized_shape))
self.eps = eps
def forward(self, x):
xf = x
xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
x = xf.type_as(x)
if self.elementwise_scale_and_shift:
x = self.weight * x
x = x + self.bias
return x
class MLPSpeculator(nn.Module):
"""An implementation of the speculative models introduced in
"Accelerating Production LLMs with Combined Token/Embedding
Speculators"
https://arxiv.org/pdf/2404.19124
Trained speculators of this type are available on HF hub at:
https://huggingface.co/ibm-ai-platform and https://huggingface.co/ibm-granite
"""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
config = vllm_config.model_config.hf_config
self.n_predict = config.n_predict
self.vocab_size = config.vocab_size
self.emb_dim = config.emb_dim
self.inner_dim = config.inner_dim if config.inner_dim != 0 else config.emb_dim
self.max_speculative_tokens = config.num_lookahead_tokens
self.tie_weights = config.tie_weights
self.scale_input = config.scale_input
if self.tie_weights:
assert self.n_predict > 1, (
"You cannot tie weights between stages when only 1 exists"
)
embedding = VocabParallelEmbedding(
config.vocab_size, self.inner_dim, org_num_embeddings=config.vocab_size
)
self.emb = nn.ModuleList([embedding] * self.max_speculative_tokens)
# the initial projection from the base model may
# have a different size, so that stays separate.
proj_first = nn.Linear(self.emb_dim, self.inner_dim, bias=False)
proj_tied = nn.Linear(self.inner_dim, self.inner_dim, bias=False)
self.proj = nn.ModuleList(
[proj_first] + [proj_tied] * (self.max_speculative_tokens - 1)
)
self.head = nn.ModuleList(
[
ParallelLMHead(
self.vocab_size,
self.inner_dim,
bias=False,
prefix=maybe_prefix(prefix, f"head.{i}"),
)
for i in range(self.max_speculative_tokens)
]
)
ln = MLPSpeculatorLayerNorm(
self.inner_dim, elementwise_scale_and_shift=True
)
self.ln = nn.ModuleList([ln] * self.max_speculative_tokens)
else:
self.emb = nn.ModuleList(
[
VocabParallelEmbedding(
config.vocab_size,
self.inner_dim,
)
for _ in range(self.max_speculative_tokens)
]
)
self.proj = nn.ModuleList(
[
nn.Linear(
(self.emb_dim if i == 0 else self.inner_dim),
self.inner_dim,
bias=False,
)
for i in range(self.max_speculative_tokens)
]
)
self.head = nn.ModuleList(
[
ParallelLMHead(
self.vocab_size,
self.inner_dim,
bias=False,
prefix=maybe_prefix(prefix, f"head.{i}"),
)
for i in range(self.max_speculative_tokens)
]
)
self.ln = nn.ModuleList(
[
MLPSpeculatorLayerNorm(
self.inner_dim, elementwise_scale_and_shift=True
)
for _ in range(self.max_speculative_tokens)
]
)
if self.scale_input:
self.ln0 = MLPSpeculatorLayerNorm(
self.emb_dim, elementwise_scale_and_shift=False
)
self.state_weight = 0.5 ** (0.5 / config.n_predict)
self.emb_weight = math.sqrt((1 - self.state_weight**2) * (self.inner_dim / 2))
self.activation = nn.GELU()
self.config = config
self.logits_processor = LogitsProcessor(
config.vocab_size, config.vocab_size, 1.0
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
for name, loaded_weight in weights:
name = name.replace("speculator.", "")
param = params_dict.get(name)
if param is not None:
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
return loaded_params
vllm/model_executor/models/registry.py
-4
View File
@@ -691,7 +691,6 @@ _SPECULATIVE_DECODING_MODELS = {
"Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"),
"GlmOcrMTPModel": ("glm_ocr_mtp", "GlmOcrMTP"),
"Glm5NextMTPModel": ("vllm.models.glm5next", "Glm5NextMTP"),
"MedusaModel": ("medusa", "Medusa"),
"OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"),
"Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),
"Qwen4ExpMTP": ("vllm.models.qwen4_exp", "Qwen4ExpMTP"),
@@ -702,9 +701,6 @@ _SPECULATIVE_DECODING_MODELS = {
"HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"),
"HYV4MTPModel": ("vllm.models.hy_v4", "HYV4MTP"),
"KimiK3MTPModel": ("vllm.models.kimi_k3", "KimiK3MTP"),
# Temporarily disabled.
# # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
# "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
}
_TRANSFORMERS_SUPPORTED_MODELS = {
vllm/transformers_utils/config.py
+1 -3
View File
@@ -145,8 +145,6 @@ _CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict(
kimi_k3="KimiK3Config",
RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct)
RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct)
mlp_speculator="MLPSpeculatorConfig",
medusa="MedusaConfig",
midashenglm="MiDashengLMConfig",
minimax_m3_vl="MiniMaxM3Config",
minimax_m3_mtp="MiniMaxM3MTPConfig",
@@ -169,7 +167,7 @@ _CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict(
inkling_model="InklingModelConfig",
)
_SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators", "medusa"}
_SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators"}
# Checkpoints whose config.json has no model_type, only remote code in auto_map
_ARCHITECTURE_TO_MODEL_TYPE: dict[str, str] = {
vllm/transformers_utils/configs/__init__.py
-4
View File
@@ -43,12 +43,10 @@ _CLASS_TO_MODULE: dict[str, str] = {
# tiiuae/falcon-7b(-instruct) models. Newer Falcon models will use the
# `FalconConfig` class from the official HuggingFace transformers library.
"RWConfig": "vllm.transformers_utils.configs.falcon",
"MedusaConfig": "vllm.transformers_utils.configs.medusa",
"MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",
"MiniMaxM3Config": "vllm.transformers_utils.configs.minimax_m3",
"MiniMaxM3MTPConfig": "vllm.transformers_utils.configs.minimax_m3",
"MiniMaxM3TextConfig": "vllm.transformers_utils.configs.minimax_m3",
"MLPSpeculatorConfig": "vllm.transformers_utils.configs.mlp_speculator",
"Moondream3Config": "vllm.transformers_utils.configs.moondream3",
"Moondream3TextConfig": "vllm.transformers_utils.configs.moondream3",
"Moondream3VisionConfig": "vllm.transformers_utils.configs.moondream3",
@@ -116,12 +114,10 @@ __all__ = [
"HyperCLOVAXConfig",
"IsaacConfig",
"RWConfig",
"MedusaConfig",
"MiDashengLMConfig",
"MiniMaxM3Config",
"MiniMaxM3MTPConfig",
"MiniMaxM3TextConfig",
"MLPSpeculatorConfig",
"Moondream3Config",
"Moondream3TextConfig",
"Moondream3VisionConfig",
vllm/transformers_utils/configs/medusa.py
-67
View File
@@ -1,67 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from transformers import PreTrainedConfig
from vllm.transformers_utils.utils import without_trust_remote_code
class MedusaConfig(PreTrainedConfig):
model_type = "medusa"
def __init__(
self,
hidden_size: int = 4096,
vocab_size: int = 32001,
num_heads: int = 5,
num_hidden_layers: int = 1,
max_paths: int = 64,
topk: int = 10,
truncated_vocab_size: int | None = None,
**kwargs,
):
self.hidden_size = hidden_size
self.vocab_size = vocab_size
self.num_heads = num_heads
self.num_hidden_layers = num_hidden_layers
self.max_paths = max_paths
self.topk = topk
self.max_seq_len = int(2**20)
self.truncated_vocab_size = (
vocab_size if truncated_vocab_size is None else truncated_vocab_size
)
if "architectures" not in kwargs:
kwargs["architectures"] = ["MedusaModel"]
super().__init__(**kwargs)
@classmethod
def from_pretrained(
cls,
pretrained_model_name_or_path: str | os.PathLike,
**kwargs,
) -> "MedusaConfig":
config_dict, kwargs = cls.get_config_dict(
pretrained_model_name_or_path, **without_trust_remote_code(kwargs)
)
for k in list(config_dict.keys()):
if "num" in k:
if "heads" in k:
config_dict["num_heads"] = config_dict.pop(k)
elif "layers" in k:
config_dict["num_hidden_layers"] = config_dict.pop(k)
return cls.from_dict(config_dict, **kwargs)
@property
def num_attention_heads(self):
return 0
@property
def num_lookahead_tokens(self):
return self.num_heads
@num_lookahead_tokens.setter
def num_lookahead_tokens(self, num_lookahead_tokens: int):
self.num_heads = num_lookahead_tokens
vllm/transformers_utils/configs/mlp_speculator.py
-69
View File
@@ -1,69 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from transformers import PreTrainedConfig
class MLPSpeculatorConfig(PreTrainedConfig):
model_type = "mlp_speculator"
attribute_map = {
"hidden_size": "emb_dim",
}
def __init__(
self,
vocab_size: int = 32000,
emb_dim: int = 4096,
inner_dim: int = 0,
n_predict: int = 3,
top_k_tokens_per_head: list[int] | None = None,
n_candidates: int = 5,
tie_weights: bool = False,
scale_input: bool = False,
**kwargs,
):
"""Initialize an MLPSpeculatorConfig.
Args:
vocab_size: int
the model vocab size
emb_dim: int
the model embedding dimension
inner_dim: int
the inner dimension of the model. If 0, will be the emb_dim.
n_predict: int
the number of lookaheads for the speculator
top_k_tokens_per_head: list[int]
Number of tokens to consider from each head when forming the
candidate tree.
For each candidate branch in the tree, head n produces topk[n]
additional sub-branches.
NOTE: This parameter is currently unused.
n_candidates: int
number of child candidates to create per sequence
tie_weights: bool
If true, use a single set of weights for every model
head/stage after the first. The initial projection
from the base model may have a different size, so that
stays separate.
scale_input: bool
if True, will scale the initial hidden states from
the base model.
"""
if top_k_tokens_per_head is None:
top_k_tokens_per_head = [5, 4, 3]
assert len(top_k_tokens_per_head) == n_predict
self.vocab_size = vocab_size
self.emb_dim = emb_dim
self.inner_dim = inner_dim
self.n_predict = n_predict
self.top_k_tokens_per_head = top_k_tokens_per_head
self.n_candidates = n_candidates
self.num_lookahead_tokens = n_predict
self.tie_weights = tie_weights
self.scale_input = scale_input
super().__init__(**kwargs)
vllm/transformers_utils/model_arch_config_convertor.py
-9
View File
@@ -496,14 +496,6 @@ class TerratorchModelArchConfigConvertor(ModelArchConfigConvertorBase):
return 0
class MedusaModelArchConfigConvertor(ModelArchConfigConvertorBase):
def get_head_size(self) -> int:
return 0
def get_total_num_kv_heads(self) -> int:
return 0
class Zamba2ModelArchConfigConvertor(ModelArchConfigConvertorBase):
def get_head_size(self) -> int:
return getattr(self.hf_text_config, "attention_head_dim", 0)
@@ -849,7 +841,6 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
"longcat_flash": LongCatFlashModelArchConfigConvertor,
"longcat_flash_mtp": LongCatFlashMTPModelArchConfigConvertor,
"mamba": MambaModelArchConfigConvertor,
"medusa": MedusaModelArchConfigConvertor,
"mimo_mtp": MimoMTPModelArchConfigConvertor,
"mimo_v2": MimoV2ModelArchConfigConvertor,
"mimo_v2_flash": MimoV2ModelArchConfigConvertor,
vllm/v1/spec_decode/medusa.py
-79
View File
@@ -1,79 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.forward_context import set_forward_context
from vllm.logger import init_logger
from vllm.model_executor.model_loader import get_model
from vllm.model_executor.models.interfaces import is_mixture_of_experts
from vllm.v1.sample.metadata import SamplingMetadata
# Initialize logger
logger = init_logger(__name__)
class MedusaProposer:
"""Medusa proposer class for generating token sequences."""
def __init__(
self,
vllm_config: VllmConfig,
device: torch.device,
):
# Save config parameters
self.vllm_config = vllm_config
assert vllm_config.speculative_config is not None, (
"Speculative config must be set"
)
self.spec_config = vllm_config.speculative_config
self.device = device
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
self.hidden_size = self.spec_config.draft_model_config.get_hidden_size()
self.dtype = vllm_config.model_config.dtype
self.num_speculative_tokens = self.spec_config.num_speculative_tokens
def propose(
self,
num_speculative_tokens: int,
target_hidden_states: torch.Tensor,
sampling_metadata: SamplingMetadata,
slot_mappings: dict[str, torch.Tensor]
| list[dict[str, torch.Tensor]]
| None = None, # unused
) -> torch.Tensor:
assert num_speculative_tokens == self.num_speculative_tokens
# Generate blocks and compute logits
blocks = self.model(target_hidden_states)
logits = self.model.compute_logits(blocks)
# 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)
return draft_tokens
def load_model(self, target_model: nn.Module) -> None:
from vllm.compilation.backends import set_model_tag
with set_model_tag("medusa_head"):
self.model = get_model(
vllm_config=self.vllm_config,
model_config=self.spec_config.draft_model_config,
)
assert not (
is_mixture_of_experts(self.model)
and self.vllm_config.parallel_config.enable_eplb
), "EPLB for Medusa is not supported"
@torch.inference_mode()
def dummy_run(self, num_tokens: int) -> None:
hidden_states = torch.zeros(
(self.max_num_tokens, self.hidden_size),
dtype=self.dtype,
device=self.device,
)
with set_forward_context(None, self.vllm_config, num_tokens=num_tokens):
self.model(hidden_states)
vllm/v1/worker/gpu_model_runner.py
-33
View File
@@ -180,7 +180,6 @@ from vllm.v1.spec_decode.draft_model import DraftModelProposer
from vllm.v1.spec_decode.eagle import EagleProposer
from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
from vllm.v1.spec_decode.gemma4 import Gemma4Proposer
from vllm.v1.spec_decode.medusa import MedusaProposer
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
from vllm.v1.spec_decode.ngram_proposer_gpu import (
NgramProposerGPU,
@@ -595,7 +594,6 @@ class GPUModelRunner(
| EagleProposer
| DFlashProposer
| DraftModelProposer
| MedusaProposer
| ExtractHiddenStatesProposer
| Gemma4Proposer
| Step3p5MTPProposer
@@ -646,10 +644,6 @@ class GPUModelRunner(
self.use_aux_hidden_state_outputs = (
self.drafter.eagle3_use_aux_hidden_state
)
elif self.speculative_config.method == "medusa":
self.drafter = MedusaProposer(
vllm_config=self.vllm_config, device=self.device
)
elif self.speculative_config.method == "extract_hidden_states":
self.drafter = ExtractHiddenStatesProposer(
vllm_config=self.vllm_config, device=self.device
@@ -4830,33 +4824,6 @@ class GPUModelRunner(
sampled_token_ids,
slot_mappings=slot_mappings,
)
elif spec_config.method == "medusa":
assert isinstance(sampled_token_ids, list)
assert isinstance(self.drafter, MedusaProposer)
if sample_hidden_states.shape[0] == len(sampled_token_ids):
# The input to the target model does not include draft tokens.
hidden_states = sample_hidden_states
else:
indices = []
offset = 0
assert spec_decode_metadata is not None, (
"No spec decode metadata for medusa"
)
for num_draft, tokens in zip(
spec_decode_metadata.num_draft_tokens, sampled_token_ids
):
indices.append(offset + len(tokens) - 1)
offset += num_draft + 1
indices = async_tensor_h2d(indices, device=self.device)
hidden_states = sample_hidden_states[indices]
draft_token_ids = self.drafter.propose(
num_speculative_tokens=num_spec_tokens_to_schedule,
target_hidden_states=hidden_states,
sampling_metadata=sampling_metadata,
slot_mappings=slot_mappings,
)
elif spec_config.uses_extract_hidden_states():
assert isinstance(self.drafter, ExtractHiddenStatesProposer)
assert isinstance(sampled_token_ids, torch.Tensor), (