feat(vllm-model): consume native Dynamo token data

[jthomson04]

Summary

Move native Dynamo response-token handling into NeMo-Gym's shared vLLM token-selection path.

• Read prompt_token_ids, completion_token_ids, and completion_logprobs from nvext.engine_data.
• Normalize them to Gym's existing prompt_token_ids, generation_token_ids, and generation_log_probs bundle.
• Select token data in this order: assistant-message metadata, Dynamo engine data, standard vLLM token IDs, then choice.logprobs plus /tokenize.
• Derive required_prefix_token_ids from the latest tokenized assistant message unless the caller supplied a non-null list.
• Reject partial or malformed native data and duplicate token IDs that disagree.
• Preserve response fields, including choice.logprobs and message extensions such as routed_experts.

Motivation

NeMo-Gym needs the model's original token IDs and log probabilities for training. Dynamo already returns these values in nvext.engine_data. Consuming that native bundle in Gym removes duplicate response conversion from NeMo-RL and avoids an unsupported /tokenize fallback.

Gym owns response interpretation. The NeMo-RL wrapper continues to render exact prompt tokens, splice the prior prefix, send nvext.token_data, and request engine_data.

Compatibility

• Existing NeMo-RL wrappers that copy matching token metadata onto the assistant message remain supported.
• Native Dynamo data does not require choice.logprobs and remains authoritative if that optional field differs.
• Standard vLLM token IDs and the existing /tokenize fallback remain unchanged when native data is absent.
• No public schema or configuration flag is added.

Validation

• uv run --extra dev pytest responses_api_models/vllm_model/tests/test_app.py (132 passed)
• uv run --extra dev pre-commit run --files responses_api_models/vllm_model/app.py responses_api_models/vllm_model/tests/test_app.py

The final managed Dynamo smoke will run from NeMo-RL after this PR merges and NeMo-RL pins the exact containing Gym main commit.

Related

• Token-selection framework: fix(vllm_model): prefer inline token metadata over prompt tokenization #1558
• NeMo-RL cleanup: to follow after this PR merges

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

[jthomson04]

Code review

14 findings from a recall-oriented pass over the diff and the enclosing functions in responses_api_models/vllm_model/app.py.

Blocking (1-3):

1. nvext is never popped from the response. Lines 817-818 deliberately strip prompt_token_ids and choice.token_ids as transport details, but the nvext.engine_data bundle read at line 879 is left in place. NeMoGymChatCompletion inherits openai's extra="allow" base model, so the full token payload ships back to every caller a second time, roughly doubling stored rollout size.
2. The message-vs-Dynamo cross-check ignores generation_log_probs. Only token IDs are compared, and the message bundle wins on a tie. Stale message-level logprobs silently override Dynamo's authoritative values, corrupting the importance-sampling ratio with no error raised.
3. required_prefix_token_ids is derived for every capture-mode server with no config gate. It now reaches both /chat/completions and /tokenize. The prior turn's generation is not a byte-prefix of the re-rendered prompt for reasoning models (historical <think> is dropped) or tool-call turns (tool_calls are re-serialized), so this regresses previously-working configurations.

Also flagged (4-14): partial engine_data hard-failing instead of falling through; Gym-injected token_id: pseudo-tokens now returned to callers on the native path (and locked in by a new test); a lower-priority source able to veto a complete Dynamo bundle; stale-older-turn prefix selection; error messages naming Gym field names rather than Dynamo's; the feature being unreachable without the NeMo-RL wrapper since Gym never requests engine_data; test gaps; and reuse/efficiency cleanups.

Two comments (#1 and #5) concern lines 817 and 814, which fall between diff hunks and cannot be anchored exactly; each is attached to the nearest changed line in the same function and states the line it refers to.

Generated with Claude Code.

[jthomson04]

1. Blocking — nvext is never stripped from the response, so the full Dynamo token payload is echoed back twice.

This comment is about line 817, which falls outside the diff hunks; anchoring to the nearest changed line in the same function.

Gym/responses_api_models/vllm_model/app.py

Lines 815 to 819 in 3e7528c

	# Top-level and choice-level token-ID fields are transport details.
	chat_completion_dict.pop("prompt_token_ids", None)
	choice_dict.pop("token_ids", None)
	choice_dict["message"] = NeMoGymChatCompletionMessageForTraining.model_validate(message_dict)

Lines 817-818 deliberately strip the transport-only token fields:

chat_completion_dict.pop("prompt_token_ids", None)
choice_dict.pop("token_ids", None)

but nvext — read at line 879 and carrying exactly the same data — is left in place. NeMoGymChatCompletion extends openai's ChatCompletion, whose base model is extra="allow", so extras survive validation and FastAPI serialization. The existing assertion at test_app.py#L4825 (assert "prompt_token_ids" not in data) only makes sense because extras survive — which is precisely why those two pop calls exist.

Concretely, for a response with nvext.engine_data = {"prompt_token_ids": [...8k ints...], "completion_token_ids": [...2k ints...], "completion_logprobs": [...2k floats...]}, the returned JSON now carries the whole bundle twice: once normalized onto choices[0].message, once raw under nvext.engine_data. Every stored rollout record roughly doubles in token-data size, and transport internals the code hides everywhere else are re-exposed to downstream consumers.

Suggested fix: chat_completion_dict.pop("nvext", None) alongside the existing pops (or strip just engine_data if other nvext fields are meant to survive).

[jthomson04]

2. Blocking — the message-vs-Dynamo cross-check ignores generation_log_probs, and the message bundle silently wins.

Gym/responses_api_models/vllm_model/app.py

Lines 763 to 772 in 3e7528c

	if message_bundle is not None and dynamo_bundle is not None:
	if (
	message_bundle["prompt_token_ids"] != dynamo_bundle["prompt_token_ids"]
	or message_bundle["generation_token_ids"] != dynamo_bundle["generation_token_ids"]
	):
	raise RuntimeError("Message-level token metadata disagrees with Dynamo engine data.")
	native_bundle = message_bundle if message_bundle is not None else dynamo_bundle
	if native_bundle is not None:

The equality check covers only prompt_token_ids and generation_token_ids. If the two sources agree on token IDs but disagree on log probabilities, no error is raised and line 771 picks message_bundle, discarding Dynamo's values.

Failure case: a NeMo-RL wrapper copies correct token IDs onto choice.message but stale or recomputed generation_log_probs (previous engine step, different sampling config), while nvext.engine_data.completion_logprobs holds the true values. Both bundles pass the ID check, the message bundle wins, and Gym emits the wrong per-token logprobs. Those drive the importance-sampling ratio in RL, so training silently consumes a corrupted signal with no error anywhere.

This also contradicts the PR description's own statement that native Dynamo data "remains authoritative". Either extend the comparison to generation_log_probs, or make dynamo_bundle win for that field.

[jthomson04]

3. Blocking — required_prefix_token_ids is now auto-derived and forwarded for every server with return_token_id_information=true, with no config gate.

Gym/responses_api_models/vllm_model/app.py

Lines 604 to 607 in 3e7528c

	if self.config.return_token_id_information:
	self._derive_required_prefix_token_ids(body_dict)
	self._apply_sampling_overrides(body_dict)
	self._validate_single_choice_token_request(body_dict)

Compare the sibling vLLM-specific token flag, which is gated:

if self.config.request_prompt_and_generation_token_ids:
    body_dict["return_token_ids"] = True

(app.py#L487-L488)

The derived value flows into create_chat_completion(**body_dict) and, via _TOKENIZE_CHAT_FIELDS (app.py#L237-L244), into /tokenize as well. Before this PR the field was opt-in through extra_body only — see #1558, which added it purely to "keep fallback tokenization aligned with the generation prompt".

The derivation assumes the previous turn's prompt_token_ids + generation_token_ids is a byte-prefix of the newly rendered prompt. Two common cases where it is not:

• Reasoning models. With uses_reasoning_parser=true, turn 1's generation_token_ids include the <think>...</think> tokens, but on turn 2 app.py#L491-L523 moves reasoning into reasoning_content and most templates drop historical reasoning entirely.
• Tool calls. vLLM re-serializes tool_calls from parsed JSON, not the original generated bytes, so whitespace and key order drift.

In both cases an engine that enforces the prefix 400s the request, and the previously-working /tokenize fallback breaks too. Recommend gating this behind a config flag rather than applying it to all capture-mode servers.

[jthomson04]

4. Partial nvext.engine_data hard-fails instead of falling through to sources 3 and 4.

Gym/responses_api_models/vllm_model/app.py

Lines 896 to 900 in 3e7528c

	bundle = {
	destination: engine_data[source] for source, destination in field_mapping.items() if source in engine_data
	}
	return cls._validate_token_bundle(bundle, "nvext.engine_data")

If Dynamo returns engine_data = {"prompt_token_ids": [...], "completion_token_ids": [...]} — logprobs not requested, or not supported on that engine build — _validate_token_bundle raises RuntimeError("nvext.engine_data returned partial token metadata; missing: generation_log_probs.") before choice.logprobs or /tokenize is ever consulted.

The pre-PR code would have handled that response fine via the choice.logprobs path, which is still present in the same response. A capability the engine only partially advertises now becomes an unrecoverable failure for the entire rollout. Consider returning None (fall through) when the bundle is empty-or-partial in a way the other sources can cover, and reserving the hard error for genuinely contradictory data.

[jthomson04]

5. On the native path Gym returns its own injected token_id:<int> pseudo-tokens to callers, and a new test locks that in.

This comment is about line 814, which falls outside the diff hunks; anchoring to the nearest changed line (the else that guards it).

Gym/responses_api_models/vllm_model/app.py

Lines 812 to 816 in 3e7528c

	# The adapter consumed this compatibility payload.
	choice_dict.pop("logprobs", None)
	# Top-level and choice-level token-ID fields are transport details.

choice_dict.pop("logprobs", None) lives only in the else branch. The capture path force-sets logprobs=True, top_logprobs=0, return_tokens_as_token_ids=True (app.py#L478-L486), so with a Dynamo bundle present native_bundle is not None, the pop is skipped, and the response ships:

"logprobs": {"content": [{"token": "token_id:123", "logprob": -9.1, ...}]}

Any consumer reading logprobs.content[*].token expecting decoded text (log viewers, per-token attribution, external eval tooling) gets a Gym-internal debug encoding. This was already true for the message-bundle path, but the new test_capture_path_uses_dynamo_logprobs_without_consuming_choice_logprobs asserts the leaked payload is preserved verbatim, cementing it as intended behaviour rather than fixing it.

[jthomson04]

10. The derivation tests bypass the real request pipeline.

Gym/responses_api_models/vllm_model/tests/test_app.py

Lines 4641 to 4670 in 3e7528c

	def test_capture_path_derives_latest_required_prefix_token_ids(self) -> None:
	model = _make_top_logprobs_model(return_token_id_information=True)
	messages = [
	{"role": "user", "content": "old"},
	{
	"role": "assistant",
	"content": "old answer",
	"prompt_token_ids": [90],
	"generation_token_ids": [91],
	"generation_log_probs": [-0.9],
	},
	{"role": "user", "content": "first"},
	{
	"role": "assistant",
	"content": "answer",
	"prompt_token_ids": [1, 2],
	"generation_token_ids": [3, 4],
	"generation_log_probs": [-0.1, -0.2],
	},
	{"role": "user", "content": "next"},
	]
	result = model._preprocess_chat_completion_create_params(
	MagicMock(),
	{"model": "dummy_model", "messages": messages},
	)
	assert result["required_prefix_token_ids"] == [1, 2, 3, 4]
	result = model._preprocess_chat_completion_create_params(

Both assertions come from calling the private _preprocess_chat_completion_create_params with a hand-built dict, so nothing exercises POST /v1/chat/completions.

That matters here specifically: NeMoGymChatCompletionCreateParamsNonStreaming is a plain pydantic BaseModel whose messages is a smart union of TypedDicts (openai_utils.py#L773-L800). Whether an inbound assistant message's prompt_token_ids / generation_token_ids survive body.model_dump(exclude_unset=True) depends entirely on pydantic picking NeMoGymChatCompletionAssistantMessageForTrainingParam over NeMoGymChatCompletionAssistantMessageParam — the exact behaviour this feature depends on, and the exact behaviour these tests skip. A pydantic upgrade that changes smart-union preference would silently disable derivation with every test still green.

Worth adding one end-to-end test that posts a multi-turn body and asserts required_prefix_token_ids reaches create_chat_completion.

[jthomson04]

11. Three new branches added by this PR have no test coverage.

Gym/responses_api_models/vllm_model/app.py

Lines 880 to 886 in 3e7528c

	if nvext is None:
	return None
	if not isinstance(nvext, dict):
	raise RuntimeError("`nvext` must be an object when present.")
	if "engine_data" not in nvext:
	return None

• Line 883, RuntimeError("nvext must be an object when present.") — test_capture_path_rejects_malformed_dynamo_engine_data only parametrizes engine_data=[], never a non-dict nvext.
• Lines 884-885, nvext present without an engine_data key returning None. This is the common case for any Dynamo response carrying other nvext fields, and a regression here would break every request.
• Lines 972-975, the ValueError in _derive_required_prefix_token_ids.

CLAUDE.md#L140: "Test coverage must be >= 96%".

[jthomson04]

12. Reuse — this re-implements token-metadata validation with hardcoded literals instead of the shared helpers.

Gym/responses_api_models/vllm_model/app.py

Lines 967 to 976 in 3e7528c

	continue
	prompt_token_ids = message.get("prompt_token_ids")
	generation_token_ids = message.get("generation_token_ids")
	if prompt_token_ids is None and generation_token_ids is None:
	continue
	if not isinstance(prompt_token_ids, list) or not isinstance(generation_token_ids, list):
	raise ValueError(
	"Assistant token metadata must include prompt_token_ids and generation_token_ids as lists."
	)
	body_dict["required_prefix_token_ids"] = [*prompt_token_ids, *generation_token_ids]

The module already imports REQUIRED_TOKEN_METADATA_FIELDS (app.py#L35-L36), and the file already has _extract_message_token_bundle / _validate_token_bundle doing this validation; _validate_atomic_token_metadata (openai_utils.py#L157-L172) enforces the same all-or-nothing rule.

Adding a fourth metadata field would now require editing this literal-keyed function too, and its bespoke ValueError diverges in both wording and exception type from the RuntimeError the other three sites raise for the identical condition.

[jthomson04]

13. field_mapping is rebuilt on every response; make it a ClassVar.

Gym/responses_api_models/vllm_model/app.py

Lines 890 to 898 in 3e7528c

	field_mapping = {
	"prompt_token_ids": "prompt_token_ids",
	"completion_token_ids": "generation_token_ids",
	"completion_logprobs": "generation_log_probs",
	}
	bundle = {
	destination: engine_data[source] for source, destination in field_mapping.items() if source in engine_data
	}

_extract_dynamo_engine_data runs once per response on the hot path of a training rollout loop, allocating a fresh three-entry dict each time. The class already uses the ClassVar pattern for exactly this kind of constant table — _TOKENIZE_CHAT_FIELDS (app.py#L237) and _AUDIO_EXT_TO_MIME (app.py#L387).

[jthomson04]

14. CLAUDE.md — validation is unit-tests-only, with the real Dynamo smoke deferred until after merge.

CLAUDE.md#L14: "For environment or agent changes: run real rollouts with a model and inspect agent and verifier behavior. Green unit tests alone are not enough."

The PR's Validation section lists only pytest responses_api_models/vllm_model/tests/test_app.py and pre-commit, and states: "The final managed Dynamo smoke will run from NeMo-RL after this PR merges and NeMo-RL pins the exact containing Gym main commit."

Every nvext.engine_data payload shape in the tests is hand-authored, so nothing confirms the real field names, the nesting level (top-level vs. per-choice — this code only reads top-level chat_completion_dict["nvext"]), or the element type of completion_logprobs. _require_log_prob_list only checks that the value is a list, so a list of per-token logprob objects would pass validation here and surface later as a pydantic 500 rather than the clean rejection the PR promises.

[jthomson04]

One finding, surfaced while team-reviewing the NeMo-RL counterpart (#3763) — filed here because the fix belongs in this file, not there.

The token-selection reordering is a clear improvement over what NeMo-RL was doing: the cross-checks against message-level metadata and against vLLM's own response token IDs are checks the NeMo-RL wrapper never performed, and consolidating the Dynamo field mapping into _DYNAMO_ENGINE_DATA_FIELD_MAPPING means adding a field like routed_experts is now a one-repo change instead of two.

The gap is on the sad path: engine_data is requested but its absence is never asserted, so a missing field degrades into a /tokenize 404 instead of an error naming the field. Verified against the pinned Dynamo 1.3.0 — details inline.

Generated by Claude Code

[jthomson04]

responses_api_models/vllm_model/app.py:772

1 action item.

TL;DR — this PR asks Dynamo for nvext.engine_data but never checks it came back; when it's absent the code silently falls through to a /tokenize call that the NeMo-RL wrapper does not serve, turning a one-line diagnostic into a bare 404.

Raised here rather than on NeMo-RL #3763 because the fix lands in this file. #3763 deletes _validate_engine_data, which raised "Dynamo response did not include nvext.engine_data." — a clean 502 naming the field. _extract_dynamo_engine_data returns None in that case instead, so the check disappears rather than moving.

Why the fallthrough terminates at a 404 rather than at an earlier, friendlier error — I checked this against the exact pinned Dynamo (NeMo-RL pins ai-dynamo[vllm]==1.3.0.post1, asserted at grpo_dynamo.sh:32; the 1.3.0 workspace tree is what's cited below):

1. native_bundle is None, so the else at line 787 runs.
2. _extract_vllm_response_token_ids already returned None — Dynamo's chat path emits token IDs only under nvext (chat_completions/delta.rs:341-343) and there is no top-level prompt_token_ids anywhere in its chat_completions module (grep → 0 hits). choice.token_ids would need vLLM's return_token_ids: true, which this app only sends when request_prompt_and_generation_token_ids is set — and vllm_model_for_training.yaml sets it false.
3. So _extract_choice_logprobs does not save us either: Dynamo honors return_tokens_as_token_ids (chat_completions.rs:123) and emits format!("token_id:{}", tid) (delta.rs:122), and this app pins return_tokens_as_token_ids=True — so the logprobs parse cleanly and execution reaches...
4. ...await client.create_tokenize(...) at line 796, which POSTs /tokenize. The NeMo-RL wrapper registers only /health and /v1/chat/completions. 404.

Both conditions are structural properties of Dynamo's chat protocol, not a coincidence — _extract_vllm_response_token_ids returns None for every Dynamo chat response by construction.

Not hypothetical: docs/guides/dynamo-generation.md:65-67 in NeMo-RL already documents that "Dynamo 1.3.0's legacy tool jail removes nvext.engine_data" for tool_choice=auto.

Action: assert the contract where you already know it was requested. body_dict is in scope here (used at line 796), and _requests_dynamo_engine_data already encodes "the caller asked for this field". Putting it here rather than inside _extract_dynamo_engine_data keeps that helper a pure extractor. This blocks nothing legitimate — it fires only when the request itself put "engine_data" in nvext.extra_fields, i.e. exactly when the response is contractually required to carry it.

Suggested change

	native_bundle = message_bundle if message_bundle is not None else dynamo_bundle
	native_bundle = message_bundle if message_bundle is not None else dynamo_bundle
	if native_bundle is None and self._requests_dynamo_engine_data(body_dict):
	raise RuntimeError(
	"Request asked for `nvext.engine_data` but the response did not include it."
	)