A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 12:15 UTC
MiniMax-M2: add Eagle3 speculative decoding support (#37512)
Signed-off-by: liuchenbing <chenliumail@163.com> Signed-off-by: liucb <liuchengbao_work@163.com> Co-authored-by: liuchenbing <chenliumail@163.com>
liuchenbing2026 committed
Apr 6, 2026 at 10:50 UTC
f6983f01de2bf2e92ab468fa735ebac39cddd670
4 files changed
+24
-5
tests/models/registry.py
+6
@@ -1246,6 +1246,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1246
use_original_num_layers=True,
1247
max_model_len=10240,
1248
),
1249
+ "Eagle3MiniMaxM2ForCausalLM": _HfExamplesInfo(
1250
+ "MiniMaxAI/MiniMax-M2",
1251
+ trust_remote_code=True,
1252
+ speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
1253
+ tokenizer="MiniMaxAI/MiniMax-M2",
1254
+ ),
1255
"EagleMistralLarge3ForCausalLM": _HfExamplesInfo(
1256
"mistralai/Mistral-Large-3-675B-Instruct-2512",
1257
speculative_model="mistralai/Mistral-Large-3-675B-Instruct-2512-Eagle",
vllm/config/speculative.py
+1
@@ -817,6 +817,7 @@ class SpeculativeConfig:
817
"deepseek_v3",
818
"kimi_k2",
819
"kimi_k25",
820
+ "minimax_m2",
821
]
822
if (
823
self.method in ("eagle3", "extract_hidden_states", "dflash")
vllm/model_executor/models/minimax_m2.py
+16
-5
@@ -24,6 +24,7 @@
24
"""Inference-only MiniMaxM2 model."""
25
26
from collections.abc import Iterable
27
+from itertools import islice
28
from typing import Any
29
30
import torch
@@ -59,7 +60,7 @@ from vllm.model_executor.model_loader.weight_utils import (
60
)
61
from vllm.sequence import IntermediateTensors
62
62
-from .interfaces import SupportsLoRA, SupportsPP
63
+from .interfaces import EagleModelMixin, SupportsEagle3, SupportsLoRA, SupportsPP
64
from .utils import (
65
AutoWeightsLoader,
66
PPMissingLayer,
@@ -313,7 +314,7 @@ class MiniMaxM2DecoderLayer(nn.Module):
314
315
316
@support_torch_compile
316
-class MiniMaxM2Model(nn.Module):
317
+class MiniMaxM2Model(nn.Module, EagleModelMixin):
318
fall_back_to_pt_during_load = False
319
320
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -366,7 +367,7 @@ class MiniMaxM2Model(nn.Module):
367
positions: torch.Tensor,
368
intermediate_tensors: IntermediateTensors | None,
369
inputs_embeds: torch.Tensor | None = None,
369
- ) -> torch.Tensor | IntermediateTensors:
370
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
371
if get_pp_group().is_first_rank:
372
if inputs_embeds is not None:
373
hidden_states = inputs_embeds
@@ -378,14 +379,24 @@ class MiniMaxM2Model(nn.Module):
379
hidden_states = intermediate_tensors["hidden_states"]
380
residual = intermediate_tensors["residual"]
381
381
- for layer in self.layers[self.start_layer : self.end_layer]:
382
+ aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
383
+ for idx, layer in enumerate(
384
+ islice(self.layers, self.start_layer, self.end_layer)
385
+ ):
386
hidden_states, residual = layer(positions, hidden_states, residual)
387
+ self._maybe_add_hidden_state(
388
+ aux_hidden_states, idx + 1, hidden_states, residual
389
+ )
390
391
if not get_pp_group().is_last_rank:
392
return IntermediateTensors(
393
{"hidden_states": hidden_states, "residual": residual}
394
)
395
hidden_states, _ = self.norm(hidden_states, residual)
396
+
397
+ if len(aux_hidden_states) > 0:
398
+ return hidden_states, aux_hidden_states
399
+
400
return hidden_states
401
402
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
@@ -496,7 +507,7 @@ class MiniMaxM2Model(nn.Module):
507
return loaded_params
508
509
499
-class MiniMaxM2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
510
+class MiniMaxM2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
511
packed_modules_mapping = {
512
"qkv_proj": [
513
"q_proj",
vllm/model_executor/models/registry.py
+1
@@ -554,6 +554,7 @@ _SPECULATIVE_DECODING_MODELS = {
554
"EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
555
"DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
556
"Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
557
+ "Eagle3MiniMaxM2ForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
558
"LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
559
"Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
560
"Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),