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

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 2,536 lines 88.1 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4	import asyncio
5	import contextlib
6	import errno
7	import socket
8	import threading
9	import time
10	from types import SimpleNamespace
11	from unittest.mock import AsyncMock, MagicMock, patch
12
13	import httpx
14	import pytest
15	import torch
16	import zmq.asyncio
17
18	from tests.v1.attention.utils import dense_kv_cache_views
19	from vllm import envs
20	from vllm.config import KVTransferConfig, set_current_vllm_config
21	from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
22	mooncake_connector,
23	mooncake_utils,
24	)
25	from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
26	_SHARED_REGION_GROUP_ID,
27	KVConnectorRole,
28	MooncakeConnector,
29	MooncakeConnectorMetadata,
30	MooncakeConnectorWorker,
31	MooncakeXferMetadata,
32	MooncakeXferResponse,
33	MooncakeXferResponseStatus,
34	PullReqMeta,
35	SendBlockMeta,
36	TransferRegion,
37	_align_transfer_regions,
38	_block_ids_for_region,
39	_coalesce_contiguous_transfer_regions,
40	_compute_sender_transfer_plan,
41	_has_opaque_packed_row,
42	_pp_mismatch_hides_packed_layers,
43	_validate_asymmetric_region_lengths,
44	get_mooncake_bootstrap_addr,
45	should_launch_bootstrap_server,
46	)
47	from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
48	MooncakeBootstrapServer,
49	RegisterWorkerPayload,
50	)
51	from vllm.v1.kv_cache_interface import (
52	FullAttentionSpec,
53	KVCacheConfig,
54	KVCacheGroupSpec,
55	KVCacheLayout,
56	MLAAttentionSpec,
57	)
58	from vllm.v1.request import RequestStatus
59
60	from .utils import create_request, create_scheduler, create_vllm_config
61
62
63	@pytest.mark.parametrize(
64	("remote_tp_size", "expected_dst_offsets"),
65	[
66	(1, [0, None, 65536, None, 131072, None, 196608, None]),
67	(2, [0, None, 65536, None, 0, None, 65536, None]),
68	],
69	)
70	@pytest.mark.parametrize("local_tp_rank", range(8))
71	def test_sender_plan_replicated_gqa_tp8_to_smaller_tp(
72	remote_tp_size,
73	expected_dst_offsets,
74	local_tp_rank,
75	):
76	expected_dst_offset = expected_dst_offsets[local_tp_rank]
77	plan = _compute_sender_transfer_plan(
78	local_tp_rank=local_tp_rank,
79	local_tp_size=8,
80	remote_tp_rank=local_tp_rank // (8 // remote_tp_size),
81	remote_tp_size=remote_tp_size,
82	local_kv_block_len=65536,
83	remote_kv_block_len=262144 // remote_tp_size,
84	producer_cache_replicated=True,
85	total_num_kv_heads=4,
86	)
87	if expected_dst_offset is None:
88	assert not plan[0]
89	else:
90	assert plan == (True, 0, expected_dst_offset, 65536)
91
92
93	@pytest.mark.parametrize(
94	("local_tp_size", "expected_src_offsets"),
95	[
96	(1, [0, 0, 65536, 65536, 131072, 131072, 196608, 196608]),
97	(2, [0, 0, 65536, 65536, 0, 0, 65536, 65536]),
98	(4, [0] * 8),
99	],
100	)
101	@pytest.mark.parametrize("remote_tp_rank", range(8))
102	def test_sender_plan_gqa_to_replicated_tp8(
103	local_tp_size,
104	expected_src_offsets,
105	remote_tp_rank,
106	):
107	local_head_count = 4 // local_tp_size
108	plan = _compute_sender_transfer_plan(
109	local_tp_rank=remote_tp_rank // (8 // local_tp_size),
110	local_tp_size=local_tp_size,
111	remote_tp_rank=remote_tp_rank,
112	remote_tp_size=8,
113	local_kv_block_len=local_head_count * 65536,
114	remote_kv_block_len=65536,
115	producer_cache_replicated=False,
116	total_num_kv_heads=4,
117	)
118	assert plan == (True, expected_src_offsets[remote_tp_rank], 0, 65536)
119
120
121	@pytest.mark.parametrize(
122	"local_tp,remote_tp,local_len,remote_len,valid",
123	[
124	(1, 8, 262144, 65536, True),
125	(8, 1, 65536, 262144, True),
126	(1, 8, 262144, 131072, False),
127	(8, 1, 131072, 262144, False),
128	],
129	)
130	def test_region_length_validation_checks_replicated_gqa_heads(
131	local_tp, remote_tp, local_len, remote_len, valid
132	):
133	"""Replicated heads must have matching, whole per-head payloads."""
134	local_region = TransferRegion("layer", 0, 0, local_len, local_len)
135	remote_region = TransferRegion("layer", 0, 0, remote_len, remote_len)
136
137	assert (
138	_validate_asymmetric_region_lengths(
139	local_regions=[local_region],
140	remote_regions=[remote_region],
141	local_tp_size=local_tp,
142	remote_tp_size=remote_tp,
143	producer_cache_replicated=local_tp > 4,
144	total_num_kv_heads=4,
145	)
146	is None
147	) == valid
148
149
150	def _make_test_kv_cache_config() -> KVCacheConfig:
151	return KVCacheConfig(
152	num_blocks=0,
153	kv_cache_tensors=[],
154	kv_cache_groups=[
155	KVCacheGroupSpec(
156	[
157	"model.layers.0.self_attn",
158	"model.layers.1.self_attn",
159	"model.layers.0.mla_attn",
160	"model.layers.1.eagle_attn",
161	],
162	FullAttentionSpec(
163	block_size=16,
164	num_kv_heads=4,
165	head_size=64,
166	dtype=torch.float16,
167	),
168	)
169	],
170	)
171
172
173	def _make_packed_mla_kv_cache_config(
174	num_blocks: int, *, num_groups: int = 1
175	) -> KVCacheConfig:
176	spec = MLAAttentionSpec(
177	block_size=16,
178	num_kv_heads=1,
179	head_size=64,
180	dtype=torch.uint8,
181	)
182	if num_groups == 1:
183	groups = [
184	KVCacheGroupSpec(
185	["model.layers.0.mla_attn", "model.layers.1.indexer"],
186	spec,
187	)
188	]
189	else:
190	groups = [
191	KVCacheGroupSpec(["model.layers.0.mla_attn"], spec),
192	KVCacheGroupSpec(["model.layers.1.indexer"], spec),
193	]
194	return KVCacheConfig(
195	num_blocks=num_blocks,
196	kv_cache_tensors=[],
197	kv_cache_groups=groups,
198	)
199
200
201	class FakeMooncakeWrapper:
202	"""Mock Mooncake TransferEngine for unit testing environments."""
203
204	def __init__(self, args, *kwargs):
205	pass
206
207	def initialize(self, local_hostname, metadata_server, protocol, device_name) -> int:
208	return 0
209
210	def get_rpc_port(self) -> int:
211	return 12345
212
213	def batch_transfer_sync_write(
214	self, target_hostname, buffers, peer_buffer_addresses, lengths
215	) -> int:
216	return 0
217
218	def batch_register_memory(self, buffer_addresses, capacities) -> int:
219	return 0
220
221
222	def test_align_transfer_regions_uses_layer_name_occurrences():
223	"""Repeated layer names should align by occurrence order."""
224	local_regions = [
225	TransferRegion(
226	layer_name="model.layers.1.self_attn",
227	layer_index=1,
228	base_addr=0x1000,
229	block_len=256,
230	kv_block_len=128,
231	),
232	TransferRegion(
233	layer_name="model.layers.1.self_attn",
234	layer_index=1,
235	base_addr=0x1100,
236	block_len=256,
237	kv_block_len=128,
238	),
239	]
240	remote_regions = [
241	TransferRegion(
242	layer_name="model.layers.0.self_attn",
243	layer_index=0,
244	base_addr=0xA000,
245	block_len=256,
246	kv_block_len=128,
247	),
248	TransferRegion(
249	layer_name="model.layers.1.self_attn",
250	layer_index=1,
251	base_addr=0xB000,
252	block_len=256,
253	kv_block_len=128,
254	),
255	TransferRegion(
256	layer_name="model.layers.1.self_attn",
257	layer_index=1,
258	base_addr=0xB100,
259	block_len=256,
260	kv_block_len=128,
261	),
262	]
263
264	aligned_local, aligned_remote, err = _align_transfer_regions(
265	local_regions, remote_regions
266	)
267
268	assert err is None
269	assert [r.base_addr for r in aligned_local] == [0x1000, 0x1100]
270	assert [r.base_addr for r in aligned_remote] == [0xB000, 0xB100]
271
272	# Hetero PP with no shared layers is an empty success. Same PP fails closed.
273	p_regions = [
274	_region(
275	0x1000,
276	block_len=2048,
277	kv_block_len=1024,
278	row_offset=0,
279	layer_name="model.layers.2.mla_attn",
280	layer_index=2,
281	)
282	]
283	early_d = [
284	_region(
285	0xA000,
286	block_len=2048,
287	kv_block_len=1024,
288	row_offset=0,
289	layer_name="model.layers.0.mla_attn",
290	layer_index=0,
291	)
292	]
293	none_p, none_d, none_err = _align_transfer_regions(
294	p_regions, early_d, allow_partial_layers=True
295	)
296	assert none_err is None
297	assert none_p == [] and none_d == []
298	_, _, mismatch_err = _align_transfer_regions(
299	p_regions, early_d, allow_partial_layers=False
300	)
301	assert mismatch_err is not None
302
303
304	@pytest.mark.asyncio
305	async def test_build_transfer_params_separates_prefill_pp_layers():
306	"""Each producer PP stage should send only its registered layer shard."""
307	worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker)
308	worker.async_zmq_ctx = MagicMock()
309	worker.is_kv_consumer = True
310	worker.is_kv_producer = True
311	worker.tp_rank = 0
312	worker.tp_size = 1
313	worker.use_mla = False
314	worker.kv_cache_config = _make_test_kv_cache_config()
315	worker._physical_blocks_per_logical_kv_block = 1
316	worker.transfer_topo = SimpleNamespace(
317	local_replicates_kv_cache=False,
318	total_num_kv_heads=4,
319	)
320
321	block_len = 256
322	remote_regions = [
323	TransferRegion(
324	layer_name=f"model.layers.{layer_index}.self_attn",
325	layer_index=layer_index,
326	base_addr=base_addr,
327	block_len=block_len,
328	kv_block_len=block_len,
329	)
330	for layer_index, base_addr in [
331	(0, 0xA000),
332	(1, 0xB000),
333	(2, 0xC000),
334	(3, 0xD000),
335	]
336	]
337	producer_pp_regions = {
338	0: [
339	TransferRegion(
340	layer_name="model.layers.0.self_attn",
341	layer_index=0,
342	base_addr=0x1000,
343	block_len=block_len,
344	kv_block_len=block_len,
345	),
346	TransferRegion(
347	layer_name="model.layers.1.self_attn",
348	layer_index=1,
349	base_addr=0x2000,
350	block_len=block_len,
351	kv_block_len=block_len,
352	),
353	],
354	1: [
355	TransferRegion(
356	layer_name="model.layers.2.self_attn",
357	layer_index=2,
358	base_addr=0x3000,
359	block_len=block_len,
360	kv_block_len=block_len,
361	),
362	TransferRegion(
363	layer_name="model.layers.3.self_attn",
364	layer_index=3,
365	base_addr=0x4000,
366	block_len=block_len,
367	kv_block_len=block_len,
368	),
369	],
370	}
371	expected_by_pp_rank = {
372	0: {
373	"layers": [0, 1],
374	"src_ptrs": [0x1000 + 10 * block_len, 0x2000 + 10 * block_len],
375	"dst_ptrs": [0xA000 + 20 * block_len, 0xB000 + 20 * block_len],
376	},
377	1: {
378	"layers": [2, 3],
379	"src_ptrs": [0x3000 + 10 * block_len, 0x4000 + 10 * block_len],
380	"dst_ptrs": [0xC000 + 20 * block_len, 0xD000 + 20 * block_len],
381	},
382	}
383
384	transfer_id = "xfer-pp-split"
385	send_meta = SendBlockMeta(
386	p_req_id="p-req-pp",
387	transfer_id=transfer_id,
388	local_block_ids=[[10, 11]],
389	ready=asyncio.Event(),
390	)
391	xfer_meta = _xfer_meta(remote_regions, {"d-req-pp": (transfer_id, [[20, 21]])})
392
393	for pp_rank, local_regions in producer_pp_regions.items():
394	aligned_local, aligned_remote, err = _align_transfer_regions(
395	local_regions, remote_regions
396	)
397
398	assert err is None
399	assert [r.layer_index for r in aligned_local] == (
400	expected_by_pp_rank[pp_rank]["layers"]
401	)
402	assert [r.layer_index for r in aligned_remote] == (
403	expected_by_pp_rank[pp_rank]["layers"]
404	)
405
406	(
407	src_ptrs,
408	dst_ptrs,
409	lengths,
410	err_reqs,
411	err_msg,
412	) = await worker._build_transfer_params(
413	ready_reqs=[("d-req-pp", send_meta)],
414	agent_meta=xfer_meta,
415	local_regions=aligned_local,
416	remote_regions=aligned_remote,
417	)
418
419	assert err_reqs == []
420	assert err_msg is None
421	assert src_ptrs == expected_by_pp_rank[pp_rank]["src_ptrs"]
422	assert dst_ptrs == expected_by_pp_rank[pp_rank]["dst_ptrs"]
423	assert lengths == [2 * block_len, 2 * block_len]
424
425
426	@pytest.mark.asyncio
427	async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
428	monkeypatch,
429	):
430	"""Producer sends its PP layer shard to the matching consumer layer address."""
431	monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
432	vllm_config = create_vllm_config(
433	kv_connector="MooncakeConnector", kv_role="kv_producer"
434	)
435
436	with set_current_vllm_config(vllm_config), patch_worker_dependencies():
437	prefill_connector = MooncakeConnector(
438	vllm_config,
439	KVConnectorRole.WORKER,
440	_make_test_kv_cache_config(),
441	)
442	prefill_worker = prefill_connector.connector_worker
443
444	block_len = 4096
445	prefill_worker.kv_caches_base_addr = [0x1000]
446	prefill_worker.block_len_per_layer = [block_len]
447	prefill_worker.kv_block_len_per_layer = [block_len]
448	prefill_worker.registered_layer_names = ["model.layers.1.self_attn"]
449	prefill_worker.registered_layer_indices = [1]
450
451	class InlineSenderLoop:
452	async def run_in_executor(self, executor, func, *args):
453	return func(*args)
454
455	origin_sender_loop = prefill_worker.sender_loop
456	prefill_worker.sender_loop = InlineSenderLoop()
457
458	transfer_id = "xfer-layer-align"
459	send_meta = SendBlockMeta(
460	p_req_id="p-req-layer-align",
461	transfer_id=transfer_id,
462	local_block_ids=[[10]],
463	ready=asyncio.Event(),
464	)
465	prefill_worker.reqs_need_send[transfer_id] = send_meta
466	send_meta.ready.set()
467
468	xfer_meta = _xfer_meta(
469	[
470	_region(
471	0xA000,
472	block_len=block_len,
473	row_offset=-1,
474	layer_name="model.layers.0.self_attn",
475	layer_index=0,
476	),
477	_region(
478	0xB000,
479	block_len=block_len,
480	row_offset=-1,
481	layer_name="model.layers.1.self_attn",
482	layer_index=1,
483	),
484	],
485	{"d-req-layer-align": (transfer_id, [[20]])},
486	)
487	mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
488	mock_socket.send_multipart = AsyncMock()
489	identity = b"consumer-layer-align"
490
491	with patch.object(
492	prefill_worker, "_send_blocks", return_value=0
493	) as mock_send_blocks:
494	await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
495
496	src_ptrs, dst_ptrs, lengths = mock_send_blocks.call_args[0][1:]
497	assert src_ptrs == [0x1000 + 10 * block_len]
498	assert dst_ptrs == [0xB000 + 20 * block_len]
499	assert lengths == [block_len]
500
501	sent_identity, sent_payload = mock_socket.send_multipart.call_args[0][0]
502	assert sent_identity == identity
503	response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
504	assert response.status == MooncakeXferResponseStatus.FINISH
505	assert response.ok_reqs == ["d-req-layer-align"]
506
507	prefill_worker.sender_loop = origin_sender_loop
508	prefill_worker.shutdown()
509
510
511	def test_basic_interface():
512	"""Unit test for basic MooncakeConnector interface functionality."""
513	vllm_config = create_vllm_config(
514	kv_connector="MooncakeConnector", kv_role="kv_consumer"
515	)
516	scheduler = create_scheduler(vllm_config)
517
518	# 2 Full Blocks and 1 Half Block.
519	BLOCK_SIZE = vllm_config.cache_config.block_size
520	NUM_EXTERNAL_FULL_BLOCKS = 2
521	NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
522
523	request = create_request(
524	request_id=1,
525	block_size=BLOCK_SIZE,
526	num_tokens=NUM_TOKENS,
527	do_remote_prefill=True,
528	)
529	request_id = request.request_id
530	request.kv_transfer_params.update(
531	{
532	"transfer_id": request_id,
533	"remote_bootstrap_addr": 54321,
534	}
535	)
536
537	scheduler.add_request(request)
538
539	# Remote Prefill, triggers NixlConnectorMetadata.
540	scheduler_output = scheduler.schedule()
541	kv_connector_metadata = scheduler_output.kv_connector_metadata
542	assert kv_connector_metadata is not None
543	assert isinstance(kv_connector_metadata, MooncakeConnectorMetadata)
544
545	assert len(kv_connector_metadata.reqs_to_recv) == 1
546	assert request_id in kv_connector_metadata.reqs_to_recv["my-engine-id"]
547	req_meta = kv_connector_metadata.reqs_to_recv["my-engine-id"][request_id]
548
549	# local_block_ids is list[list[int]] (per-group); flatten for comparison.
550	all_block_ids = [bid for group in req_meta.local_block_ids for bid in group]
551	for block_id, block in zip(
552	all_block_ids,
553	scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
554	request_id
555	],
556	):
557	assert block_id == block.block_id
558
559
560	def test_prompt_less_than_block_size():
561	"""Test that we can handle case where prompt is < block."""
562	vllm_config = create_vllm_config(
563	kv_connector="MooncakeConnector", kv_role="kv_consumer"
564	)
565	scheduler = create_scheduler(vllm_config)
566
567	# Half of a block.
568	BLOCK_SIZE = vllm_config.cache_config.block_size
569	NUM_TOKENS = int(BLOCK_SIZE * 0.5)
570
571	# Request will have 1 partial remote block.
572	request = create_request(
573	request_id=1,
574	block_size=BLOCK_SIZE,
575	num_tokens=NUM_TOKENS,
576	do_remote_prefill=True,
577	num_remote_blocks=1,
578	)
579	request.kv_transfer_params.update(
580	{
581	"transfer_id": request.request_id,
582	"remote_bootstrap_addr": 54321,
583	}
584	)
585
586	scheduler.add_request(request)
587	scheduler_output = scheduler.schedule()
588
589	# This request will read async.
590	kv_connector_metadata = scheduler_output.kv_connector_metadata
591	assert kv_connector_metadata is not None
592	assert isinstance(kv_connector_metadata, MooncakeConnectorMetadata)
593	assert len(kv_connector_metadata.reqs_to_recv["my-engine-id"]) == 1
594	assert len(scheduler_output.scheduled_new_reqs) == 0
595
596
597	@pytest.fixture
598	def bootstrap_server():
599	"""Fixture to launch and cleanup a Mooncake Bootstrap HTTP Server."""
600	server = MooncakeBootstrapServer("127.0.0.1", 0)
601	try:
602	server.start()
603	yield server
604	finally:
605	server.shutdown()
606
607
608	def test_bootstrap_server_occupied_port_raises_to_caller():
609	with socket.create_server(("127.0.0.1", 0)) as listener:
610	server = MooncakeBootstrapServer("127.0.0.1", listener.getsockname()[1])
611	try:
612	with pytest.raises(OSError) as exc_info:
613	server.start(timeout=0.05)
614	assert exc_info.value.errno == errno.EADDRINUSE
615	finally:
616	server.shutdown()
617
618
619	def test_bootstrap_server_retains_and_releases_assigned_port(bootstrap_server):
620	port = bootstrap_server.port
621	assert port > 0
622	with (
623	pytest.raises(OSError) as exc_info,
624	socket.create_server(("127.0.0.1", port)),
625	):
626	pass
627	assert exc_info.value.errno == errno.EADDRINUSE
628
629	bootstrap_server.shutdown()
630	with socket.create_server(("127.0.0.1", port)):
631	pass
632
633
634	@pytest.mark.parametrize("failure", ["exit", "error", "timeout"])
635	def test_bootstrap_server_startup_failure_releases_port(monkeypatch, failure):
636	stop = threading.Event()
637
638	def run_without_starting(server, *, sockets):
639	if failure == "error":
640	raise SystemExit(1)
641	if failure == "timeout":
642	while not server.should_exit and not stop.wait(0.01):
643	pass
644
645	monkeypatch.setattr(mooncake_utils.uvicorn.Server, "run", run_without_starting)
646	server = MooncakeBootstrapServer("127.0.0.1", 0)
647	try:
648	error = TimeoutError if failure == "timeout" else RuntimeError
649	with pytest.raises(error):
650	server.start(timeout=0.05 if failure == "timeout" else 1.0)
651	assert server.port > 0
652	with socket.create_server(("127.0.0.1", server.port)):
653	pass
654	finally:
655	stop.set()
656	server.shutdown()
657
658
659	@pytest.mark.parametrize("state", ["stopping", "exited"])
660	def test_bootstrap_server_rejects_stale_started_state(monkeypatch, state):
661	release = threading.Event()
662
663	def run_until_released(server, *, sockets):
664	server.started = True
665	release.wait()
666
667	monkeypatch.setattr(mooncake_utils.uvicorn.Server, "run", run_until_released)
668	server = MooncakeBootstrapServer("127.0.0.1", 0)
669	try:
670	server.start()
671	thread = server.server_thread
672	assert thread is not None
673	if state == "stopping":
674	with patch.object(thread, "join", return_value=None):
675	server.shutdown()
676	else:
677	release.set()
678	thread.join(timeout=5)
679	assert not thread.is_alive()
680
681	with pytest.raises(RuntimeError, match="still starting or stopping"):
682	server.start()
683	finally:
684	release.set()
685	server.shutdown()
686
687
688	@pytest.mark.asyncio
689	@pytest.mark.parametrize("external", [False, True])
690	async def test_register_worker_recovers_from_slow_bootstrap_server(
691	monkeypatch, external
692	):
693	"""End-to-end: a real HTTP server that responds slower than the configured
694	timeout must be retried on the wire, not treated as fatal."""
695	original = MooncakeBootstrapServer.register_worker
696	state = {"done": False, "calls": 0}
697
698	async def slow_register(self, payload: RegisterWorkerPayload):
699	state["calls"] += 1
700	if not state["done"]:
701	state["done"] = True
702	# Record the worker first, then stall past the client timeout, so
703	# the retry exercises the duplicate-registration path.
704	response = await original(self, payload)
705	await asyncio.sleep(1.5)
706	return response
707	return await original(self, payload)
708
709	monkeypatch.setattr(MooncakeBootstrapServer, "register_worker", slow_register)
710
711	monkeypatch.setenv("VLLM_MOONCAKE_CONNECTOR_TIMEOUT", "0.5")
712	monkeypatch.setattr(mooncake_connector, "_BOOTSTRAP_MAX_ATTEMPTS", 5)
713
714	server = MooncakeBootstrapServer("127.0.0.1", 0)
715	server.start()
716	try:
717	monkeypatch.setenv("VLLM_MOONCAKE_BOOTSTRAP_PORT", str(server.port))
718	worker = _make_local_register_worker_stub()
719	if external:
720	worker.vllm_config = _make_bootstrap_vllm_config(
721	bootstrap_server_address=f"127.0.0.1:{server.port}"
722	)
723	assert not should_launch_bootstrap_server(worker.vllm_config)
724	await MooncakeConnectorWorker.register_worker_with_bootstrap(worker)
725	assert state["calls"] == 2
726
727	async with httpx.AsyncClient() as client:
728	response = await client.get(f"http://127.0.0.1:{server.port}/query")
729	assert response.status_code == 200
730	assert response.json()["0"]["engine_id"] == "eng-1"
731	finally:
732	server.shutdown()
733
734
735	@pytest.mark.asyncio
736	async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):
737	"""Tests the bootstrap server's api for worker registration and querying.
738
739	Validates DP/TP/PP rank indexing and error handling for duplicate registrations.
740	"""
741	import httpx
742
743	base_url = f"http://127.0.0.1:{bootstrap_server.port}"
744
745	# Query when empty
746	async with httpx.AsyncClient() as client:
747	response = await client.get(f"{base_url}/query")
748	assert response.status_code == 200
749	assert response.json() == {}
750
751	# Register multiple PP workers from the same producer engine.
752	payload1 = {
753	"engine_id": "eng-1",
754	"dp_rank": 0,
755	"tp_rank": 0,
756	"pp_rank": 0,
757	"addr": "tcp://1.1.1.1:1111",
758	}
759	async with httpx.AsyncClient() as client:
760	response = await client.post(f"{base_url}/register", json=payload1)
761	assert response.status_code == 200
762	assert response.json() == {"status": "ok"}
763
764	payload2 = {
765	"engine_id": "eng-1",
766	"dp_rank": 0,
767	"tp_rank": 0,
768	"pp_rank": 1,
769	"addr": "tcp://2.2.2.2:2222",
770	}
771	async with httpx.AsyncClient() as client:
772	response = await client.post(f"{base_url}/register", json=payload2)
773	assert response.status_code == 200
774	assert response.json() == {"status": "ok"}
775
776	# Query after registration should preserve the PP dimension.
777	async with httpx.AsyncClient() as client:
778	response = await client.get(f"{base_url}/query")
779	assert response.status_code == 200
780	data = response.json()
781	assert "0" in data
782	assert data["0"]["engine_id"] == "eng-1"
783	assert data["0"]["worker_addr"]["0"]["0"] == "tcp://1.1.1.1:1111"
784	assert data["0"]["worker_addr"]["0"]["1"] == "tcp://2.2.2.2:2222"
785
786	# Re-registering the identical payload is idempotent.
787	async with httpx.AsyncClient() as client:
788	response = await client.post(f"{base_url}/register", json=payload1)
789	assert response.status_code == 200
790	assert response.json() == {"status": "ok"}
791
792	# Test failure: same ranks, conflicting address
793	payload1_conflict = dict(payload1, addr="tcp://9.9.9.9:9999")
794	async with httpx.AsyncClient() as client:
795	response = await client.post(f"{base_url}/register", json=payload1_conflict)
796	assert response.status_code == 400
797	assert "is already registered" in response.text
798
799	# Test failure: engine_id mismatch for same dp_rank
800	payload3_fail = {
801	"engine_id": "eng-2",
802	"dp_rank": 0,
803	"tp_rank": 1,
804	"pp_rank": 0,
805	"addr": "tcp://3.3.3.3:3333",
806	}
807	async with httpx.AsyncClient() as client:
808	response = await client.post(f"{base_url}/register", json=payload3_fail)
809	assert response.status_code == 400
810	assert "Engine ID mismatch" in response.text
811
812
813	def _make_bootstrap_vllm_config(
814	*,
815	local_engines_only: bool = False,
816	data_parallel_rank_local: int = 0,
817	data_parallel_index: int = 0,
818	nnodes_within_dp: int = 1,
819	bootstrap_server_address: str \| None = None,
820	) -> SimpleNamespace:
821	return SimpleNamespace(
822	kv_transfer_config=KVTransferConfig(
823	kv_connector="MooncakeConnector",
824	kv_role="kv_producer",
825	kv_connector_extra_config={
826	"bootstrap_server_address": bootstrap_server_address
827	},
828	),
829	parallel_config=SimpleNamespace(
830	local_engines_only=local_engines_only,
831	data_parallel_rank_local=data_parallel_rank_local,
832	data_parallel_index=data_parallel_index,
833	nnodes_within_dp=nnodes_within_dp,
834	master_addr="model-parallel-master",
835	data_parallel_master_ip="data-parallel-master",
836	),
837	)
838
839
840	def _make_local_register_worker_stub():
841	return SimpleNamespace(
842	vllm_config=SimpleNamespace(
843	kv_transfer_config=None,
844	parallel_config=SimpleNamespace(
845	local_engines_only=False,
846	data_parallel_rank_local=0,
847	data_parallel_index=0,
848	nnodes_within_dp=1,
849	master_addr="127.0.0.1",
850	data_parallel_master_ip="127.0.0.1",
851	),
852	),
853	hostname="127.0.0.1",
854	side_channel_port=1234,
855	engine_id="eng-1",
856	dp_rank=0,
857	tp_rank=0,
858	pp_rank=0,
859	)
860
861
862	def _make_register_worker_stub():
863	return SimpleNamespace(
864	vllm_config=_make_bootstrap_vllm_config(),
865	hostname="127.0.0.1",
866	side_channel_port=1234,
867	engine_id="eng-1",
868	dp_rank=0,
869	tp_rank=0,
870	pp_rank=0,
871	)
872
873
874	class _FlakyAsyncClient:
875	"""httpx.AsyncClient stub that read-times-out `failures` times first."""
876
877	def __init__(self, failures: int, calls: list[int], **kwargs):
878	self.failures = failures
879	self.calls = calls
880
881	async def __aenter__(self):
882	return self
883
884	async def __aexit__(self, *exc_info):
885	return False
886
887	async def post(self, url, json=None):
888	self.calls.append(1)
889	if len(self.calls) <= self.failures:
890	raise httpx.ReadTimeout("simulated slow bootstrap server")
891	response = MagicMock()
892	response.raise_for_status = MagicMock()
893	return response
894
895
896	@pytest.mark.asyncio
897	async def test_register_worker_retries_on_read_timeout(monkeypatch):
898	"""A slow rank-0 bootstrap response must be retried, not treated as fatal."""
899	monkeypatch.setattr(mooncake_connector, "_BOOTSTRAP_MAX_ATTEMPTS", 5)
900	monkeypatch.setattr(asyncio, "sleep", AsyncMock())
901
902	calls: list[int] = []
903	monkeypatch.setattr(
904	httpx,
905	"AsyncClient",
906	lambda **kwargs: _FlakyAsyncClient(2, calls),
907	)
908
909	worker = _make_register_worker_stub()
910	await MooncakeConnectorWorker.register_worker_with_bootstrap(worker)
911
912	assert len(calls) == 3
913
914
915	@pytest.mark.asyncio
916	async def test_register_worker_raises_after_max_attempts(monkeypatch):
917	"""Terminal registration failure must raise, not loop forever."""
918	monkeypatch.setattr(mooncake_connector, "_BOOTSTRAP_MAX_ATTEMPTS", 3)
919	monkeypatch.setattr(asyncio, "sleep", AsyncMock())
920
921	calls: list[int] = []
922	monkeypatch.setattr(
923	httpx,
924	"AsyncClient",
925	lambda **kwargs: _FlakyAsyncClient(99, calls),
926	)
927
928	worker = _make_register_worker_stub()
929	with pytest.raises(RuntimeError, match="after 3 attempts"):
930	await MooncakeConnectorWorker.register_worker_with_bootstrap(worker)
931
932	assert len(calls) == 3
933
934
935	@pytest.mark.asyncio
936	async def test_sender_listener_failure_propagates_to_caller(monkeypatch):
937	"""A terminal registration failure must surface via the Future rather than
938	leaving register_kv_caches blocked on its ready event."""
939	monkeypatch.setenv("VLLM_MOONCAKE_CONNECTOR_TIMEOUT", "0.1")
940	monkeypatch.setattr(mooncake_connector, "_BOOTSTRAP_MAX_ATTEMPTS", 1)
941
942	loop = asyncio.new_event_loop()
943	thread = threading.Thread(target=loop.run_forever, daemon=True)
944	thread.start()
945
946	async def failing_listener(ready_event):
947	raise RuntimeError("simulated terminal registration failure")
948
949	try:
950	ready_event = threading.Event()
951	fut = asyncio.run_coroutine_threadsafe(failing_listener(ready_event), loop)
952
953	deadline = time.monotonic() + 10.0
954	raised = None
955	while not ready_event.wait(timeout=0.1):
956	if fut.done():
957	try:
958	fut.result()
959	except RuntimeError as e:
960	raised = e
961	break
962	assert time.monotonic() < deadline, "caller blocked past its deadline"
963
964	assert raised is not None
965	assert "simulated terminal registration failure" in str(raised)
966	finally:
967	loop.call_soon_threadsafe(loop.stop)
968	thread.join(timeout=5)
969	loop.close()
970
971
972	@pytest.mark.parametrize(
973	(
974	"tp_rank",
975	"pp_rank",
976	"local_engines_only",
977	"data_parallel_rank_local",
978	"data_parallel_index",
979	"expected",
980	),
981	[
982	(1, 0, False, 0, 0, False),
983	(0, 1, False, 0, 0, False),
984	(0, 0, True, 0, 1, True),
985	(0, 0, True, 1, 0, False),
986	(0, 0, False, 0, 0, True),
987	(0, 0, False, 0, 1, False),
988	],
989	ids=[
990	"nonzero_tp_rank",
991	"nonzero_pp_rank",
992	"local_engine_rank_zero",
993	"local_engine_nonzero_rank",
994	"internal_lb_first_dp_engine",
995	"internal_lb_nonzero_dp_engine",
996	],
997	)
998	def test_should_launch_bootstrap_server_selects_single_owner(
999	tp_rank: int,
1000	pp_rank: int,
1001	local_engines_only: bool,
1002	data_parallel_rank_local: int,
1003	data_parallel_index: int,
1004	expected: bool,
1005	):
1006	vllm_config = _make_bootstrap_vllm_config(
1007	local_engines_only=local_engines_only,
1008	data_parallel_rank_local=data_parallel_rank_local,
1009	data_parallel_index=data_parallel_index,
1010	)
1011	with (
1012	patch(
1013	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
1014	"mooncake_connector.get_tensor_model_parallel_rank",
1015	return_value=tp_rank,
1016	),
1017	patch(
1018	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
1019	"mooncake_connector.get_pp_group"
1020	) as mock_pp_group,
1021	):
1022	mock_pp_group.return_value.rank_in_group = pp_rank
1023	assert should_launch_bootstrap_server(vllm_config) is expected
1024
1025
1026	@pytest.mark.parametrize(
1027	("local_engines_only", "nnodes_within_dp", "expected_host"),
1028	[
1029	(True, 2, "127.0.0.1"),
1030	(False, 2, "model-parallel-master"),
1031	(False, 1, "data-parallel-master"),
1032	],
1033	ids=["local_engine", "multi_node_tp_or_pp", "single_node_internal_lb"],
1034	)
1035	def test_get_mooncake_bootstrap_addr_selects_expected_host(
1036	local_engines_only: bool,
1037	nnodes_within_dp: int,
1038	expected_host: str,
1039	):
1040	vllm_config = _make_bootstrap_vllm_config(
1041	local_engines_only=local_engines_only,
1042	nnodes_within_dp=nnodes_within_dp,
1043	)
1044
1045	assert get_mooncake_bootstrap_addr(vllm_config) == (
1046	expected_host,
1047	envs.VLLM_MOONCAKE_BOOTSTRAP_PORT,
1048	)
1049
1050
1051	@pytest.mark.parametrize(
1052	("address", "expected"),
1053	[("external:12345", ("external", 12345)), ("[::1]:12345", ("::1", 12345))],
1054	)
1055	def test_get_mooncake_bootstrap_addr_uses_external_endpoint(address, expected):
1056	config = _make_bootstrap_vllm_config(bootstrap_server_address=address)
1057	assert get_mooncake_bootstrap_addr(config) == expected
1058
1059
1060	@pytest.mark.parametrize("address", [":12345", "host:0", "host:65536"])
1061	def test_get_mooncake_bootstrap_addr_rejects_invalid_external_address(address):
1062	config = _make_bootstrap_vllm_config(bootstrap_server_address=address)
1063	with pytest.raises(ValueError, match="bootstrap_server_address"):
1064	get_mooncake_bootstrap_addr(config)
1065
1066
1067	def test_scheduler_request_finished():
1068	"""Tests the scheduler-side logic when a request finishes.
1069
1070	Differentiates between 'Finished' (requires transfer)
1071	and 'Aborted' (immediate free).
1072	"""
1073	vllm_config = create_vllm_config(
1074	kv_connector="MooncakeConnector", kv_role="kv_producer"
1075	)
1076	scheduler = create_scheduler(vllm_config)
1077	scheduler_connector = scheduler.get_kv_connector().connector_scheduler
1078
1079	request = create_request(request_id=1, do_remote_decode=True)
1080	request.kv_transfer_params["transfer_id"] = request.request_id
1081
1082	# Case: Capped length (Successful prefill, need to send to decoder)
1083	request.status = RequestStatus.FINISHED_LENGTH_CAPPED
1084	delay_free, _ = scheduler_connector.request_finished(request, block_ids=([10, 11],))
1085	assert delay_free is True
1086	assert "id-1" in scheduler_connector._reqs_need_send
1087	assert scheduler_connector._reqs_need_send["id-1"][1] == [[10, 11]]
1088
1089	# Case: Aborted (No need to transfer, free blocks immediately)
1090	scheduler_connector._reqs_need_send.clear()
1091	request.status = RequestStatus.FINISHED_ABORTED
1092	delay_free, _ = scheduler_connector.request_finished(request, block_ids=([12],))
1093	assert delay_free is False
1094	assert len(scheduler_connector._reqs_need_send) == 0
1095	assert "id-1" in scheduler_connector._reqs_not_processed
1096
1097
1098	@contextlib.contextmanager
1099	def patch_worker_dependencies():
1100	"""Helper to mock all distributed and network dependencies for Worker tests."""
1101	with (
1102	patch(
1103	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
1104	FakeMooncakeWrapper,
1105	),
1106	patch(
1107	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_ip",
1108	return_value="127.0.0.1",
1109	),
1110	patch(
1111	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_tensor_model_parallel_rank",
1112	return_value=0,
1113	),
1114	patch(
1115	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_tensor_model_parallel_world_size",
1116	return_value=1,
1117	),
1118	patch(
1119	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_pp_group"
1120	) as mock_pp,
1121	patch(
1122	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.should_launch_bootstrap_server",
1123	return_value=False,
1124	),
1125	patch(
1126	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.make_zmq_socket"
1127	) as mock_make_zmq,
1128	patch("httpx.AsyncClient") as mock_async_client,
1129	):
1130	# Mock PP group
1131	mock_pp_group = MagicMock()
1132	mock_pp_group.rank_in_group = 0
1133	mock_pp.return_value = mock_pp_group
1134
1135	# Mock ZMQ socket
1136	mock_socket_object = AsyncMock()
1137	mock_socket_object.setsockopt = MagicMock()
1138	mock_socket_ctx = MagicMock()
1139	mock_socket_ctx.__enter__.return_value = mock_socket_object
1140	mock_make_zmq.return_value = mock_socket_ctx
1141
1142	# Mock httpx client
1143	mock_http_client_instance = AsyncMock()
1144	mock_async_client.return_value = mock_http_client_instance
1145
1146	yield {
1147	"mock_make_zmq": mock_make_zmq,
1148	"mock_socket_object": mock_socket_object,
1149	"mock_async_client": mock_async_client,
1150	"mock_http_client": mock_http_client_instance,
1151	}
1152
1153
1154	def _region(
1155	base: int,
1156	*,
1157	block_len: int = 256,
1158	kv_block_len: int \| None = None,
1159	row_offset: int \| None = None,
1160	layer_name: str = "layer",
1161	layer_index: int = 0,
1162	group_index: int = 0,
1163	shared_group_ids: tuple[int, ...] = (),
1164	) -> TransferRegion:
1165	"""One transfer slice. `row_offset` defaults to `base` for in-row addresses."""
1166	return TransferRegion(
1167	layer_name=layer_name,
1168	layer_index=layer_index,
1169	base_addr=base,
1170	block_len=block_len,
1171	kv_block_len=block_len if kv_block_len is None else kv_block_len,
1172	group_index=group_index,
1173	shared_group_ids=shared_group_ids,
1174	row_offset=base if row_offset is None else row_offset,
1175	)
1176
1177
1178	def _xfer_meta(
1179	regions: list[TransferRegion],
1180	req_blocks: dict,
1181	**overrides,
1182	) -> MooncakeXferMetadata:
1183	"""Handshake metadata derived from the regions under test."""
1184	fields = dict(
1185	remote_hostname="consumer-host",
1186	remote_port=54321,
1187	remote_tp_size=1,
1188	remote_tp_rank=0,
1189	req_blocks=req_blocks,
1190	kv_caches_base_addr=[region.base_addr for region in regions],
1191	block_lens=[region.block_len for region in regions],
1192	kv_block_lens=[region.kv_block_len for region in regions],
1193	registered_layer_names=[region.layer_name for region in regions],
1194	registered_layer_indices=[region.layer_index for region in regions],
1195	registered_group_indices=[region.group_index for region in regions],
1196	registered_shared_group_ids=[
1197	list(region.shared_group_ids) for region in regions
1198	],
1199	registered_row_offsets=[region.row_offset for region in regions],
1200	)
1201	fields.update(overrides)
1202	return MooncakeXferMetadata(**fields)
1203
1204
1205	@contextlib.contextmanager
1206	def mooncake_register_worker(kv_cache_config, *, kv_role: str = "kv_consumer"):
1207	"""Build a worker whose engine registration and sender thread are patched."""
1208	vllm_config = create_vllm_config(kv_connector="MooncakeConnector", kv_role=kv_role)
1209	connector_mod = (
1210	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector"
1211	)
1212	with (
1213	set_current_vllm_config(vllm_config),
1214	patch_worker_dependencies(),
1215	patch(f"{connector_mod}.threading.Event"),
1216	patch(f"{connector_mod}.threading.Thread") as mock_thread,
1217	):
1218	connector = MooncakeConnector(
1219	vllm_config, KVConnectorRole.WORKER, kv_cache_config
1220	)
1221	worker = connector.connector_worker
1222	mock_thread.return_value.is_alive.return_value = False
1223	with patch.object(
1224	worker.engine, "batch_register_memory", return_value=0
1225	) as reg:
1226	yield worker, reg
1227
1228
1229	@pytest.mark.asyncio
1230	@pytest.mark.parametrize(
1231	("local_pp_size", "local_pp_rank", "expected_addrs"),
1232	[
1233	(1, 0, ["tcp://producer-pp0:1234", "tcp://producer-pp1:1234"]),
1234	(2, 1, ["tcp://producer-pp1:1234"]),
1235	],
1236	ids=["heterogeneous_pp_pulls_all_remote_pp", "matching_pp_pulls_same_rank"],
1237	)
1238	async def test_receive_kv_selects_remote_pp_workers(
1239	local_pp_size: int,
1240	local_pp_rank: int,
1241	expected_addrs: list[str],
1242	):
1243	"""Decode workers should not hard-code producer pp_rank 0."""
1244	vllm_config = create_vllm_config(
1245	kv_connector="MooncakeConnector", kv_role="kv_consumer"
1246	)
1247
1248	with set_current_vllm_config(vllm_config), patch_worker_dependencies():
1249	decode_connector = MooncakeConnector(
1250	vllm_config,
1251	KVConnectorRole.WORKER,
1252	_make_test_kv_cache_config(),
1253	)
1254	decode_worker = decode_connector.connector_worker
1255	decode_worker.pp_size = local_pp_size
1256	decode_worker.pp_rank = local_pp_rank
1257	decode_worker._remote_agents = {
1258	"p-engine": {
1259	0: {
1260	0: "tcp://producer-pp0:1234",
1261	1: "tcp://producer-pp1:1234",
1262	}
1263	}
1264	}
1265	decode_worker._tp_size["p-engine"] = 1
1266
1267	pull_metas = {
1268	"d-req-1": PullReqMeta(
1269	d_req_id="d-req-1",
1270	transfer_id="xfer-req-1",
1271	local_block_ids=[[100, 101]],
1272	remote_engine_id="p-engine",
1273	remote_bootstrap_addr="http://bootstrap:33333",
1274	)
1275	}
1276	seen_addrs: list[str] = []
1277
1278	async def fake_receive(worker_addr: str, metas: dict[str, PullReqMeta]):
1279	seen_addrs.append(worker_addr)
1280	for meta in metas.values():
1281	meta.pull_tasks_count -= 1
1282
1283	with patch.object(
1284	decode_worker,
1285	"receive_kv_from_single_worker",
1286	side_effect=fake_receive,
1287	):
1288	decode_worker.receive_kv("p-engine", pull_metas)
1289	await asyncio.sleep(0)
1290
1291	assert seen_addrs == expected_addrs
1292	assert pull_metas["d-req-1"].pull_tasks_count == 0
1293	decode_worker.shutdown()
1294
1295
1296	def test_resolve_need_send_accounts_for_remote_tp_fanout():
1297	"""Producer-side completion waits for every paired consumer TP pull."""
1298	worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker)
1299	worker.async_zmq_ctx = MagicMock()
1300	worker.is_kv_consumer = True
1301	worker.is_kv_producer = True
1302	send_meta = SendBlockMeta(
1303	p_req_id="p-req-1",
1304	transfer_id="xfer-req-1",
1305	local_block_ids=[[1]],
1306	ready=asyncio.Event(),
1307	)
1308
1309	worker.resolve_need_send(send_meta, remote_tp_ranks=[0, 1])
1310
1311	assert send_meta.need_send == 2
1312
1313
1314	@pytest.mark.asyncio
1315	async def test_heterogeneous_pp_waits_for_consumer_without_shared_layers():
1316	"""P4/D2 must retain KV until both D stages finish, including a no-op pull."""
1317	config = create_vllm_config(kv_connector="MooncakeConnector", kv_role="kv_producer")
1318	with (
1319	set_current_vllm_config(config),
1320	patch_worker_dependencies(),
1321	patch.object(MooncakeConnectorWorker, "_sync_block_size_with_kernel"),
1322	):
1323	worker = MooncakeConnector(
1324	config, KVConnectorRole.WORKER, _make_test_kv_cache_config()
1325	).connector_worker
1326	try:
1327	worker.pp_size = 4
1328	worker.kv_caches_base_addr = [0x1000]
1329	worker.block_len_per_layer = [256]
1330	worker.kv_block_len_per_layer = [256]
1331	worker.registered_layer_names = ["model.layers.0.self_attn"]
1332	worker.registered_layer_indices = [0]
1333	send_meta = SendBlockMeta(
1334	p_req_id="p-req",
1335	transfer_id="transfer",
1336	local_block_ids=[[1]],
1337	ready=asyncio.Event(),
1338	)
1339	send_meta.ready.set()
1340	worker.reqs_need_send["transfer"] = send_meta
1341	sock = AsyncMock(spec=zmq.asyncio.Socket, send_multipart=AsyncMock())
1342	with (
1343	patch.object(worker, "sender_loop", asyncio.get_running_loop()),
1344	patch.object(worker, "_send_blocks", return_value=0) as send,
1345	):
1346	for pp_rank in range(2):
1347	metadata = _xfer_meta(
1348	[
1349	_region(
1350	0x2000,
1351	row_offset=-1,
1352	layer_name=f"model.layers.{pp_rank}.self_attn",
1353	layer_index=pp_rank,
1354	)
1355	],
1356	{"d-req": ("transfer", [[2]])},
1357	remote_hostname="consumer",
1358	remote_port=1234 + pp_rank,
1359	remote_pp_size=2,
1360	)
1361	await worker.send_kv_to_decode(b"consumer", sock, metadata)
1362	response = worker._xfer_resp_decoder.decode(
1363	sock.send_multipart.call_args.args[0][1]
1364	)
1365	assert response.status == MooncakeXferResponseStatus.FINISH
1366	assert response.ok_reqs == ["d-req"]
1367	assert ("transfer" in worker.reqs_need_send) == (pp_rank == 0)
1368	assert worker.finished_sending_reqs == (
1369	set() if pp_rank == 0 else {"p-req"}
1370	)
1371	send.assert_called_once_with("consumer:1234", [0x1100], [0x2200], [256])
1372	finally:
1373	worker.shutdown()
1374	worker.is_kv_consumer = True
1375
1376
1377	@pytest.mark.asyncio
1378	@patch(
1379	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
1380	FakeMooncakeWrapper,
1381	)
1382	async def test_kv_producer(monkeypatch):
1383	"""Simulates a Producer Worker (Prefiller) receiving a transfer request
1384	from a Consumer (Decoder).
1385
1386	Verifies memory offset calculation: ptr = base_addr + block_id * block_len.
1387	"""
1388	monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
1389	vllm_config = create_vllm_config(
1390	kv_connector="MooncakeConnector", kv_role="kv_producer"
1391	)
1392
1393	with set_current_vllm_config(vllm_config), patch_worker_dependencies():
1394	prefill_connector = MooncakeConnector(
1395	vllm_config,
1396	KVConnectorRole.WORKER,
1397	_make_test_kv_cache_config(),
1398	)
1399	prefill_worker = prefill_connector.connector_worker
1400	prefill_worker.kv_caches_base_addr = [0x1000]
1401	block_len = 4096
1402	prefill_worker.block_len_per_layer = [block_len]
1403	prefill_worker.kv_block_len_per_layer = [block_len]
1404	prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
1405	prefill_worker.registered_layer_indices = [0]
1406
1407	# Override loop to use current test loop
1408	origin_sender_loop = prefill_worker.sender_loop
1409	prefill_worker.sender_loop = asyncio.get_event_loop()
1410
1411	# A request is finished on Producer and ready to be sent.
1412	transfer_id = "xfer-req-1"
1413	send_meta = SendBlockMeta(
1414	p_req_id="p-req-1",
1415	transfer_id=transfer_id,
1416	local_block_ids=[[10, 11]],
1417	ready=asyncio.Event(),
1418	)
1419	prefill_worker.reqs_need_send[transfer_id] = send_meta
1420	send_meta.ready.set()
1421
1422	# Remote consumer request metadata
1423	xfer_meta = _xfer_meta(
1424	[
1425	_region(
1426	0x2000,
1427	block_len=block_len,
1428	row_offset=-1,
1429	layer_name="model.layers.0.self_attn",
1430	layer_index=0,
1431	)
1432	],
1433	{"d-req-1": (transfer_id, [[20, 21]])},
1434	)
1435
1436	mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
1437	mock_socket.send_multipart = AsyncMock()
1438	identity = b"consumer-id"
1439
1440	with patch.object(
1441	prefill_worker, "_send_blocks", return_value=0
1442	) as mock_send_blocks:
1443
1444	def expected_transfers(src_base, dst_base, src_blocks, dst_blocks):
1445	n = len(src_blocks)
1446	return (
1447	[src_base + src_blocks[0] * block_len],
1448	[dst_base + dst_blocks[0] * block_len],
1449	[n * block_len],
1450	)
1451
1452	# Normal case: 2 blocks to 2 blocks
1453	await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
1454	src, dst, lens = expected_transfers(0x1000, 0x2000, [10, 11], [20, 21])
1455	mock_send_blocks.assert_called_once_with(
1456	"consumer-host:54321",
1457	src,
1458	dst,
1459	lens,
1460	)
1461	mock_socket.send_multipart.assert_called_once()
1462
1463	# Verify the response sent back to the consumer
1464	sent_call = mock_socket.send_multipart.call_args[0][0]
1465	sent_identity, sent_payload = sent_call
1466	assert sent_identity == identity
1467	response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
1468	assert response.status == MooncakeXferResponseStatus.FINISH
1469	assert response.ok_reqs == ["d-req-1"]
1470
1471	# Verify internal state cleanup
1472	assert transfer_id not in prefill_worker.reqs_need_send
1473	assert "p-req-1" in prefill_worker.finished_sending_reqs
1474
1475	# More cases:
1476	# Consumer only needs 1 block (less than P)
1477	mock_send_blocks.reset_mock()
1478	mock_socket.send_multipart.reset_mock()
1479	prefill_worker.reqs_need_send[transfer_id] = send_meta
1480	send_meta.sent = 0
1481	send_meta.ready.set()
1482	xfer_meta.req_blocks["d-req-1"] = (transfer_id, [[20]])
1483	# Worker processes the consumer's request
1484	await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
1485	# Verify transfer parameters are correct: 11 to 20
1486	src, dst, lens = expected_transfers(0x1000, 0x2000, [11], [20])
1487	mock_send_blocks.assert_called_once_with(
1488	"consumer-host:54321",
1489	src,
1490	dst,
1491	lens,
1492	)
1493	mock_socket.send_multipart.assert_called_once()
1494
1495	# Consumer needs 3 blocks (more than P, error case)
1496	mock_send_blocks.reset_mock()
1497	mock_socket.send_multipart.reset_mock()
1498	prefill_worker.reqs_need_send[transfer_id] = send_meta
1499	send_meta.sent = 0
1500	send_meta.ready.set()
1501	xfer_meta.req_blocks["d-req-1"] = (transfer_id, [[20, 21, 22]])
1502	# Worker processes the consumer's request
1503	await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
1504	# This should not be called because error.
1505	mock_send_blocks.assert_not_called()
1506	mock_socket.send_multipart.assert_called_once()
1507	_, sent_payload = mock_socket.send_multipart.call_args[0][0]
1508	response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
1509	assert response.err_msg == "P num blocks less than D"
1510	assert response.err_reqs == ["d-req-1"]
1511
1512	# Timeout
1513	mock_send_blocks.reset_mock()
1514	mock_socket.send_multipart.reset_mock()
1515	prefill_worker.reqs_need_send[transfer_id] = send_meta
1516	send_meta.sent = 0
1517	send_meta.ready.clear()
1518	xfer_meta.req_blocks["d-req-1"] = (transfer_id, [[20, 21]])
1519	# Worker processes the consumer's request
1520	await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
1521	# This should not be called because timeout.
1522	mock_send_blocks.assert_not_called()
1523	mock_socket.send_multipart.assert_called_once()
1524	_, sent_payload = mock_socket.send_multipart.call_args[0][0]
1525	response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
1526	assert response.err_msg == "Timeout waiting for P side ready."
1527	assert response.err_reqs == ["d-req-1"]
1528
1529	# Transfer error
1530	with patch.object(
1531	prefill_worker, "_send_blocks", return_value=123
1532	) as mock_send_blocks:
1533	mock_socket.send_multipart.reset_mock()
1534	prefill_worker.reqs_need_send[transfer_id] = send_meta
1535	send_meta.sent = 0
1536	send_meta.ready.set()
1537	xfer_meta.req_blocks["d-req-1"] = (transfer_id, [[20, 21]])
1538	# Worker processes the consumer's request
1539	await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
1540	mock_send_blocks.assert_called_once()
1541	mock_socket.send_multipart.assert_called_once()
1542	_, sent_payload = mock_socket.send_multipart.call_args[0][0]
1543	response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
1544	assert response.err_msg == "Mooncake transfer engine returned 123"
1545	assert response.err_reqs == ["d-req-1"]
1546
1547	# Clean up
1548	prefill_worker.sender_loop = origin_sender_loop
1549	prefill_worker.shutdown()
1550
1551
1552	@pytest.mark.asyncio
1553	async def test_kv_consumuer(monkeypatch):
1554	"""Simulates a Consumer Worker (Decoder) initiating a pull from a Producer.
1555
1556	Verifies that MooncakeXferMetadata is correctly serialized and sent via ZMQ.
1557	"""
1558	vllm_config = create_vllm_config(
1559	kv_connector="MooncakeConnector", kv_role="kv_consumer"
1560	)
1561
1562	with set_current_vllm_config(vllm_config), patch_worker_dependencies() as mocks:
1563	decode_connector = MooncakeConnector(
1564	vllm_config,
1565	KVConnectorRole.WORKER,
1566	_make_test_kv_cache_config(),
1567	)
1568	decode_worker = decode_connector.connector_worker
1569	decode_worker.kv_caches_base_addr = [0x1000]
1570	decode_worker.block_len_per_layer = [4096]
1571	decode_worker.kv_block_len_per_layer = [4096]
1572	decode_worker.registered_layer_names = ["model.layers.0.self_attn"]
1573	decode_worker.registered_layer_indices = [0]
1574	decode_worker.rpc_port = 54321
1575
1576	# A request to pull data arrives.
1577	pull_metas = {
1578	"d-req-1": PullReqMeta(
1579	d_req_id="d-req-1",
1580	transfer_id="xfer-req-1",
1581	local_block_ids=[[100, 101]],
1582	remote_engine_id="p-engine",
1583	remote_bootstrap_addr="http://bootstrap:33333",
1584	pull_tasks_count=1,
1585	)
1586	}
1587	decode_worker._remote_agents = {"p-engine": {0: {0: "tcp://producer:1234"}}}
1588	decode_worker._tp_size["p-engine"] = 1
1589
1590	# Mock the response from the producer.
1591	mock_response = MooncakeXferResponse(
1592	status=MooncakeXferResponseStatus.FINISH, ok_reqs=["d-req-1"]
1593	)
1594	encoded_response = decode_worker._encoder.encode(mock_response)
1595	mocks["mock_socket_object"].recv.return_value = encoded_response
1596
1597	# Trigger the receive logic.
1598	decode_worker.receive_kv("p-engine", pull_metas)
1599	await asyncio.sleep(1) # Allow async task to run
1600
1601	# Verify the metadata sent to the producer.
1602	mocks["mock_make_zmq"].assert_called_with(
1603	decode_worker.async_zmq_ctx,
1604	"tcp://producer:1234",
1605	zmq.DEALER,
1606	bind=False,
1607	linger=0,
1608	)
1609	sent_payload = mocks["mock_socket_object"].send.call_args[0][0]
1610	sent_meta = decode_worker._xfer_meta_decoder.decode(sent_payload)
1611
1612	assert sent_meta.remote_hostname == "127.0.0.1"
1613	assert sent_meta.remote_port == 54321
1614	assert sent_meta.req_blocks["d-req-1"] == ("xfer-req-1", [[100, 101]])
1615	assert sent_meta.kv_caches_base_addr == [0x1000]
1616	assert sent_meta.block_lens == [4096]
1617	assert sent_meta.registered_layer_names == ["model.layers.0.self_attn"]
1618	assert sent_meta.registered_layer_indices == [0]
1619
1620	# Verify internal state is updated correctly.
1621	assert "d-req-1" in decode_worker.finished_recving_reqs
1622
1623	# Clean up
1624	decode_worker.shutdown()
1625
1626
1627	@pytest.mark.asyncio
1628	async def test_worker_get_finished_timeout(monkeypatch):
1629	"""Tests the cleanup mechanism for requests."""
1630	vllm_config = create_vllm_config(
1631	kv_connector="MooncakeConnector", kv_role="kv_producer"
1632	)
1633	with set_current_vllm_config(vllm_config), patch_worker_dependencies():
1634	prefill_connector = MooncakeConnector(
1635	vllm_config,
1636	KVConnectorRole.WORKER,
1637	_make_test_kv_cache_config(),
1638	)
1639	prefill_worker = prefill_connector.connector_worker
1640
1641	# Add an expired request (expire_time is in the past).
1642	prefill_worker.reqs_need_send["tx-expired"] = SendBlockMeta(
1643	p_req_id="p-req-expired",
1644	transfer_id="tx-expired",
1645	local_block_ids=[[1, 2]],
1646	ready=MagicMock(),
1647	expire_time=time.perf_counter() - 100,
1648	)
1649
1650	# Add a non-expired request.
1651	prefill_worker.reqs_need_send["tx-active"] = SendBlockMeta(
1652	p_req_id="p-req-active",
1653	transfer_id="tx-active",
1654	local_block_ids=[[3, 4]],
1655	ready=MagicMock(),
1656	expire_time=time.perf_counter() + 100,
1657	)
1658
1659	finished_reqs = await prefill_worker.fetch_finished_sending_reqs()
1660
1661	assert "p-req-expired" in finished_reqs
1662	assert "p-req-active" not in finished_reqs
1663	assert "tx-expired" not in prefill_worker.reqs_need_send
1664	assert "tx-active" in prefill_worker.reqs_need_send
1665
1666
1667	@pytest.mark.parametrize(
1668	("layout", "separate_kv_head_groups"),
1669	[
1670	(KVCacheLayout.LBHNC, False),
1671	(KVCacheLayout.BLHNC, False),
1672	(KVCacheLayout.BHLNC, False),
1673	# LHBNC gives each head group its own region; the K/V split doubles the
1674	# head count but the registration shape is driven by the layout.
1675	(KVCacheLayout.LHBNC, False),
1676	(KVCacheLayout.LHBNC, True),
1677	],
1678	)
1679	def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool):
1680	"""Tests the memory registration logic with the underlying Mooncake engine."""
1681	spec = FullAttentionSpec(
1682	block_size=16,
1683	num_kv_heads=4,
1684	head_size=64,
1685	dtype=torch.float16,
1686	num_head_slots=2 if separate_kv_head_groups else None,
1687	state_content_bytes=4 * 64 * 2 if separate_kv_head_groups else None,
1688	)
1689	layer_names = [
1690	"model.layers.0.self_attn",
1691	"model.layers.1.self_attn",
1692	]
1693	with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, reg):
1694	for layer_name in layer_names:
1695	worker._layer_specs[layer_name] = spec
1696	raw = torch.zeros(2 * 2 * spec.page_size_bytes, dtype=torch.int8)
1697	tensor1, tensor2 = dense_kv_cache_views(raw, spec, 2, 2, layout)
1698	kv_caches = dict(zip(layer_names, (tensor1, tensor2)))
1699	worker.register_kv_caches(kv_caches)
1700
1701	reg.assert_called_once()
1702	registered_ptrs, registered_lens = reg.call_args[0]
1703	assert registered_ptrs == [raw.data_ptr()]
1704	assert registered_lens == [raw.nbytes]
1705
1706	n_blocks = tensor1.shape[0] if tensor1.ndim > 1 else 1
1707	storage_end = raw.data_ptr() + raw.nbytes
1708	for base, stride, kv_len in zip(
1709	worker.kv_caches_base_addr,
1710	worker.block_len_per_layer,
1711	worker.kv_block_len_per_layer,
1712	):
1713	assert base + (n_blocks - 1) * stride + kv_len <= storage_end
1714
1715	# Expected shape is the layout, not a recomputation of the pack predicate.
1716	kind = {
1717	KVCacheLayout.LBHNC: "page",
1718	KVCacheLayout.BLHNC: "page",
1719	KVCacheLayout.BHLNC: "opaque",
1720	KVCacheLayout.LHBNC: "per_head",
1721	}[layout]
1722	if kind == "opaque":
1723	packed_row = raw.nbytes // n_blocks
1724	assert worker.kv_caches_base_addr == [raw.data_ptr()]
1725	assert worker.block_len_per_layer == [packed_row]
1726	assert worker.kv_block_len_per_layer == [packed_row]
1727	assert worker.registered_layer_names == [layer_names[0]]
1728	elif kind == "per_head":
1729	expected_addrs = [
1730	cache[:, head_idx].data_ptr()
1731	for cache in (tensor1, tensor2)
1732	for head_idx in range(cache.shape[1])
1733	]
1734	head_block_bytes = tensor1.stride(0) * tensor1.element_size()
1735	assert worker.kv_caches_base_addr == expected_addrs
1736	assert worker.block_len_per_layer == [head_block_bytes] * len(
1737	expected_addrs
1738	)
1739	assert worker.kv_block_len_per_layer == [head_block_bytes] * len(
1740	expected_addrs
1741	)
1742	assert worker.registered_layer_names == [
1743	layer_name
1744	for layer_name in layer_names
1745	for _ in range(tensor1.shape[1])
1746	]
1747	else:
1748	page_stride = tensor1.stride(0) * tensor1.element_size()
1749	assert len(worker.block_len_per_layer) == len(kv_caches)
1750	assert worker.block_len_per_layer == [page_stride, page_stride]
1751	assert worker.kv_block_len_per_layer == [spec.page_size_bytes] * 2
1752	assert worker.registered_layer_names == list(kv_caches)
1753	assert worker.registered_layer_indices == [0, 1]
1754
1755
1756	def test_register_noncontiguous_packed_allows_matching_pp():
1757	"""Same PP on both sides can still register a non-contiguous packed row.
1758
1759	A remote PP mismatch is rejected at send time, once the peer PP size is
1760	known. Rejecting every local PP>1 layout at startup is too strict.
1761	"""
1762	spec = FullAttentionSpec(
1763	block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
1764	)
1765	layer_names = ["model.layers.0.self_attn", "model.layers.1.self_attn"]
1766	with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, _reg):
1767	worker.pp_size = 2
1768	for layer_name in layer_names:
1769	worker._layer_specs[layer_name] = spec
1770	raw = torch.zeros(2 * 2 * spec.page_size_bytes, dtype=torch.int8)
1771	tensor1, tensor2 = dense_kv_cache_views(raw, spec, 2, 2, KVCacheLayout.BHLNC)
1772	assert not tensor1[0].is_contiguous()
1773	worker.register_kv_caches(dict(zip(layer_names, (tensor1, tensor2))))
1774	assert worker.kv_caches_base_addr == [raw.data_ptr()]
1775	assert worker.region_row_offsets == [0]
1776
1777
1778	def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
1779	"""Mixed MLA+Eagle caches should register by byte length, not shape."""
1780	with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, reg):
1781	worker.use_mla = True
1782	worker.transfer_topo.is_mla = True
1783	mla_cache = torch.zeros((2, 16, 512), dtype=torch.float16)
1784	eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16)
1785	kv_caches = {
1786	"model.layers.0.mla_attn": mla_cache,
1787	"model.layers.1.eagle_attn": eagle_cache,
1788	}
1789	worker.register_kv_caches(kv_caches)
1790
1791	reg.assert_called_once()
1792	registered_ptrs, registered_lens = reg.call_args[0]
1793	assert registered_ptrs == [mla_cache.data_ptr(), eagle_cache.data_ptr()]
1794	assert registered_lens == [mla_cache.nbytes, eagle_cache.nbytes]
1795	assert worker.block_len_per_layer == [
1796	mla_cache.nbytes // mla_cache.shape[0],
1797	eagle_cache.nbytes // eagle_cache.shape[0],
1798	]
1799	assert worker.registered_layer_names == [
1800	"model.layers.0.mla_attn",
1801	"model.layers.1.eagle_attn",
1802	]
1803	assert worker.registered_layer_indices == [0, 1]
1804	assert len(worker.kv_caches_base_addr) == 2
1805
1806
1807	def test_packed_and_unpacked_region_lengths_fail_homogeneous_tp_handshake():
1808	packed = TransferRegion(
1809	layer_name="model.layers.0.mla_attn",
1810	layer_index=0,
1811	base_addr=0x1000,
1812	block_len=2048,
1813	kv_block_len=2048,
1814	group_index=0,
1815	)
1816	unpacked = TransferRegion(
1817	layer_name="model.layers.0.mla_attn",
1818	layer_index=0,
1819	base_addr=0x1000,
1820	block_len=64,
1821	kv_block_len=64,
1822	group_index=0,
1823	)
1824	err = _validate_asymmetric_region_lengths(
1825	local_regions=[packed],
1826	remote_regions=[unpacked],
1827	local_tp_size=1,
1828	remote_tp_size=1,
1829	producer_cache_replicated=False,
1830	)
1831	assert err is not None
1832	assert "length mismatch" in err
1833
1834	count_err = _validate_asymmetric_region_lengths(
1835	local_regions=[packed],
1836	remote_regions=[unpacked, unpacked],
1837	local_tp_size=1,
1838	remote_tp_size=1,
1839	producer_cache_replicated=False,
1840	)
1841	assert count_err is not None
1842	assert "region counts" in count_err
1843
1844
1845	@pytest.mark.parametrize(
1846	("num_groups", "expected_group_index"),
1847	[(1, 0), (2, _SHARED_REGION_GROUP_ID)],
1848	)
1849	def test_register_kv_caches_collapses_shared_mla_storage(
1850	num_groups, expected_group_index
1851	):
1852	"""Shared backing collapses to one packed region; two groups mark it -1."""
1853	num_blocks = 4
1854	packed_row = 256
1855	with mooncake_register_worker(
1856	_make_packed_mla_kv_cache_config(num_blocks, num_groups=num_groups)
1857	) as (worker, reg):
1858	backing = torch.zeros((num_blocks, packed_row), dtype=torch.uint8)
1859	worker.register_kv_caches(
1860	{
1861	"model.layers.0.mla_attn": backing,
1862	"model.layers.1.indexer": backing,
1863	}
1864	)
1865
1866	reg.assert_called_once()
1867	registered_ptrs, registered_lens = reg.call_args[0]
1868	assert registered_ptrs == [backing.data_ptr()]
1869	assert registered_lens == [backing.untyped_storage().nbytes()]
1870	assert worker.kv_caches_base_addr == [backing.data_ptr()]
1871	assert worker.block_len_per_layer == [packed_row]
1872	assert worker.kv_block_len_per_layer == [packed_row]
1873	assert worker.registered_layer_names == ["model.layers.0.mla_attn"]
1874	assert worker.registered_group_indices == [expected_group_index]
1875	if expected_group_index == _SHARED_REGION_GROUP_ID:
1876	assert worker.region_shared_groups == [(0, 1)]
1877
1878
1879	def test_block_ids_for_region_flattens_shared_groups():
1880	ids = [[10, 11], [12]]
1881	assert _block_ids_for_region(ids, 0) == [10, 11]
1882	assert _block_ids_for_region(ids, 1) == [12]
1883	assert _block_ids_for_region(ids, _SHARED_REGION_GROUP_ID, (0, 1)) == [
1884	10,
1885	11,
1886	12,
1887	]
1888	# A group that does not share the allocation stays out.
1889	assert _block_ids_for_region(
1890	[[10, 11], [12], [99]], _SHARED_REGION_GROUP_ID, (0, 1)
1891	) == [10, 11, 12]
1892
1893
1894	@pytest.mark.asyncio
1895	@pytest.mark.parametrize(
1896	("group_index", "local_ids", "remote_ids", "n_blocks"),
1897	[
1898	(0, [[10, 11]], [[20, 21]], 2),
1899	(_SHARED_REGION_GROUP_ID, [[10, 11], [12]], [[20, 21], [22]], 3),
1900	],
1901	)
1902	async def test_build_transfer_params_sends_packed_region_once(
1903	group_index, local_ids, remote_ids, n_blocks
1904	):
1905	"""Packed region coalesces contiguous blocks, including shared-group flatten."""
1906	worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker)
1907	worker.async_zmq_ctx = MagicMock()
1908	worker.is_kv_consumer = True
1909	worker.is_kv_producer = True
1910	worker.tp_rank = 0
1911	worker.tp_size = 1
1912	worker.use_mla = True
1913	worker.kv_cache_config = _make_packed_mla_kv_cache_config(
1914	num_blocks=4, num_groups=len(local_ids)
1915	)
1916	worker._physical_blocks_per_logical_kv_block = 1
1917	worker.transfer_topo = SimpleNamespace(
1918	local_replicates_kv_cache=False,
1919	total_num_kv_heads=1,
1920	)
1921
1922	block_len = 256
1923	shared = (0, 1) if group_index == _SHARED_REGION_GROUP_ID else ()
1924	local_region = _region(
1925	0x1000,
1926	block_len=block_len,
1927	row_offset=0,
1928	layer_name="model.layers.0.mla_attn",
1929	group_index=group_index,
1930	shared_group_ids=shared,
1931	)
1932	remote_region = _region(
1933	0xA000,
1934	block_len=block_len,
1935	row_offset=0,
1936	layer_name="model.layers.0.mla_attn",
1937	group_index=group_index,
1938	shared_group_ids=shared,
1939	)
1940	transfer_id = "xfer-packed"
1941	send_meta = SendBlockMeta(
1942	p_req_id="p-packed",
1943	transfer_id=transfer_id,
1944	local_block_ids=local_ids,
1945	ready=asyncio.Event(),
1946	)
1947	xfer_meta = _xfer_meta([remote_region], {"d-packed": (transfer_id, remote_ids)})
1948
1949	(
1950	src_ptrs,
1951	dst_ptrs,
1952	lengths,
1953	err_reqs,
1954	err_msg,
1955	) = await worker._build_transfer_params(
1956	ready_reqs=[("d-packed", send_meta)],
1957	agent_meta=xfer_meta,
1958	local_regions=[local_region],
1959	remote_regions=[remote_region],
1960	)
1961
1962	assert err_reqs == []
1963	assert err_msg is None
1964	assert src_ptrs == [0x1000 + 10 * block_len]
1965	assert dst_ptrs == [0xA000 + 20 * block_len]
1966	assert lengths == [n_blocks * block_len]
1967
1968
1969	def test_coalesce_promotes_padding_only_for_a_full_row():
1970	"""A full-row merge includes the padding tail; a partial row does not."""
1971	page, row = 1024, 2560
1972	local = [
1973	_region(0, block_len=row, kv_block_len=page, row_offset=0),
1974	_region(page, block_len=row, kv_block_len=page, row_offset=page),
1975	]
1976	remote = [
1977	_region(100, block_len=row, kv_block_len=page, row_offset=0),
1978	_region(100 + page, block_len=row, kv_block_len=page, row_offset=page),
1979	]
1980	full_local, full_remote = _coalesce_contiguous_transfer_regions(
1981	local, remote, promote_full_row=True
1982	)
1983	assert len(full_local) == 1
1984	assert full_local[0].kv_block_len == full_remote[0].kv_block_len == row
1985
1986	partial_local, _ = _coalesce_contiguous_transfer_regions(local, remote)
1987	assert partial_local[0].kv_block_len == 2 * page
1988
1989
1990	def test_coalesce_spans_padding_between_pages():
1991	"""Unpadded pages merge across the padding gap without a full-row lift."""
1992	row = 2560
1993	merged, _ = _coalesce_contiguous_transfer_regions(
1994	[
1995	_region(0, block_len=row, kv_block_len=1000),
1996	_region(1024, block_len=row, kv_block_len=1000),
1997	],
1998	[
1999	_region(5000, block_len=row, kv_block_len=1000, row_offset=0),
2000	_region(6024, block_len=row, kv_block_len=1000, row_offset=1024),
2001	],
2002	)
2003	assert len(merged) == 1
2004	assert merged[0].kv_block_len == 2024
2005	assert merged[0].row_offset == 0
2006
2007
2008	def test_coalesce_promotes_each_group_that_starts_at_row_zero():
2009	"""Groups overlay one row, so each row-start run is its own full-row copy."""
2010	page, gap, row = 1000, 1024, 4096
2011
2012	def side(origin: int) -> list[TransferRegion]:
2013	return [
2014	_region(
2015	origin + base,
2016	block_len=row,
2017	kv_block_len=page,
2018	row_offset=offset,
2019	layer_name=f"group-{group}",
2020	layer_index=group,
2021	group_index=group,
2022	)
2023	for group, base, offset in (
2024	(0, 0, 0),
2025	(0, gap, gap),
2026	(1, 0, 0),
2027	(1, gap, gap),
2028	)
2029	]
2030
2031	merged_local, merged_remote = _coalesce_contiguous_transfer_regions(
2032	side(0), side(8000), promote_full_row=True
2033	)
2034	assert [region.group_index for region in merged_local] == [0, 1]
2035	assert [region.kv_block_len for region in merged_local] == [row, row]
2036	assert [region.kv_block_len for region in merged_remote] == [row, row]
2037	assert merged_local[1].base_addr == 0
2038	assert merged_remote[1].base_addr == 8000
2039
2040
2041	def test_coalesce_does_not_promote_a_mid_row_run():
2042	"""A matched PP slice that does not start at row offset 0 stays partial."""
2043	page, row = 1000, 4096
2044	merged, _ = _coalesce_contiguous_transfer_regions(
2045	[
2046	_region(
2047	1024, block_len=row, kv_block_len=page, row_offset=1024, layer_index=2
2048	),
2049	_region(
2050	2048, block_len=row, kv_block_len=page, row_offset=2048, layer_index=2
2051	),
2052	],
2053	[
2054	_region(
2055	9000 + 1024,
2056	block_len=row,
2057	kv_block_len=page,
2058	row_offset=1024,
2059	layer_index=2,
2060	),
2061	_region(
2062	9000 + 2048,
2063	block_len=row,
2064	kv_block_len=page,
2065	row_offset=2048,
2066	layer_index=2,
2067	),
2068	],
2069	promote_full_row=True,
2070	)
2071	assert len(merged) == 1
2072	assert merged[0].row_offset == 1024
2073	assert merged[0].kv_block_len == 1024 + page
2074
2075
2076	def test_opaque_packed_row_is_full_stride_at_offset_zero():
2077	"""Branch B is row-sized at offset 0. A page view at offset 0 is not."""
2078	assert _has_opaque_packed_row([0], [4096], [4096])
2079	assert not _has_opaque_packed_row([0, 1024], [1000, 1000], [4096, 4096])
2080	assert not _has_opaque_packed_row([], [4096], [4096])
2081
2082
2083	def test_pp_mismatch_rejects_opaque_row_in_both_directions():
2084	"""The smaller PP side can be either the producer or the consumer."""
2085	page_view = dict(
2086	remote_row_offsets=[0, 1024],
2087	remote_kv_block_lens=[1000, 1000],
2088	remote_block_lens=[4096, 4096],
2089	)
2090	opaque_row = dict(
2091	remote_row_offsets=[0],
2092	remote_kv_block_lens=[4096],
2093	remote_block_lens=[4096],
2094	)
2095	# P pp=2, D pp=1, local row is opaque.
2096	assert _pp_mismatch_hides_packed_layers(
2097	2, 1, local_has_opaque_row=True, **page_view
2098	)
2099	# P pp=1, D pp=2, remote row is opaque.
2100	assert _pp_mismatch_hides_packed_layers(
2101	1, 2, local_has_opaque_row=False, **opaque_row
2102	)
2103	# Matching PP still aligns by the shared first layer name.
2104	assert not _pp_mismatch_hides_packed_layers(
2105	2, 2, local_has_opaque_row=True, **opaque_row
2106	)
2107	# Hetero PP with page-contiguous views is not this rejection.
2108	assert not _pp_mismatch_hides_packed_layers(
2109	2, 1, local_has_opaque_row=False, **page_view
2110	)
2111
2112
2113	def test_prepare_transfer_regions_reuses_success_and_error():
2114	"""A peer layout is built once. Re-registering drops the cache."""
2115	worker, _, _, _ = _register_sliced_packed_mla([0, 1])
2116	regions = worker._get_transfer_regions(
2117	worker.kv_caches_base_addr,
2118	worker.block_len_per_layer,
2119	worker.kv_block_len_per_layer,
2120	worker.registered_layer_names,
2121	worker.registered_layer_indices,
2122	worker.registered_group_indices,
2123	[tuple(groups) for groups in worker.region_shared_groups],
2124	worker.region_row_offsets,
2125	)
2126	meta = _xfer_meta(
2127	regions,
2128	{},
2129	remote_hostname="peer",
2130	remote_port=1,
2131	remote_tp_size=worker.tp_size,
2132	remote_pp_size=worker.pp_size,
2133	)
2134	calls = 0
2135	original = worker._get_transfer_regions
2136
2137	def counting(args, *kwargs):
2138	nonlocal calls
2139	calls += 1
2140	return original(args, *kwargs)
2141
2142	worker._get_transfer_regions = counting
2143	prepared = worker._prepare_transfer_regions(meta)
2144	assert prepared[2] is None
2145	assert calls == 2
2146	assert worker._prepare_transfer_regions(meta) == prepared
2147	assert calls == 2
2148
2149	mismatch = _xfer_meta(
2150	[
2151	_region(
2152	0x2000,
2153	block_len=1024,
2154	row_offset=0,
2155	layer_name="model.layers.9.mla_attn",
2156	layer_index=9,
2157	)
2158	],
2159	{},
2160	remote_hostname="peer",
2161	remote_port=2,
2162	remote_tp_size=worker.tp_size,
2163	remote_pp_size=worker.pp_size,
2164	)
2165	failed = worker._prepare_transfer_regions(mismatch)
2166	assert failed[2] is not None
2167	assert calls == 4
2168	assert worker._prepare_transfer_regions(mismatch)[2] == failed[2]
2169	assert calls == 4
2170
2171	page = 1024
2172	backing = torch.zeros((4, page * 2), dtype=torch.uint8)
2173	kv_caches = {
2174	"model.layers.0.mla_attn": backing[:, :page],
2175	"model.layers.1.mla_attn": backing[:, page:],
2176	}
2177	worker.is_kv_consumer = True
2178	# The first register built a producer. Its thread patch is already gone,
2179	# so stay on the consumer path and return before the sender thread starts.
2180	with patch.object(worker.engine, "batch_register_memory", return_value=0):
2181	worker.register_kv_caches(kv_caches)
2182	assert worker._prepared_transfer_regions == {}
2183	worker._prepare_transfer_regions(meta)
2184	assert calls == 6
2185
2186
2187	def test_coalesce_keeps_a_run_inside_its_starting_row():
2188	"""The row bound is the run start, not the previous slice."""
2189	row = 2500
2190	slices = [
2191	_region(base, block_len=row, kv_block_len=1000) for base in (0, 1000, 2000)
2192	]
2193	merged, _ = _coalesce_contiguous_transfer_regions(slices, slices)
2194	assert len(merged) == 2
2195	assert merged[0].kv_block_len == 2000
2196	assert merged[1].base_addr == 2000
2197	assert merged[1].kv_block_len == 1000
2198
2199
2200	def _layer_name(idx: int) -> str:
2201	return f"model.layers.{idx}.mla_attn"
2202
2203
2204	def _register_sliced_packed_mla(
2205	layer_idxs: list[int], *, page: int = 1024, num_blocks: int = 4
2206	):
2207	"""Register column slices of one MLA packed row. Returns worker fields."""
2208	spec = MLAAttentionSpec(
2209	block_size=16,
2210	num_kv_heads=1,
2211	head_size=64,
2212	dtype=torch.uint8,
2213	)
2214	kv_cache_config = KVCacheConfig(
2215	num_blocks=num_blocks,
2216	kv_cache_tensors=[],
2217	kv_cache_groups=[
2218	KVCacheGroupSpec([_layer_name(i) for i in layer_idxs], spec),
2219	],
2220	)
2221	backing = torch.zeros((num_blocks, page * len(layer_idxs)), dtype=torch.uint8)
2222	kv_caches = {
2223	_layer_name(idx): backing[:, i * page : (i + 1) * page]
2224	for i, idx in enumerate(layer_idxs)
2225	}
2226	with mooncake_register_worker(kv_cache_config, kv_role="kv_producer") as (
2227	worker,
2228	_,
2229	):
2230	worker.register_kv_caches(kv_caches)
2231	regions = worker._get_transfer_regions(
2232	worker.kv_caches_base_addr,
2233	worker.block_len_per_layer,
2234	worker.kv_block_len_per_layer,
2235	worker.registered_layer_names,
2236	worker.registered_layer_indices,
2237	worker.registered_group_indices,
2238	row_offsets=worker.region_row_offsets,
2239	)
2240	return worker, kv_cache_config, regions, backing.data_ptr()
2241
2242
2243	def test_sliced_packed_mla_registers_each_layer_view():
2244	"""Each column of a packed MLA row stays its own view, in row order."""
2245	page = 1024
2246	layer_idxs = [0, 1]
2247	with mooncake_register_worker(
2248	KVCacheConfig(
2249	num_blocks=4,
2250	kv_cache_tensors=[],
2251	kv_cache_groups=[
2252	KVCacheGroupSpec(
2253	[_layer_name(i) for i in layer_idxs],
2254	MLAAttentionSpec(
2255	block_size=16,
2256	num_kv_heads=1,
2257	head_size=64,
2258	dtype=torch.uint8,
2259	),
2260	)
2261	],
2262	),
2263	kv_role="kv_producer",
2264	) as (worker, reg):
2265	backing = torch.zeros((4, page * len(layer_idxs)), dtype=torch.uint8)
2266	worker.register_kv_caches(
2267	{
2268	_layer_name(idx): backing[:, i * page : (i + 1) * page]
2269	for i, idx in enumerate(layer_idxs)
2270	}
2271	)
2272	reg.assert_called_once()
2273	assert reg.call_args[0][0] == [backing.data_ptr()]
2274	assert worker.region_row_offsets == [0, page]
2275
2276
2277	async def _copy_packed_slices(
2278	p_layers: list[int],
2279	d_layers: list[int],
2280	*,
2281	allow_partial: bool,
2282	local_ids: list[list[int]],
2283	remote_ids: list[list[int]],
2284	page: int = 1024,
2285	):
2286	p_worker, p_config, p_regions, p_base = _register_sliced_packed_mla(
2287	p_layers, page=page
2288	)
2289	_, _, d_regions, d_base = _register_sliced_packed_mla(d_layers, page=page)
2290	aligned_p, aligned_d, err = _align_transfer_regions(
2291	p_regions, d_regions, allow_partial_layers=allow_partial
2292	)
2293	assert err is None
2294	coalesced_p, coalesced_d = _coalesce_contiguous_transfer_regions(
2295	aligned_p, aligned_d
2296	)
2297	transfer_id = "xfer-packed-slices"
2298	send_meta = SendBlockMeta(
2299	p_req_id="p-slices",
2300	transfer_id=transfer_id,
2301	local_block_ids=local_ids,
2302	ready=asyncio.Event(),
2303	)
2304	p_worker.kv_cache_config = p_config
2305	src, dst, lengths, err_reqs, err_msg = await p_worker._build_transfer_params(
2306	ready_reqs=[("d-slices", send_meta)],
2307	agent_meta=_xfer_meta(coalesced_d, {"d-slices": (transfer_id, remote_ids)}),
2308	local_regions=coalesced_p,
2309	remote_regions=coalesced_d,
2310	)
2311	return SimpleNamespace(
2312	p_regions=p_regions,
2313	d_regions=d_regions,
2314	coalesced_p=coalesced_p,
2315	coalesced_d=coalesced_d,
2316	p_base=p_base,
2317	d_base=d_base,
2318	src=src,
2319	dst=dst,
2320	lengths=lengths,
2321	err_reqs=err_reqs,
2322	err_msg=err_msg,
2323	)
2324
2325
2326	@pytest.mark.asyncio
2327	async def test_packed_hetero_pp_keeps_shared_layers():
2328	"""P-PP4/D-PP2 copies the shared MLA span, not either side's full row."""
2329	page = 1024
2330	copied = await _copy_packed_slices(
2331	[2, 3],
2332	[0, 1, 2, 3],
2333	allow_partial=True,
2334	local_ids=[[10, 11]],
2335	remote_ids=[[10, 11]],
2336	page=page,
2337	)
2338	assert [region.layer_index for region in copied.p_regions] == [2, 3]
2339	assert [region.kv_block_len for region in copied.p_regions] == [page, page]
2340	assert copied.p_regions[0].block_len == 2 * page
2341	assert copied.d_regions[0].block_len == 4 * page
2342	assert len(copied.coalesced_p) == 1
2343	assert copied.coalesced_p[0].kv_block_len == 2 * page
2344	assert copied.coalesced_d[0].base_addr == copied.d_base + 2 * page
2345	assert copied.err_reqs == []
2346	assert copied.err_msg is None
2347	assert copied.lengths == [2 * page, 2 * page]
2348	assert copied.src == [
2349	copied.p_base + 10 * (2 * page),
2350	copied.p_base + 11 * (2 * page),
2351	]
2352	assert copied.dst == [
2353	copied.d_base + 2 * page + 10 * (4 * page),
2354	copied.d_base + 2 * page + 11 * (4 * page),
2355	]
2356
2357
2358	@pytest.mark.asyncio
2359	async def test_packed_matching_layers_coalesce_to_one_copy():
2360	"""Identical packed layer sets still emit one copy for contiguous blocks."""
2361	page = 1024
2362	copied = await _copy_packed_slices(
2363	[0, 1],
2364	[0, 1],
2365	allow_partial=False,
2366	local_ids=[[10, 11]],
2367	remote_ids=[[20, 21]],
2368	page=page,
2369	)
2370	assert copied.err_reqs == []
2371	assert copied.err_msg is None
2372	assert len(copied.coalesced_p) == 1
2373	assert copied.coalesced_p[0].kv_block_len == copied.coalesced_p[0].block_len
2374	block_len = 2 * page
2375	assert copied.src == [copied.p_base + 10 * block_len]
2376	assert copied.dst == [copied.d_base + 20 * block_len]
2377	assert copied.lengths == [2 * block_len]
2378
2379
2380	@pytest.mark.asyncio
2381	@patch(
2382	"vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
2383	"mooncake_connector.TransferEngine",
2384	FakeMooncakeWrapper,
2385	)
2386	@pytest.mark.parametrize("d_tp_size", [1, 4], ids=["p_tp2_d_tp1", "p_tp2_d_tp4"])
2387	async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
2388	"""Tests heterogeneous TP support in the producer transfer path.
2389
2390	Verifies correct pointer and offset calculation when producer TP=2
2391	sends to consumer with TP=1 (P>D) or TP=4 (P<D).
2392
2393	Parametrized cases:
2394	- P TP=2 > D TP=1: one D rank receives; dst_offset based on P rank
2395	- P TP=2 < D TP=4: two D ranks receive; src_offset based on D rank
2396	"""
2397	P_TP_SIZE = 2
2398	P_TP_RANK = 0
2399	LOCAL_BLOCK_LEN = 4096
2400
2401	local_block_len = LOCAL_BLOCK_LEN
2402	remote_block_len = LOCAL_BLOCK_LEN * P_TP_SIZE // d_tp_size
2403
2404	monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
2405	vllm_config = create_vllm_config(
2406	kv_connector="MooncakeConnector", kv_role="kv_producer"
2407	)
2408
2409	with set_current_vllm_config(vllm_config), patch_worker_dependencies():
2410	prefill_connector = MooncakeConnector(
2411	vllm_config,
2412	KVConnectorRole.WORKER,
2413	_make_test_kv_cache_config(),
2414	)
2415	prefill_worker = prefill_connector.connector_worker
2416
2417	# Override TP rank/size to simulate P TP=2
2418	prefill_worker.tp_rank = P_TP_RANK
2419	prefill_worker.tp_size = P_TP_SIZE
2420	prefill_worker._tp_size[prefill_worker.engine_id] = P_TP_SIZE
2421	prefill_worker.transfer_topo.tp_rank = P_TP_RANK
2422	prefill_worker.transfer_topo.tp_size = P_TP_SIZE
2423
2424	prefill_worker.kv_caches_base_addr = [0x1000]
2425	prefill_worker.block_len_per_layer = [local_block_len]
2426	prefill_worker.kv_block_len_per_layer = [local_block_len]
2427	prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
2428	prefill_worker.registered_layer_indices = [0]
2429
2430	origin_sender_loop = prefill_worker.sender_loop
2431	prefill_worker.sender_loop = asyncio.get_event_loop()
2432
2433	transfer_id = "xfer-hetero-1"
2434	local_block_ids = [[10, 11]]
2435	send_meta = SendBlockMeta(
2436	p_req_id="p-req-h1",
2437	transfer_id=transfer_id,
2438	local_block_ids=local_block_ids,
2439	ready=asyncio.Event(),
2440	)
2441	prefill_worker.reqs_need_send[transfer_id] = send_meta
2442	send_meta.ready.set()
2443
2444	# Compute target D ranks using the production code path
2445	target_d_ranks = prefill_worker.transfer_topo.handshake_target_ranks(d_tp_size)
2446
2447	mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
2448	mock_socket.send_multipart = AsyncMock()
2449	identity = b"consumer-hetero"
2450
2451	# Assign different remote block IDs per D rank (nested per-group)
2452	d_rank_remote_blocks = {
2453	rank: [[20 + i * 10, 21 + i * 10]] for i, rank in enumerate(target_d_ranks)
2454	}
2455
2456	with patch.object(
2457	prefill_worker, "_send_blocks", return_value=0
2458	) as mock_send_blocks:
2459	for d_rank in target_d_ranks:
2460	remote_block_ids = d_rank_remote_blocks[d_rank]
2461	xfer_meta = _xfer_meta(
2462	[
2463	_region(
2464	0x2000,
2465	block_len=remote_block_len,
2466	row_offset=-1,
2467	layer_name="model.layers.0.self_attn",
2468	layer_index=0,
2469	)
2470	],
2471	{
2472	f"d-req-h1-r{d_rank}": (
2473	transfer_id,
2474	remote_block_ids,
2475	)
2476	},
2477	remote_tp_size=d_tp_size,
2478	remote_tp_rank=d_rank,
2479	)
2480
2481	mock_send_blocks.reset_mock()
2482	mock_socket.send_multipart.reset_mock()
2483
2484	await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
2485
2486	# Verify _send_blocks was called
2487	mock_send_blocks.assert_called_once()
2488	call_args = mock_send_blocks.call_args[0]
2489	src_ptrs = call_args[1]
2490	dst_ptrs = call_args[2]
2491	lengths = call_args[3]
2492
2493	# Flatten nested per-group block IDs for assertions
2494	flat_local = [b for g in local_block_ids for b in g]
2495	flat_remote = [b for g in remote_block_ids for b in g]
2496	num_blocks = len(flat_local)
2497
2498	assert len(src_ptrs) == num_blocks
2499	assert len(dst_ptrs) == num_blocks
2500	assert len(lengths) == num_blocks
2501
2502	if d_tp_size <= P_TP_SIZE:
2503	tp_ratio = P_TP_SIZE // d_tp_size
2504	expected_src_off = 0
2505	expected_dst_off = (P_TP_RANK % tp_ratio) * local_block_len
2506	expected_xfer_len = local_block_len
2507	else:
2508	ratio_abs = d_tp_size // P_TP_SIZE
2509	expected_src_off = (d_rank % ratio_abs) * remote_block_len
2510	expected_dst_off = 0
2511	expected_xfer_len = remote_block_len
2512
2513	local_region_base = 0x1000
2514	remote_region_base = 0x2000
2515	for blk_idx, (lblk, rblk) in enumerate(zip(flat_local, flat_remote)):
2516	assert src_ptrs[blk_idx] == (
2517	local_region_base + lblk * local_block_len + expected_src_off
2518	)
2519	assert dst_ptrs[blk_idx] == (
2520	remote_region_base + rblk * remote_block_len + expected_dst_off
2521	)
2522	assert lengths[blk_idx] == expected_xfer_len
2523
2524	# Verify successful response sent back to consumer
2525	mock_socket.send_multipart.assert_called_once()
2526	_, sent_payload = mock_socket.send_multipart.call_args[0][0]
2527	response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
2528	assert response.status == MooncakeXferResponseStatus.FINISH
2529	assert response.ok_reqs == [f"d-req-h1-r{d_rank}"]
2530
2531	# After serving all D ranks, the request should be complete
2532	assert transfer_id not in prefill_worker.reqs_need_send
2533	assert "p-req-h1" in prefill_worker.finished_sending_reqs
2534
2535	prefill_worker.sender_loop = origin_sender_loop
2536	prefill_worker.shutdown()