From 770e9bd6b3519e8a6fa5738877974af274ec946e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 11 May 2026 11:27:30 +0200 Subject: [PATCH] [Nixl][PD] Lease renewal TTL KV blocks on P (#41383) Signed-off-by: NickLucche --- docs/features/nixl_connector_usage.md | 10 +- .../kv_connector/unit/test_multi_connector.py | 19 +- .../kv_connector/unit/test_nixl_connector.py | 6 +- .../kv_connector/unit/test_nixl_heartbeat.py | 165 ++++++++++++++++++ tests/v1/kv_connector/unit/utils.py | 30 +++- .../kv_transfer/kv_connector/v1/base.py | 8 + .../kv_connector/v1/multi_connector.py | 4 + .../kv_connector/v1/nixl/connector.py | 9 + .../kv_connector/v1/nixl/metadata.py | 17 +- .../kv_connector/v1/nixl/scheduler.py | 96 +++++++++- .../kv_connector/v1/nixl/worker.py | 136 ++++++++++++--- vllm/envs.py | 8 - vllm/v1/core/sched/scheduler.py | 2 + 13 files changed, 462 insertions(+), 48 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_nixl_heartbeat.py diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index e7ee3a6400..960108b689 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -126,9 +126,15 @@ python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \ - Set when prefiller and decoder are on different machines - Connection info is passed via KVTransferParams from prefiller to decoder for handshake -- `VLLM_NIXL_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefiller’s KV cache for a particular request. (Optional) +- `kv_lease_duration` (via `kv_connector_extra_config`): Lease duration (in seconds) for the prefiller's KV cache blocks. (Optional) + - Default: 30 + - When a prefill request finishes, its KV blocks are held for this duration waiting for the decoder to read them. While the request is queued on the decoder, periodic heartbeats automatically extend the lease. If neither a heartbeat nor a read notification arrives before the lease expires, the blocks are freed. The heartbeat interval and extension amount are derived automatically from this value. + - Example: `--kv-transfer-config '{"kv_connector_extra_config": {"kv_lease_duration": 60}}'` + +- `decoder_kv_blocks_ttl` (via `kv_connector_extra_config`): TTL (in seconds) for KV blocks cached on the decoder in bidirectional transfer mode. (Optional) - Default: 480 - - If a request is aborted and the decoder has not yet read the KV-cache blocks through the nixl channel, the prefill instance will release its KV-cache blocks after this timeout to avoid holding them indefinitely. + - In bidirectional mode, the decoder caches KV blocks for multi-turn conversations. This TTL controls how long those blocks are held before being released. Unlike the prefiller lease, this TTL is not renewed via heartbeats. + - Example: `--kv-transfer-config '{"kv_connector_extra_config": {"decoder_kv_blocks_ttl": 600}}'` ## Multi-Instance Setup diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 51024fb921..f0aa4e260f 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -262,9 +262,11 @@ def test_multi_example_connector_consistency(): events = get_connector_events() # First event is set_xfer_handshake_metadata from initialization, then + # on_new_request when the request is enqueued, then # get_num_new_matched_tokens and update_state_after_alloc from generate(). - assert events["storage1-SCHEDULER"][:4] == [ + assert events["storage1-SCHEDULER"][:5] == [ "set_xfer_handshake_metadata", + "on_new_request", "get_num_new_matched_tokens 0", "update_state_after_alloc num_blocks=[0] 0", "build_connector_meta", @@ -281,8 +283,9 @@ def test_multi_example_connector_consistency(): "wait_for_layer_load", "save_kv_layer", ] - assert events["storage2-SCHEDULER"][:4] == [ + assert events["storage2-SCHEDULER"][:5] == [ "set_xfer_handshake_metadata", + "on_new_request", "get_num_new_matched_tokens 0", "update_state_after_alloc num_blocks=[0] 0", "build_connector_meta", @@ -310,12 +313,14 @@ def test_multi_example_connector_consistency(): # connector so update_state_after_alloc will be with allocated blocks # on that one but with zero blocks for others (first nonzero match is # chosen). - assert events["storage1-SCHEDULER"][:3] == [ + assert events["storage1-SCHEDULER"][:4] == [ + "on_new_request", "get_num_new_matched_tokens 0", "update_state_after_alloc num_blocks=[7] 96", "build_connector_meta", ] - assert events["storage2-SCHEDULER"][:3] == [ + assert events["storage2-SCHEDULER"][:4] == [ + "on_new_request", "get_num_new_matched_tokens 0", "update_state_after_alloc num_blocks=[0] 0", "build_connector_meta", @@ -336,12 +341,14 @@ def test_multi_example_connector_consistency(): # return 0 from the first connector, but the second connector should have # a hit, so update_state_after_alloc will only be called with allocated # blocks for the second connector. - assert events["storage1-SCHEDULER"][:3] == [ + assert events["storage1-SCHEDULER"][:4] == [ + "on_new_request", "get_num_new_matched_tokens 0", "update_state_after_alloc num_blocks=[0] 0", "build_connector_meta", ] - assert events["storage2-SCHEDULER"][:3] == [ + assert events["storage2-SCHEDULER"][:4] == [ + "on_new_request", "get_num_new_matched_tokens 0", "update_state_after_alloc num_blocks=[7] 96", "build_connector_meta", diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 1a8612cd3f..35370efa03 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1374,9 +1374,11 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend): | {eventually free blocks} """ model_name = "Qwen/Qwen3-0.6B" + timeout = 6 kv_transfer_config = KVTransferConfig( kv_connector="NixlConnector", kv_role="kv_both", + kv_connector_extra_config={"kv_lease_duration": timeout}, ) llm_kwargs = { "model": model_name, @@ -1386,9 +1388,7 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend): "distributed_executor_backend": distributed_executor_backend, } - timeout = 6 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") - monkeypatch.setenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", str(timeout)) def run_test_and_cleanup(): llm = LLM(**llm_kwargs) @@ -1403,8 +1403,6 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend): runtime_env = { "working_dir": working_dir, # ship fake nixl package "env_vars": { - "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": str(timeout), - # TODO: for ray to carry over, remove once we set "NIXL_TELEMETRY_ENABLE": "1", }, } diff --git a/tests/v1/kv_connector/unit/test_nixl_heartbeat.py b/tests/v1/kv_connector/unit/test_nixl_heartbeat.py new file mode 100644 index 0000000000..345e10621f --- /dev/null +++ b/tests/v1/kv_connector/unit/test_nixl_heartbeat.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the scheduler-driven heartbeat / lease-renewal system.""" + +import time +from unittest.mock import MagicMock + +import pytest + +from vllm.v1.outputs import KVConnectorOutput + +from .utils import create_request, make_nixl_scheduler + +_ENGINE_A = "my-engine-id" + + +def _sched(kv_lease_duration: int = 30): + return make_nixl_scheduler(heartbeat=True, kv_lease_duration=kv_lease_duration) + + +def _req(request_id: int = 1): + return create_request(request_id=request_id, do_remote_prefill=True) + + +def _worker_stub(): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + w = object.__new__(NixlConnectorWorker) + w._reqs_to_send = {} + w._lease_extension = 20 + return w + + +# =================================================================== +# Scheduler: on_new_request +# =================================================================== + + +def test_on_new_request_tracks_and_groups(): + """Add two reqs to same engine, one to another; verify grouping.""" + s = _sched() + s.on_new_request(_req(1)) + s.on_new_request(_req(2)) + + assert s._heartbeat_by_engine[_ENGINE_A].req_ids == {"prefill-1", "prefill-2"} + info = s._heartbeat_by_engine[_ENGINE_A] + assert (info.host, info.port, info.tp_size) == ("my-host", 1234, 1) + assert s._heartbeat_req_engine["id-1"] == (_ENGINE_A, "prefill-1") + + # Different engine. + r3 = _req(3) + r3.kv_transfer_params["remote_engine_id"] = "engine-b" + s.on_new_request(r3) + assert len(s._heartbeat_by_engine) == 2 + + +@pytest.mark.parametrize( + "make_req", + [ + lambda: create_request(request_id=2, do_remote_decode=True), + lambda: create_request(request_id=3), # no kv_transfer_params + ], + ids=["decode", "plain"], +) +def test_on_new_request_ignores_non_prefill(make_req): + s = _sched() + s.on_new_request(make_req()) + assert len(s._heartbeat_by_engine) == 0 + + +# =================================================================== +# Scheduler: _stop_heartbeat +# =================================================================== + + +def test_stop_heartbeat_partial_and_full(): + """Stop one of two reqs on same engine, then stop the other.""" + s = _sched() + s.on_new_request(_req(1)) + s.on_new_request(_req(2)) + + s._stop_heartbeat("id-1") + assert s._heartbeat_by_engine[_ENGINE_A].req_ids == {"prefill-2"} + assert "id-1" not in s._heartbeat_req_engine + + s._stop_heartbeat("id-2") + assert len(s._heartbeat_by_engine) == 0 + assert len(s._heartbeat_req_engine) == 0 + + +# =================================================================== +# Scheduler: build_connector_meta throttling +# =================================================================== + + +def test_build_connector_meta_heartbeat_throttling(): + # kv_lease_duration=30 => _heartbeat_interval = 30 // 6 = 5 + s = _sched(kv_lease_duration=30) + s.on_new_request(_req(1)) + + # Ensure the first call triggers by placing last_heartbeat far in the past. + s._last_heartbeat_time = time.perf_counter() - 10 + meta1 = s.build_connector_meta(MagicMock()) + assert _ENGINE_A in meta1.heartbeat_by_engine + + # Immediate second call is throttled (< 5s since last). + meta2 = s.build_connector_meta(MagicMock()) + assert len(meta2.heartbeat_by_engine) == 0 + + +# =================================================================== +# Scheduler: cleanup paths (update_connector_output / request_finished) +# =================================================================== + + +def test_update_connector_output_stops_heartbeat(): + s = _sched() + s.on_new_request(_req(1)) + + s.update_connector_output( + KVConnectorOutput( + finished_sending=None, + finished_recving={"id-1"}, + invalid_block_ids=set(), + ) + ) + + assert len(s._heartbeat_by_engine) == 0 + assert len(s._heartbeat_req_engine) == 0 + + +def test_request_finished_stops_heartbeat(): + s = _sched() + r = _req(1) + s.on_new_request(r) + + # Simulate update_state_after_alloc having consumed do_remote_prefill. + r.kv_transfer_params["do_remote_prefill"] = False + s.request_finished(r, block_ids=()) + + assert len(s._heartbeat_by_engine) == 0 + assert len(s._heartbeat_req_engine) == 0 + + +# =================================================================== +# Worker: _handle_heartbeat +# =================================================================== + + +def test_handle_heartbeat(): + w = _worker_stub() + far_future = time.perf_counter() + 99999 + w._reqs_to_send = {"req-a": 100.0, "req-b": far_future} + + before = time.perf_counter() + w._handle_heartbeat("req-a,req-b,req-unknown") + + # req-a: pushed forward to ~now+20. + assert w._reqs_to_send["req-a"] >= before + 20 + # req-b: already far out, max() keeps it. + assert w._reqs_to_send["req-b"] >= far_future + # req-unknown: not added. + assert "req-unknown" not in w._reqs_to_send diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 2c46538618..5db2f7e7a9 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -223,6 +223,7 @@ def create_request( remote_block_ids=list(range(num_remote_blocks)), remote_host="my-host", remote_port=1234, + tp_size=1, ) max_tokens = 1 if do_remote_decode else max_tokens @@ -482,10 +483,16 @@ def make_kv_cache_config( ) -def make_nixl_scheduler(has_mamba: bool = False, is_hma_required: bool = False): +def make_nixl_scheduler( + has_mamba: bool = False, + is_hma_required: bool = False, + heartbeat: bool = False, + kv_lease_duration: int = 30, +): """Create a NixlConnectorScheduler via __new__ (skipping __init__). - Only sets the two flags needed by the N-1 prefill logic. + Only sets the flags needed by the tests. When *heartbeat=True* the + scheduler-side heartbeat bookkeeping fields are also initialised. """ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( NixlConnectorScheduler, @@ -494,4 +501,23 @@ def make_nixl_scheduler(has_mamba: bool = False, is_hma_required: bool = False): sched = object.__new__(NixlConnectorScheduler) sched._has_mamba = has_mamba sched._is_hma_required = is_hma_required + + if heartbeat: + sched._heartbeat_by_engine = {} + sched._heartbeat_req_engine = {} + sched._last_heartbeat_time = 0.0 + sched._kv_lease_duration = kv_lease_duration + sched._heartbeat_interval = kv_lease_duration // 6 + # Fields touched by build_connector_meta / request_finished: + sched._reqs_need_recv = {} + sched._reqs_need_send = {} + sched._reqs_in_batch = set() + sched._reqs_not_processed = set() + sched._reqs_need_save = {} + sched.use_host_buffer = False + sched.engine_id = "test-engine" + sched.side_channel_host = "localhost" + sched.side_channel_port = 5555 + sched.blocks_per_sw = [] + sched.is_bidirectional_kv_xfer_enabled = False return sched diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index a2fa9fa6c6..17eb591dad 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -510,6 +510,14 @@ class KVConnectorBase_V1(ABC): """ pass + def on_new_request(self, request: "Request") -> None: + """Called by the scheduler when a new request is added. + + Connectors can override this to inspect the request and perform + bookkeeping. The default implementation is a no-op. + """ + return + def update_connector_output(self, connector_output: KVConnectorOutput): """ Update KVConnector state from worker-side connectors output. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index a340f313e0..aaf8f5fe6b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -389,6 +389,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): # Call with empty blocks for other connectors. c.update_state_after_alloc(request, empty_blocks, 0) + def on_new_request(self, request: "Request") -> None: + for c in self._connectors: + c.on_new_request(request) + def build_connector_meta( self, scheduler_output: SchedulerOutput ) -> MultiKVConnectorMetadata: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 53ad031a4c..187322b4ae 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -44,6 +44,7 @@ from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import MambaSpec +from vllm.v1.outputs import KVConnectorOutput if TYPE_CHECKING: from vllm.v1.core.kv_cache_manager import KVCacheBlocks @@ -156,6 +157,14 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA): assert self.connector_scheduler is not None return self.connector_scheduler.build_connector_meta(scheduler_output) + def on_new_request(self, request: "Request") -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.on_new_request(request) + + def update_connector_output(self, connector_output: KVConnectorOutput): + assert self.connector_scheduler is not None + self.connector_scheduler.update_connector_output(connector_output) + def request_finished( self, request: "Request", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index c56e373ba9..b9e3436f50 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -6,7 +6,7 @@ from dataclasses import dataclass from typing import Any from vllm.config import VllmConfig -from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds, EngineId from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorHandshakeMetadata, KVConnectorMetadata, @@ -33,8 +33,9 @@ GET_META_MSG = b"get_meta_msg" # 1: Initial version with compatibility checking # 2: Add remote_request_id to kv_transfer_params # 3: Add physical_blocks_per_logical_kv_block to NixlAgentMetadata +# 4: Add KV block lease renewal through heartbeats # -NIXL_CONNECTOR_VERSION: int = 3 +NIXL_CONNECTOR_VERSION: int = 4 @dataclass @@ -133,6 +134,16 @@ def compute_nixl_compatibility_hash( return compat_hash +@dataclass +class HeartbeatInfo: + """Heartbeat data for a single remote engine, sent from D worker to P.""" + + req_ids: set[ReqId] + host: str + port: int + tp_size: int + + @dataclass class RemoteMeta: block_ids: BlockIds @@ -158,6 +169,8 @@ class NixlConnectorMetadata(KVConnectorMetadata): self.reqs_to_send: dict[ReqId, float] = {} self.reqs_in_batch: set[ReqId] = set() self.reqs_not_processed: set[ReqId] = set() + # Heartbeat data grouped by remote engine, sent by D worker to P. + self.heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {} def _add_new_req( self, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py index 285e054ddd..b2122ed0d3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py @@ -21,6 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( GET_META_MSG, + HeartbeatInfo, NixlConnectorMetadata, NixlHandshakePayload, ReqId, @@ -41,6 +42,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request logger = init_logger(__name__) @@ -65,6 +67,13 @@ class NixlConnectorScheduler: + vllm_config.parallel_config.data_parallel_index ) assert vllm_config.kv_transfer_config is not None + self._kv_lease_duration: int = ( + vllm_config.kv_transfer_config.get_from_extra_config( + "kv_lease_duration", 30 + ) + ) + # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface. + self._heartbeat_interval = self._kv_lease_duration // 6 if current_platform.device_type == "cpu": self.use_host_buffer = False else: @@ -104,6 +113,13 @@ class NixlConnectorScheduler: # remote prefill or aborted. self._reqs_not_processed: set[ReqId] = set() + # Heartbeat tracking: requests needing periodic lease-renewal heartbeats to + # remote P-side, stored as ready-to-send HeartbeatInfo grouped by remote engine + self._heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {} + # Reverse lookup: local req_id -> (engine_id, remote_req_id) for O(1) removal + self._heartbeat_req_engine: dict[ReqId, tuple[EngineId, ReqId]] = {} + self._last_heartbeat_time: float = 0.0 + # Gather Sliding Window sizes for each kv cache group (if any) in number of # blocks per KV cache group. This is used to clip the local attention window. sw_sizes_tokens: list[tuple[int, int]] = [ @@ -135,12 +151,19 @@ class NixlConnectorScheduler: "bidirectional_kv_xfer", False ) ) + self.decoder_kv_blocks_ttl = ( + vllm_config.kv_transfer_config.get_from_extra_config( + "decoder_kv_blocks_ttl", 480 + ) + ) if self.is_bidirectional_kv_xfer_enabled and self.kv_recompute_threshold > 0: logger.info( "Bidirectional KV transfer is enabled and the kv " - "recompute threshold is set to %d tokens", + "recompute threshold is set to %d tokens." + "KV blocks on D are released after a TTL of %d seconds.", self.kv_recompute_threshold, + self.decoder_kv_blocks_ttl, ) def shutdown(self): @@ -149,6 +172,50 @@ class NixlConnectorScheduler: self._nixl_handshake_listener_t.join() self._nixl_handshake_listener_t = None + def on_new_request(self, request: "Request") -> None: + """Track a request that may need heartbeats.""" + params = request.kv_transfer_params + # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are + # effectively disabled for Bidirectional KV transfer. + if params is None or not params.get("do_remote_prefill"): + return + # Only track if all required remote fields are present. + remote_engine_id = params.get("remote_engine_id") + remote_request_id = params.get("remote_request_id") + host = params.get("remote_host") + port = params.get("remote_port") + tp_size = params.get("tp_size") + if ( + remote_engine_id is None + or remote_request_id is None + or host is None + or port is None + or tp_size is None + ): + return + if remote_engine_id not in self._heartbeat_by_engine: + self._heartbeat_by_engine[remote_engine_id] = HeartbeatInfo( + req_ids=set(), + host=host, + port=port, + tp_size=tp_size, + ) + self._heartbeat_by_engine[remote_engine_id].req_ids.add(remote_request_id) + self._heartbeat_req_engine[request.request_id] = ( + remote_engine_id, + remote_request_id, + ) + + def _stop_heartbeat(self, req_id: ReqId) -> None: + """Remove *req_id* from heartbeat tracking (if tracked).""" + if key := self._heartbeat_req_engine.pop(req_id, None): + engine_id, remote_id = key + if info := self._heartbeat_by_engine.get(engine_id): + info.req_ids.discard(remote_id) + if not info.req_ids: + # Clean up empty engines so we don't leak a key when remote dies. + del self._heartbeat_by_engine[engine_id] + def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds: """ Clip the number of blocks to the sliding window size for each kv cache group @@ -490,6 +557,13 @@ class NixlConnectorScheduler: meta.reqs_in_batch = self._reqs_in_batch meta.reqs_not_processed = self._reqs_not_processed + # Package heartbeats, throttled by heartbeat_interval. + if self._heartbeat_by_engine: + now = time.perf_counter() + if now - self._last_heartbeat_time >= self._heartbeat_interval: + self._last_heartbeat_time = now + meta.heartbeat_by_engine = self._heartbeat_by_engine + # Clear the list once workers start the transfers self._reqs_need_recv.clear() self._reqs_in_batch = set() @@ -498,6 +572,11 @@ class NixlConnectorScheduler: return meta + def update_connector_output(self, connector_output: "KVConnectorOutput") -> None: + """Stop heartbeating for requests whose KV transfer completed.""" + for req_id in connector_output.finished_recving or (): + self._stop_heartbeat(req_id) + def request_finished( self, request: "Request", @@ -523,6 +602,10 @@ class NixlConnectorScheduler: is_p_node = bool(params.get("do_remote_decode")) is_d_node = not is_p_node + # Stop heartbeating for aborted requests that never reached finished_recving: + # normal path cleans up in update_connector_output. + self._stop_heartbeat(request.request_id) + if params.get("do_remote_prefill"): # If do_remote_prefill is still True when the request is finished, # update_state_after_alloc must not have been called (the request @@ -556,14 +639,19 @@ class NixlConnectorScheduler: remote_num_tokens = 0 if delay_free_blocks: # Prefill request on remote. It will be read from D upon completion + request_kv_blocks_ttl = self._kv_lease_duration + if is_d_node: + # For blocks pinned on D, use a simpler timeout for now instead of a + # lease mechanism as turn2 request is client-driven. + request_kv_blocks_ttl = self.decoder_kv_blocks_ttl logger.debug( "NIXLConnector request_finished(%s) waiting for %d seconds " - "for remote decode to fetch blocks", + "before releasing blocks", request.request_id, - envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT, + request_kv_blocks_ttl, ) self._reqs_need_send[request.request_id] = ( - time.perf_counter() + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT + time.perf_counter() + request_kv_blocks_ttl ) # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones), # trimming down after allocating for the whole sequence length. Empty diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py index 18144f2ef5..a46367a9ab 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py @@ -18,7 +18,6 @@ import numpy as np import torch import zmq -from vllm import envs from vllm.distributed.kv_transfer.kv_connector.utils import ( BlockIds, EngineId, @@ -217,6 +216,12 @@ class NixlConnectorWorker: self.nixl_backends = vllm_config.kv_transfer_config.get_from_extra_config( "backends", ["UCX"] ) + kv_lease_duration: int = vllm_config.kv_transfer_config.get_from_extra_config( + "kv_lease_duration", 30 + ) + # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface. + self._lease_extension = kv_lease_duration * 2 // 3 + self._is_hma_required = ( not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager and any( @@ -688,23 +693,38 @@ class NixlConnectorWorker: stacklevel=2, ) - def _background_nixl_handshake( - self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta - ): - # Do NIXL handshake in background and add to _ready_requests when done. - fut = self._handshake_futures.get(remote_engine_id) - if fut is None: - assert meta.remote is not None + def _ensure_handshake( + self, + engine_id: EngineId, + host: str, + port: int, + tp_size: int, + ) -> Future[dict[int, str]] | None: + """ + Ensure a handshake is in-flight (or already done) for *engine_id*. + + Returns the ``Future`` if a handshake is pending (or was just + started), or ``None`` if the handshake already completed + successfully. Callers can attach per-request callbacks to the + returned future. + Failures to handshake are logged and the request is marked as failed. + """ + with self._handshake_lock: + if engine_id in self._remote_agents: + return None + fut = self._handshake_futures.get(engine_id) + if fut is not None: + return fut fut = self._handshake_initiation_executor.submit( self._nixl_handshake, - meta.remote.host, - meta.remote.port, - meta.tp_size, - remote_engine_id, + host, + port, + tp_size, + engine_id, ) - self._handshake_futures[remote_engine_id] = fut + self._handshake_futures[engine_id] = fut - def done_callback(f: Future[dict[int, str]], eid=remote_engine_id): + def done_callback(f: Future[dict[int, str]], eid=engine_id): with self._handshake_lock: del self._handshake_futures[eid] try: @@ -718,15 +738,30 @@ class NixlConnectorWorker: ) fut.add_done_callback(done_callback) + return fut - # check handshake success before proceeding with request + def _background_nixl_handshake( + self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta + ): + # Do NIXL handshake in background and add to _ready_requests when done. + assert meta.remote is not None + fut = self._ensure_handshake( + remote_engine_id, + meta.remote.host, + meta.remote.port, + meta.tp_size, + ) + if fut is None: + # Already handshaked — only happens if caller does not pre-check. + self._ready_requests.put((req_id, meta)) + return + + # Check handshake success before proceeding with request. def request_ready(f: Future[Any], entry=(req_id, meta)): try: - # check if handshake succeeded f.result() self._ready_requests.put(entry) except Exception as e: - # handshake failed - mark blocks as invalid self._log_failure( failure_type="handshake_failed", req_id=req_id, @@ -1721,10 +1756,9 @@ class NixlConnectorWorker: self.xfer_stats.record_kv_expired_req() logger.warning( "Releasing expired KV blocks for request %s which were " - "retrieved by %d decode worker(s) within %d seconds.", + "retrieved by %d remote worker(s) before lease expired.", req_id, count, - envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT, ) self._reqs_to_process.remove(req_id) del self._reqs_to_send[req_id] @@ -1737,12 +1771,22 @@ class NixlConnectorWorker: Get req_ids which got a remote xfer message. When multiple consumers are reading from the same producer (heterogeneous TP scenario), wait for all consumers to be done pulling. + + Also handles heartbeat notifications ("HB:req1,req2,...") by + extending the lease on the referenced requests. """ assert self.transfer_topo is not None notified_req_ids: set[str] = set() for notifs in self.nixl_wrapper.get_new_notifs().values(): for notif in notifs: - req_id, tp_size = notif.decode("utf-8").rsplit(":", 1) + msg = notif.decode("utf-8") + + # Handle heartbeat messages from D-side. + if msg.startswith("HB:"): + self._handle_heartbeat(msg[3:]) + continue + + req_id, tp_size = msg.rsplit(":", 1) if ( req_id not in self._reqs_to_send and req_id not in self._reqs_to_process @@ -1777,6 +1821,27 @@ class NixlConnectorWorker: self._reqs_to_send.pop(req_id, None) return notified_req_ids + def _handle_heartbeat(self, payload: str) -> None: + """Extend leases for requests referenced in a heartbeat. + + Args: + payload: comma-separated P-side request IDs, e.g. + "req_abc,req_def". + """ + new_expiry = time.perf_counter() + self._lease_extension + for req_id in payload.split(","): + if req_id in self._reqs_to_send: + old = self._reqs_to_send[req_id] + self._reqs_to_send[req_id] = max(old, new_expiry) + logger.debug( + "Heartbeat extended lease for request %s " + "by %ds (old_expiry=%.1f, new_expiry=%.1f)", + req_id, + self._lease_extension, + old, + new_expiry, + ) + def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: """ Pop completed xfers by checking for DONE state. @@ -1897,6 +1962,37 @@ class NixlConnectorWorker: if req_id in self._reqs_to_process: self._reqs_to_send[req_id] = expiration_time + # Send heartbeats to P-side engines to keep KV blocks alive while + # requests sit in the D scheduler WAITING queue. + self._send_heartbeats(metadata) + + def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None: + """ + Send heartbeat notifications to remote engines, extending lease on KV blocks. + """ + for engine_id, hb_info in metadata.heartbeat_by_engine.items(): + # Proactive handshake (this request may still be in waiting queue) so + # the **next** heartbeat for this remote can go through. + if ( + self._ensure_handshake( + engine_id, hb_info.host, hb_info.port, hb_info.tp_size + ) + is not None + ): + continue # handshake is still pending + + # Build the heartbeat message: "HB:req1,req2,..." + hb_msg = ("HB:" + ",".join(hb_info.req_ids)).encode() + for agent_name in self._remote_agents[engine_id].values(): + try: + self.nixl_wrapper.send_notif(agent_name, notif_msg=hb_msg) + except Exception: + logger.debug( + "Failed to send heartbeat to engine %s", + engine_id, + exc_info=True, + ) + def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): assert meta.remote is not None and self.transfer_topo is not None engine_id = meta.remote.engine_id diff --git a/vllm/envs.py b/vllm/envs.py index 0919957229..66e2a33bc1 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -202,7 +202,6 @@ if TYPE_CHECKING: ] = "NONE" VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None - VLLM_NIXL_ABORT_REQUEST_TIMEOUT: int = 480 VLLM_MORIIO_CONNECTOR_READ_MODE: bool = False VLLM_MORIIO_QP_PER_TRANSFER: int = 1 VLLM_MORIIO_POST_BATCH_SIZE: int = -1 @@ -1465,13 +1464,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_USE_NVFP4_CT_EMULATIONS": lambda: bool( int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0")) ), - # Time (in seconds) after which the KV cache on the producer side is - # automatically cleared if no READ notification is received from the - # consumer. This is only applicable when using NixlConnector in a - # disaggregated decode-prefill setup. - "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": lambda: int( - os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "480") - ), # Controls the read mode for the Mori-IO connector "VLLM_MORIIO_CONNECTOR_READ_MODE": lambda: ( os.getenv("VLLM_MORIIO_CONNECTOR_READ_MODE", "False").lower() in ("true", "1") diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 8aaeb39700..39e2f966ad 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1681,6 +1681,8 @@ class Scheduler(SchedulerInterface): request.streaming_queue = deque() self._enqueue_waiting_request(request) self.requests[request.request_id] = request + if self.connector is not None: + self.connector.on_new_request(request) if self.log_stats: request.record_event(EngineCoreEventType.QUEUED)