[Bugfix][NIXL] Avoid receive reports for notification-only requests (#57570)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
Nicolò Lucchesi committed
Sep 18, 2026 at 22:33 UTC
49e0427ca5183a56ccb112e5be1300feb2873ca9
4 files changed
+91
-2
.buildkite/test_areas/disaggregated.yaml
+3
@@ -311,6 +311,9 @@ steps:
311
PREFILLER_TP_SIZE: "2"
312
DECODER_TP_SIZE: "2"
313
MODEL_NAMES: "moonshotai/Kimi-Linear-48B-A3B-Instruct"
314
+ # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima.
315
+ PREFILL_BLOCK_SIZE: "2048"
316
+ DECODE_BLOCK_SIZE: "2048"
317
VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768"
318
commands:
319
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
+33
@@ -35,6 +35,39 @@ def scheduler():
35
return create_scheduler(vllm_config)
36
37
38
+@pytest.mark.parametrize("policy", ["fail", "recompute"])
39
+def test_failed_receive_completion_honors_failure_policy(policy):
40
+ """Failed receives still need completion to free blocks or resume computation."""
41
+ scheduler = create_scheduler(create_vllm_config(kv_load_failure_policy=policy))
42
+ request = create_request(num_tokens=3 * scheduler.block_size)
43
+ scheduler.add_request(request)
44
+ scheduler.connector = Mock()
45
+ scheduler.connector.get_num_new_matched_tokens.side_effect = [
46
+ (2 * scheduler.block_size, True),
47
+ (0, False),
48
+ ]
49
+ scheduler.connector.request_finished.return_value = (False, None)
50
+ scheduler.connector.take_events.return_value = ()
51
+ block_pool = scheduler.kv_cache_manager.block_pool
52
+ free_blocks_before = block_pool.get_num_free_blocks()
53
+ scheduler_output = scheduler.schedule()
54
+ assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
55
+
56
+ output = create_model_runner_output([], finished_recving={request.request_id})
57
+ assert output.kv_connector_output is not None
58
+ output.kv_connector_output.failed_recving = {request.request_id}
59
+ scheduler.update_from_output(scheduler_output, output)
60
+
61
+ if policy == "fail":
62
+ assert request.status == RequestStatus.FINISHED_ERROR
63
+ assert request.request_id not in scheduler.requests
64
+ assert block_pool.get_num_free_blocks() == free_blocks_before
65
+ else:
66
+ assert request.num_computed_tokens == 0
67
+ next_output = scheduler.schedule()
68
+ assert request.request_id in next_output.num_scheduled_tokens
69
+
70
+
71
@pytest.mark.parametrize(
72
"num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs",
73
[
tests/v1/kv_connector/unit/test_nixl_connector.py
+47
@@ -10,6 +10,7 @@ import textwrap
10
import time
11
import uuid
12
from collections import defaultdict
13
+from concurrent.futures import Future
14
from types import SimpleNamespace
15
from typing import Any, cast
16
from unittest.mock import MagicMock, patch
@@ -2899,6 +2900,52 @@ def test_empty_recv_is_reported_only_when_awaited(
2900
assert (request_id in done_recving) is awaiting_kvs
2901
2902
2903
+@pytest.mark.parametrize("is_hma", [False, True])
2904
+@pytest.mark.parametrize(
2905
+ "local_block_ids,awaiting_kvs",
2906
+ [((), False), (([],), False), (([],), True), (([1, 2, 3],), True)],
2907
+)
2908
+def test_handshake_failure_reports_only_awaited_recvs(
2909
+ recv_worker, is_hma, local_block_ids, awaiting_kvs
2910
+):
2911
+ """A failed cleanup handshake must not complete an already-aborted request."""
2912
+ worker = recv_worker
2913
+ worker._is_hma_required = is_hma
2914
+ worker._recving_metadata.clear()
2915
+ worker._remote_agents = {}
2916
+ worker._handshake_lock = contextlib.nullcontext()
2917
+ worker._ready_requests = queue.Queue()
2918
+ worker._reqs_to_process = set()
2919
+ worker.pcp_rank = 0
2920
+ metadata = NixlConnectorMetadata()
2921
+ metadata.add_new_req_to_recv(
2922
+ request_id="request",
2923
+ local_block_ids=local_block_ids,
2924
+ kv_transfer_params={
2925
+ "remote_block_ids": ([4, 5, 6],),
2926
+ "remote_engine_id": "prefill",
2927
+ "remote_request_id": "prefill-request",
2928
+ "remote_host": "localhost",
2929
+ "remote_port": 1234,
2930
+ },
2931
+ awaiting_kvs=awaiting_kvs,
2932
+ )
2933
+ handshake = Future[None]()
2934
+ handshake.set_exception(RuntimeError("handshake failed"))
2935
+ with patch.object(worker, "_ensure_handshake", return_value=handshake):
2936
+ worker.start_load_kv(metadata)
2937
+
2938
+ results = worker.get_transfer_results()
2939
+ expected = {"request"} if awaiting_kvs else set()
2940
+ assert results.finished_recving == results.failed_recving == expected
2941
+ assert worker.get_block_ids_with_load_errors() == (
2942
+ set(local_block_ids[0]) if awaiting_kvs and not is_hma else set()
2943
+ )
2944
+ assert not worker._recving_metadata
2945
+ assert not worker._recv_failures
2946
+ assert not worker._pending_recv_notifs
2947
+
2948
+
2949
@patch(
2950
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2951
FailingNixlWrapper,
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+8
-2
@@ -64,8 +64,14 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
64
len(meta.local_physical_block_ids),
65
len(meta.remote.block_ids),
66
)
67
- # always store metadata for failure recovery
68
- self._recving_metadata[req_id] = meta
67
+ # Full local hits and aborted cleanup only notify P; no recv is awaited.
68
+ # On a full local hit:
69
+ # - Notification failure must not fail the request: its KV is local.
70
+ # - Receive completion must not be reported: num_external_tokens == 0,
71
+ # so the scheduler never entered WAITING_FOR_REMOTE_KVS to begin with.
72
+ # Aborted cleanup requests have already been removed from the scheduler.
73
+ if meta.awaiting_kvs or any(meta.local_block_ids):
74
+ self._recving_metadata[req_id] = meta
75
if remote_engine_id not in self._remote_agents:
76
# Initiate handshake with remote engine to exchange metadata.
77
with self._handshake_lock: