@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 17ee641c45

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 07:00 UTC

[Responses API] Add kv_transfer_params for PD disaggregation (#37424)

Signed-off-by: bongwoobak <bongwoobak@gmail.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>

Bongwoo Bak committed Mar 21, 2026 at 14:48 UTC 17ee641c4510c93d8d2b826b19daa9f86126894e

3 files changed +29 -2

vllm/entrypoints/openai/responses/context.py

+12 -1

@@ -9,7 +9,7 @@ from abc import ABC, abstractmethod

9 from collections.abc import Callable

10 from contextlib import AsyncExitStack

11 from dataclasses import replace

12 -from typing import TYPE_CHECKING, Final, Union

12 +from typing import TYPE_CHECKING, Any, Final, Union

14 from openai.types.responses.response_function_tool_call_output_item import (

15 ResponseFunctionToolCallOutputItem,

@@ -182,6 +182,7 @@ class SimpleContext(ConversationContext):

182 self.all_turn_metrics = []

183

184 self.input_messages: list[ResponseRawMessageAndToken] = []

185 + self.kv_transfer_params: dict[str, Any] | None = None

186

187 def append_output(self, output) -> None:

188 self.last_output = output

@@ -190,6 +191,8 @@ class SimpleContext(ConversationContext):

191 self.num_prompt_tokens = len(output.prompt_token_ids or [])

192 self.num_cached_tokens = output.num_cached_tokens or 0

193 self.num_output_tokens += len(output.outputs[0].token_ids or [])

194 + if output.kv_transfer_params is not None:

195 + self.kv_transfer_params = output.kv_transfer_params

196

197 # Accumulate text, token_ids, and logprobs for streaming mode

198 delta_output = output.outputs[0]

@@ -308,11 +311,14 @@ class ParsableContext(ConversationContext):

311 self.input_messages: list[ResponseRawMessageAndToken] = []

312 self.output_messages: list[ResponseRawMessageAndToken] = []

313 self._accumulated_token_ids: list[int] = []

314 + self.kv_transfer_params: dict[str, Any] | None = None

315

316 def append_output(self, output: RequestOutput) -> None:

317 self.num_prompt_tokens = len(output.prompt_token_ids or [])

318 self.num_cached_tokens = output.num_cached_tokens or 0

319 self.num_output_tokens += len(output.outputs[0].token_ids or [])

320 + if output.kv_transfer_params is not None:

321 + self.kv_transfer_params = output.kv_transfer_params

322 self.parser.process(output.outputs[0])

323 output_token_ids = output.outputs[0].token_ids or []

324 self._accumulated_token_ids.extend(output_token_ids)

@@ -538,6 +544,7 @@ class HarmonyContext(ConversationContext):

544 self.all_turn_metrics: list[TurnMetrics] = []

545 self.is_first_turn = True

546 self.first_tok_of_message = True # For streaming support

547 + self.kv_transfer_params: dict[str, Any] | None = None

548

549 def _update_num_reasoning_tokens(self):

550 channel = self.parser.current_channel

@@ -557,6 +564,8 @@ class HarmonyContext(ConversationContext):

564 self._update_num_reasoning_tokens()

565 self._update_prefill_token_usage(output)

566 self._update_decode_token_usage(output)

567 + if output.kv_transfer_params is not None:

568 + self.kv_transfer_params = output.kv_transfer_params

569 # Append current turn to all turn list for next turn's calculations

570 self.all_turn_metrics.append(self.current_turn_metrics.copy())

571 self.current_turn_metrics.reset()

@@ -868,6 +877,8 @@ class StreamingHarmonyContext(HarmonyContext):

877 if last_delta_text:

878 self.last_content_delta = last_delta_text

879 self._update_decode_token_usage(output)

880 + if output.kv_transfer_params is not None:

881 + self.kv_transfer_params = output.kv_transfer_params

882

883 # For streaming, update previous turn when message is complete

884 if output.finished:

vllm/entrypoints/openai/responses/protocol.py

+16 -1

@@ -252,6 +252,10 @@ class ResponsesRequest(OpenAIBaseModel):

252 "numeric values, used by custom extensions."

253 ),

254 )

255 + kv_transfer_params: dict[str, Any] | None = Field(

256 + default=None,

257 + description="KVTransfer parameters used for disaggregated serving.",

258 + )

259 # --8<-- [end:responses-extra-params]

260

261 def build_chat_params(

@@ -351,6 +355,10 @@ class ResponsesRequest(OpenAIBaseModel):

355 if isinstance(stop, str):

356 stop = [stop]

357

358 + extra_args: dict[str, Any] = self.vllm_xargs if self.vllm_xargs else {}

359 + if self.kv_transfer_params:

360 + extra_args["kv_transfer_params"] = self.kv_transfer_params

361 +

362 return SamplingParams.from_optional(

363 temperature=temperature,

364 top_p=top_p,

@@ -367,7 +375,7 @@ class ResponsesRequest(OpenAIBaseModel):

375 ),

376 structured_outputs=structured_outputs,

377 logit_bias=self.logit_bias,

370 - extra_args=self.vllm_xargs or {},

378 + extra_args=extra_args,

379 skip_clone=True, # Created fresh per request, safe to skip clone

380 skip_special_tokens=self.skip_special_tokens,

381 include_stop_str_in_output=self.include_stop_str_in_output,

@@ -488,6 +496,11 @@ class ResponsesResponse(OpenAIBaseModel):

496 usage: ResponseUsage | None = None

497 user: str | None = None

498

499 + # vLLM-specific fields that are not in OpenAI spec

500 + kv_transfer_params: dict[str, Any] | None = Field(

501 + default=None, description="KVTransfer parameters."

502 + )

503 +

504 # --8<-- [start:responses-response-extra-params]

505 # These are populated when enable_response_messages is set to True

506 # NOTE: custom serialization is needed

@@ -531,6 +544,7 @@ class ResponsesResponse(OpenAIBaseModel):

544 usage: ResponseUsage | None = None,

545 input_messages: ResponseInputOutputMessage | None = None,

546 output_messages: ResponseInputOutputMessage | None = None,

547 + kv_transfer_params: dict[str, Any] | None = None,

548 ) -> "ResponsesResponse":

549 incomplete_details: IncompleteDetails | None = None

550 if status == "incomplete":

@@ -566,6 +580,7 @@ class ResponsesResponse(OpenAIBaseModel):

580 truncation=request.truncation,

581 user=request.user,

582 usage=usage,

583 + kv_transfer_params=kv_transfer_params,

584 )

585

586

vllm/entrypoints/openai/responses/serving.py

@@ -873,6 +873,7 @@ class OpenAIServingResponses(OpenAIServing):

873 output=output,

874 status=status,

875 usage=usage,

876 + kv_transfer_params=context.kv_transfer_params,

877 )

878

879 if request.store: