[Frontend] Reuse prefill token ids on the decode chat path for disaggregated serving

[eicherseiji]

Why

In prefill and decode disaggregation, the prefill stage renders the prompt from messages and tokenizes it. The router forwards the same chat request to the decode stage, which renders and tokenizes it a second time. For long prompts, that repeated render and tokenize adds latency on the decode critical path.

The decode stage does not need to redo it. The prefill response already contains the token ids, and the router already forwards state to the decode request. This change makes the decode chat path use the forwarded ids and skip the render and tokenize.

The ids are carried on the chat path rather than /generate because the decode output must still be chat-shaped: detokenized text, tool and reasoning parsing, and streaming. /generate in #24261 returns raw tokens without that parsing. Running on the chat path also reuses the chat generator's existing stateful parser instead of adding a second one.

What changed

• On decode, the renderer reads the forwarded ids from kv_transfer_params, builds the engine input from them, and skips chat templating and tokenization. It removes the key after reading it, so the ids are not copied into engine sampling metadata.
• The reuse branch is in preprocess_chat and the Harmony path, not at the top of render_chat. Tool-choice validation and adjust_request still run, so grammar, structured output, and reasoning constraints still apply.
• No public request field is added. protocol.py is unchanged, messages stays required, and a request without the ids renders as before.

Usage

The prefill response returns prompt_token_ids when return_token_ids is set. The router already writes kv_transfer_params on the decode request to coordinate the transfer. Reuse adds one more key to that dict.

1. Send the prefill request with return_token_ids set. Read prompt_token_ids from the response.
2. Set kv_transfer_params["prompt_token_ids"] on the decode request to those ids, next to the transfer flags.
3. The decode stage returns the same chat-formatted, streamed output without re-tokenizing.

prefill = client.chat.completions.create(
    model=model, messages=messages,
    extra_body={"return_token_ids": True, "kv_transfer_params": {"do_remote_decode": True}},
)
ids = prefill.prompt_token_ids

decode = client.chat.completions.create(
    model=model, messages=messages, stream=True,
    extra_body={"kv_transfer_params": {"do_remote_prefill": True, "prompt_token_ids": ids}},
)

Why this is not a duplicate

Related work	Relationship
#39756 skip decode-side re-tokenization	Direct predecessor. Added a public prompt_token_ids field gated on transfer params and was abandoned in rebase. This version carries the ids in kv_transfer_params and adds no field.
#24261 /generate tokens in and out	Merged. Returns raw tokens with no chat tool or reasoning parsing. This change reuses ids on the chat endpoint instead, so that parsing and detokenization still run.
#47161 RFC, #47301 streaming derender	Open. The derender formats as a separate stateless service for scale-out setups such as llm-d and Dynamo, with client-carried state. #47301 is Part 1 and streams detokenized text only. Reasoning and tool parsing are stubbed for a later Part 2 that is not open yet. This change targets in-process prefill and decode and reuses the decode server's existing stateful parser, so parsed streaming chat output works today. Whether one parser can serve both is an open question, since this path runs in-process and stateful and the derender runs stateless across a fleet.
#22587 return_token_ids	Merged. The token-out half. The prefill ids come from it.
#22817 Disaggregated Everything	Open RFC. The parent token-in and token-out direction.

No open PR adds decode-side chat token reuse this way.

Open questions

• Whether kv_transfer_params is the right channel for the ids, or whether a dedicated internal field is cleaner. The renderer reads a key from a dict that otherwise holds connector parameters.
• Whether the streaming reasoning and tool parser should be shared with the derender work rather than written twice. The two run under different models, one stateful in-process and one stateless client-carried, so sharing is a design question for the derender and chat-serving owners, not a drop-in.

Model evaluation

This is a serving-path change and needs an accuracy check on a disaggregated setup, not only the mechanism tests below. Reusing the prefill ids should produce the same result as tokenizing on decode, and the round-trip test confirms the engine receives exactly the forwarded ids. The end-to-end accuracy run needs a prefill and decode deployment, which the single-GPU tests here do not cover.

• GSM8K on a 1P1D setup, reuse on versus off, before merge.

Test plan

Added two tests to tests/entrypoints/openai/chat_completion/test_chat_completion.py. The decode request carries different messages from the reused ids, so a response whose prompt_token_ids match the forwarded ids confirms the ids were used and not the request's own messages. Generated text is not compared across requests because vLLM greedy decoding is not bitwise-reproducible.

pytest tests/entrypoints/openai/chat_completion/test_chat_completion.py -k prompt_token_ids -v

Result on 1x A10G: 2 passed, round-trip and streaming. ruff check and ruff format pass.

Note

AI assistance from Claude Code was used to prepare this change. The submitter reviewed it.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[NickLucche]

Hey @eicherseiji thanks a lot for the PR!
I am also not fully sure where we should be placing the token_ids, as the original tokens-in-out API entails the final "completions wrapping" to be delegated to the router/coordinator.
However, I agree we should avoid re-tokenization on D even for completions API and I agree it makes sense to do that while also allowing D to stream back completions chunks (instead of "raw" GenerateResponse).
Stuffing everything inside kv_transfer_params makes the interface opaque and basically undiscoverable, but I think we can still get this merged for now until tokens-in-out becomes more standard.

I would suggest adding the token_ids path to the docs, but other than that this LGTM.

[mergify]

Documentation preview: https://vllm--48145.org.readthedocs.build/en/48145/

[mergify]

Hi @eicherseiji, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.