@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 71d1d75b7a

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 15, 03:30 UTC

[PD][Nixl] Remote consumer READ timeout for clearing request blocks (#20139)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Jul 8, 2025 at 09:56 UTC 71d1d75b7abb0e2dfa70443e0c3f111de0a6be22

3 files changed +115 -10

tests/v1/kv_connector/unit/test_nixl_connector.py

+74 -4

@@ -9,10 +9,13 @@ from unittest.mock import patch

10 import pytest

12 +from vllm import LLM

13 +from vllm.config import KVTransferConfig

14 from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (

15 KVConnectorRole, NixlAgentMetadata, NixlConnector, NixlConnectorMetadata,

16 NixlConnectorWorker)

17 from vllm.forward_context import ForwardContext

18 +from vllm.sampling_params import SamplingParams

20 from .utils import create_request, create_scheduler, create_vllm_config

@@ -41,9 +44,9 @@ def test_basic_interface():

44 assert kv_connector_metadata is not None

45 assert isinstance(kv_connector_metadata, NixlConnectorMetadata)

44 - assert len(kv_connector_metadata.requests) == 1

45 - assert request_id in kv_connector_metadata.requests

46 - req_meta = kv_connector_metadata.requests[request_id]

47 + assert len(kv_connector_metadata.reqs_to_recv) == 1

48 + assert request_id in kv_connector_metadata.reqs_to_recv

49 + req_meta = kv_connector_metadata.reqs_to_recv[request_id]

51 for block_id, block in zip(

52 req_meta.local_block_ids, scheduler.kv_cache_manager.coordinator.

@@ -78,7 +81,7 @@ def test_prompt_less_than_block_size():

81 kv_connector_metadata = scheduler_output.kv_connector_metadata

82 assert kv_connector_metadata is not None

83 assert isinstance(kv_connector_metadata, NixlConnectorMetadata)

81 - assert len(kv_connector_metadata.requests) == 0

84 + assert len(kv_connector_metadata.reqs_to_recv) == 0

86 # This request should be scheduled regularly.

87 assert len(scheduler_output.scheduled_new_reqs) == 1

@@ -371,3 +374,70 @@ class TestNixlHandshake:

374 if cnt_finished_reqs == total_reqs:

375 return

376 raise TimeoutError("Took too long to complete async handshake.")

377 +

378 +

379 +@patch(

380 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

381 + FakeNixlWrapper)

382 +def test_abort_timeout_on_prefiller(monkeypatch):

383 + """

384 + Test lifecycle of an aborted Remote Prefill request hitting the timeout.

385 + -----> P

386 + | {process request}

387 + <-\--- | {result is NOT delivered, eg proxy is down}

388 + |

389 + |

390 + | {eventually free blocks}

391 + """

392 + model_name = "Qwen/Qwen3-0.6B"

393 + kv_transfer_config = KVTransferConfig(

394 + kv_connector="NixlConnector",

395 + kv_role="kv_both",

396 + )

397 + timeout = 6

398 + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")

399 + monkeypatch.setenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", str(timeout))

400 + llm = LLM(

401 + model=model_name,

402 + enforce_eager=True,

403 + gpu_memory_utilization=0.5,

404 + kv_transfer_config=kv_transfer_config,

405 + )

406 + remote_prefill_opts = {

407 + "do_remote_decode": True,

408 + "do_remote_prefill": False,

409 + "remote_engine_id": None,

410 + "remote_block_ids": None,

411 + "remote_host": None,

412 + "remote_port": None,

413 + }

414 + # Simulate sidecar request

415 + sampling_params = SamplingParams(

416 + temperature=0.0,

417 + max_tokens=1,

418 + extra_args={"kv_transfer_params": remote_prefill_opts})

419 + scheduler = llm.llm_engine.engine_core.engine_core.scheduler

420 + req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[

421 + 0].req_to_blocks

422 +

423 + padding = "Just making this request a little longer so that we're sure "

424 + "we're not hitting the small-request lower bound beneath which we don't "

425 + "actually trigger the whole kv transfer, but rather just recompute the "

426 + "blocks on D."

427 + _ = llm.generate([f"What is the capital of Japan? {padding}"],

428 + sampling_params)

429 +

430 + # Request finished but not freed

431 + assert '0' in scheduler.finished_req_ids and '0' in req_to_blocks

432 + # Some other request, 0 still not freed

433 + _ = llm.generate([f"What is the capital of Italy? {padding}"],

434 + sampling_params)

435 + assert '0' in req_to_blocks

436 + assert '1' in scheduler.finished_req_ids and '1' in req_to_blocks

437 +

438 + # Wait for timeout and trigger another scheduler loop

439 + time.sleep(timeout)

440 + _ = llm.generate([f"What is the capital of France? {padding}"],

441 + sampling_params)

442 + # Request-0 times out and is cleared!

443 + assert '0' not in req_to_blocks

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+32 -5

@@ -79,7 +79,8 @@ class ReqMeta:

79 class NixlConnectorMetadata(KVConnectorMetadata):

81 def __init__(self):

82 - self.requests: dict[ReqId, ReqMeta] = {}

82 + self.reqs_to_recv: dict[ReqId, ReqMeta] = {}

83 + self.reqs_to_send: dict[ReqId, float] = {}

85 def add_new_req(

86 self,

@@ -87,7 +88,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):

88 local_block_ids: list[int],

89 kv_transfer_params: dict[str, Any],

90 ):

90 - self.requests[request_id] = ReqMeta(

91 + self.reqs_to_recv[request_id] = ReqMeta(

92 local_block_ids=local_block_ids,

93 remote_block_ids=kv_transfer_params["remote_block_ids"],

94 remote_engine_id=kv_transfer_params["remote_engine_id"],

@@ -194,10 +195,12 @@ class NixlConnectorScheduler:

195 vllm_config.parallel_config.tensor_parallel_size)

196 logger.info("Initializing NIXL Scheduler %s", engine_id)

197

197 - # Requests that need to start recv.

198 + # Requests that need to start recv/send.

199 # New requests are added by update_state_after_alloc in

200 # the scheduler. Used to make metadata passed to Worker.

201 self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}

202 + # Reqs to send and their expiration time

203 + self._reqs_need_send: dict[ReqId, float] = {}

204

205 def get_num_new_matched_tokens(

206 self, request: "Request",

@@ -284,6 +287,9 @@ class NixlConnectorScheduler:

287 # Clear the list once workers start the transfers

288 self._reqs_need_recv.clear()

289

290 + meta.reqs_to_send = self._reqs_need_send

291 + self._reqs_need_send = {}

292 +

293 return meta

294

295 def request_finished(

@@ -325,6 +331,11 @@ class NixlConnectorScheduler:

331 # If prompt < block_size, no xfer so free blocks immediately.

332 delay_free_blocks = len(computed_block_ids) > 0

333

334 + if delay_free_blocks:

335 + # Prefill request on remote. It will be read from D upon completion

336 + self._reqs_need_send[request.request_id] = time.perf_counter(

337 + ) + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT

338 +

339 return delay_free_blocks, dict(

340 do_remote_prefill=True,

341 do_remote_decode=False,

@@ -394,6 +405,8 @@ class NixlConnectorWorker:

405 # In progress transfers.

406 # [req_id -> list[handle]]

407 self._recving_transfers = defaultdict[ReqId, list[Transfer]](list)

408 + # Track the expiration time of requests that are waiting to be sent.

409 + self._reqs_to_send: dict[ReqId, float] = {}

410

411 # Complete transfer tracker. Used by the rank 0 to track finished

412 # transactions on ranks 1 to N-1.

@@ -826,6 +839,16 @@ class NixlConnectorWorker:

839 "and %s requests done recving", self.tp_rank,

840 len(done_sending), len(done_recving))

841

842 + # Handle timeout to avoid stranding blocks on remote.

843 + now = time.perf_counter()

844 + while self._reqs_to_send:

845 + req_id, expires = next(iter(self._reqs_to_send.items()))

846 + # Sorted dict, oldest requests are put first so we can exit early.

847 + if now < expires:

848 + break

849 + del self._reqs_to_send[req_id]

850 + done_sending.add(req_id)

851 +

852 if self.world_size == 1:

853 return done_sending, done_recving

854

@@ -857,7 +880,7 @@ class NixlConnectorWorker:

880

881 all_done_sending: set[str] = set()

882 for req_id in list(self._done_sending_count.keys()):

860 - if self._done_sending_count[req_id] == self.world_size:

883 + if self._done_sending_count[req_id] >= self.world_size:

884 del self._done_sending_count[req_id]

885 all_done_sending.add(req_id)

886

@@ -887,6 +910,7 @@ class NixlConnectorWorker:

910 tp_ratio):

911 notified_req_ids.add(req_id)

912 del self.consumer_notification_counts_by_req[req_id]

913 + del self._reqs_to_send[req_id]

914 return notified_req_ids

915

916 def _pop_done_transfers(

@@ -921,7 +945,7 @@ class NixlConnectorWorker:

945 Start loading by triggering non-blocking nixl_xfer.

946 We check for these trnxs to complete in each step().

947 """

924 - for req_id, meta in metadata.requests.items():

948 + for req_id, meta in metadata.reqs_to_recv.items():

949 remote_engine_id = meta.remote_engine_id

950 logger.debug(

951 "start_load_kv for request %s from remote engine %s. "

@@ -943,6 +967,9 @@ class NixlConnectorWorker:

967 while not self._ready_requests.empty():

968 self._read_blocks_for_req(*self._ready_requests.get_nowait())

969

970 + # Add to requests that are waiting to be read and track expiration.

971 + self._reqs_to_send.update(metadata.reqs_to_send)

972 +

973 def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):

974 logger.debug(

975 "Remote agent %s available, calling _read_blocks for req %s",

vllm/envs.py

+9 -1

@@ -138,6 +138,7 @@ if TYPE_CHECKING:

138 VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: str = "NONE"

139 VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True

140 VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: Optional[int] = None

141 + VLLM_NIXL_ABORT_REQUEST_TIMEOUT: int = 120

142

143

144 def get_default_cache_root():

@@ -953,7 +954,14 @@ environment_variables: dict[str, Callable[[], Any]] = {

954 # generations on machines < 100 for compressed-tensors

955 # models

956 "VLLM_USE_NVFP4_CT_EMULATIONS":

956 - lambda: bool(int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0")))

957 + lambda: bool(int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0"))),

958 +

959 + # Time (in seconds) after which the KV cache on the producer side is

960 + # automatically cleared if no READ notification is received from the

961 + # consumer. This is only applicable when using NixlConnector in a

962 + # disaggregated decode-prefill setup.

963 + "VLLM_NIXL_ABORT_REQUEST_TIMEOUT":

964 + lambda: int(os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "120"))

965 }

966

967 # --8<-- [end:env-vars-definition]