@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3f5bd482f5

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 4, 00:00 UTC

[Bugfix][KV Transfer][NIXL] Notify P node on pre-admission rejection to free stranded KV blocks (#41269)

Dao007forever committed May 9, 2026 at 22:52 UTC 3f5bd482f5c1a5dbdffbbf68d624e20bb7032013

11 files changed +156 -2

tests/v1/kv_connector/unit/test_nixl_connector.py

+28

@@ -321,6 +321,34 @@ def test_prompt_less_than_block_size():

321 assert len(scheduler_output.scheduled_new_reqs) == 0

322

323

324 +def test_abort_immediately_remote_prefill_enqueues_empty_recv():

325 + """A remote-prefill request added with abort_immediately=True should

326 + be added to the scheduler's waiting queue then immediately aborted, so the

327 + NIXL connector's request_finished hook enqueues an empty recv to notify

328 + the prefill instance to free its blocks."""

329 + from vllm.v1.request import RequestStatus

330 +

331 + scheduler = create_scheduler(create_vllm_config())

332 +

333 + request = create_request(request_id=42, num_tokens=10, do_remote_prefill=True)

334 + assert request.kv_transfer_params is not None

335 + assert request.kv_transfer_params["do_remote_prefill"] is True

336 +

337 + # Mimic the EngineCore.add_request path for an abort-immediately req.

338 + scheduler.add_request(request)

339 + scheduler.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED)

340 +

341 + scheduler_output = scheduler.schedule()

342 + meta = scheduler_output.kv_connector_metadata

343 + assert isinstance(meta, NixlConnectorMetadata)

344 + assert set(meta.reqs_to_recv) == {request.request_id}

345 + req_meta = meta.reqs_to_recv[request.request_id]

346 + assert req_meta.local_block_ids == []

347 + assert req_meta.remote.request_id == f"prefill-{42}"

348 + # do_remote_prefill is consumed by request_finished to prevent re-issuing.

349 + assert request.kv_transfer_params["do_remote_prefill"] is False

350 +

351 +

352 @patch(

353 "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",

354 FakeNixlWrapper,

vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py

+3 -1

@@ -525,7 +525,9 @@ class NixlConnectorScheduler:

525 if params.get("do_remote_prefill"):

526 # If do_remote_prefill is still True when the request is finished,

527 # update_state_after_alloc must not have been called (the request

528 - # must have been aborted before it was scheduled).

528 + # must have been aborted before it was scheduled, e.g. via the

529 + # abort_immediately path used to clean up KV-transfer requests

530 + # rejected at the D-side serving layer).

531 # To avoid stranding the prefill blocks in the prefill instance,

532 # we must add empty block_ids to _reqs_need_recv so that our

533 # worker side will notify and free blocks in the prefill instance.

vllm/engine/protocol.py

+14

@@ -108,6 +108,20 @@ class EngineClient(ABC):

108 """

109 ...

110

111 + @abstractmethod

112 + async def notify_kv_transfer_request_rejected(

113 + self,

114 + request_id: str,

115 + kv_transfer_params: dict[str, Any],

116 + *,

117 + data_parallel_rank: int | None = None,

118 + ) -> None:

119 + """Notify the engine that a KV-transfer request was rejected before

120 + engine admission, so connector-side cleanup can run (e.g. free

121 + prefill blocks pinned on the P node).

122 + """

123 + ...

124 +

125 @abstractmethod

126 async def is_tracing_enabled(self) -> bool: ...

127

vllm/entrypoints/openai/chat_completion/serving.py

@@ -234,6 +234,15 @@ class OpenAIServingChat(OpenAIServing):

234 for the API specification. This API mimics the OpenAI

235 Chat Completion API.

236 """

237 + return await self._with_kv_transfer_rejection_cleanup(

238 + self._create_chat_completion(request, raw_request), request, raw_request

239 + )

240 +

241 + async def _create_chat_completion(

242 + self,

243 + request: ChatCompletionRequest,

244 + raw_request: Request | None = None,

245 + ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:

246 # Streaming response

247 tokenizer = self.renderer.tokenizer

248 assert tokenizer is not None

vllm/entrypoints/openai/completion/serving.py

@@ -118,6 +118,15 @@ class OpenAIServingCompletion(OpenAIServing):

118 - suffix (the language models we currently support do not support

119 suffix)

120 """

121 + return await self._with_kv_transfer_rejection_cleanup(

122 + self._create_completion(request, raw_request), request, raw_request

123 + )

124 +

125 + async def _create_completion(

126 + self,

127 + request: CompletionRequest,

128 + raw_request: Request | None = None,

129 + ) -> AsyncGenerator[str, None] | CompletionResponse | ErrorResponse:

130 if request.stream and request.use_beam_search:

131 return self.create_error_response(

132 "Streaming is not currently supported with beam search"

vllm/entrypoints/openai/engine/serving.py

+39 -1

@@ -4,7 +4,7 @@ import asyncio

4 import contextlib

5 import json

6 import time

7 -from collections.abc import AsyncGenerator, Mapping

7 +from collections.abc import AsyncGenerator, Awaitable, Mapping

8 from dataclasses import dataclass, field

9 from http import HTTPStatus

10 from typing import Any, ClassVar, Generic, Protocol, TypeAlias, TypeVar

@@ -118,6 +118,7 @@ AnyResponse: TypeAlias = (

118 )

119

120 RequestT = TypeVar("RequestT", bound=AnyRequest)

121 +_T = TypeVar("_T")

122

123

124 @dataclass(kw_only=True)

@@ -156,6 +157,9 @@ class OpenAIServing:

157 self.model_config = engine_client.model_config

158 self.renderer = engine_client.renderer

159 self.input_processor = engine_client.input_processor

160 + vllm_config = getattr(engine_client, "vllm_config", None)

161 + kv_transfer_config = getattr(vllm_config, "kv_transfer_config", None)

162 + self.has_kv_connector = kv_transfer_config is not None

163

164 # Computed once at startup (cached by ``vllm_config`` identity) and

165 # stamped on non-streaming responses. Streaming chunks deliberately

@@ -616,6 +620,40 @@ class OpenAIServing:

620 except ValueError:

621 return None

622

623 + async def _with_kv_transfer_rejection_cleanup(

624 + self,

625 + awaitable: Awaitable[_T],

626 + request: ChatCompletionRequest | CompletionRequest | ResponsesRequest,

627 + raw_request: Request | None,

628 + ) -> _T:

629 + """Wrap a `create_*` coroutine so that, if it raises or returns an

630 + ErrorResponse (i.e. the request never reached the engine), the KV

631 + connector is notified to free any pinned remote-prefill blocks."""

632 + kv_transfer_params = self.has_kv_connector and request.kv_transfer_params

633 + if not kv_transfer_params or not kv_transfer_params.get("do_remote_prefill"):

634 + return await awaitable

635 +

636 + notify = True

637 + try:

638 + result = await awaitable

639 + if not isinstance(result, ErrorResponse):

640 + notify = False

641 + return result

642 + finally:

643 + if notify:

644 + try:

645 + await self.engine_client.notify_kv_transfer_request_rejected(

646 + request.request_id,

647 + kv_transfer_params,

648 + data_parallel_rank=self._get_data_parallel_rank(raw_request),

649 + )

650 + except Exception:

651 + logger.warning(

652 + "Failed to notify KV connector about rejected request %s",

653 + request.request_id,

654 + exc_info=True,

655 + )

656 +

657 @staticmethod

658 def _parse_tool_calls_from_content(

659 request: ResponsesRequest | ChatCompletionRequest,

vllm/entrypoints/openai/responses/serving.py

+11

@@ -323,6 +323,17 @@ class OpenAIServingResponses(OpenAIServing):

323 AsyncGenerator[StreamingResponsesResponse, None]

324 | ResponsesResponse

325 | ErrorResponse

326 + ):

327 + return await self._with_kv_transfer_rejection_cleanup(

328 + self._create_responses(request, raw_request), request, raw_request

329 + )

330 +

331 + async def _create_responses(

332 + self, request: ResponsesRequest, raw_request: Request | None = None

333 + ) -> (

334 + AsyncGenerator[StreamingResponsesResponse, None]

335 + | ResponsesResponse

336 + | ErrorResponse

337 ):

338 error_check_ret = await self._check_model(request)

339 if error_check_ret is not None:

vllm/v1/engine/__init__.py

@@ -122,6 +122,12 @@ class EngineCoreRequest(

122 reasoning_ended: bool | None = None

123 reasoning_parser_kwargs: dict[str, Any] | None = None

124

125 + # If True, the request should be added to the scheduler's waiting queue

126 + # and immediately aborted, so connector-side cleanup runs via the standard

127 + # request_finished hook. Used to free P-side prefill blocks when a

128 + # KV-transfer request is rejected on the D node before engine admission.

129 + abort_immediately: bool = False

130 +

131 @property

132 def params(self) -> SamplingParams | PoolingParams:

133 """Return the processed params (sampling or pooling)."""

vllm/v1/engine/async_llm.py

+27

@@ -720,6 +720,33 @@ class AsyncLLM(EngineClient):

720 if self.log_requests:

721 logger.info("Aborted request(s) %s.", ",".join(request_ids))

722

723 + async def notify_kv_transfer_request_rejected(

724 + self,

725 + request_id: str,

726 + kv_transfer_params: dict[str, Any],

727 + *,

728 + data_parallel_rank: int | None = None,

729 + ) -> None:

730 + """Submit a pre-aborted request so the connector's request_finished

731 + hook runs to free any pre-admission KV-transfer resources (e.g. NIXL

732 + prefill blocks pinned on the P node)."""

733 + request = EngineCoreRequest(

734 + request_id=request_id,

735 + prompt_token_ids=[0],

736 + mm_features=None,

737 + sampling_params=SamplingParams(

738 + max_tokens=1,

739 + extra_args={"kv_transfer_params": dict(kv_transfer_params)},

740 + ),

741 + pooling_params=None,

742 + arrival_time=time.time(),

743 + lora_request=None,

744 + cache_salt=None,

745 + data_parallel_rank=data_parallel_rank,

746 + abort_immediately=True,

747 + )

748 + await self.engine_core.add_request_async(request)

749 +

750 async def pause_generation(

751 self,

752 *,

vllm/v1/engine/core.py

@@ -344,6 +344,10 @@ class EngineCore:

344 )

345

346 self.scheduler.add_request(request)

347 + if request.abort_immediately:

348 + # Immediately abort so the connector's request_finished hook runs

349 + # to free any pre-admission KV-transfer resources.

350 + self.abort_requests([request.request_id])

351

352 def abort_requests(self, request_ids: list[str]):

353 """Abort requests from the scheduler."""

vllm/v1/request.py

@@ -76,6 +76,7 @@ class Request:

76 resumable: bool = False,

77 reasoning_ended: bool | None = None,

78 reasoning_parser_kwargs: dict[str, Any] | None = None,

79 + abort_immediately: bool = False,

80 ) -> None:

81 self.request_id = request_id

82 self.client_index = client_index

@@ -182,6 +183,10 @@ class Request:

183 # None entry in the queue means finished.

184 self.streaming_queue: deque[StreamingUpdate | None] | None = None

185

186 + # If True, request should be aborted immediately after being added to

187 + # the scheduler so the connector's request_finished hook runs.

188 + self.abort_immediately = abort_immediately

189 +

190 @classmethod

191 def from_engine_core_request(

192 cls,

@@ -206,6 +211,7 @@ class Request:

211 resumable=request.resumable,

212 reasoning_ended=request.reasoning_ended,

213 reasoning_parser_kwargs=request.reasoning_parser_kwargs,

214 + abort_immediately=request.abort_immediately,

215 )

216

217 def append_output_token_ids(