[Frontend] Reuse prefill token ids on the decode chat path for disaggregated serving (#48145)
Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Seiji Eicher committed
Jul 29, 2026 at 01:02 UTC
6370e53f246c97649615ccd7fca5366ee73f54aa
6 files changed
+198
-13
.buildkite/test-amd.yaml
+1
-1
@@ -2289,7 +2289,7 @@ steps:
2289
- export VLLM_USE_RUST_FRONTEND=1
2290
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
2291
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
2292
- - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
2292
+ - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
2293
- pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
2294
- pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
2295
- pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"
.buildkite/test_areas/rust_frontend.yaml
+1
-1
@@ -26,7 +26,7 @@ steps:
26
- export VLLM_USE_RUST_FRONTEND=1
27
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
28
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
29
- - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
29
+ - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
30
- pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
31
32
# - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
docs/features/disagg_prefill.md
+28
@@ -49,6 +49,34 @@ Now supports 9 types of connectors:
49
--kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
50
```
51
52
+## Reusing prefill token ids on decode
53
+
54
+!!! note
55
+ This applies to disaggregated prefill and decode serving on the `/v1/chat/completions` endpoint, using a KV connector configured as in the Usage example above. It is experimental and subject to change.
56
+
57
+In disaggregated serving, the prefill and decode stages both render the chat prompt from `messages` and tokenize it. Because the prefill stage has already produced the token ids, the decode stage can reuse them and skip its own templating and tokenization. The output is otherwise identical to a normal chat completion: it is detokenized to text, and tool and reasoning parsing, streaming, and structured output constraints all still apply.
58
+
59
+The token ids are passed to the decode stage through `kv_transfer_params`, the dict already attached to the decode request to coordinate the transfer:
60
+
61
+1. Send the prefill request with `return_token_ids` enabled, and read `prompt_token_ids` from the response.
62
+2. Set `kv_transfer_params["prompt_token_ids"]` to those ids on the decode request. `messages` is still required, but its content is not tokenized when the ids are present.
63
+
64
+```python
65
+prefill = client.chat.completions.create(
66
+ model=model,
67
+ messages=messages,
68
+ extra_body={"return_token_ids": True, "kv_transfer_params": {"do_remote_decode": True}},
69
+)
70
+ids = prefill.prompt_token_ids
71
+
72
+decode = client.chat.completions.create(
73
+ model=model,
74
+ messages=messages,
75
+ stream=True,
76
+ extra_body={"kv_transfer_params": {"do_remote_prefill": True, "prompt_token_ids": ids}},
77
+)
78
+```
79
+
80
## Development
81
82
We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance.
tests/entrypoints/openai/chat_completion/test_chat_completion.py
+89
@@ -158,3 +158,92 @@ async def test_empty_grammar(client: openai.AsyncOpenAI, model_name: str) -> Non
158
],
159
extra_body={"structured_outputs": {"grammar": ""}},
160
)
161
+
162
+
163
+# Decode-side token reuse for disaggregated serving. The router forwards the
164
+# prefill stage's prompt token ids in kv_transfer_params so the decode stage
165
+# skips re-tokenizing.
166
+
167
+TOKEN_IN_MESSAGES = [{"role": "user", "content": "Hello, how are you today?"}]
168
+DECODE_MESSAGES = [{"role": "user", "content": "unrelated decode-side text"}]
169
+
170
+
171
+@pytest.mark.asyncio
172
+async def test_kv_transfer_prompt_token_ids_round_trip(client: openai.AsyncOpenAI):
173
+ """Ids forwarded in kv_transfer_params are used verbatim, skipping tokenize.
174
+
175
+ The decode request carries different messages, so a response whose
176
+ prompt_token_ids match the forwarded ids proves the ids were used rather
177
+ than the request's own messages. Generated text is not compared across
178
+ requests because vLLM greedy decoding is not bitwise-reproducible.
179
+ """
180
+ baseline = await client.chat.completions.create(
181
+ model=MODEL_NAME,
182
+ messages=TOKEN_IN_MESSAGES,
183
+ max_completion_tokens=16,
184
+ temperature=0,
185
+ extra_body={"return_token_ids": True},
186
+ )
187
+ reused_ids = baseline.prompt_token_ids
188
+ assert reused_ids
189
+
190
+ decode = await client.chat.completions.create(
191
+ model=MODEL_NAME,
192
+ messages=DECODE_MESSAGES,
193
+ max_completion_tokens=16,
194
+ temperature=0,
195
+ extra_body={
196
+ "kv_transfer_params": {"prompt_token_ids": reused_ids},
197
+ "return_token_ids": True,
198
+ },
199
+ )
200
+
201
+ # The engine saw the forwarded ids, not the decode request's own messages.
202
+ assert decode.prompt_token_ids == reused_ids
203
+ # text-out: reuse still yields a detokenized message.
204
+ assert decode.choices[0].message.content
205
+
206
+
207
+@pytest.mark.asyncio
208
+async def test_kv_transfer_prompt_token_ids_streaming(client: openai.AsyncOpenAI):
209
+ """Decode-side token reuse streams chat-formatted text-out."""
210
+ baseline = await client.chat.completions.create(
211
+ model=MODEL_NAME,
212
+ messages=TOKEN_IN_MESSAGES,
213
+ max_completion_tokens=16,
214
+ temperature=0,
215
+ extra_body={"return_token_ids": True},
216
+ )
217
+ reused_ids = baseline.prompt_token_ids
218
+ assert reused_ids
219
+
220
+ stream = await client.chat.completions.create(
221
+ model=MODEL_NAME,
222
+ messages=DECODE_MESSAGES,
223
+ max_completion_tokens=16,
224
+ temperature=0,
225
+ stream=True,
226
+ extra_body={
227
+ "kv_transfer_params": {"prompt_token_ids": reused_ids},
228
+ "return_token_ids": True,
229
+ },
230
+ )
231
+
232
+ content = ""
233
+ delta_token_ids: list[int] = []
234
+ first_chunk = True
235
+ async for chunk in stream:
236
+ if first_chunk:
237
+ # prompt_token_ids arrives once, on the first chunk.
238
+ assert chunk.prompt_token_ids == reused_ids
239
+ first_chunk = False
240
+ if not chunk.choices:
241
+ continue
242
+ if chunk.choices[0].delta.content:
243
+ content += chunk.choices[0].delta.content
244
+ if tids := getattr(chunk.choices[0], "token_ids", None):
245
+ delta_token_ids.extend(tids)
246
+
247
+ # streamed text-out, reconstructed from deltas, with generated token ids.
248
+ assert content
249
+ assert delta_token_ids
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+31
@@ -2314,3 +2314,34 @@ async def test_streaming_n_gt1_independent_tool_parsers():
2314
f"Choice {choice_idx}: expected finish_reason='tool_calls', "
2315
f"got '{reasons[0]}'"
2316
)
2317
+
2318
+
2319
+def test_make_request_with_harmony_reuses_kv_transfer_prompt_token_ids():
2320
+ """The Harmony reuse branch honors ids forwarded in kv_transfer_params.
2321
+
2322
+ A GPT-OSS server is impractical to stand up here, so this exercises the
2323
+ branch directly on a harmony-configured renderer.
2324
+ """
2325
+ engine = MockEngine()
2326
+ engine.model_config.hf_config = MockHFConfig(model_type="gpt_oss")
2327
+ models = OpenAIServingModels(engine, BASE_MODEL_PATHS)
2328
+ online_renderer = _build_online_renderer(engine, models.registry)
2329
+ assert online_renderer.use_harmony
2330
+
2331
+ request = ChatCompletionRequest(
2332
+ model=MODEL_NAME,
2333
+ messages=[{"role": "user", "content": "hi"}],
2334
+ kv_transfer_params={
2335
+ "prompt_token_ids": [10, 20, 30],
2336
+ "do_remote_prefill": True,
2337
+ },
2338
+ )
2339
+ conversation, engine_inputs = online_renderer._make_request_with_harmony(request)
2340
+
2341
+ assert conversation == []
2342
+ assert len(engine_inputs) == 1
2343
+ engine_input = engine_inputs[0]
2344
+ assert engine_input["type"] == "token"
2345
+ assert engine_input["prompt_token_ids"] == [10, 20, 30]
2346
+ # The reuse key is consumed and other kv_transfer_params are preserved.
2347
+ assert request.kv_transfer_params == {"do_remote_prefill": True}
vllm/renderers/online_renderer.py
+48
-11
@@ -47,6 +47,19 @@ from vllm.utils.mistral import mt as _mt
47
logger = init_logger(__name__)
48
49
50
+def _reused_prompt_token_ids(request: Any) -> list[int] | None:
51
+ """Pop prompt token ids forwarded for decode-side reuse, if any.
52
+
53
+ Disaggregated serving carries the prefill stage's ids in
54
+ ``kv_transfer_params`` so the decode stage can skip re-tokenizing. Removing
55
+ the key keeps the id list out of the engine's sampling metadata.
56
+ """
57
+ kv = getattr(request, "kv_transfer_params", None)
58
+ if not isinstance(kv, dict):
59
+ return None
60
+ return kv.pop("prompt_token_ids", None) or None
61
+
62
+
63
class OnlineRenderer:
64
def __init__(
65
self,
@@ -102,6 +115,12 @@ class OnlineRenderer:
115
116
Called directly by render_chat_request and delegated to by
117
OpenAIServingChat.render_chat_request after its engine-aware checks.
118
+
119
+ Decode-side token reuse (ids forwarded in ``kv_transfer_params``) is
120
+ handled deeper, in ``preprocess_chat`` / ``_make_request_with_harmony``,
121
+ so it skips only templating and tokenization while tool-choice
122
+ validation and ``adjust_request`` still run and the output is
123
+ detokenized (text-out).
124
"""
125
tokenizer = self.renderer.tokenizer
126
@@ -186,6 +205,13 @@ class OnlineRenderer:
205
should_include_tools: bool = True,
206
):
207
"""Build Harmony (GPT-OSS) messages and engine prompt from a chat request."""
208
+ reuse_ids = _reused_prompt_token_ids(request)
209
+ if reuse_ids:
210
+ # Decode-side token reuse: feed the forwarded ids straight to the
211
+ # engine. Harmony has no adjust_request hook to preserve.
212
+ engine_input = tokens_input(reuse_ids, cache_salt=request.cache_salt)
213
+ return [], [engine_input]
214
+
215
messages: list[OpenAIMessage] = []
216
217
# because of issues with pydantic we need to potentially
@@ -368,17 +394,28 @@ class OnlineRenderer:
394
default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
395
)
396
371
- (conversation,), (engine_input,) = await renderer.render_chat_async(
372
- [messages],
373
- chat_params,
374
- tok_params,
375
- prompt_extras={
376
- k: v
377
- for k in ("mm_processor_kwargs", "cache_salt")
378
- if (v := getattr(request, k, None)) is not None
379
- },
380
- skip_mm_cache=skip_mm_cache,
381
- )
397
+ reuse_ids = _reused_prompt_token_ids(request)
398
+ if reuse_ids:
399
+ # Decode-side token reuse: feed the forwarded ids straight to the
400
+ # engine, skipping templating and tokenization. ``messages`` are not
401
+ # tokenized, so conversation is empty. The adjust_request tail below
402
+ # still runs.
403
+ conversation: list[ConversationMessage] = []
404
+ engine_input = tokens_input(
405
+ reuse_ids, cache_salt=getattr(request, "cache_salt", None)
406
+ )
407
+ else:
408
+ (conversation,), (engine_input,) = await renderer.render_chat_async(
409
+ [messages],
410
+ chat_params,
411
+ tok_params,
412
+ prompt_extras={
413
+ k: v
414
+ for k in ("mm_processor_kwargs", "cache_salt")
415
+ if (v := getattr(request, k, None)) is not None
416
+ },
417
+ skip_mm_cache=skip_mm_cache,
418
+ )
419
420
# tool parsing is done only if a tool_parser has been set and if
421
# tool_choice is not "none" (if tool_choice is "none" but a tool_parser