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 14, 05:30 UTC
[NIXL] Ignore abort on already-finished request (#25067)
Signed-off-by: Mark McLoughlin <markmc@redhat.com>
Mark McLoughlin committed
Oct 10, 2025 at 11:21 UTC
784c231151d3bdb361cc10a1753f2f28258f1600
4 files changed
+64
-16
tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py
+48
-1
@@ -43,6 +43,7 @@ def test_basic_lifecycle():
43
# STEP (1): Prefill.
44
# (1a): schedule()
45
scheduler_output = scheduler.schedule()
46
+ assert len(scheduler.requests) == 1
47
assert len(scheduler.running) == 1
48
assert len(scheduler_output.scheduled_new_reqs) == 1
49
@@ -67,6 +68,7 @@ def test_basic_lifecycle():
68
assert len(scheduler.waiting) == 0
69
70
# ... but blocks should not be freed.
71
+ assert len(scheduler.requests) == 1
72
blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
73
0
74
].req_to_blocks[request_id]
@@ -76,6 +78,7 @@ def test_basic_lifecycle():
78
# STEP (2): Send Finished to PB.
79
# (2a): schedule() - pass finished request to PB.
80
scheduler_output = scheduler.schedule()
81
+ assert len(scheduler.requests) == 1
82
assert len(scheduler.running) == 0
83
assert len(scheduler_output.finished_req_ids) == 1
84
assert request_id in scheduler_output.finished_req_ids
@@ -92,6 +95,7 @@ def test_basic_lifecycle():
95
# STEP (3): Finished sending.
96
# (3a): schedule() - pass finished request to PB.
97
scheduler_output = scheduler.schedule()
98
+ assert len(scheduler.requests) == 1
99
assert len(scheduler.running) == 0
100
assert len(scheduler_output.finished_req_ids) == 0
101
assert len(scheduler_output.scheduled_new_reqs) == 0
@@ -133,6 +137,7 @@ def test_short_prompt_lifecycle():
137
# STEP (1): Prefill.
138
# (1a): schedule()
139
scheduler_output = scheduler.schedule()
140
+ assert len(scheduler.requests) == 1
141
assert len(scheduler.running) == 1
142
assert len(scheduler_output.scheduled_new_reqs) == 1
143
@@ -178,7 +183,7 @@ def test_prefix_cache_lifecycle():
183
reqs=[request_normal], use_eos=True
184
)
185
scheduler.update_from_output(scheduler_output, model_runner_output)
181
- scheduler.schedule()
186
+ scheduler_output = scheduler.schedule()
187
scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
188
189
#####################
@@ -213,3 +218,45 @@ def test_prefix_cache_lifecycle():
218
)
219
scheduler.update_from_output(scheduler_output, model_runner_output)
220
assert_scheduler_empty(scheduler)
221
+
222
+
223
+def test_abort_during_kv_transfer():
224
+ """Test aborting request does not release blocks for remote decode."""
225
+
226
+ vllm_config = create_vllm_config()
227
+ scheduler = create_scheduler(vllm_config)
228
+
229
+ # Prime the KVCache.
230
+ BLOCK_SIZE = vllm_config.cache_config.block_size
231
+ NUM_EXTERNAL_FULL_BLOCKS = 2
232
+ NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
233
+
234
+ request = create_request(
235
+ request_id=1,
236
+ block_size=BLOCK_SIZE,
237
+ num_tokens=NUM_TOKENS,
238
+ do_remote_decode=True,
239
+ )
240
+
241
+ scheduler.add_request(request)
242
+ scheduler_output = scheduler.schedule()
243
+ model_runner_output = create_model_runner_output(reqs=[request])
244
+ scheduler.update_from_output(scheduler_output, model_runner_output)
245
+ scheduler_output = scheduler.schedule()
246
+ scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
247
+
248
+ # Request removed from PB but blocks should not be freed.
249
+ assert len(scheduler.requests) == 1
250
+
251
+ # Abort the request, and check the blocks are still not freed
252
+ scheduler.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED)
253
+ assert len(scheduler.requests) == 1
254
+
255
+ # Simulate a finished sending notification
256
+ scheduler_output = scheduler.schedule()
257
+ model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
258
+ model_runner_output.kv_connector_output = KVConnectorOutput(
259
+ finished_sending=[request.request_id]
260
+ )
261
+ scheduler.update_from_output(scheduler_output, model_runner_output)
262
+ assert_scheduler_empty(scheduler)
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+11
-6
@@ -14,11 +14,12 @@ The class provides the following primitives:
14
temporary buffer alloc by the CacheManager.
15
update_connector_output() - update KVConnector state after
16
output is received from worker-side connectors.
17
- request_finished() - called when a request is finished, with
18
- the computed kv cache blocks for the request.
19
- Returns whether KV cache should be freed now or will be
20
- freed asynchronously and optionally returns KV transfer
21
- params.
17
+ request_finished() - called once when a request is finished,
18
+ with the computed kv cache blocks for the request.
19
+ Returns whether KV cache should be freed now or if the
20
+ connector now assumes responsibility for freeing the
21
+ the blocks asynchronously. Also optionally returns KV
22
+ transfer params.
23
take_events() - returns new KV events that were collected
24
by the connector since the last call.
25
@@ -362,7 +363,11 @@ class KVConnectorBase_V1(ABC):
363
block_ids: list[int],
364
) -> tuple[bool, Optional[dict[str, Any]]]:
365
"""
365
- Called when a request has finished, before its blocks are freed.
366
+ Called exactly once when a request has finished, before its blocks are
367
+ freed.
368
+
369
+ The connector may assumes responsibility for freeing the the blocks
370
+ asynchronously by returning True.
371
372
Returns:
373
True if the request is being saved/sent asynchronously and blocks
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+2
@@ -1345,6 +1345,8 @@ class NixlConnectorWorker:
1345
# Remove all requests that are not to be processed (eg aborted).
1346
for req_id in metadata.reqs_not_processed:
1347
self._reqs_to_process.discard(req_id)
1348
+ # We should never get an abort after setting an expiry timer
1349
+ assert req_id not in self._reqs_to_send
1350
1351
# Add to requests that are waiting to be read and track expiration.
1352
for req_id, expiration_time in metadata.reqs_to_send.items():
vllm/v1/core/sched/scheduler.py
+3
-9
@@ -1187,7 +1187,7 @@ class Scheduler(SchedulerInterface):
1187
# First pass: collect requests to remove from queues
1188
for req_id in request_ids:
1189
request = self.requests.get(req_id)
1190
- if request is None:
1190
+ if request is None or request.is_finished():
1191
# Invalid request ID.
1192
continue
1193
@@ -1365,14 +1365,8 @@ class Scheduler(SchedulerInterface):
1365
self.finished_recving_kv_req_ids.add(req_id)
1366
for req_id in kv_connector_output.finished_sending or ():
1367
logger.debug("Finished sending KV transfer for request %s", req_id)
1368
- if req_id not in self.requests:
1369
- logger.warning(
1370
- "Got finished sending KV transfer for request %s,"
1371
- "but the request is already freed.",
1372
- req_id,
1373
- )
1374
- else:
1375
- self._free_blocks(self.requests[req_id])
1368
+ assert req_id in self.requests
1369
+ self._free_blocks(self.requests[req_id])
1370
1371
def _update_requests_with_invalid_blocks(
1372
self, requests: Iterable[Request], invalid_block_ids: set[int]