Repository navigation
perf(vllm): return token metadata with chat completions - #3390
Conversation
|
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. |
|
/ok to test 5fc345f |
jthomson04
left a comment
jthomson04
left a comment
There was a problem hiding this comment.
Nice change — the restructuring that decouples routed-experts attachment from token attachment is a clear improvement, and attach_token_information_to_chat_response_choices mirrors its routed-experts analog closely.
One blocking-ish issue (the gate on return_tokens_as_token_ids makes the feature a silent no-op for a request shape Gym itself supports), plus a new 500 path, some missing negative-test coverage, and a nit. Details inline.
One cross-repo consequence worth calling out here since it isn't visible in this diff: at the pinned Gym commit 473f446, choice_dict.pop("logprobs") lives inside if self.config.return_token_id_information and "prompt_token_ids" not in choice_dict["message"]: (responses_api_models/vllm_model/app.py:446, :499). Now that the server always supplies prompt_token_ids, that block is skipped entirely, so the duplicated per-token logprobs payload (token string + logprob + bytes + top_logprobs) stays on every returned message and flows into logged trajectories — the data the comment there calls "duplicated information". It doesn't break validation (NeMoGymChoice inherits Choice, which declares logprobs), but it partially offsets the payload win. Worth a companion Gym change, or at least a note in the PR description.
|
/ok to test c312907 |
|
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. |
|
/ok to test e5fa780 |
e5fa780 to
a7d20e9
Compare
|
/ok to test a7d20e9 |
a7d20e9 to
8e67f82
Compare
|
/ok to test 8e67f82 |
There was a problem hiding this comment.
Overall LGTM on the code — self-review pass already resolved the substantive issues (server-default fallback for return_tokens_as_token_ids, HTTP 400 for top_logprobs=null instead of 500, native CompletionOutput.logprobs instead of the token_id: string parse, and comprehensive error-branch coverage on the new helper).
The 3 attached fields (prompt_token_ids, generation_token_ids, generation_log_probs) map to the existing Gym consumer contract at responses_converter.py:166 and :431. Until now those fields were synthesized upstream by Gym's vllm_model server (app.py:462-496) via a per-turn /tokenize round-trip + token_id: prefix parsing on choice.logprobs.content, guarded by a START TODO / END TODO block whose exit condition is exactly "NeMo-RL vLLM natively supports prompt token ids". This PR is that upgrade — the producer now attaches natively, enabling a follow-up Gym change to delete the TODO block.
[BLOCKER] perf gain requires the paired Gym change. The currently pinned Gym (473f446) unconditionally calls /tokenize and .update()-overwrites the natively-attached fields on message_dict, so this PR by itself is perf-neutral (negligible extra serialization — see breakdown below). The 14.6% / 392.1s / "zero /tokenize calls" numbers land only after app.py:462-496 is deleted on the Gym side and the submodule pin is bumped — currently the PR description reads as if this PR is standalone-shippable for perf.
Payload size breakdown (per response, ~1000-token prompt + ~1000-token generation, for context on the "extra serialization" above):
| Component | Size | Share of total |
|---|---|---|
| Plain text response (no logprobs, no metadata) | ~2 KB | 0.1% |
+ native choice.logprobs.content (1000 positions × ~50 B) |
+50 KB | 2.4% |
+ prompt_token_ids (1000 int, this PR) |
+5 KB | 0.2% |
+ generation_token_ids (1000 int, this PR) |
+5 KB | 0.2% |
+ generation_log_probs (1000 float, this PR) |
+10-15 KB | 0.6% |
+ routed_experts for Qwen3-30BA3B (48 MoE layers × topk=8 × int16, base64 envelope) |
+2000 KB | 96.5% |
| Total | ~2.1 MB | 100% |
The 3 new fields sum to ~20 KB — <1% of a Qwen3-30BA3B response and ~40% growth on top of a non-MoE baseline — both dwarfed by the per-turn /tokenize round-trip latency they replace.
Two nits attached separately.
|
hi @jthomson04 , nice feature!
for this one, want to confirm are you using a different version nemo-gym, or am I missing something? |
|
Thanks for checking this. The benchmark did use Gym At the pinned Gym revision if self.config.return_token_id_information and "prompt_token_ids" not in choice_dict["message"]:Both The standard I updated the description to disclose the benchmark Gym revision and make the pinned-Gym behavior explicit. No Gym PR or submodule bump is needed. |
|
/ok to test 9d9676f |
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
9d9676f to
60c006d
Compare
|
/ok to test 60c006d |
yuki-97
left a comment
yuki-97
left a comment
There was a problem hiding this comment.
LGTM, thanks for the improvement and the clear clarify!
…probs, real logprobs - Strip trailing whitespace from content before asserting tool parser contracts on both backends; canonical expected value is now "before" - Add top_logprobs=0 to all three vLLM scenario request bodies to satisfy the NVIDIA-NeMo#3390 precondition check (VLLMValidationError when top_logprobs is None) - Replace _extract_generation_token_ids string-parsing with msg field reads; use actual generation_log_probs from the response instead of fabricated zeros Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
…probs, real logprobs - Strip trailing whitespace from content before asserting tool parser contracts on both backends; canonical expected value is now "before" - Add top_logprobs=0 to all three vLLM scenario request bodies to satisfy the #3390 precondition check (VLLMValidationError when top_logprobs is None) - Replace _extract_generation_token_ids string-parsing with msg field reads; use actual generation_log_probs from the response instead of fabricated zeros Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
#1558) ## Summary - Consume a complete token bundle from `choice.message` without making a second `/tokenize` request. - Support opt-in vLLM response token IDs for endpoints that return top-level `prompt_token_ids` and choice-level `token_ids`. - Retain `/tokenize` for integrations that return neither inline representation. - Reject partial or conflicting sources. Token capture supports one completion choice per request (`n` omitted or set to `1`). ## Why Gym needs the prompt token IDs, generated token IDs, and selected-token log probabilities used during generation. The previous adapter reconstructed prompt IDs through a second awaited `/tokenize` request. That request adds inter-turn latency and can reproduce a different prompt when prompt-affecting inputs are omitted. NeMo RL [#3390](NVIDIA-NeMo/RL#3390) attaches the token bundle to the assistant message. Gym reads that bundle directly and does not request duplicate response-level IDs by default. Gym cannot require that message shape from every training integration. Other integrations, including Verl, can request vLLM's response token IDs or continue using the compatibility fallback. ## Source precedence 1. A complete token bundle on `choice.message`. 2. Top-level `prompt_token_ids`, choice-level `token_ids`, and `choice.logprobs` from the vLLM response. 3. Generation metadata from `choice.logprobs` and prompt IDs from `/tokenize`. `request_prompt_and_generation_token_ids: true` enables the second source by sending `return_token_ids=true` to compatible vLLM endpoints. The option defaults to `false`. If both inline sources are present, their token IDs must agree. The selected training message is validated once to avoid repeated list scans and copies at long sequence lengths. The fallback forwards `model`, `messages`, `tools`, `chat_template_kwargs`, `mm_processor_kwargs`, and `required_prefix_token_ids`. This keeps fallback tokenization aligned with the generation prompt. `prompt_logprobs` is not requested. vLLM returns prompt IDs independently when `return_token_ids` is enabled, and Gym does not consume prompt-token log probabilities. ## Related work - Gym [#700](#700) is superseded by this change. - Gym [#2576](#2576) is superseded by this change. - NeMo RL [#3390](NVIDIA-NeMo/RL#3390) produces the message-level token bundle consumed by the first source. - NeMo RL [#3581](NVIDIA-NeMo/RL#3581) reuses replayed token metadata across turns. - Gym [#1784](#1784) can add the Dynamo `nvext.engine_data` source independently. - Gym [#2324](#2324) can retain its video-specific schema, conversion, and adapter behavior independently. --------- Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…probs, real logprobs - Strip trailing whitespace from content before asserting tool parser contracts on both backends; canonical expected value is now "before" - Add top_logprobs=0 to all three vLLM scenario request bodies to satisfy the #3390 precondition check (VLLMValidationError when top_logprobs is None) - Replace _extract_generation_token_ids string-parsing with msg field reads; use actual generation_log_probs from the response instead of fabricated zeros Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
…probs, real logprobs - Strip trailing whitespace from content before asserting tool parser contracts on both backends; canonical expected value is now "before" - Add top_logprobs=0 to all three vLLM scenario request bodies to satisfy the NVIDIA-NeMo#3390 precondition check (VLLMValidationError when top_logprobs is None) - Replace _extract_generation_token_ids string-parsing with msg field reads; use actual generation_log_probs from the response instead of fabricated zeros Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
…probs, real logprobs - Strip trailing whitespace from content before asserting tool parser contracts on both backends; canonical expected value is now "before" - Add top_logprobs=0 to all three vLLM scenario request bodies to satisfy the #3390 precondition check (VLLMValidationError when top_logprobs is None) - Replace _extract_generation_token_ids string-parsing with msg field reads; use actual generation_log_probs from the response instead of fabricated zeros Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
…probs, real logprobs - Strip trailing whitespace from content before asserting tool parser contracts on both backends; canonical expected value is now "before" - Add top_logprobs=0 to all three vLLM scenario request bodies to satisfy the NVIDIA-NeMo#3390 precondition check (VLLMValidationError when top_logprobs is None) - Replace _extract_generation_token_ids string-parsing with msg field reads; use actual generation_log_probs from the response instead of fabricated zeros Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Summary
This PR removes a redundant, awaited
/tokenizeround trip from every raw-vLLM model turn. Under a 64-concurrent multi-turn workload, that round trip was on the session critical path and accumulated into a significant model-serving bottleneck./tokenizecompatibility fallback unchangedPerformance impact
In a matched vLLM 0.23.0 benchmark using Gym
eddd5e98, this change:/tokenizeentirely across 18,078 successful model callsBecause Gym awaited
/tokenizebefore starting the next turn, this was not background bookkeeping overhead: its latency directly serialized multi-turn sessions. Returning metadata already available in vLLM's finalRequestOutputremoves that critical-path work without changing generation behavior.Root cause
The raw-vLLM backend already had the exact prompt and generated token IDs in the final vLLM
RequestOutput, but discarded them while serializing the chat-completion response. Gym therefore issued a separate awaited/tokenizerequest on every model turn to reconstruct prompt token IDs, adding inter-turn latency under concurrency.Gym compatibility
No paired Gym change or submodule bump is required. NeMo-RL's pinned Gym revision,
473f446, already guards the compatibility fallback on the absence of native metadata:Both
client.create_tokenize(...)and the subsequentmessage_dict.update(...)are inside that block. Because this PR serializesprompt_token_idsonchoice.message, pinned Gym skips the entire fallback. The fallback remains available for other servers that do not return native token metadata.Validation
/tokenizecallsThe whole-run first-request-to-last-response span was 3,290.9s because this run accumulated 999.6s of between-wave environment delay, dominated by one 1,008s OpenHands trajectory. The benchmark therefore supports model-serving parity, not environment-inclusive wall-time parity.
Known Gym payload behavior
Gym's native-metadata path skips its compatibility cleanup, so the standard OpenAI
choice.logprobsblock remains alongsidechoice.message.generation_log_probs. NeMo-RL intentionally preserveschoice.logprobsto maintain the endpoint's OpenAI response contract.