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 3, 22:45 UTC
[Nixl][PD] Lease renewal TTL KV blocks on P (#41383)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
May 11, 2026 at 11:27 UTC
770e9bd6b3519e8a6fa5738877974af274ec946e
13 files changed
+462
-48
docs/features/nixl_connector_usage.md
+8
-2
@@ -126,9 +126,15 @@ python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \
126
- Set when prefiller and decoder are on different machines
127
- Connection info is passed via KVTransferParams from prefiller to decoder for handshake
128
129
-- `VLLM_NIXL_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefiller’s KV cache for a particular request. (Optional)
129
+- `kv_lease_duration` (via `kv_connector_extra_config`): Lease duration (in seconds) for the prefiller's KV cache blocks. (Optional)
130
+ - Default: 30
131
+ - 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.
132
+ - Example: `--kv-transfer-config '{"kv_connector_extra_config": {"kv_lease_duration": 60}}'`
133
+
134
+- `decoder_kv_blocks_ttl` (via `kv_connector_extra_config`): TTL (in seconds) for KV blocks cached on the decoder in bidirectional transfer mode. (Optional)
135
- Default: 480
131
- - 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.
136
+ - 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.
137
+ - Example: `--kv-transfer-config '{"kv_connector_extra_config": {"decoder_kv_blocks_ttl": 600}}'`
138
139
## Multi-Instance Setup
140
tests/v1/kv_connector/unit/test_multi_connector.py
+13
-6
@@ -262,9 +262,11 @@ def test_multi_example_connector_consistency():
262
263
events = get_connector_events()
264
# First event is set_xfer_handshake_metadata from initialization, then
265
+ # on_new_request when the request is enqueued, then
266
# get_num_new_matched_tokens and update_state_after_alloc from generate().
266
- assert events["storage1-SCHEDULER"][:4] == [
267
+ assert events["storage1-SCHEDULER"][:5] == [
268
"set_xfer_handshake_metadata",
269
+ "on_new_request",
270
"get_num_new_matched_tokens 0",
271
"update_state_after_alloc num_blocks=[0] 0",
272
"build_connector_meta",
@@ -281,8 +283,9 @@ def test_multi_example_connector_consistency():
283
"wait_for_layer_load",
284
"save_kv_layer",
285
]
284
- assert events["storage2-SCHEDULER"][:4] == [
286
+ assert events["storage2-SCHEDULER"][:5] == [
287
"set_xfer_handshake_metadata",
288
+ "on_new_request",
289
"get_num_new_matched_tokens 0",
290
"update_state_after_alloc num_blocks=[0] 0",
291
"build_connector_meta",
@@ -310,12 +313,14 @@ def test_multi_example_connector_consistency():
313
# connector so update_state_after_alloc will be with allocated blocks
314
# on that one but with zero blocks for others (first nonzero match is
315
# chosen).
313
- assert events["storage1-SCHEDULER"][:3] == [
316
+ assert events["storage1-SCHEDULER"][:4] == [
317
+ "on_new_request",
318
"get_num_new_matched_tokens 0",
319
"update_state_after_alloc num_blocks=[7] 96",
320
"build_connector_meta",
321
]
318
- assert events["storage2-SCHEDULER"][:3] == [
322
+ assert events["storage2-SCHEDULER"][:4] == [
323
+ "on_new_request",
324
"get_num_new_matched_tokens 0",
325
"update_state_after_alloc num_blocks=[0] 0",
326
"build_connector_meta",
@@ -336,12 +341,14 @@ def test_multi_example_connector_consistency():
341
# return 0 from the first connector, but the second connector should have
342
# a hit, so update_state_after_alloc will only be called with allocated
343
# blocks for the second connector.
339
- assert events["storage1-SCHEDULER"][:3] == [
344
+ assert events["storage1-SCHEDULER"][:4] == [
345
+ "on_new_request",
346
"get_num_new_matched_tokens 0",
347
"update_state_after_alloc num_blocks=[0] 0",
348
"build_connector_meta",
349
]
344
- assert events["storage2-SCHEDULER"][:3] == [
350
+ assert events["storage2-SCHEDULER"][:4] == [
351
+ "on_new_request",
352
"get_num_new_matched_tokens 0",
353
"update_state_after_alloc num_blocks=[7] 96",
354
"build_connector_meta",
tests/v1/kv_connector/unit/test_nixl_connector.py
+2
-4
@@ -1374,9 +1374,11 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
1374
| {eventually free blocks}
1375
"""
1376
model_name = "Qwen/Qwen3-0.6B"
1377
+ timeout = 6
1378
kv_transfer_config = KVTransferConfig(
1379
kv_connector="NixlConnector",
1380
kv_role="kv_both",
1381
+ kv_connector_extra_config={"kv_lease_duration": timeout},
1382
)
1383
llm_kwargs = {
1384
"model": model_name,
@@ -1386,9 +1388,7 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
1388
"distributed_executor_backend": distributed_executor_backend,
1389
}
1390
1389
- timeout = 6
1391
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
1391
- monkeypatch.setenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", str(timeout))
1392
1393
def run_test_and_cleanup():
1394
llm = LLM(**llm_kwargs)
@@ -1403,8 +1403,6 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
1403
runtime_env = {
1404
"working_dir": working_dir, # ship fake nixl package
1405
"env_vars": {
1406
- "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": str(timeout),
1407
- # TODO: for ray to carry over, remove once we set
1406
"NIXL_TELEMETRY_ENABLE": "1",
1407
},
1408
}
tests/v1/kv_connector/unit/test_nixl_heartbeat.py
new
+165
@@ -0,0 +1,165 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the scheduler-driven heartbeat / lease-renewal system."""
4
+
5
+import time
6
+from unittest.mock import MagicMock
7
+
8
+import pytest
9
+
10
+from vllm.v1.outputs import KVConnectorOutput
11
+
12
+from .utils import create_request, make_nixl_scheduler
13
+
14
+_ENGINE_A = "my-engine-id"
15
+
16
+
17
+def _sched(kv_lease_duration: int = 30):
18
+ return make_nixl_scheduler(heartbeat=True, kv_lease_duration=kv_lease_duration)
19
+
20
+
21
+def _req(request_id: int = 1):
22
+ return create_request(request_id=request_id, do_remote_prefill=True)
23
+
24
+
25
+def _worker_stub():
26
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
27
+ NixlConnectorWorker,
28
+ )
29
+
30
+ w = object.__new__(NixlConnectorWorker)
31
+ w._reqs_to_send = {}
32
+ w._lease_extension = 20
33
+ return w
34
+
35
+
36
+# ===================================================================
37
+# Scheduler: on_new_request
38
+# ===================================================================
39
+
40
+
41
+def test_on_new_request_tracks_and_groups():
42
+ """Add two reqs to same engine, one to another; verify grouping."""
43
+ s = _sched()
44
+ s.on_new_request(_req(1))
45
+ s.on_new_request(_req(2))
46
+
47
+ assert s._heartbeat_by_engine[_ENGINE_A].req_ids == {"prefill-1", "prefill-2"}
48
+ info = s._heartbeat_by_engine[_ENGINE_A]
49
+ assert (info.host, info.port, info.tp_size) == ("my-host", 1234, 1)
50
+ assert s._heartbeat_req_engine["id-1"] == (_ENGINE_A, "prefill-1")
51
+
52
+ # Different engine.
53
+ r3 = _req(3)
54
+ r3.kv_transfer_params["remote_engine_id"] = "engine-b"
55
+ s.on_new_request(r3)
56
+ assert len(s._heartbeat_by_engine) == 2
57
+
58
+
59
+@pytest.mark.parametrize(
60
+ "make_req",
61
+ [
62
+ lambda: create_request(request_id=2, do_remote_decode=True),
63
+ lambda: create_request(request_id=3), # no kv_transfer_params
64
+ ],
65
+ ids=["decode", "plain"],
66
+)
67
+def test_on_new_request_ignores_non_prefill(make_req):
68
+ s = _sched()
69
+ s.on_new_request(make_req())
70
+ assert len(s._heartbeat_by_engine) == 0
71
+
72
+
73
+# ===================================================================
74
+# Scheduler: _stop_heartbeat
75
+# ===================================================================
76
+
77
+
78
+def test_stop_heartbeat_partial_and_full():
79
+ """Stop one of two reqs on same engine, then stop the other."""
80
+ s = _sched()
81
+ s.on_new_request(_req(1))
82
+ s.on_new_request(_req(2))
83
+
84
+ s._stop_heartbeat("id-1")
85
+ assert s._heartbeat_by_engine[_ENGINE_A].req_ids == {"prefill-2"}
86
+ assert "id-1" not in s._heartbeat_req_engine
87
+
88
+ s._stop_heartbeat("id-2")
89
+ assert len(s._heartbeat_by_engine) == 0
90
+ assert len(s._heartbeat_req_engine) == 0
91
+
92
+
93
+# ===================================================================
94
+# Scheduler: build_connector_meta throttling
95
+# ===================================================================
96
+
97
+
98
+def test_build_connector_meta_heartbeat_throttling():
99
+ # kv_lease_duration=30 => _heartbeat_interval = 30 // 6 = 5
100
+ s = _sched(kv_lease_duration=30)
101
+ s.on_new_request(_req(1))
102
+
103
+ # Ensure the first call triggers by placing last_heartbeat far in the past.
104
+ s._last_heartbeat_time = time.perf_counter() - 10
105
+ meta1 = s.build_connector_meta(MagicMock())
106
+ assert _ENGINE_A in meta1.heartbeat_by_engine
107
+
108
+ # Immediate second call is throttled (< 5s since last).
109
+ meta2 = s.build_connector_meta(MagicMock())
110
+ assert len(meta2.heartbeat_by_engine) == 0
111
+
112
+
113
+# ===================================================================
114
+# Scheduler: cleanup paths (update_connector_output / request_finished)
115
+# ===================================================================
116
+
117
+
118
+def test_update_connector_output_stops_heartbeat():
119
+ s = _sched()
120
+ s.on_new_request(_req(1))
121
+
122
+ s.update_connector_output(
123
+ KVConnectorOutput(
124
+ finished_sending=None,
125
+ finished_recving={"id-1"},
126
+ invalid_block_ids=set(),
127
+ )
128
+ )
129
+
130
+ assert len(s._heartbeat_by_engine) == 0
131
+ assert len(s._heartbeat_req_engine) == 0
132
+
133
+
134
+def test_request_finished_stops_heartbeat():
135
+ s = _sched()
136
+ r = _req(1)
137
+ s.on_new_request(r)
138
+
139
+ # Simulate update_state_after_alloc having consumed do_remote_prefill.
140
+ r.kv_transfer_params["do_remote_prefill"] = False
141
+ s.request_finished(r, block_ids=())
142
+
143
+ assert len(s._heartbeat_by_engine) == 0
144
+ assert len(s._heartbeat_req_engine) == 0
145
+
146
+
147
+# ===================================================================
148
+# Worker: _handle_heartbeat
149
+# ===================================================================
150
+
151
+
152
+def test_handle_heartbeat():
153
+ w = _worker_stub()
154
+ far_future = time.perf_counter() + 99999
155
+ w._reqs_to_send = {"req-a": 100.0, "req-b": far_future}
156
+
157
+ before = time.perf_counter()
158
+ w._handle_heartbeat("req-a,req-b,req-unknown")
159
+
160
+ # req-a: pushed forward to ~now+20.
161
+ assert w._reqs_to_send["req-a"] >= before + 20
162
+ # req-b: already far out, max() keeps it.
163
+ assert w._reqs_to_send["req-b"] >= far_future
164
+ # req-unknown: not added.
165
+ assert "req-unknown" not in w._reqs_to_send
tests/v1/kv_connector/unit/utils.py
+28
-2
@@ -223,6 +223,7 @@ def create_request(
223
remote_block_ids=list(range(num_remote_blocks)),
224
remote_host="my-host",
225
remote_port=1234,
226
+ tp_size=1,
227
)
228
229
max_tokens = 1 if do_remote_decode else max_tokens
@@ -482,10 +483,16 @@ def make_kv_cache_config(
483
)
484
485
485
-def make_nixl_scheduler(has_mamba: bool = False, is_hma_required: bool = False):
486
+def make_nixl_scheduler(
487
+ has_mamba: bool = False,
488
+ is_hma_required: bool = False,
489
+ heartbeat: bool = False,
490
+ kv_lease_duration: int = 30,
491
+):
492
"""Create a NixlConnectorScheduler via __new__ (skipping __init__).
493
488
- Only sets the two flags needed by the N-1 prefill logic.
494
+ Only sets the flags needed by the tests. When *heartbeat=True* the
495
+ scheduler-side heartbeat bookkeeping fields are also initialised.
496
"""
497
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
498
NixlConnectorScheduler,
@@ -494,4 +501,23 @@ def make_nixl_scheduler(has_mamba: bool = False, is_hma_required: bool = False):
501
sched = object.__new__(NixlConnectorScheduler)
502
sched._has_mamba = has_mamba
503
sched._is_hma_required = is_hma_required
504
+
505
+ if heartbeat:
506
+ sched._heartbeat_by_engine = {}
507
+ sched._heartbeat_req_engine = {}
508
+ sched._last_heartbeat_time = 0.0
509
+ sched._kv_lease_duration = kv_lease_duration
510
+ sched._heartbeat_interval = kv_lease_duration // 6
511
+ # Fields touched by build_connector_meta / request_finished:
512
+ sched._reqs_need_recv = {}
513
+ sched._reqs_need_send = {}
514
+ sched._reqs_in_batch = set()
515
+ sched._reqs_not_processed = set()
516
+ sched._reqs_need_save = {}
517
+ sched.use_host_buffer = False
518
+ sched.engine_id = "test-engine"
519
+ sched.side_channel_host = "localhost"
520
+ sched.side_channel_port = 5555
521
+ sched.blocks_per_sw = []
522
+ sched.is_bidirectional_kv_xfer_enabled = False
523
return sched
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+8
@@ -510,6 +510,14 @@ class KVConnectorBase_V1(ABC):
510
"""
511
pass
512
513
+ def on_new_request(self, request: "Request") -> None:
514
+ """Called by the scheduler when a new request is added.
515
+
516
+ Connectors can override this to inspect the request and perform
517
+ bookkeeping. The default implementation is a no-op.
518
+ """
519
+ return
520
+
521
def update_connector_output(self, connector_output: KVConnectorOutput):
522
"""
523
Update KVConnector state from worker-side connectors output.
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+4
@@ -389,6 +389,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
389
# Call with empty blocks for other connectors.
390
c.update_state_after_alloc(request, empty_blocks, 0)
391
392
+ def on_new_request(self, request: "Request") -> None:
393
+ for c in self._connectors:
394
+ c.on_new_request(request)
395
+
396
def build_connector_meta(
397
self, scheduler_output: SchedulerOutput
398
) -> MultiKVConnectorMetadata:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+9
@@ -44,6 +44,7 @@ from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
44
from vllm.v1.attention.backends.utils import get_kv_cache_layout
45
from vllm.v1.core.sched.output import SchedulerOutput
46
from vllm.v1.kv_cache_interface import MambaSpec
47
+from vllm.v1.outputs import KVConnectorOutput
48
49
if TYPE_CHECKING:
50
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
@@ -156,6 +157,14 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
157
assert self.connector_scheduler is not None
158
return self.connector_scheduler.build_connector_meta(scheduler_output)
159
160
+ def on_new_request(self, request: "Request") -> None:
161
+ assert self.connector_scheduler is not None
162
+ self.connector_scheduler.on_new_request(request)
163
+
164
+ def update_connector_output(self, connector_output: KVConnectorOutput):
165
+ assert self.connector_scheduler is not None
166
+ self.connector_scheduler.update_connector_output(connector_output)
167
+
168
def request_finished(
169
self,
170
request: "Request",
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+15
-2
@@ -6,7 +6,7 @@ from dataclasses import dataclass
6
from typing import Any
7
8
from vllm.config import VllmConfig
9
-from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
9
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds, EngineId
10
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
11
KVConnectorHandshakeMetadata,
12
KVConnectorMetadata,
@@ -33,8 +33,9 @@ GET_META_MSG = b"get_meta_msg"
33
# 1: Initial version with compatibility checking
34
# 2: Add remote_request_id to kv_transfer_params
35
# 3: Add physical_blocks_per_logical_kv_block to NixlAgentMetadata
36
+# 4: Add KV block lease renewal through heartbeats
37
#
37
-NIXL_CONNECTOR_VERSION: int = 3
38
+NIXL_CONNECTOR_VERSION: int = 4
39
40
41
@dataclass
@@ -133,6 +134,16 @@ def compute_nixl_compatibility_hash(
134
return compat_hash
135
136
137
+@dataclass
138
+class HeartbeatInfo:
139
+ """Heartbeat data for a single remote engine, sent from D worker to P."""
140
+
141
+ req_ids: set[ReqId]
142
+ host: str
143
+ port: int
144
+ tp_size: int
145
+
146
+
147
@dataclass
148
class RemoteMeta:
149
block_ids: BlockIds
@@ -158,6 +169,8 @@ class NixlConnectorMetadata(KVConnectorMetadata):
169
self.reqs_to_send: dict[ReqId, float] = {}
170
self.reqs_in_batch: set[ReqId] = set()
171
self.reqs_not_processed: set[ReqId] = set()
172
+ # Heartbeat data grouped by remote engine, sent by D worker to P.
173
+ self.heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {}
174
175
def _add_new_req(
176
self,
vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py
+92
-4
@@ -21,6 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
21
)
22
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
23
GET_META_MSG,
24
+ HeartbeatInfo,
25
NixlConnectorMetadata,
26
NixlHandshakePayload,
27
ReqId,
@@ -41,6 +42,7 @@ if TYPE_CHECKING:
42
from vllm.config import VllmConfig
43
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
44
from vllm.v1.kv_cache_interface import KVCacheConfig
45
+ from vllm.v1.outputs import KVConnectorOutput
46
from vllm.v1.request import Request
47
48
logger = init_logger(__name__)
@@ -65,6 +67,13 @@ class NixlConnectorScheduler:
67
+ vllm_config.parallel_config.data_parallel_index
68
)
69
assert vllm_config.kv_transfer_config is not None
70
+ self._kv_lease_duration: int = (
71
+ vllm_config.kv_transfer_config.get_from_extra_config(
72
+ "kv_lease_duration", 30
73
+ )
74
+ )
75
+ # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface.
76
+ self._heartbeat_interval = self._kv_lease_duration // 6
77
if current_platform.device_type == "cpu":
78
self.use_host_buffer = False
79
else:
@@ -104,6 +113,13 @@ class NixlConnectorScheduler:
113
# remote prefill or aborted.
114
self._reqs_not_processed: set[ReqId] = set()
115
116
+ # Heartbeat tracking: requests needing periodic lease-renewal heartbeats to
117
+ # remote P-side, stored as ready-to-send HeartbeatInfo grouped by remote engine
118
+ self._heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {}
119
+ # Reverse lookup: local req_id -> (engine_id, remote_req_id) for O(1) removal
120
+ self._heartbeat_req_engine: dict[ReqId, tuple[EngineId, ReqId]] = {}
121
+ self._last_heartbeat_time: float = 0.0
122
+
123
# Gather Sliding Window sizes for each kv cache group (if any) in number of
124
# blocks per KV cache group. This is used to clip the local attention window.
125
sw_sizes_tokens: list[tuple[int, int]] = [
@@ -135,12 +151,19 @@ class NixlConnectorScheduler:
151
"bidirectional_kv_xfer", False
152
)
153
)
154
+ self.decoder_kv_blocks_ttl = (
155
+ vllm_config.kv_transfer_config.get_from_extra_config(
156
+ "decoder_kv_blocks_ttl", 480
157
+ )
158
+ )
159
160
if self.is_bidirectional_kv_xfer_enabled and self.kv_recompute_threshold > 0:
161
logger.info(
162
"Bidirectional KV transfer is enabled and the kv "
142
- "recompute threshold is set to %d tokens",
163
+ "recompute threshold is set to %d tokens."
164
+ "KV blocks on D are released after a TTL of %d seconds.",
165
self.kv_recompute_threshold,
166
+ self.decoder_kv_blocks_ttl,
167
)
168
169
def shutdown(self):
@@ -149,6 +172,50 @@ class NixlConnectorScheduler:
172
self._nixl_handshake_listener_t.join()
173
self._nixl_handshake_listener_t = None
174
175
+ def on_new_request(self, request: "Request") -> None:
176
+ """Track a request that may need heartbeats."""
177
+ params = request.kv_transfer_params
178
+ # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are
179
+ # effectively disabled for Bidirectional KV transfer.
180
+ if params is None or not params.get("do_remote_prefill"):
181
+ return
182
+ # Only track if all required remote fields are present.
183
+ remote_engine_id = params.get("remote_engine_id")
184
+ remote_request_id = params.get("remote_request_id")
185
+ host = params.get("remote_host")
186
+ port = params.get("remote_port")
187
+ tp_size = params.get("tp_size")
188
+ if (
189
+ remote_engine_id is None
190
+ or remote_request_id is None
191
+ or host is None
192
+ or port is None
193
+ or tp_size is None
194
+ ):
195
+ return
196
+ if remote_engine_id not in self._heartbeat_by_engine:
197
+ self._heartbeat_by_engine[remote_engine_id] = HeartbeatInfo(
198
+ req_ids=set(),
199
+ host=host,
200
+ port=port,
201
+ tp_size=tp_size,
202
+ )
203
+ self._heartbeat_by_engine[remote_engine_id].req_ids.add(remote_request_id)
204
+ self._heartbeat_req_engine[request.request_id] = (
205
+ remote_engine_id,
206
+ remote_request_id,
207
+ )
208
+
209
+ def _stop_heartbeat(self, req_id: ReqId) -> None:
210
+ """Remove *req_id* from heartbeat tracking (if tracked)."""
211
+ if key := self._heartbeat_req_engine.pop(req_id, None):
212
+ engine_id, remote_id = key
213
+ if info := self._heartbeat_by_engine.get(engine_id):
214
+ info.req_ids.discard(remote_id)
215
+ if not info.req_ids:
216
+ # Clean up empty engines so we don't leak a key when remote dies.
217
+ del self._heartbeat_by_engine[engine_id]
218
+
219
def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
220
"""
221
Clip the number of blocks to the sliding window size for each kv cache group
@@ -490,6 +557,13 @@ class NixlConnectorScheduler:
557
meta.reqs_in_batch = self._reqs_in_batch
558
meta.reqs_not_processed = self._reqs_not_processed
559
560
+ # Package heartbeats, throttled by heartbeat_interval.
561
+ if self._heartbeat_by_engine:
562
+ now = time.perf_counter()
563
+ if now - self._last_heartbeat_time >= self._heartbeat_interval:
564
+ self._last_heartbeat_time = now
565
+ meta.heartbeat_by_engine = self._heartbeat_by_engine
566
+
567
# Clear the list once workers start the transfers
568
self._reqs_need_recv.clear()
569
self._reqs_in_batch = set()
@@ -498,6 +572,11 @@ class NixlConnectorScheduler:
572
573
return meta
574
575
+ def update_connector_output(self, connector_output: "KVConnectorOutput") -> None:
576
+ """Stop heartbeating for requests whose KV transfer completed."""
577
+ for req_id in connector_output.finished_recving or ():
578
+ self._stop_heartbeat(req_id)
579
+
580
def request_finished(
581
self,
582
request: "Request",
@@ -523,6 +602,10 @@ class NixlConnectorScheduler:
602
is_p_node = bool(params.get("do_remote_decode"))
603
is_d_node = not is_p_node
604
605
+ # Stop heartbeating for aborted requests that never reached finished_recving:
606
+ # normal path cleans up in update_connector_output.
607
+ self._stop_heartbeat(request.request_id)
608
+
609
if params.get("do_remote_prefill"):
610
# If do_remote_prefill is still True when the request is finished,
611
# update_state_after_alloc must not have been called (the request
@@ -556,14 +639,19 @@ class NixlConnectorScheduler:
639
remote_num_tokens = 0
640
if delay_free_blocks:
641
# Prefill request on remote. It will be read from D upon completion
642
+ request_kv_blocks_ttl = self._kv_lease_duration
643
+ if is_d_node:
644
+ # For blocks pinned on D, use a simpler timeout for now instead of a
645
+ # lease mechanism as turn2 request is client-driven.
646
+ request_kv_blocks_ttl = self.decoder_kv_blocks_ttl
647
logger.debug(
648
"NIXLConnector request_finished(%s) waiting for %d seconds "
561
- "for remote decode to fetch blocks",
649
+ "before releasing blocks",
650
request.request_id,
563
- envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT,
651
+ request_kv_blocks_ttl,
652
)
653
self._reqs_need_send[request.request_id] = (
566
- time.perf_counter() + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT
654
+ time.perf_counter() + request_kv_blocks_ttl
655
)
656
# NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
657
# trimming down after allocating for the whole sequence length. Empty
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+116
-20
@@ -18,7 +18,6 @@ import numpy as np
18
import torch
19
import zmq
20
21
-from vllm import envs
21
from vllm.distributed.kv_transfer.kv_connector.utils import (
22
BlockIds,
23
EngineId,
@@ -217,6 +216,12 @@ class NixlConnectorWorker:
216
self.nixl_backends = vllm_config.kv_transfer_config.get_from_extra_config(
217
"backends", ["UCX"]
218
)
219
+ kv_lease_duration: int = vllm_config.kv_transfer_config.get_from_extra_config(
220
+ "kv_lease_duration", 30
221
+ )
222
+ # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface.
223
+ self._lease_extension = kv_lease_duration * 2 // 3
224
+
225
self._is_hma_required = (
226
not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
227
and any(
@@ -688,23 +693,38 @@ class NixlConnectorWorker:
693
stacklevel=2,
694
)
695
691
- def _background_nixl_handshake(
692
- self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta
693
- ):
694
- # Do NIXL handshake in background and add to _ready_requests when done.
695
- fut = self._handshake_futures.get(remote_engine_id)
696
- if fut is None:
697
- assert meta.remote is not None
696
+ def _ensure_handshake(
697
+ self,
698
+ engine_id: EngineId,
699
+ host: str,
700
+ port: int,
701
+ tp_size: int,
702
+ ) -> Future[dict[int, str]] | None:
703
+ """
704
+ Ensure a handshake is in-flight (or already done) for *engine_id*.
705
+
706
+ Returns the ``Future`` if a handshake is pending (or was just
707
+ started), or ``None`` if the handshake already completed
708
+ successfully. Callers can attach per-request callbacks to the
709
+ returned future.
710
+ Failures to handshake are logged and the request is marked as failed.
711
+ """
712
+ with self._handshake_lock:
713
+ if engine_id in self._remote_agents:
714
+ return None
715
+ fut = self._handshake_futures.get(engine_id)
716
+ if fut is not None:
717
+ return fut
718
fut = self._handshake_initiation_executor.submit(
719
self._nixl_handshake,
700
- meta.remote.host,
701
- meta.remote.port,
702
- meta.tp_size,
703
- remote_engine_id,
720
+ host,
721
+ port,
722
+ tp_size,
723
+ engine_id,
724
)
705
- self._handshake_futures[remote_engine_id] = fut
725
+ self._handshake_futures[engine_id] = fut
726
707
- def done_callback(f: Future[dict[int, str]], eid=remote_engine_id):
727
+ def done_callback(f: Future[dict[int, str]], eid=engine_id):
728
with self._handshake_lock:
729
del self._handshake_futures[eid]
730
try:
@@ -718,15 +738,30 @@ class NixlConnectorWorker:
738
)
739
740
fut.add_done_callback(done_callback)
741
+ return fut
742
722
- # check handshake success before proceeding with request
743
+ def _background_nixl_handshake(
744
+ self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta
745
+ ):
746
+ # Do NIXL handshake in background and add to _ready_requests when done.
747
+ assert meta.remote is not None
748
+ fut = self._ensure_handshake(
749
+ remote_engine_id,
750
+ meta.remote.host,
751
+ meta.remote.port,
752
+ meta.tp_size,
753
+ )
754
+ if fut is None:
755
+ # Already handshaked — only happens if caller does not pre-check.
756
+ self._ready_requests.put((req_id, meta))
757
+ return
758
+
759
+ # Check handshake success before proceeding with request.
760
def request_ready(f: Future[Any], entry=(req_id, meta)):
761
try:
725
- # check if handshake succeeded
762
f.result()
763
self._ready_requests.put(entry)
764
except Exception as e:
729
- # handshake failed - mark blocks as invalid
765
self._log_failure(
766
failure_type="handshake_failed",
767
req_id=req_id,
@@ -1721,10 +1756,9 @@ class NixlConnectorWorker:
1756
self.xfer_stats.record_kv_expired_req()
1757
logger.warning(
1758
"Releasing expired KV blocks for request %s which were "
1724
- "retrieved by %d decode worker(s) within %d seconds.",
1759
+ "retrieved by %d remote worker(s) before lease expired.",
1760
req_id,
1761
count,
1727
- envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT,
1762
)
1763
self._reqs_to_process.remove(req_id)
1764
del self._reqs_to_send[req_id]
@@ -1737,12 +1771,22 @@ class NixlConnectorWorker:
1771
Get req_ids which got a remote xfer message. When multiple consumers
1772
are reading from the same producer (heterogeneous TP scenario), wait
1773
for all consumers to be done pulling.
1774
+
1775
+ Also handles heartbeat notifications ("HB:req1,req2,...") by
1776
+ extending the lease on the referenced requests.
1777
"""
1778
assert self.transfer_topo is not None
1779
notified_req_ids: set[str] = set()
1780
for notifs in self.nixl_wrapper.get_new_notifs().values():
1781
for notif in notifs:
1745
- req_id, tp_size = notif.decode("utf-8").rsplit(":", 1)
1782
+ msg = notif.decode("utf-8")
1783
+
1784
+ # Handle heartbeat messages from D-side.
1785
+ if msg.startswith("HB:"):
1786
+ self._handle_heartbeat(msg[3:])
1787
+ continue
1788
+
1789
+ req_id, tp_size = msg.rsplit(":", 1)
1790
if (
1791
req_id not in self._reqs_to_send
1792
and req_id not in self._reqs_to_process
@@ -1777,6 +1821,27 @@ class NixlConnectorWorker:
1821
self._reqs_to_send.pop(req_id, None)
1822
return notified_req_ids
1823
1824
+ def _handle_heartbeat(self, payload: str) -> None:
1825
+ """Extend leases for requests referenced in a heartbeat.
1826
+
1827
+ Args:
1828
+ payload: comma-separated P-side request IDs, e.g.
1829
+ "req_abc,req_def".
1830
+ """
1831
+ new_expiry = time.perf_counter() + self._lease_extension
1832
+ for req_id in payload.split(","):
1833
+ if req_id in self._reqs_to_send:
1834
+ old = self._reqs_to_send[req_id]
1835
+ self._reqs_to_send[req_id] = max(old, new_expiry)
1836
+ logger.debug(
1837
+ "Heartbeat extended lease for request %s "
1838
+ "by %ds (old_expiry=%.1f, new_expiry=%.1f)",
1839
+ req_id,
1840
+ self._lease_extension,
1841
+ old,
1842
+ new_expiry,
1843
+ )
1844
+
1845
def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]:
1846
"""
1847
Pop completed xfers by checking for DONE state.
@@ -1897,6 +1962,37 @@ class NixlConnectorWorker:
1962
if req_id in self._reqs_to_process:
1963
self._reqs_to_send[req_id] = expiration_time
1964
1965
+ # Send heartbeats to P-side engines to keep KV blocks alive while
1966
+ # requests sit in the D scheduler WAITING queue.
1967
+ self._send_heartbeats(metadata)
1968
+
1969
+ def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
1970
+ """
1971
+ Send heartbeat notifications to remote engines, extending lease on KV blocks.
1972
+ """
1973
+ for engine_id, hb_info in metadata.heartbeat_by_engine.items():
1974
+ # Proactive handshake (this request may still be in waiting queue) so
1975
+ # the **next** heartbeat for this remote can go through.
1976
+ if (
1977
+ self._ensure_handshake(
1978
+ engine_id, hb_info.host, hb_info.port, hb_info.tp_size
1979
+ )
1980
+ is not None
1981
+ ):
1982
+ continue # handshake is still pending
1983
+
1984
+ # Build the heartbeat message: "HB:req1,req2,..."
1985
+ hb_msg = ("HB:" + ",".join(hb_info.req_ids)).encode()
1986
+ for agent_name in self._remote_agents[engine_id].values():
1987
+ try:
1988
+ self.nixl_wrapper.send_notif(agent_name, notif_msg=hb_msg)
1989
+ except Exception:
1990
+ logger.debug(
1991
+ "Failed to send heartbeat to engine %s",
1992
+ engine_id,
1993
+ exc_info=True,
1994
+ )
1995
+
1996
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
1997
assert meta.remote is not None and self.transfer_topo is not None
1998
engine_id = meta.remote.engine_id
vllm/envs.py
-8
@@ -202,7 +202,6 @@ if TYPE_CHECKING:
202
] = "NONE"
203
VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True
204
VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None
205
- VLLM_NIXL_ABORT_REQUEST_TIMEOUT: int = 480
205
VLLM_MORIIO_CONNECTOR_READ_MODE: bool = False
206
VLLM_MORIIO_QP_PER_TRANSFER: int = 1
207
VLLM_MORIIO_POST_BATCH_SIZE: int = -1
@@ -1465,13 +1464,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
1464
"VLLM_USE_NVFP4_CT_EMULATIONS": lambda: bool(
1465
int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0"))
1466
),
1468
- # Time (in seconds) after which the KV cache on the producer side is
1469
- # automatically cleared if no READ notification is received from the
1470
- # consumer. This is only applicable when using NixlConnector in a
1471
- # disaggregated decode-prefill setup.
1472
- "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": lambda: int(
1473
- os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "480")
1474
- ),
1467
# Controls the read mode for the Mori-IO connector
1468
"VLLM_MORIIO_CONNECTOR_READ_MODE": lambda: (
1469
os.getenv("VLLM_MORIIO_CONNECTOR_READ_MODE", "False").lower() in ("true", "1")
vllm/v1/core/sched/scheduler.py
+2
@@ -1681,6 +1681,8 @@ class Scheduler(SchedulerInterface):
1681
request.streaming_queue = deque()
1682
self._enqueue_waiting_request(request)
1683
self.requests[request.request_id] = request
1684
+ if self.connector is not None:
1685
+ self.connector.on_new_request(request)
1686
if self.log_stats:
1687
request.record_event(EngineCoreEventType.QUEUED)
1688