test_nixl_connector.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 4,260 lines 163 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4	import contextlib
5	import inspect
6	import os
7	import queue
8	import tempfile
9	import textwrap
10	import threading
11	import time
12	import uuid
13	from collections import defaultdict
14	from concurrent.futures import Future
15	from types import SimpleNamespace
16	from typing import Any, cast
17	from unittest.mock import MagicMock, patch
18
19	import msgspec
20	import numpy as np
21	import pytest
22	import ray
23	import torch
24
25	from tests.v1.attention.utils import dense_kv_cache_views
26	from vllm import LLM
27	from vllm.config import KVTransferConfig, set_current_vllm_config
28	from vllm.distributed.kv_transfer.kv_connector.utils import (
29	EngineTransferInfo,
30	KVOutputAggregator,
31	TransferTopology,
32	get_current_attn_backend,
33	)
34	from vllm.distributed.kv_transfer.kv_connector.v1 import nixl
35	from vllm.distributed.kv_transfer.kv_connector.v1.base import (
36	KVConnectorRole,
37	KVConnectorTransferResults,
38	)
39	from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
40	from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
41	MultiKVConnectorStats,
42	)
43	from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
44	NixlAgentMetadata,
45	NixlConnector,
46	NixlConnectorMetadata,
47	NixlConnectorScheduler,
48	NixlConnectorWorker,
49	NixlHandshakePayload,
50	NixlKVConnectorStats,
51	)
52	from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
53	RemoteMeta,
54	ReqMeta,
55	compute_nixl_compatibility_hash,
56	)
57	from vllm.distributed.kv_transfer.kv_transfer_state import (
58	ensure_kv_transfer_shutdown,
59	has_kv_transfer_group,
60	)
61	from vllm.forward_context import ForwardContext
62	from vllm.outputs import RequestOutput
63	from vllm.platforms import current_platform
64	from vllm.platforms.interface import Platform
65	from vllm.sampling_params import SamplingParams
66	from vllm.v1.engine import EngineCoreRequest
67	from vllm.v1.engine.output_processor import OutputProcessor
68	from vllm.v1.kv_cache_interface import (
69	AttentionSpec,
70	FullAttentionSpec,
71	KVCacheConfig,
72	KVCacheGroupSpec,
73	KVCacheLayout,
74	MLAAttentionSpec,
75	compute_layer_kv_cache_shape_bytes,
76	)
77	from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
78	from vllm.v1.request import RequestStatus
79
80	from .utils import (
81	create_model_runner_output,
82	create_request,
83	create_scheduler,
84	create_vllm_config,
85	make_kv_cache_config,
86	)
87
88
89	@pytest.fixture(scope="module", autouse=True)
90	def clear_kv_transfer():
91	"""The test cases in this file use `VLLM_ENABLE_V1_MULTIPROCESSING=0`,
92	causing the global variable `_KV_CONNECTOR_AGENT`
93	to be assigned but never deleted.
94
95	Since the current pytest process does not terminate and instead
96	continues running tests from other files,
97	this global variable remains in memory and interferes
98	with test cases in other modules.
99
100	So we use this fixture to ensure that the global variable
101	`_KV_CONNECTOR_AGENT` is properly cleaned up after each test.
102	"""
103	yield
104	if has_kv_transfer_group():
105	ensure_kv_transfer_shutdown()
106
107
108	def get_default_xfer_telemetry(
109	xferDurationS: float = 1,
110	postDurationS: float = 1,
111	totalBytes: int = 1,
112	descCount: int = 1,
113	) -> dict:
114	class AttributeDict(dict):
115	__slots__ = ()
116	__getattr__ = dict.__getitem__
117	__setattr__ = dict.__setitem__
118
119	# We can't instantiate nixlXferTelemetry because it's read only and
120	# ray env does not have NIXL, so we must fake it
121	return AttributeDict(
122	xferDuration=xferDurationS * 1e6, # in us
123	postDuration=postDurationS * 1e6, # in us
124	totalBytes=totalBytes,
125	descCount=descCount,
126	)
127
128
129	class FakeNixlWrapper:
130	"""Mock implementation of NixlWrapper for testing.
131
132	We don't inherit from nixl._api.nixl_agent because nixl may not be
133	installed.
134
135	Note: The complete source of this class is also used in the
136	`_make_fake_nixl_pkg` function to create a fake nixl package
137	for Ray workers.
138	"""
139
140	AGENT_METADATA = b"fake_agent_metadata"
141	REMOTE_AGENT_NAME = "remote_agent"
142
143	def __init__(self, agent_name: str, args, *kwargs):
144	self._cycles_before_xfer_done = 0
145	self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(lambda: 0)
146
147	def get_reg_descs(self, caches_data, memory_type: str) -> list:
148	return [str(uuid.uuid4()) for _ in caches_data]
149
150	def register_memory(self, descs, backends) -> None:
151	pass
152
153	def deregister_memory(self, descs) -> None:
154	pass
155
156	def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
157	return [str(uuid.uuid4()) for _ in blocks_data]
158
159	def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
160	return uuid.uuid4().int
161
162	def get_agent_metadata(self) -> bytes:
163	return self.AGENT_METADATA
164
165	def add_remote_agent(self, agent_metadata: bytes) -> str:
166	return self.REMOTE_AGENT_NAME
167
168	def get_new_notifs(self) -> dict[str, list[bytes]]:
169	# Used to collect done_sending, which we don't test yet.
170	return {}
171
172	def check_xfer_state(self, handle: int) -> str:
173	if self._check_xfer_state_cycles[handle] >= self._cycles_before_xfer_done:
174	return "DONE"
175	self._check_xfer_state_cycles[handle] += 1
176	return "PROC"
177
178	def release_xfer_handle(self, handle: int) -> None:
179	pass
180
181	def release_dlist_handle(self, handle: int) -> None:
182	pass
183
184	def remove_remote_agent(self, agent: str) -> None:
185	pass
186
187	def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
188	pass
189
190	def make_prepped_xfer(
191	self,
192	xfer_type: str,
193	local_xfer_side_handle: int,
194	local_block_descs_ids: list[int],
195	remote_xfer_side_handle: int,
196	remote_block_descs_ids: list[int],
197	notif_msg: bytes \| None = None,
198	) -> int:
199	return uuid.uuid4().int
200
201	def transfer(self, handle: int) -> str:
202	return "PROC"
203
204	def get_xfer_telemetry(self, handle: int) -> dict:
205	return get_default_xfer_telemetry()
206
207
208	@contextlib.contextmanager
209	def _make_fake_nixl_pkg():
210	"""Context manager that creates a temporary package making
211	`from nixl._api import nixl_agent` resolve to our FakeNixlWrapper.
212	Also creates the ROCm NIXL packages.
213
214	Automatically cleans up the temporary directory when done.
215	"""
216	with tempfile.TemporaryDirectory() as td:
217	for pkg_name in ["nixl", "nixl_rocm"]:
218	pkg_root = os.path.join(td, pkg_name, "_api")
219	os.makedirs(pkg_root, exist_ok=True)
220
221	# Get the source code of FakeNixlWrapper class and dedent it
222	fake_nixl_source = inspect.getsource(FakeNixlWrapper)
223	fake_nixl_source = textwrap.dedent(fake_nixl_source)
224
225	stub = f"""\
226	# Copy of FakeNixlWrapper implementation for Ray workers
227	import uuid
228	from collections import defaultdict
229
230	{fake_nixl_source}
231
232	# Export as nixl_agent
233	nixl_agent = FakeNixlWrapper
234	"""
235	with open(os.path.join(pkg_root, "__init__.py"), "w") as f:
236	f.write(stub)
237
238	# Mock nixlXferTelemetry class
239	pkg_root2 = os.path.join(td, pkg_name, "_bindings")
240	os.makedirs(pkg_root2, exist_ok=True)
241	with open(os.path.join(pkg_root2, "__init__.py"), "w") as f:
242	f.write("class nixlXferTelemetry: pass")
243	# touch parent package
244	open(os.path.join(td, pkg_name, "__init__.py"), "w").close()
245
246	yield td
247
248
249	def test_basic_interface():
250	"""Unit test for basic NixlConnector interface functionality."""
251	vllm_config = create_vllm_config()
252	scheduler = create_scheduler(vllm_config)
253
254	# 2 Full Blocks and 1 Half Block.
255	BLOCK_SIZE = vllm_config.cache_config.block_size
256	NUM_EXTERNAL_FULL_BLOCKS = 2
257	NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
258
259	request = create_request(
260	request_id=1,
261	block_size=BLOCK_SIZE,
262	num_tokens=NUM_TOKENS,
263	do_remote_prefill=True,
264	)
265	request_id = request.request_id
266
267	scheduler.add_request(request)
268
269	# Remote Prefill, triggers NixlConnectorMetadata.
270	scheduler_output = scheduler.schedule()
271	kv_connector_metadata = scheduler_output.kv_connector_metadata
272	assert kv_connector_metadata is not None
273	assert isinstance(kv_connector_metadata, NixlConnectorMetadata)
274
275	assert len(kv_connector_metadata.reqs_to_recv) == 1
276	assert request_id in kv_connector_metadata.reqs_to_recv
277	req_meta = kv_connector_metadata.reqs_to_recv[request_id]
278
279	for block_id, block in zip(
280	req_meta.local_block_ids[0],
281	scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
282	request_id
283	],
284	):
285	assert block_id == block.block_id
286
287
288	def test_prompt_less_than_block_size():
289	"""Test that we can handle case where prompt is < block.
290
291	In this case, the P worker will still send remote_block_ids of the
292	partial block. The D worker should schedule an async read
293	in this case.
294	"""
295	vllm_config = create_vllm_config()
296	scheduler = create_scheduler(vllm_config)
297
298	# Half of a block.
299	BLOCK_SIZE = vllm_config.cache_config.block_size
300	NUM_TOKENS = int(BLOCK_SIZE * 0.5)
301
302	# Request will have 1 partial remote block.
303	request = create_request(
304	request_id=1,
305	block_size=BLOCK_SIZE,
306	num_tokens=NUM_TOKENS,
307	do_remote_prefill=True,
308	num_remote_blocks=1,
309	)
310	scheduler.add_request(request)
311	scheduler_output = scheduler.schedule()
312
313	# This request will read async.
314	kv_connector_metadata = scheduler_output.kv_connector_metadata
315	assert kv_connector_metadata is not None
316	assert isinstance(kv_connector_metadata, NixlConnectorMetadata)
317	assert len(kv_connector_metadata.reqs_to_recv) == 1
318	assert len(scheduler_output.scheduled_new_reqs) == 0
319
320
321	def test_abort_immediately_remote_prefill_enqueues_empty_recv():
322	"""A remote-prefill request added with abort_immediately=True should
323	be added to the scheduler's waiting queue then immediately aborted, so the
324	NIXL connector's request_finished hook enqueues an empty recv to notify
325	the prefill instance to free its blocks."""
326	from vllm.v1.request import RequestStatus
327
328	scheduler = create_scheduler(create_vllm_config())
329
330	request = create_request(request_id=42, num_tokens=10, do_remote_prefill=True)
331	assert request.kv_transfer_params is not None
332	assert request.kv_transfer_params["do_remote_prefill"] is True
333
334	# Mimic the EngineCore.add_request path for an abort-immediately req.
335	scheduler.add_request(request)
336	scheduler.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED)
337
338	scheduler_output = scheduler.schedule()
339	meta = scheduler_output.kv_connector_metadata
340	assert isinstance(meta, NixlConnectorMetadata)
341	assert set(meta.reqs_to_recv) == {request.request_id}
342	req_meta = meta.reqs_to_recv[request.request_id]
343	assert req_meta.local_block_ids == []
344	assert req_meta.remote.request_id == f"prefill-{42}"
345	# do_remote_prefill is consumed by request_finished to prevent re-issuing.
346	assert request.kv_transfer_params["do_remote_prefill"] is False
347	# The scheduler is not waiting on this recv -- the request is already gone
348	# from self.requests, so reporting it would trip `assert req_id in
349	# self.requests` in _update_from_kv_xfer_finished.
350	assert req_meta.awaiting_kvs is False
351
352
353	def test_prefill_exports_cached_tokens_in_kv_transfer_params():
354	"""The P worker reports its own prefix-cache hits in the returned
355	kv_transfer_params so the D worker can surface them in
356	prompt_tokens_details instead of the ~100% local hit it measures
357	when pulling the KVs from the remote.
358	"""
359	vllm_config = create_vllm_config()
360	scheduler = create_scheduler(vllm_config)
361
362	BLOCK_SIZE = vllm_config.cache_config.block_size
363	NUM_TOKENS = BLOCK_SIZE * 3
364
365	# Warm the prefix cache with a request sharing the full prompt.
366	warmup = create_request(
367	request_id=1,
368	num_tokens=NUM_TOKENS,
369	common_prefix_len=NUM_TOKENS,
370	block_size=BLOCK_SIZE,
371	)
372	scheduler.add_request(warmup)
373	scheduler_output = scheduler.schedule()
374	model_runner_output = create_model_runner_output([warmup], use_eos=True)
375	scheduler.update_from_output(scheduler_output, model_runner_output)
376
377	# P-side request with the same prompt hits the local cache on all but
378	# the last block (recomputed to obtain logits).
379	request = create_request(
380	request_id=2,
381	num_tokens=NUM_TOKENS,
382	common_prefix_len=NUM_TOKENS,
383	block_size=BLOCK_SIZE,
384	do_remote_decode=True,
385	)
386	scheduler.add_request(request)
387	scheduler_output = scheduler.schedule()
388	assert request.prefill_stats is not None
389	assert request.prefill_stats.num_cached_tokens == NUM_TOKENS - BLOCK_SIZE
390
391	# max_tokens=1, so the request finishes and returns kv_transfer_params.
392	model_runner_output = create_model_runner_output([request])
393	engine_core_outputs = scheduler.update_from_output(
394	scheduler_output, model_runner_output
395	)
396
397	output = engine_core_outputs[0].outputs[0]
398	assert output.finish_reason is not None
399	assert output.kv_transfer_params is not None
400	assert (
401	output.kv_transfer_params["remote_prefill_cached_tokens"]
402	== NUM_TOKENS - BLOCK_SIZE
403	)
404
405
406	@patch(
407	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
408	FakeNixlWrapper,
409	)
410	def test_kv_transfer_handshake(dist_init):
411	"""Unit test for basic NixlConnector interface functionality."""
412	# Test setup, we creates a scheduler that contains a NixlConnector
413	# of role SCHEDULER, and expect it to be serving NixlAgentMetadata from
414	# all workers of the instance.
415	vllm_config = create_vllm_config()
416	vllm_config.cache_config.kv_cache_layout = "BLHNC"
417	# in case the test runs on non-GPU machine
418	vllm_config.kv_transfer_config.kv_buffer_device = "cpu"
419	scheduler = create_scheduler(vllm_config)
420
421	with set_current_vllm_config(vllm_config):
422	# Create two NixlConnector of role WORKER, one is the worker of
423	# the scheduler (prefill), the other is a worker of decode instance.
424
425	# Prefill connector will register KV cache to populate proper handshake
426	# metadata.
427	kv_cache_groups = [
428	KVCacheGroupSpec(
429	["layer0", "layer1", "layer2"],
430	FullAttentionSpec(
431	block_size=16,
432	num_kv_heads=4,
433	head_size=16,
434	dtype=torch.float16,
435	),
436	)
437	]
438	kv_cache_config = KVCacheConfig(
439	num_blocks=2, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
440	)
441	prefill_connector = NixlConnector(
442	vllm_config, KVConnectorRole.WORKER, kv_cache_config
443	)
444	kv_cache_spec = cast(
445	AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
446	)
447	raw = torch.zeros(
448	kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 3,
449	dtype=torch.int8,
450	)
451	caches = dense_kv_cache_views(
452	raw,
453	kv_cache_spec,
454	kv_cache_config.num_blocks,
455	num_layers=3,
456	layout=KVCacheLayout.BLHNC,
457	)
458	kv_caches = {
459	f"layer{layer_idx}": cache for layer_idx, cache in enumerate(caches)
460	}
461	prefill_connector.register_kv_caches(kv_caches)
462
463	# Simulate EngineCore initialization that would gather connector
464	# metadata from all workers
465	metadata = prefill_connector.get_handshake_metadata()
466
467	# metadata is a NixlHandshakePayload, decode it to get NixlAgentMetadata
468	decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
469	expected_agent_metadata = decoder.decode(metadata.agent_metadata_bytes)
470
471	# The scheduler connector expects metadata keyed by
472	# (pp_rank, tp_rank).
473	scheduler_connector = scheduler.get_kv_connector()
474	scheduler_connector.set_xfer_handshake_metadata_pp_aware({(0, 0): metadata})
475
476	# Simulate a request that finishes prefill, which returns
477	# corresponding NixlConnectorMetadata for decode instance.
478	BLOCK_SIZE = vllm_config.cache_config.block_size
479	NUM_EXTERNAL_FULL_BLOCKS = 2
480	NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
481
482	request = create_request(
483	request_id=1,
484	block_size=BLOCK_SIZE,
485	num_tokens=NUM_TOKENS,
486	do_remote_decode=True,
487	)
488	request.status = RequestStatus.FINISHED_LENGTH_CAPPED
489	delay, kv_connector_metadata = (
490	scheduler.get_kv_connector().request_finished_all_groups(
491	request, ([0, 1, 2],)
492	)
493	)
494	assert delay
495	# Pull connector advertises its transfer mode in kv_transfer_params so
496	# an external router can distinguish it from a push producer.
497	assert kv_connector_metadata["transfer_mode"] == "pull"
498
499	# Decode connector will be able to create handshake with the prefill connector.
500	decode_connector = NixlConnector(
501	vllm_config, KVConnectorRole.WORKER, kv_cache_config
502	)
503	decode_connector.register_kv_caches(kv_caches)
504
505	# Here we are testing the retrieval of NIXLAgentMetadata.
506	# Knowing the implementation detail, we override the add_remote_agent
507	# to validate the metadata received is the same as the one in prefill_connector.
508	with patch.object(
509	decode_connector.connector_worker, "add_remote_agent"
510	) as mock_add_remote_agent:
511	mock_add_remote_agent.return_type = "remote_agent"
512
513	decode_connector.connector_worker._nixl_handshake(
514	kv_connector_metadata["remote_host"],
515	kv_connector_metadata["remote_port"],
516	kv_connector_metadata["tp_size"],
517	kv_connector_metadata["remote_engine_id"],
518	)
519
520	received_metadata = mock_add_remote_agent.call_args.args
521	assert received_metadata[0] == expected_agent_metadata
522	assert received_metadata[1] == 0 # remote_tp_rank
523	assert received_metadata[2] == 1 # remote_tp_size
524
525	# Need to shutdown the background thread to release NIXL side channel port
526	scheduler_connector.shutdown()
527
528
529	class FakeNixlConnectorWorker(NixlConnectorWorker):
530	REMOTE_ENGINE_ID = "remote_engine"
531
532	def __init__(
533	self,
534	*args,
535	hand_shake_latency: float = 1.8,
536	kv_cache_layout="LBHNC",
537	kv_cache_config=None,
538	**kwargs,
539	):
540	if kv_cache_config is None:
541	kv_cache_config = make_kv_cache_config(block_size=16)
542	super().__init__(args, kv_cache_config=kv_cache_config, *kwargs)
543	self._hand_shake_latency = hand_shake_latency
544	self.kv_cache_layout = kv_cache_layout
545	# Mock register_kv_caches attributes needed for tests that do not call it.
546	self.src_xfer_handles_by_block_size = {self.block_size: 1}
547	self.src_blocks_data = np.empty((0, 3), dtype=np.uint64)
548	rep_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
549	test_shape = compute_layer_kv_cache_shape_bytes(rep_spec, 1)
550	self.transfer_topo = TransferTopology(
551	tp_rank=self.tp_rank,
552	tp_size=self.world_size,
553	block_size=self.block_size,
554	engine_id=self.engine_id,
555	is_mla=self.use_mla,
556	is_mamba=False,
557	total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
558	attn_backends=self.attn_backends,
559	tensor_shape=test_shape,
560	)
561
562	self.compat_hash = compute_nixl_compatibility_hash(
563	self.vllm_config, self.backend_name
564	)
565
566	def _nixl_handshake(
567	self,
568	host: str,
569	port: int,
570	remote_tp_size: int,
571	expected_engine_id: str,
572	remote_dcp_size: int = 1,
573	remote_pp_size: int = 1,
574	notif_agents_only: bool = False,
575	) -> tuple[dict[tuple[int, int], str], float]:
576	# Mimic slow _nixl_handshake, as well as bypass zmq communication.
577	time.sleep(self._hand_shake_latency)
578	# These should've been done in register_kv_caches(), called by
579	# gpu_model_runner. Here we just hardcode some dummy values.
580	slot_size_bytes = 4096
581	self.slot_size_per_layer = [slot_size_bytes]
582	self.block_len_per_layer = [slot_size_bytes * self.block_size]
583	self.num_blocks = self.kv_cache_config.num_blocks
584	self.num_regions = 1
585	self.block_stride_per_layer = list(self.block_len_per_layer)
586	self.region_num_blocks = [self.num_blocks]
587	self.region_group_ids = [0]
588	self.region_mem_types = [self.nixl_memory_type]
589	self.dst_num_blocks[self.engine_id] = self.num_blocks
590	self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks
591	self.dst_region_group_ids[self.engine_id] = self.region_group_ids
592	self.dst_region_mem_types[self.engine_id] = self.region_mem_types
593
594	assert expected_engine_id == self.REMOTE_ENGINE_ID
595
596	# Adjust remote block length metadata to satisfy heterogeneous TP
597	# invariants enforced during handshake validation. Use per-rank
598	# head ratio (not tp_ratio) to account for GQA replication capping.
599	remote_block_lens = list(self.block_len_per_layer)
600	tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
601	total_kv = self.transfer_topo.total_num_kv_heads
602	local_heads = self.transfer_topo.local_physical_heads
603	remote_heads = max(1, total_kv // remote_tp_size)
604	if remote_tp_size != self.world_size:
605	remote_block_lens = [
606	block_len * remote_heads // local_heads
607	for block_len in remote_block_lens
608	]
609
610	# When remote tp_size > local tp_size, handshake with multiple
611	# remote ranks.
612	num_handshakes = 1 if tp_ratio > 0 else -tp_ratio
613	remote_agents: dict[tuple[int, int], str] = {}
614	for remote_tp_rank in range(num_handshakes):
615	remote_agent_name = self.add_remote_agent(
616	NixlAgentMetadata(
617	engine_id=self.REMOTE_ENGINE_ID,
618	agent_metadata=FakeNixlWrapper.AGENT_METADATA,
619	kv_caches_base_addr=[0],
620	device_id=remote_tp_rank,
621	num_blocks=self.num_blocks,
622	block_lens=remote_block_lens,
623	block_strides=remote_block_lens,
624	kv_cache_layout="LBHNC",
625	block_size=self.block_size,
626	ssm_sizes=(0, 0),
627	attn_backend_name=self.backend_name,
628	physical_blocks_per_logical_kv_block=1,
629	region_num_blocks=self.region_num_blocks,
630	region_group_ids=self.region_group_ids,
631	region_mem_types=self.region_mem_types,
632	),
633	remote_tp_rank=remote_tp_rank,
634	remote_tp_size=remote_tp_size,
635	)
636	remote_agents[(0, remote_tp_rank)] = remote_agent_name
637	# Handshake bypasses zmq, so report a zero clock offset to the peer.
638	return remote_agents, 0.0
639
640
641	class TestNixlHandshake:
642	@pytest.mark.parametrize(
643	("pcp_rank", "pcp_size", "dcp_size", "expected_tracked"),
644	[
645	(0, 2, 1, True),
646	(1, 2, 1, False),
647	(0, 2, 2, True),
648	(1, 2, 2, True),
649	(0, 4, 4, True),
650	(1, 4, 4, True),
651	(2, 4, 4, True),
652	(3, 4, 4, True),
653	],
654	)
655	@patch(
656	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
657	FakeNixlWrapper,
658	)
659	def test_pcp_producer_exposes_dcp_shards_or_canonical_replica(
660	self,
661	default_vllm_config,
662	dist_init,
663	pcp_rank,
664	pcp_size,
665	dcp_size,
666	expected_tracked,
667	):
668	"""Replicated PCP is canonicalized; PCP-DCP publishes every shard."""
669	from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
670
671	vllm_config = create_vllm_config(kv_role="kv_producer")
672	vllm_config.parallel_config.prefill_context_parallel_size = pcp_size
673	vllm_config.parallel_config.decode_context_parallel_size = dcp_size
674	with (
675	patch(
676	"vllm.distributed.kv_transfer.kv_connector.v1.nixl."
677	"base_worker.get_current_attn_backends",
678	return_value=[FlashAttentionBackend],
679	),
680	patch(
681	"vllm.distributed.kv_transfer.kv_connector.v1.nixl."
682	"base_worker.get_pcp_group"
683	) as mock_get_pcp_group,
684	):
685	mock_get_pcp_group.return_value.rank_in_group = pcp_rank
686	connector = NixlConnector(
687	vllm_config,
688	KVConnectorRole.WORKER,
689	make_kv_cache_config(block_size=16),
690	)
691
692	worker = connector.connector_worker
693	assert worker is not None
694	assert worker.pcp_rank == pcp_rank
695	assert worker.pcp_dcp_sharded is (dcp_size > 1)
696	assert worker.transfer_tp_rank == (pcp_rank if dcp_size > 1 else 0)
697	assert worker.transfer_tp_size == (pcp_size if dcp_size > 1 else 1)
698
699	req_id = "req"
700	metadata = NixlConnectorMetadata()
701	metadata.reqs_in_batch.add(req_id)
702	metadata.reqs_to_send[req_id] = time.perf_counter() + 10
703	worker.start_load_kv(metadata)
704	assert (req_id in worker._reqs_to_process) == expected_tracked
705	assert (req_id in worker._reqs_to_send) == expected_tracked
706
707	payload = MagicMock(spec=NixlHandshakePayload)
708	worker.xfer_handshake_metadata = payload
709	worker.transfer_topo = MagicMock()
710	worker._get_new_notifs = MagicMock(
711	side_effect=lambda: {"sent"} if expected_tracked else set()
712	)
713
714	expected_payload = payload if expected_tracked else None
715	assert connector.get_handshake_metadata() is expected_payload
716	done_sending, done_recving = connector.get_finished(set())
717	assert done_sending == ({"sent"} if expected_tracked else {req_id})
718	assert done_recving == set()
719	if not expected_tracked:
720	assert connector.get_finished(set()) == (set(), set())
721
722	worker.get_transfer_results = MagicMock(
723	return_value=KVConnectorTransferResults(finished_sending={"sent"})
724	)
725	results = connector.get_transfer_results(set())
726	assert results.finished_sending == ({"sent"} if expected_tracked else set())
727
728	@patch(
729	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
730	FakeNixlWrapper,
731	)
732	def test_multi_xfer_one_engine(
733	self,
734	default_vllm_config,
735	# dist_init is a fixture that initializes the distributed environment.
736	dist_init,
737	):
738	"""Test case where multiple xfers are initiated to the same engine.
739
740	This test triggers the connector to load remote KV for the same
741	`request_id`.
742	"""
743	vllm_config = create_vllm_config()
744
745	request_id = "req_id"
746
747	# Test worker role in decode server.
748	kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10)
749	connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
750	connector.connector_worker = FakeNixlConnectorWorker(
751	vllm_config,
752	connector.engine_id,
753	hand_shake_latency=0,
754	kv_cache_config=kv_cache_config,
755	)
756	assert isinstance(connector.connector_worker.nixl_wrapper, FakeNixlWrapper)
757	worker = connector.connector_worker
758	# simulate handshake
759	worker.dst_xfer_side_handles = {
760	FakeNixlConnectorWorker.REMOTE_ENGINE_ID: {0: 1}
761	}
762	worker.kv_cache_layout = "LBHNC"
763	num_xfers = 4
764	while True:
765	# For the same request_id, initiate multiple xfers across different
766	# round of `execute_model` calls.
767	metadata = NixlConnectorMetadata()
768	if num_xfers > 0:
769	num_xfers -= 1
770	metadata.add_new_req_to_recv(
771	request_id=request_id,
772	local_block_ids=([num_xfers + 1, num_xfers + 2, num_xfers + 3],),
773	kv_transfer_params={
774	"remote_block_ids": (
775	[
776	num_xfers + 4,
777	num_xfers + 5,
778	num_xfers + 6,
779	],
780	),
781	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
782	"remote_request_id": f"prefill-{request_id}",
783	"remote_host": "localhost",
784	"remote_port": 1234,
785	"remote_tp_size": 1,
786	},
787	)
788	connector.bind_connector_metadata(metadata)
789
790	# Mimic logic in KVConnectorModelRunnerMixin._get_kv_connector_output.
791	dummy_ctx = ForwardContext(
792	no_compile_layers={},
793	attn_metadata={},
794	slot_mapping={},
795	)
796	_before_load = time.perf_counter()
797	connector.start_load_kv(dummy_ctx)
798	_after_load = time.perf_counter()
799	assert _after_load - _before_load < 0.1, (
800	f"start_load_kv took {_after_load - _before_load} seconds"
801	)
802
803	# Mimic logic in KVConnectorModelRunnerMixin._get_kv_connector_output.
804	_, done_recving = connector.get_finished(finished_req_ids=set())
805	if len(done_recving) > 0:
806	assert request_id in done_recving
807	break
808
809	connector.clear_connector_metadata()
810
811	@patch(
812	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
813	FakeNixlWrapper,
814	)
815	@pytest.mark.parametrize(
816	"decode_tp_size, prefill_tp_size",
817	[
818	(1, 1),
819	(2, 1),
820	(4, 2),
821	(4, 4),
822	],
823	)
824	def test_async_load_kv(
825	self,
826	default_vllm_config,
827	# Fixture that initializes the distributed environment.
828	dist_init,
829	# Simulate consumer-producer TP sizes.
830	decode_tp_size,
831	prefill_tp_size,
832	):
833	"""Test that NixlConnector's start_load_kv should be non-blocking."""
834	vllm_config = create_vllm_config()
835	vllm_config.parallel_config.tensor_parallel_size = decode_tp_size
836
837	# Test worker role in decode server.
838	connector = NixlConnector(
839	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
840	)
841	connector.connector_worker = FakeNixlConnectorWorker(
842	vllm_config, connector.engine_id
843	)
844	metadata = NixlConnectorMetadata()
845	metadata.add_new_req_to_recv(
846	request_id="id",
847	local_block_ids=([1, 2, 3],),
848	kv_transfer_params={
849	"remote_block_ids": ([4, 5, 6],),
850	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
851	"remote_request_id": "prefill-id",
852	"remote_host": "localhost",
853	"remote_port": 1234,
854	"remote_tp_size": prefill_tp_size,
855	},
856	)
857	connector.bind_connector_metadata(metadata)
858
859	timeout = 2.5
860	start = time.perf_counter()
861	while time.perf_counter() - start < timeout:
862	dummy_ctx = ForwardContext(
863	no_compile_layers={},
864	attn_metadata={},
865	slot_mapping={},
866	)
867	_before_load = time.perf_counter()
868	connector.start_load_kv(dummy_ctx)
869	_after_load = time.perf_counter()
870	assert _after_load - _before_load < 0.1, (
871	f"start_load_kv took {_after_load - _before_load} seconds"
872	)
873	time.sleep(0.5) # backoff for the async handshake to complete.
874	connector.bind_connector_metadata(NixlConnectorMetadata())
875	_, done_recving = connector.get_finished(finished_req_ids=set())
876	if len(done_recving) > 0:
877	return
878	raise TimeoutError("Took too long to complete async handshake.")
879
880	@patch(
881	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
882	FakeNixlWrapper,
883	)
884	@pytest.mark.parametrize("local_tp_size", [1, 2])
885	def test_prefill_tp_size_greater_than_decode_tp_size(
886	self, local_tp_size: int, default_vllm_config, dist_init, monkeypatch
887	):
888	"""Verify remote TP > local TP handshake succeeds with different
889	remote configurations.
890	"""
891	monkeypatch.setattr(
892	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size",
893	lambda: local_tp_size,
894	)
895
896	vllm_config = create_vllm_config()
897
898	connector = NixlConnector(
899	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
900	)
901	connector.connector_worker = FakeNixlConnectorWorker(
902	vllm_config, connector.engine_id, hand_shake_latency=0
903	)
904	worker = connector.connector_worker
905
906	# Minimal local registration params used by add_remote_agent
907	worker.slot_size_per_layer = [4096]
908	worker.block_len_per_layer = [4096 * worker.block_size]
909	worker.num_blocks = 1
910	worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
911	worker.src_blocks_data = np.array(
912	[(0, worker.block_len_per_layer[0], worker.tp_rank)],
913	dtype=np.uint64,
914	)
915	worker.num_descs = len(worker.src_blocks_data)
916
917	def check_handshake(remote_tp_size: int):
918	tp_ratio = remote_tp_size // local_tp_size
919	assert set(remote_agents.keys()) == {(0, r) for r in range(tp_ratio)}
920
921	remote_engine_id = worker.REMOTE_ENGINE_ID
922	remote_info = worker.transfer_topo.get_engine_info(remote_engine_id)
923	assert remote_info.remote_tp_size == remote_tp_size
924	assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size)
925	# ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks
926	split_key = (-tp_ratio, worker.block_size)
927	assert split_key in worker.src_xfer_handles_by_tp_ratio
928	assert len(worker.src_xfer_handles_by_tp_ratio[split_key]) == tp_ratio
929	assert remote_engine_id in worker.dst_xfer_side_handles
930	assert set(worker.dst_xfer_side_handles[remote_engine_id].keys()) == set(
931	range(tp_ratio)
932	)
933
934	remote_agents, _ = worker._nixl_handshake(
935	host="localhost",
936	port=1234,
937	remote_tp_size=4,
938	expected_engine_id=worker.REMOTE_ENGINE_ID,
939	)
940	check_handshake(4)
941
942	# NOTE flexibility: a second remote with higher number of ranks is
943	# discovered. This is not a scenario we actively support right now, but
944	# the connector allows it.
945	worker.REMOTE_ENGINE_ID = "remote_engine_2"
946	remote_agents, _ = worker._nixl_handshake(
947	host="localhost",
948	port=1234,
949	remote_tp_size=6,
950	expected_engine_id=worker.REMOTE_ENGINE_ID,
951	)
952	check_handshake(6)
953
954	@patch(
955	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
956	FakeNixlWrapper,
957	)
958	def test_prefill_tp_size_greater_than_decode_tp_size_mla(
959	self, default_vllm_config, dist_init
960	):
961	"""Verify remote TP > local TP handshake succeeds with different
962	remote configurations for an MLA model.
963	"""
964	vllm_config = create_vllm_config()
965	d_tp_size = 1
966	p_tp_size = 2
967
968	# Build two separate connectors/workers to emulate P TP=2 ranks.
969	conn_p0 = NixlConnector(
970	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
971	)
972	conn_p1 = NixlConnector(
973	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
974	)
975	conn_p0.connector_worker = FakeNixlConnectorWorker(
976	vllm_config, conn_p0.engine_id, hand_shake_latency=0
977	)
978	conn_p1.connector_worker = FakeNixlConnectorWorker(
979	vllm_config, conn_p1.engine_id, hand_shake_latency=0
980	)
981
982	# Force P world size to 2 for both workers and emulate distinct tp_ranks.
983	# Also enable MLA path so that expected_finished_count is updated.
984	for rank, worker in enumerate(
985	(conn_p0.connector_worker, conn_p1.connector_worker)
986	):
987	worker.world_size = p_tp_size
988	worker.transfer_topo.tp_size = p_tp_size
989	worker.tp_rank = rank
990	worker.use_mla = True
991
992	req_id = "req-ep-dp2-p0"
993	now = time.perf_counter()
994	# Register a request on P that is waiting for consumers to read
995	# (both workers track it).
996	conn_p0.connector_worker._reqs_to_send[req_id] = now + 10.0
997	conn_p0.connector_worker._reqs_to_process.add(req_id)
998	conn_p1.connector_worker._reqs_to_send[req_id] = now + 10.0
999	conn_p1.connector_worker._reqs_to_process.add(req_id)
1000
1001	# Simulate a read notification coming from D with (tp=1, dp=2).
1002	notif = f"{req_id}:{d_tp_size}".encode()
1003	# D0-0->P0 notif
1004	conn_p0.connector_worker.nixl_wrapper.get_new_notifs = lambda: {
1005	"agent": [notif]
1006	} # type: ignore[method-assign]
1007	conn_p1.connector_worker.nixl_wrapper.get_new_notifs = lambda: {
1008	"agent": [notif]
1009	} # type: ignore[method-assign]
1010
1011	# Trigger notification processing via get_finished().
1012	done_sending0, _ = conn_p0.get_finished(finished_req_ids=set())
1013	done_sending1, _ = conn_p1.get_finished(finished_req_ids=set())
1014	assert req_id in done_sending0 and req_id in done_sending1
1015
1016	# E2E aggregation: ensure the aggregated output marks the request
1017	# as finished using the connector's expected_finished_count.
1018	from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
1019
1020	aggregator = KVOutputAggregator.from_connector(conn_p0, world_size=2)
1021
1022	out0 = ModelRunnerOutput(
1023	req_ids=[req_id],
1024	req_id_to_index={req_id: 0},
1025	sampled_token_ids=[[0]],
1026	logprobs=None,
1027	prompt_logprobs_dict={},
1028	pooler_output=[None],
1029	kv_connector_output=KVConnectorOutput(
1030	finished_sending=done_sending0,
1031	finished_recving=None,
1032	),
1033	)
1034	out1 = ModelRunnerOutput(
1035	req_ids=[req_id],
1036	req_id_to_index={req_id: 0},
1037	sampled_token_ids=[[0]],
1038	logprobs=None,
1039	prompt_logprobs_dict={},
1040	pooler_output=[None],
1041	kv_connector_output=KVConnectorOutput(
1042	finished_sending=done_sending1,
1043	finished_recving=None,
1044	),
1045	)
1046	aggregated = aggregator.aggregate([out0, out1], output_rank=0)
1047	assert aggregated.kv_connector_output is not None
1048	assert aggregated.kv_connector_output.finished_sending == {req_id}
1049
1050	# Producers cleaned up state for the finished request.
1051	assert req_id not in conn_p0.connector_worker._reqs_to_send
1052	assert req_id not in conn_p0.connector_worker._reqs_to_process
1053	assert req_id not in conn_p1.connector_worker._reqs_to_send
1054	assert req_id not in conn_p1.connector_worker._reqs_to_process
1055
1056	@patch(
1057	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1058	FakeNixlWrapper,
1059	)
1060	def test_concurrent_load_kv(
1061	self,
1062	default_vllm_config,
1063	# dist_init is a fixture that initializes the distributed environment.
1064	dist_init,
1065	):
1066	"""Test that multiple start_load_kv calls should occur concurrently."""
1067	vllm_config = create_vllm_config()
1068
1069	# Test worker role in decode server.
1070	connector = NixlConnector(
1071	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1072	)
1073	connector.connector_worker = FakeNixlConnectorWorker(
1074	vllm_config, connector.engine_id
1075	)
1076	# Register (mocked) local xfer handler
1077	# worker = connector.connector_worker
1078	# worker.src_xfer_handles_by_block_size = {worker.block_size: 1}
1079	metadata = NixlConnectorMetadata()
1080	total_reqs = 5
1081	for i in range(total_reqs):
1082	metadata.add_new_req_to_recv(
1083	request_id=f"id_{i}",
1084	local_block_ids=([1, 2, 3],),
1085	kv_transfer_params={
1086	"remote_block_ids": ([4, 5, 6],),
1087	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1088	"remote_request_id": f"prefill-id-{i}",
1089	"remote_host": "localhost",
1090	"remote_port": 1234,
1091	"remote_tp_size": 1,
1092	},
1093	)
1094	connector.bind_connector_metadata(metadata)
1095
1096	timeout = 2.5 * total_reqs
1097	cnt_finished_reqs = 0
1098	start = time.perf_counter()
1099	while time.perf_counter() - start < timeout:
1100	dummy_ctx = ForwardContext(
1101	no_compile_layers={},
1102	attn_metadata={},
1103	slot_mapping={},
1104	)
1105	_before_load = time.perf_counter()
1106	connector.start_load_kv(dummy_ctx)
1107	_after_load = time.perf_counter()
1108	assert _after_load - _before_load < 0.1, (
1109	f"start_load_kv took {_after_load - _before_load} seconds"
1110	)
1111	time.sleep(0.5) # backoff for the async handshake to complete.
1112	connector.bind_connector_metadata(NixlConnectorMetadata())
1113	_, done_recving = connector.get_finished(finished_req_ids=set())
1114	if len(done_recving) > 0:
1115	cnt_finished_reqs += len(done_recving)
1116	if cnt_finished_reqs == total_reqs:
1117	return
1118	raise TimeoutError("Took too long to complete async handshake.")
1119
1120	@patch(
1121	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1122	FakeNixlWrapper,
1123	)
1124	def test_handshake_fails_on_kv_cache_layout_mismatch(
1125	self, default_vllm_config, dist_init
1126	):
1127	"""Verify that adding a remote agent fails if kv_cache_layout differs.
1128	This test is only relevant for heterogeneous TP.
1129	"""
1130	vllm_config = create_vllm_config()
1131
1132	# Mock TP world size to 2 to force heterogeneous TP when
1133	# remote_tp_size=1
1134	with patch(
1135	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
1136	return_value=2,
1137	):
1138	# Initialize connector and worker (with fake NIXL wrapper)
1139	connector = NixlConnector(
1140	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1141	)
1142	connector.connector_worker = FakeNixlConnectorWorker(
1143	vllm_config, connector.engine_id, hand_shake_latency=0
1144	)
1145	worker = connector.connector_worker
1146
1147	# Minimal local registration params used by add_remote_agent
1148	worker.slot_size_per_layer = [4096]
1149	worker.block_len_per_layer = [4096 * worker.block_size]
1150	worker.num_blocks = 1
1151	worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
1152
1153	# Metadata with different kv_cache_layout than local worker
1154	mismatched_layout = (
1155	"LBHNC" if worker.kv_cache_layout != "LBHNC" else "LBNHC"
1156	)
1157	meta = NixlAgentMetadata(
1158	engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1159	agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1160	kv_caches_base_addr=[0],
1161	device_id=0,
1162	num_blocks=1,
1163	block_lens=worker.block_len_per_layer,
1164	block_strides=worker.block_len_per_layer,
1165	kv_cache_layout=mismatched_layout,
1166	block_size=worker.block_size,
1167	ssm_sizes=(0, 0),
1168	attn_backend_name=worker.backend_name,
1169	physical_blocks_per_logical_kv_block=1,
1170	)
1171
1172	with pytest.raises(RuntimeError):
1173	# mismatched layout is expected to fail
1174	worker.add_remote_agent(meta, remote_tp_size=2)
1175	worker.add_remote_agent(meta, remote_tp_size=1)
1176
1177	@patch(
1178	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1179	FakeNixlWrapper,
1180	)
1181	def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental(
1182	self, default_vllm_config, dist_init
1183	):
1184	"""Verify that adding a remote agent fails if kv_cache_layout differs.
1185	This test is only relevant for heterogeneous TP.
1186	"""
1187	vllm_config = create_vllm_config(enable_permute_local_kv=True)
1188
1189	# Mock TP world size to 2 to force heterogeneous TP when
1190	# remote_tp_size=1
1191	with patch(
1192	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
1193	return_value=2,
1194	):
1195	# Initialize connector and worker (with fake NIXL wrapper)
1196	connector = NixlConnector(
1197	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1198	)
1199	connector.connector_worker = FakeNixlConnectorWorker(
1200	vllm_config,
1201	connector.engine_id,
1202	hand_shake_latency=0,
1203	kv_cache_layout="LBNHC",
1204	)
1205	worker = connector.connector_worker
1206
1207	# Minimal local registration params used by add_remote_agent
1208	worker.slot_size_per_layer = [2048]
1209	worker.block_len_per_layer = [2048 * worker.block_size]
1210	worker.num_blocks = 1
1211	worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
1212
1213	# Metadata with different kv_cache_layout than local worker
1214	# prefill TP=1, decode TP=2, remote block_lens is double to local
1215	remote_block_lens = [i * 2 for i in worker.block_len_per_layer]
1216	meta = NixlAgentMetadata(
1217	engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1218	agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1219	kv_caches_base_addr=[0],
1220	device_id=0,
1221	num_blocks=1,
1222	block_lens=remote_block_lens,
1223	block_strides=remote_block_lens,
1224	kv_cache_layout="LBHNC",
1225	block_size=worker.block_size,
1226	ssm_sizes=(0, 0),
1227	attn_backend_name=worker.backend_name,
1228	physical_blocks_per_logical_kv_block=1,
1229	)
1230
1231	# We don't check layout for homogeneous TP and MLA for now, as the
1232	# whole block is moved.
1233	worker.add_remote_agent(meta, remote_tp_size=1)
1234
1235	@patch(
1236	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1237	FakeNixlWrapper,
1238	)
1239	def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices(
1240	self, default_vllm_config, dist_init
1241	):
1242	worker = FakeNixlConnectorWorker(
1243	create_vllm_config(), "engine", hand_shake_latency=0
1244	)
1245
1246	remote_block_len = 2048
1247	local_block_len = remote_block_len // 2
1248	worker.block_len_per_layer = [local_block_len]
1249	worker._region_is_mla = [False]
1250	worker.num_blocks = 1
1251	worker.num_regions = 1
1252	worker._has_mamba = True
1253	worker._mamba_ssm_size = (128, 256)
1254	worker.transfer_topo = TransferTopology(
1255	tp_rank=1,
1256	tp_size=2,
1257	block_size=worker.block_size,
1258	engine_id=worker.engine_id,
1259	is_mla=False,
1260	is_mamba=True,
1261	total_num_kv_heads=2,
1262	attn_backends=worker.attn_backends,
1263	tensor_shape=None,
1264	)
1265	plan = MagicMock(
1266	source_ranks_per_group=((0,), (0,)),
1267	rank_offset_factor=1,
1268	)
1269	meta = NixlAgentMetadata(
1270	engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1271	agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1272	kv_caches_base_addr=[0x1000],
1273	device_id=0,
1274	num_blocks=1,
1275	block_lens=[remote_block_len],
1276	block_strides=[remote_block_len],
1277	kv_cache_layout="HND",
1278	block_size=worker.block_size,
1279	ssm_sizes=(0, 0),
1280	attn_backend_name=worker.backend_name,
1281	physical_blocks_per_logical_kv_block=1,
1282	region_num_blocks=None,
1283	)
1284
1285	assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
1286	[0x1000 + local_block_len, local_block_len, 0]
1287	]
1288
1289	@patch(
1290	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1291	FakeNixlWrapper,
1292	)
1293	def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init):
1294	"""Mixed full-attn (SPLIT) + MLA (REPLICATE) single KV group under
1295	heterogeneous TP must NOT raise (previously a NotImplementedError),
1296	and the per-region gate must still reject a wrong block_len.
1297	"""
1298	vllm_config = create_vllm_config()
1299	with patch(
1300	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
1301	return_value=2,
1302	):
1303	connector = NixlConnector(
1304	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1305	)
1306	connector.connector_worker = FakeNixlConnectorWorker(
1307	vllm_config, connector.engine_id, hand_shake_latency=0
1308	)
1309	worker = connector.connector_worker
1310
1311	# Region 0: full-attn (SPLIT). Region 1: MLA (REPLICATE).
1312	fa_len = 4096 * worker.block_size
1313	idx_len = 512 * worker.block_size
1314	worker.slot_size_per_layer = [4096, 512]
1315	worker.block_len_per_layer = [fa_len, idx_len]
1316	worker._region_is_mla = [False, True]
1317	worker.num_blocks = 1
1318	worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
1319	worker.src_blocks_data = np.array(
1320	[
1321	(0, fa_len, worker.tp_rank),
1322	(0, idx_len, worker.tp_rank),
1323	],
1324	dtype=np.uint64,
1325	)
1326	worker.num_descs = len(worker.src_blocks_data)
1327
1328	# D_TP=2, P_TP=1 -> tp_ratio=2. SPLIT region scales by tp_ratio;
1329	# REPLICATE region is unchanged.
1330	tp_ratio = 2
1331	meta = NixlAgentMetadata(
1332	engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1333	agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1334	kv_caches_base_addr=[0, 0],
1335	device_id=0,
1336	num_blocks=1,
1337	block_lens=[fa_len * tp_ratio, idx_len],
1338	block_strides=[fa_len * tp_ratio, idx_len],
1339	kv_cache_layout=worker.kv_cache_layout,
1340	block_size=worker.block_size,
1341	ssm_sizes=(0, 0),
1342	attn_backend_name=worker.backend_name,
1343	physical_blocks_per_logical_kv_block=1,
1344	)
1345	worker.add_remote_agent(meta, remote_tp_size=1)
1346	assert (
1347	FakeNixlConnectorWorker.REMOTE_ENGINE_ID in worker.dst_xfer_side_handles
1348	)
1349	# Gate rejects an MLA region wrongly scaled by tp_ratio.
1350	worker2 = FakeNixlConnectorWorker(
1351	vllm_config, connector.engine_id, hand_shake_latency=0
1352	)
1353	worker2.block_len_per_layer = [fa_len, idx_len]
1354	worker2._region_is_mla = [False, True]
1355	worker2.num_blocks = 1
1356	worker2.dst_num_blocks[worker2.engine_id] = worker2.num_blocks
1357	bad_meta = NixlAgentMetadata(
1358	engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1359	agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1360	kv_caches_base_addr=[0, 0],
1361	device_id=0,
1362	num_blocks=1,
1363	# WRONG: MLA region scaled by tp_ratio (it should be replicated).
1364	block_lens=[fa_len * tp_ratio, idx_len * tp_ratio],
1365	block_strides=[fa_len * tp_ratio, idx_len * tp_ratio],
1366	kv_cache_layout=worker2.kv_cache_layout,
1367	block_size=worker2.block_size,
1368	ssm_sizes=(0, 0),
1369	attn_backend_name=worker2.backend_name,
1370	physical_blocks_per_logical_kv_block=1,
1371	)
1372	with pytest.raises(AssertionError):
1373	worker2.add_remote_agent(bad_meta, remote_tp_size=1)
1374
1375	@patch(
1376	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1377	FakeNixlWrapper,
1378	)
1379	def test_handshake_validates_gqa_replicated_block_len(
1380	self, default_vllm_config, dist_init
1381	):
1382	"""Regression test for #45330.
1383
1384	When tp_size > total_num_kv_heads, GQA replication caps per-rank
1385	KV heads at 1, so block_len stops scaling with 1/tp. With 8 KV
1386	heads and D_TP=16 pulling from P_TP=8, both sides hold one head
1387	per rank and report the same block_len; the old validation
1388	expected local_block_len * tp_ratio and rejected the valid
1389	handshake.
1390	"""
1391	vllm_config = create_vllm_config()
1392
1393	with patch(
1394	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
1395	return_value=16,
1396	):
1397	connector = NixlConnector(
1398	vllm_config,
1399	KVConnectorRole.WORKER,
1400	make_kv_cache_config(block_size=16),
1401	)
1402	connector.connector_worker = FakeNixlConnectorWorker(
1403	vllm_config, connector.engine_id, hand_shake_latency=0
1404	)
1405	worker = connector.connector_worker
1406
1407	worker.transfer_topo.total_num_kv_heads = 8
1408	worker.transfer_topo.local_physical_heads = 1
1409	worker.kv_cache_layout = "LBHNC"
1410
1411	worker.slot_size_per_layer = [4096]
1412	worker.block_len_per_layer = [4096 * worker.block_size]
1413	worker.num_blocks = 1
1414	worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
1415
1416	# Remote P with TP=8 also has 1 head/rank -> identical
1417	# block_len despite tp_ratio == 2.
1418	meta = NixlAgentMetadata(
1419	engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1420	agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1421	kv_caches_base_addr=[0],
1422	device_id=0,
1423	num_blocks=1,
1424	block_lens=list(worker.block_len_per_layer),
1425	block_strides=list(worker.block_len_per_layer),
1426	kv_cache_layout="LBHNC",
1427	block_size=worker.block_size,
1428	ssm_sizes=(0, 0),
1429	attn_backend_name=worker.backend_name,
1430	physical_blocks_per_logical_kv_block=1,
1431	)
1432
1433	# Must validate cleanly (used to raise AssertionError).
1434	worker.add_remote_agent(meta, remote_tp_size=8)
1435
1436	@patch(
1437	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1438	FakeNixlWrapper,
1439	)
1440	def test_handshake_rejects_wrong_block_len_without_gqa_replication(
1441	self, default_vllm_config, dist_init
1442	):
1443	"""Ensure the head-ratio validation still rejects genuinely wrong
1444	block_lens when GQA replication is NOT in effect (32 KV heads,
1445	D_TP=4, P_TP=2: head_ratio=4, both sides have >1 head/rank).
1446	"""
1447	vllm_config = create_vllm_config()
1448
1449	with patch(
1450	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
1451	return_value=4,
1452	):
1453	connector = NixlConnector(
1454	vllm_config,
1455	KVConnectorRole.WORKER,
1456	make_kv_cache_config(block_size=16),
1457	)
1458	connector.connector_worker = FakeNixlConnectorWorker(
1459	vllm_config, connector.engine_id, hand_shake_latency=0
1460	)
1461	worker = connector.connector_worker
1462
1463	worker.transfer_topo.total_num_kv_heads = 32
1464	worker.transfer_topo.local_physical_heads = 8 # 32 // 4
1465	worker.kv_cache_layout = "LBHNC"
1466
1467	slot_size = 4096
1468	worker.slot_size_per_layer = [slot_size]
1469	worker.block_len_per_layer = [slot_size * worker.block_size]
1470	worker.num_blocks = 1
1471	worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
1472
1473	# Remote P_TP=2 has 16 heads/rank -> head_ratio = 16/8 = 2.
1474	# Correct remote block_len = local * 2. Send local * 1
1475	# (wrong) to verify rejection.
1476	bad_meta = NixlAgentMetadata(
1477	engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1478	agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1479	kv_caches_base_addr=[0],
1480	device_id=0,
1481	num_blocks=1,
1482	block_lens=list(worker.block_len_per_layer),
1483	block_strides=list(worker.block_len_per_layer),
1484	kv_cache_layout="LBHNC",
1485	block_size=worker.block_size,
1486	ssm_sizes=(0, 0),
1487	attn_backend_name=worker.backend_name,
1488	physical_blocks_per_logical_kv_block=1,
1489	)
1490
1491	with pytest.raises(AssertionError):
1492	worker.add_remote_agent(bad_meta, remote_tp_size=2)
1493
1494
1495	# NOTE: resource cleanup in mp backend is a bit finicky, so the order in which
1496	# we put here is important. First run ray, it will clean up the resources, then
1497	# the rest of the tests.
1498	@patch(
1499	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1500	FakeNixlWrapper,
1501	)
1502	def test_kv_connector_stats(default_vllm_config, dist_init):
1503	"""Test that KV transfer stats are properly recorded and retrieved."""
1504	vllm_config = create_vllm_config()
1505
1506	# Test worker role in decode server.
1507	connector = NixlConnector(
1508	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1509	)
1510	connector.connector_worker = FakeNixlConnectorWorker(
1511	vllm_config, connector.engine_id, hand_shake_latency=0
1512	)
1513
1514	# Verify that xfer_stats starts empty
1515	initial_stats = connector.get_kv_connector_stats()
1516	assert initial_stats is None
1517
1518	# Create transfer metadata
1519	request_id = "test_req_for_stats"
1520	metadata = NixlConnectorMetadata()
1521	metadata.add_new_req_to_recv(
1522	request_id=request_id,
1523	local_block_ids=([0],),
1524	kv_transfer_params={
1525	"remote_block_ids": ([0],),
1526	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1527	"remote_request_id": f"prefill-{request_id}",
1528	"remote_host": "localhost",
1529	"remote_port": 1234,
1530	"remote_tp_size": 1,
1531	},
1532	)
1533	connector.bind_connector_metadata(metadata)
1534
1535	# Start the transfer
1536	dummy_ctx = ForwardContext(
1537	no_compile_layers={},
1538	attn_metadata={},
1539	slot_mapping={},
1540	)
1541	connector.start_load_kv(dummy_ctx)
1542
1543	# Verify stats are recorded after transfer is complete
1544	max_iterations = 2
1545	# Clear metadata before start_load_kv to prevent reprocessing same request
1546	connector.bind_connector_metadata(NixlConnectorMetadata())
1547	for _ in range(max_iterations):
1548	# Need to call start_load_kv to process completed handshakes
1549	connector.start_load_kv(dummy_ctx)
1550	_, done_recving = connector.get_finished(finished_req_ids=set())
1551	if len(done_recving) > 0 and request_id in done_recving:
1552	break
1553	time.sleep(0.1) # Small delay to allow background handshake to complete
1554	else:
1555	assert "Transfer did not complete within expected iterations"
1556
1557	# Now check that stats were recorded
1558	stats_after_transfer = connector.get_kv_connector_stats()
1559	assert isinstance(stats_after_transfer, NixlKVConnectorStats)
1560
1561	# Verify stats values are recorded
1562	assert not stats_after_transfer.is_empty()
1563	assert stats_after_transfer.num_successful_transfers == 1
1564
1565	# Verify stats are reset after retrieval
1566	stats_after_reset = connector.get_kv_connector_stats()
1567	assert stats_after_reset is None
1568
1569
1570	@patch(
1571	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1572	FakeNixlWrapper,
1573	)
1574	def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init):
1575	"""reqs_to_send deadlines are stamped with the scheduler process's
1576	perf_counter, whose epoch differs across processes and (by boot-time
1577	deltas) across nodes. Without rebasing, a P worker on a node whose
1578	monotonic clock is ahead of the scheduler's by more than the TTL
1579	expires the lease on arrival and reports done_sending before D has
1580	read the blocks — the freed blocks can then be reallocated and the
1581	remote read pulls another request's data (silent accuracy corruption).
1582	The worker must anchor the remaining TTL to its own clock.
1583	"""
1584	vllm_config = create_vllm_config()
1585	connector = NixlConnector(
1586	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1587	)
1588	connector.connector_worker = FakeNixlConnectorWorker(
1589	vllm_config, connector.engine_id, hand_shake_latency=0
1590	)
1591	worker = connector.connector_worker
1592
1593	req_id = "req-lease-clock"
1594	ttl = 480.0
1595	# Simulate a scheduler whose monotonic clock is 10,000 s behind this
1596	# worker's (e.g. its node booted much later): the raw deadline is
1597	# then already far in the past in this worker's clock domain.
1598	scheduler_clock = time.perf_counter() - 10_000.0
1599
1600	metadata = NixlConnectorMetadata()
1601	metadata.reqs_in_batch = {req_id}
1602	metadata.reqs_to_send = {req_id: scheduler_clock + ttl}
1603	metadata.scheduler_clock = scheduler_clock
1604	connector.bind_connector_metadata(metadata)
1605	dummy_ctx = ForwardContext(
1606	no_compile_layers={},
1607	attn_metadata={},
1608	slot_mapping={},
1609	)
1610	connector.start_load_kv(dummy_ctx)
1611
1612	remaining = worker._reqs_to_send[req_id] - time.perf_counter()
1613	assert ttl - 5.0 < remaining <= ttl + 5.0
1614
1615	# The expiry sweep must not release the request.
1616	done_sending, _ = worker.get_finished()
1617	assert req_id not in done_sending
1618	assert req_id in worker._reqs_to_process
1619
1620
1621	def test_kv_connector_stats_aggregation():
1622	"""Test KV transfer stats aggregation across TP ranks using
1623	KVOutputAggregator (used by MultiprocExecutor).
1624	"""
1625	# Create KVOutputAggregator for 3 workers (simulating TP=3), same thing
1626	# done in MultiprocExecutor.execute_model
1627	aggregator = KVOutputAggregator(expected_finished_count=3)
1628
1629	# Create stats for multiple workers with different transfer patterns
1630	worker1_stats = NixlKVConnectorStats()
1631	worker2_stats = NixlKVConnectorStats()
1632	worker3_stats = NixlKVConnectorStats()
1633
1634	# Record different transfers on each worker
1635	# Worker 1: 2 transfers
1636	stats = get_default_xfer_telemetry()
1637	worker1_stats.record_transfer(stats)
1638	worker1_stats.record_transfer(stats)
1639
1640	# Worker 2: 1 transfer
1641	worker2_stats.record_transfer(stats)
1642
1643	# Worker 3: 3 transfers
1644	stats = get_default_xfer_telemetry(
1645	xferDurationS=2, postDurationS=2, totalBytes=2, descCount=2
1646	)
1647	worker3_stats.record_transfer(stats)
1648	worker3_stats.record_transfer(stats)
1649	worker3_stats.record_transfer(stats)
1650
1651	# Create ModelRunnerOutput instances for each worker
1652	worker_outputs = []
1653	for i, worker_stats in enumerate([worker1_stats, worker2_stats, worker3_stats]):
1654	output = ModelRunnerOutput(
1655	req_ids=[f"req_{i}"],
1656	req_id_to_index={f"req_{i}": 0},
1657	sampled_token_ids=[[123]], # dummy token
1658	logprobs=None,
1659	prompt_logprobs_dict={},
1660	pooler_output=[None],
1661	kv_connector_output=KVConnectorOutput(
1662	finished_sending=set([f"req_{i}_send"])
1663	if i < 2
1664	else None, # Workers 0,1 finished sending
1665	finished_recving=set([f"req_{i}_recv"])
1666	if i > 0
1667	else None, # Workers 1,2 finished receiving
1668	kv_connector_stats=worker_stats,
1669	),
1670	)
1671	worker_outputs.append(output)
1672
1673	# Use the real aggregation mechanism (like MultiprocExecutor.execute_model)
1674	aggregated_output = aggregator.aggregate(worker_outputs, output_rank=0)
1675	kv_connector_stats = aggregated_output.kv_connector_output.kv_connector_stats
1676	assert isinstance(kv_connector_stats, NixlKVConnectorStats)
1677	# Number of total transfers across all workers.
1678	assert kv_connector_stats.num_successful_transfers == 6
1679	# Logging proc, call reduce() to get CLI-friendly stats.
1680	cli_stats = kv_connector_stats.reduce()
1681	assert cli_stats["Avg xfer time (ms)"] == 1500.0
1682	assert cli_stats["Avg post time (ms)"] == 1500.0
1683	assert cli_stats["Avg number of descriptors"] == 1.5
1684	# Reduced values must be plain Python scalars so CLI logging renders
1685	# them without numpy reprs (eg np.float64(...)).
1686	assert all(not isinstance(v, np.generic) for v in cli_stats.values())
1687
1688
1689	def test_kv_connector_stats_failure_grouping():
1690	"""Transfer, handshake and notification failures are reported as one
1691	transport-failure count, while KV expiry is reported separately: the
1692	former are sporadic lower-transport-layer events, the latter an
1693	autoscaler signal."""
1694	stats = NixlKVConnectorStats()
1695	assert stats.is_empty()
1696
1697	stats.record_failed_transfer()
1698	stats.record_failed_handshake()
1699	stats.record_failed_notification()
1700	stats.record_kv_expired_req()
1701	stats.record_notification_after_expiry()
1702	assert not stats.is_empty()
1703
1704	# No successful transfers: latency stats are zero but the failure
1705	# counts still surface.
1706	reduced = stats.reduce()
1707	assert reduced["Num successful transfers"] == 0
1708	assert reduced["Num failed transfers"] == 3
1709	assert reduced["Num KV expired reqs"] == 1
1710	assert reduced["Num notifs after expiry"] == 1
1711
1712
1713	def test_nixl_prom_metrics_group_handshake_with_transfer_failures():
1714	"""vllm:nixl_num_failed_transfers counts handshake and notification
1715	failures too, while vllm:nixl_num_kv_expired_reqs stays a separate
1716	counter."""
1717	from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram
1718
1719	from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
1720	NixlPromMetrics,
1721	)
1722
1723	registry = CollectorRegistry()
1724
1725	class RegistryGauge(Gauge):
1726	def __init__(self, args, *kwargs):
1727	super().__init__(args, registry=registry, *kwargs)
1728
1729	class RegistryCounter(Counter):
1730	def __init__(self, args, *kwargs):
1731	super().__init__(args, registry=registry, *kwargs)
1732
1733	class RegistryHistogram(Histogram):
1734	def __init__(self, args, *kwargs):
1735	super().__init__(args, registry=registry, *kwargs)
1736
1737	vllm_config = create_vllm_config()
1738	metric_types = {
1739	Gauge: RegistryGauge,
1740	Counter: RegistryCounter,
1741	Histogram: RegistryHistogram,
1742	}
1743	prom = NixlPromMetrics(
1744	vllm_config,
1745	metric_types,
1746	labelnames=["engine"],
1747	per_engine_labelvalues={0: ["engine-0"]},
1748	)
1749
1750	stats = NixlKVConnectorStats()
1751	stats.record_failed_transfer()
1752	stats.record_failed_handshake()
1753	stats.record_failed_notification()
1754	stats.record_kv_expired_req()
1755	stats.record_notification_after_expiry()
1756	prom.observe(stats.data, engine_idx=0)
1757
1758	def counter_value(name: str) -> float:
1759	for metric in registry.collect():
1760	for sample in metric.samples:
1761	if sample.name == name:
1762	return sample.value
1763	raise AssertionError(f"metric {name} not found in registry")
1764
1765	assert counter_value("vllm:nixl_num_failed_transfers_total") == 3.0
1766	assert counter_value("vllm:nixl_num_kv_expired_reqs_total") == 1.0
1767	assert counter_value("vllm:nixl_num_notifications_after_expiry_total") == 1.0
1768
1769
1770	@patch(
1771	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1772	FakeNixlWrapper,
1773	)
1774	def test_notification_after_expiry_is_counted(default_vllm_config, dist_init):
1775	vllm_config = create_vllm_config()
1776	connector = NixlConnector(
1777	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1778	)
1779	connector.connector_worker = FakeNixlConnectorWorker(
1780	vllm_config, connector.engine_id, hand_shake_latency=0
1781	)
1782	worker = connector.connector_worker
1783	worker._reqs_to_process.add("known")
1784	worker._reqs_to_send["known"] = time.perf_counter() + 10
1785	worker.nixl_wrapper.get_new_notifs = MagicMock(
1786	return_value={"decode-agent": [b"known:1", b"unknown:1"]}
1787	)
1788
1789	assert worker._get_new_notifs() == {"known"}
1790
1791	stats = connector.get_kv_connector_stats()
1792	assert isinstance(stats, NixlKVConnectorStats)
1793	assert stats.data["num_notifications_after_expiry"] == [1]
1794
1795
1796	def test_multi_kv_connector_stats_aggregation():
1797	"""Test MultiKVConnectorStats aggregation across TP ranks using
1798	KVOutputAggregator (used by MultiprocExecutor).
1799	"""
1800	aggregator = KVOutputAggregator(expected_finished_count=3)
1801
1802	from dataclasses import dataclass
1803
1804	# Mock a KVConnectorStats class for testing aggregation over connectors.
1805	@dataclass
1806	class FooKVConnectorStats(KVConnectorStats):
1807	def reset(self):
1808	self.data = {"num_foo_transfers": 0}
1809
1810	def record_transfer(self):
1811	if "num_foo_transfers" not in self.data:
1812	self.data["num_foo_transfers"] = 0
1813	self.data["num_foo_transfers"] += 1
1814
1815	def is_empty(self) -> bool:
1816	return self.data["num_foo_transfers"] == 0
1817
1818	def aggregate(self, other: "FooKVConnectorStats") -> "FooKVConnectorStats":
1819	if not other.is_empty():
1820	self.data["num_foo_transfers"] += other.data["num_foo_transfers"]
1821	return self
1822
1823	def make_multi_stats(nixl_count: int, foo_count: int) -> MultiKVConnectorStats:
1824	data: dict[str, KVConnectorStats] = {}
1825	if nixl_count > 0:
1826	nixl_stats = NixlKVConnectorStats()
1827	for _ in range(nixl_count):
1828	nixl_stats.record_transfer(get_default_xfer_telemetry())
1829	data["NixlConnector"] = nixl_stats
1830	if foo_count > 0:
1831	foo_stats = FooKVConnectorStats()
1832	for _ in range(foo_count):
1833	foo_stats.record_transfer()
1834	data["FooConnector"] = foo_stats
1835	return MultiKVConnectorStats(data=data)
1836
1837	# Create heterogeneous stats across 3 workers
1838	worker_patterns = [(2, 1), (3, 0), (0, 5)] # (Nixl, Foo)
1839
1840	worker_outputs: list[ModelRunnerOutput] = []
1841	for i, (nixl_count, foo) in enumerate(worker_patterns):
1842	stats = make_multi_stats(nixl_count, foo)
1843	output = ModelRunnerOutput(
1844	req_ids=[f"req_{i}"],
1845	req_id_to_index={f"req_{i}": 0},
1846	sampled_token_ids=[[123]],
1847	logprobs=None,
1848	prompt_logprobs_dict={},
1849	pooler_output=[None],
1850	kv_connector_output=KVConnectorOutput(
1851	finished_sending=set([f"req_{i}_send"]) if i < 2 else None,
1852	finished_recving=set([f"req_{i}_recv"]) if i > 0 else None,
1853	kv_connector_stats=stats,
1854	),
1855	)
1856	worker_outputs.append(output)
1857
1858	aggregated_output = aggregator.aggregate(worker_outputs, output_rank=0)
1859	kv_connector_stats = aggregated_output.kv_connector_output.kv_connector_stats
1860	assert isinstance(kv_connector_stats, MultiKVConnectorStats)
1861
1862	# Validate per-connector totals across workers
1863	assert isinstance(kv_connector_stats["NixlConnector"], NixlKVConnectorStats)
1864	assert kv_connector_stats["NixlConnector"].num_successful_transfers == 5
1865	assert isinstance(kv_connector_stats["FooConnector"], FooKVConnectorStats)
1866	assert kv_connector_stats["FooConnector"].data["num_foo_transfers"] == 6
1867
1868
1869	@patch(
1870	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1871	FakeNixlWrapper,
1872	)
1873	def test_scheduler_kv_connector_stats_aggregation():
1874	"""Test scheduler and worker KV connector stats aggregation."""
1875	from vllm.v1.core.sched.output import SchedulerOutput
1876
1877	scheduler = create_scheduler(create_vllm_config())
1878
1879	# Worker stats with transfer metrics
1880	worker_stats = NixlKVConnectorStats()
1881	worker_stats.record_transfer(get_default_xfer_telemetry())
1882
1883	# Scheduler stats with custom metric (needs dummy transfer to avoid being skipped)
1884	scheduler_stats = NixlKVConnectorStats()
1885	scheduler_stats.data.update(
1886	{ # dummy transfer just for testing, to bypass is_empty() check
1887	"transfer_duration": [0],
1888	"post_duration": [0],
1889	"bytes_transferred": [0],
1890	"num_descriptors": [0],
1891	}
1892	)
1893
1894	# Mock the scheduler connector's stats method
1895	scheduler.connector.get_kv_connector_stats = lambda: MultiKVConnectorStats(
1896	data={"NixlConnector": scheduler_stats}
1897	)
1898
1899	model_output = ModelRunnerOutput(
1900	req_ids=["req_0"],
1901	req_id_to_index={"req_0": 0},
1902	sampled_token_ids=[[123]],
1903	logprobs=None,
1904	prompt_logprobs_dict={},
1905	pooler_output=[None],
1906	kv_connector_output=KVConnectorOutput(
1907	kv_connector_stats=MultiKVConnectorStats(
1908	data={"NixlConnector": worker_stats}
1909	)
1910	),
1911	)
1912	scheduler_output = SchedulerOutput(
1913	scheduled_new_reqs=[],
1914	scheduled_cached_reqs=None,
1915	num_scheduled_tokens={"req_0": 1},
1916	total_num_scheduled_tokens=1,
1917	scheduled_spec_decode_tokens={},
1918	scheduled_encoder_inputs={},
1919	num_common_prefix_blocks=[0],
1920	finished_req_ids=set(),
1921	free_encoder_mm_hashes=[],
1922	)
1923
1924	engine_core_outputs = scheduler.update_from_output(scheduler_output, model_output)
1925
1926	final_stats = next(
1927	iter(engine_core_outputs.values())
1928	).scheduler_stats.kv_connector_stats
1929	# The scheduler stats payload carries serialized per-connector dicts.
1930	nixl_stats = final_stats["NixlConnector"]
1931	assert len(nixl_stats["transfer_duration"]) == 2
1932
1933
1934	@pytest.mark.parametrize("distributed_executor_backend", ["ray", None])
1935	@patch(
1936	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1937	FakeNixlWrapper,
1938	)
1939	def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
1940	"""Test lifecycle of an aborted Remote Prefill request hitting the timeout.
1941	-----> P
1942	\| {process request}
1943	<-/--- \| {result is NOT delivered, eg proxy is down}
1944	\|
1945	\|
1946	\| {eventually free blocks}
1947	"""
1948	model_name = "Qwen/Qwen3-0.6B"
1949	timeout = 6
1950	kv_transfer_config = KVTransferConfig(
1951	kv_connector="NixlConnector",
1952	kv_role="kv_consumer",
1953	kv_connector_extra_config={"kv_lease_duration": timeout},
1954	)
1955	llm_kwargs = {
1956	"model": model_name,
1957	"enforce_eager": True,
1958	"gpu_memory_utilization": 0.5,
1959	"kv_transfer_config": kv_transfer_config,
1960	"distributed_executor_backend": distributed_executor_backend,
1961	}
1962
1963	monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
1964
1965	def run_test_and_cleanup():
1966	llm = LLM(**llm_kwargs)
1967	try:
1968	_run_abort_timeout_test(llm, timeout)
1969	finally:
1970	llm.llm_engine.engine_core.shutdown()
1971
1972	# Build runtime_env only if we're using Ray
1973	if distributed_executor_backend == "ray":
1974	with _make_fake_nixl_pkg() as working_dir:
1975	runtime_env = {
1976	"working_dir": working_dir, # ship fake nixl package
1977	"env_vars": {
1978	"NIXL_TELEMETRY_ENABLE": "1",
1979	},
1980	}
1981	# On XPU/ROCm, vLLM expects Ray's device key to be "GPU".
1982	# Explicitly reserving GPU resources here prevents false negatives
1983	# when Ray cannot auto-detect accelerator resources in test envs.
1984	ray_init_kwargs: dict[str, Any] = {"runtime_env": runtime_env}
1985	if not current_platform.is_cuda():
1986	ray_init_kwargs["num_gpus"] = 1
1987	ray.init(**ray_init_kwargs)
1988	try:
1989	run_test_and_cleanup()
1990	finally:
1991	ray.shutdown()
1992	else:
1993	run_test_and_cleanup()
1994
1995
1996	class RequestIdMapper:
1997	"""Helper class to map external request IDs to internal request IDs."""
1998
1999	def __init__(self, output_processor: OutputProcessor):
2000	self.req_id_mapping: dict[str, str] = {}
2001	self.original_add_request = output_processor.add_request
2002	output_processor.add_request = self._add_request
2003
2004	def _add_request(self, request: EngineCoreRequest, args, *kwargs):
2005	self.req_id_mapping[request.external_req_id] = request.request_id
2006	return self.original_add_request(request, args, *kwargs)
2007
2008	def __call__(self, external_req_id: str) -> str:
2009	return self.req_id_mapping[external_req_id]
2010
2011
2012	def _run_abort_timeout_test(llm: LLM, timeout: int):
2013	"""Helper function to run the abort timeout test logic."""
2014	remote_prefill_opts = {
2015	"do_remote_decode": True,
2016	"do_remote_prefill": False,
2017	"remote_engine_id": None,
2018	"remote_block_ids": None,
2019	"remote_host": None,
2020	"remote_port": None,
2021	}
2022	# Simulate sidecar request
2023	sampling_params = SamplingParams(
2024	temperature=0.0,
2025	max_tokens=1,
2026	extra_args={"kv_transfer_params": remote_prefill_opts},
2027	)
2028	scheduler = llm.llm_engine.engine_core.engine_core.scheduler
2029	req_to_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
2030	0
2031	].req_to_blocks
2032
2033	id_mapper = RequestIdMapper(llm.llm_engine.output_processor)
2034
2035	def req_id(outputs: list[RequestOutput]) -> str:
2036	assert len(outputs) == 1
2037	return id_mapper(outputs[0].request_id)
2038
2039	padding = "Just making this request a little longer so that we're sure "
2040	"we're not hitting the small-request lower bound beneath which we don't "
2041	"actually trigger the whole kv transfer, but rather just recompute the "
2042	"blocks on D."
2043	req0_id = req_id(
2044	llm.generate([f"What is the capital of Japan? {padding}"], sampling_params)
2045	)
2046
2047	# Request finished but not freed
2048	assert req0_id in scheduler.finished_req_ids and req0_id in req_to_blocks
2049	# Some other request, 0 still not freed
2050	req1_id = req_id(
2051	llm.generate([f"What is the capital of Italy? {padding}"], sampling_params)
2052	)
2053	assert req0_id in req_to_blocks
2054	assert req1_id in scheduler.finished_req_ids and req1_id in req_to_blocks
2055
2056	# Wait for timeout and trigger another scheduler loop
2057	time.sleep(timeout)
2058	_ = llm.generate([f"What is the capital of France? {padding}"], sampling_params)
2059	# Request-0 times out and is cleared!
2060	assert req0_id not in req_to_blocks
2061	# Need to shutdown the background thread to release NIXL side channel port
2062	llm.llm_engine.engine_core.shutdown()
2063
2064
2065	def test_mixed_memory_local_descriptors_split_by_memory_type():
2066	worker = object.__new__(NixlConnectorWorker)
2067	worker.transfer_topo = MagicMock()
2068	worker.block_size = 16
2069	worker.engine_id = "local"
2070	worker.tp_rank = 0
2071	worker.device_id = 3
2072	worker.kv_caches_base_addr = {"local": {0: [100, 200]}}
2073	worker._has_mamba = False
2074	worker._mixed_mem_types = True
2075	worker.region_mem_types = ["DRAM", "VRAM"]
2076	worker.region_num_blocks = [2, 2]
2077	worker._desc_is_dram_by_block_size = {}
2078	worker._desc_pos_by_block_size = {}
2079	worker._dram_src_handles_by_block_size = {}
2080	worker.nixl_memory_type = "VRAM"
2081	worker._build_fa_local = MagicMock( # type: ignore[method-assign]
2082	return_value=np.array(
2083	[
2084	[100, 10, 3],
2085	[110, 10, 3],
2086	[200, 10, 3],
2087	[210, 10, 3],
2088	]
2089	)
2090	)
2091	worker.nixl_wrapper = MagicMock()
2092	worker.nixl_wrapper.get_xfer_descs.side_effect = lambda blocks, memory_type: (
2093	memory_type,
2094	blocks,
2095	)
2096	worker.nixl_wrapper.prep_xfer_dlist.side_effect = [11, 22]
2097
2098	handle, blocks = worker.register_local_xfer_handler(worker.block_size)
2099
2100	assert handle == 22
2101	assert worker._dram_src_handles_by_block_size[worker.block_size] == 11
2102	assert [block[2] for block in blocks] == [0, 0, 3, 3]
2103	memory_types = [
2104	call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list
2105	]
2106	assert memory_types == ["DRAM", "VRAM"]
2107
2108
2109	@pytest.fixture
2110	def recv_worker():
2111	"""Receive lifecycle state without distributed or device initialization."""
2112	worker = object.__new__(NixlConnectorWorker)
2113	worker.transfer_topo = MagicMock()
2114	worker.transfer_topo.block_size_ratio.return_value = 1
2115	worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(
2116	remote_block_size=16, remote_physical_blocks_per_logical=1
2117	)
2118	worker._physical_blocks_per_logical_kv_block = 1
2119	worker._recving_metadata = {"request": MagicMock(local_block_ids=([1, 2, 3],))}
2120	worker._recving_transfers = defaultdict(list)
2121	worker._failed_recv_reqs = queue.Queue()
2122	worker._recv_failures = set()
2123	worker._handshake_lock = threading.RLock()
2124	worker._handshake_futures = {}
2125	worker._remote_agents = {}
2126	worker._engine_by_address = {}
2127	worker._replicated_pcp_done_sending = set()
2128	worker._invalid_block_ids = queue.Queue()
2129	worker._pending_recv_notifs = {}
2130	worker._reqs_to_send = {}
2131	worker._replicated_pcp_done_sending = set()
2132	worker._is_hma_required = False
2133	worker._has_mamba = False
2134	worker.use_host_buffer = False
2135	worker.enable_permute_local_kv = False
2136	worker.enable_heterogeneous_attn_post_process = False
2137	worker.tp_rank = 0
2138	worker._log_failure = MagicMock() # type: ignore[method-assign]
2139	worker.xfer_stats = NixlKVConnectorStats()
2140	worker.nixl_wrapper = MagicMock()
2141	worker.nixl_wrapper.get_xfer_telemetry.return_value = get_default_xfer_telemetry()
2142	return worker
2143
2144
2145	def test_mixed_memory_read_notifies_after_both_transfers_finish(recv_worker):
2146	worker = recv_worker
2147	worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])}
2148	worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])}
2149	worker._dram_src_handles_by_block_size = {16: 10}
2150	worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102]
2151	worker.nixl_wrapper.check_xfer_state.return_value = "DONE"
2152
2153	worker._read_blocks_mixed(
2154	request_id="request",
2155	local_block_size_key=16,
2156	local_device_handle=20,
2157	local_dram_handle=10,
2158	remote_xfer_side_handle=30,
2159	local_block_descs_ids=np.array([0, 2]),
2160	remote_block_descs_ids=np.array([5, 7]),
2161	notif_agent="prefill",
2162	notif_id=b"request:1",
2163	)
2164
2165	dram_read, device_read = worker.nixl_wrapper.make_prepped_xfer.call_args_list
2166	assert dram_read.args[:2] == ("READ", 10)
2167	assert dram_read.args[3] == 30
2168	np.testing.assert_array_equal(dram_read.args[2], [0])
2169	np.testing.assert_array_equal(dram_read.args[4], [5])
2170	assert device_read.args[:2] == ("READ", 20)
2171	assert device_read.args[3] == 30
2172	np.testing.assert_array_equal(device_read.args[2], [0])
2173	np.testing.assert_array_equal(device_read.args[4], [7])
2174	worker.nixl_wrapper.send_notif.assert_not_called()
2175
2176	assert worker.get_finished() == (set(), {"request"})
2177	worker.nixl_wrapper.send_notif.assert_called_once_with(
2178	"prefill", notif_msg=b"request:1"
2179	)
2180
2181
2182	def test_mixed_memory_read_failure_does_not_notify_producer(recv_worker):
2183	"""A failed half of a split READ must suppress its success notification."""
2184	worker = recv_worker
2185	worker._recving_transfers = {"request": [101, 102]}
2186	worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}
2187	worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"]
2188
2189	assert worker.get_finished() == (set(), {"request"})
2190
2191	assert "request" not in worker._pending_recv_notifs
2192	worker.nixl_wrapper.send_notif.assert_not_called()
2193
2194
2195	def element_byte_addrs(view: torch.Tensor) -> list[int]:
2196	"""Absolute byte addresses of every element of a (possibly strided) view."""
2197	offsets = torch.zeros(view.shape, dtype=torch.int64)
2198	for dim, (size, stride) in enumerate(zip(view.shape, view.stride())):
2199	shape = [1] * view.ndim
2200	shape[dim] = size
2201	offsets = offsets + torch.arange(size, dtype=torch.int64).view(shape) * stride
2202	esize = view.element_size()
2203	byte_offsets = (offsets.flatten() * esize).tolist()
2204	base = view.data_ptr()
2205	return [base + offset + i for offset in byte_offsets for i in range(esize)]
2206
2207
2208	@pytest.mark.parametrize(
2209	"attn_backend",
2210	[
2211	pytest.param(
2212	"FLASH_ATTN",
2213	marks=pytest.mark.skipif(
2214	current_platform.is_rocm(),
2215	reason="Attention backend FLASH_ATTN is not supported on ROCm",
2216	),
2217	),
2218	"TRITON_ATTN",
2219	],
2220	)
2221	@pytest.mark.parametrize("layout", [layout.name for layout in KVCacheLayout])
2222	@pytest.mark.parametrize("separate_kv_head_groups", [False, True])
2223	def test_register_kv_caches(
2224	default_vllm_config,
2225	dist_init,
2226	attn_backend,
2227	layout,
2228	separate_kv_head_groups,
2229	):
2230	"""Test that register_kv_caches() properly calls nixl_wrapper methods with
2231	correct data.
2232
2233	This test verifies:
2234	1. nixl_wrapper.get_reg_descs() is called with caches_data containing
2235	tensor metadata
2236	2. nixl_wrapper.get_xfer_descs() is called with blocks_data containing
2237	block layout info
2238	"""
2239	vllm_config = create_vllm_config(attention_backend=attn_backend)
2240	vllm_config.cache_config.kv_cache_layout = layout
2241
2242	# Import the appropriate backend based on the parameter
2243	if attn_backend == "FLASH_ATTN":
2244	from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
2245
2246	backend_cls = FlashAttentionBackend
2247	elif attn_backend == "ROCM_ATTN":
2248	from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
2249
2250	backend_cls = RocmAttentionBackend
2251	else: # TRITON_ATTN
2252	from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
2253
2254	backend_cls = TritonAttentionBackend
2255
2256	nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
2257	with (
2258	patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper,
2259	patch(f"{nixl_worker}.threading.Event"),
2260	patch(f"{nixl_worker}.threading.Thread") as mock_thread,
2261	patch(f"{nixl_worker}.get_current_attn_backends") as mock_get_attn_backends,
2262	):
2263	mock_get_attn_backends.return_value = [backend_cls]
2264	block_size = 16
2265	num_blocks = 8
2266	num_heads = 4
2267	head_size = 16
2268
2269	kv_cache_spec = FullAttentionSpec(
2270	block_size=block_size,
2271	num_kv_heads=num_heads,
2272	head_size=head_size,
2273	dtype=torch.float16,
2274	num_head_slots=2 if separate_kv_head_groups else None,
2275	state_content_bytes=num_heads * head_size * 2
2276	if separate_kv_head_groups
2277	else None,
2278	)
2279	kv_cache_config = KVCacheConfig(
2280	num_blocks=num_blocks,
2281	kv_cache_tensors=[],
2282	kv_cache_groups=[
2283	KVCacheGroupSpec(
2284	["layer0", "layer1", "layer2", "layer3"], kv_cache_spec
2285	)
2286	],
2287	)
2288	# Create connector
2289	connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
2290	connector.connector_worker = FakeNixlConnectorWorker(
2291	vllm_config,
2292	connector.engine_id,
2293	hand_shake_latency=0,
2294	kv_cache_config=kv_cache_config,
2295	)
2296
2297	# Get the mock instance
2298	mock_wrapper_instance = mock_nixl_wrapper.return_value
2299	connector.connector_worker.nixl_wrapper = mock_wrapper_instance
2300
2301	# Appease NixlHandshakePayload encoding with some bytes
2302	mock_wrapper_instance.get_agent_metadata.return_value = b"fake_agent_metadata"
2303
2304	# Reassure the shutdown() check that the thread is terminated
2305	mock_thread.return_value.is_alive.return_value = False
2306
2307	raw0 = torch.zeros(
2308	kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 2,
2309	dtype=torch.int8,
2310	device=current_platform.device_type,
2311	)
2312	raw1 = torch.zeros(
2313	kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks,
2314	dtype=torch.int8,
2315	device=current_platform.device_type,
2316	)
2317	tensor0, tensor1 = dense_kv_cache_views(
2318	raw0,
2319	kv_cache_spec,
2320	kv_cache_config.num_blocks,
2321	num_layers=2,
2322	layout=KVCacheLayout[layout],
2323	)
2324	(tensor2,) = dense_kv_cache_views(
2325	raw1,
2326	kv_cache_spec,
2327	kv_cache_config.num_blocks,
2328	num_layers=1,
2329	layout=KVCacheLayout[layout],
2330	)
2331	kv_caches = {
2332	"layer0": tensor0,
2333	"layer1": tensor1,
2334	"layer2": tensor2,
2335	"layer3": tensor0,
2336	}
2337
2338	# Execute register_kv_caches
2339	connector.register_kv_caches(kv_caches)
2340
2341	# Verify get_reg_descs was called with caches_data
2342	assert mock_wrapper_instance.get_reg_descs.called
2343	caches_data, _ = mock_wrapper_instance.get_reg_descs.call_args[0]
2344	assert len(caches_data) == 2
2345
2346	for cache_entry, raw in zip(caches_data, (raw0, raw1)):
2347	base_addr, size, _tp_rank, _ = cache_entry
2348	assert size == raw.nbytes
2349	assert base_addr == raw.data_ptr()
2350
2351	# Verify get_xfer_descs was called with blocks_data
2352	assert mock_wrapper_instance.get_xfer_descs.called
2353	blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
2354
2355	# Layout-blind contract: whatever regions the worker carves out,
2356	# transferring "block b" must move exactly logical block b's bytes for
2357	# every layer. Map each registered byte to the block that owns it, then
2358	# require every descriptor window to hold bytes of a single block and
2359	# the windows of block b to cover exactly block b's bytes.
2360	owner: dict[int, int] = {}
2361	for cache in (tensor0, tensor1, tensor2):
2362	for blk in range(num_blocks):
2363	for addr in element_byte_addrs(cache[blk]):
2364	assert owner.setdefault(addr, blk) == blk
2365	block_bytes = defaultdict(set)
2366	for addr, blk in owner.items():
2367	block_bytes[blk].add(addr)
2368
2369	covered: defaultdict[int, set[int]] = defaultdict(set)
2370	for block_start_addr, block_len, _tp_rank in blocks_data:
2371	window = range(block_start_addr, block_start_addr + block_len)
2372	owners = {owner[addr] for addr in window}
2373	assert len(owners) == 1, "descriptor window spans logical blocks"
2374	covered[owners.pop()].update(window)
2375	assert covered == block_bytes
2376
2377	# Region bases are exactly the block-0 window starts, one per region.
2378	base_addrs = connector.connector_worker.kv_caches_base_addr[
2379	connector.connector_worker.engine_id
2380	][0]
2381	if layout == "BLHNC":
2382	assert len(base_addrs) == 3
2383	assert set(base_addrs) == {
2384	start for start, _len, _tp in blocks_data if owner[start] == 0
2385	}
2386	assert len(blocks_data) == num_blocks * len(base_addrs)
2387
2388	assert connector.connector_worker.block_size == 16
2389
2390
2391	def test_register_packed_dsv4_mla_cache_as_single_region(
2392	default_vllm_config, dist_init
2393	):
2394	from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
2395
2396	nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
2397	with (
2398	patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper,
2399	patch(f"{nixl_worker}.threading.Event"),
2400	patch(f"{nixl_worker}.threading.Thread") as mock_thread,
2401	patch(f"{nixl_worker}.get_current_attn_backends") as mock_backends,
2402	):
2403	mock_backends.return_value = [TritonAttentionBackend]
2404	num_blocks = 2
2405	num_layers = 4
2406	block_size = 256
2407	spec = MLAAttentionSpec(
2408	block_size=block_size,
2409	num_kv_heads=1,
2410	head_size=512,
2411	dtype=torch.uint8,
2412	cache_dtype_str="fp8_ds_mla",
2413	tokens_per_state=4,
2414	alignment=576,
2415	model_version="deepseek_v4",
2416	state_content_bytes=584,
2417	)
2418	layer_names = [f"layer{idx}" for idx in range(num_layers)]
2419	kv_cache_config = KVCacheConfig(
2420	num_blocks=num_blocks,
2421	kv_cache_tensors=[],
2422	kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)],
2423	)
2424	vllm_config = create_vllm_config(attention_backend="TRITON_ATTN")
2425	vllm_config.cache_config.block_size = block_size
2426	vllm_config.cache_config.kv_cache_layout = "BLHNC"
2427	connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
2428	connector.connector_worker = FakeNixlConnectorWorker(
2429	vllm_config,
2430	connector.engine_id,
2431	hand_shake_latency=0,
2432	kv_cache_layout="BLHNC",
2433	kv_cache_config=kv_cache_config,
2434	)
2435	wrapper = mock_nixl_wrapper.return_value
2436	connector.connector_worker.nixl_wrapper = wrapper
2437	wrapper.get_agent_metadata.return_value = b"fake_agent_metadata"
2438	mock_thread.return_value.is_alive.return_value = False
2439
2440	raw = torch.zeros(
2441	spec.page_size_bytes * num_blocks * num_layers,
2442	dtype=torch.int8,
2443	device=current_platform.device_type,
2444	)
2445	views = dense_kv_cache_views(
2446	raw, spec, num_blocks, num_layers, KVCacheLayout.BLHNC
2447	)
2448	connector.register_kv_caches(dict(zip(layer_names, views)))
2449
2450	blocks_data, _ = wrapper.get_xfer_descs.call_args[0]
2451	packed_block_len = num_layers * spec.page_size_bytes
2452	assert connector.connector_worker.kv_caches_base_addr[
2453	connector.connector_worker.engine_id
2454	][0] == [raw.data_ptr()]
2455	assert blocks_data.tolist() == [
2456	[raw.data_ptr() + block_idx * packed_block_len, packed_block_len, 0]
2457	for block_idx in range(num_blocks)
2458	]
2459
2460
2461	class FakePlatform(Platform):
2462	device_type: str = "oot"
2463
2464	@classmethod
2465	def get_nixl_supported_devices(cls) -> dict[str, tuple[str, ...]]:
2466	"""Returns a mapping from device_type to a tuple of supported
2467	kv_buffer_device for nixl.
2468	"""
2469	return {"oot": ("oot",)}
2470
2471	@classmethod
2472	def get_nixl_memory_type(cls) -> str \| None:
2473	"""Returns the nixl memory type for the current platform."""
2474	return "VRAM"
2475
2476
2477	@pytest.mark.parametrize(
2478	"kv_buffer_device, nixl_memory_type",
2479	[
2480	("oot", "VRAM"),
2481	],
2482	)
2483	def test_kv_buffer_to_nixl_memory_types(
2484	default_vllm_config, dist_init, kv_buffer_device, nixl_memory_type
2485	):
2486	"""Test that register_kv_caches() passes the correct memory types from the
2487	config to the nixl_wrapper.
2488	"""
2489	vllm_config = create_vllm_config()
2490	# Override the default memory types in the config
2491	vllm_config.kv_transfer_config.kv_buffer_device = kv_buffer_device
2492	from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
2493	_NIXL_SUPPORTED_DEVICE,
2494	)
2495
2496	_NIXL_SUPPORTED_DEVICE.update(FakePlatform.get_nixl_supported_devices())
2497
2498	with (
2499	patch(
2500	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper"
2501	),
2502	patch(
2503	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.threading.Event"
2504	),
2505	patch(
2506	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.threading.Thread"
2507	),
2508	patch(
2509	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.current_platform",
2510	FakePlatform,
2511	),
2512	patch(
2513	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils._NIXL_SUPPORTED_DEVICE",
2514	_NIXL_SUPPORTED_DEVICE,
2515	),
2516	): # noqa: E501
2517	# Create connector and replace its worker with a fake one for isolation
2518	connector = NixlConnector(
2519	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
2520	)
2521
2522	# Verify get_reg_descs was called with the correct memory_type
2523	assert connector.connector_worker.kv_buffer_device == kv_buffer_device
2524	assert connector.connector_worker.nixl_memory_type == nixl_memory_type
2525
2526
2527	@patch(
2528	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2529	FakeNixlWrapper,
2530	)
2531	def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
2532	"""Test that shutdown() properly cleans up all resources."""
2533	vllm_config = create_vllm_config()
2534
2535	scheduler = NixlConnectorScheduler(
2536	vllm_config,
2537	vllm_config.kv_transfer_config.engine_id,
2538	make_kv_cache_config(block_size=16),
2539	)
2540	worker = NixlConnectorWorker(
2541	vllm_config,
2542	vllm_config.kv_transfer_config.engine_id,
2543	make_kv_cache_config(block_size=16),
2544	)
2545	nixl_wrapper = worker.nixl_wrapper
2546
2547	with (
2548	patch.object(worker, "_handshake_initiation_executor") as mock_exec,
2549	patch.object(scheduler, "_nixl_handshake_listener_t") as mock_listener,
2550	patch.object(nixl_wrapper, "release_xfer_handle") as mock_rel_xfer,
2551	patch.object(nixl_wrapper, "release_dlist_handle") as mock_rel_dlist,
2552	patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem_agent,
2553	patch.object(nixl_wrapper, "deregister_memory") as mock_dereg,
2554	):
2555	worker._recving_transfers = {"req1": [123]}
2556	# Mock register_kv_cache which registers local handle
2557	worker.src_xfer_handles_by_block_size = {worker.block_size: 455}
2558	# P TP = 2 * D TP case, we should register 2 local handles
2559	worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]}
2560	worker.dst_xfer_side_handles = {"engine1": {0: 789}}
2561	worker._remote_agents = {"engine1": {(0, 0): "agent1"}}
2562	# _cleanup_remote_engine (called by shutdown) also clears these:
2563	worker.kv_caches_base_addr["engine1"] = {0: [0xABC]}
2564	worker.dst_num_blocks["engine1"] = 50
2565	worker.tp_mappings["engine1"] = MagicMock()
2566	worker._engine_last_active["engine1"] = time.perf_counter()
2567	worker._registered_descs = ["desc1", "desc2"]
2568
2569	mock_listener.is_alive.return_value = False
2570
2571	worker.shutdown()
2572
2573	# Test idempotency
2574	worker.shutdown()
2575	worker.shutdown()
2576
2577	mock_exec.shutdown.assert_called_with(wait=False)
2578
2579	# Same sequence on scheduler.shutdown()
2580	scheduler.shutdown()
2581	scheduler.shutdown()
2582	scheduler.shutdown()
2583	mock_listener.join.assert_called_once()
2584
2585	mock_rel_xfer.assert_called_once_with(123)
2586	assert mock_rel_dlist.call_count == 4
2587	mock_rel_dlist.assert_any_call(455) # src handle (whole region)
2588	mock_rel_dlist.assert_any_call(456) # src handle (1st chunk)
2589	mock_rel_dlist.assert_any_call(457) # src handle (2nd chunk)
2590	mock_rel_dlist.assert_any_call(789) # dst handle
2591	mock_rem_agent.assert_called_once_with("agent1")
2592	assert mock_dereg.call_count == 2
2593	mock_dereg.assert_any_call("desc1")
2594	mock_dereg.assert_any_call("desc2")
2595
2596
2597	# ── TTL-based remote engine eviction tests ──────────────────────────
2598
2599
2600	def _setup_worker_with_remote_engine(
2601	engine_ttl: float = 10.0,
2602	) -> tuple[Any, str]:
2603	"""Create a worker with one remote engine registered."""
2604	vllm_config = create_vllm_config(
2605	kv_connector_extra_config={"engine_ttl": engine_ttl},
2606	)
2607	worker = NixlConnectorWorker(
2608	vllm_config,
2609	vllm_config.kv_transfer_config.engine_id,
2610	make_kv_cache_config(block_size=16),
2611	)
2612
2613	engine_id = "remote-engine-1"
2614	worker._remote_agents[engine_id] = {(0, 0): "agent_0", (0, 1): "agent_1"}
2615	worker.dst_xfer_side_handles[engine_id] = {0: 100, 1: 200}
2616	worker.kv_caches_base_addr[engine_id] = {0: [0xABC]}
2617	worker.dst_num_blocks[engine_id] = 50
2618	worker.tp_mappings[engine_id] = MagicMock()
2619	worker._engine_last_active[engine_id] = time.perf_counter()
2620
2621	worker.transfer_topo = MagicMock()
2622
2623	return worker, engine_id
2624
2625
2626	@patch(
2627	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2628	FakeNixlWrapper,
2629	)
2630	def test_engine_ttl_eviction(default_vllm_config, dist_init):
2631	"""Stale engines are evicted when TTL expires."""
2632	worker, engine_id = _setup_worker_with_remote_engine(engine_ttl=10.0)
2633	nixl_wrapper = worker.nixl_wrapper
2634
2635	with (
2636	patch.object(nixl_wrapper, "release_dlist_handle") as mock_rel,
2637	patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem,
2638	):
2639	# Make the engine stale.
2640	worker._engine_last_active[engine_id] = time.perf_counter() - 20.0
2641
2642	worker._evict_stale_engines()
2643
2644	assert engine_id not in worker._remote_agents
2645	assert engine_id not in worker.dst_xfer_side_handles
2646	assert engine_id not in worker.kv_caches_base_addr
2647	assert engine_id not in worker.dst_num_blocks
2648	assert engine_id not in worker.tp_mappings
2649	assert engine_id not in worker._engine_last_active
2650	worker.transfer_topo.unregister_remote_engine.assert_called_with(engine_id)
2651
2652	assert mock_rel.call_count == 2
2653	mock_rel.assert_any_call(100)
2654	mock_rel.assert_any_call(200)
2655
2656	assert mock_rem.call_count == 2
2657	mock_rem.assert_any_call("agent_0")
2658	mock_rem.assert_any_call("agent_1")
2659
2660
2661	@patch(
2662	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2663	FakeNixlWrapper,
2664	)
2665	def test_engine_with_inflight_transfer_is_not_evicted(default_vllm_config, dist_init):
2666	"""A transfer that outlives the TTL must keep its engine registered.
2667
2668	_engine_last_active is stamped when a read is issued and not refreshed
2669	while it runs, so a stalled transfer -- a peer that has lost its NIC holds
2670	one indefinitely -- leaves its engine looking idle. Evicting it releases
2671	the dlist handle and remote agent the transfer is still reading through.
2672	"""
2673	worker, engine_id = _setup_worker_with_remote_engine(engine_ttl=10.0)
2674	nixl_wrapper = worker.nixl_wrapper
2675
2676	request_id = "req-still-reading"
2677	worker._recving_transfers[request_id] = [MagicMock()]
2678	worker._recving_metadata[request_id] = MagicMock(
2679	remote=MagicMock(engine_id=engine_id)
2680	)
2681
2682	with (
2683	patch.object(nixl_wrapper, "release_dlist_handle") as mock_rel,
2684	patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem,
2685	):
2686	worker._engine_last_active[engine_id] = time.perf_counter() - 20.0
2687
2688	worker._evict_stale_engines()
2689
2690	assert engine_id in worker._remote_agents
2691	assert engine_id in worker.dst_xfer_side_handles
2692	mock_rel.assert_not_called()
2693	mock_rem.assert_not_called()
2694
2695	# Once the transfer is done the engine is stale like any other.
2696	del worker._recving_transfers[request_id]
2697	worker._evict_stale_engines()
2698
2699	assert engine_id not in worker._remote_agents
2700	assert mock_rem.call_count == 2
2701
2702
2703	@patch(
2704	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2705	FakeNixlWrapper,
2706	)
2707	def test_engine_ttl_disabled(default_vllm_config, dist_init):
2708	"""Eviction is disabled when engine_ttl <= 0."""
2709	worker, engine_id = _setup_worker_with_remote_engine(engine_ttl=0.0)
2710
2711	# Make the engine stale.
2712	worker._engine_last_active[engine_id] = time.perf_counter() - 9999.0
2713
2714	worker._evict_stale_engines()
2715
2716	# Nothing should be evicted.
2717	assert engine_id in worker._remote_agents
2718	assert engine_id in worker.dst_xfer_side_handles
2719
2720
2721	@pytest.mark.cpu_test
2722	class TestPeerReplacement:
2723	@pytest.fixture(autouse=True)
2724	def setup(self, monkeypatch):
2725	from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
2726	from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
2727
2728	config = create_vllm_config(kv_connector_extra_config={"engine_ttl": 0})
2729	config.kv_transfer_config.kv_buffer_device = "cpu"
2730	platform = SimpleNamespace(
2731	device_type="cpu",
2732	discover_numa_topology=lambda: [],
2733	get_nixl_memory_type=lambda: "DRAM",
2734	is_rocm=lambda: False,
2735	)
2736	with (
2737	patch.object(bw, "NixlWrapper", FakeNixlWrapper),
2738	patch.object(bw, "current_platform", platform),
2739	patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
2740	patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
2741	patch.object(
2742	bw, "get_current_attn_backends", return_value=[FlashAttentionBackend]
2743	),
2744	):
2745	self.worker = FakeNixlConnectorWorker(config, "local", hand_shake_latency=0)
2746	self.nixl = self.worker.nixl_wrapper
2747	self.transport = MagicMock(wraps=self.nixl)
2748	self.worker.nixl_wrapper = self.transport
2749	monkeypatch.setattr(
2750	self.worker._handshake_initiation_executor,
2751	"submit",
2752	lambda *args: Future(),
2753	)
2754	self._connect("old")
2755	yield
2756	self.worker.shutdown()
2757
2758	def _connect(self, engine_id="new", host="localhost", port=1234):
2759	future = self.worker._ensure_handshake(engine_id, host, port, 1)
2760	self.worker.REMOTE_ENGINE_ID = engine_id
2761	self.nixl.REMOTE_AGENT_NAME = engine_id
2762	future.set_result(self.worker._nixl_handshake(host, port, 1, engine_id))
2763
2764	def _request(self, engine_id="new", blocks=(1,), awaiting_kvs=True):
2765	metadata = NixlConnectorMetadata()
2766	metadata.reqs_to_recv[f"{engine_id}-req"] = ReqMeta(
2767	local_block_ids=(blocks,),
2768	local_physical_block_ids=(blocks,),
2769	tp_size=1,
2770	remote=RemoteMeta(
2771	([2],), "localhost", 1234, engine_id, f"{engine_id}-prefill"
2772	),
2773	awaiting_kvs=awaiting_kvs,
2774	)
2775	return metadata
2776
2777	def test_cleans_old_peer_after_confirmed_handshake(self):
2778	self._connect("healthy", host="other-host")
2779	old_handle = self.worker.dst_xfer_side_handles["old"][0]
2780	self.worker.start_load_kv(self._request())
2781	self.worker.get_transfer_results()
2782	self.transport.remove_remote_agent.assert_not_called()
2783
2784	self._connect()
2785	self.worker.get_transfer_results()
2786	self.transport.release_dlist_handle.assert_called_once_with(old_handle)
2787	self.transport.remove_remote_agent.assert_called_once_with("old")
2788	assert set(self.worker._remote_agents) == {"new", "healthy"}
2789	assert self.worker._engine_by_address == {
2790	("localhost", 1234): "new",
2791	("other-host", 1234): "healthy",
2792	}
2793	with pytest.raises(KeyError):
2794	self.worker.transfer_topo.get_engine_info("old")
2795
2796	self.worker.start_load_kv(NixlConnectorMetadata())
2797	assert self.worker.get_transfer_results().finished_recving == {"new-req"}
2798	self.transport.remove_remote_agent.assert_called_once_with("old")
2799
2800	@pytest.mark.parametrize("host,port", [("other-host", 1234), ("localhost", 5678)])
2801	def test_preserves_other_addresses(self, host, port):
2802	self._connect(host=host, port=port)
2803	self.worker.get_transfer_results()
2804	assert "old" in self.worker._remote_agents
2805	self.transport.remove_remote_agent.assert_not_called()
2806
2807	def test_preserves_old_peer_when_handshake_fails(self):
2808	self.worker.start_load_kv(self._request())
2809	self.worker._handshake_futures["new"].set_exception(
2810	RuntimeError("handshake failed")
2811	)
2812	assert self.worker.get_transfer_results().failed_recving == {"new-req"}
2813	assert "old" in self.worker._remote_agents
2814	assert self.worker._engine_by_address == {("localhost", 1234): "old"}
2815	self.transport.remove_remote_agent.assert_not_called()
2816
2817	@pytest.mark.parametrize(
2818	"states,failed",
2819	[
2820	(["PROC", "DONE"], False),
2821	(["ERR", "PROC", "DONE"], True),
2822	(["ERR", "PROC", "ERR"], True),
2823	],
2824	ids=["healthy", "failed-sibling-completes", "failed-sibling-fails"],
2825	)
2826	def test_waits_for_reads_before_cleanup_and_failure_reporting(self, states, failed):
2827	self.worker._recving_metadata.update(self._request("old").reqs_to_recv)
2828	self.worker._recving_transfers["old-req"] = [1, 2] if failed else [1]
2829	self.transport.check_xfer_state.side_effect = states
2830	self._connect()
2831
2832	result = self.worker.get_transfer_results()
2833	assert result.finished_recving == result.failed_recving == set()
2834	assert self.worker.get_block_ids_with_load_errors() == set()
2835	self.transport.remove_remote_agent.assert_not_called()
2836
2837	result = self.worker.get_transfer_results()
2838	assert result.finished_recving == {"old-req"}
2839	assert result.failed_recving == ({"old-req"} if failed else set())
2840	assert self.worker.get_block_ids_with_load_errors() == (
2841	{1} if failed else set()
2842	)
2843	assert self.transport.release_xfer_handle.call_count == (2 if failed else 1)
2844	self.transport.remove_remote_agent.assert_called_once_with("old")
2845
2846	def test_waits_for_queued_notification(self):
2847	metadata = self._request("old", blocks=())
2848	self.worker._recving_metadata.update(metadata.reqs_to_recv)
2849	self.worker._background_nixl_handshake(
2850	"old-req", "old", metadata.reqs_to_recv["old-req"]
2851	)
2852	self._connect()
2853	self.worker.get_transfer_results()
2854	self.transport.remove_remote_agent.assert_not_called()
2855
2856	self.worker.start_load_kv(NixlConnectorMetadata())
2857	self.transport.send_notif.assert_called_once_with(
2858	"old", notif_msg=b"old-prefill:1"
2859	)
2860	result = self.worker.get_transfer_results()
2861	assert result.finished_recving == {"old-req"}
2862	assert result.failed_recving == set()
2863	self.transport.remove_remote_agent.assert_called_once_with("old")
2864
2865	def test_waits_for_other_handshakes(self):
2866	self.worker._ensure_handshake("other", "other-host", 1234, 1)
2867	self._connect()
2868	self.worker.get_transfer_results()
2869	self.transport.remove_remote_agent.assert_not_called()
2870
2871	self._connect("other", host="other-host")
2872	self.worker.get_transfer_results()
2873	self.transport.remove_remote_agent.assert_called_once_with("old")
2874
2875	def test_rehandshakes_queued_request_for_released_engine(self):
2876	# A request queued after its engine's handshake, but drained only after
2877	# that engine was replaced and released, must not read from it.
2878	metadata = self._request("old", blocks=(), awaiting_kvs=False)
2879	self.worker._ready_requests.put(("old-req", metadata.reqs_to_recv["old-req"]))
2880	self._connect()
2881	self.worker.get_transfer_results()
2882	self.transport.remove_remote_agent.assert_called_once_with("old")
2883
2884	self.worker.start_load_kv(NixlConnectorMetadata())
2885	assert "old" in self.worker._handshake_futures
2886	self.transport.send_notif.assert_not_called()
2887
2888
2889	def test_transfer_topology_unregister():
2890	"""TransferTopology.unregister_remote_engine removes the engine."""
2891	from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
2892
2893	topo = TransferTopology(
2894	tp_rank=0,
2895	tp_size=1,
2896	block_size=16,
2897	engine_id="local",
2898	is_mla=False,
2899	is_mamba=False,
2900	total_num_kv_heads=4,
2901	attn_backends=[FlashAttentionBackend],
2902	)
2903
2904	info = EngineTransferInfo(
2905	remote_tp_size=1,
2906	remote_block_size=16,
2907	remote_block_len=64,
2908	remote_physical_blocks_per_logical=1,
2909	)
2910	topo.register_remote_engine("remote-1", info)
2911	assert topo.get_engine_info("remote-1") is info
2912
2913	topo.unregister_remote_engine("remote-1")
2914	with pytest.raises(KeyError):
2915	topo.get_engine_info("remote-1")
2916
2917	# Idempotent: no error on double-unregister
2918	topo.unregister_remote_engine("remote-1")
2919
2920
2921	@patch(
2922	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2923	FakeNixlWrapper,
2924	)
2925	def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_init):
2926	"""Create and schedule a request so that P adds it to in-batch tracking via
2927	the real scheduler, then simulate an abort (request not in next scheduler
2928	iteration) and verify the worker no longer tracks it as in-batch.
2929	"""
2930	vllm_config = create_vllm_config()
2931
2932	scheduler = create_scheduler(vllm_config)
2933	# KVConnector Worker in P
2934	connector = NixlConnector(
2935	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
2936	)
2937	connector.connector_worker = FakeNixlConnectorWorker(
2938	vllm_config, connector.engine_id, hand_shake_latency=0
2939	)
2940
2941	# Create a request that triggers do_remote_decode so that
2942	# the scheduler adds it to reqs_in_batch
2943	req = create_request(request_id=1, do_remote_decode=True, max_tokens=1)
2944	scheduler.add_request(req)
2945
2946	# First scheduling pass - examine build_connector_meta output
2947	sched_out = scheduler.schedule()
2948	kv_meta = sched_out.kv_connector_metadata
2949	assert kv_meta is not None
2950	assert isinstance(kv_meta, NixlConnectorMetadata)
2951	assert req.request_id in kv_meta.reqs_in_batch
2952
2953	#### Model Runner start ####
2954	# Bind scheduler-produced metadata and start worker processing.
2955	connector.bind_connector_metadata(kv_meta)
2956
2957	dummy_ctx = ForwardContext(
2958	no_compile_layers={},
2959	attn_metadata={},
2960	slot_mapping={},
2961	)
2962	connector.start_load_kv(dummy_ctx)
2963
2964	# Ensure it was tracked by the worker
2965	assert req.request_id in connector.connector_worker._reqs_to_process
2966
2967	#### Model Runner end ####
2968
2969	# Abort request - request_finished call in connector scheduler
2970	scheduler.finish_requests(req.request_id, RequestStatus.FINISHED_ABORTED)
2971	# Second scheduling pass - build metadata with aborted request
2972	sched_out2 = scheduler.schedule()
2973	kv_meta2 = sched_out2.kv_connector_metadata
2974	assert kv_meta2 is not None
2975	assert isinstance(kv_meta2, NixlConnectorMetadata)
2976	assert req.request_id not in kv_meta2.reqs_in_batch
2977
2978	# Bind empty/abort metadata and run worker step
2979	#### Model Runner start ####
2980	connector.bind_connector_metadata(kv_meta2)
2981	connector.start_load_kv(dummy_ctx)
2982
2983	# After abort, the worker should not keep tracking it as "in-batch"
2984	assert req.request_id not in connector.connector_worker._reqs_to_process
2985	#### Model Runner end ####
2986
2987
2988	class FailingNixlWrapper(FakeNixlWrapper):
2989	"""Mock NixlWrapper that fails on specific operations."""
2990
2991	def __init__(self, args, *kwargs):
2992	super().__init__(args, *kwargs)
2993	self.fail_handshake = False
2994	self.fail_transfer_setup = False
2995	self.fail_send_notif = False
2996	self.fail_transfer_state = False # Returns "ERR" state
2997	self.fail_transfer_exception = False # Raises exception in check_xfer_state
2998
2999	def add_remote_agent(self, agent_metadata: bytes) -> str:
3000	if self.fail_handshake:
3001	from zmq.error import Again
3002
3003	raise Again("Simulated timeout failure")
3004	return super().add_remote_agent(agent_metadata)
3005
3006	def make_prepped_xfer(
3007	self,
3008	xfer_type: str,
3009	local_xfer_side_handle: int,
3010	local_block_descs_ids: list[int],
3011	remote_xfer_side_handle: int,
3012	remote_block_descs_ids: list[int],
3013	notif_msg: bytes \| None = None,
3014	) -> int:
3015	if self.fail_transfer_setup:
3016	# classic RuntimeError to simulate failure
3017	raise RuntimeError("BAD STATUS")
3018	return super().make_prepped_xfer(
3019	xfer_type,
3020	local_xfer_side_handle,
3021	local_block_descs_ids,
3022	remote_xfer_side_handle,
3023	remote_block_descs_ids,
3024	notif_msg,
3025	)
3026
3027	def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
3028	if self.fail_send_notif:
3029	raise RuntimeError("Simulated send_notif failure")
3030	return super().send_notif(agent_name, notif_msg)
3031
3032	def check_xfer_state(self, handle: int) -> str:
3033	if self.fail_transfer_exception:
3034	raise RuntimeError("Simulated check_xfer_state exception")
3035	if self.fail_transfer_state:
3036	return "ERR" # Bad transfer state
3037	return super().check_xfer_state(handle)
3038
3039
3040	@patch(
3041	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
3042	FailingNixlWrapper,
3043	)
3044	@pytest.mark.parametrize("awaiting_kvs", [True, False])
3045	@pytest.mark.parametrize("notification_fails", [False, True])
3046	def test_empty_recv_is_reported_only_when_awaited(
3047	default_vllm_config, dist_init, awaiting_kvs, notification_fails
3048	):
3049	"""An empty recv is reported iff the scheduler parked the request on it.
3050
3051	_read_blocks returns early when the local block list is empty -- D already
3052	holds the KV and only the producer needs telling. Two kinds of caller reach
3053	that path, and they need opposite handling:
3054
3055	Parked (awaiting_kvs=True, a full prefix cache hit). Without a report the
3056	request is named in neither _recving_transfers nor _failed_recv_reqs and
3057	never reaches finished_recving. The scheduler has no other way to release a
3058	WAITING_FOR_REMOTE_KVS request, and there is no timeout, so it sits in
3059	kv_holding_waiting holding its blocks for the life of the process.
3060
3061	Notify-only (awaiting_kvs=False): request_finished seeding an empty recv to
3062	free P's blocks for a request aborted before it was scheduled, or a readback
3063	on a request that stays RUNNING. Reporting either one crashes the engine
3064	core -- _update_from_kv_xfer_finished asserts the id is still in
3065	self.requests and that the request is parked or finished.
3066
3067	A failed notification changes neither: the KV is local either way, and the
3068	producer frees its own blocks on a timeout.
3069	"""
3070	vllm_config = create_vllm_config()
3071	connector = NixlConnector(
3072	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
3073	)
3074	connector.connector_worker = FakeNixlConnectorWorker(
3075	vllm_config,
3076	connector.engine_id,
3077	hand_shake_latency=0.0,
3078	kv_cache_config=connector._kv_cache_config,
3079	)
3080	connector.connector_worker.nixl_wrapper.fail_send_notif = notification_fails
3081
3082	request_id = f"test_empty_recv_awaited_{awaiting_kvs}"
3083	metadata = NixlConnectorMetadata()
3084	metadata.add_new_req_to_recv(
3085	request_id=request_id,
3086	local_block_ids=(), # empty: the whole prompt is already cached locally
3087	kv_transfer_params={
3088	"remote_block_ids": [[20, 21, 22]],
3089	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3090	"remote_request_id": f"prefill-{request_id}",
3091	"remote_host": "localhost",
3092	"remote_port": 1234,
3093	"remote_tp_size": 1,
3094	},
3095	awaiting_kvs=awaiting_kvs,
3096	)
3097	connector.bind_connector_metadata(metadata)
3098	dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
3099	connector.start_load_kv(dummy_ctx)
3100	connector.bind_connector_metadata(NixlConnectorMetadata())
3101	time.sleep(0.1)
3102	connector.start_load_kv(dummy_ctx)
3103
3104	_, done_recving = connector.get_finished(finished_req_ids=set())
3105	assert (request_id in done_recving) is awaiting_kvs
3106
3107
3108	@pytest.mark.parametrize("is_hma", [False, True])
3109	@pytest.mark.parametrize(
3110	"local_block_ids,awaiting_kvs",
3111	[((), False), (([],), False), (([],), True), (([1, 2, 3],), True)],
3112	)
3113	def test_handshake_failure_reports_only_awaited_recvs(
3114	recv_worker, is_hma, local_block_ids, awaiting_kvs
3115	):
3116	"""A failed cleanup handshake must not complete an already-aborted request."""
3117	worker = recv_worker
3118	worker._is_hma_required = is_hma
3119	worker._recving_metadata.clear()
3120	worker._remote_agents = {}
3121	worker._handshake_lock = contextlib.nullcontext()
3122	worker._ready_requests = queue.Queue()
3123	worker._reqs_to_process = set()
3124	worker.pcp_rank = 0
3125	metadata = NixlConnectorMetadata()
3126	metadata.add_new_req_to_recv(
3127	request_id="request",
3128	local_block_ids=local_block_ids,
3129	kv_transfer_params={
3130	"remote_block_ids": ([4, 5, 6],),
3131	"remote_engine_id": "prefill",
3132	"remote_request_id": "prefill-request",
3133	"remote_host": "localhost",
3134	"remote_port": 1234,
3135	},
3136	awaiting_kvs=awaiting_kvs,
3137	)
3138	handshake = Future[None]()
3139	handshake.set_exception(RuntimeError("handshake failed"))
3140	with patch.object(worker, "_ensure_handshake", return_value=handshake):
3141	worker.start_load_kv(metadata)
3142
3143	results = worker.get_transfer_results()
3144	expected = {"request"} if awaiting_kvs else set()
3145	assert results.finished_recving == results.failed_recving == expected
3146	assert worker.get_block_ids_with_load_errors() == (
3147	set(local_block_ids[0]) if awaiting_kvs and not is_hma else set()
3148	)
3149	assert not worker._recving_metadata
3150	assert not worker._recv_failures
3151	assert not worker._pending_recv_notifs
3152
3153
3154	@patch(
3155	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
3156	FailingNixlWrapper,
3157	)
3158	@pytest.mark.parametrize(
3159	"failure_type,wrapper_config,needs_get_finished",
3160	[
3161	("transfer_setup_failed", {"fail_transfer_setup": True}, False),
3162	("handshake_failed", {"fail_handshake": True}, False),
3163	("notification_failed", {"fail_send_notif": True}, False),
3164	("transfer_failed", {"fail_transfer_state": True}, True),
3165	("transfer_exception", {"fail_transfer_exception": True}, True),
3166	],
3167	)
3168	@pytest.mark.parametrize("enable_hma", [False, True])
3169	def test_transfer_failure_logging(
3170	default_vllm_config,
3171	dist_init,
3172	failure_type,
3173	wrapper_config,
3174	needs_get_finished,
3175	enable_hma,
3176	):
3177	"""Test that transfer failures are logged with structured context.
3178
3179	Run with `pytest -sv` to see the log output.
3180
3181	Covers failure types:
3182	- transfer_setup_failed: make_prepped_xfer fails
3183	- handshake_failed: add_remote_agent fails during request handshake
3184	- notification_failed: send_notif fails
3185	- transfer_failed: check_xfer_state returns bad state (e.g., "ERR")
3186	- transfer_exception: check_xfer_state raises exception
3187	"""
3188	import logging
3189
3190	vllm_config = create_vllm_config()
3191
3192	connector = NixlConnector(
3193	vllm_config,
3194	KVConnectorRole.WORKER,
3195	make_kv_cache_config(block_size=16, swa_enabled=enable_hma),
3196	)
3197	connector.connector_worker = FakeNixlConnectorWorker(
3198	vllm_config,
3199	connector.engine_id,
3200	hand_shake_latency=0.0,
3201	kv_cache_config=connector._kv_cache_config,
3202	)
3203
3204	# Configure FailingNixlWrapper to fail in the specified way
3205	for key, value in wrapper_config.items():
3206	setattr(connector.connector_worker.nixl_wrapper, key, value)
3207
3208	request_id = f"test_{failure_type}_req"
3209
3210	# For notification_failed, we need empty local blocks
3211	# (full cache hit path to trigger send_notif)
3212	local_blocks: tuple[()] \| tuple[list[int], ...]
3213	if enable_hma:
3214	# HMA enabled: multiple groups (FA + SW)
3215	local_blocks = (
3216	() if failure_type == "notification_failed" else ([10, 11, 12], [13, 14])
3217	)
3218	remote_blocks = [[20, 21, 22], [23, 24]]
3219	else:
3220	# HMA disabled: single group
3221	local_blocks = () if failure_type == "notification_failed" else ([10, 11, 12],)
3222	remote_blocks = [[20, 21, 22]]
3223
3224	metadata = NixlConnectorMetadata()
3225	metadata.add_new_req_to_recv(
3226	request_id=request_id,
3227	local_block_ids=local_blocks,
3228	kv_transfer_params={
3229	"remote_block_ids": remote_blocks,
3230	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3231	"remote_request_id": f"prefill-{request_id}",
3232	"remote_host": "localhost",
3233	"remote_port": 1234,
3234	"remote_tp_size": 1,
3235	},
3236	)
3237	connector.bind_connector_metadata(metadata)
3238
3239	dummy_ctx = ForwardContext(
3240	no_compile_layers={},
3241	attn_metadata={},
3242	slot_mapping={},
3243	)
3244
3245	# Capture logs from the nixl connector loggers
3246	# vLLM loggers have propagate=False, so we need to capture directly
3247	nixl_logger = logging.getLogger(
3248	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
3249	)
3250	pull_logger = logging.getLogger(
3251	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker"
3252	)
3253	captured_logs: list[logging.LogRecord] = []
3254
3255	class LogCapture(logging.Handler):
3256	def emit(self, record):
3257	captured_logs.append(record)
3258
3259	handler = LogCapture()
3260	handler.setLevel(logging.ERROR)
3261	nixl_logger.addHandler(handler)
3262	pull_logger.addHandler(handler)
3263
3264	try:
3265	connector.start_load_kv(dummy_ctx)
3266	# Process the ready_requests queue (for async handshake)
3267	connector.bind_connector_metadata(NixlConnectorMetadata())
3268	# Wait for async handshake to complete
3269	time.sleep(0.2)
3270	connector.start_load_kv(dummy_ctx)
3271
3272	# For transfer_failed/transfer_exception, the error happens in
3273	# get_finished() when checking transfer state
3274	if needs_get_finished:
3275	connector.get_finished(finished_req_ids=set())
3276	finally:
3277	nixl_logger.removeHandler(handler)
3278	pull_logger.removeHandler(handler)
3279
3280	# Print logs for manual comparison between commits
3281	error_logs = [r for r in captured_logs if r.levelno >= logging.ERROR]
3282	print("\n" + "=" * 60)
3283	print(f"CAPTURED ERROR LOGS for {failure_type}:")
3284	print("=" * 60)
3285	for i, record in enumerate(error_logs):
3286	print(f"\n--- Log {i + 1} ---")
3287	print(f"Message: {record.message}")
3288	print("=" * 60 + "\n")
3289
3290	assert len(error_logs) >= 1, f"Expected at least one error log for {failure_type}"
3291
3292	# Verify structured logging output (new format)
3293	# Check that at least one log matches the expected format
3294	all_messages = [r.message for r in error_logs]
3295	combined_logs = "\n".join(all_messages)
3296
3297	assert any("NIXL transfer failure" in msg for msg in all_messages), (
3298	f"Expected structured log format with 'NIXL transfer failure' prefix "
3299	f"for {failure_type}. Got: {all_messages}"
3300	)
3301	assert any("failure_type" in msg for msg in all_messages), (
3302	f"Expected 'failure_type' in logs. Got: {all_messages}"
3303	)
3304	assert any("Context:" in msg for msg in all_messages), (
3305	f"Expected 'Context:' in logs. Got: {all_messages}"
3306	)
3307	# Check that the expected failure_type appears in at least one log
3308	# Note: handshake_failed also triggers handshake_setup_failed
3309	assert failure_type in combined_logs or (
3310	failure_type == "handshake_failed" and "handshake_setup_failed" in combined_logs
3311	), f"Expected '{failure_type}' in logs. Got: {all_messages}"
3312
3313
3314	@patch(
3315	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
3316	FailingNixlWrapper,
3317	)
3318	def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
3319	"""Test that handshake failures mark blocks invalid and return via get_finished."""
3320	vllm_config = create_vllm_config()
3321
3322	connector = NixlConnector(
3323	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
3324	)
3325	connector.connector_worker = FakeNixlConnectorWorker(
3326	vllm_config, connector.engine_id, hand_shake_latency=0.1
3327	)
3328	connector.connector_worker.nixl_wrapper.fail_handshake = True
3329
3330	request_id = "test_handshake_fail"
3331	metadata = NixlConnectorMetadata()
3332	metadata.add_new_req_to_recv(
3333	request_id=request_id,
3334	local_block_ids=([1, 2, 3],),
3335	kv_transfer_params={
3336	"remote_block_ids": ([4, 5, 6],),
3337	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3338	"remote_request_id": f"prefill-{request_id}",
3339	"remote_host": "localhost",
3340	"remote_port": 1234,
3341	"remote_tp_size": 1,
3342	},
3343	)
3344	connector.bind_connector_metadata(metadata)
3345
3346	dummy_ctx = ForwardContext(
3347	no_compile_layers={},
3348	attn_metadata={},
3349	slot_mapping={},
3350	)
3351	connector.start_load_kv(dummy_ctx)
3352
3353	# Wait for handshake to fail
3354	time.sleep(0.3)
3355
3356	assert connector.get_block_ids_with_load_errors() == set()
3357	_, done_recving = connector.get_finished(finished_req_ids=set())
3358	assert request_id in done_recving
3359	assert connector.get_block_ids_with_load_errors() == {1, 2, 3}
3360
3361	# Handshake failures are recorded as transport failures, separately
3362	# from KV expiry.
3363	assert connector.connector_worker.xfer_stats.data["num_failed_handshakes"]
3364	assert connector.connector_worker.xfer_stats.data["num_failed_transfers"] == []
3365
3366
3367	@patch(
3368	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
3369	FailingNixlWrapper,
3370	)
3371	@pytest.mark.parametrize("is_hma", [False, True])
3372	def test_transfer_setup_failure_returns_finished(
3373	default_vllm_config, dist_init, is_hma
3374	):
3375	"""Setup failures report the request; only non-HMA reports block IDs."""
3376	vllm_config = create_vllm_config()
3377
3378	connector = NixlConnector(
3379	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
3380	)
3381	connector.connector_worker = FakeNixlConnectorWorker(
3382	vllm_config, connector.engine_id, hand_shake_latency=0
3383	)
3384	connector.connector_worker.nixl_wrapper.fail_transfer_setup = True
3385	connector.connector_worker._is_hma_required = is_hma
3386
3387	request_id = "test_transfer_fail"
3388	metadata = NixlConnectorMetadata()
3389	metadata.add_new_req_to_recv(
3390	request_id=request_id,
3391	local_block_ids=([7, 8, 9],),
3392	kv_transfer_params={
3393	"remote_block_ids": ([10, 11, 12],),
3394	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3395	"remote_request_id": f"prefill-{request_id}",
3396	"remote_host": "localhost",
3397	"remote_port": 1234,
3398	"remote_tp_size": 1,
3399	},
3400	)
3401	connector.bind_connector_metadata(metadata)
3402
3403	dummy_ctx = ForwardContext(
3404	no_compile_layers={},
3405	attn_metadata={},
3406	slot_mapping={},
3407	)
3408	connector.start_load_kv(dummy_ctx)
3409
3410	# Wait for handshake to complete and process ready_requests
3411	connector.bind_connector_metadata(NixlConnectorMetadata())
3412	time.sleep(0.1)
3413	connector.start_load_kv(dummy_ctx)
3414
3415	assert connector.get_block_ids_with_load_errors() == set()
3416	results = connector.get_transfer_results(finished_req_ids=set())
3417	assert request_id in results.finished_recving
3418	assert results.failed_recving == {request_id}
3419	invalid_blocks = connector.get_block_ids_with_load_errors()
3420	assert invalid_blocks == (set() if is_hma else {7, 8, 9})
3421
3422
3423	class _ScriptedXferWrapper(FakeNixlWrapper):
3424	"""Scripts per-handle xfer states; forbids releasing in-flight handles."""
3425
3426	def __init__(self, args, *kwargs):
3427	super().__init__(args, *kwargs)
3428	# handle -> successive states; the last one repeats.
3429	self.xfer_states: dict[int, list[str]] = {}
3430	self.released: list[int] = []
3431
3432	def check_xfer_state(self, handle: int) -> str:
3433	states = self.xfer_states[handle]
3434	return states.pop(0) if len(states) > 1 else states[0]
3435
3436	def release_xfer_handle(self, handle: int) -> None:
3437	# A posted-but-unfinished transfer cannot be aborted: releasing it
3438	# leaves the RDMA READ armed. Production code must never do this.
3439	assert self.xfer_states.get(handle, ["DONE"])[0] != "PROC", (
3440	f"released in-flight handle {handle}"
3441	)
3442	self.released.append(handle)
3443
3444
3445	def _make_split_read_connector(vllm_config, request_id, states):
3446	"""Seed a request with scripted in-flight xfer handles + recv metadata."""
3447	connector = NixlConnector(
3448	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
3449	)
3450	connector.connector_worker = FakeNixlConnectorWorker(
3451	vllm_config, connector.engine_id, hand_shake_latency=0
3452	)
3453	worker = connector.connector_worker
3454	wrapper = _ScriptedXferWrapper("agent")
3455	worker.nixl_wrapper = wrapper
3456	metadata = NixlConnectorMetadata()
3457	metadata.add_new_req_to_recv(
3458	request_id=request_id,
3459	local_block_ids=([7, 8, 9],),
3460	kv_transfer_params={
3461	"remote_block_ids": ([10, 11, 12],),
3462	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3463	"remote_request_id": f"prefill-{request_id}",
3464	"remote_host": "localhost",
3465	"remote_port": 1234,
3466	"remote_tp_size": 1,
3467	},
3468	)
3469	worker._recving_metadata[request_id] = metadata.reqs_to_recv[request_id]
3470	wrapper.xfer_states = dict(states)
3471	worker._recving_transfers[request_id] = list(states)
3472	return connector, worker, wrapper
3473
3474
3475	@patch(
3476	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
3477	FakeNixlWrapper,
3478	)
3479	def test_split_read_failure_defers_report_until_last_handle(
3480	default_vllm_config, dist_init
3481	):
3482	"""One half of a split (mixed DRAM/VRAM) read failing must not report the
3483	request — nor invalidate its blocks — while the sibling xfer is still in
3484	flight: a posted READ cannot be aborted and would DMA into blocks the
3485	scheduler could free and reuse. The report happens exactly once, when the
3486	last handle is terminal."""
3487	request_id = "split_read_partial_failure"
3488	err_handle, live_handle = 11, 22
3489	connector, worker, wrapper = _make_split_read_connector(
3490	create_vllm_config(),
3491	request_id,
3492	{err_handle: ["ERR"], live_handle: ["PROC", "PROC", "DONE"]},
3493	)
3494
3495	# Poll 1: one half fails; the sibling is in flight -> nothing reported.
3496	_, done_recving = connector.get_finished(finished_req_ids=set())
3497	assert done_recving == set()
3498	assert connector.get_block_ids_with_load_errors() == set()
3499	assert request_id in worker._recving_metadata
3500	assert wrapper.released == [err_handle]
3501
3502	# Poll 2: sibling still in flight.
3503	_, done_recving = connector.get_finished(finished_req_ids=set())
3504	assert done_recving == set()
3505
3506	# Poll 3: sibling terminal -> reported exactly once, blocks invalidated.
3507	_, done_recving = connector.get_finished(finished_req_ids=set())
3508	assert done_recving == {request_id}
3509	assert connector.get_block_ids_with_load_errors() == {7, 8, 9}
3510	assert request_id not in worker._recving_metadata
3511	assert wrapper.released == [err_handle, live_handle]
3512
3513	# Poll 4: nothing left; no double report.
3514	_, done_recving = connector.get_finished(finished_req_ids=set())
3515	assert done_recving == set()
3516
3517
3518	@patch(
3519	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
3520	FailingNixlWrapper,
3521	)
3522	@pytest.mark.parametrize(
3523	"failure_mode",
3524	[
3525	"handshake",
3526	"transfer_setup",
3527	"transfer_failed",
3528	"transfer_exception",
3529	],
3530	)
3531	def test_failed_request_skips_kv_postprocessing(
3532	default_vllm_config, dist_init, failure_mode
3533	):
3534	"""Test that failed requests skip KV sync and post-processing in
3535	get_transfer_results().
3536
3537	This is the core safety behavior: when a KV transfer fails at any stage,
3538	the request must still appear in done_recving (so the scheduler can apply
3539	kv_load_failure_policy), but sync_recved_kv_to_device and post-processing
3540	must NOT be called since no valid KV data was received.
3541
3542	Covers all failure paths that involve an actual (attempted) KV transfer:
3543	- handshake: add_remote_agent raises during async handshake
3544	- transfer_setup: make_prepped_xfer raises before handle is in transfers
3545	- transfer_failed: check_xfer_state returns bad state ("ERR") in
3546	_pop_done_transfers — this is the path that previously had the bug
3547	where post-processing was NOT skipped
3548	- transfer_exception: check_xfer_state raises in _pop_done_transfers
3549
3550	Note: notification_failed (send_notif raises on the full-cache-hit path)
3551	is intentionally excluded. That path is a best-effort D→P courtesy
3552	notification; the blocks are already in D's cache, so no KV transfer
3553	was attempted and done_recving is correctly empty.
3554	"""
3555	# Map each failure mode to the FailingNixlWrapper attribute to set.
3556	_WRAPPER_CONFIG: dict[str, str] = {
3557	"handshake": "fail_handshake",
3558	"transfer_setup": "fail_transfer_setup",
3559	"transfer_failed": "fail_transfer_state",
3560	"transfer_exception": "fail_transfer_exception",
3561	}
3562
3563	# Use enable_permute_local_kv=True so that
3564	# post_process_device_kv_on_receive would be called on the success path,
3565	# making the assertion meaningful (not trivially true).
3566	vllm_config = create_vllm_config(enable_permute_local_kv=True)
3567
3568	connector = NixlConnector(
3569	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
3570	)
3571	connector.connector_worker = FakeNixlConnectorWorker(
3572	vllm_config,
3573	connector.engine_id,
3574	hand_shake_latency=0.1 if failure_mode == "handshake" else 0,
3575	)
3576	worker = connector.connector_worker
3577	setattr(worker.nixl_wrapper, _WRAPPER_CONFIG[failure_mode], True)
3578
3579	request_id = f"test_{failure_mode}_skip_postprocess"
3580	metadata = NixlConnectorMetadata()
3581	metadata.add_new_req_to_recv(
3582	request_id=request_id,
3583	local_block_ids=([1, 2, 3],),
3584	kv_transfer_params={
3585	"remote_block_ids": ([4, 5, 6],),
3586	"remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3587	"remote_request_id": f"prefill-{request_id}",
3588	"remote_host": "localhost",
3589	"remote_port": 1234,
3590	"remote_tp_size": 1,
3591	},
3592	)
3593	connector.bind_connector_metadata(metadata)
3594
3595	dummy_ctx = ForwardContext(
3596	no_compile_layers={},
3597	attn_metadata={},
3598	slot_mapping={},
3599	)
3600	connector.start_load_kv(dummy_ctx)
3601
3602	if failure_mode == "handshake":
3603	# Wait for async handshake to fail.
3604	time.sleep(0.3)
3605	else:
3606	# All other modes: let the handshake complete, then process the
3607	# ready_requests queue. For transfer_failed / transfer_exception the
3608	# handle ends up in _recving_transfers; the failure surfaces in
3609	# get_finished() via _pop_done_transfers below.
3610	connector.bind_connector_metadata(NixlConnectorMetadata())
3611	time.sleep(0.1)
3612	connector.start_load_kv(dummy_ctx)
3613
3614	# Spy on sync_recved_kv_to_device and post_process_device_kv_on_receive
3615	# to verify they are NOT called for the failed request.
3616	with (
3617	patch.object(worker, "sync_recved_kv_to_device") as mock_sync,
3618	patch.object(worker, "post_process_device_kv_on_receive") as mock_postprocess,
3619	):
3620	results = connector.get_transfer_results(finished_req_ids=set())
3621
3622	assert request_id in results.finished_recving
3623	assert results.failed_recving == {request_id}
3624
3625	# Critical: KV sync and post-processing must NOT have been called
3626	# since no valid KV data was received for the failed request.
3627	mock_sync.assert_not_called()
3628	mock_postprocess.assert_not_called()
3629
3630	# Metadata for the request should have been cleaned up.
3631	assert request_id not in worker._recving_metadata
3632
3633	# Blocks should have been marked as invalid.
3634	invalid_blocks = connector.get_block_ids_with_load_errors()
3635	assert invalid_blocks == {1, 2, 3}
3636
3637
3638	@pytest.mark.parametrize("sibling_state", ["ERR", "DONE"])
3639	def test_recv_failure_waits_for_sibling_transfer(recv_worker, sibling_state):
3640	"""Failure must not release blocks while a sibling can still write to them."""
3641	worker = recv_worker
3642	worker._recving_transfers["request"] = [101, 102]
3643	worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}
3644	worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "PROC", sibling_state]
3645
3646	results = worker.get_transfer_results()
3647	assert results.finished_recving == results.failed_recving == set()
3648	assert worker.get_block_ids_with_load_errors() == set()
3649	assert "request" in worker._recving_metadata
3650	worker.nixl_wrapper.release_xfer_handle.assert_called_once_with(101)
3651
3652	results = worker.get_transfer_results()
3653	assert results.finished_recving == results.failed_recving == {"request"}
3654	assert worker.get_block_ids_with_load_errors() == {1, 2, 3}
3655	worker.nixl_wrapper.send_notif.assert_not_called()
3656
3657	results = worker.get_transfer_results()
3658	assert results.finished_recving == results.failed_recving == set()
3659	assert worker.get_block_ids_with_load_errors() == set()
3660
3661
3662	def test_recv_failure_waits_for_unpollable_handle_to_be_released(recv_worker):
3663	"""A polling exception and failed release do not prove that DMA has stopped."""
3664	worker = recv_worker
3665	worker._recving_transfers["request"] = [101]
3666	worker.nixl_wrapper.check_xfer_state.side_effect = [RuntimeError("poll"), "DONE"]
3667	worker.nixl_wrapper.release_xfer_handle.side_effect = [
3668	RuntimeError("release"),
3669	None,
3670	]
3671
3672	results = worker.get_transfer_results()
3673	assert results.finished_recving == results.failed_recving == set()
3674	assert worker.get_block_ids_with_load_errors() == set()
3675	assert "request" in worker._recving_metadata
3676
3677	results = worker.get_transfer_results()
3678	assert results.finished_recving == results.failed_recving == {"request"}
3679	assert worker.get_block_ids_with_load_errors() == {1, 2, 3}
3680	assert worker.nixl_wrapper.release_xfer_handle.call_count == 2
3681
3682
3683	def _set_test_speculative_config(
3684	vllm_config,
3685	*,
3686	method: str = "eagle3",
3687	model: str = "test/eagle3-drafter",
3688	revision: str \| None = None,
3689	code_revision: str \| None = None,
3690	num_speculative_tokens: int = 1,
3691	parallel_drafting: bool = False,
3692	kv_cache_dtype: str \| None = None,
3693	attention_backend: str \| None = None,
3694	auxiliary_layer_ids: tuple[int, ...] = (2, 16, 29),
3695	) -> None:
3696	draft_model_config = SimpleNamespace(
3697	model=model,
3698	revision=revision,
3699	code_revision=code_revision,
3700	hf_config=SimpleNamespace(eagle_aux_hidden_state_layer_ids=auxiliary_layer_ids),
3701	)
3702	vllm_config.speculative_config = SimpleNamespace(
3703	method=method,
3704	draft_model_config=draft_model_config,
3705	num_speculative_tokens=num_speculative_tokens,
3706	parallel_drafting=parallel_drafting,
3707	kv_cache_dtype=kv_cache_dtype,
3708	attention_backend=attention_backend,
3709	use_eagle=lambda: True,
3710	)
3711
3712
3713	@pytest.mark.parametrize(
3714	"remote_overrides,should_match",
3715	[
3716	({}, True),
3717	({"num_speculative_tokens": 2}, True),
3718	({"method": "mtp"}, False),
3719	({"model": "test/different-drafter"}, False),
3720	({"revision": "different-revision"}, False),
3721	({"parallel_drafting": True}, False),
3722	({"kv_cache_dtype": "fp8"}, False),
3723	# attention_backend is intentionally not part of the compat hash
3724	# (see _get_speculative_compatibility_factors); overriding it must
3725	# not change the hash.
3726	({"attention_backend": "FLASHINFER"}, True),
3727	({"auxiliary_layer_ids": (2, 16, 30)}, False),
3728	],
3729	)
3730	@pytest.mark.skip_global_cleanup
3731	def test_speculative_config_compatibility_hash(
3732	remote_overrides: dict[str, Any], should_match: bool
3733	):
3734	local_config = create_vllm_config()
3735	remote_config = create_vllm_config()
3736	_set_test_speculative_config(local_config)
3737	_set_test_speculative_config(remote_config, **remote_overrides)
3738
3739	local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN")
3740	remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN")
3741
3742	assert (local_hash == remote_hash) is should_match
3743
3744
3745	@pytest.mark.skip_global_cleanup
3746	def test_missing_speculative_config_changes_compatibility_hash():
3747	regular_config = create_vllm_config()
3748	speculative_config = create_vllm_config()
3749	_set_test_speculative_config(speculative_config)
3750
3751	regular_hash = compute_nixl_compatibility_hash(regular_config, "FLASH_ATTN")
3752	speculative_hash = compute_nixl_compatibility_hash(speculative_config, "FLASH_ATTN")
3753
3754	assert regular_hash != speculative_hash
3755
3756
3757	@pytest.mark.skip_global_cleanup
3758	def test_speculative_kv_cache_dtype_resolves_to_target():
3759	# The draft kv_cache_dtype override defaults to None ("inherit the target's
3760	# --kv-cache-dtype"). An explicit setting on one side that matches the
3761	# other side's inherited (resolved) dtype must not spuriously mismatch.
3762	local_config = create_vllm_config(cache_dtype="fp8")
3763	remote_config = create_vllm_config(cache_dtype="fp8")
3764	_set_test_speculative_config(local_config, kv_cache_dtype="fp8") # explicit
3765	_set_test_speculative_config(remote_config, kv_cache_dtype=None) # inherits
3766
3767	local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN")
3768	remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN")
3769
3770	assert local_hash == remote_hash
3771
3772
3773	@pytest.mark.skip_global_cleanup
3774	def test_speculative_attention_backend_not_in_compatibility_hash():
3775	# The draft attention_backend is intentionally excluded from the hash: the
3776	# connector only has the raw override (auto-select), and its transfer-
3777	# relevant effect is validated per region at runtime. Differing overrides
3778	# must not change the hash.
3779	local_config = create_vllm_config()
3780	remote_config = create_vllm_config()
3781	_set_test_speculative_config(local_config, attention_backend=None)
3782	_set_test_speculative_config(remote_config, attention_backend="FLASHINFER")
3783
3784	local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN")
3785	remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN")
3786
3787	assert local_hash == remote_hash
3788
3789
3790	@pytest.mark.skip_global_cleanup
3791	def test_transfer_mode_changes_compatibility_hash():
3792	# push (WRITE) and pull (READ) connectors use incompatible transfer
3793	# protocols, so their compatibility hashes must differ; identical modes
3794	# must match. The default mode is pull.
3795	config = create_vllm_config()
3796
3797	pull_hash = compute_nixl_compatibility_hash(
3798	config, "FLASH_ATTN", transfer_mode="pull"
3799	)
3800	push_hash = compute_nixl_compatibility_hash(
3801	config, "FLASH_ATTN", transfer_mode="push"
3802	)
3803
3804	assert pull_hash != push_hash
3805	assert pull_hash == compute_nixl_compatibility_hash(
3806	config, "FLASH_ATTN", transfer_mode="pull"
3807	)
3808	assert compute_nixl_compatibility_hash(config, "FLASH_ATTN") == pull_hash
3809
3810
3811	@pytest.mark.skip_global_cleanup
3812	@pytest.mark.parametrize("mode", ["pull", "push"])
3813	@pytest.mark.parametrize(
3814	"tp_size,pcp_size,dcp_size,transfer_tp_size",
3815	[(1, 1, 1, 1), (4, 1, 1, 4), (4, 1, 4, 4), (1, 4, 1, 1), (1, 4, 4, 4)],
3816	)
3817	def test_scheduler_advertises_transfer_topology(
3818	mode, tp_size, pcp_size, dcp_size, transfer_tp_size
3819	):
3820	"""The consumer must address every distinct producer KV shard."""
3821	from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
3822	NixlPullConnectorScheduler,
3823	)
3824	from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import (
3825	NixlPushConnectorScheduler,
3826	)
3827
3828	config = create_vllm_config(kv_role="kv_producer")
3829	config.parallel_config.tensor_parallel_size = tp_size
3830	config.parallel_config.prefill_context_parallel_size = pcp_size
3831	config.parallel_config.decode_context_parallel_size = dcp_size
3832	cls = NixlPullConnectorScheduler if mode == "pull" else NixlPushConnectorScheduler
3833	scheduler = cls(config, "prefiller", make_kv_cache_config(block_size=16))
3834	request = create_request(request_id=1, num_tokens=32, do_remote_decode=True)
3835	request.status = RequestStatus.FINISHED_LENGTH_CAPPED
3836	try:
3837	delay, params = scheduler.request_finished(request, ([0, 1],))
3838	assert delay
3839	assert params["transfer_mode"] == mode
3840	assert params["tp_size"] == transfer_tp_size
3841	assert params["dcp_size"] == dcp_size
3842	finally:
3843	scheduler.shutdown()
3844
3845
3846	@pytest.mark.parametrize(
3847	"mismatch_type,config_overrides,version_override,should_fail,enforce_handshake_compat",
3848	[
3849	("vllm_version", {}, {"vllm_version": "0.6.1"}, True, True),
3850	("nixl_connector_version", {}, {"connector_version": 37}, True, True),
3851	("model_name", {"model": "facebook/opt-350m"}, {}, True, True),
3852	("dtype", {"dtype": "bfloat16"}, {}, True, True),
3853	("cache_dtype", {"cache_dtype": "fp8"}, {}, True, True),
3854	("num_kv_heads", {"hf_overrides": {"num_key_value_heads": 8}}, {}, True, True),
3855	(
3856	"num_hidden_layers",
3857	{"hf_overrides": {"num_hidden_layers": 24}},
3858	{},
3859	True,
3860	True,
3861	),
3862	("hidden_size", {"hf_overrides": {"hidden_size": 1536}}, {}, True, True),
3863	("block_size", {"block_size": 8}, {}, False, True),
3864	("matching_config", {}, {}, False, True),
3865	("escape_hatch", {"model": "facebook/opt-350m"}, {}, False, False),
3866	],
3867	)
3868	@patch(
3869	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
3870	FakeNixlWrapper,
3871	)
3872	def test_compatibility_hash_validation(
3873	default_vllm_config,
3874	dist_init,
3875	mismatch_type,
3876	config_overrides,
3877	version_override,
3878	should_fail,
3879	enforce_handshake_compat,
3880	):
3881	"""Test NIXL compatibility hash validation during handshake.
3882
3883	Parameters
3884	----------
3885	mismatch_type: description of what is being tested
3886	config_overrides: dict of config to override for the remote instance
3887	version_override: version dict e.g. {"vllm_version": "0.6.1"}
3888	should_fail: whether the handshake should fail
3889	enforce_handshake_compat: whether to enforce compatibility checking
3890
3891	"""
3892	local_vllm_config = create_vllm_config(
3893	model="facebook/opt-125m",
3894	block_size=16,
3895	kv_connector_extra_config={
3896	"enforce_handshake_compat": enforce_handshake_compat
3897	},
3898	)
3899	kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
3900	decode_connector = NixlConnector(
3901	local_vllm_config, KVConnectorRole.WORKER, kv_cache_config
3902	)
3903	decode_worker = decode_connector.connector_worker
3904	kv_cache_spec = cast(
3905	AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
3906	)
3907	shape = compute_layer_kv_cache_shape_bytes(
3908	kv_cache_spec, kv_cache_config.num_blocks
3909	)
3910	shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype)
3911	unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype)
3912	# Build kv_caches from the actual layer names in kv_cache_config so that
3913	# _layer_specs lookups in register_kv_caches always find a matching key.
3914	layer_names = [
3915	name for group in kv_cache_config.kv_cache_groups for name in group.layer_names
3916	]
3917	kv_caches = {
3918	name: shared_tensor if i % 2 == 0 else unique_tensor
3919	for i, name in enumerate(layer_names)
3920	}
3921	decode_connector.register_kv_caches(kv_caches)
3922
3923	remote_config_params: dict[str, Any] = {
3924	"model": "facebook/opt-125m",
3925	"block_size": 16,
3926	**config_overrides,
3927	}
3928	remote_vllm_config = create_vllm_config(**remote_config_params)
3929
3930	with contextlib.ExitStack() as stack:
3931	if "vllm_version" in version_override:
3932	stack.enter_context(
3933	patch("vllm.__version__", version_override["vllm_version"])
3934	)
3935	elif "connector_version" in version_override:
3936	stack.enter_context(
3937	patch.object(
3938	nixl.metadata,
3939	"NIXL_CONNECTOR_VERSION",
3940	version_override["connector_version"],
3941	)
3942	)
3943	remote_hash = compute_nixl_compatibility_hash(
3944	remote_vllm_config,
3945	decode_worker.backend_name,
3946	)
3947
3948	prefill_block_size = config_overrides.get("block_size", 16)
3949	prefill_block_lens = [4096 * prefill_block_size]
3950	prefill_metadata = NixlAgentMetadata(
3951	engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3952	agent_metadata=FakeNixlWrapper.AGENT_METADATA,
3953	kv_caches_base_addr=[0],
3954	device_id=0,
3955	num_blocks=1,
3956	block_lens=prefill_block_lens,
3957	block_strides=prefill_block_lens,
3958	kv_cache_layout="LBHNC",
3959	block_size=prefill_block_size,
3960	ssm_sizes=(0, 0),
3961	attn_backend_name=decode_worker.backend_name,
3962	physical_blocks_per_logical_kv_block=1,
3963	)
3964	handshake_payload = NixlHandshakePayload(
3965	compatibility_hash=remote_hash,
3966	agent_metadata_bytes=msgspec.msgpack.encode(prefill_metadata),
3967	)
3968
3969	# Mock ZMQ socket to return our handshake payload
3970	mock_socket = MagicMock()
3971	mock_socket.recv_multipart.return_value = [
3972	msgspec.msgpack.encode(handshake_payload),
3973	msgspec.msgpack.encode(time.perf_counter()),
3974	]
3975
3976	# Mock add_remote_agent to avoid actual NIXL operations
3977	# Patch zmq_ctx to return our mock socket
3978	with (
3979	patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"),
3980	patch.object(nixl.base_worker, "zmq_ctx") as mock_zmq_ctx,
3981	):
3982	mock_zmq_ctx.return_value.__enter__.return_value = mock_socket
3983
3984	if should_fail:
3985	with pytest.raises(RuntimeError, match="compatibility hash mismatch"):
3986	decode_worker._nixl_handshake(
3987	host="localhost",
3988	port=1234,
3989	remote_tp_size=1,
3990	expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3991	)
3992	else:
3993	result, _ = decode_worker._nixl_handshake(
3994	host="localhost",
3995	port=1234,
3996	remote_tp_size=1,
3997	expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
3998	)
3999	# Verify handshake returned agent mapping
4000	assert isinstance(result, dict)
4001	assert len(result) == 1
4002
4003
4004	@pytest.mark.parametrize(
4005	"error_scenario",
4006	[
4007	"handshake_decode_error",
4008	"handshake_validation_error",
4009	"metadata_decode_error",
4010	"metadata_validation_error",
4011	],
4012	)
4013	@patch(
4014	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
4015	FakeNixlWrapper,
4016	)
4017	def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario):
4018	"""Test that msgspec decode errors are properly handled during handshake.
4019
4020	Tests both DecodeError and ValidationError for both decoders:
4021	- NixlHandshakePayload decoder
4022	- NixlAgentMetadata decoder
4023	"""
4024	local_vllm_config = create_vllm_config(
4025	model="facebook/opt-125m",
4026	block_size=16,
4027	)
4028	decode_connector = NixlConnector(
4029	local_vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
4030	)
4031	decode_worker = decode_connector.connector_worker
4032
4033	backend = get_current_attn_backend(local_vllm_config)
4034	probe_spec = decode_worker.kv_cache_config.kv_cache_groups[0].kv_cache_spec
4035	test_shape = compute_layer_kv_cache_shape_bytes(probe_spec, 1)
4036	decode_worker.transfer_topo = TransferTopology(
4037	tp_rank=decode_worker.tp_rank,
4038	tp_size=decode_worker.world_size,
4039	block_size=decode_worker.block_size,
4040	engine_id=decode_worker.engine_id,
4041	is_mla=decode_worker.use_mla,
4042	is_mamba=False,
4043	total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
4044	attn_backends=[backend],
4045	tensor_shape=test_shape,
4046	)
4047
4048	decode_worker.compat_hash = compute_nixl_compatibility_hash(
4049	decode_worker.vllm_config,
4050	decode_worker.backend_name,
4051	)
4052
4053	if error_scenario == "handshake_decode_error":
4054	msg_bytes = b"this is not valid msgpack data"
4055	elif error_scenario == "handshake_validation_error":
4056	msg_bytes = msgspec.msgpack.encode({"wrong_field": "value"})
4057	elif error_scenario == "metadata_decode_error":
4058	valid_handshake = NixlHandshakePayload(
4059	compatibility_hash=decode_worker.compat_hash,
4060	agent_metadata_bytes=b"invalid msgpack for metadata",
4061	)
4062	msg_bytes = msgspec.msgpack.encode(valid_handshake)
4063
4064	elif error_scenario == "metadata_validation_error":
4065	valid_handshake = NixlHandshakePayload(
4066	compatibility_hash=decode_worker.compat_hash,
4067	agent_metadata_bytes=msgspec.msgpack.encode({"missing": "fields"}),
4068	)
4069	msg_bytes = msgspec.msgpack.encode(valid_handshake)
4070	else:
4071	raise AssertionError(f"{error_scenario} not a valid scenario")
4072
4073	mock_socket = MagicMock()
4074	mock_socket.recv_multipart.return_value = [
4075	msg_bytes,
4076	msgspec.msgpack.encode(time.perf_counter()),
4077	]
4078	with (
4079	patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"),
4080	patch.object(nixl.base_worker, "zmq_ctx") as mock_zmq_ctx,
4081	):
4082	mock_zmq_ctx.return_value.__enter__.return_value = mock_socket
4083
4084	with pytest.raises(RuntimeError):
4085	decode_worker._nixl_handshake(
4086	host="localhost",
4087	port=1234,
4088	remote_tp_size=1,
4089	expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
4090	)
4091
4092	@patch(
4093	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
4094	FakeNixlWrapper,
4095	)
4096	def test_mla_broadcast_notif_uses_remote_request_id(
4097	self, default_vllm_config, dist_init
4098	):
4099	"""MLA + remote TP > local TP: the broadcast notification sent to
4100	non-read prefill ranks must be keyed by the prefill-side request
4101	id (``meta.remote.request_id``), not the local decode request id.
4102
4103	Prefill ranks key ``_reqs_to_send`` by their own request id, so a
4104	broadcast keyed by the decode id is rejected in
4105	``_get_new_notifs`` with "Potentially invalid KV blocks for
4106	unrecognized request" and the blocks only release via the abort
4107	timeout. See ``_read_blocks_for_req`` in
4108	``vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py``.
4109	"""
4110	decode_tp_size = 1
4111	prefill_tp_size = 4
4112
4113	vllm_config = create_vllm_config()
4114	vllm_config.parallel_config.tensor_parallel_size = decode_tp_size
4115
4116	connector = NixlConnector(
4117	vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
4118	)
4119	connector.connector_worker = FakeNixlConnectorWorker(
4120	vllm_config, connector.engine_id, hand_shake_latency=0
4121	)
4122	worker = connector.connector_worker
4123
4124	# Force the MLA path; only `self.use_mla` gates the branches we
4125	# exercise inside `_read_blocks_for_req`.
4126	worker.use_mla = True
4127
4128	# Manually register the remote (P) engine and pre-populate the
4129	# per-rank state the handshake would normally fill in. The real
4130	# `_nixl_handshake` is unnecessary here — we only need
4131	# `transfer_topo` to know `remote_tp_size`, and `_remote_agents`
4132	# / `dst_xfer_side_handles` to be keyed by remote rank.
4133	remote_engine_id = "remote_engine"
4134	worker.transfer_topo.register_remote_engine(
4135	remote_engine_id=remote_engine_id,
4136	remote_tp_size=prefill_tp_size,
4137	remote_block_size=worker.block_size,
4138	remote_block_len=worker.block_size * 4096,
4139	remote_physical_blocks_per_logical=1,
4140	local_block_len=worker.block_size * 4096,
4141	)
4142	worker._remote_agents[remote_engine_id] = {
4143	(0, rank): f"agent_p{rank}" for rank in range(prefill_tp_size)
4144	}
4145	worker.dst_xfer_side_handles = {
4146	remote_engine_id: {rank: 100 + rank for rank in range(prefill_tp_size)}
4147	}
4148	# Sanity: D TP=1, P TP=4 => tp_ratio = -4 (P > D).
4149	assert worker.transfer_topo.tp_ratio(prefill_tp_size) == -prefill_tp_size
4150
4151	# Distinct ids on each side — that's the whole point of the bug.
4152	decode_req_id = "decode-req-AAAA"
4153	prefill_req_id = "prefill-req-BBBB"
4154	assert decode_req_id != prefill_req_id
4155
4156	metadata = NixlConnectorMetadata()
4157	metadata.add_new_req_to_recv(
4158	request_id=decode_req_id,
4159	local_block_ids=([0, 1, 2],),
4160	kv_transfer_params={
4161	"remote_block_ids": ([10, 11, 12],),
4162	"remote_engine_id": remote_engine_id,
4163	"remote_request_id": prefill_req_id,
4164	"remote_host": "localhost",
4165	"remote_port": 1234,
4166	"remote_tp_size": prefill_tp_size,
4167	},
4168	)
4169	meta = metadata.reqs_to_recv[decode_req_id]
4170
4171	# Capture broadcast send_notif calls; stub `_read_blocks` so we
4172	# don't need a working xfer path. Real `_read_blocks` emits its
4173	# auto-notif via `make_prepped_xfer`, not via `send_notif`, so
4174	# any captured `send_notif` here is a broadcast.
4175	send_notif_calls: list[tuple[str, bytes]] = []
4176	worker.nixl_wrapper.send_notif = ( # type: ignore[method-assign]
4177	lambda agent_name, notif_msg: send_notif_calls.append(
4178	(agent_name, notif_msg)
4179	)
4180	)
4181	worker._read_blocks = MagicMock() # type: ignore[method-assign]
4182
4183	worker._read_blocks_for_req(decode_req_id, meta)
4184
4185	# MLA: read once from rank 0 and broadcast to the other ranks.
4186	worker._read_blocks.assert_called_once()
4187	assert worker._read_blocks.call_args.kwargs["remote_rank"] == 0
4188	assert (
4189	worker._read_blocks.call_args.kwargs["remote_request_id"] == prefill_req_id
4190	)
4191
4192	# Broadcast goes to ranks {1, 2, 3} only, never to the read target.
4193	expected_recipients = {
4194	worker._remote_agents[remote_engine_id][(0, r)]
4195	for r in range(1, prefill_tp_size)
4196	}
4197	assert {agent for agent, _ in send_notif_calls} == expected_recipients
4198
4199	# Every broadcast notif must be keyed by the prefill request id.
4200	# Pre-fix this used the decode request id, which prefill ranks
4201	# didn't recognize.
4202	expected_notif = f"{prefill_req_id}:{decode_tp_size}".encode()
4203	bad_notif = f"{decode_req_id}:{decode_tp_size}".encode()
4204	for agent, notif in send_notif_calls:
4205	assert notif == expected_notif, (
4206	f"Broadcast notif to {agent!r} must use prefill_req_id; "
4207	f"got {notif!r} (expected {expected_notif!r}, "
4208	f"buggy form would be {bad_notif!r})"
4209	)
4210
4211
4212	@patch(
4213	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
4214	FakeNixlWrapper,
4215	)
4216	def test_kv_both_deprecation_warning(default_vllm_config, dist_init):
4217	"""kv_role='kv_both' should emit a deprecation log warning."""
4218	from vllm.logger import _print_warning_once
4219
4220	_print_warning_once.cache_clear()
4221
4222	vllm_config = create_vllm_config(kv_role="kv_both")
4223
4224	with patch(
4225	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector.logger"
4226	) as mock_logger:
4227	mock_logger.warning_once = mock_logger.warning_once
4228	NixlConnector(
4229	vllm_config,
4230	KVConnectorRole.WORKER,
4231	make_kv_cache_config(block_size=16),
4232	)
4233
4234	mock_logger.warning_once.assert_called_once()
4235	msg = mock_logger.warning_once.call_args[0][0]
4236	assert "kv_role='kv_both'" in msg
4237	assert "deprecated" in msg
4238
4239
4240	@patch(
4241	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
4242	FakeNixlWrapper,
4243	)
4244	def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init):
4245	"""kv_role='kv_consumer' or 'kv_producer' should NOT emit a warning."""
4246	for role in ("kv_consumer", "kv_producer"):
4247	vllm_config = create_vllm_config(kv_role=role)
4248	with patch(
4249	"vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector.logger"
4250	) as mock_logger:
4251	NixlConnector(
4252	vllm_config,
4253	KVConnectorRole.WORKER,
4254	make_kv_cache_config(block_size=16),
4255	)
4256
4257	(
4258	mock_logger.warning_once.assert_not_called(),
4259	(f"kv_role={role!r} should not emit deprecation warning"),
4260	)