[Serve][LLM] Avoid double prompt tokenization in prefill-decode disaggregation

[eicherseiji]

Summary

In prefill-decode disaggregation, the same chat prompt is tokenized twice. Each stage drives the engine through vLLM's chat-completions path, which renders and tokenizes messages. Prefill tokenizes the prompt and decode tokenizes it again. For large prompts the tokenizer pass dominates TTFT.

With this change, decode reuses the prompt token ids that prefill already produced and skips the redundant encode. Decode stays on the chat path. Chat-template rendering, tool and reasoning parsing, sampling, and response shaping are unchanged. Only the tokenizer encode is skipped.

The feature is opt-in and off by default. Enable it with experimental_configs["pd_tokenize_once"].

Measured impact

Numbers come from a 4-prefill, 4-decode agentic benchmark with microsoft/Phi-tiny-MoE-instruct, roughly 37k input tokens, and 256 concurrent clients. The off and on runs used the same build. Median client TTFT drops 16%.

Metric	Off	On
Client TTFT median	989 ms	833 ms
Client TTFT mean	1030 ms	902 ms
Decode tokenize median	164 ms	~0 ms
Decode dispatch to first token median	400 ms	188 ms

Prefill still tokenizes once. TPOT is unaffected.

How it works

vLLM's renderer already skips the tokenizer encode when a rendered prompt arrives with prompt_token_ids set. It still preserves multi_modal_data and runs detokenization, truncation, and validation. So decode does not need a new code path. It needs a way to get prefill's ids to that one spot inside vLLM.

That spot is several layers away. The orchestrator _pd_handle_request holds the ids. Tokenization happens inside vLLM's renderer, reached through super().chat(), the engine, and vLLM's serving layer. The chat request has no field for pre-tokenized input, and the intervening layers are vLLM code, so no signature can thread the ids through.

A contextvars.ContextVar bridges that gap. The value is task-local, so each request's decode reads only the ids its orchestrator set and concurrent requests stay isolated. The orchestrator sets the ids before driving decode, and the tokenize call deep inside vLLM reads them when it runs.

Concretely:

• The orchestrator sets the existing return_token_ids=True field on the prefill request and reads the echoed ids off the prefill response. Chat carries them top-level as prompt_token_ids. Completions carry them on the first choice as CompletionResponseChoice.prompt_token_ids. The orchestrator then wraps the local decode call in with reuse_prompt_token_ids(ids), which sets the contextvar for that request's decode and clears it after.
• install() runs once per decode replica from PDDecodeServer.__init__ and wraps BaseRenderer.tokenize_prompts_async. When the contextvar holds ids and the request rendered to a single ordinary text prompt, the wrapper adds prompt_token_ids to that prompt and delegates to the real method, which then skips the encode. Every other case delegates unchanged: no ids set, batched, embeds, encoder-decoder, or already tokenized.
• The install fails safe. When vLLM's renderer does not match, install() returns False, the feature stays off, and tokenization is untouched.

The wrap lives in ray/llm/_internal/common/patches/vllm/, a package for temporary vLLM patches.

Scope and correctness

Only the two sequential handoff paths reuse tokens. The concurrent_handoff path starts decode before prefill returns the ids, so it does not request the prefill echo.

Correctness relies on decode tokenization running on the orchestrator's async task or a task spawned from it. If a future change breaks that, the contextvar is not visible at tokenize time and decode falls back to normal tokenization. The result is slower and still correct.

The wrap patches a vLLM internal, BaseRenderer.tokenize_prompts_async, because the chat-completions path has no public way to accept pre-tokenized input. The intended end state is native pre-tokenized chat input, which removes this wrap. See vllm-project/vllm#22817 for the upstream direction.

Test plan

test_tokenize_once.py covers the following:

• Id injection that preserves other prompt fields, and fall-through for batched, embeds, encoder-decoder, already-tokenized, and disabled prompts.
• Idempotent install, no-op install without vLLM, and install returning False when the renderer method is absent.
• Cross-context teardown, and concurrent tasks that do not cross-contaminate the reuse contextvar.
• Application against the real vLLM renderer. The wrap lands on BaseRenderer.tokenize_prompts_async, and an injected prompt_token_ids reaches per-prompt tokenization so the encode is skipped.
• Orchestrator reuse helpers. _decode_reuse_ids reads ids top-level for chat and from the first choice for completions and returns None when disabled. _request_prefill_token_ids sets return_token_ids only when enabled.

The benchmark above covers performance.

[gemini-code-assist]

Code Review

This pull request introduces a "tokenize-once" optimization for prefill/decode (P/D) disaggregation in Ray LLM, allowing the decode stage to reuse prompt token IDs generated during the prefill stage and bypass redundant tokenization. The feedback focuses on improving the robustness and safety of this optimization. Specifically, it is recommended to wrap the monkey-patching logic in a try-except block with safe attribute access to prevent server startup crashes on incompatible vLLM versions, add a defensive check for non-iterable token IDs, and disable the optimization flag if the hook installation fails.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[kouroshHakha]

this works, but is there a better way that is native to vllm?

[kouroshHakha]

big idea: can't we use servingtokens implementation for token-in text-out implementation instead of this brittle patching:

https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/serve/disagg/serving.py

I am looking for some way to natively implement token-in, text-out (or maybe token-out).

[eicherseiji]

yeah token-in, text-out would be ideal. A similar idea and is being implemented in this series as part of RL support, e.g. /render and /derender with ServingTokens in the middle.

• [Render] Add reasoning and tool call parsing to /derender vllm-project/vllm#45045
• [Frontend] Add GPU-less render serving path (vllm launch render) vllm-project/vllm#36166
• [RFC]: Derender Endpoints to Provide Detokenization for Disaggregated Serving vllm-project/vllm#42729
• [RFC]: Disaggregated Frontend — Separating Online Serving from Engine vllm-project/vllm#34407
• [RFC]: Unified Parser for reasoning, tool calling vllm-project/vllm#32713

However, this /render + /derender explicitly does not support streaming. After some research, we're going to need some vLLM side patches no matter what.

[kouroshHakha]

running release tests to be sure: https://buildkite.com/ray-project/release/builds/97389