fix(vllm): realign multimodal placeholders after prefix replacement

[aroshanghias-nvd]

What does this PR do?

Realigns vLLM multimodal placeholder offsets after NeMo-RL replaces re-tokenized conversation history with the exact model-generated token prefix.

In the OpenAI-compatible multi-turn generation path, vLLM first renders the chat template and computes mm_placeholders against that token sequence. NeMo-RL then calls replace_prefix_tokens() so prior turns use the exact token IDs originally consumed and sampled by the model. When the exact and re-tokenized prefixes have different lengths, replacing only prompt_token_ids leaves the multimodal offsets in the old template coordinates. vLLM can then apply media embeddings at incorrect token positions.

This PR:

• relocates each existing media-token span into the final exact-token prompt;
• processes image, audio, and video ranges in global prompt order so per-item shifts are preserved;
• preserves mapping and frozen-dataclass range representations;
• fails closed if an expected media span cannot be located;
• keeps an inexpensive no-op path when the prompt tokens are unchanged or there are no multimodal placeholders.

Issues

None.

Usage

No user-facing configuration or API changes are required.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added focused unit coverage for accumulated per-item shifts, mixed modalities, mapping/dataclass preservation, and missing-span failure.
• Focused unit validation completed: 3 passed, 42 deselected.
• No documentation changes are required because this is an internal token/media-alignment correction with no public API or configuration changes.

Additional Information

The repair belongs in NeMo-RL rather than Gym or vLLM: vLLM's offsets are correct for the prompt it originally preprocesses, and NeMo-RL subsequently changes that token sequence to preserve exact multi-turn rollout tokens.

A 20-turn multimodal no-compaction causal validation after applying the repair completed with TMPE 1.030905 (W&B run).

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

[aroshanghias-nvd]

Team review of #3858 — 5 agents (RL expert, bug finder, test agent, design reviewer, devil's advocate), each finding independently challenged before staging.

The change is needed and correctly layered. mm_placeholders carries offsets into the chat-template token sequence, and the existing code then replaces prompt_token_ids with a different-length splice — so media embeddings land on the wrong tokens with no error at all. Nice catch, and the fail-closed instinct is the right one. The two correctness findings below are both about the search technique, not the premise.

One CI blocker: the new test file isn't ruff-formatted, so pre-commit will fail.

Verified and cleared, so nobody re-litigates them: the call-site statement ordering is correct as written; the walrus guard's empty/missing cases are right; out-of-order ranges within a modality write back correctly; dataclasses.replace on the frozen PlaceholderRange handles the cached_property cleanly; the mm_placeholders key and its consumption by InputProcessor are correct at the pinned vLLM 0.25.1; and the three other replace_prefix_tokens call sites (TRT-LLM, Dynamo, the sync worker) genuinely need no remap. mm_hashes correctly needs no remapping either — hashes are over media content, not positions. No performance evidence is needed here: the call-site guard, the first-token short-circuit, and the monotonic search_start bound this to a single left-to-right pass.

Two things one field over, both latent, neither for this PR. assistant_tokens_mask on the same MultiModalInput is a per-token mask over prompt_token_ids and goes stale on the same splice — NeMo-RL never sets return_assistant_tokens_mask, so it's dormant. And super().preprocess_chat is called twice in this override (L487 and L544), the second re-decoding every image in the history; pre-existing, but it's the real cost on this path.

Generated by Claude Code

[aroshanghias-nvd]

1 action item.

tests/unit/models/generation/test_vllm_utils.py:77 — PR-introduced. This file is not ruff-formatted, so the pre-commit gate will fail CI. Verified with the repo-pinned ruff==0.9.9 (.pre-commit-config.yaml): ruff format --diff reports 1 file would be reformatted. ruff check, import-sort, and both source files are clean.

AI-1

Suggested change

	with pytest.raises(
	ValueError, match="Could not locate image placeholder range 0"
	):
	with pytest.raises(ValueError, match="Could not locate image placeholder range 0"):

[aroshanghias-nvd]

1 action item.

TL;DR — Qwen2.5-Omni emits an audio range with byte-identical (offset, length) to its video range, and the monotonic search_start makes the second one unfindable, so the request dies with a spurious ValueError.

PR-introduced. Not reachable today, but it blocks the backlog item already recorded as BL-20260428-omni-use-audio-in-video (intent.py:26-29).

How it breaks. With use_audio_in_video=True, vLLM derives the audio placeholder from the video placeholder using the same start_idx and the same tokens, differing only in is_embed — qwen2_5_omni_thinker.py L552-566 — and to_range() sets offset=start_idx, length=len(tokens). So both PlaceholderRanges are identical.

1. Both entries sort to the same old_offset.
2. The first resolves and sets search_start = new_offset + length — past the only occurrence.
3. The second searches from there, finds nothing, and raises.

Reproduced against this commit: template=[1,2,VS,S,VE], final=[1,VS,S,VE], {"video":[PR(2,3)],"audio":[PR(2,3)]} → ValueError: Could not locate audio placeholder range 0 from template offset 2.

It fails closed, so no silent corruption — but the ValueError is raised outside the except (ValueError, VLLMValidationError) above, and vLLM's route doesn't guard the handler, so it surfaces as an opaque HTTP 500 rather than a structured 400. Per the comment at nemo_rl/models/generation/vllm/vllm_worker_async.py that hits /tokenize as well as /v1/chat/completions.

AI-1

Memoize resolved offsets and advance search_start only for newly-resolved spans. Verified to fix both the one-video and two-video coincident cases:

Suggested change

	search_start = 0
	for old_offset, modality, item_index, placeholder_range, length in sorted(
	entries, key=lambda entry: entry[0]
	):
	expected = template_token_ids[old_offset : old_offset + length]
	max_start = len(final_token_ids) - length
	new_offset = next(
	(
	candidate
	for candidate in range(search_start, max_start + 1)
	if final_token_ids[candidate] == expected[0]
	and final_token_ids[candidate : candidate + length] == expected
	),
	None,
	)
	if new_offset is None:
	raise ValueError(
	f"Could not locate {modality} placeholder range {item_index} "
	f"from template offset {old_offset} in the final exact-token prompt"
	)
	if isinstance(placeholder_range, Mapping):
	updated_range = dict(placeholder_range)
	updated_range["offset"] = new_offset
	elif is_dataclass(placeholder_range):
	updated_range = replace(placeholder_range, offset=new_offset)
	else:
	raise TypeError(
	"Multimodal placeholder ranges must be mappings or dataclass "
	f"instances, got {type(placeholder_range).__name__}"
	)
	remapped[modality][item_index] = updated_range
	search_start = new_offset + length
	search_start = 0
	resolved: dict[tuple[int, int], int] = {}
	for old_offset, modality, item_index, placeholder_range, length in sorted(
	entries, key=lambda entry: entry[0]
	):
	# Modalities can share an identical span -- Qwen2.5-Omni derives the audio
	# range from the video range, differing only in `is_embed` -- so reuse a
	# resolved offset instead of scanning past it.
	cached_offset = resolved.get((old_offset, length))
	if cached_offset is not None:
	new_offset = cached_offset
	else:
	expected = template_token_ids[old_offset : old_offset + length]
	max_start = len(final_token_ids) - length
	new_offset = next(
	(
	candidate
	for candidate in range(search_start, max_start + 1)
	if final_token_ids[candidate] == expected[0]
	and final_token_ids[candidate : candidate + length] == expected
	),
	None,
	)
	if new_offset is None:
	raise ValueError(
	f"Could not locate {modality} placeholder range {item_index} "
	f"from template offset {old_offset} in the final exact-token prompt"
	)
	resolved[(old_offset, length)] = new_offset
	if isinstance(placeholder_range, Mapping):
	updated_range = dict(placeholder_range)
	updated_range["offset"] = new_offset
	elif is_dataclass(placeholder_range):
	updated_range = replace(placeholder_range, offset=new_offset)
	else:
	raise TypeError(
	"Multimodal placeholder ranges must be mappings or dataclass "
	f"instances, got {type(placeholder_range).__name__}"
	)
	remapped[modality][item_index] = updated_range
	search_start = max(search_start, new_offset + length)

[aroshanghias-nvd]

1 action item.

TL;DR — "fails closed" is not guaranteed: if a media run is longer in the exact prefix than in the template, the next item silently matches inside the previous item's run.

PR-introduced. Distinct from the coincident-range issue above — that one is a duplicate-key collision, this one is an unanchored first-match. Confirmed by measurement: the memoization fix above resolves the coincident case and leaves this one bit-identical.

Repro on this commit (executed, no exception raised):

template_token_ids = [1, *[pad] * 4, 2, *[pad] * 4, 3]
final_token_ids    = [1, *[pad] * 8, 2, *[pad] * 4, 3]
mm_placeholders    = {"video": [PR(offset=1, length=4), PR(offset=6, length=4)]}
# -> offsets [1, 5];  correct answer [1, 10]

Item 0 matches at the run start, but search_start = new_offset + length stops inside the grown run, so item 1 matches item 0's own pad tokens. The general form is the same: Qwen2-VL placeholder spans are bare runs of one repeated token with no <|vision_start|>/<|vision_end|> bracketing, so a first-match search has nothing anchoring it.

I could not demonstrate a trigger on any in-repo path — still images have deterministic pad counts, nemotron_vl samples deterministically, and the other video path sends video_as_images=True. So this is hardening plus docstring accuracy, not a live break.

AI-1

Either add a run-boundary check (reject a match whose neighbouring tokens disagree with the template) or drop the fail-closed claim. One way to do it — of several:

Suggested change

	Ranges are matched in global prompt order because different media items can
	accumulate different shifts. A missing span fails closed rather than
	submitting token IDs with incorrect multimodal positions.
	Ranges are matched in global prompt order because different media items can
	accumulate different shifts. A span that cannot be located fails closed, but
	note that an ambiguous match is not currently detected: if a media run is
	longer in ``final_token_ids`` than in the template, a later item can match
	inside an earlier item's run.

Please don't take the ambiguity-rejection route (raise if a second occurrence exists) — two identical images legitimately produce two identical spans, so that would fail closed on the most ordinary multi-image prompt.

[aroshanghias-nvd]

1 action item.

PR-introduced. This block has a load-bearing order dependency that nothing encodes: remap_multimodal_placeholders must read engine_prompt["prompt_token_ids"] while it still holds the template ids, and line 572 overwrites that key with final_prompt_token_ids.

Swap the two statements — the natural "assign the result where you compute it" refactor — and template_token_ids == final_token_ids, so the function takes its early-return fast path, placeholders stay in template coordinates, and no exception fires. That silently reintroduces the exact bug this PR fixes. preprocess_chat is defined inside a method body (L456) so it can't be imported, and no unit test references it — nothing would go red.

AI-1

Suggested change

	if mm_placeholders := engine_prompt.get("mm_placeholders"):
	engine_prompt["mm_placeholders"] = remap_multimodal_placeholders(
	template_token_ids=engine_prompt["prompt_token_ids"],
	final_token_ids=final_prompt_token_ids,
	mm_placeholders=mm_placeholders,
	)
	# Must run before `prompt_token_ids` is reassigned below: the remap
	# reads the current value as the template-coordinate token sequence.
	if mm_placeholders := engine_prompt.get("mm_placeholders"):
	engine_prompt["mm_placeholders"] = remap_multimodal_placeholders(
	template_token_ids=engine_prompt["prompt_token_ids"],
	final_token_ids=final_prompt_token_ids,
	mm_placeholders=mm_placeholders,
	)

Follow-up

Optionally fold both statements into one apply_exact_prefix_token_ids(engine_prompt, final_token_ids) helper in vllm/utils.py, which would make the contract testable against a plain dict. Worth it for the test, though it moves the ordering hazard inside the helper rather than removing it — your call.

[aroshanghias-nvd]

1 action item.

Four of the new function's branches are unreached by any test. All four tests below were run against this commit and pass. SimpleNamespace is already imported at line 17, so no new imports are needed.

Worth noting one asymmetry they expose: the fast path returns before range validation, so an invalid range is silently accepted when the prompt is unchanged but rejected when it changed.

AI-1

def test_remap_multimodal_placeholders_leaves_an_unchanged_prompt_alone():
    ranges = [_PlaceholderRange(offset=1, length=2)]
    placeholders = {"image": ranges}
    token_ids = [1, 18, 18, 2]

    result = remap_multimodal_placeholders(
        template_token_ids=token_ids,
        final_token_ids=list(token_ids),
        mm_placeholders=placeholders,
    )

    assert result == placeholders
    assert result is not placeholders
    assert result["image"] is not ranges


def test_remap_multimodal_placeholders_without_media_returns_empty():
    assert (
        remap_multimodal_placeholders(
            template_token_ids=[1, 2, 3],
            final_token_ids=[9, 9, 1, 2, 3],
            mm_placeholders={},
        )
        == {}
    )


@pytest.mark.parametrize("offset, length", [(-1, 2), (1, 0), (1, -2), (2, 3)])
def test_remap_multimodal_placeholders_rejects_out_of_range_input(offset, length):
    with pytest.raises(ValueError, match="Invalid image placeholder range 0"):
        remap_multimodal_placeholders(
            template_token_ids=[1, 18, 18, 2],
            final_token_ids=[9, 1, 18, 18, 2],
            mm_placeholders={
                "image": [_PlaceholderRange(offset=offset, length=length)]
            },
        )


def test_remap_multimodal_placeholders_rejects_an_unknown_range_type():
    with pytest.raises(TypeError, match="got SimpleNamespace"):
        remap_multimodal_placeholders(
            template_token_ids=[1, 18, 18, 2],
            final_token_ids=[9, 1, 18, 18, 2],
            mm_placeholders={"image": [SimpleNamespace(offset=1, length=2)]},
        )

Follow-up

A red test documenting the grown-run case from the docstring comment (currently returns [1, 5] instead of [1, 10]). Land it with whichever fix you choose, and assert bare pytest.raises(ValueError) rather than matching a message — the natural fix would emit a different string than the existing "could not locate" template.

[aroshanghias-nvd]

1 action item. Low severity.

PR-introduced. This assertion certifies an output shape the engine cannot consume. vLLM only ever emits PlaceholderRange — the engine-prompt TypedDict pins it as Mapping[str, Sequence[PlaceholderRange]] — and every consumer uses attribute access: x[2].offset in argsort_mm_positions, mm_position.get_num_embeds(). A plain dict would AttributeError inside the engine.

AI-1

Change the audio fixture at line 65 to a _PlaceholderRange and assert its .offset, so the test only certifies shapes the engine accepts.

Context — no action. The Mapping branch itself is harmless: it costs four lines and cannot fire at the pinned version. Keeping it for cross-version tolerance is defensible — the ask here is only about the test.

[aroshanghias-nvd]

No action needed in this PR — raising it as a follow-up because it subsumes the two search-related findings above.

The remap searches for information the caller already has exactly. replace_prefix_tokens is a splice at a known boundary — it computes template_cut_start and model_cut_end, returns only the concatenation, and discards both indices. Given them, every span in the suffix region (old_offset >= template_cut_start) is pure arithmetic:

new_offset = old_offset - template_cut_start + model_cut_end

with a slice-equality assertion turning "fails closed" from a hope into a verified postcondition. Only spans landing inside the replaced prefix would still need a search, bounded to final[:model_cut_end]. That removes the ambiguity class entirely rather than hardening the scan against it.

It touches a second file and widens a shared signature, so it's genuinely a follow-up rather than something to do here.