perf(vllm): avoid repeated multi-turn tokenization

[jthomson04]

What does this PR do ?

Avoid repeated full-history tokenization for supported multi-turn vLLM chat requests. When an earlier assistant message includes exact token metadata, NeMo-RL reuses that model prefix and tokenizes only the contextual suffix for the next turn.

The guarded path keeps vLLM chat parsing, template rendering, tool and reasoning adjustment, validation, and engine preparation. First turns and unsupported requests continue to use native vLLM preprocessing. This PR does not change the public API or require a Gym change.

Performance

A matched end-to-end validation used 320 rollouts per arm and approximately 18,000 model calls:

The baseline used commit 0ded4a9f510f55efb154d8005519d2b7e4ed463e. The optimized arm used commit 56bb18ece647029ec49ad477e2047ac934578cab.

Metric	Before	After	Reduction
Server preprocessing, mean	0.938 s	0.038 s	96.0%
Server preprocessing, p50	0.700 s	0.026 s	96.3%
Server preprocessing, p95	2.497 s	0.093 s	96.3%
Observed per-turn response latency, mean	3.984 s	3.327 s	16.5%
Observed per-turn response latency, p50	2.956 s	2.047 s	30.8%

The response-latency measurements are end-to-end model-call latencies observed by the agent. They include inference and request-processing time.

The validation used a Nemotron Nano v3.5 SWE checkpoint with:

• NemotronHForCausalLM
• PreTrainedTokenizerFast
• <|im_end|> as the EOS token
• enable_thinking=true
• vLLM 0.25.1
• chat-template SHA-256 82753bef5cedc4932c1ed509b5c9a12be680fd86d1adb65bc3f7398d11c8eebc

Issues

None.

Usage

No configuration change is required. NeMo-RL selects the optimization only when complete token metadata is available and all safety guards pass.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added focused unit tests for prompt-token construction, guards, fallbacks, and HTTP behavior.
• Ran the focused unit and end-to-end functional validation.
• No documentation update is required because there is no public API or configuration change.

Additional Information

Validation included:

• pytest -q tests/unit/models/generation/test_openai_server_utils.py
• pytest -q -s tests/unit/models/generation/test_vllm_generation.py::test_vllm_http_server_correct_merged_tokens_matches_baseline
• Exact native-versus-injected prompt comparison across 3,709 subsequent-turn requests, with zero mismatches.
• Full end-to-end validation with 320 valid rollouts per arm, zero injection failures, and zero unexpected fallbacks.
• Head-level shadow validation at commit 8790130218435a6d350d3c0b089f67660266e333 completed 64 valid rollouts. It compared 3,639 injected prompts with native vLLM preprocessing and found zero mismatches, zero injection failures, zero unexpected fallbacks, zero /tokenize calls, and zero HTTP errors.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[jthomson04]

/ok to test 56bb18e

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[jthomson04]

/ok to test ace567b

[jthomson04]

Reviewed with a team of agents (RL codebase, vLLM API, bug scan, tests, plus an adversarial pass), each verifying claims against the pinned vLLM v0.25.1 source rather than the local checkout. Nice result on the numbers — the 96% preprocessing reduction with an exact native-vs-injected comparison over 3,709 requests is the right correctness experiment, and the fail-safe posture is well built: nearly every hazard probed routes to CompleteTokenInjectionError and native preprocessing.

Verdict: this is worth landing. The win is measured, large, and on the right axis, and I went looking for a simpler design without finding one — rendering the truncated conversation to locate the boundary is what the fallback already does (two full renders, i.e. the cost being eliminated), and a prefix cache would need a conversation key Gym doesn't send. The marker trick buys one render plus a suffix-only encode while preserving context-sensitive template behavior after the spliced turn.

Nothing below is a blocker. The three I'd most want addressed are small and concrete: the missing empty-prefix guard, the untested load-bearing deepcopy, and exposing the injection counters. The rest is hardening, tests, and docs.

An adversarial pass downgraded several of my own initial findings, and I've rewritten those comments rather than leave them overstated:

• The empty-prefix and prefix/ordinal items are defensive hardening, not live bugs — I searched for a producer of empty token lists and found none, and the derived-prefix path is provably safe because upstream passes request.messages through untouched.
• My suggestion to wrap the fast path in a broad except Exception was wrong and is retracted in-thread; it conflicts with this repo's error-handling guidance and would make an upstream break undetectable. Two targeted fixes replace it.
• A whitespace-splice concern I raised was refuted: the behavior is byte-identical to replace_prefix_tokens, so it's consistent by design. Only a test rename and a clarifying comment remain.

Items recorded as verified clean, so they don't get re-litigated:

• The seven getattr(request, X, False) gates are all correct for v0.25.1. return_prompt_text is the only true tri-state (bool | None = None), but its only two consumers use plain truthiness with no server-level default, so None genuinely means off.
• isinstance(self.renderer, HfRenderer) is exact — v0.25.1 has zero HfRenderer subclasses, so MistralRenderer and friends are excluded and the hardcoded tokenize=False matches native. kimi_audio also maps to HfRenderer but is excluded by the multimodal gate.
• The mm_data/mm_uuids guard does not fire on text-only requests (resolve_items() returns (None, None)), so the optimization is live rather than dead code.
• The assistant_ordinal <-> parsed-conversation mapping holds: conversation.extend(sub_messages) looks like it could fan out, but _parse_chat_message_content_parts returns a single-element list on both branches, and _postprocess_messages never adds or removes messages.
• Max-model-len validation is not bypassed — apply_post_tokenization still runs _token_len_check for token-id prompts.

Two open questions rather than defects: whether the fast path engages for Gym's Responses-API agents given their list-shaped content, and the fast-vs-fallback boundary divergence. On the latter — the fallback selects the last assistant while the fast path selects the latest tokenized one, so for a conversation ending in an untokenized assistant they produce different prompts (the fallback drops an intervening user turn). That is pre-existing and out of scope here, and the new path is the correct one; I mention it only because it means a silent fallback can change the answer, which is why the unread counters matter.

Finally, the L0_Unit_Tests_Automodel failure is test_value_worker_train_decreases_loss[2gpu_dp2] — value-worker/DTensor, unrelated to this change.

Generated by Claude Code

[jthomson04]

/ok to test 706674b

[jthomson04]

/ok to test a05de87

[yuki-97]

thanks @jthomson04 , nice optimization!

[jthomson04]

/ok to test 71cbbcc

[jthomson04]

/ok to test 8eebc8e

[jthomson04]

/ok to test 8790130

[jthomson04]

/ok to test 3db4b3d