From c6b772cf6be783b1a0dbf6046570e17c3bc85ec6 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Thu, 15 Oct 2026 13:34:15 +0100 Subject: [PATCH] [Bugfix][Frontend] Emit kv/ec transfer params at the top level of render output (#61241) Signed-off-by: Martin Hickey --- .../scale_out/render/test_render.py | 50 +++++++++++++++++++ vllm/entrypoints/scale_out/render/serving.py | 32 ++++++++++-- 2 files changed, 79 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/scale_out/render/test_render.py b/tests/entrypoints/scale_out/render/test_render.py index 37f4b8eecf..53bdbf423d 100644 --- a/tests/entrypoints/scale_out/render/test_render.py +++ b/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(): diff --git a/vllm/entrypoints/scale_out/render/serving.py b/vllm/entrypoints/scale_out/render/serving.py index b8d8ab7e1b..54d6606427 100644 --- a/vllm/entrypoints/scale_out/render/serving.py +++ b/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"), )