[Perf][Parser] Stream MiniMax M2 tool argument values incrementally (16 KB call: 5.9 s to 26 ms) (#62035)

Signed-off-by: Afonso Ribeiro <afonso.diamantino2@gmail.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>

tests/parser/engine/test_minimax_m2.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import dataclasses
 import json
 
 import pytest
@@ -15,11 +16,14 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionToolsParam,
     FunctionDefinition,
 )
+from vllm.parser.engine.parser_engine import ParserEngine
 from vllm.parser.minimax_m2 import (
     THINK_END,
     TOOL_CALL_END,
     TOOL_CALL_START,
     MinimaxM2Parser,
+    _minimax_m2_arg_value_extender,
+    minimax_m2_config,
 )
 
 
@@ -238,6 +242,123 @@ class TestStreaming:
         assert collect_tool_arguments(results) == ""
 
 
+WRITE_CALL_HEAD = (
+    '<minimax:tool_call>\n<invoke name="write_file">\n'
+    '<parameter name="path">/repo/app.py</parameter>\n'
+    '<parameter name="mode">644</parameter>\n'
+    '<parameter name="content">'
+)
+WRITE_CALL_TAIL = "</parameter>\n</invoke>\n</minimax:tool_call>"
+# Escapes, non-ASCII, blank lines and tag-like text that must stay inside
+# the value, including prefixes of the closing and next-parameter tags.
+CONTENT_LINE = 'if a < b > c: print("x\\y")  # </p> </param <parameter\té\n\n'
+
+
+class TestArgValueExtender:
+    """The open parameter value streams without re-running the converter."""
+
+    @pytest.fixture
+    def write_tool(self):
+        return ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="write_file",
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        "path": {"type": "string"},
+                        "mode": {"type": "integer"},
+                        "content": {"type": "string"},
+                    },
+                },
+            ),
+        )
+
+    @pytest.mark.parametrize("use_schema", [False, True], ids=["no_schema", "schema"])
+    def test_long_call_streams_non_streaming_arguments(
+        self, mock_tokenizer, mock_request, write_tool, use_schema
+    ):
+        mock_request.tools = [write_tool] if use_schema else []
+        text = WRITE_CALL_HEAD + CONTENT_LINE * 200 + WRITE_CALL_TAIL
+
+        results = simulate_tool_streaming(
+            MinimaxM2Parser(mock_tokenizer),
+            mock_request,
+            [text[i : i + 7] for i in range(0, len(text), 7)],
+        )
+        result = MinimaxM2Parser(mock_tokenizer).extract_tool_calls(text, mock_request)
+
+        arguments = collect_tool_arguments(results)
+        assert arguments == result.tool_calls[0].function.arguments
+        assert json.loads(arguments)["content"] == CONTENT_LINE * 200
+        assert ('"mode": 644,' if use_schema else '"mode": "644",') in arguments
+
+    def test_conversions_do_not_grow_with_argument_length(
+        self, mock_tokenizer, mock_request, write_tool
+    ):
+        mock_request.tools = [write_tool]
+        base = minimax_m2_config()
+        converter = base.arg_converter
+        assert converter is not None
+
+        def streaming_conversions(lines: int) -> int:
+            calls = 0
+
+            def counting_converter(raw_args: str, partial: bool) -> str:
+                nonlocal calls
+                calls += partial
+                return converter(raw_args, partial)
+
+            parser = ParserEngine(
+                mock_tokenizer,
+                parser_engine_config=dataclasses.replace(
+                    base, arg_converter=counting_converter
+                ),
+            )
+            text = WRITE_CALL_HEAD + CONTENT_LINE * lines + WRITE_CALL_TAIL
+            simulate_tool_streaming(
+                parser,
+                mock_request,
+                [text[i : i + 4] for i in range(0, len(text), 4)],
+            )
+            return calls
+
+        assert streaming_conversions(lines=20) == streaming_conversions(lines=200)
+
+    def test_extender_stops_at_split_end_tag(self):
+        extend = _minimax_m2_arg_value_extender('<parameter name="k">')
+        assert extend is not None
+        deltas = ['a"\n', "</param", "eter>"]
+        assert [extend(d) for d in deltas] == ['a\\"\\n', "</param", None]
+
+    def test_extender_stops_at_split_next_parameter(self):
+        extend = _minimax_m2_arg_value_extender("<parameter name='k'>x")
+        assert extend is not None
+        assert [extend(d) for d in ["< parameter", " name ="]] == ["< parameter", None]
+
+    @pytest.mark.parametrize(
+        "raw_args",
+        [
+            '<parameter name="k">x</parameter>',
+            '<parameter name="k',
+            (
+                '<parameter name="k">1</parameter><parameter name="j">2</parameter>'
+                '<parameter name="k">3'
+            ),
+            '<parameter name="a<parameter name=k>x',
+            '<parameter name="k>x',
+        ],
+        ids=[
+            "closed_value",
+            "open_name",
+            "repeated_key_not_last",
+            "earlier_unmatched_opener",
+            "bare_name_opens_quote",
+        ],
+    )
+    def test_no_extender_unless_open_value_is_stable(self, raw_args):
+        assert _minimax_m2_arg_value_extender(raw_args) is None
+
+
 class TestReasoning:
     def test_extract_reasoning_without_start_token(self, parser, mock_request):
         reasoning, content = parser.extract_reasoning(

vllm/parser/minimax_m2.py

@@ -22,6 +22,7 @@ import regex as re
 from vllm.parser.engine.events import EventType
 from vllm.parser.engine.parser_engine import ParserEngine
 from vllm.parser.engine.parser_engine_config import (
+    ArgValueExtender,
     ParserEngineConfig,
     ParserState,
     Transition,
@@ -57,37 +58,85 @@ _PARTIAL_PARAM_RE = re.compile(
 )
 
 
-def _minimax_m2_arg_converter(raw_args: str, partial: bool) -> str:
+# What ends an open value in ``_PARAM_RE``: a closing tag or the next opener.
+_PARAM_VALUE_END_RE = re.compile(r"<\s*/\s*parameter\s*>|<\s*parameter\s+name\s*=")
+
+
+def _param_name(match: re.Match) -> str:
+    return (
+        match.group("dq_name")
+        or match.group("sq_name")
+        or match.group("bare_name")
+        or ""
+    ).strip()
+
+
+def _minimax_m2_parse_args(
+    raw_args: str, partial: bool
+) -> tuple[dict[str, object], re.Match | None]:
+    """Return the arguments and the match of the open trailing value, if any."""
     params: dict[str, object] = {}
 
     for match in _PARAM_RE.finditer(raw_args):
-        name = (
-            match.group("dq_name")
-            or match.group("sq_name")
-            or match.group("bare_name")
-            or ""
-        ).strip()
+        name = _param_name(match)
         if not name:
             continue
         # Keep the value verbatim (like the glm47 converter): it is inline
         # between `>` and `</parameter>`, so surrounding whitespace is data.
         params[name] = match.group("value")
 
+    open_match = None
     if partial:
         remaining = _PARAM_RE.sub("", raw_args)
         match = _PARTIAL_PARAM_RE.search(remaining)
         if match:
-            name = (
-                match.group("dq_name")
-                or match.group("sq_name")
-                or match.group("bare_name")
-                or ""
-            ).strip()
+            name = _param_name(match)
             if name:
                 # Verbatim, same as the complete-match loop above.
                 params[name] = match.group("value")
+                open_match = match
 
-    return json.dumps(params, ensure_ascii=False)
+    return params, open_match
+
+
+def _minimax_m2_arg_converter(raw_args: str, partial: bool) -> str:
+    return json.dumps(_minimax_m2_parse_args(raw_args, partial)[0], ensure_ascii=False)
+
+
+def _value_end_window(text: str) -> str:
+    """Return the suffix of *text* that may still grow into a value end."""
+    start = text.rfind("<")
+    if start == -1:
+        return ""
+    tail = text[start:]
+    return tail if _PARAM_VALUE_END_RE.fullmatch(tail, partial=True) else ""
+
+
+def _minimax_m2_arg_value_extender(raw_args: str) -> ArgValueExtender | None:
+    # The open value runs to the end of raw_args and is the last JSON value,
+    # so a delta appends its JSON-escaped text unless it closes the value.
+    params, match = _minimax_m2_parse_args(raw_args, partial=True)
+    if match is None or next(reversed(params)) != _param_name(match):
+        return None
+    # Later text could still re-parse the opener: an earlier unmatched one, or
+    # a bare name that starts a quote, may close into a different name.
+    bare_name = match.group("bare_name")
+    if _PARAM_VALUE_END_RE.search(match.string, 0, match.start()) or (
+        bare_name is not None and bare_name[0] in "\"'"
+    ):
+        return None
+
+    window = _value_end_window(match.group("value"))
+
+    def extend(delta: str) -> str | None:
+        nonlocal window
+        text = window + delta
+        if _PARAM_VALUE_END_RE.search(text):
+            return None
+        window = _value_end_window(text)
+        return json.dumps(delta, ensure_ascii=False)[1:-1]
+
+    return extend
 
 
 @functools.cache
@@ -191,6 +240,7 @@ def minimax_m2_config() -> ParserEngineConfig:
             },
         },
         arg_converter=_minimax_m2_arg_converter,
+        arg_value_extender=_minimax_m2_arg_value_extender,
         stream_arg_deltas=True,
         tool_args_json=False,
         validate_tool_names=True,