Skip to content

fix(template): emit Reasoning Effort unconditionally to keep the prefix cache - #63

Merged
plotarmordev merged 1 commit into
MiaAI-Lab:mainfrom
sovereignbrah:fix/reasoning-effort-head-prefix-cache
Sep 1, 2026
Merged

plotarmordev merged 1 commit into
MiaAI-Lab:mainfrom
sovereignbrah:fix/reasoning-effort-head-prefix-cache

Conversation

@sovereignbrah

Copy link
Copy Markdown
Contributor

Problem

The reasoning-effort head line is gated on thinking_enabled:

{%- if thinking_enabled and effective_reasoning_effort is not none -%}<|system|>Reasoning Effort: {{ effective_reasoning_effort | capitalize }}{%- endif -%}

It renders at roughly token 2, immediately after [gMASK]<sop> and before the tools block. So toggling thinking adds or removes ~7 tokens at the very head of the prompt.

vLLM chains prefix-cache block hashes forward from token 0 over token ids, so a divergence that early invalidates the entire prompt. On this kit the MLA KpoolTailManager only counts block-aligned hits, so there is no partial credit either — it is a total miss, not a degraded one.

Rendering the current template on a 4-message conversation with a system prompt, thinking on vs off:

shared prefix
without tools 22 / 219 chars
with tools 22 / 855 chars

Agent traffic always carries a system prompt and tools, so in practice every thinking toggle is a full re-prefill.

Change

Drop the thinking_enabled guard so the line is always emitted. Both shapes then share an identical head, and the thinking-off prompt becomes a strict extension of the thinking-on prompt, differing only by the trailing </think>. After the change both cases above share 100% of the on-prompt.

This also restores the behaviour of the upstream zai-org/GLM-5.3-Flash template, which emits this line unconditionally — the guard is a local addition here.

The control itself is unchanged: thinking is still disabled by the closed <think></think> block in the generation prompt. Only its cache cost changes. With thinking off the prompt still reads Reasoning Effort: Max, which is inert — the model cannot reason through an already-closed block.

Measured

2× GB10, TP=2, EXL3 4bpw, DFlash2 k=7, 1M context, prefix caching on. 21,642-token prompt, max_tokens=1:

step effort time
cold, first sight high 26.54 s
warm, same effort high 0.62 s
toggle high → none none 0.42 s
toggle none → high high 0.62 s
toggle high → low low 0.60 s

Every toggle now hits warm. Before the change, row 3 cost a full re-prefill.

Tests

Two existing tests in tests/test_chat_template.py asserted the line was absent when thinking was off, which encoded the divergence. They now assert the head stays stable while the tail carries the toggle.

Added PrefixStabilityTests: the off-shape must be a strict extension of the on-shape, with and without tools. The with-tools case is the one that actually regressed, since the tools block renders after the effort line.

Also registers a tojson filter accepting ensure_ascii in the test environment, so the template renders under bare Jinja2 the way vLLM's renderer supplies it. Without it these tests error out on the tools path.

$ python3 -m unittest tests.test_chat_template
Ran 7 tests — OK

Known limitation, not addressed here

Changing effort level (low ↔ high) still rewrites the head, since the level word is in that line. Making that cache-stable too would mean moving the instruction to the last user turn, the way NVIDIA's Nemotron 3 templates do — a larger change with a real question attached about whether effort adherence survives the move, given the model was trained with it at the head. Deliberately out of scope for this PR.

🤖 Generated with Claude Code

…ix 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 `</think>`. 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
`<think></think>` 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 <noreply@anthropic.com>

@plotarmordev plotarmordev left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exact-head validation passed: clean merge onto post-PR86 main, 7 focused template tests, deterministic byte-prefix checks, and real GLM-5.3 tokenizer verification with and without tools. Thinking-off is exactly thinking-on plus one token; parser semantics are unchanged and the effort-line behavior restores upstream template parity. Non-blocking documentation follow-up: README's statement that thinking-off omits the reasoning-effort hint becomes stale and should be corrected.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

reviewed Our source review or technical verdict is recorded; not merge approval or runtime qualification.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants