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 14, 05:30 UTC
[spec decode] Fix MTP inference path for MiMo-7B model (#25136)
Signed-off-by: zixi-qi <qizixi@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
qizixi committed
Sep 18, 2025 at 09:12 UTC
c4cb0af98a8e39950fa9b99acf7c241959a14ac8
3 files changed
+20
-6
examples/offline_inference/spec_decode.py
+5
-1
@@ -53,7 +53,6 @@ def parse_args():
53
"--method",
54
type=str,
55
default="eagle",
56
- choices=["ngram", "eagle", "eagle3", "mtp"],
56
)
57
parser.add_argument("--num-spec-tokens", type=int, default=2)
58
parser.add_argument("--prompt-lookup-max", type=int, default=5)
@@ -118,6 +117,11 @@ def main():
117
"prompt_lookup_max": args.prompt_lookup_max,
118
"prompt_lookup_min": args.prompt_lookup_min,
119
}
120
+ elif args.method.endswith("mtp"):
121
+ speculative_config = {
122
+ "method": args.method,
123
+ "num_speculative_tokens": args.num_spec_tokens,
124
+ }
125
else:
126
raise ValueError(f"unknown method: {args.method}")
127
vllm/config/speculative.py
+1
-1
@@ -31,7 +31,7 @@ logger = init_logger(__name__)
31
32
SpeculativeMethod = Literal["ngram", "eagle", "eagle3", "medusa",
33
"mlp_speculator", "draft_model", "deepseek_mtp",
34
- "ernie_mtp", "qwen3_next_mtp"]
34
+ "ernie_mtp", "qwen3_next_mtp", "mimo_mtp"]
35
36
37
@config
vllm/model_executor/models/mimo_mtp.py
+14
-4
@@ -241,6 +241,15 @@ class MiMoMTP(nn.Module):
241
242
def map_model_name_to_mtp_param_name(self, name: str) -> str:
243
import regex as re
244
+
245
+ # append mtp_start_layer_idx
246
+ pattern = r"(model\.mtp_layers\.)(\d+)(\.)"
247
+ match = re.match(pattern, name)
248
+ if match:
249
+ original_num = int(match.group(2))
250
+ new_num = original_num + self.config.num_hidden_layers
251
+ name = name.replace(match.group(), f"{match.group(1)}{new_num}.")
252
+ # check for early turn
253
name_without_prefix = [
254
"token_layernorm", "hidden_layernorm", "input_proj",
255
"final_layernorm"
@@ -248,10 +257,11 @@ class MiMoMTP(nn.Module):
257
for sub_name in name_without_prefix:
258
if sub_name in name:
259
return name
251
- pattern = r"model.mtp_layers.(\d+)."
252
- group = re.match(pattern, name)
253
- if group is not None:
254
- name = name.replace(group.group(), group.group() + "mtp_block.")
260
+ # add mtp_block
261
+ pattern = r"(model\.mtp_layers\.\d+\.)"
262
+ match = re.match(pattern, name)
263
+ if match:
264
+ name = name.replace(match.group(), match.group() + "mtp_block.")
265
return name
266
267
def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: