@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / acee8f48aa
Fork 0

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 15, 11:00 UTC

[Model] Support MiMo-7B inference with MTP (#17433)

Signed-off-by: wp-alpha <wangpeng66@xiaomi.com> Co-authored-by: wangpeng66 <wangpeng66@xiaomi.com>

bwshen-mi committed May 13, 2025 at 07:25 UTC acee8f48aa9c580d0a54d00117ab772b6fe68650
7 files changed +507 -4
docs/source/models/supported_models.md
+5
@@ -592,6 +592,11 @@ Specified using `--task generate`.
592 * `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc.
593 *
594 *
595 +- * `MiMoForCausalLM`
596 + * MiMo
597 + * `XiaomiMiMo/MiMo-7B-RL`, etc.
598 + *
599 + *
600 :::
601
602 :::{note}
tests/models/registry.py
+5
@@ -242,6 +242,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
242 is_available_online=False,
243 trust_remote_code=True),
244 "Zamba2ForCausalLM": _HfExamplesInfo("Zyphra/Zamba2-7B-instruct"),
245 + "MiMoForCausalLM": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL",
246 + trust_remote_code=True),
247 # [Encoder-decoder]
248 "BartModel": _HfExamplesInfo("facebook/bart-base"),
249 "BartForConditionalGeneration": _HfExamplesInfo("facebook/bart-large-cnn"),
@@ -403,6 +405,9 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
405 trust_remote_code=True,
406 speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
407 tokenizer="meta-llama/Llama-3.1-8B-Instruct"),
408 + "MiMoMTPModel": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL",
409 + trust_remote_code=True,
410 + speculative_model="XiaomiMiMo/MiMo-7B-RL")
411 }
412
413 _TRANSFORMERS_MODELS = {
vllm/config.py
+17 -3
@@ -1139,7 +1139,8 @@ class ModelConfig:
1139 def get_layers_start_end_indices(
1140 self, parallel_config: "ParallelConfig") -> tuple[int, int]:
1141 from vllm.distributed.utils import get_pp_indices
1142 - if self.hf_text_config.model_type == "deepseek_mtp":
1142 + if (self.hf_text_config.model_type == "deepseek_mtp"
1143 + or self.hf_config.model_type == "mimo_mtp"):
1144 total_num_hidden_layers = getattr(self.hf_text_config,
1145 "num_nextn_predict_layers", 0)
1146 else:
@@ -2357,6 +2358,17 @@ class SpeculativeConfig:
2358 "n_predict": n_predict,
2359 "architectures": ["DeepSeekMTPModel"]
2360 })
2361 +
2362 + if hf_config.architectures[0] == "MiMoForCausalLM":
2363 + hf_config.model_type = "mimo_mtp"
2364 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
2365 + hf_config.update({
2366 + "num_hidden_layers": 0,
2367 + "n_predict": n_predict,
2368 + "architectures": ["MiMoMTPModel"]
2369 + })
2370 + return hf_config
2371 +
2372 return hf_config
2373
2374 def __post_init__(self):
@@ -2373,8 +2385,10 @@ class SpeculativeConfig:
2385 # TODO(Shangming): Refactor mtp configuration logic when supporting
2386 # mtp acceleration for more models besides deepseek_v3
2387 if self.target_model_config and \
2376 - self.target_model_config.hf_text_config.model_type \
2377 - == "deepseek_v3":
2388 + (self.target_model_config.hf_text_config.model_type \
2389 + == "deepseek_v3" or
2390 + self.target_model_config.hf_text_config.model_type \
2391 + == "mimo"):
2392 # use the draft model from the same model:
2393 self.model = self.target_model_config.model
2394 elif self.method in ("ngram", "[ngram]"):
vllm/model_executor/models/mimo.py new
+190
@@ -0,0 +1,190 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Adapted from
4 +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py
5 +# Copyright 2025 Xiaomi Corporation.
6 +# Copyright 2024 The Qwen team.
7 +# Copyright 2023 The vLLM team.
8 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
9 +#
10 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
11 +# and OPT implementations in this library. It has been modified from its
12 +# original forms to accommodate minor architectural differences compared
13 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
14 +#
15 +# Licensed under the Apache License, Version 2.0 (the "License");
16 +# you may not use this file except in compliance with the License.
17 +# You may obtain a copy of the License at
18 +#
19 +# http://www.apache.org/licenses/LICENSE-2.0
20 +#
21 +# Unless required by applicable law or agreed to in writing, software
22 +# distributed under the License is distributed on an "AS IS" BASIS,
23 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
24 +# See the License for the specific language governing permissions and
25 +# limitations under the License.
26 +"""Inference-only MiMo model compatible with HuggingFace weights."""
27 +from typing import Iterable, Optional, Set, Tuple, Union
28 +
29 +import torch
30 +import torch.nn as nn
31 +
32 +from vllm.compilation.decorators import support_torch_compile
33 +from vllm.config import VllmConfig
34 +from vllm.distributed import get_pp_group
35 +from vllm.logger import init_logger
36 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
37 +from vllm.model_executor.layers.sampler import get_sampler
38 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
39 +from vllm.model_executor.model_loader.weight_utils import (
40 + default_weight_loader, maybe_remap_kv_scale_name)
41 +from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM, Qwen2Model
42 +from vllm.model_executor.sampling_metadata import SamplingMetadata
43 +from vllm.sequence import IntermediateTensors
44 +
45 +from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix
46 +
47 +logger = init_logger(__name__)
48 +
49 +
50 +@support_torch_compile(
51 + dynamic_arg_dims={
52 + "input_ids": 0,
53 + "positions": -1,
54 + "intermediate_tensors": 0,
55 + "inputs_embeds": 0,
56 + })
57 +class MiMoModel(Qwen2Model):
58 +
59 + def forward(
60 + self,
61 + input_ids: torch.Tensor,
62 + positions: torch.Tensor,
63 + intermediate_tensors: Optional[IntermediateTensors] = None,
64 + inputs_embeds: Optional[torch.Tensor] = None,
65 + ) -> Union[torch.Tensor, IntermediateTensors]:
66 + if get_pp_group().is_first_rank:
67 + if inputs_embeds is not None:
68 + hidden_states = inputs_embeds
69 + else:
70 + hidden_states = self.get_input_embeddings(input_ids)
71 + residual = None
72 + else:
73 + assert intermediate_tensors is not None
74 + hidden_states = intermediate_tensors["hidden_states"]
75 + residual = intermediate_tensors["residual"]
76 + for layer in self.layers[self.start_layer:self.end_layer]:
77 + hidden_states, residual = layer(
78 + positions,
79 + hidden_states,
80 + residual,
81 + )
82 + if not get_pp_group().is_last_rank:
83 + return IntermediateTensors({
84 + "hidden_states": hidden_states,
85 + "residual": residual
86 + })
87 + hidden_states = hidden_states + residual
88 + return hidden_states
89 +
90 + def load_weights(self, weights: Iterable[Tuple[str,
91 + torch.Tensor]]) -> Set[str]:
92 + stacked_params_mapping = [
93 + ("qkv_proj", "q_proj", "q"),
94 + ("qkv_proj", "k_proj", "k"),
95 + ("qkv_proj", "v_proj", "v"),
96 + ("gate_up_proj", "gate_proj", 0),
97 + ("gate_up_proj", "up_proj", 1),
98 + ]
99 + params_dict = dict(self.named_parameters(remove_duplicate=False))
100 + loaded_params: Set[str] = set()
101 + for name, loaded_weight in weights:
102 + if "mtp_layers" in name:
103 + continue
104 + if "rotary_emb.inv_freq" in name:
105 + continue
106 + if (self.quant_config is not None and
107 + (scale_name := self.quant_config.get_cache_scale(name))):
108 + # Loading kv cache quantization scales
109 + param = params_dict[scale_name]
110 + weight_loader = getattr(param, "weight_loader",
111 + default_weight_loader)
112 + loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
113 + loaded_weight[0])
114 + weight_loader(param, loaded_weight)
115 + loaded_params.add(scale_name)
116 + continue
117 + for (param_name, weight_name, shard_id) in stacked_params_mapping:
118 + if weight_name not in name:
119 + continue
120 + name = name.replace(weight_name, param_name)
121 + # Skip loading extra bias for GPTQ models.
122 + if name.endswith(".bias") and name not in params_dict:
123 + continue
124 + if is_pp_missing_parameter(name, self):
125 + continue
126 + param = params_dict[name]
127 + weight_loader = param.weight_loader
128 + weight_loader(param, loaded_weight, shard_id)
129 + break
130 + else:
131 + # Skip loading extra bias for GPTQ models.
132 + if name.endswith(".bias") and name not in params_dict:
133 + continue
134 + # Remapping the name of FP8 kv-scale.
135 + name = maybe_remap_kv_scale_name(name, params_dict)
136 + if name is None:
137 + continue
138 + if is_pp_missing_parameter(name, self):
139 + continue
140 + param = params_dict[name]
141 + weight_loader = getattr(param, "weight_loader",
142 + default_weight_loader)
143 + weight_loader(param, loaded_weight)
144 + loaded_params.add(name)
145 + return loaded_params
146 +
147 +
148 +class MiMoForCausalLM(Qwen2ForCausalLM, nn.Module):
149 +
150 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
151 + nn.Module.__init__(self)
152 + config = vllm_config.model_config.hf_config
153 + quant_config = vllm_config.quant_config
154 + lora_config = vllm_config.lora_config
155 +
156 + self.config = config
157 + self.lora_config = lora_config
158 +
159 + self.quant_config = quant_config
160 +
161 + self.model = MiMoModel(vllm_config=vllm_config,
162 + prefix=maybe_prefix(prefix, "model"))
163 +
164 + if get_pp_group().is_last_rank:
165 + if config.tie_word_embeddings:
166 + self.lm_head = self.model.embed_tokens
167 + else:
168 + self.lm_head = ParallelLMHead(config.vocab_size,
169 + config.hidden_size,
170 + quant_config=quant_config,
171 + prefix=maybe_prefix(
172 + prefix, "lm_head"))
173 + else:
174 + self.lm_head = PPMissingLayer()
175 +
176 + self.logits_processor = LogitsProcessor(config.vocab_size)
177 + self.sampler = get_sampler()
178 +
179 + self.make_empty_intermediate_tensors = (
180 + self.model.make_empty_intermediate_tensors)
181 +
182 + def compute_logits(
183 + self,
184 + hidden_states: torch.Tensor,
185 + sampling_metadata: SamplingMetadata,
186 + ) -> Optional[torch.Tensor]:
187 + hidden_states = self.model.norm(hidden_states)
188 + logits = self.logits_processor(self.lm_head, hidden_states,
189 + sampling_metadata)
190 + return logits
vllm/model_executor/models/mimo_mtp.py new
+283
@@ -0,0 +1,283 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# Adapted from
4 +# https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/deepseek_mtp.py
5 +# Copyright 2025 Xiaomi Corporation.
6 +# Copyright 2023 The vLLM team.
7 +# Copyright 2024 DeepSeek-AI team.
8 +
9 +# Licensed under the Apache License, Version 2.0 (the "License");
10 +# you may not use this file except in compliance with the License.
11 +# You may obtain a copy of the License at
12 +#
13 +# http://www.apache.org/licenses/LICENSE-2.0
14 +#
15 +# Unless required by applicable law or agreed to in writing, software
16 +# distributed under the License is distributed on an "AS IS" BASIS,
17 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18 +# See the License for the specific language governing permissions and
19 +# limitations under the License.
20 +"""Inference-only MiMo-MTP model."""
21 +from typing import Iterable, Optional, Set, Tuple
22 +
23 +import torch
24 +import torch.nn as nn
25 +from transformers import PretrainedConfig
26 +
27 +from vllm.config import CacheConfig, ModelConfig, VllmConfig
28 +from vllm.model_executor.layers.layernorm import RMSNorm
29 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
30 +from vllm.model_executor.layers.quantization import QuantizationConfig
31 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
32 +from vllm.model_executor.layers.vocab_parallel_embedding import (
33 + ParallelLMHead, VocabParallelEmbedding)
34 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
35 +from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer
36 +from vllm.model_executor.sampling_metadata import SamplingMetadata
37 +from vllm.sequence import IntermediateTensors
38 +
39 +from .utils import maybe_prefix
40 +
41 +
42 +class MiMoMultiTokenPredictorLayer(nn.Module):
43 +
44 + def __init__(
45 + self,
46 + config: PretrainedConfig,
47 + prefix: str,
48 + model_config: ModelConfig,
49 + cache_config: Optional[CacheConfig] = None,
50 + quant_config: Optional[QuantizationConfig] = None,
51 + ) -> None:
52 + super().__init__()
53 +
54 + self.token_layernorm = RMSNorm(config.hidden_size,
55 + eps=config.rms_norm_eps)
56 + self.hidden_layernorm = RMSNorm(config.hidden_size,
57 + eps=config.rms_norm_eps)
58 + self.input_proj = nn.Linear(config.hidden_size * 2,
59 + config.hidden_size,
60 + bias=False)
61 + self.mtp_block = Qwen2DecoderLayer(config=config,
62 + cache_config=cache_config,
63 + quant_config=quant_config,
64 + prefix=prefix)
65 + self.final_layernorm = RMSNorm(config.hidden_size,
66 + eps=config.rms_norm_eps)
67 +
68 + def forward(
69 + self,
70 + inputs_embeds: torch.Tensor,
71 + positions: torch.Tensor,
72 + previous_hidden_states: torch.Tensor,
73 + spec_step_index: int = 0,
74 + ) -> torch.Tensor:
75 + assert inputs_embeds is not None
76 + # masking inputs at position 0, as not needed by MTP
77 + inputs_embeds[positions == 0] = 0
78 + inputs_embeds = self.token_layernorm(inputs_embeds)
79 + previous_hidden_states = self.hidden_layernorm(previous_hidden_states)
80 +
81 + hidden_states = self.input_proj(
82 + torch.cat([previous_hidden_states, inputs_embeds], dim=-1))
83 +
84 + hidden_states, residual = self.mtp_block(positions=positions,
85 + hidden_states=hidden_states,
86 + residual=None)
87 + hidden_states = residual + hidden_states
88 + return self.final_layernorm(hidden_states)
89 +
90 +
91 +class MiMoMultiTokenPredictor(nn.Module):
92 +
93 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
94 + super().__init__()
95 +
96 + config = vllm_config.model_config.hf_config
97 + self.mtp_start_layer_idx = config.num_hidden_layers
98 + self.num_mtp_layers = config.num_nextn_predict_layers
99 +
100 + self.embed_tokens = VocabParallelEmbedding(
101 + config.vocab_size,
102 + config.hidden_size,
103 + )
104 +
105 + self.mtp_layers = torch.nn.ModuleDict({
106 + str(idx):
107 + MiMoMultiTokenPredictorLayer(
108 + config,
109 + f"{prefix}.layers.{idx}",
110 + model_config=vllm_config.model_config,
111 + cache_config=vllm_config.cache_config,
112 + quant_config=vllm_config.quant_config,
113 + )
114 + for idx in range(self.mtp_start_layer_idx,
115 + self.mtp_start_layer_idx + self.num_mtp_layers)
116 + })
117 +
118 + self.logits_processor = LogitsProcessor(config.vocab_size)
119 +
120 + def forward(
121 + self,
122 + input_ids: torch.Tensor,
123 + positions: torch.Tensor,
124 + previous_hidden_states: torch.Tensor,
125 + inputs_embeds: Optional[torch.Tensor] = None,
126 + spec_step_idx: int = 0,
127 + ) -> torch.Tensor:
128 +
129 + if inputs_embeds is None:
130 + inputs_embeds = self.embed_tokens(input_ids)
131 + return self.mtp_layers[str(self.mtp_start_layer_idx + spec_step_idx)](
132 + inputs_embeds,
133 + positions,
134 + previous_hidden_states,
135 + spec_step_idx,
136 + )
137 +
138 + def compute_logits(
139 + self,
140 + hidden_states: torch.Tensor,
141 + lm_head: ParallelLMHead,
142 + sampling_metadata: SamplingMetadata,
143 + spec_step_idx: int = 0,
144 + ) -> torch.Tensor:
145 + self.mtp_layers[str(self.mtp_start_layer_idx + spec_step_idx)]
146 + logits = self.logits_processor(lm_head, hidden_states,
147 + sampling_metadata)
148 + return logits
149 +
150 +
151 +class MiMoMTP(nn.Module):
152 +
153 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
154 + super().__init__()
155 + self.config = vllm_config.model_config.hf_config
156 + self.model = MiMoMultiTokenPredictor(vllm_config=vllm_config,
157 + prefix=maybe_prefix(
158 + prefix, "model"))
159 + self.lm_head = ParallelLMHead(self.config.vocab_size,
160 + self.config.hidden_size)
161 +
162 + self.sampler = get_sampler()
163 +
164 + def forward(
165 + self,
166 + input_ids: torch.Tensor,
167 + positions: torch.Tensor,
168 + previous_hidden_states: torch.Tensor,
169 + intermediate_tensors: Optional[IntermediateTensors] = None,
170 + inputs_embeds: Optional[torch.Tensor] = None,
171 + spec_step_idx: int = 0,
172 + ) -> torch.Tensor:
173 + assert spec_step_idx == 0, "mimo_mtp only support predict one token now"
174 + hidden_states = self.model(input_ids, positions,
175 + previous_hidden_states, inputs_embeds,
176 + spec_step_idx)
177 + return hidden_states
178 +
179 + def compute_logits(
180 + self,
181 + hidden_states: torch.Tensor,
182 + sampling_metadata: SamplingMetadata,
183 + spec_step_idx: int = 0,
184 + ) -> Optional[torch.Tensor]:
185 + return self.model.compute_logits(hidden_states, self.lm_head,
186 + sampling_metadata, spec_step_idx)
187 +
188 + def sample(
189 + self,
190 + logits: torch.Tensor,
191 + sampling_metadata: SamplingMetadata,
192 + ) -> Optional[SamplerOutput]:
193 + next_tokens = self.sampler(logits, sampling_metadata)
194 + return next_tokens
195 +
196 + def load_weights(self, weights: Iterable[Tuple[str,
197 + torch.Tensor]]) -> Set[str]:
198 + stacked_params_mapping = [
199 + ("qkv_proj", "q_proj", "q"),
200 + ("qkv_proj", "k_proj", "k"),
201 + ("qkv_proj", "v_proj", "v"),
202 + ("gate_up_proj", "gate_proj", 0),
203 + ("gate_up_proj", "up_proj", 1),
204 + ]
205 +
206 + params_dict = dict(self.named_parameters())
207 + loaded_params: Set[str] = set()
208 + for name, loaded_weight in weights:
209 +
210 + if "rotary_emb.inv_freq" in name:
211 + continue
212 + name = self.map_model_name_to_mtp_param_name(name)
213 +
214 + for (param_name, weight_name, shard_id) in stacked_params_mapping:
215 + # Skip non-stacked layers and experts (experts handled below).
216 + if weight_name not in name:
217 + continue
218 + if "mtp_layers" not in name:
219 + break
220 + # We have mlp.experts[0].gate_proj in the checkpoint.
221 + # Since we handle the experts below in expert_params_mapping,
222 + # we need to skip here BEFORE we update the name, otherwise
223 + # name will be updated to mlp.experts[0].gate_up_proj, which
224 + # will then be updated below in expert_params_mapping
225 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.
226 + if (("mlp.experts." in name) and name not in params_dict):
227 + continue
228 + name = name.replace(weight_name, param_name)
229 + # Skip loading extra bias for GPTQ models.
230 + if name.endswith(".bias") and name not in params_dict:
231 + continue
232 +
233 + param = params_dict[name]
234 + weight_loader = param.weight_loader
235 + weight_loader(param, loaded_weight, shard_id)
236 + break
237 + else:
238 + # Skip loading extra bias for GPTQ models.
239 + if name.endswith(".bias") and name not in params_dict:
240 + continue
241 + if "mtp_layers" not in name and ("embed_tokens" not in name
242 + and "lm_head" not in name):
243 + continue
244 + param = params_dict[name]
245 + weight_loader = getattr(param, "weight_loader",
246 + default_weight_loader)
247 + weight_loader(param, loaded_weight)
248 + loaded_params.add(name)
249 + return loaded_params
250 +
251 + def map_model_name_to_mtp_param_name(self, name: str) -> str:
252 + import re
253 + name_without_prefix = [
254 + "token_layernorm", "hidden_layernorm", "input_proj",
255 + "final_layernorm"
256 + ]
257 + for sub_name in name_without_prefix:
258 + if sub_name in name:
259 + return name
260 + pattern = r"model.mtp_layers.(\d+)."
261 + group = re.match(pattern, name)
262 + if group is not None:
263 + name = name.replace(group.group(), group.group() + "mtp_block.")
264 + return name
265 +
266 + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
267 + """
268 + Rewrite the weight name to match the format of the original model.
269 + Add .mtp_block for modules in transformer layer block for spec layer
270 + """
271 + spec_layer_weight_names = [
272 + "embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head"
273 + ]
274 + spec_layer_weight = False
275 + for weight_name in spec_layer_weight_names:
276 + if weight_name in name:
277 + spec_layer_weight = True
278 + break
279 + if not spec_layer_weight:
280 + # treat rest weights as weights for transformer layer block
281 + name = name.replace(f"model.layers.{spec_layer}.",
282 + f"model.layers.{spec_layer}.mtp_block.")
283 + return name
vllm/model_executor/models/registry.py
+2
@@ -88,6 +88,7 @@ _TEXT_GENERATION_MODELS = {
88 # transformers's mpt class has lower case
89 "MptForCausalLM": ("mpt", "MPTForCausalLM"),
90 "MPTForCausalLM": ("mpt", "MPTForCausalLM"),
91 + "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),
92 "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
93 "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
94 "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
@@ -215,6 +216,7 @@ _MULTIMODAL_MODELS = {
216 }
217
218 _SPECULATIVE_DECODING_MODELS = {
219 + "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
220 "EAGLEModel": ("eagle", "EAGLE"),
221 "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
222 "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
vllm/worker/worker.py
+5 -1
@@ -71,7 +71,11 @@ class Worker(LocalOrDistributedWorkerBase):
71 or (speculative_config.draft_model_config.hf_config.model_type ==
72 model_config.hf_config.model_type) \
73 or (speculative_config.draft_model_config.hf_config.model_type
74 - not in ("medusa", "mlp_speculator", "eagle", "deepseek_mtp")) \
74 + not in ("medusa",
75 + "mlp_speculator",
76 + "eagle",
77 + "deepseek_mtp",
78 + "mimo_mtp")) \
79 else {"return_hidden_states": True}
80
81 ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner