Add MTP speculative decoding via MTPCandidateGenerator#45618
Add MTP speculative decoding via MTPCandidateGenerator#45618ArthurZucker wants to merge 2 commits into
Conversation
Wires MTP speculative decoding into `generate()` for DeepSeek-V3 and GLM-4 MoE
checkpoints that ship MTP modules (DeepSeek-V3 at `model.layers.61`,
GLM-4 MoE at `model.layers.46`/`.92` β previously hidden by
`_keys_to_ignore_on_load_unexpected`).
**Model side**
- New `num_nextn_predict_layers: int = 0` on `DeepseekV3Config` /
`Glm4MoeConfig` (propagates to downstream variants). Default keeps the
existing no-op behavior.
- `DeepseekV3MTPLayer` / `Glm4MoeMTPLayer` modules mirror the DeepSeek-V3
spec as implemented in vLLM: `enorm` + `hnorm` RMSNorms β concat β linear
`eh_proj(2H β H)` β a full decoder block β `shared_head (norm + lm_head)`.
- `DeepseekV3Model` / `Glm4MoeModel` extend `self.layers` past
`num_hidden_layers` with MTP modules; the base `forward` still iterates
only `self.layers[: num_hidden_layers]`. MTP is reached exclusively via a
new `model.forward_mtp(input_ids, previous_hidden_state, past_key_values,
position_ids, mtp_depth)` helper (lazily extends the KV cache for MTP
layer indices).
**Generation side**
- `GenerationConfig.use_mtp: bool = False` and a new
`GenerationMode.MTP_DECODING` routed from `get_generation_mode` whenever
the base mode is greedy or sample.
- `_mtp_decoding` in `generation/utils.py`: main forward β sample `x_{t+1}`
β chain K MTP depths for draft tokens β single verify forward β reuses
`_speculative_sampling` for accept/reject β `past_key_values.crop`.
Batch size 1, dynamic cache; leaves `_assisted_decoding` untouched.
- `ContinuousBatchingManager` refuses `use_mtp=True` for now β
paged-attention slot reservation + per-request accept/reject is tracked
separately and will come as a follow-up.
**Tests**
- `tests/generation/test_mtp.py` covers: mode dispatch, greedy
token-for-token parity vs plain `_sample` for K=1/2/3 on both models,
`num_nextn_predict_layers=0` rejection, layer extension, base-forward
equivalence when MTP layers are added, `forward_mtp` shapes, and the
`generate_batch` `NotImplementedError`.
All 9 MTP tests pass locally. `make style` clean. `make fix-repo` clean
apart from the pre-existing `mlinter._using_rule_specs` env mismatch in
`check_modeling_rules_doc.py` / `check_modeling_structure.py` that also
fails on an unmodified checkout.
- MTP modules no longer live on DeepseekV3Model / Glm4MoeModel (no more layer-list extension or forward_mtp method); configs still expose num_nextn_predict_layers as metadata. - New transformers.generation.candidate_generators.MTPCandidateGenerator (nn.Module, implements CandidateGenerator) owns the MTP layers and introspects the base model's decoder + RMSNorm classes to build them. from_pretrained pulls MTP-specific keys out of the main checkpoint. - _mtp_decoding: use self.model(...) -> last_hidden_state + self.lm_head instead of forcing output_hidden_states=True. - Tests updated for the new architecture; all 9 MTP tests + 55 tests in the generation suite pass.
|
[For maintainers] Suggested jobs to run (before merge) run-slow: deepseek_v3, glm4_moe, glm4v_moe, solar_open, youtu |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
Hi @ArthurZucker, really excited about this PR! MTP speculative decoding is a huge win for inference throughput. I noticed the branch currently has merge conflicts β would it help if I (or someone from the community) rebased it against latest main? Also, for awareness, I'm working on MTP training support for Qwen3.5 in #45638, which could complement this nicely if we align on the config naming (num_nextn_predict_layers vs mtp_num_hidden_layers). No pressure, just offering help if it speeds things up! π |
|
This was taken over in 46229 i |
Adds
use_mtp=Truetogeneratefor DeepSeek-V3 / GLM-4 MoE. MTP modules live in a newgeneration.candidate_generators.MTPCandidateGenerator(loaded via itsfrom_pretrained) β the base model stays clean.