chat template: speed up tool-result validation and reordering

#31

The tool result block checks ids and reorders results with nested loops over the block: has_dup_tool_result_id
and tc_id_exists per result, the pairwise tool_calls id check, then one scan of the block per tool call. That is
about 4 x N^2 macro calls for N parallel tool calls in one turn, so render time grows quadratically: 7.8 s at 800
parallel calls on an M4 and about 13 s on a Sapphire Rapids core. Past about 1,200 calls it exceeds vLLM's default
30 s render timeout, and the abandoned render thread keeps the API server's CPU busy, so other requests time out
too.

This change collects the result ids and the tool call ids once, then runs the same checks with
unique(case_sensitive=true) and in, and finds each result with index. Same checks, same order, same fallback
when ids are missing or duplicated.

Rendered output is byte-identical. 15,000 randomized conversations render the same through the current and the new
template (transformers 5.19.0, jinja2 3.1.6): duplicate, missing, case-variant and integer ids; shuffled, missing
and extra results; list-of-outputs tool messages; tool_reference outputs; tools and thinking on and off.

Parallel tool calls Before After
100 0.13 s 0.004 s
400 1.95 s 0.015 s
800 7.81 s 0.034 s
1,200 - 0.057 s

(M4, one core.) has_dup_tool_result_id and tc_id_exists are left defined but unused, to keep the diff to the
one block.

vMaroon changed pull request title from chat template: linear tool result id checks and reordering to chat template: speed up tool-result validation and reordering

Hi, @vMaroon . Thanks for the optimization! The performance improvement looks great.

One potential compatibility concern: the new implementation relies on Python's list.index() method:

{%- set p = ns_chk.res.index(tc_id) -%}

While this works with Python Jinja2, other Jinja-compatible renderers, such as Rust MiniJinja (used by NVIDIA Dynamo), may not support list.index() even with their Python compatibility extensions. This could cause template rendering failures when processing tool responses.

Could we avoid relying on this Python-specific method and use a more portable Jinja-compatible approach instead? For example, we could locate the index through an explicit loop:

{%- set ns_idx = namespace(value=-1) -%}
{%- for i in range(ns_chk.res | length) -%}
    {%- if ns_idx.value == -1 and ns_chk.res[i] == tc_id -%}
        {%- set ns_idx.value = i -%}
    {%- endif -%}
{%- endfor -%}
{%- set p = ns_idx.value -%}

It would also be helpful to add a cross-renderer regression test covering parallel tool calls and their corresponding responses (No need in this PR. Testing by yourself).

This should help preserve the optimization while avoiding potential compatibility regressions across different inference engines. Thanks!

Confirmed: under Dynamo's renderer the PR head fails with unknown method: iterator has no method named index in every conversation with a tool block it reorders. Updated the PR (ffc0d7c6) so it no longer needs index.

The explicit loop renders correctly but keeps one scan of the results per tool call, so the quadratic cost comes back, and under MiniJinja it ends up slower than current main: 2.9 s against 1.3 s at 800 parallel calls, 13.9 s against 3.0 s at 1,200. The update pairs ids with a sort-merge instead. It tags the result ids and the tool call ids with their positions, sorts both with sort(attribute='id', case_sensitive=true), walks both lists once to pair each tool call with its result, and emits the pairs with sort(attribute='i') so results keep tool call order. The checks before it are unchanged, so in this branch ids are non-empty unique strings and every result id matches a tool call id, which is all the merge relies on. The block only uses sort, unique, in, length and namespace.

Cross-renderer test, through Dynamo's own formatter (dynamo-renderer 8.0.1 with minijinja and minijinja-contrib 2.24.0, the versions in Dynamo main's Cargo.lock) and through transformers 5.19.0 with jinja2 3.1.6:

  • 15,000 randomized conversations and 9,000 targeted parallel-call blocks (case-variant, prefix and non-ASCII ids; shuffled and missing results; per-message and list-of-outputs results). Under Dynamo the updated template matches current main in all 24,000, including the same error where main raises. Under jinja2 it matches in 15,000 random and 3,000 targeted.
  • The previous PR head fails under Dynamo in every case with a reorderable block: 4,536 of the 8,000 cases of one seed.
Parallel tool calls main, jinja2 main, Dynamo explicit loop, Dynamo this update, jinja2 this update, Dynamo
100 0.11 s 0.034 s 0.006 s 0.005 s 0.002 s
400 1.64 s 0.31 s 0.23 s 0.019 s 0.011 s
800 6.56 s 1.26 s 2.94 s 0.040 s 0.047 s
1,200 15.4 s 3.01 s 13.9 s 0.065 s 0.14 s

(M4, one core, results in call order.) Under MiniJinja the list appends (ns.x = ns.x + [...]) start to show past a few thousand calls: 1.2 s at 2,400 and 3.7 s at 4,800.

Two things in main came up in the test that this PR leaves as they are; happy to send them separately:

  • A mapping tool result ("content": {"k": "v"}) renders as {'k': 'v'} in jinja2 and {"k": "v"} in Dynamo.
  • tool_to_json and tool_references_to_response are defined inside {% if tools %}, so a tool_reference result in a request without tools raises in both renderers.

Thanks for the quick fix and the detailed cross-renderer validation!

The sort-merge approach looks much better, especially since it avoids the list.index() compatibility issue without reintroducing the quadratic matching overhead.

The additional Dynamo and Jinja2 regression tests are reassuring.

Thanks again for addressing this so thoroughly!

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

Sign up or log in to comment