[Frontend] Reuse prefill token ids on the decode chat path for disaggregated serving (#48145)

Signed-off-by: Seiji Eicher <[email protected]>

.buildkite/test-amd.yaml

@@ -2289,7 +2289,7 @@ steps:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
-  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
+  - 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"
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
   - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
   - pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"

.buildkite/test_areas/rust_frontend.yaml

@@ -26,7 +26,7 @@ steps:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
-  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
+  - 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"
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
 
   # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"

docs/features/disagg_prefill.md

@@ -49,6 +49,34 @@ Now supports 9 types of connectors:
   --kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
   ```
 
+## Reusing prefill token ids on decode
+
+!!! note
+    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.
+
+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.
+
+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:
+
+1. Send the prefill request with `return_token_ids` enabled, and read `prompt_token_ids` from the response.
+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.
+
+```python
+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}},
+)
+```
+
 ## Development
 
 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

@@ -158,3 +158,92 @@ async def test_empty_grammar(client: openai.AsyncOpenAI, model_name: str) -> Non
             ],
             extra_body={"structured_outputs": {"grammar": ""}},
         )
+
+
+# Decode-side token reuse for disaggregated serving. The router forwards the
+# prefill stage's prompt token ids in kv_transfer_params so the decode stage
+# skips re-tokenizing.
+
+TOKEN_IN_MESSAGES = [{"role": "user", "content": "Hello, how are you today?"}]
+DECODE_MESSAGES = [{"role": "user", "content": "unrelated decode-side text"}]
+
+
+@pytest.mark.asyncio
+async def test_kv_transfer_prompt_token_ids_round_trip(client: openai.AsyncOpenAI):
+    """Ids forwarded in kv_transfer_params are used verbatim, skipping tokenize.
+
+    The decode request carries different messages, so a response whose
+    prompt_token_ids match the forwarded ids proves the ids were used rather
+    than the request's own messages. Generated text is not compared across
+    requests because vLLM greedy decoding is not bitwise-reproducible.
+    """
+    baseline = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=TOKEN_IN_MESSAGES,
+        max_completion_tokens=16,
+        temperature=0,
+        extra_body={"return_token_ids": True},
+    )
+    reused_ids = baseline.prompt_token_ids
+    assert reused_ids
+
+    decode = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=DECODE_MESSAGES,
+        max_completion_tokens=16,
+        temperature=0,
+        extra_body={
+            "kv_transfer_params": {"prompt_token_ids": reused_ids},
+            "return_token_ids": True,
+        },
+    )
+
+    # The engine saw the forwarded ids, not the decode request's own messages.
+    assert decode.prompt_token_ids == reused_ids
+    # text-out: reuse still yields a detokenized message.
+    assert decode.choices[0].message.content
+
+
+@pytest.mark.asyncio
+async def test_kv_transfer_prompt_token_ids_streaming(client: openai.AsyncOpenAI):
+    """Decode-side token reuse streams chat-formatted text-out."""
+    baseline = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=TOKEN_IN_MESSAGES,
+        max_completion_tokens=16,
+        temperature=0,
+        extra_body={"return_token_ids": True},
+    )
+    reused_ids = baseline.prompt_token_ids
+    assert reused_ids
+
+    stream = await client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=DECODE_MESSAGES,
+        max_completion_tokens=16,
+        temperature=0,
+        stream=True,
+        extra_body={
+            "kv_transfer_params": {"prompt_token_ids": reused_ids},
+            "return_token_ids": True,
+        },
+    )
+
+    content = ""
+    delta_token_ids: list[int] = []
+    first_chunk = True
+    async for chunk in stream:
+        if first_chunk:
+            # prompt_token_ids arrives once, on the first chunk.
+            assert chunk.prompt_token_ids == reused_ids
+            first_chunk = False
+        if not chunk.choices:
+            continue
+        if chunk.choices[0].delta.content:
+            content += chunk.choices[0].delta.content
+        if tids := getattr(chunk.choices[0], "token_ids", None):
+            delta_token_ids.extend(tids)
+
+    # streamed text-out, reconstructed from deltas, with generated token ids.
+    assert content
+    assert delta_token_ids

tests/entrypoints/openai/chat_completion/test_serving_chat.py

@@ -2314,3 +2314,34 @@ async def test_streaming_n_gt1_independent_tool_parsers():
             f"Choice {choice_idx}: expected finish_reason='tool_calls', "
             f"got '{reasons[0]}'"
         )
+
+
+def test_make_request_with_harmony_reuses_kv_transfer_prompt_token_ids():
+    """The Harmony reuse branch honors ids forwarded in kv_transfer_params.
+
+    A GPT-OSS server is impractical to stand up here, so this exercises the
+    branch directly on a harmony-configured renderer.
+    """
+    engine = MockEngine()
+    engine.model_config.hf_config = MockHFConfig(model_type="gpt_oss")
+    models = OpenAIServingModels(engine, BASE_MODEL_PATHS)
+    online_renderer = _build_online_renderer(engine, models.registry)
+    assert online_renderer.use_harmony
+
+    request = ChatCompletionRequest(
+        model=MODEL_NAME,
+        messages=[{"role": "user", "content": "hi"}],
+        kv_transfer_params={
+            "prompt_token_ids": [10, 20, 30],
+            "do_remote_prefill": True,
+        },
+    )
+    conversation, engine_inputs = online_renderer._make_request_with_harmony(request)
+
+    assert conversation == []
+    assert len(engine_inputs) == 1
+    engine_input = engine_inputs[0]
+    assert engine_input["type"] == "token"
+    assert engine_input["prompt_token_ids"] == [10, 20, 30]
+    # The reuse key is consumed and other kv_transfer_params are preserved.
+    assert request.kv_transfer_params == {"do_remote_prefill": True}

vllm/renderers/online_renderer.py

@@ -47,6 +47,19 @@ from vllm.utils.mistral import mt as _mt
 logger = init_logger(__name__)
 
 
+def _reused_prompt_token_ids(request: Any) -> list[int] | None:
+    """Pop prompt token ids forwarded for decode-side reuse, if any.
+
+    Disaggregated serving carries the prefill stage's ids in
+    ``kv_transfer_params`` so the decode stage can skip re-tokenizing. Removing
+    the key keeps the id list out of the engine's sampling metadata.
+    """
+    kv = getattr(request, "kv_transfer_params", None)
+    if not isinstance(kv, dict):
+        return None
+    return kv.pop("prompt_token_ids", None) or None
+
+
 class OnlineRenderer:
     def __init__(
         self,
@@ -102,6 +115,12 @@ class OnlineRenderer:
 
         Called directly by render_chat_request and delegated to by
         OpenAIServingChat.render_chat_request after its engine-aware checks.
+
+        Decode-side token reuse (ids forwarded in ``kv_transfer_params``) is
+        handled deeper, in ``preprocess_chat`` / ``_make_request_with_harmony``,
+        so it skips only templating and tokenization while tool-choice
+        validation and ``adjust_request`` still run and the output is
+        detokenized (text-out).
         """
         tokenizer = self.renderer.tokenizer
 
@@ -186,6 +205,13 @@ class OnlineRenderer:
         should_include_tools: bool = True,
     ):
         """Build Harmony (GPT-OSS) messages and engine prompt from a chat request."""
+        reuse_ids = _reused_prompt_token_ids(request)
+        if reuse_ids:
+            # Decode-side token reuse: feed the forwarded ids straight to the
+            # engine. Harmony has no adjust_request hook to preserve.
+            engine_input = tokens_input(reuse_ids, cache_salt=request.cache_salt)
+            return [], [engine_input]
+
         messages: list[OpenAIMessage] = []
 
         # because of issues with pydantic we need to potentially
@@ -368,17 +394,28 @@ class OnlineRenderer:
             default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
         )
 
-        (conversation,), (engine_input,) = await renderer.render_chat_async(
-            [messages],
-            chat_params,
-            tok_params,
-            prompt_extras={
-                k: v
-                for k in ("mm_processor_kwargs", "cache_salt")
-                if (v := getattr(request, k, None)) is not None
-            },
-            skip_mm_cache=skip_mm_cache,
-        )
+        reuse_ids = _reused_prompt_token_ids(request)
+        if reuse_ids:
+            # Decode-side token reuse: feed the forwarded ids straight to the
+            # engine, skipping templating and tokenization. ``messages`` are not
+            # tokenized, so conversation is empty. The adjust_request tail below
+            # still runs.
+            conversation: list[ConversationMessage] = []
+            engine_input = tokens_input(
+                reuse_ids, cache_salt=getattr(request, "cache_salt", None)
+            )
+        else:
+            (conversation,), (engine_input,) = await renderer.render_chat_async(
+                [messages],
+                chat_params,
+                tok_params,
+                prompt_extras={
+                    k: v
+                    for k in ("mm_processor_kwargs", "cache_salt")
+                    if (v := getattr(request, k, None)) is not None
+                },
+                skip_mm_cache=skip_mm_cache,
+            )
 
         # tool parsing is done only if a tool_parser has been set and if
         # tool_choice is not "none" (if tool_choice is "none" but a tool_parser