From 8530aed2b78a19eb00e242031c22bb78d615eea2 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 14 Oct 2026 17:29:23 -0700 Subject: [PATCH] [Spec Decode] Remove Medusa and MLP speculator support (#61701) Signed-off-by: Woosuk Kwon Co-authored-by: Claude Opus 5.5 --- .buildkite/test-amd.yaml | 2 - .buildkite/test_areas/misc.yaml | 2 - .buildkite/test_areas/spec_decode.yaml | 1 - .github/mergify.yml | 3 +- .github/workflows/issue_autolabel.yml | 1 - docs/features/speculative_decoding/README.md | 4 +- docs/features/speculative_decoding/mlp.md | 48 ----- docs/mkdocs/gen_files/generate_examples.py | 1 - .../config/draft_model_arch_groundtruth.json | 17 -- tests/config/test_model_arch_config.py | 27 +-- tests/models/registry.py | 10 - tests/models/test_registry.py | 3 - .../acceptance_rates/mtp_other/test_medusa.py | 50 ----- tests/v1/worker/test_gpu_warmup_blocks.py | 2 - vllm/config/speculative.py | 61 +----- vllm/config/vllm.py | 7 +- vllm/model_executor/models/interfaces_base.py | 20 +- vllm/model_executor/models/medusa.py | 202 ------------------ vllm/model_executor/models/mlp_speculator.py | 182 ---------------- vllm/model_executor/models/registry.py | 4 - vllm/transformers_utils/config.py | 4 +- vllm/transformers_utils/configs/__init__.py | 4 - vllm/transformers_utils/configs/medusa.py | 67 ------ .../configs/mlp_speculator.py | 69 ------ .../model_arch_config_convertor.py | 9 - vllm/v1/spec_decode/medusa.py | 79 ------- vllm/v1/worker/gpu_model_runner.py | 33 --- 27 files changed, 27 insertions(+), 885 deletions(-) delete mode 100644 docs/features/speculative_decoding/mlp.md delete mode 100644 tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py delete mode 100644 vllm/model_executor/models/medusa.py delete mode 100644 vllm/model_executor/models/mlp_speculator.py delete mode 100644 vllm/transformers_utils/configs/medusa.py delete mode 100644 vllm/transformers_utils/configs/mlp_speculator.py delete mode 100644 vllm/v1/spec_decode/medusa.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 0dc52c87ba..e70b2410b4 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -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 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 198def78e0..64f5d8047f 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -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 diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index d031d4fdea..eceb914c86 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -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 diff --git a/.github/mergify.yml b/.github/mergify.yml index 98c7d0327a..60661a4e10 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -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: diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index c3190f8f18..1852bd73b6 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -279,7 +279,6 @@ jobs: "speculative-decoding": { keywords: [ { term: "eagle", searchIn: "title" }, - { term: "medusa", searchIn: "title" }, ], substrings: [ { term: "speculative decoding", searchIn: "title" }, diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index 7e2c36d293..ddf3610946 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -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). | diff --git a/docs/features/speculative_decoding/mlp.md b/docs/features/speculative_decoding/mlp.md deleted file mode 100644 index 1af678a860..0000000000 --- a/docs/features/speculative_decoding/mlp.md +++ /dev/null @@ -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) diff --git a/docs/mkdocs/gen_files/generate_examples.py b/docs/mkdocs/gen_files/generate_examples.py index a514dcf879..e91f7c1d95 100644 --- a/docs/mkdocs/gen_files/generate_examples.py +++ b/docs/mkdocs/gen_files/generate_examples.py @@ -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 } diff --git a/tests/config/draft_model_arch_groundtruth.json b/tests/config/draft_model_arch_groundtruth.json index 5ea8136e9b..108a407a31 100644 --- a/tests/config/draft_model_arch_groundtruth.json +++ b/tests/config/draft_model_arch_groundtruth.json @@ -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" diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index 7d3063046e..65b4f440d8 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -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(): diff --git a/tests/models/registry.py b/tests/models/registry.py index c0568becd2..76e49aec33 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -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", diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index cffc728043..37a95d9f7f 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -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), ], diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py deleted file mode 100644 index 1e10d34527..0000000000 --- a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py +++ /dev/null @@ -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}" - ) diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index e0a99d5e8f..b48f8249e6 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -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), ], diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index b916d7da2b..c80016fd36 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -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, diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 58fe85c50e..06614f8326 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -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). diff --git a/vllm/model_executor/models/interfaces_base.py b/vllm/model_executor/models/interfaces_base.py index e1a6b650c6..7c7b291905 100644 --- a/vllm/model_executor/models/interfaces_base.py +++ b/vllm/model_executor/models/interfaces_base.py @@ -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 diff --git a/vllm/model_executor/models/medusa.py b/vllm/model_executor/models/medusa.py deleted file mode 100644 index 6adf3b7669..0000000000 --- a/vllm/model_executor/models/medusa.py +++ /dev/null @@ -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 diff --git a/vllm/model_executor/models/mlp_speculator.py b/vllm/model_executor/models/mlp_speculator.py deleted file mode 100644 index 0e58952102..0000000000 --- a/vllm/model_executor/models/mlp_speculator.py +++ /dev/null @@ -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 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 12599791e2..9bb811f203 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -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 = { diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 5b6586269c..892e00c43d 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -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] = { diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 336b96c1a3..cf1d08d67c 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -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", diff --git a/vllm/transformers_utils/configs/medusa.py b/vllm/transformers_utils/configs/medusa.py deleted file mode 100644 index e8cef7a35f..0000000000 --- a/vllm/transformers_utils/configs/medusa.py +++ /dev/null @@ -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 diff --git a/vllm/transformers_utils/configs/mlp_speculator.py b/vllm/transformers_utils/configs/mlp_speculator.py deleted file mode 100644 index f668e5ef52..0000000000 --- a/vllm/transformers_utils/configs/mlp_speculator.py +++ /dev/null @@ -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) diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index e3b19f3735..b16e0321c2 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -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, diff --git a/vllm/v1/spec_decode/medusa.py b/vllm/v1/spec_decode/medusa.py deleted file mode 100644 index 7c926ba837..0000000000 --- a/vllm/v1/spec_decode/medusa.py +++ /dev/null @@ -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) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 90059fd39f..4bb39c80cb 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -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), (