From 4753ac2e28a4053a2c7039c1f8956454fd94f106 Mon Sep 17 00:00:00 2001 From: sovereignbrah Date: Sun, 30 Aug 2026 19:13:34 -0700 Subject: [PATCH] fix(template): emit Reasoning Effort unconditionally to keep the prefix cache The head line is currently gated on `thinking_enabled`, so toggling thinking adds or removes `<|system|>Reasoning Effort: X` at roughly token 2. vLLM chains prefix-cache block hashes forward from token 0 over token ids, so a divergence that early invalidates the entire prompt, and MLA KpoolTailManager only counts block-aligned hits, so there is no partial credit either. Rendering the real template on a 4-message conversation, thinking on vs off: tools=no shared prefix 22 / 219 chars tools=yes shared prefix 22 / 855 chars Agent traffic always carries tools and a system prompt, so in practice every thinking toggle is a full re-prefill. Dropping the `thinking_enabled` guard makes both shapes share an identical head; the off-shape becomes a strict extension of the on-shape, differing only by the trailing ``. After the change both cases share 100% of the on-prompt. Measured on 2x GB10 (TP=2, EXL3 4bpw, DFlash2 k=7, 1M context) with a 21,642-token prompt and max_tokens=1: cold, first sight 26.54 s warm, same effort 0.62 s toggle high -> none 0.42 s toggle none -> high 0.62 s toggle high -> low 0.60 s Every toggle now hits warm; previously the third row cost a full re-prefill. This also restores the upstream zai-org/GLM-5.3-Flash behaviour, whose template emits the line unconditionally. Thinking is still disabled by the closed `` block, so the control is unchanged; only its cache cost is. Two existing tests asserted the line was absent when thinking was off, which encoded the divergence. They now assert the head is stable and the tail carries the toggle. Adds prefix-stability tests with and without tools, and registers a `tojson` filter accepting `ensure_ascii` so the template renders under bare Jinja2 the way vLLM's renderer supplies it. Co-Authored-By: Claude Opus 5 --- files/chat_template.jinja | 2 +- tests/test_chat_template.py | 100 ++++++++++++++++++++++++++++++++++-- 2 files changed, 96 insertions(+), 6 deletions(-) diff --git a/files/chat_template.jinja b/files/chat_template.jinja index 83c1083e..8847fa99 100644 --- a/files/chat_template.jinja +++ b/files/chat_template.jinja @@ -5,7 +5,7 @@ {%- set thinking_enabled = true -%} {%- endif -%} {%- set effective_reasoning_effort = reasoning_effort if reasoning_effort is defined and reasoning_effort in ['low', 'high'] else 'max' -%} -{%- if thinking_enabled and effective_reasoning_effort is not none -%}<|system|>Reasoning Effort: {{ effective_reasoning_effort | capitalize }}{%- endif -%} +{%- if effective_reasoning_effort is not none -%}<|system|>Reasoning Effort: {{ effective_reasoning_effort | capitalize }}{%- endif -%} {%- set clear_thinking = clear_thinking if clear_thinking is defined else false -%} {%- if tools -%} {%- macro tool_to_json(tool) -%} diff --git a/tests/test_chat_template.py b/tests/test_chat_template.py index f387d16a..427f1c42 100644 --- a/tests/test_chat_template.py +++ b/tests/test_chat_template.py @@ -1,19 +1,52 @@ #!/usr/bin/env python3 """Regression checks for GLM-5.3 chat-template reasoning controls.""" +import json import unittest from pathlib import Path from jinja2 import Environment +def _tojson(value, ensure_ascii=False, indent=None, **_kwargs): + """vLLM's renderer supplies a tojson accepting ensure_ascii; bare Jinja2 + does not. Register the same shape so the template renders standalone.""" + return json.dumps(value, ensure_ascii=ensure_ascii, indent=indent) + + +def _environment() -> Environment: + env = Environment(extensions=["jinja2.ext.loopcontrols"]) + env.filters["tojson"] = _tojson + return env + + TEMPLATE = Path(__file__).parents[1] / "files" / "chat_template.jinja" +TOOLS = [ + { + "type": "function", + "function": { + "name": "multiply", + "description": "Multiply two numbers", + "parameters": { + "type": "object", + "properties": {"a": {"type": "number"}, "b": {"type": "number"}}, + "required": ["a", "b"], + }, + }, + } +] + +CONVERSATION = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "Hi. How can I help?"}, + {"role": "user", "content": "and 3+3?"}, +] + def render_generation_prompt(**kwargs: object) -> str: - template = Environment(extensions=["jinja2.ext.loopcontrols"]).from_string( - TEMPLATE.read_text() - ) + template = _environment().from_string(TEMPLATE.read_text()) return template.render( messages=[{"role": "user", "content": "hello"}], tools=None, @@ -22,6 +55,11 @@ def render_generation_prompt(**kwargs: object) -> str: ) +def render_conversation(**kwargs: object) -> str: + template = _environment().from_string(TEMPLATE.read_text()) + return template.render(add_generation_prompt=True, **kwargs) + + class ChatTemplateTests(unittest.TestCase): def test_thinking_defaults_on(self) -> None: rendered = render_generation_prompt() @@ -30,12 +68,14 @@ def test_thinking_defaults_on(self) -> None: def test_thinking_can_be_disabled(self) -> None: rendered = render_generation_prompt(enable_thinking=False) - self.assertNotIn("Reasoning Effort:", rendered) + # The head line stays: it is what keeps the cached prefix stable across + # a thinking toggle. Thinking is disabled by the closed block instead. + self.assertIn("<|system|>Reasoning Effort: Max", rendered) self.assertTrue(rendered.endswith("<|assistant|>"), rendered) def test_thinking_alias_matches_parser_behavior(self) -> None: rendered = render_generation_prompt(thinking=False) - self.assertNotIn("Reasoning Effort:", rendered) + self.assertIn("<|system|>Reasoning Effort: Max", rendered) self.assertTrue(rendered.endswith("<|assistant|>"), rendered) def test_explicit_thinking_preserves_reasoning_effort(self) -> None: @@ -47,5 +87,55 @@ def test_explicit_thinking_preserves_reasoning_effort(self) -> None: self.assertTrue(rendered.endswith("<|assistant|>"), rendered) +class PrefixStabilityTests(unittest.TestCase): + """Toggling thinking must not change the prompt before the final token. + + vLLM chains prefix-cache block hashes forward from token 0, so any + divergence near the head invalidates the whole prompt. The off-shape must + therefore be a strict extension of the on-shape. + """ + + def _assert_strict_extension(self, tools) -> None: + on = render_conversation( + messages=CONVERSATION, tools=tools, enable_thinking=True + ) + off = render_conversation( + messages=CONVERSATION, tools=tools, enable_thinking=False + ) + self.assertTrue( + off.startswith(on), + "thinking-off prompt must extend thinking-on prompt, but they " + f"diverge at char {len(_common_prefix(on, off))} of {len(on)}", + ) + self.assertEqual(off[len(on) :], "") + + def test_toggle_is_prefix_stable_without_tools(self) -> None: + self._assert_strict_extension(None) + + def test_toggle_is_prefix_stable_with_tools(self) -> None: + # Agent traffic always carries tools; the tools block renders after the + # reasoning-effort line, so this is the case that actually regressed. + self._assert_strict_extension(TOOLS) + + def test_effort_levels_share_the_prompt_up_to_the_effort_word(self) -> None: + low = render_conversation( + messages=CONVERSATION, tools=TOOLS, enable_thinking=True, + reasoning_effort="low", + ) + high = render_conversation( + messages=CONVERSATION, tools=TOOLS, enable_thinking=True, + reasoning_effort="high", + ) + self.assertEqual(len(_common_prefix(low, high)), len("[gMASK]<|system|>Reasoning Effort: ")) + + +def _common_prefix(a: str, b: str) -> str: + limit = min(len(a), len(b)) + index = 0 + while index < limit and a[index] == b[index]: + index += 1 + return a[:index] + + if __name__ == "__main__": unittest.main()