Skip to content

[rollout, perf] fix: cache the generation-prompt delta in incremental tokenization - #7619

Closed
ruiling-smartbear wants to merge 2 commits into
verl-project:mainfrom
ruiling-smartbear:fix/bounded-generation-prompt-delta
Closed

ruiling-smartbear wants to merge 2 commits into
verl-project:mainfrom
ruiling-smartbear:fix/bounded-generation-prompt-delta

Conversation

@ruiling-smartbear

@ruiling-smartbear ruiling-smartbear commented Aug 30, 2026

Copy link
Copy Markdown
Contributor

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 that add_generation_prompt=True adds. 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

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).
  • Real tokenizers, this branch vs main, 100-turn tool loops with 800-char tool results through the public incremental API — token ids identical on every turn:
builder model main this PR last-turn cost
qwen3 Qwen3-8B 6.31s 0.19s 147 ms -> 1.5 ms
qwen35 Qwen3.5-9B 6.16s 0.25s 117 ms -> 1.9 ms
glm47 GLM-4.7 4.50s 0.30s 91 ms -> 1.4 ms
minimaxm2 MiniMax-M2 4.94s 0.31s 98 ms -> 1.4 ms
gptoss gpt-oss-20b 5.55s 0.22s 111 ms -> 0.3 ms
deepseek DeepSeek-V3 4.80s 0.22s 96 ms -> 0.7 ms
default Qwen3-8B 6.14s 0.35s 112 ms -> 1.6 ms

Tokenization only, CPU. Rollout-level numbers are the issue's, not mine.

  • The cache key has to include the final role: DeepSeek-V3.2-Exp emits <|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; Gemma4ContinuousTokenBuilder keeps its existing last-message render there.
  • cache_generation_prompt_delta class flag: VLContinuousTokenMixin sets it to False, so processor-backed builders keep today's behaviour until someone checks them.

Checklist Before Submitting

  • Read the Contribute Guide.
  • ruff check / ruff format pass on the touched files.
  • Docs: no user-facing change (internal tokenizer utility).
  • Unit tests added to tests/utils/test_continuous_token_on_cpu.py, covered by the CPU unit-test workflow.
  • CI request in ci-request once ready.

@CLAassistant

CLAassistant commented Aug 30, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

… 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
ruiling-smartbear force-pushed the fix/bounded-generation-prompt-delta branch from 57c93d9 to 1c59f25 Compare August 31, 2026 04:02
@ruiling-smartbear ruiling-smartbear changed the title [rollout, perf] fix: bound the generation-prompt delta render in incremental tokenization [rollout, perf] fix: cache the generation-prompt delta in incremental tokenization Aug 31, 2026
…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.
@wuxibin89

Copy link
Copy Markdown
Collaborator

Fixed in #7628 (comment)

@wuxibin89 wuxibin89 closed this Aug 31, 2026
@ruiling-smartbear

Copy link
Copy Markdown
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.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

# [rollout] perf: _tokenize_generation_prompt_delta re-tokenizes the full conversation after every tool turn

3 participants