@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6607a80dab

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Bugfix][Gemma4] Fix offline parser truncation, adjust_request token leak, and chat template sync (#45553)

Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>

Luciano Martins committed Jun 16, 2026 at 01:31 UTC 6607a80dabfa03932515808895b016d2666b0a55

5 files changed +94 -74

examples/tool_chat_template_gemma4.jinja

+64 -40

@@ -116,7 +116,9 @@

116 }

117 {%- endmacro -%}

118 {%- macro format_argument(argument, escape_keys=True) -%}

119 - {%- if argument is string -%}

119 + {%- if argument is none -%}

120 + {{- 'null' -}}

121 + {%- elif argument is string -%}

122 {{- '<|"|>' + argument + '<|"|>' -}}

123 {%- elif argument is boolean -%}

124 {{- 'true' if argument else 'false' -}}

@@ -172,18 +174,21 @@

174 {{- '<tool_response|>' -}}

175 {%- endmacro -%}

176

175 -{%- set ns = namespace(prev_message_type=None) -%}

177 +{#- ===== SETUP ===== -#}

178 +{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}

179 {%- set loop_messages = messages -%}

180 +{%- set enable_thinking = enable_thinking | default(false) -%}

181 +{%- set preserve_thinking = preserve_thinking | default(false) -%}

182 {{- bos_token -}}

183 {#- Handle System/Tool Definitions Block -#}

179 -{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}

184 +{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}

185 {{- '<|turn>system\n' -}}

186 {#- Inject Thinking token at the very top of the FIRST system turn -#}

182 - {%- if enable_thinking is defined and enable_thinking -%}

187 + {%- if enable_thinking -%}

188 {{- '<|think|>\n' -}}

189 {%- set ns.prev_message_type = 'think' -%}

190 {%- endif -%}

186 - {%- if messages[0]['role'] in ['system', 'developer'] -%}

191 + {%- if messages and messages[0]['role'] in ['system', 'developer'] -%}

192 {%- if messages[0]['content'] is string -%}

193 {{- messages[0]['content'] | trim -}}

194 {%- elif messages[0]['content'] is sequence -%}

@@ -217,31 +222,24 @@

222 {%- if message['role'] != 'tool' -%}

223 {%- set ns.prev_message_type = None -%}

224 {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}

220 - {#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}

221 - {%- set prev_nt = namespace(role=None, found=false) -%}

222 - {%- if loop.index0 > 0 -%}

223 - {%- for j in range(loop.index0 - 1, -1, -1) -%}

224 - {%- if not prev_nt.found -%}

225 - {%- if loop_messages[j]['role'] != 'tool' -%}

226 - {%- set prev_nt.role = loop_messages[j]['role'] -%}

227 - {%- set prev_nt.found = true -%}

228 - {%- endif -%}

229 - {%- endif -%}

230 - {%- endfor -%}

231 - {%- endif -%}

232 - {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}

225 + {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#}

226 + {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}

227 {%- if not continue_same_model_turn -%}

228 {{- '<|turn>' + role + '\n' }}

229 + {%- if role == 'model' and not enable_thinking and not (message.get('reasoning') or message.get('reasoning_content')) -%}

230 + {{- '<|channel>thought\n<channel|>' -}}

231 + {%- endif -%}

232 {%- endif -%}

233

237 - {#- Render reasoning/reasoning_content as thinking channel -#}

234 + {#- Render reasoning/reasoning_content as thinking channel (tool-call turns only) -#}

235 {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}

239 - {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}

236 + {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%}

237 + {%- if thinking_text and thinking_gate and message.get('tool_calls') -%}

238 {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}

239 {%- endif -%}

240

243 - {%- if message['tool_calls'] -%}

244 - {%- for tool_call in message['tool_calls'] -%}

241 + {%- if message.get('tool_calls') -%}

242 + {%- for tool_call in message.get('tool_calls') -%}

243 {%- set function = tool_call['function'] -%}

244 {{- '<|tool_call>call:' + function['name'] + '{' -}}

245 {%- if function['arguments'] is mapping -%}

@@ -251,8 +249,13 @@

249 {%- set ns_args.found_first = true -%}

250 {{- key -}}:{{- format_argument(value, escape_keys=False) -}}

251 {%- endfor -%}

254 - {%- elif function['arguments'] is string -%}

255 - {{- function['arguments'] -}}

252 + {%- elif function['arguments'] is none -%}

253 + {%- else -%}

254 + {{- raise_exception(

255 + "chat_template: tool_calls[].function.arguments must be a "

256 + "JSON object (mapping), not a string. Deserialize arguments "

257 + "before passing to the template."

258 + ) -}}

259 {%- endif -%}

260 {{- '}<tool_call|>' -}}

261 {%- endfor -%}

@@ -262,7 +265,7 @@

265 {%- set ns_tr_out = namespace(flag=false) -%}

266 {%- if message.get('tool_responses') -%}

267 {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}

265 - {%- for tool_response in message['tool_responses'] -%}

268 + {%- for tool_response in message.get('tool_responses') -%}

269 {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}

270 {%- set ns_tr_out.flag = true -%}

271 {%- set ns.prev_message_type = 'tool_response' -%}

@@ -277,8 +280,8 @@

280 {%- else -%}

281 {%- set follow = loop_messages[k] -%}

282 {#- Resolve tool_call_id to function name -#}

280 - {%- set ns_tname = namespace(name=follow.get('name') | default('unknown', true)) -%}

281 - {%- for tc in message['tool_calls'] -%}

283 + {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}

284 + {%- for tc in message.get('tool_calls') -%}

285 {%- if tc.get('id') == follow.get('tool_call_id') -%}

286 {%- set ns_tname.name = tc['function']['name'] -%}

287 {%- endif -%}

@@ -296,9 +299,9 @@

299 {%- endfor -%}

300 {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}

301 {%- for part in tool_body -%}

299 - {%- if part.get('type') == 'image' -%}

302 + {%- if part.get('type') in ['image', 'image_url'] -%}

303 {{- '<|image|>' -}}

301 - {%- elif part.get('type') == 'audio' -%}

304 + {%- elif part.get('type') in ['audio', 'input_audio'] -%}

305 {{- '<|audio|>' -}}

306 {%- elif part.get('type') == 'video' -%}

307 {{- '<|video|>' -}}

@@ -314,29 +317,26 @@

317 {%- endif -%}

318

319 {%- set captured_content -%}

317 - {%- if message['content'] is string -%}

320 + {%- if message.get('content') is string -%}

321 {%- if role == 'model' -%}

322 {{- strip_thinking(message['content']) -}}

323 {%- else -%}

324 {{- message['content'] | trim -}}

325 {%- endif -%}

323 - {%- elif message['content'] is sequence -%}

326 + {%- elif message.get('content') is sequence -%}

327 {%- for item in message['content'] -%}

325 - {%- if item['type'] == 'text' -%}

328 + {%- if item.get('type') == 'text' -%}

329 {%- if role == 'model' -%}

330 {{- strip_thinking(item['text']) -}}

331 {%- else -%}

332 {{- item['text'] | trim -}}

333 {%- endif -%}

331 - {%- elif item['type'] == 'image' -%}

334 + {%- elif item.get('type') in ['image', 'image_url'] -%}

335 {{- '<|image|>' -}}

333 - {%- set ns.prev_message_type = 'image' -%}

334 - {%- elif item['type'] == 'audio' -%}

336 + {%- elif item.get('type') in ['audio', 'input_audio'] -%}

337 {{- '<|audio|>' -}}

336 - {%- set ns.prev_message_type = 'audio' -%}

337 - {%- elif item['type'] == 'video' -%}

338 + {%- elif item.get('type') == 'video' -%}

339 {{- '<|video|>' -}}

339 - {%- set ns.prev_message_type = 'video' -%}

340 {%- endif -%}

341 {%- endfor -%}

342 {%- endif -%}

@@ -345,19 +345,43 @@

345 {{- captured_content -}}

346 {%- set has_content = captured_content | trim | length > 0 -%}

347

348 + {#- Forward-scan: find next non-tool message role for continuation detection -#}

349 + {%- set next_nt = namespace(role=None, found=false) -%}

350 + {%- for j in range(loop.index0 + 1, loop_messages | length) -%}

351 + {%- if not next_nt.found -%}

352 + {%- if loop_messages[j]['role'] != 'tool' -%}

353 + {%- set next_nt.role = loop_messages[j]['role'] -%}

354 + {%- set next_nt.found = true -%}

355 + {%- endif -%}

356 + {%- endif -%}

357 + {%- endfor -%}

358 +

359 + {%- set continues_into_next = (

360 + role == 'model'

361 + and next_nt.role == 'assistant'

362 + and (not message.get('tool_calls') or ns_tr_out.flag)

363 + ) -%}

364 +

365 {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}

366 {{- '<|tool_response>' -}}

367 + {%- elif continues_into_next -%}

368 + {{- '\n' -}}

369 {%- elif not (ns_tr_out.flag and not has_content) -%}

370 {{- '<turn|>\n' -}}

371 {%- endif -%}

372 +

373 + {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}

374 + {%- set ns.prev_non_tool_role = message['role'] -%}

375 {%- endif -%}

376 {%- endfor -%}

377

378 {%- if add_generation_prompt -%}

379 {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}

380 {{- '<|turn>model\n' -}}

359 - {%- if not enable_thinking | default(false) -%}

381 + {%- if not enable_thinking -%}

382 {{- '<|channel>thought\n<channel|>' -}}

383 {%- endif -%}

384 + {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%}

385 + {{- '<|channel>thought\n' -}}

386 {%- endif -%}

363 -{%- endif -%}

\ No newline at end of file

387 +{%- endif -%}

tests/reasoning/test_gemma4_reasoning_parser.py

+1 -1

@@ -54,7 +54,7 @@ NO_REASONING = {

54 "output": "This is content",

55 "reasoning": None,

56 "content": "This is content",

57 - "is_reasoning_end": False,

57 + "is_reasoning_end": True,

58 }

59 REASONING_WITH_CHANNEL = {

60 "output": "<|channel>This is a reasoning section<channel|>This is the rest",

tests/renderers/test_gemma4_chat_template.py

+2 -2

@@ -358,7 +358,7 @@ class TestGemma4ChatTemplate:

358 "type": "function",

359 "function": {

360 "name": "download_image",

361 - "arguments": '{"url": "https://example.com/x.png"}',

361 + "arguments": {"url": "https://example.com/x.png"},

362 },

363 },

364 ],

@@ -392,7 +392,7 @@ class TestGemma4ChatTemplate:

392 "type": "function",

393 "function": {

394 "name": "process",

395 - "arguments": "{}",

395 + "arguments": {},

396 },

397 },

398 ],

vllm/parser/gemma4.py

+20 -3

@@ -423,6 +423,8 @@ class Gemma4Parser(ParserEngine):

423 tools: list[Tool] | None = None,

424 **kwargs,

425 ) -> None:

426 + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}

427 + self._thinking_enabled = chat_kwargs.get("enable_thinking", True)

428 super().__init__(

429 tokenizer,

430 tools,

@@ -437,6 +439,21 @@ class Gemma4Parser(ParserEngine):

439 self._prefix_stripped: bool = False

440 self._is_first_feed: bool = True

441

442 + def adjust_request(

443 + self,

444 + request: ChatCompletionRequest | ResponsesRequest,

445 + ) -> ChatCompletionRequest | ResponsesRequest:

446 + """Skip ``skip_special_tokens=False`` when thinking is disabled.

447 +

448 + When there are no reasoning channel tokens to preserve,

449 + keeping the default prevents tool-call delimiter tokens

450 + from leaking into content (e.g. with ``tool_choice="none"``).

451 + """

452 + chat_template_kwargs = getattr(request, "chat_template_kwargs", None) or {}

453 + if not chat_template_kwargs.get("enable_thinking", True):

454 + return request

455 + return super().adjust_request(request)

456 +

457 def _reset(self, initial_state=None) -> None:

458 super()._reset(initial_state=initial_state)

459 self._reasoning_text = ""

@@ -494,12 +511,12 @@ class Gemma4Parser(ParserEngine):

511 if tool_call_id is not None and tid == tool_call_id:

512 return True

513 if new_turn_id is not None and tid == new_turn_id:

497 - return False

514 + return not self._thinking_enabled

515 if tool_response_id is not None and tid == tool_response_id:

499 - return False

516 + return not self._thinking_enabled

517 if end_id is not None and tid == end_id:

518 return True

502 - return self._reasoning_ended

519 + return True

520

521 def _events_to_delta(

522 self,

vllm/tool_parsers/gemma4_utils.py

+7 -28

@@ -35,8 +35,6 @@ Ported from ``transformers.models.gemma4.utils_gemma4`` so that vLLM users

35 do not need a transformers dependency for output parsing.

36 """

38 -import json

39 -

38 import regex as re

40 # Tool call delimiter tokens as they appear in decoded text.

@@ -52,42 +50,23 @@ _ESCAPE_TOKEN = '<|"|>'

50 def _parse_tool_arguments(args_str: str) -> dict[str, str]:

51 """Parse tool call arguments from the Gemma4 compact format.

55 - Handles the ``key:<|"|>value<|"|>`` format used by Gemma4, with fallback

56 - to heuristic key-value extraction. Also tolerates the slightly different

57 - ``key: "value"`` format (space + plain quotes) that some chat templates

58 - produce.

53 + Delegates to the native ``<|"|>``-aware parser from

54 + ``vllm.parser.gemma4``, which handles internal quotes, nested

55 + objects, arrays, and all Gemma4 value types correctly.

57 Args:

58 args_str: Raw argument string from inside ``call:name{...}``.

60 Returns:

64 - Dictionary of argument name → value.

61 + Dictionary of argument name → string value.

62 """

63 if not args_str or not args_str.strip():

64 return {}

69 - # Replace Gemma4 escape tokens with standard quotes.

70 - cleaned = args_str.replace(_ESCAPE_TOKEN, '"')

71 -

72 - # Try JSON parsing first (handles nested values, arrays, etc.).

73 - try:

74 - parsed = json.loads("{" + cleaned + "}")

75 - # Ensure all values are strings for consistency.

76 - return {k: str(v) if not isinstance(v, str) else v for k, v in parsed.items()}

77 - except (json.JSONDecodeError, ValueError):

78 - pass

79 -

80 - # Fallback: extract key:"value" pairs (allow optional space after colon).

81 - arguments = {}

82 - for key, value in re.findall(r'(\w+):\s*"([^"]*)"', cleaned):

83 - arguments[key] = value

84 -

85 - if not arguments:

86 - # Last resort: extract key:value pairs (unquoted).

87 - for key, value in re.findall(r"(\w+):\s*([^,}]+)", args_str):

88 - arguments[key] = value.strip().strip('"').replace(_ESCAPE_TOKEN, "")

66 + from vllm.parser.gemma4 import _parse_gemma4_args

90 - return arguments

68 + parsed = _parse_gemma4_args(args_str)

69 + return {k: str(v) if not isinstance(v, str) else v for k, v in parsed.items()}

72 def parse_tool_calls(text: str, *, strict: bool = False) -> list[dict]: