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

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

@@ -2307,3 +2307,34 @@ async def result_generator():
             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 @@
 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 @@ async def render_chat(
 
         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 @@ def _make_request_with_harmony(
         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 @@ async def preprocess_chat(
             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