Add Medusa speculative decoding e2e test (#41396)
Signed-off-by: Anshika Ojha <anshikao@nvidia.com> Signed-off-by: Rishi Puri <riship@nvidia.com> Signed-off-by: Rishi Puri <puririshi98@berkeley.edu> Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: Anshika Ojha <215760622+ojhaanshika@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Stefano Castagnetta <scastagnetta@nvidia.com>
Rishi Puri committed
Jun 30, 2026 at 15:02 UTC
727971f1c1056f7b8a3ea89713526756e6f7c8cf
4 files changed
+85
-3
tests/v1/e2e/spec_decode/test_spec_decode.py
+41
@@ -718,6 +718,47 @@ def test_eagle_correctness_heavy(
718
)
719
720
721
+@large_gpu_mark(min_gb=24)
722
+def test_medusa_acceptance_rate(
723
+ sampling_config: SamplingParams,
724
+):
725
+ """Verify a trained Medusa checkpoint achieves nonzero acceptance rate.
726
+
727
+ Uses the canonical FasterDecoding vicuna-7b checkpoint to confirm the
728
+ speculation path actually accepts tokens — unlike test_medusa_correctness,
729
+ which uses a random head and only validates output correctness.
730
+ """
731
+ target_model = "lmsys/vicuna-7b-v1.3"
732
+ medusa_model = "FasterDecoding/medusa-vicuna-7b-v1.3"
733
+ prompts = _build_gsm8k_prompts(num_questions=10, num_shots=1)[0]
734
+
735
+ spec_llm = LLM(
736
+ model=target_model,
737
+ speculative_config={
738
+ "method": "medusa",
739
+ "model": medusa_model,
740
+ "num_speculative_tokens": 3,
741
+ },
742
+ max_model_len=1024,
743
+ enforce_eager=True,
744
+ disable_log_stats=False,
745
+ )
746
+ spec_llm.generate(prompts, sampling_config)
747
+ metrics = spec_llm.get_metrics()
748
+ acceptance_rate = compute_acceptance_rate(metrics)
749
+ del spec_llm
750
+ torch.accelerator.empty_cache()
751
+ cleanup_dist_env_and_memory()
752
+
753
+ min_acceptance_rate = 0.198
754
+ print(f"Medusa acceptance rate: {acceptance_rate:.4f} (min {min_acceptance_rate})")
755
+
756
+ # Regression guard at 90% of the measured baseline.
757
+ assert acceptance_rate >= min_acceptance_rate, (
758
+ f"Medusa acceptance rate {acceptance_rate:.4f} below min {min_acceptance_rate}"
759
+ )
760
+
761
+
762
@pytest.mark.parametrize(
763
["model_setup", "mm_enabled", "expected_accuracy_threshold"],
764
[
vllm/config/speculative.py
+22
-2
@@ -9,7 +9,7 @@ from typing_extensions import Self
9
10
from vllm.config import LoadConfig
11
from vllm.config.kernel import MoEBackend
12
-from vllm.config.model import ModelConfig
12
+from vllm.config.model import HfOverrides, ModelConfig
13
from vllm.config.parallel import ParallelConfig
14
from vllm.config.utils import config
15
from vllm.logger import init_logger
@@ -712,6 +712,15 @@ class SpeculativeConfig:
712
self.prompt_lookup_min = 0
713
714
if self.model is not None:
715
+ # Old-format Medusa checkpoints (e.g. FasterDecoding/medusa-*)
716
+ # lack a model_type key in config.json, so AutoConfig cannot
717
+ # detect them. When the method is explicitly "medusa", inject
718
+ # model_type so MedusaConfig.from_pretrained is used instead.
719
+ draft_hf_overrides: HfOverrides
720
+ if self.method == "medusa":
721
+ draft_hf_overrides = {"model_type": "medusa"}
722
+ else:
723
+ draft_hf_overrides = SpeculativeConfig.hf_config_override
724
self.draft_model_config = ModelConfig(
725
model=self.model,
726
runner="draft",
@@ -730,10 +739,21 @@ class SpeculativeConfig:
739
quantization=self.quantization,
740
enforce_eager=self.target_model_config.enforce_eager,
741
max_logprobs=self.target_model_config.max_logprobs,
733
- hf_overrides=SpeculativeConfig.hf_config_override,
742
+ hf_overrides=draft_hf_overrides,
743
config_format=self.target_model_config.config_format,
744
)
745
746
+ # Old-format Medusa checkpoints (e.g. FasterDecoding/medusa-*)
747
+ # omit vocab_size in config.json, so MedusaConfig falls back to
748
+ # its default (32001). Align with the target model's vocab size
749
+ # to avoid shape mismatches when loading LM-head weights.
750
+ if self.method == "medusa":
751
+ target_vocab = self.target_model_config.hf_config.vocab_size
752
+ draft_hf = self.draft_model_config.hf_config
753
+ if draft_hf.vocab_size != target_vocab:
754
+ draft_hf.vocab_size = target_vocab
755
+ draft_hf.truncated_vocab_size = target_vocab
756
+
757
# Automatically detect the method
758
if self.method in ("eagle", "eagle3", "dflash"):
759
pass
vllm/model_executor/models/medusa.py
+21
@@ -136,6 +136,26 @@ class Medusa(nn.Module):
136
137
return logits_lst
138
139
+ @staticmethod
140
+ def _remap_old_checkpoint_key(name: str) -> str:
141
+ """Map old FasterDecoding checkpoint keys to vLLM parameter names.
142
+
143
+ Old format uses bare numeric prefixes:
144
+ '{head}.{layer}.linear.weight' -> 'blocks.{head}.layers.{layer}.weight'
145
+ '{head}.{layer}.linear.bias' -> 'blocks.{head}.layers.{layer}.bias'
146
+ '{head}.{N}.weight' -> 'lm_heads.{head}.weight'
147
+ """
148
+ parts = name.split(".")
149
+ if len(parts) >= 3 and parts[0].isdigit() and parts[1].isdigit():
150
+ head, layer = parts[0], parts[1]
151
+ rest = parts[2:]
152
+ if "linear" in rest:
153
+ param = ".".join(rest[rest.index("linear") + 1 :])
154
+ return f"blocks.{head}.layers.{layer}.{param}"
155
+ else:
156
+ return f"lm_heads.{head}.{'.'.join(rest)}"
157
+ return name
158
+
159
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
160
params_dict = dict(self.named_parameters())
161
loaded_params: set[str] = set()
@@ -144,6 +164,7 @@ class Medusa(nn.Module):
164
165
for name, loaded_weight in weights:
166
name = name.replace("medusa_heads.", "")
167
+ name = self._remap_old_checkpoint_key(name)
168
169
if name == "token_map":
170
if self.truncated_vocab_size < self.orig_vocab_size:
vllm/transformers_utils/config.py
+1
-1
@@ -126,7 +126,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
126
**{"unlimited-ocr": "UnlimitedOCRConfig"},
127
)
128
129
-_SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators"}
129
+_SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators", "medusa"}
130
131
_PATCH_HF_VALIDATE_ROPE: set[str] = {"sarvam_mla"}
132