[Bugfix][Frontend] Emit kv/ec transfer params at the top level of render output (#61241)

Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>

tests/entrypoints/scale_out/render/test_render.py

@@ -12,6 +12,7 @@ import pytest
 import pytest_asyncio
 
 from tests.utils import RemoteLaunchRenderServer
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.entrypoints.scale_out.render.api_router import router
@@ -115,6 +116,7 @@ async def test_render_responses_returns_generate_request_without_stored_state():
     assert response.priority == 3
     assert response.kv_transfer_params == {"do_remote_prefill": True}
     assert response.ec_transfer_params == {"remote": True}
+    assert response.sampling_params.extra_args is None
     serving.online_renderer.render_responses.assert_awaited_once_with(
         request,
         previous_messages=None,
@@ -124,6 +126,54 @@ async def test_render_responses_returns_generate_request_without_stored_state():
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.skip_global_cleanup
+async def test_render_chat_hoists_transfer_params_out_of_extra_args():
+    serving = _build_serving_render()
+    serving.online_renderer.render_chat = AsyncMock(
+        return_value=([], [{"prompt_token_ids": [7, 8, 9]}])
+    )
+    request = ChatCompletionRequest(
+        model=MODEL_NAME,
+        messages=[{"role": "user", "content": "Test prompt"}],
+        kv_transfer_params={"do_remote_prefill": True},
+        ec_transfer_params={"remote": True},
+        vllm_xargs={"foo": 1},
+    )
+
+    response = await serving.render_chat_request(request)
+
+    assert response.kv_transfer_params == {"do_remote_prefill": True}
+    assert response.ec_transfer_params == {"remote": True}
+    assert response.sampling_params.extra_args == {"foo": 1}
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip_global_cleanup
+async def test_render_completion_hoists_transfer_params_for_every_prompt():
+    serving = _build_serving_render()
+    serving.online_renderer.render_completion = AsyncMock(
+        return_value=[
+            {"prompt_token_ids": [7, 8, 9]},
+            {"prompt_token_ids": [4, 5]},
+        ]
+    )
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt=["a", "b"],
+        kv_transfer_params={"do_remote_prefill": True},
+        ec_transfer_params={"remote": True},
+    )
+
+    responses = await serving.render_completion_request(request)
+
+    assert len(responses) == 2
+    for response in responses:
+        assert response.kv_transfer_params == {"do_remote_prefill": True}
+        assert response.ec_transfer_params == {"remote": True}
+        assert response.sampling_params.extra_args is None
+
+
 @pytest.mark.asyncio
 @pytest.mark.skip_global_cleanup
 async def test_render_responses_rejects_previous_response_id():

vllm/entrypoints/scale_out/render/serving.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from fastapi import Request
 
@@ -33,6 +33,7 @@ from vllm.renderers.inputs.preprocess import (
     extract_prompt_len,
 )
 from vllm.renderers.online_renderer import OnlineRenderer
+from vllm.sampling_params import SamplingParams
 
 logger = init_logger(__name__)
 
@@ -40,6 +41,24 @@ if TYPE_CHECKING:
     from vllm.entrypoints.mcp.tool_server import ToolServer
 
 
+def _hoist_transfer_params(
+    params: SamplingParams,
+) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
+    """Move KV/EC transfer params out of `extra_args`.
+
+    Render emits them only as top-level `GenerateRequest` fields. Generate
+    copies them back into `extra_args` for the engine.
+    """
+    if not params.extra_args:
+        return None, None
+    # to_sampling_params aliases request.vllm_xargs, so copy before popping.
+    extra = dict(params.extra_args)
+    kv = extra.pop("kv_transfer_params", None)
+    ec = extra.pop("ec_transfer_params", None)
+    params.extra_args = extra or None
+    return kv, ec
+
+
 class ServingRender(BaseServing):
     def __init__(
         self,
@@ -126,6 +145,7 @@ class ServingRender(BaseServing):
             truncate_prompt_tokens=request.truncate_prompt_tokens,
         )
         params = request.to_sampling_params(max_tokens, self.default_sampling_params)
+        kv_transfer_params, ec_transfer_params = _hoist_transfer_params(params)
 
         # Resolve here what OpenAIServingChat passes to the engine:
         # `_grammar_from_parser` (set by `adjust_request`) is not serialized.
@@ -152,6 +172,8 @@ class ServingRender(BaseServing):
                 engine_input, include_mm_kwargs=request.return_mm_kwargs
             ),
             sampling_params=params,
+            kv_transfer_params=kv_transfer_params,
+            ec_transfer_params=ec_transfer_params,
             model=request.model,
             reasoning_ended=reasoning_ended,
             reasoning_parser_kwargs=reasoning_parser_kwargs,
@@ -226,6 +248,7 @@ class ServingRender(BaseServing):
             params = request.to_sampling_params(
                 max_tokens, self.default_sampling_params
             )
+            kv_transfer_params, ec_transfer_params = _hoist_transfer_params(params)
 
             request_id = base_id if len(result) == 1 else f"{base_id}-{i}"
 
@@ -237,6 +260,8 @@ class ServingRender(BaseServing):
                         engine_input, include_mm_kwargs=request.return_mm_kwargs
                     ),
                     sampling_params=params,
+                    kv_transfer_params=kv_transfer_params,
+                    ec_transfer_params=ec_transfer_params,
                     model=request.model,
                     stream=bool(request.stream),
                     stream_options=(request.stream_options if request.stream else None),
@@ -291,6 +316,7 @@ class ServingRender(BaseServing):
             truncate_prompt_tokens=(-1 if request.truncation != "disabled" else None),
         )
         params = request.to_sampling_params(max_tokens, self.default_sampling_params)
+        kv_transfer_params, ec_transfer_params = _hoist_transfer_params(params)
 
         return GenerateRequest(
             request_id=request.request_id,
@@ -304,8 +330,8 @@ class ServingRender(BaseServing):
             stream=bool(request.stream),
             cache_salt=request.cache_salt,
             priority=request.priority,
-            kv_transfer_params=request.kv_transfer_params,
-            ec_transfer_params=request.ec_transfer_params,
+            kv_transfer_params=kv_transfer_params,
+            ec_transfer_params=ec_transfer_params,
             token_offsets=engine_input.get("prompt_token_offsets"),
         )