@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9e25706560

[Bugfix] Tolerate misspelled DSML tool_calls wrapper (#56141)

Signed-off-by: sfeng33 <4florafeng@gmail.com>

Flora Feng committed Sep 10, 2026 at 03:04 UTC 9e2570656013a90cb601b391723af4bf9b3207d4
9 files changed +262 -14
tests/parser/engine/test_deepseek_v4.py
+162
@@ -30,6 +30,7 @@ from vllm.parser.deepseek_v4 import (
30 DSML_THINK_START,
31 DSML_TOOL_END,
32 DSML_TOOL_START,
33 + DSML_TOOL_START_VARIANTS,
34 DeepSeekV4Parser,
35 _dsml_arg_converter,
36 _unwrap_wrapper_args,
@@ -1229,3 +1230,164 @@ class TestDelegatingParserLargeDelta:
1230 assert eos_text not in output.reasoning
1231 assert output.content == ""
1232 assert output.tool_calls == []
1233 +
1234 +
1235 +# ── Malformed DSML markup must not leak into content (#51914) ─────────
1236 +
1237 +_WRAPPER_VARIANTS = list(DSML_TOOL_START_VARIANTS)
1238 +
1239 +
1240 +class TestMalformedDsmlNoise:
1241 + """DeepSeek-V4-Flash intermittently misspells the tool-call opener, e.g.
1242 + ``<｜DSML｜toolcalls>``. Known corrupted openers open a tool block like
1243 + the real wrapper so the markup never surfaces as assistant content.
1244 + """
1245 +
1246 + _INNER = (
1247 + f"{DSML_INVOKE_PREFIX}get_weather{DSML_INVOKE_NAME_END}\n"
1248 + f"{_param('city', 'true', 'Seoul')}\n"
1249 + f"{DSML_INVOKE_END}\n"
1250 + f"{DSML_TOOL_END}"
1251 + )
1252 +
1253 + @staticmethod
1254 + def _parser(mock_tokenizer, thinking: bool = False):
1255 + return DeepSeekV4Parser(
1256 + mock_tokenizer, chat_template_kwargs={"thinking": thinking}
1257 + )
1258 +
1259 + @pytest.mark.parametrize("wrapper", _WRAPPER_VARIANTS)
1260 + @pytest.mark.parametrize("preamble", ["", "Let me check.\n"])
1261 + def test_corrupted_wrapper_behaves_like_real_one(
1262 + self, mock_tokenizer, mock_request, wrapper, preamble
1263 + ):
1264 + text = f"{preamble}{wrapper}\n{self._INNER}"
1265 + reference = f"{preamble}{DSML_TOOL_START}\n{self._INNER}"
1266 +
1267 + result = self._parser(mock_tokenizer).extract_tool_calls(text, mock_request)
1268 + expected = self._parser(mock_tokenizer).extract_tool_calls(
1269 + reference, mock_request
1270 + )
1271 + assert [tc.function.name for tc in result.tool_calls] == ["get_weather"]
1272 + assert json.loads(result.tool_calls[0].function.arguments) == {"city": "Seoul"}
1273 + assert result.content == expected.content
1274 + assert "DSML" not in (result.content or "")
1275 +
1276 + results = simulate_tool_streaming(
1277 + self._parser(mock_tokenizer), mock_request, list(text)
1278 + )
1279 + assert collect_function_name(results) == "get_weather"
1280 + assert json.loads(collect_tool_arguments(results)) == {"city": "Seoul"}
1281 + assert collect_content(results) == (expected.content or "")
1282 +
1283 + @pytest.mark.parametrize("wrapper", _WRAPPER_VARIANTS)
1284 + def test_corrupted_wrapper_ends_reasoning(
1285 + self, mock_tokenizer, mock_request, wrapper
1286 + ):
1287 + """Like the real wrapper, a corrupted one inside ``<think>`` closes
1288 + the reasoning block and starts the tool call."""
1289 + text = f"thinking{wrapper}\n{self._INNER}"
1290 + parser = self._parser(mock_tokenizer, thinking=True)
1291 + reasoning, content = parser.extract_reasoning(text, mock_request)
1292 + assert reasoning == "thinking"
1293 + assert "DSML" not in (content or "")
1294 +
1295 + result = self._parser(mock_tokenizer, thinking=True).extract_tool_calls(
1296 + text, mock_request
1297 + )
1298 + assert [tc.function.name for tc in result.tool_calls] == ["get_weather"]
1299 +
1300 + @pytest.mark.parametrize(
1301 + "text",
1302 + [
1303 + # Fullwidth bars, DSML namespace, but not a registered spelling.
1304 + "see <｜DSML｜foo> and <｜DSML｜tool_call> here",
1305 + # Unclosed prefix.
1306 + "see <｜DSML｜ here\nnext line",
1307 + # ASCII pipes are not the DSML sigil.
1308 + "see <|DSML|tool> and <|DSML|toolcalls> there",
1309 + ],
1310 + )
1311 + def test_unregistered_markup_is_ordinary_content(
1312 + self, mock_tokenizer, mock_request, text
1313 + ):
1314 + """Only the registered spellings are treated as wrappers."""
1315 + result = self._parser(mock_tokenizer).extract_tool_calls(text, mock_request)
1316 + assert result.tools_called is False
1317 + assert result.content == text
1318 +
1319 + results = simulate_tool_streaming(
1320 + self._parser(mock_tokenizer), mock_request, list(text)
1321 + )
1322 + assert collect_content(results) == text
1323 +
1324 + @pytest.mark.parametrize("wrapper", _WRAPPER_VARIANTS)
1325 + def test_prose_mention_matches_canonical_wrapper(
1326 + self, mock_tokenizer, mock_request, wrapper
1327 + ):
1328 + """A registered spelling mentioned in prose behaves exactly like a
1329 + prose mention of the real wrapper: the opener starts a tool block and
1330 + the rest of the message is dropped. That is pre-existing behavior of
1331 + ``<｜DSML｜tool_calls>``; the variants must not differ from it either
1332 + way."""
1333 + template = "The opener looks like {} and then params follow. Done."
1334 + text = template.format(wrapper)
1335 + reference = template.format(DSML_TOOL_START)
1336 +
1337 + result = self._parser(mock_tokenizer).extract_tool_calls(text, mock_request)
1338 + expected = self._parser(mock_tokenizer).extract_tool_calls(
1339 + reference, mock_request
1340 + )
1341 + assert result.tools_called == expected.tools_called
1342 + assert result.content == expected.content
1343 +
1344 + results = simulate_tool_streaming(
1345 + self._parser(mock_tokenizer), mock_request, list(text)
1346 + )
1347 + expected_results = simulate_tool_streaming(
1348 + self._parser(mock_tokenizer), mock_request, list(reference)
1349 + )
1350 + assert collect_content(results) == collect_content(expected_results)
1351 +
1352 + @pytest.mark.parametrize("wrapper", _WRAPPER_VARIANTS)
1353 + @pytest.mark.parametrize("chunk_size", [1, 3, None], ids=lambda c: f"chunk={c}")
1354 + def test_delegating_parser_corrupted_wrapper(self, wrapper, chunk_size):
1355 + """The serving-layer shape: reasoning and tool adapters on separate
1356 + engines, with the corrupted opener arriving as plain text."""
1357 + tokens = _dsv4_tokens(
1358 + reasoning="Checking the weather.",
1359 + tool_name="get_weather",
1360 + params=[("city", "true", "Seoul")],
1361 + )
1362 + start_id = _DSV4_FULL_VOCAB[DSML_TOOL_START]
1363 + tokens = [
1364 + (900, wrapper) if tid == start_id else (tid, text) for tid, text in tokens
1365 + ]
1366 + # Only the think markers are single tokens in the real
1367 + # DeepSeek-V4-Flash vocab; the DSML wrappers are ordinary text and
1368 + # must be matched by the lexer. Mirror that here.
1369 + vocab = {
1370 + DSML_THINK_START: _DSV4_FULL_VOCAB[DSML_THINK_START],
1371 + DSML_THINK_END: _DSV4_FULL_VOCAB[DSML_THINK_END],
1372 + }
1373 + tokenizer = MockTokenizer(vocab=vocab, tokens=tokens)
1374 + parser = _DeepSeekV4Delegating(
1375 + tokenizer, chat_template_kwargs={"thinking": True}
1376 + )
1377 +
1378 + deltas = replay_streaming(
1379 + parser,
1380 + tokens,
1381 + chunk_size=chunk_size,
1382 + finished_on_last=True,
1383 + tools=DUMMY_TOOLS,
1384 + )
1385 + output = collect_output(deltas)
1386 +
1387 + assert "Checking" in output.reasoning
1388 + assert len(output.tool_calls) == 1, (
1389 + f"reasoning={output.reasoning!r}, content={output.content!r}"
1390 + )
1391 + assert output.tool_calls[0]["name"] == "get_weather"
1392 + assert json.loads(output.tool_calls[0]["arguments"]) == {"city": "Seoul"}
1393 + assert "DSML" not in (output.content or "")
tests/parser/engine/test_delegating_replay.py
+1 -1
@@ -197,7 +197,7 @@ def test_delegating_parse_tool_choice_none(parser_cls, parser_name, sample):
197 cfg = parser._tool_parser._parser_engine.parser_engine_config
198 terminals = sorted(
199 v
200 - for v in set(cfg.terminals.values()) | set(cfg.token_id_terminals.values())
200 + for v in cfg.terminal_literals | set(cfg.token_id_terminals.values())
201 if len(v) > 1
202 )
203 assert_no_terminal_leakage(
tests/parser/engine/test_engine.py
+58
@@ -953,3 +953,61 @@ class TestSkipToolParsingFromMessageHeader:
953 assert types[0] == EventType.TEXT_CHUNK
954 assert events[0].value == "<tool_call>"
955 assert engine._message_header_buffer == ""
956 +
957 +
958 +# ── Terminals with several spellings ──────────────────────────────────
959 +
960 +
961 +def _alias_config() -> ParserEngineConfig:
962 + """``TOOL_START`` has a canonical spelling and one corrupted alias."""
963 + return ParserEngineConfig(
964 + name="alias_test",
965 + terminals={"TOOL_START": ("<tc>", "<tcx>"), "TOOL_END": "</tc>"},
966 + transitions={
967 + (ParserState.CONTENT, "TOOL_START"): Transition(
968 + ParserState.TOOL_ARGS,
969 + (EventType.TOOL_CALL_START,),
970 + ),
971 + (ParserState.TOOL_ARGS, "TOOL_END"): Transition(
972 + ParserState.CONTENT,
973 + (EventType.TOOL_CALL_END,),
974 + ),
975 + },
976 + content_events={
977 + ParserState.CONTENT: EventType.TEXT_CHUNK,
978 + ParserState.TOOL_ARGS: EventType.ARG_VALUE_CHUNK,
979 + },
980 + tool_args_json=False,
981 + )
982 +
983 +
984 +class TestTerminalAliases:
985 + def test_terminal_defs_expand_every_spelling(self):
986 + defs = terminals_from_literals({"A": ("x", "y"), "B": "z"})
987 + assert [(d.name, d.literal) for d in defs] == [
988 + ("A", "x"),
989 + ("A", "y"),
990 + ("B", "z"),
991 + ]
992 +
993 + def test_terminal_literal_is_canonical_spelling(self):
994 + cfg = _alias_config()
995 + assert cfg.terminal_literal("TOOL_START") == "<tc>"
996 + assert cfg.terminal_literal("TOOL_END") == "</tc>"
997 + assert cfg.terminal_literal("MISSING") is None
998 + assert cfg.terminal_literals == {"<tc>", "<tcx>", "</tc>"}
999 +
1000 + @pytest.mark.parametrize("spelling", ["<tc>", "<tcx>"])
1001 + @pytest.mark.parametrize("char_by_char", [False, True])
1002 + def test_every_spelling_shares_transitions(self, spelling, char_by_char):
1003 + engine = StreamingParserEngine(_alias_config(), None)
1004 + text = f"{spelling}a</tc>"
1005 + events = []
1006 + for chunk in list(text) if char_by_char else [text]:
1007 + events.extend(engine.feed(chunk, []))
1008 + events.extend(engine.finish())
1009 + assert [e.type for e in events] == [
1010 + EventType.TOOL_CALL_START,
1011 + EventType.ARG_VALUE_CHUNK,
1012 + EventType.TOOL_CALL_END,
1013 + ]
tests/parser/engine/test_replay.py
+6 -5
@@ -83,7 +83,7 @@ def _discover_parsers() -> list[_ParserInfo]:
83 raise RuntimeError(
84 f"{obj.__name__} config missing 'TOOL_END' in token_id_terminals"
85 )
86 - all_vals = set(cfg.terminals.values()) | set(cfg.token_id_terminals.values())
86 + all_vals = cfg.terminal_literals | set(cfg.token_id_terminals.values())
87 found.append(
88 _ParserInfo(
89 parser_cls=obj,
@@ -91,12 +91,13 @@ def _discover_parsers() -> list[_ParserInfo]:
91 samples=build_samples(cfg.name),
92 terminals=sorted(v for v in all_vals if len(v) > 1),
93 tool_end=tool_end,
94 - think_end=cfg.terminals.get("THINK_END", ""),
94 + think_end=cfg.terminal_literal("THINK_END") or "",
95 tool_start=(
96 - cfg.terminals["TOOL_SECTION_START"]
96 + cfg.terminal_literal("TOOL_SECTION_START")
97 if (ParserState.CONTENT, "TOOL_SECTION_START") in cfg.transitions
98 - else cfg.terminals.get("TOOL_START", "")
99 - ),
98 + else cfg.terminal_literal("TOOL_START")
99 + )
100 + or "",
101 )
102 )
103 if missing_builders:
vllm/parser/deepseek_v4.py
+7 -1
@@ -50,6 +50,12 @@ DSML_INVOKE_END = f"</{_DSML}invoke>"
50 DSML_PARAM_START = f"<{_DSML}parameter"
51 DSML_PARAM_CLOSE = f"</{_DSML}parameter>"
52
53 +# Spellings variants of ``DSML_TOOL_START`` observed in production.
54 +DSML_TOOL_START_VARIANTS: tuple[str, ...] = (
55 + f"<{_DSML}toolcalls>",
56 + f"<{_DSML}tool>",
57 +)
58 +
59 _ESCAPED_DSML = re.escape(_DSML)
60 _PARAM_RE = re.compile(
61 rf'<{_ESCAPED_DSML}parameter\s+name="([^"]+)"\s+string="(true|false)">'
@@ -131,7 +137,7 @@ def deepseek_v4_config(thinking: bool = False) -> ParserEngineConfig:
137 terminals={
138 "THINK_START": DSML_THINK_START,
139 "THINK_END": DSML_THINK_END,
134 - "TOOL_START": DSML_TOOL_START,
140 + "TOOL_START": (DSML_TOOL_START, *DSML_TOOL_START_VARIANTS),
141 "TOOL_END": DSML_TOOL_END,
142 "INVOKE_PREFIX": DSML_INVOKE_PREFIX,
143 "INVOKE_NAME_END": DSML_INVOKE_NAME_END,
vllm/parser/engine/incremental_lexer.py
+6 -2
@@ -5,6 +5,7 @@ tokens, with prefix-match buffering for ambiguous boundaries."""
5
6 from __future__ import annotations
7
8 +from collections.abc import Mapping
9 from dataclasses import dataclass
10
11 import regex as re
@@ -284,7 +285,9 @@ class IncrementalLexer:
285 return n
286
287
287 -def terminals_from_literals(literals: dict[str, str]) -> list[TerminalDef]:
288 +def terminals_from_literals(
289 + literals: Mapping[str, str | tuple[str, ...]],
290 +) -> list[TerminalDef]:
291 return [
292 TerminalDef(
293 name=name,
@@ -292,5 +295,6 @@ def terminals_from_literals(literals: dict[str, str]) -> list[TerminalDef]:
295 is_literal=True,
296 literal=lit,
297 )
295 - for name, lit in literals.items()
298 + for name, value in literals.items()
299 + for lit in ((value,) if isinstance(value, str) else value)
300 ]
vllm/parser/engine/parser_engine.py
+2 -2
@@ -158,11 +158,11 @@ class ParserEngine(Parser):
158
159 @property
160 def reasoning_start_str(self) -> str | None:
161 - return self.parser_engine_config.terminals.get("THINK_START")
161 + return self.parser_engine_config.terminal_literal("THINK_START")
162
163 @property
164 def reasoning_end_str(self) -> str | None:
165 - return self.parser_engine_config.terminals.get("THINK_END")
165 + return self.parser_engine_config.terminal_literal("THINK_END")
166
167 @cached_property
168 def vocab(self) -> dict[str, int]:
vllm/parser/engine/parser_engine_config.py
+19 -2
@@ -17,7 +17,7 @@ Each model format is described by a :class:`ParserEngineConfig` that specifies:
17
18 from __future__ import annotations
19
20 -from collections.abc import Callable
20 +from collections.abc import Callable, Mapping
21 from dataclasses import dataclass, field
22 from enum import Enum, auto
23 from functools import cached_property
@@ -54,7 +54,8 @@ class ParserEngineConfig:
54
55 name: str
56
57 - terminals: dict[str, str] = field(default_factory=dict)
57 + # A terminal may carry several spellings; the first spelling is the canonical one.
58 + terminals: Mapping[str, str | tuple[str, ...]] = field(default_factory=dict)
59
60 token_id_terminals: dict[str, str] = field(default_factory=dict)
61
@@ -100,6 +101,22 @@ class ParserEngineConfig:
101 # Reject tool calls whose names are absent from the request tools.
102 validate_tool_names: bool = False
103
104 + def terminal_literal(self, name: str) -> str | None:
105 + """Canonical spelling of terminal *name*, or ``None`` if undeclared."""
106 + value = self.terminals.get(name)
107 + if isinstance(value, tuple):
108 + return value[0] if value else None
109 + return value
110 +
111 + @property
112 + def terminal_literals(self) -> frozenset[str]:
113 + """Every declared spelling across all terminals."""
114 + return frozenset(
115 + lit
116 + for value in self.terminals.values()
117 + for lit in ((value,) if isinstance(value, str) else value)
118 + )
119 +
120 @cached_property
121 def terminal_defs(self):
122 from vllm.parser.engine.incremental_lexer import terminals_from_literals
vllm/parser/engine/streaming_parser_engine.py
+1 -1
@@ -51,7 +51,7 @@ def _build_drop_info(
51
52 configured_texts = (
53 set(config.token_id_terminals.values())
54 - | set(config.terminals.values())
54 + | config.terminal_literals
55 | config.preserve_tokens
56 )
57