[RL+VLM] Avoid retokenization drift for pre-tokenized (token-id) VLM requests

[ByronHsu]

TL;DR

For RL, we want token-in, token-out on both the inference and training sides so the two see the exact same token ids. Today the VLM path breaks this: when a request comes in as input_ids, the server decodes them back to a string and then re-tokenizes that string. Decode → re-tokenize is not guaranteed to round-trip, so the prompt the model actually runs on can differ from what the client sent ("retokenization drift").

This is harmless for plain text, but in RL it can cause a catastrophic trainer–inference mismatch — especially with router replay, where the recorded expert-routing decisions are indexed per token and must line up exactly.

Example:

trainer    prompt = [1, 2]
inference  prompt = [1, 2]  ->  decoded to a string  ->  re-tokenized to [3]   (drift!)

# routing is recorded per token:
inference produced routing for 1 token,
but the trainer replays 2 tokens  ->  length mismatch, replay is corrupted

Modification

• BaseMultimodalProcessor.process_and_combine_mm_data: for a pre-tokenized (list[int]) prompt, rebuild the final input_ids from the user's original tokens and expand only the image placeholders, instead of adopting the HF re-tokenized prompt. Non-media tokens can no longer drift. The whole path is guarded by SGLANG_MM_AVOID_RETOKENIZE (default on).
• resolve_image_token_counts: per-image expanded token counts computed without re-tokenizing. The base implementation uses the transformers in-tree convention _get_num_multimodal_tokens(image_sizes=...) (Qwen-VL, Gemma3, GLM4V, ...); Kimi overrides it with media_tokens_calculator. On genuine failure it raises; the caller logs the cause and falls back to decode+retokenize.
• _expand_input_ids: copies the original tokens verbatim and expands the i-th image placeholder into counts[i] copies of the image token id. Raises ValueError when the number of placeholders does not match the number of images.

Test

test/registered/vlm/test_token_id_retokenize_e2e.py launches a real Qwen2.5-VL server twice with SGLANG_MM_AVOID_RETOKENIZE off then on, sending the same predefined non-canonical prompt (the word "Describe" split into "D"+"escribe", which decodes to the same text but re-encodes to the single merged token) plus one image:

• flag OFF → the prompt re-tokenizes (drift): prompt_tokens shrinks by the drift delta.
• flag ON → no drift: prompt_tokens keeps the original length (image placeholder expanded).

Asserts prompt_tokens[on] - prompt_tokens[off] == drift_delta.

---

CI States

Latest PR Test (Base): ⏳ Run #26699777289
Latest PR Test (Extra): ❌ Run #26699777237

[gemini-code-assist]

Code Review

This pull request introduces direct token-id space expansion for Kimi VLM prompts to resolve off-by-one length drift issues caused by decoding and re-tokenizing. It adds a dedicated processor module, integrates it into the Kimi K2.5 and VL processors, and includes both a reproduction script and unit tests. The reviewer suggested an optimization in kimi_token_ids.py to calculate token counts directly from existing offsets instead of redundantly calling the media token calculator.

[gemini-code-assist]

The media_tokens_calculator is called again for each image here, which is redundant because the token counts have already been calculated and are encoded in image_offsets as (start, end) ranges. We can optimize this by calculating num_tokens directly from image_offsets using end - start + 1, avoiding redundant CPU overhead.

Suggested change

	for image in images:
	num_tokens = processor._processor.media_processor.media_tokens_calculator(
	{"type": "image", "image": image}
	)
	media_text_parts.append(processor.mm_tokens.image_token * num_tokens)
	mediums.append({"type": "image", "image": image})
	for image, (start, end) in zip(images, image_offsets):
	num_tokens = end - start + 1
	media_text_parts.append(processor.mm_tokens.image_token * num_tokens)
	mediums.append({"type": "image", "image": image})

[ByronHsu]

/tag-and-rerun-ci

[mickqian]

do we need to take care of video/audio placeholder in this PR?

[ByronHsu]

we can do it later. We don't have use cases for video/audio RL for now.

[yhyang201]

/tag-and-rerun-ci