Skip to content

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

Merged
WoosukKwon merged 3 commits into
mainfrom
woosuk/remove-medusa-mlp-speculator
Oct 15, 2026
Merged

WoosukKwon merged 3 commits into
mainfrom
woosuk/remove-medusa-mlp-speculator

Conversation

@WoosukKwon

@WoosukKwon WoosukKwon commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

Overview

Removes Medusa and MLP speculator support and simplifies the code paths, configs, tests, docs, and CI rules that existed only for them. Net −858 lines, no new code.

Claims

  • method="medusa" and method="mlp_speculator" are removed. They are now rejected by SpeculativeConfig validation, which lists the valid methods.
  • The MLP speculator was already dead code: its registry entry was commented out ("temporarily disabled") until V1 support, which never landed. Medusa is not supported by Model Runner V2, and EAGLE/MTP/DFlash/DSpark cover its use case.
  • Other spec-decode methods are unchanged.

Validation

  • pre-commit (ruff, mypy, Buildkite/autolabel rule checks) passes.
  • pytest tests/models/test_registry.py tests/config/test_model_arch_config.py tests/config/test_speculative_draft_hf_overrides.py tests/v1/worker/test_gpu_warmup_blocks.py: 504 passed, 11 skipped, 2 failed. The 2 failures (test_registry_imports[Dots3Note*]) are unrelated: my local DeepGEMM build is older than the tag bump in [Build] Update deep_gemm git tag for locality_domain update #61491 (cannot import name 'get_block_m_for_nvfp4_mega_moe').
  • E2E smoke on the V1 model runner (VLLM_USE_V2_MODEL_RUNNER=0, Llama-3.1-8B-Instruct, GB200, 16 prompts, greedy, 128 tokens). Results are identical before and after:
Method drafts accepted mean accept len
ngram (k=3) 151 192 2.27
eagle3 (k=3, yuhuili/EAGLE3-LLaMA3.1-Instruct-8B) 1480 568 1.38

No model eval is needed: no surviving method's output changes.

Details

Removed:

  • vllm/model_executor/models/{medusa,mlp_speculator}.py, vllm/transformers_utils/configs/{medusa,mlp_speculator}.py, vllm/v1/spec_decode/medusa.py, and their registry and config-map entries.
  • In SpeculativeConfig:
    • the medusa/mlp_speculator method literals and their auto-detection;
    • the old-format Medusa checkpoint workarounds (model_type hf_override injection and vocab-size patching);
    • the mlp_speculator draft-TP=1 override, including the now-unused draft_hf_config parameter of _verify_and_get_draft_tp;
    • the hf_config.num_lookahead_tokens sync (only Medusa/MLP speculator configs defined that attribute).
  • The MedusaProposer branches in the V1 gpu_model_runner.py.
  • The T/T_co generic on VllmModel/VllmModelForTextGeneration/VllmModelForPooling. It existed only for Medusa's list[torch.Tensor] hidden states; the default was already torch.Tensor, so typing for every other model is unchanged.
  • MedusaModelArchConfigConvertor, plus the check_head_size escape hatch in tests/config/test_model_arch_config.py, which only existed for Medusa's zero head size.
  • The Medusa e2e acceptance test, MLP docs page, MRV2-unsupported-method entries, Buildkite source deps, mergify/autolabel rules, and docs table rows.

Breaking change: users still passing method="medusa" or a Medusa/MLP-speculator draft checkpoint will get a config validation error.

Duplicate check: no open PR removes these. Draft #61357 migrates Medusa to the flat model layout; this PR would make it unnecessary.


Pull Request Checklist
  • I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

  • AI assistance was used during the creation of this PR.

  • Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

  • Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

  • Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

  • Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

Medusa and MLP speculator are superseded by EAGLE/MTP/DFlash/DSpark.
The MLP speculator was already unregistered ("temporarily disabled")
and Medusa is unsupported by Model Runner V2.

Removes the models, HF configs, V1 proposer, spec method literals,
auto-detection, Medusa checkpoint hacks, mlp_speculator draft-TP
override, the num_lookahead_tokens sync, the Medusa-only generic on
the VllmModel protocols, tests, docs, and CI/mergify references.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

Tip: disable this comment in your organization's Code Review settings.

@mergify

mergify Bot commented Oct 14, 2026

Copy link
Copy Markdown
Contributor

Documentation preview: https://vllm--61701.org.readthedocs.build/en/61701/

@mergify mergify Bot added documentation Improvements or additions to documentation ci/build speculative-decoding labels Oct 14, 2026
@WoosukKwon

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #94942 for commit f4fa28231af9.

@WoosukKwon

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #94958 for commit da3b7e75cc40.

@WoosukKwon
WoosukKwon enabled auto-merge (squash) October 14, 2026 22:35
@github-actions github-actions Bot added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 14, 2026
@WoosukKwon
WoosukKwon disabled auto-merge October 15, 2026 00:29
@WoosukKwon
WoosukKwon merged commit 8530aed into main Oct 15, 2026
265 of 267 checks passed
@WoosukKwon
WoosukKwon deleted the woosuk/remove-medusa-mlp-speculator branch October 15, 2026 00:29
mjbrn added a commit to brntech/vllm that referenced this pull request Oct 15, 2026
…lative-decoding README method list

One conflict, docs/features/speculative_decoding/README.md: upstream vllm-project#61701 removed
the MLP speculator line next to the Uno entry. Keep Uno, drop the MLP line.

Signed-off-by: MJ <121681560+mjbrn@users.noreply.github.com>
mjbrn added a commit to brntech/vllm that referenced this pull request Oct 15, 2026
Upstream vllm-project#61701 removed the Medusa and MLP speculator branches from the
method lists and the speculator paths. Uno's factory branch is the trailing
`elif` of init_speculator, and no test built a Uno config through it, so a
broken import or an earlier branch claiming method='uno' would surface only
at engine start on a GPU. The new test builds a real Uno SpeculativeConfig
and checks init_speculator hands it to UnoSpeculator.

Assisted-by: Claude <noreply@anthropic.com>
Signed-off-by: MJ <121681560+mjbrn@users.noreply.github.com>
michael-gpu pushed a commit to practical-inference/vllm-ampere-extended that referenced this pull request Oct 20, 2026
…61701)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
zzzzwwjj pushed a commit to vllm-project/vllm-ascend that referenced this pull request Oct 23, 2026
### What this PR does / why we need it?

Port `vllm-ascend` to vLLM main `c0b6f74` (bumped from `a98247a`) while
keeping the
`v0.31.0` release lane green. The upstream window moved the KV-cache
planning
helpers into a new module, made EAGLE draft groups explicit, typed
DeepSeek-V4
configs through upstream Transformers, removed Medusa/MLP speculator
support, and
changed several spec-decode/KV-cache/EPLB signatures. Every gate below
is either a
`vllm_version_is("0.31.0")` split or an attribute/argument probe where
two main dev
builds share one version string.

| Files | Upstream vLLM change (why it breaks vllm-ascend) | vllm-ascend
adaptation |

|-------|--------------------------------------------------|------------------------|
| `.github/vllm-main-verified.commit` | — (main-lane pin only) |
Verified main commit bumped `a98247a` -> `c0b6f74` |
| `tests/e2e/pull_request/four_card/test_minimax_m3.py` | The main-lane
pin re-JITs the static kernels, so the first generate on a cold cache
shifts the near-uniform dummy logprobs (~3e-3) enough to flip a greedy
token | On main run a throwaway `generate_greedy_logprobs` before the
baseline/replay pair so both measure the warmed steady state and share
the cached prefix; v0.31.0 keeps the single-shot comparison |
| `tests/ut/core/test_circular_buffer.py` |
[vllm#61126](vllm-project/vllm#61126)
(cascade-attention removal) moved the prefix-block query up to
`KVCacheManager` and dropped it from single-type managers | On main
assert only the manager's own `get_num_skipped_tokens(65536)`; keep the
combined `get_num_common_prefix_blocks(...)` assert on v0.31.0 |
| `tests/ut/core/test_kv_cache_interface.py` |
[vllm#61183](vllm-project/vllm#61183) moved
`is_kv_cache_spec_uniform` into `vllm.v1.core.kv_cache_planning` |
Version-aware import (`kv_cache_utils` on v0.31.0 / `kv_cache_planning`
on main) |
| `tests/ut/distributed/ascend_store/_mock_deps.py` |
[vllm#61183](vllm-project/vllm#61183) added the
`vllm.v1.core.kv_cache_planning` module | Add it to the mocked-vLLM
module list |
| `tests/ut/distributed/ascend_store/test_request_lifecycle.py` |
[vllm#58668](vllm-project/vllm#58668) trims
`num_computed_tokens` to the tokens actually kept when a request stops
mid token-list | On main subtract `(accepted + 1) - <newly appended
output tokens>` from the expected computed count; keep the old
expectation on v0.31.0 |
| `tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py` |
[vllm#61183](vllm-project/vllm#61183) moved
`get_kv_cache_config_from_groups` / `is_kv_cache_spec_uniform` |
Version-aware import |
| `tests/ut/models/test_deepseek_v41_cache.py` |
[vllm#61183](vllm-project/vllm#61183) relocated
the planning helpers;
[vllm#53558](vllm-project/vllm#53558) gave
`_max_memory_usage_bytes_from_groups` a leading `builder` argument |
Import from the live planning module and probe the upstream signature to
call the 2-arg (v0.31.0) vs 3-arg (main) form, feeding an Ascend
`_get_pool_bytes_per_block` builder |
| `tests/ut/ops/test_gdn_attn_builder.py` |
[vllm#58737](vllm-project/vllm#58737) derives
`num_reqs` from the block table and asserts its row count equals
`seq_lens`; Ascend's `unpadded()` keeps the graph-padded rows | New UT:
a padded 4-row block table with 2 real requests must still yield a 2-row
`cache_indices` |
| `tests/ut/patch/platform/test_packed_kv_cache_compat.py` |
[vllm#61183](vllm-project/vllm#61183) relocated
the patched module;
[vllm#61455](vllm-project/vllm#61455) removed
`_annotate_eagle_groups` from main | Patch
`cache_patch._kv_cache_planning`; new UT asserts the packed hook
tolerates an upstream module without the EAGLE-annotation API |
| `tests/ut/patch/platform/test_patch_eplb.py` |
[vllm#57549](vllm-project/vllm#57549) emits the
async EPLB cycle-committed log from the async worker wrapper on main |
Expect the workspace-move log only on v0.31.0 |
| `tests/ut/patch/platform/test_patch_kv_cache_coordinator.py` |
[vllm#61183](vllm-project/vllm#61183) relocated
the patch entry point | Assert against `kcu_patch._kv_cache_planning`
instead of `vllm.v1.core.kv_cache_utils` |
| `tests/ut/patch/platform/test_patch_pp_mtp.py` |
[vllm#56322](vllm-project/vllm#56322) makes
`ModelConfig.verify_with_parallel_config` read
`parallel_config.enable_hc_sp` | Add `enable_hc_sp=False` to the
constructed parallel config (ignored on v0.31.0) |
| `tests/ut/patch/platform/test_prefix_cache_cp_patches.py` |
[vllm#60043](vllm-project/vllm#60043) removed
the `mamba_cache_mode="all"` leftover and requires `"align"` with prefix
caching; [vllm#53558](vllm-project/vllm#53558)
added the `builder` argument | Version-select the mamba cache mode;
probe the upstream admission signature for the `builder` argument |
| `tests/ut/spec_decode/test_mla_dspark.py` |
[vllm#60692](vllm-project/vllm#60692) renamed
the capture hook `precompute_context_kv(num_reqs)` ->
`prepare_forward(num_reqs, attn_metadata)` | Version-split the expected
capture args/kwargs |
| `tests/ut/test_compressed_prefix_cache.py` |
[vllm#61183](vllm-project/vllm#61183) moved
`is_kv_cache_spec_uniform` | Version-aware import |
| `tests/ut/worker/test_attn_utils_v2.py` |
[vllm#61455](vllm-project/vllm#61455) annotates
EAGLE groups from `spec.is_draft`;
[vllm#60766](vllm-project/vllm#60766) resolves
pin memory lazily via `is_pin_memory_available()` | New UT that
draft-layer specs get tagged; on main monkeypatch
`is_pin_memory_available` to `False` |
| `tests/ut/worker/test_model_runner_v2.py` |
[vllm#56716](vllm-project/vllm#56716) suppresses
EPLB even for `skip_eplb` dummy runs;
[vllm#60946](vllm-project/vllm#60946) renamed
the PCP staging hook | Add a `runner.eplb` stub; call the
version-appropriate staging method |
| `tests/ut/worker/test_mtp_pcp_speculator_v2.py` |
[vllm#60885](vllm-project/vllm#60885) rejects a
mismatched explicit float16/bfloat16 KV-cache dtype | Build the model
config with `bfloat16` instead of `float32` |
| `tests/ut/worker/test_pcp_manager_v2.py` |
[vllm#60946](vllm-project/vllm#60946) renamed
`prepare_inputs_to_capture` -> `prepare_inputs_for_capture_or_profile` |
Patch and call the live method name per lane |
| `tests/ut/worker/test_v2_dummy_padding.py` |
[vllm#61767](vllm-project/vllm#61767) collects
DP votes through `get_dp_group().cpu_all_reduce_()` instead of
`dist.all_reduce(..., group=cpu_group)` | Version-split the mocked
collective |
| `tests/ut/worker/v2/test_eplb_controller.py` |
[vllm#57549](vllm-project/vllm#57549) gives
`add_model` an optional `model_name` and names the draft entry | Expect
the 3-arg call (`model_name="draft (draft)"`) on main, 2-arg on v0.31.0
|
| `vllm_ascend/_310p/worker_310p.py` |
[vllm#61183](vllm-project/vllm#61183) moved
`get_kv_cache_groups` | Version-aware import |
| `vllm_ascend/distributed/ec_transfer/ec_memcache_connector.py` |
[vllm#57904](vllm-project/vllm#57904) dropped
`_get_encoder_cache_hidden_dim` from `cpu.common` (encoder widths are
measured at startup) | On main size the receive buffer as
`out_hidden_size * (1 + len(deepstack_visual_indexes))` (Qwen3-VL
DeepStack) else `get_inputs_embeds_size()`; keep the helper on v0.31.0 |
| `vllm_ascend/models/deepseek_v4/compressor.py` |
[vllm#61392](vllm-project/vllm#61392) switched
`DeepseekV4Config` to the upstream Transformers class | Import from
`transformers` on main, from
`vllm.transformers_utils.configs.deepseek_v4` on v0.31.0 |
| `vllm_ascend/models/deepseek_v4/indexer.py` | Same
[vllm#61392](vllm-project/vllm#61392) | Same
version-aware import |
| `vllm_ascend/models/deepseek_v4/mm_preprocess.py` | Same
[vllm#61392](vllm-project/vllm#61392) | Same
version-aware import |
| `vllm_ascend/models/deepseek_v4/model.py` |
[vllm#61392](vllm-project/vllm#61392) moves the
config to Transformers, folds `compress_ratios`/`num_hash_layers` into
`layer_types`/`mlp_layer_types`, and keys YaRN under the compressor
(`compress`) group | Restore the legacy fields on main from the new
lists, read the `compress`-scoped rope params with safe defaults, and
keep the legacy path on v0.31.0 |
| `vllm_ascend/models/deepseek_v41/cache_config.py` |
[vllm#61183](vllm-project/vllm#61183) moved
`may_override_num_blocks` | Version-aware import |
| `vllm_ascend/models/glm5next/cache_config.py` |
[vllm#61183](vllm-project/vllm#61183) moved
`create_kv_cache_group_specs` / `may_override_num_blocks` |
Version-aware import |
| `vllm_ascend/ops/gdn_attn_builder.py` |
[vllm#58737](vllm-project/vllm#58737) derives
`num_reqs` from the block table and asserts row parity with `seq_lens` |
On main trim the padded block table to `seq_lens.shape[0]` before
`mamba_get_block_table_tensor` when the sizes differ |
| `vllm_ascend/patch/__init__.py` |
[vllm#61183](vllm-project/vllm#61183) moved
`get_kv_cache_config_from_groups` to `kv_cache_planning` | Update the
patch-registry doc |
| `vllm_ascend/patch/platform/__init__.py` |
[vllm#61183](vllm-project/vllm#61183) moved the
KV-cache planning/patcher module | Update the patch-registry doc |
| `vllm_ascend/patch/platform/patch_eplb.py` |
[vllm#57549](vllm-project/vllm#57549) moves the
async EPLB cycle-committed log onto the async worker path on main |
Report completion from the wrapped async worker (rank 0) on main; keep
the workspace-move log on v0.31.0 |
| `vllm_ascend/patch/platform/patch_kv_cache_utils.py` |
[vllm#61183](vllm-project/vllm#61183) moved the
planning helpers;
[vllm#53558](vllm-project/vllm#53558) added the
`builder` argument;
[vllm#61455](vllm-project/vllm#61455) dropped
`_annotate_eagle_groups` | Resolve the live planning module by version,
patch its symbols, gate the EAGLE annotation on attribute presence, and
dispatch the admission helper by argument count |
| `vllm_ascend/quantization/methods/kv_cache/turboquant/cache.py` |
[vllm#61183](vllm-project/vllm#61183) moved
`may_override_num_blocks` | Version-aware import |
| `vllm_ascend/spec_decode/__init__.py` |
[vllm#61701](vllm-project/vllm#61701) removed
Medusa/MLP speculator support | Import `AscendMedusaProposer` only on
v0.31.0; raise a clear error for `method="medusa"` on main |
| `vllm_ascend/spec_decode/medusa_proposer.py` |
[vllm#61701](vllm-project/vllm#61701) removed
`vllm.v1.spec_decode.medusa` | Add `# type: ignore[import-not-found]`;
the module only loads on the lane that still ships it |
| `vllm_ascend/worker/v1/model_runner.py` | Same
[vllm#61701](vllm-project/vllm#61701) | Gate the
Medusa import and the `isinstance(self.drafter, AscendMedusaProposer)`
branch on v0.31.0 |
| `vllm_ascend/worker/v2/attn_utils.py` |
[vllm#61455](vllm-project/vllm#61455) annotates
EAGLE groups from `spec.is_draft` | On main tag the spec with `is_draft`
from `attn_module.is_draft_layer` |
| `vllm_ascend/worker/v2/pcp_manager.py` |
[vllm#60946](vllm-project/vllm#60946) added the
`moe_gathers_tokens` argument and renamed `prepare_inputs_to_capture` ->
`prepare_inputs_for_capture_or_profile` | Forward the new argument;
version-split the staging hook |
| `vllm_ascend/worker/v2/sample/apply_top_k_top_p.py` |
[vllm#61996](vllm-project/vllm#61996) added
`has_topp_only_rows` to the split top-p pipeline | Accept and ignore the
extra argument on main (the PyTorch fallback handles top-p-only rows) |
| `vllm_ascend/worker/v2/spec_decode/dflash/aclgraph.py` |
[vllm#60692](vllm-project/vllm#60692) renamed
`precompute_context_kv` -> `prepare_forward` | Pass the capture hook the
installed vLLM expects |
| `vllm_ascend/worker/v2/spec_decode/lmhead_tp_utils.py` |
[vllm#59554](vllm-project/vllm#59554) added
`num_steps` to `Speculator.sample_draft` | Version-split the override
signature and forward `num_steps` on main |
| `vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py` |
[vllm#61451](vllm-project/vllm#61451) renamed
the logits-stats kernel argument `num_speculative_steps` -> `num_logits`
and moved bonus-token detection into `_is_bonus_logit` | Pass
`num_logits` on main; keep `num_speculative_steps` on v0.31.0 |
| `vllm_ascend/worker/worker.py` |
[vllm#61183](vllm-project/vllm#61183) moved
`get_kv_cache_groups`/`unify_hybrid_kv_cache_specs`;
[vllm#57516](vllm-project/vllm#57516) reclaims
cached memory via `torch.accelerator.empty_cache()` inside `load_model`
| Version-aware imports; on main neutralize `torch.npu.empty_cache`
while loading weights under the sleep-mode pluggable memory pool |

### Does this PR introduce _any_ user-facing change?

No. This is an internal main2main compatibility update; on the `v0.31.0`
release lane
behavior is unchanged, and on main the only behavioral note is that
Medusa
speculative decoding now raises a clear error (upstream removed it in
vllm#61701).

### How was this patch tested?

- Unit tests: the changed UT files above (KV-cache planning/coordinator,
EPLB,
PCP/spec-decode, DeepSeek-V4/V4.1 cache, mamba align,
`worker/worker.py`).
- E2E: `tests/e2e/pull_request/four_card/test_minimax_m3.py`.
- Run the main lane against the pinned `c0b6f74` and the `v0.31.0` lane
via
  `vllm_version_is`.

---

### Checklist

#### Testing

- [ ] **Boundary / edge-case coverage**: I reviewed whether this change
involves
boundary conditions (e.g. empty input, single/zero element, min/max
values,
off-by-one, sequence-length limits, integer overflow). Where applicable,
I
added unit tests covering these edge cases in `tests/ut/` or
`tests/e2e/`.
- [ ] **Regression test**: If this PR fixes a bug, I added a regression
test
  that would fail without the fix.
- [ ] **Error / failure path coverage**: I covered abnormal paths
(exceptions,
timeouts, out-of-range, division by zero) in addition to the happy path.
- [ ] **Meaningful & reliable tests**: New tests assert real behavior
and are
  not written merely to satisfy CI.
- [ ] **Performance impact**: If this PR touches a performance-sensitive
path,
  I attached benchmark results and added performance regression tests.


- vLLM main:
vllm-project/vllm@a98247a

Signed-off-by: zhangxinyuehfad <starmoon_zhang@163.com>
Morrowmake pushed a commit to Morrowmake/vllm-cmp170hx that referenced this pull request Oct 26, 2026
… meta test

Upstream dropped its draft_hf_config parameter with the Medusa and MLP
speculator removal (vllm-project#61701); the GLM wrapper meta cases passed it and failed
with a TypeError before building anything.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci/build documentation Improvements or additions to documentation ready ONLY add when PR is ready to merge/full CI is needed speculative-decoding

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants