fix(vllm_model): prefer inline token metadata over prompt tokenization

[ananthsub]

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 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 is superseded by this change.
• Gym #2576 is superseded by this change.
• NeMo RL #3390 produces the message-level token bundle consumed by the first source.
• NeMo RL #3581 reuses replayed token metadata across turns.
• Gym #1784 can add the Dynamo nvext.engine_data source independently.
• Gym #2324 can retain its video-specific schema, conversion, and adapter behavior independently.

[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.

[cmunley1]

lgtm if its working in benchmarks and training validation

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.