@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 720af6ab79
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 14, 05:30 UTC

[Model][MiniMax-M2] Support MiniMax-M2 Model (#27535)

Signed-off-by: xuebi <xuebi@minimaxi.com> Co-authored-by: xuebi <xuebi@minimaxi.com>

Roger Young committed Oct 27, 2025 at 00:59 UTC 720af6ab791164175eca32c67de7cfe2994642fc
7 files changed +1306
tests/models/registry.py
+3
@@ -341,6 +341,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
341 "MiniMaxM1ForCausalLM": _HfExamplesInfo(
342 "MiniMaxAI/MiniMax-M1-40k", trust_remote_code=True
343 ),
344 + "MiniMaxM2ForCausalLM": _HfExamplesInfo(
345 + "MiniMaxAI/MiniMax-M2", trust_remote_code=True
346 + ),
347 "MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"),
348 "MixtralForCausalLM": _HfExamplesInfo(
349 "mistralai/Mixtral-8x7B-Instruct-v0.1",
vllm/entrypoints/openai/tool_parsers/__init__.py
+2
@@ -16,6 +16,7 @@ from .kimi_k2_tool_parser import KimiK2ToolParser
16 from .llama4_pythonic_tool_parser import Llama4PythonicToolParser
17 from .llama_tool_parser import Llama3JsonToolParser
18 from .longcat_tool_parser import LongcatFlashToolParser
19 +from .minimax_m2_tool_parser import MinimaxM2ToolParser
20 from .minimax_tool_parser import MinimaxToolParser
21 from .mistral_tool_parser import MistralToolParser
22 from .olmo3_tool_parser import Olmo3PythonicToolParser
@@ -56,4 +57,5 @@ __all__ = [
57 "SeedOssToolParser",
58 "Step3ToolParser",
59 "OpenAIToolParser",
60 + "MinimaxM2ToolParser",
61 ]
vllm/entrypoints/openai/tool_parsers/minimax_m2_tool_parser.py new
+644
@@ -0,0 +1,644 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +import re
6 +import uuid
7 +from collections.abc import Sequence
8 +from typing import Any
9 +
10 +from vllm.entrypoints.openai.protocol import (
11 + ChatCompletionRequest,
12 + DeltaFunctionCall,
13 + DeltaMessage,
14 + DeltaToolCall,
15 + ExtractedToolCallInformation,
16 + FunctionCall,
17 + ToolCall,
18 +)
19 +from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
20 + ToolParser,
21 + ToolParserManager,
22 +)
23 +from vllm.logger import init_logger
24 +from vllm.transformers_utils.tokenizer import AnyTokenizer
25 +
26 +logger = init_logger(__name__)
27 +
28 +
29 +@ToolParserManager.register_module("minimax_m2")
30 +class MinimaxM2ToolParser(ToolParser):
31 + def __init__(self, tokenizer: AnyTokenizer):
32 + super().__init__(tokenizer)
33 +
34 + self.prev_tool_call_arr: list[dict] = []
35 +
36 + # Sentinel tokens
37 + self.tool_call_start_token: str = "<minimax:tool_call>"
38 + self.tool_call_end_token: str = "</minimax:tool_call>"
39 + self.invoke_start_prefix: str = "<invoke name="
40 + self.invoke_end_token: str = "</invoke>"
41 + self.parameter_prefix: str = "<parameter name="
42 + self.parameter_end_token: str = "</parameter>"
43 +
44 + # Streaming state variables
45 + self.current_tool_name_sent: bool = False
46 + # Override base class type - we use string IDs for tool calls
47 + self.current_tool_id: str | None = None # type: ignore
48 + self.streamed_args_for_tool: list[str] = []
49 + self.is_tool_call_started: bool = False
50 + self.failed_count: int = 0
51 +
52 + # Initialize streaming state variables
53 + self.current_tool_index: int = 0
54 + self.invoke_index: int = 0
55 + self.header_sent: bool = False
56 + self.current_function_name: str | None = None
57 + self.current_param_name: str | None = None
58 + self.current_param_value: str = ""
59 + self.param_count: int = 0
60 + self.in_param: bool = False
61 + self.in_function: bool = False
62 + self.accumulated_text: str = ""
63 + self.json_started: bool = False
64 + self.json_closed: bool = False
65 + self.accumulated_params: dict = {}
66 + self.streaming_request: ChatCompletionRequest | None = None
67 +
68 + # Enhanced streaming state - reset for each new message
69 + self._reset_streaming_state()
70 +
71 + # Regex patterns for complete parsing
72 + self.tool_call_complete_regex = re.compile(
73 + r"<minimax:tool_call>(.*?)</minimax:tool_call>", re.DOTALL
74 + )
75 + self.invoke_complete_regex = re.compile(
76 + r"<invoke name=(.*?)</invoke>", re.DOTALL
77 + )
78 + self.parameter_complete_regex = re.compile(
79 + r"<parameter name=(.*?)</parameter>", re.DOTALL
80 + )
81 +
82 + if not self.model_tokenizer:
83 + raise ValueError(
84 + "The model tokenizer must be passed to the ToolParser "
85 + "constructor during construction."
86 + )
87 +
88 + self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)
89 + self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
90 +
91 + if self.tool_call_start_token_id is None or self.tool_call_end_token_id is None:
92 + raise RuntimeError(
93 + "MiniMax M2 Tool parser could not locate tool call start/end "
94 + "tokens in the tokenizer!"
95 + )
96 +
97 + logger.info(
98 + "vLLM Successfully import tool parser %s !", self.__class__.__name__
99 + )
100 +
101 + def _generate_tool_call_id(self) -> str:
102 + """Generate a unique tool call ID."""
103 + return f"call_{uuid.uuid4().hex[:24]}"
104 +
105 + def _reset_streaming_state(self):
106 + """Reset all streaming state."""
107 + self.current_tool_index = 0
108 + self.invoke_index = 0
109 + self.is_tool_call_started = False
110 + self.header_sent = False
111 + self.current_tool_id = None
112 + self.current_function_name = None
113 + self.current_param_name = None
114 + self.current_param_value = ""
115 + self.param_count = 0
116 + self.in_param = False
117 + self.in_function = False
118 + self.accumulated_text = ""
119 + self.json_started = False
120 + self.json_closed = False
121 + # Store accumulated parameters for type conversion
122 + self.accumulated_params = {}
123 + self.streaming_request = None
124 + # Clear previous tool call history to avoid state pollution
125 + self.prev_tool_call_arr.clear()
126 +
127 + def _extract_name(self, name_str: str) -> str:
128 + """Extract name from quoted string."""
129 + name_str = name_str.strip()
130 + if (
131 + name_str.startswith('"')
132 + and name_str.endswith('"')
133 + or name_str.startswith("'")
134 + and name_str.endswith("'")
135 + ):
136 + return name_str[1:-1]
137 + return name_str
138 +
139 + def _convert_param_value(self, value: str, param_type: str) -> Any:
140 + """Convert parameter value to the correct type."""
141 + if value.lower() == "null":
142 + return None
143 +
144 + param_type = param_type.lower()
145 + if param_type in ["string", "str", "text"]:
146 + return value
147 + elif param_type in ["integer", "int"]:
148 + try:
149 + return int(value)
150 + except (ValueError, TypeError):
151 + return value
152 + elif param_type in ["number", "float"]:
153 + try:
154 + val = float(value)
155 + return val if val != int(val) else int(val)
156 + except (ValueError, TypeError):
157 + return value
158 + elif param_type in ["boolean", "bool"]:
159 + return value.lower() in ["true", "1"]
160 + elif param_type in ["object", "array"]:
161 + try:
162 + return json.loads(value)
163 + except json.JSONDecodeError:
164 + return value
165 + else:
166 + # Try JSON parse first, fallback to string
167 + try:
168 + return json.loads(value)
169 + except json.JSONDecodeError:
170 + return value
171 +
172 + def _parse_single_invoke(
173 + self, invoke_str: str, tools: list | None
174 + ) -> ToolCall | None:
175 + """Parse a single <invoke> block."""
176 + # Extract function name
177 + name_match = re.search(r"^([^>]+)", invoke_str)
178 + if not name_match:
179 + return None
180 +
181 + function_name = self._extract_name(name_match.group(1))
182 +
183 + # Get parameter configuration
184 + param_config = {}
185 + if tools:
186 + for tool in tools:
187 + if (
188 + hasattr(tool, "function")
189 + and tool.function.name == function_name
190 + and hasattr(tool.function, "parameters")
191 + ):
192 + params = tool.function.parameters
193 + if isinstance(params, dict) and "properties" in params:
194 + param_config = params["properties"]
195 + break
196 +
197 + # Extract parameters
198 + param_dict = {}
199 + for match in self.parameter_complete_regex.findall(invoke_str):
200 + param_match = re.search(r"^([^>]+)>(.*)", match, re.DOTALL)
201 + if param_match:
202 + param_name = self._extract_name(param_match.group(1))
203 + param_value = param_match.group(2).strip()
204 + if param_value.startswith("\n"):
205 + param_value = param_value[1:]
206 + if param_value.endswith("\n"):
207 + param_value = param_value[:-1]
208 +
209 + # Get parameter type
210 + param_type = "string"
211 + if (
212 + param_name in param_config
213 + and isinstance(param_config[param_name], dict)
214 + and "type" in param_config[param_name]
215 + ):
216 + param_type = param_config[param_name]["type"]
217 +
218 + # Convert value
219 + param_dict[param_name] = self._convert_param_value(
220 + param_value, param_type
221 + )
222 +
223 + return ToolCall(
224 + type="function",
225 + function=FunctionCall(
226 + name=function_name,
227 + arguments=json.dumps(param_dict, ensure_ascii=False),
228 + ),
229 + )
230 +
231 + def extract_tool_calls(
232 + self,
233 + model_output: str,
234 + request: ChatCompletionRequest,
235 + ) -> ExtractedToolCallInformation:
236 + """Extract tool calls from complete model output (non-streaming)."""
237 + # Quick check
238 + if self.tool_call_start_token not in model_output:
239 + return ExtractedToolCallInformation(
240 + tools_called=False, tool_calls=[], content=model_output
241 + )
242 +
243 + try:
244 + tool_calls = []
245 +
246 + # Find all complete tool_call blocks
247 + for tool_call_match in self.tool_call_complete_regex.findall(model_output):
248 + # Find all invokes within this tool_call
249 + for invoke_match in self.invoke_complete_regex.findall(tool_call_match):
250 + tool_call = self._parse_single_invoke(
251 + invoke_match, request.tools if request else None
252 + )
253 + if tool_call:
254 + tool_calls.append(tool_call)
255 +
256 + if not tool_calls:
257 + return ExtractedToolCallInformation(
258 + tools_called=False, tool_calls=[], content=model_output
259 + )
260 +
261 + # Update prev_tool_call_arr
262 + self.prev_tool_call_arr.clear()
263 + for tool_call in tool_calls:
264 + self.prev_tool_call_arr.append(
265 + {
266 + "name": tool_call.function.name,
267 + "arguments": tool_call.function.arguments,
268 + }
269 + )
270 +
271 + # Extract content before first tool call
272 + first_tool_idx = model_output.find(self.tool_call_start_token)
273 + content = model_output[:first_tool_idx] if first_tool_idx > 0 else None
274 +
275 + return ExtractedToolCallInformation(
276 + tools_called=True, tool_calls=tool_calls, content=content
277 + )
278 +
279 + except Exception:
280 + logger.exception("Error extracting tool calls")
281 + return ExtractedToolCallInformation(
282 + tools_called=False, tool_calls=[], content=model_output
283 + )
284 +
285 + def extract_tool_calls_streaming(
286 + self,
287 + previous_text: str,
288 + current_text: str,
289 + delta_text: str,
290 + previous_token_ids: Sequence[int], # pylint: disable=unused-argument
291 + current_token_ids: Sequence[int], # pylint: disable=unused-argument
292 + delta_token_ids: Sequence[int],
293 + request: ChatCompletionRequest,
294 + ) -> DeltaMessage | None:
295 + """Extract tool calls from streaming model output."""
296 +
297 + # Store request for type conversion
298 + if not previous_text or self.tool_call_start_token in delta_text:
299 + self._reset_streaming_state()
300 + self.streaming_request = request
301 +
302 + # If no delta text, return None unless it's an EOS token after tools
303 + if not delta_text:
304 + # Check if this is an EOS token after all tool calls are complete
305 + if delta_token_ids and self.tool_call_end_token_id not in delta_token_ids:
306 + # Count complete tool calls
307 + complete_calls = len(
308 + self.tool_call_complete_regex.findall(current_text)
309 + )
310 +
311 + # If we have completed tool calls and populated prev_tool_call_arr
312 + if complete_calls > 0 and len(self.prev_tool_call_arr) > 0:
313 + # Check if all tool calls are closed
314 + open_calls = current_text.count(
315 + self.tool_call_start_token
316 + ) - current_text.count(self.tool_call_end_token)
317 + if open_calls == 0:
318 + # Return empty delta for finish_reason processing
319 + return DeltaMessage(content="")
320 + elif not self.is_tool_call_started and current_text:
321 + # This is a regular content response that's now complete
322 + return DeltaMessage(content="")
323 + return None
324 +
325 + # Update accumulated text
326 + self.accumulated_text = current_text
327 +
328 + # Check if we need to advance to next tool
329 + if self.json_closed and not self.in_function:
330 + # Check if this tool call has ended
331 + invoke_ends = current_text.count(self.invoke_end_token)
332 + if invoke_ends > self.current_tool_index:
333 + # This tool has ended, advance to next
334 + self.current_tool_index += 1
335 + self.header_sent = False
336 + self.param_count = 0
337 + self.json_started = False
338 + self.json_closed = False
339 + self.in_function = False # Now we can safely set this to False
340 + self.accumulated_params = {}
341 + # Continue processing next tool
342 + return None
343 +
344 + # Handle normal content before tool calls
345 + if not self.is_tool_call_started:
346 + # Check if tool call is starting
347 + if (
348 + self.tool_call_start_token_id in delta_token_ids
349 + or self.tool_call_start_token in delta_text
350 + ):
351 + self.is_tool_call_started = True
352 + # Return any content before the tool call
353 + if self.tool_call_start_token in delta_text:
354 + content_before = delta_text[
355 + : delta_text.index(self.tool_call_start_token)
356 + ]
357 + if content_before:
358 + return DeltaMessage(content=content_before)
359 + return None
360 + else:
361 + # Check if we're between tool calls - skip whitespace
362 + if (
363 + current_text.rstrip().endswith(self.tool_call_end_token)
364 + and delta_text.strip() == ""
365 + ):
366 + # We just ended a tool call, skip whitespace
367 + return None
368 + # Normal content, no tool call
369 + return DeltaMessage(content=delta_text)
370 +
371 + # Check if we're between tool calls (waiting for next one)
372 + invoke_starts_count = current_text.count(self.invoke_start_prefix)
373 + if self.current_tool_index >= invoke_starts_count:
374 + # We're past all tool calls, shouldn't be here
375 + return None
376 +
377 + # Find the current tool call portion
378 + invoke_start_positions: list[int] = []
379 + idx = 0
380 + while True:
381 + idx = current_text.find(self.invoke_start_prefix, idx)
382 + if idx == -1:
383 + break
384 + invoke_start_positions.append(idx)
385 + idx += len(self.invoke_start_prefix)
386 +
387 + if self.current_tool_index >= len(invoke_start_positions):
388 + # No more tool calls to process yet
389 + return None
390 +
391 + invoke_start_idx = invoke_start_positions[self.current_tool_index]
392 + # Find where this tool call ends (or current position if not ended yet)
393 + invoke_end_idx = current_text.find(self.invoke_end_token, invoke_start_idx)
394 + if invoke_end_idx == -1:
395 + tool_text = current_text[invoke_start_idx:]
396 + else:
397 + tool_text = current_text[
398 + invoke_start_idx : invoke_end_idx + len(self.invoke_end_token)
399 + ]
400 +
401 + # Looking for function header
402 + if not self.header_sent:
403 + if self.invoke_start_prefix in tool_text:
404 + func_start = tool_text.find(self.invoke_start_prefix) + len(
405 + self.invoke_start_prefix
406 + )
407 + # Find the end quote for the function name
408 + func_end = tool_text.find(">", func_start)
409 +
410 + if func_end != -1:
411 + # Found complete function name
412 + function_name_raw = tool_text[func_start:func_end]
413 + self.current_function_name = self._extract_name(function_name_raw)
414 + self.current_tool_id = self._generate_tool_call_id()
415 + self.header_sent = True
416 + self.in_function = True
417 +
418 + # Add to prev_tool_call_arr immediately when we detect a tool call
419 + # Each tool call should be recorded regardless of function name
420 + # Ensure we don't add the same tool call index multiple times
421 + if len(self.prev_tool_call_arr) <= self.current_tool_index:
422 + self.prev_tool_call_arr.append(
423 + {
424 + "name": self.current_function_name,
425 + "arguments": "{}", # Placeholder, will be updated later
426 + }
427 + )
428 +
429 + # Send header with function info
430 + return DeltaMessage(
431 + tool_calls=[
432 + DeltaToolCall(
433 + index=self.current_tool_index,
434 + id=self.current_tool_id,
435 + function=DeltaFunctionCall(
436 + name=self.current_function_name, arguments=""
437 + ),
438 + type="function",
439 + )
440 + ]
441 + )
442 + return None
443 +
444 + # We've sent header, now handle function body
445 + if self.in_function:
446 + # Send opening brace if not sent yet
447 + if self.in_function and not self.json_started:
448 + self.json_started = True
449 + return DeltaMessage(
450 + tool_calls=[
451 + DeltaToolCall(
452 + index=self.current_tool_index,
453 + function=DeltaFunctionCall(arguments="{"),
454 + )
455 + ]
456 + )
457 +
458 + # Make sure json_started is set if we're processing parameters
459 + if not self.json_started:
460 + self.json_started = True
461 +
462 + # Check for function end in accumulated text
463 + if not self.json_closed and self.invoke_end_token in tool_text:
464 + # Count total parameters in the tool text
465 + total_param_count = tool_text.count(self.parameter_prefix)
466 +
467 + # Only close JSON if all parameters have been processed
468 + if self.param_count >= total_param_count:
469 + # Close JSON
470 + self.json_closed = True
471 +
472 + # Extract complete tool call
473 + # Find the invoke content
474 + invoke_start = tool_text.find(self.invoke_start_prefix) + len(
475 + self.invoke_start_prefix
476 + )
477 + invoke_content_end = tool_text.find(
478 + self.invoke_end_token, invoke_start
479 + )
480 + if invoke_content_end != -1:
481 + invoke_content = tool_text[invoke_start:invoke_content_end]
482 + # Parse to get the complete arguments
483 + try:
484 + parsed_tool = self._parse_single_invoke(
485 + invoke_content,
486 + self.streaming_request.tools
487 + if self.streaming_request
488 + else None,
489 + )
490 + if parsed_tool and self.current_tool_index < len(
491 + self.prev_tool_call_arr
492 + ):
493 + # Update existing entry in prev_tool_call_arr
494 + args = parsed_tool.function.arguments
495 + self.prev_tool_call_arr[self.current_tool_index][
496 + "arguments"
497 + ] = args
498 + except Exception:
499 + pass # Ignore parsing errors during streaming
500 +
501 + result = DeltaMessage(
502 + tool_calls=[
503 + DeltaToolCall(
504 + index=self.current_tool_index,
505 + function=DeltaFunctionCall(arguments="}"),
506 + )
507 + ]
508 + )
509 +
510 + # Reset state for next tool
511 + self.json_closed = True
512 + self.in_function = False
513 + self.accumulated_params = {}
514 +
515 + logger.debug("[M2_STREAMING] Tool call completed")
516 +
517 + return result
518 + else:
519 + # Don't close JSON yet, continue processing parameters
520 + return None
521 +
522 + # Look for parameters
523 + # Find all parameter starts
524 + param_starts = []
525 + idx = 0
526 + while True:
527 + idx = tool_text.find(self.parameter_prefix, idx)
528 + if idx == -1:
529 + break
530 + param_starts.append(idx)
531 + idx += len(self.parameter_prefix)
532 +
533 + # Check if we should start a new parameter
534 + if (
535 + not self.in_param
536 + and self.param_count < len(param_starts)
537 + and len(param_starts) > self.param_count
538 + ):
539 + # Process the next parameter
540 + param_idx = param_starts[self.param_count]
541 + param_start = param_idx + len(self.parameter_prefix)
542 + remaining = tool_text[param_start:]
543 +
544 + if ">" in remaining:
545 + # We have the complete parameter name
546 + name_end = remaining.find(">")
547 + param_name_raw = remaining[:name_end]
548 + self.current_param_name = self._extract_name(param_name_raw)
549 +
550 + # Find the parameter value
551 + value_start = param_start + name_end + 1
552 + value_text = tool_text[value_start:]
553 + if value_text.startswith("\n"):
554 + value_text = value_text[1:]
555 +
556 + # Find where this parameter ends
557 + param_end_idx = value_text.find(self.parameter_end_token)
558 + if param_end_idx == -1:
559 + # No closing tag, look for next parameter or function end
560 + next_param_idx = value_text.find(self.parameter_prefix)
561 + func_end_idx = value_text.find(self.invoke_end_token)
562 +
563 + if next_param_idx != -1 and (
564 + func_end_idx == -1 or next_param_idx < func_end_idx
565 + ):
566 + param_end_idx = next_param_idx
567 + elif func_end_idx != -1:
568 + param_end_idx = func_end_idx
569 + else:
570 + # Neither found, check if tool call is complete
571 + if self.invoke_end_token in tool_text:
572 + # Tool call and parameter is complete
573 + param_end_idx = len(value_text)
574 + else:
575 + # Still streaming, wait for more content
576 + return None
577 +
578 + if param_end_idx != -1:
579 + # Complete parameter found
580 + param_value = value_text[:param_end_idx]
581 + if param_value.endswith("\n"):
582 + param_value = param_value[:-1]
583 +
584 + # Store raw value for later processing
585 + self.accumulated_params[self.current_param_name] = param_value
586 +
587 + # Get parameter configuration for type conversion
588 + param_config = {}
589 + if self.streaming_request and self.streaming_request.tools:
590 + for tool in self.streaming_request.tools:
591 + if (
592 + hasattr(tool, "function")
593 + and tool.function.name == self.current_function_name
594 + and hasattr(tool.function, "parameters")
595 + ):
596 + params = tool.function.parameters
597 + if (
598 + isinstance(params, dict)
599 + and "properties" in params
600 + ):
601 + param_config = params["properties"]
602 + break
603 +
604 + # Get parameter type
605 + param_type = "string"
606 + if (
607 + self.current_param_name in param_config
608 + and isinstance(param_config[self.current_param_name], dict)
609 + and "type" in param_config[self.current_param_name]
610 + ):
611 + param_type = param_config[self.current_param_name]["type"]
612 +
613 + # Convert param value to appropriate type
614 + converted_value = self._convert_param_value(
615 + param_value, param_type
616 + )
617 +
618 + # Build JSON fragment based on the converted type
619 + # Use json.dumps to properly serialize the value
620 + serialized_value = json.dumps(
621 + converted_value, ensure_ascii=False
622 + )
623 +
624 + if self.param_count == 0:
625 + json_fragment = (
626 + f'"{self.current_param_name}": {serialized_value}'
627 + )
628 + else:
629 + json_fragment = (
630 + f', "{self.current_param_name}": {serialized_value}'
631 + )
632 +
633 + self.param_count += 1
634 +
635 + return DeltaMessage(
636 + tool_calls=[
637 + DeltaToolCall(
638 + index=self.current_tool_index,
639 + function=DeltaFunctionCall(arguments=json_fragment),
640 + )
641 + ]
642 + )
643 +
644 + return None
vllm/model_executor/models/minimax_m2.py new
+585
@@ -0,0 +1,585 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Adapted from
5 +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
6 +# Copyright 2023 The vLLM team.
7 +# Copyright 2023 DeepSeek-AI and the HuggingFace Inc. team. All rights reserved.
8 +#
9 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10 +# and OPT implementations in this library. It has been modified from its
11 +# original forms to accommodate minor architectural differences compared
12 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13 +#
14 +# Licensed under the Apache License, Version 2.0 (the "License");
15 +# you may not use this file except in compliance with the License.
16 +# You may obtain a copy of the License at
17 +#
18 +# http://www.apache.org/licenses/LICENSE-2.0
19 +#
20 +# Unless required by applicable law or agreed to in writing, software
21 +# distributed under the License is distributed on an "AS IS" BASIS,
22 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23 +# See the License for the specific language governing permissions and
24 +# limitations under the License.
25 +"""Inference-only MiniMaxM2 model."""
26 +
27 +from collections.abc import Iterable
28 +from typing import Any
29 +
30 +import torch
31 +from torch import nn
32 +from transformers import PretrainedConfig
33 +
34 +from vllm.attention import Attention
35 +from vllm.compilation.decorators import support_torch_compile
36 +from vllm.config import CacheConfig, ModelConfig, VllmConfig
37 +from vllm.distributed import (
38 + get_pp_group,
39 + get_tensor_model_parallel_world_size,
40 + tensor_model_parallel_all_reduce,
41 +)
42 +from vllm.model_executor.layers.fused_moe import FusedMoE
43 +from vllm.model_executor.layers.layernorm import RMSNorm
44 +from vllm.model_executor.layers.linear import (
45 + QKVParallelLinear,
46 + ReplicatedLinear,
47 + RowParallelLinear,
48 +)
49 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
50 +from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01RMSNormTP
51 +from vllm.model_executor.layers.quantization import QuantizationConfig
52 +from vllm.model_executor.layers.rotary_embedding import get_rope
53 +from vllm.model_executor.layers.vocab_parallel_embedding import (
54 + ParallelLMHead,
55 + VocabParallelEmbedding,
56 +)
57 +from vllm.model_executor.model_loader.weight_utils import (
58 + default_weight_loader,
59 + maybe_remap_kv_scale_name,
60 +)
61 +from vllm.sequence import IntermediateTensors
62 +
63 +from .interfaces import SupportsPP
64 +from .utils import (
65 + AutoWeightsLoader,
66 + PPMissingLayer,
67 + is_pp_missing_parameter,
68 + make_empty_intermediate_tensors_factory,
69 + make_layers,
70 + maybe_prefix,
71 +)
72 +
73 +
74 +class MiniMaxM2MoE(nn.Module):
75 + def __init__(
76 + self,
77 + config: PretrainedConfig,
78 + quant_config: QuantizationConfig | None = None,
79 + prefix: str = "",
80 + ):
81 + super().__init__()
82 + self.tp_size = get_tensor_model_parallel_world_size()
83 +
84 + if self.tp_size > config.num_local_experts:
85 + raise ValueError(
86 + f"Tensor parallel size {self.tp_size} is greater than "
87 + f"the number of experts {config.num_local_experts}."
88 + )
89 + self.use_routing_bias = getattr(config, "use_routing_bias", False)
90 + if self.use_routing_bias:
91 + self.e_score_correction_bias = nn.Parameter(
92 + torch.empty(config.num_local_experts, dtype=torch.float32)
93 + )
94 + self.e_score_correction_bias.weight_loader = (
95 + MiniMaxM2MoE.ebias_weight_loader
96 + )
97 + else:
98 + self.e_score_correction_bias = None
99 +
100 + self.experts = FusedMoE(
101 + num_experts=config.num_local_experts,
102 + top_k=config.num_experts_per_tok,
103 + scoring_func=config.scoring_func,
104 + use_grouped_topk=True,
105 + num_expert_group=1,
106 + topk_group=1,
107 + e_score_correction_bias=self.e_score_correction_bias,
108 + hidden_size=config.hidden_size,
109 + intermediate_size=config.intermediate_size,
110 + reduce_results=False,
111 + renormalize=True,
112 + quant_config=quant_config,
113 + prefix=f"{prefix}.experts",
114 + )
115 +
116 + self.gate = ReplicatedLinear(
117 + config.hidden_size,
118 + config.num_local_experts,
119 + bias=False,
120 + params_dtype=torch.float32,
121 + quant_config=None,
122 + prefix=f"{prefix}.gate",
123 + )
124 +
125 + @staticmethod
126 + def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None:
127 + assert param.size() == loaded_weight.size()
128 + param.data.copy_(loaded_weight.to(torch.float32))
129 +
130 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
131 + num_tokens, hidden_dim = hidden_states.shape
132 + hidden_states = hidden_states.view(-1, hidden_dim)
133 +
134 + # router_logits: (num_tokens, n_experts)
135 + router_logits, _ = self.gate(hidden_states.to(torch.float32))
136 + final_hidden_states = self.experts(
137 + hidden_states=hidden_states, router_logits=router_logits
138 + )
139 + final_hidden_states = final_hidden_states
140 + if self.tp_size > 1:
141 + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
142 +
143 + return final_hidden_states.view(num_tokens, hidden_dim)
144 +
145 +
146 +class MiniMaxM2Attention(nn.Module):
147 + def __init__(
148 + self,
149 + hidden_size: int,
150 + num_heads: int,
151 + num_kv_heads: int,
152 + rotary_dim: int,
153 + rope_theta: float = 10000,
154 + rope_scaling: dict[str, Any] | None = None,
155 + attn_window_size: int | None = None,
156 + max_position_embeddings: int = 8192,
157 + head_dim: int | None = None,
158 + rms_norm_eps: float = 1e-06,
159 + qkv_bias: bool = False,
160 + cache_config: CacheConfig | None = None,
161 + quant_config: QuantizationConfig | None = None,
162 + prefix: str = "",
163 + ) -> None:
164 + super().__init__()
165 + self.hidden_size = hidden_size
166 + tp_size = get_tensor_model_parallel_world_size()
167 + self.total_num_heads = num_heads
168 + assert self.total_num_heads % tp_size == 0
169 + self.num_heads = self.total_num_heads // tp_size
170 + self.total_num_kv_heads = num_kv_heads
171 + if self.total_num_kv_heads >= tp_size:
172 + # Number of KV heads is greater than TP size, so we partition
173 + # the KV heads across multiple tensor parallel GPUs.
174 + assert self.total_num_kv_heads % tp_size == 0
175 + else:
176 + # Number of KV heads is less than TP size, so we replicate
177 + # the KV heads across multiple tensor parallel GPUs.
178 + assert tp_size % self.total_num_kv_heads == 0
179 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
180 + self.head_dim = head_dim or (hidden_size // self.total_num_heads)
181 + self.q_size = self.num_heads * self.head_dim
182 + self.kv_size = self.num_kv_heads * self.head_dim
183 + self.scaling = self.head_dim**-0.5
184 + self.rope_theta = rope_theta
185 + self.max_position_embeddings = max_position_embeddings
186 +
187 + self.qkv_proj = QKVParallelLinear(
188 + hidden_size,
189 + self.head_dim,
190 + self.total_num_heads,
191 + self.total_num_kv_heads,
192 + bias=qkv_bias,
193 + quant_config=quant_config,
194 + prefix=f"{prefix}.qkv_proj",
195 + )
196 +
197 + self.o_proj = RowParallelLinear(
198 + self.total_num_heads * self.head_dim,
199 + hidden_size,
200 + bias=False,
201 + quant_config=quant_config,
202 + prefix=f"{prefix}.o_proj",
203 + )
204 +
205 + self.rotary_emb = get_rope(
206 + self.head_dim,
207 + rotary_dim=rotary_dim,
208 + max_position=max_position_embeddings,
209 + base=rope_theta,
210 + rope_scaling=rope_scaling,
211 + )
212 + self.attn = Attention(
213 + self.num_heads,
214 + self.head_dim,
215 + self.scaling,
216 + num_kv_heads=self.num_kv_heads,
217 + per_layer_sliding_window=attn_window_size,
218 + cache_config=cache_config,
219 + quant_config=quant_config,
220 + prefix=f"{prefix}.attn",
221 + )
222 +
223 + self.q_norm = MiniMaxText01RMSNormTP(
224 + self.head_dim * self.total_num_heads, eps=rms_norm_eps
225 + )
226 + self.k_norm = MiniMaxText01RMSNormTP(
227 + self.head_dim * self.total_num_kv_heads, eps=rms_norm_eps
228 + )
229 +
230 + def forward(
231 + self,
232 + positions: torch.Tensor,
233 + hidden_states: torch.Tensor,
234 + ) -> torch.Tensor:
235 + qkv, _ = self.qkv_proj(hidden_states)
236 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
237 + q = self.q_norm(q)
238 + k = self.k_norm(k)
239 + q, k = self.rotary_emb(positions, q, k)
240 + attn_output = self.attn(q, k, v)
241 + output, _ = self.o_proj(attn_output)
242 + return output
243 +
244 +
245 +class MiniMaxM2DecoderLayer(nn.Module):
246 + def __init__(
247 + self,
248 + config: PretrainedConfig,
249 + prefix: str,
250 + model_config: ModelConfig,
251 + cache_config: CacheConfig | None = None,
252 + quant_config: QuantizationConfig | None = None,
253 + ) -> None:
254 + super().__init__()
255 + self.hidden_size = config.hidden_size
256 + rope_theta = getattr(config, "rope_theta", 10000)
257 + rope_scaling = getattr(config, "rope_scaling", None)
258 + max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
259 + if hasattr(config, "max_model_len") and isinstance(config.max_model_len, int):
260 + max_position_embeddings = max(
261 + config.max_position_embeddings, config.max_model_len
262 + )
263 + # DecoderLayers are created with `make_layers` which passes the prefix
264 + # with the layer's index.
265 + layer_idx = int(prefix.split(sep=".")[-1])
266 +
267 + # TODO: support MTP
268 + attn_window_size = getattr(config, "attn_window_size", None)
269 + if attn_window_size is not None:
270 + if isinstance(attn_window_size, list):
271 + attn_window_size = attn_window_size[layer_idx]
272 + elif isinstance(attn_window_size, int):
273 + attn_window_size = attn_window_size
274 + else:
275 + raise ValueError(f"Invalid attn_window_size: {attn_window_size}")
276 + attn_window_size = None if attn_window_size <= 0 else attn_window_size
277 +
278 + # different rope theta for full layer and swa layer
279 + swa_rope_theta = getattr(config, "swa_rope_theta", -1)
280 + # default to full rope theta
281 + swa_rope_theta = rope_theta if swa_rope_theta <= 0 else swa_rope_theta
282 + rope_theta = swa_rope_theta if attn_window_size is not None else rope_theta
283 +
284 + self.layer_idx = layer_idx
285 + self.self_attn = MiniMaxM2Attention(
286 + hidden_size=self.hidden_size,
287 + num_heads=config.num_attention_heads,
288 + num_kv_heads=config.num_key_value_heads,
289 + rotary_dim=config.rotary_dim,
290 + rope_theta=rope_theta,
291 + rope_scaling=rope_scaling,
292 + attn_window_size=attn_window_size,
293 + max_position_embeddings=max_position_embeddings,
294 + rms_norm_eps=config.rms_norm_eps,
295 + qkv_bias=getattr(config, "attention_bias", False),
296 + head_dim=getattr(config, "head_dim", None),
297 + cache_config=cache_config,
298 + quant_config=quant_config,
299 + prefix=f"{prefix}.self_attn",
300 + )
301 +
302 + self.block_sparse_moe = MiniMaxM2MoE(
303 + config=config,
304 + quant_config=quant_config,
305 + prefix=f"{prefix}.mlp",
306 + )
307 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
308 + self.post_attention_layernorm = RMSNorm(
309 + config.hidden_size, eps=config.rms_norm_eps
310 + )
311 +
312 + def forward(
313 + self,
314 + positions: torch.Tensor,
315 + hidden_states: torch.Tensor,
316 + residual: torch.Tensor | None,
317 + ) -> torch.Tensor:
318 + # Self Attention
319 + if residual is None:
320 + residual = hidden_states
321 + hidden_states = self.input_layernorm(hidden_states)
322 + else:
323 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
324 + hidden_states = self.self_attn(
325 + positions=positions,
326 + hidden_states=hidden_states,
327 + )
328 +
329 + # Fully Connected
330 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
331 +
332 + hidden_states = self.block_sparse_moe(hidden_states)
333 +
334 + return hidden_states, residual
335 +
336 +
337 +@support_torch_compile
338 +class MiniMaxM2Model(nn.Module):
339 + fall_back_to_pt_during_load = False
340 +
341 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
342 + super().__init__()
343 +
344 + config = vllm_config.model_config.hf_config
345 + model_config = vllm_config.model_config
346 + cache_config = vllm_config.cache_config
347 + quant_config = vllm_config.quant_config
348 + self.config = config
349 +
350 + self.vocab_size = config.vocab_size
351 +
352 + if get_pp_group().is_first_rank:
353 + self.embed_tokens = VocabParallelEmbedding(
354 + config.vocab_size,
355 + config.hidden_size,
356 + quant_config=None,
357 + prefix=f"{prefix}.embed_tokens",
358 + )
359 + else:
360 + self.embed_tokens = PPMissingLayer()
361 +
362 + self.start_layer, self.end_layer, self.layers = make_layers(
363 + config.num_hidden_layers,
364 + lambda prefix: MiniMaxM2DecoderLayer(
365 + config,
366 + prefix,
367 + model_config=model_config,
368 + cache_config=cache_config,
369 + quant_config=quant_config,
370 + ),
371 + prefix=f"{prefix}.layers",
372 + )
373 +
374 + if get_pp_group().is_last_rank:
375 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
376 + else:
377 + self.norm = PPMissingLayer()
378 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
379 + ["hidden_states", "residual"], config.hidden_size
380 + )
381 +
382 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
383 + return self.embed_tokens(input_ids)
384 +
385 + def forward(
386 + self,
387 + input_ids: torch.Tensor,
388 + positions: torch.Tensor,
389 + intermediate_tensors: IntermediateTensors | None,
390 + inputs_embeds: torch.Tensor | None = None,
391 + ) -> torch.Tensor | IntermediateTensors:
392 + if get_pp_group().is_first_rank:
393 + if inputs_embeds is not None:
394 + hidden_states = inputs_embeds
395 + else:
396 + hidden_states = self.get_input_embeddings(input_ids)
397 + residual = None
398 + else:
399 + assert intermediate_tensors is not None
400 + hidden_states = intermediate_tensors["hidden_states"]
401 + residual = intermediate_tensors["residual"]
402 +
403 + for layer in self.layers[self.start_layer : self.end_layer]:
404 + hidden_states, residual = layer(positions, hidden_states, residual)
405 +
406 + if not get_pp_group().is_last_rank:
407 + return IntermediateTensors(
408 + {"hidden_states": hidden_states, "residual": residual}
409 + )
410 + hidden_states, _ = self.norm(hidden_states, residual)
411 + return hidden_states
412 +
413 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
414 + return FusedMoE.make_expert_params_mapping(
415 + ckpt_gate_proj_name="w1",
416 + ckpt_down_proj_name="w2",
417 + ckpt_up_proj_name="w3",
418 + num_experts=self.config.num_local_experts,
419 + )
420 +
421 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
422 + stacked_params_mapping = [
423 + # (param_name, shard_name, shard_id)
424 + ("qkv_proj", "q_proj", "q"),
425 + ("qkv_proj", "k_proj", "k"),
426 + ("qkv_proj", "v_proj", "v"),
427 + ]
428 +
429 + # Params for weights, fp8 weight scales, fp8 activation scales
430 + # (param_name, weight_name, expert_id, shard_id)
431 + expert_params_mapping = self.get_expert_mapping()
432 +
433 + params_dict = dict(self.named_parameters())
434 + loaded_params: set[str] = set()
435 + for name, loaded_weight in weights:
436 + if "rotary_emb.inv_freq" in name:
437 + continue
438 +
439 + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
440 + if spec_layer is not None:
441 + continue # skip spec decode layers for main model
442 +
443 + for param_name, weight_name, shard_id in stacked_params_mapping:
444 + # Skip non-stacked layers and experts (experts handled below).
445 + if weight_name not in name:
446 + continue
447 + # We have mlp.experts[0].gate_proj in the checkpoint.
448 + # Since we handle the experts below in expert_params_mapping,
449 + # we need to skip here BEFORE we update the name, otherwise
450 + # name will be updated to mlp.experts[0].gate_up_proj, which
451 + # will then be updated below in expert_params_mapping
452 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.
453 + if ("mlp.experts." in name) and name not in params_dict:
454 + continue
455 + name = name.replace(weight_name, param_name)
456 + # Skip loading extra bias for GPTQ models.
457 + if name.endswith(".bias") and name not in params_dict:
458 + continue
459 +
460 + if is_pp_missing_parameter(name, self):
461 + continue
462 +
463 + param = params_dict[name]
464 + weight_loader = param.weight_loader
465 + weight_loader(param, loaded_weight, shard_id)
466 + break
467 + else:
468 + for mapping in expert_params_mapping:
469 + param_name, weight_name, expert_id, shard_id = mapping
470 + if weight_name not in name:
471 + continue
472 + name = name.replace(weight_name, param_name)
473 +
474 + if is_pp_missing_parameter(name, self):
475 + continue
476 +
477 + param = params_dict[name]
478 + weight_loader = param.weight_loader
479 + weight_loader(
480 + param,
481 + loaded_weight,
482 + name,
483 + shard_id=shard_id,
484 + expert_id=expert_id,
485 + )
486 + break
487 + else:
488 + # Skip loading extra bias for GPTQ models.
489 + if name.endswith(".bias") and name not in params_dict:
490 + continue
491 +
492 + # Remapping the name of FP8 kv-scale.
493 + name = maybe_remap_kv_scale_name(name, params_dict)
494 + if name is None:
495 + continue
496 +
497 + if is_pp_missing_parameter(name, self):
498 + continue
499 +
500 + param = params_dict[name]
501 + weight_loader = getattr(
502 + param, "weight_loader", default_weight_loader
503 + )
504 + weight_loader(param, loaded_weight)
505 + loaded_params.add(name)
506 + return loaded_params
507 +
508 +
509 +class MiniMaxM2ForCausalLM(nn.Module, SupportsPP):
510 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
511 + super().__init__()
512 + config = vllm_config.model_config.hf_config
513 + quant_config = vllm_config.quant_config
514 + self.config = config
515 + self.quant_config = quant_config
516 + if hasattr(vllm_config.model_config, "max_model_len"):
517 + self.config.max_model_len = vllm_config.model_config.max_model_len
518 + self.model = MiniMaxM2Model(
519 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
520 + )
521 + if get_pp_group().is_last_rank:
522 + self.lm_head = ParallelLMHead(
523 + config.vocab_size, config.hidden_size, quant_config=None
524 + )
525 + else:
526 + self.lm_head = PPMissingLayer()
527 + self.logits_processor = LogitsProcessor(config.vocab_size)
528 + self.make_empty_intermediate_tensors = (
529 + self.model.make_empty_intermediate_tensors
530 + )
531 +
532 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
533 + return self.model.get_input_embeddings(input_ids)
534 +
535 + def forward(
536 + self,
537 + input_ids: torch.Tensor,
538 + positions: torch.Tensor,
539 + intermediate_tensors: IntermediateTensors | None = None,
540 + inputs_embeds: torch.Tensor | None = None,
541 + **kwargs,
542 + ) -> torch.Tensor | IntermediateTensors:
543 + hidden_states = self.model(
544 + input_ids, positions, intermediate_tensors, inputs_embeds
545 + )
546 + return hidden_states
547 +
548 + def compute_logits(
549 + self,
550 + hidden_states: torch.Tensor,
551 + ) -> torch.Tensor | None:
552 + logits = self.logits_processor(self.lm_head, hidden_states)
553 + return logits
554 +
555 + def make_empty_intermediate_tensors(
556 + self, batch_size: int, dtype: torch.dtype, device: torch.device
557 + ) -> IntermediateTensors:
558 + return IntermediateTensors(
559 + {
560 + "hidden_states": torch.zeros(
561 + (batch_size, self.config.hidden_size), dtype=dtype, device=device
562 + ),
563 + "residual": torch.zeros(
564 + (batch_size, self.config.hidden_size), dtype=dtype, device=device
565 + ),
566 + }
567 + )
568 +
569 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
570 + loader = AutoWeightsLoader(self)
571 + return loader.load_weights(weights)
572 +
573 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
574 + return self.model.get_expert_mapping()
575 +
576 +
577 +def get_spec_layer_idx_from_weight_name(
578 + config: PretrainedConfig, weight_name: str
579 +) -> int | None:
580 + if hasattr(config, "num_mtp_modules") and (config.num_mtp_modules > 0):
581 + layer_idx = config.num_hidden_layers
582 + for i in range(config.num_mtp_modules):
583 + if weight_name.startswith(f"model.layers.{layer_idx + i}."):
584 + return layer_idx + i
585 + return None
vllm/model_executor/models/registry.py
+1
@@ -131,6 +131,7 @@ _TEXT_GENERATION_MODELS = {
131 "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
132 "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
133 "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
134 + "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
135 "MistralForCausalLM": ("llama", "LlamaForCausalLM"),
136 "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
137 # transformers's mpt class has lower case
vllm/reasoning/__init__.py
+2
@@ -11,6 +11,7 @@ from .gptoss_reasoning_parser import GptOssReasoningParser
11 from .granite_reasoning_parser import GraniteReasoningParser
12 from .hunyuan_a13b_reasoning_parser import HunyuanA13BReasoningParser
13 from .identity_reasoning_parser import IdentityReasoningParser
14 +from .minimax_m2_reasoning_parser import MiniMaxM2ReasoningParser
15 from .mistral_reasoning_parser import MistralReasoningParser
16 from .olmo3_reasoning_parser import Olmo3ReasoningParser
17 from .qwen3_reasoning_parser import Qwen3ReasoningParser
@@ -34,4 +35,5 @@ __all__ = [
35 "Step3ReasoningParser",
36 "GptOssReasoningParser",
37 "SeedOSSReasoningParser",
38 + "MiniMaxM2ReasoningParser",
39 ]
vllm/reasoning/minimax_m2_reasoning_parser.py new
+69
@@ -0,0 +1,69 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Sequence
5 +
6 +from vllm.entrypoints.openai.protocol import (
7 + ChatCompletionRequest,
8 + DeltaMessage,
9 + ResponsesRequest,
10 +)
11 +from vllm.logger import init_logger
12 +from vllm.reasoning.abs_reasoning_parsers import ReasoningParser, ReasoningParserManager
13 +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
14 +from vllm.transformers_utils.tokenizer import AnyTokenizer
15 +
16 +logger = init_logger(__name__)
17 +
18 +
19 +@ReasoningParserManager.register_module("minimax_m2")
20 +class MiniMaxM2ReasoningParser(BaseThinkingReasoningParser):
21 + """
22 + Reasoning parser for MiniMax M2 model.
23 + """
24 +
25 + @property
26 + def start_token(self) -> str:
27 + """The token that starts reasoning content."""
28 + return "<think>"
29 +
30 + @property
31 + def end_token(self) -> str:
32 + """The token that ends reasoning content."""
33 + return "</think>"
34 +
35 +
36 +@ReasoningParserManager.register_module("minimax_m2_append_think")
37 +class MiniMaxM2AppendThinkReasoningParser(ReasoningParser):
38 + """
39 + Reasoning parser for MiniMax M2 model.
40 + """
41 +
42 + def __init__(self, tokenizer: AnyTokenizer, *args, **kwargs):
43 + super().__init__(tokenizer, *args, **kwargs)
44 + self.end_token_id = self.vocab.get("</think>")
45 +
46 + def is_reasoning_end(self, input_ids: list[int]) -> bool:
47 + end_token_id = self.end_token_id
48 + return any(input_id == end_token_id for input_id in reversed(input_ids))
49 +
50 + def extract_content_ids(self, input_ids: list[int]) -> list[int]:
51 + return input_ids
52 +
53 + def extract_reasoning_content_streaming(
54 + self,
55 + previous_text: str,
56 + current_text: str,
57 + delta_text: str,
58 + previous_token_ids: Sequence[int],
59 + current_token_ids: Sequence[int],
60 + delta_token_ids: Sequence[int],
61 + ) -> DeltaMessage | None:
62 + if len(previous_token_ids) == 0:
63 + delta_text = "<think>" + delta_text
64 + return DeltaMessage(content=delta_text)
65 +
66 + def extract_reasoning_content(
67 + self, model_output: str, request: ChatCompletionRequest | ResponsesRequest
68 + ) -> tuple[str | None, str | None]:
69 + return None, "<think>" + model_output