[rollout, perf] fix: cache the generation-prompt delta in incremental tokenization - #7619
Closed
ruiling-smartbear wants to merge 2 commits into
Closed
ruiling-smartbear wants to merge 2 commits into
ruiling-smartbear wants to merge 2 commits into
Conversation
4 tasks
… tokenization tokenize_non_assistant_incremental_messages rendered the full conversation twice after every appended tool/user/system group just to obtain the tokens that add_generation_prompt=True adds. On a 200-turn tool loop that is a quadratic amount of chat-template rendering and it dominated tokenization time (verl-project#7617). Cache that delta on the builder, keyed by the role of the final message and the tool schemas, filled by the existing full-history render. The first four uses of a key re-render and compare, then re-validation happens at powers of two; a disagreement marks the key unstable and that key renders in full for the rest of the rollout. Builders keep template-specific behaviour by overriding _render_generation_prompt_delta (Gemma4 does); VL builders opt out of the cache until their processor renders have been checked. Measured with real tokenizers against main: token ids identical for Qwen2.5, Qwen3, Qwen3.5, GLM-4.7, MiniMax-M2, gpt-oss and DeepSeek-V3 over 100-turn tool loops, with tokenization dropping from 4.5-6.3s to 0.2-0.35s.
ruiling-smartbear
force-pushed
the
fix/bounded-generation-prompt-delta
branch
from
August 31, 2026 04:02
57c93d9 to
1c59f25
Compare
…he cache Gemma4 keeps its own last-message render as the renderer hook; check that the cache serves exactly that (nothing after a tool response, the bounded render after a user turn), never a full-history render, and that later uses hit the cache.
Collaborator
|
Fixed in #7628 (comment) |
Contributor
Author
|
Sounds good, #7628 does the job. I ran the merge commit against its parent on the same trajectories (token ids identical on every row) and left the notes plus a DeepSeek caveat over there. Thanks for taking a look. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #7617.
tokenize_non_assistant_incremental_messages()rendered the full conversation twice after every appended tool/user/system group, only to obtain the tokens thatadd_generation_prompt=Trueadds. Over a long tool loop that is a quadratic amount of chat-template rendering; the issue measured it at ~34% of rollout wall time on a 200-turn Qwen3.5-9B agent loop.Following the discussion on the issue (@gxlvera, @wuxibin89): the generation-prompt delta is now cached on the builder by default, keyed by the role of the final message and the tool schemas, and filled by the existing full-history render. Builders keep template-specific behaviour through an override hook.
Checklist Before Starting
[{modules}] {type}: {description}Test
tests/utils/test_continuous_token_on_cpu.py: 128 passed (119 existing + 9 new: a Qwen tool loop no longer renders the full history every turn, the cached delta equals the rendered one turn by turn, a template whose generation prompt drifts is detected and falls back, drift after the warm-up is still caught, roles and tool schemas get separate entries, the default builder still renders the full history on first use, VL builders stay uncached, Gemma4's own last-message render is what the cache serves).main, 100-turn tool loops with 800-char tool results through the public incremental API — token ids identical on every turn:Tokenization only, CPU. Rollout-level numbers are the issue's, not mine.
<|Assistant|></think>after a user turn and nothing after a tool output. Keyed on (role, tools), every template measured is stable over 20 turns.API and Usage Example
No API change.
Design & Code Changes
ContinuousTokenBuilder._tokenize_generation_prompt_delta: cache keyed by(final message role, tools fingerprint). The first 4 uses of a key re-render and compare, then re-validation happens at powers of two (8, 16, 32, ...). One disagreement marks the key unstable and it renders in full for the rest of the rollout, with a warning. O(log n) full renders per rollout instead of O(n)._render_generation_prompt_delta: the uncached renderer, which is the override hook for template-specific behaviour. Base = full-history render;Gemma4ContinuousTokenBuilderkeeps its existing last-message render there.cache_generation_prompt_deltaclass flag:VLContinuousTokenMixinsets it toFalse, so processor-backed builders keep today's behaviour until someone checks them.Checklist Before Submitting
ruff check/ruff formatpass on the touched files.tests/utils/test_continuous_token_on_cpu.py, covered by the CPU unit-test workflow.ci-requestonce ready.