Make the Gemma 4 chat template render in linear time

#57

The chat template takes quadratic time in the number of messages. A long tool-calling conversation spends seconds in template rendering on every request, before the model does any work.

I noticed it while profiling a coding agent that makes 50 to 100 tool calls in one turn. Each round resends the whole conversation, and in llama.cpp the gap between the request arriving and the model's own timers starting grew with every round: 0.35 s at 25k tokens of context, 0.87 s at 39k and about 3.5 s at 64k. The prompt itself came from the KV cache. POST /apply-template on llama-server reproduces it with no model involved. On the July 9 template (llama.cpp build 11033, Vulkan, Gemma 4 E4B):

Messages /apply-template
90 254 ms
176 905 ms
350 3,108 ms
698 12,186 ms

Every doubling costs about four times as much. Qwen3-Coder's template on the same conversations takes 39 ms at 350 messages and 78 ms at 698, so llama.cpp's own message handling is fine and the template is the cause.

The cause is two forward scans inside the message loop. One looks for the next non-tool message (next_nt) and one collects the tool results that follow an assistant message with tool calls. Both walk every later message for each message, and Jinja has no way to stop a loop early, so each message costs O(n).

The patch finds the position of the next non-tool message for every message in one backward pass before the main loop. Both scans then use that position. The output is the same. I only changed how the position is found.

What I checked:

  1. llama.cpp (--jinja --chat-template-file): the original and the patched template render byte-identical prompts for 16 cases, with and without tool definitions: plain chat, consecutive assistant messages, a 5-call loop (also followed by a final answer and another user turn), parallel calls with two results, a call followed by two assistant messages, a conversation ending on a tool result, and reasoning content. compare-llamacpp.mjs does this.
  2. Jinja2 3.1.6 in ImmutableSandboxedEnvironment, which is what transformers uses: 36 cases (the same ones, plus an empty list, each with and without tools and with and without add_generation_prompt), no differences. check_jinja2.py does this. I used string concatenation and split to build the positions, because list.append is blocked in the sandbox.
  3. Time to render a tool loop with 10 tool definitions, in llama.cpp:
Tool calls (messages) Original Patched
50 (102) 323 ms 88 ms
100 (202) 1,120 ms 177 ms
200 (402) 4,131 ms 353 ms

The same loops in Jinja2 take 62 ms, 234 ms and 895 ms on the original for 202, 402 and 802 messages, and 11, 21 and 44 ms patched.

Doubling the messages now about doubles the time.

What I did not check: other Gemma 4 sizes (the file is the same text, but I only ran E4B), image and audio content parts in tool results, and the <|think|> paths beyond the reasoning case above. If you have a conformance suite for the template, please run it against this.

Files: chat_template.diff is the change against the July 9 template, chat_template.jinja the patched file, and the two scripts reproduce the numbers on any machine with llama-server or Python.

Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment