@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6287e7fa20
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 22:15 UTC

[P/D] Mooncake: Add unit tests and minor fixes for mooncake connector (#36946)

Signed-off-by: Tianchen Ding <dtcccc@linux.alibaba.com>

dtc committed Mar 27, 2026 at 16:26 UTC 6287e7fa208199da2dc2869646291a902c0d5caa
4 files changed +774 -13
tests/v1/kv_connector/unit/test_mooncake_connector.py new
+756
@@ -0,0 +1,756 @@
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 time
7 +from unittest.mock import AsyncMock, MagicMock, patch
8 +
9 +import pytest
10 +import torch
11 +import zmq.asyncio
12 +
13 +from vllm.config import set_current_vllm_config
14 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
15 + KVConnectorRole,
16 + MooncakeConnector,
17 + MooncakeConnectorMetadata,
18 + MooncakeXferMetadata,
19 + MooncakeXferResponse,
20 + MooncakeXferResponseStatus,
21 + PullReqMeta,
22 + SendBlockMeta,
23 +)
24 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
25 + MooncakeBootstrapServer,
26 +)
27 +from vllm.utils.network_utils import get_open_port
28 +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
29 +from vllm.v1.request import RequestStatus
30 +
31 +from .utils import create_request, create_scheduler, create_vllm_config
32 +
33 +
34 +class FakeMooncakeWrapper:
35 + """Mock Mooncake TransferEngine for unit testing environments."""
36 +
37 + def __init__(self, *args, **kwargs):
38 + pass
39 +
40 + def initialize(self, local_hostname, metadata_server, protocol, device_name) -> int:
41 + return 0
42 +
43 + def get_rpc_port(self) -> int:
44 + return 12345
45 +
46 + def batch_transfer_sync_write(
47 + self, target_hostname, buffers, peer_buffer_addresses, lengths
48 + ) -> int:
49 + return 0
50 +
51 + def batch_register_memory(self, buffer_addresses, capacities) -> int:
52 + return 0
53 +
54 +
55 +def test_basic_interface():
56 + """Unit test for basic MooncakeConnector interface functionality."""
57 +
58 + vllm_config = create_vllm_config(
59 + kv_connector="MooncakeConnector", kv_role="kv_consumer"
60 + )
61 + scheduler = create_scheduler(vllm_config)
62 +
63 + # 2 Full Blocks and 1 Half Block.
64 + BLOCK_SIZE = vllm_config.cache_config.block_size
65 + NUM_EXTERNAL_FULL_BLOCKS = 2
66 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
67 +
68 + request = create_request(
69 + request_id=1,
70 + block_size=BLOCK_SIZE,
71 + num_tokens=NUM_TOKENS,
72 + do_remote_prefill=True,
73 + )
74 + request_id = request.request_id
75 + request.kv_transfer_params.update(
76 + {
77 + "transfer_id": request_id,
78 + "remote_bootstrap_addr": 54321,
79 + }
80 + )
81 +
82 + scheduler.add_request(request)
83 +
84 + # Remote Prefill, triggers NixlConnectorMetadata.
85 + scheduler_output = scheduler.schedule()
86 + kv_connector_metadata = scheduler_output.kv_connector_metadata
87 + assert kv_connector_metadata is not None
88 + assert isinstance(kv_connector_metadata, MooncakeConnectorMetadata)
89 +
90 + assert len(kv_connector_metadata.reqs_to_recv) == 1
91 + assert request_id in kv_connector_metadata.reqs_to_recv["my-engine-id"]
92 + req_meta = kv_connector_metadata.reqs_to_recv["my-engine-id"][request_id]
93 +
94 + for block_id, block in zip(
95 + req_meta.local_block_ids,
96 + scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
97 + request_id
98 + ],
99 + ):
100 + assert block_id == block.block_id
101 +
102 +
103 +def test_prompt_less_than_block_size():
104 + """Test that we can handle case where prompt is < block."""
105 +
106 + vllm_config = create_vllm_config(
107 + kv_connector="MooncakeConnector", kv_role="kv_consumer"
108 + )
109 + scheduler = create_scheduler(vllm_config)
110 +
111 + # Half of a block.
112 + BLOCK_SIZE = vllm_config.cache_config.block_size
113 + NUM_TOKENS = int(BLOCK_SIZE * 0.5)
114 +
115 + # Request will have 1 partial remote block.
116 + request = create_request(
117 + request_id=1,
118 + block_size=BLOCK_SIZE,
119 + num_tokens=NUM_TOKENS,
120 + do_remote_prefill=True,
121 + num_remote_blocks=1,
122 + )
123 + request.kv_transfer_params.update(
124 + {
125 + "transfer_id": request.request_id,
126 + "remote_bootstrap_addr": 54321,
127 + }
128 + )
129 +
130 + scheduler.add_request(request)
131 + scheduler_output = scheduler.schedule()
132 +
133 + # This request will read async.
134 + kv_connector_metadata = scheduler_output.kv_connector_metadata
135 + assert kv_connector_metadata is not None
136 + assert isinstance(kv_connector_metadata, MooncakeConnectorMetadata)
137 + assert len(kv_connector_metadata.reqs_to_recv["my-engine-id"]) == 1
138 + assert len(scheduler_output.scheduled_new_reqs) == 0
139 +
140 +
141 +@pytest.fixture
142 +def bootstrap_server():
143 + """Fixture to launch and cleanup a Mooncake Bootstrap HTTP Server."""
144 +
145 + port = get_open_port()
146 + server = MooncakeBootstrapServer("127.0.0.1", port)
147 + server.start()
148 + yield server
149 + server.shutdown()
150 +
151 +
152 +@pytest.mark.asyncio
153 +async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):
154 + """
155 + Tests the bootstrap server's api for worker registration and querying.
156 +
157 + Validates DP/TP/PP rank indexing and error handling for duplicate registrations.
158 + """
159 +
160 + import httpx
161 +
162 + base_url = f"http://127.0.0.1:{bootstrap_server.port}"
163 +
164 + # Query when empty
165 + async with httpx.AsyncClient() as client:
166 + response = await client.get(f"{base_url}/query")
167 + assert response.status_code == 200
168 + assert response.json() == {}
169 +
170 + # Register a worker
171 + payload1 = {
172 + "engine_id": "eng-1",
173 + "dp_rank": 0,
174 + "tp_rank": 0,
175 + "pp_rank": 0,
176 + "addr": "tcp://1.1.1.1:1111",
177 + }
178 + async with httpx.AsyncClient() as client:
179 + response = await client.post(f"{base_url}/register", json=payload1)
180 + assert response.status_code == 200
181 + assert response.json() == {"status": "ok"}
182 +
183 + # Query after registration
184 + async with httpx.AsyncClient() as client:
185 + response = await client.get(f"{base_url}/query")
186 + assert response.status_code == 200
187 + data = response.json()
188 + assert "0" in data
189 + assert data["0"]["engine_id"] == "eng-1"
190 + assert data["0"]["worker_addr"]["0"]["0"] == "tcp://1.1.1.1:1111"
191 +
192 + # Test failure: re-registering the same worker
193 + async with httpx.AsyncClient() as client:
194 + response = await client.post(f"{base_url}/register", json=payload1)
195 + assert response.status_code == 400
196 + assert "is already registered" in response.text
197 +
198 + # Test failure: engine_id mismatch for same dp_rank
199 + payload3_fail = {
200 + "engine_id": "eng-2",
201 + "dp_rank": 0,
202 + "tp_rank": 1,
203 + "pp_rank": 0,
204 + "addr": "tcp://3.3.3.3:3333",
205 + }
206 + async with httpx.AsyncClient() as client:
207 + response = await client.post(f"{base_url}/register", json=payload3_fail)
208 + assert response.status_code == 400
209 + assert "Engine ID mismatch" in response.text
210 +
211 +
212 +def test_scheduler_request_finished():
213 + """
214 + Tests the scheduler-side logic when a request finishes.
215 +
216 + Differentiates between 'Finished' (requires transfer)
217 + and 'Aborted' (immediate free).
218 + """
219 +
220 + vllm_config = create_vllm_config(
221 + kv_connector="MooncakeConnector", kv_role="kv_producer"
222 + )
223 + scheduler = create_scheduler(vllm_config)
224 + scheduler_connector = scheduler.get_kv_connector().connector_scheduler
225 +
226 + request = create_request(request_id=1, do_remote_decode=True)
227 + request.kv_transfer_params["transfer_id"] = request.request_id
228 +
229 + # Case: Capped length (Successful prefill, need to send to decoder)
230 + request.status = RequestStatus.FINISHED_LENGTH_CAPPED
231 + delay_free, _ = scheduler_connector.request_finished(request, block_ids=[10, 11])
232 + assert delay_free is True
233 + assert "id-1" in scheduler_connector._reqs_need_send
234 + assert scheduler_connector._reqs_need_send["id-1"][1] == [10, 11]
235 +
236 + # Case: Aborted (No need to transfer, free blocks immediately)
237 + scheduler_connector._reqs_need_send.clear()
238 + request.status = RequestStatus.FINISHED_ABORTED
239 + delay_free, _ = scheduler_connector.request_finished(request, block_ids=[12])
240 + assert delay_free is False
241 + assert len(scheduler_connector._reqs_need_send) == 0
242 + assert "id-1" in scheduler_connector._reqs_not_processed
243 +
244 +
245 +@contextlib.contextmanager
246 +def patch_worker_dependencies():
247 + """Helper to mock all distributed and network dependencies for Worker tests."""
248 +
249 + with (
250 + patch(
251 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
252 + FakeMooncakeWrapper,
253 + ),
254 + patch(
255 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_ip",
256 + return_value="127.0.0.1",
257 + ),
258 + patch(
259 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_tensor_model_parallel_rank",
260 + return_value=0,
261 + ),
262 + patch(
263 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_tensor_model_parallel_world_size",
264 + return_value=1,
265 + ),
266 + patch(
267 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_pp_group"
268 + ) as mock_pp,
269 + patch("vllm.distributed.parallel_state.is_local_first_rank", return_value=True),
270 + patch(
271 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.should_launch_bootstrap_server",
272 + return_value=False,
273 + ),
274 + patch(
275 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.make_zmq_socket"
276 + ) as mock_make_zmq,
277 + patch("httpx.AsyncClient") as mock_async_client,
278 + ):
279 + # Mock PP group
280 + mock_pp_group = MagicMock()
281 + mock_pp_group.rank_in_group = 0
282 + mock_pp.return_value = mock_pp_group
283 +
284 + # Mock ZMQ socket
285 + mock_socket_object = AsyncMock()
286 + mock_socket_object.setsockopt = MagicMock()
287 + mock_socket_ctx = MagicMock()
288 + mock_socket_ctx.__enter__.return_value = mock_socket_object
289 + mock_make_zmq.return_value = mock_socket_ctx
290 +
291 + # Mock httpx client
292 + mock_http_client_instance = AsyncMock()
293 + mock_async_client.return_value = mock_http_client_instance
294 +
295 + yield {
296 + "mock_make_zmq": mock_make_zmq,
297 + "mock_socket_object": mock_socket_object,
298 + "mock_async_client": mock_async_client,
299 + "mock_http_client": mock_http_client_instance,
300 + }
301 +
302 +
303 +@pytest.mark.asyncio
304 +@patch(
305 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
306 + FakeMooncakeWrapper,
307 +)
308 +async def test_kv_producer(monkeypatch):
309 + """
310 + Simulates a Producer Worker (Prefiller) receiving a transfer request
311 + from a Consumer (Decoder).
312 +
313 + Verifies memory offset calculation: ptr = base_addr + block_id * block_len.
314 + """
315 +
316 + monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
317 + vllm_config = create_vllm_config(
318 + kv_connector="MooncakeConnector", kv_role="kv_producer"
319 + )
320 +
321 + with set_current_vllm_config(vllm_config), patch_worker_dependencies():
322 + prefill_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
323 + prefill_worker = prefill_connector.connector_worker
324 + prefill_worker.kv_caches_base_addr = [0x1000]
325 + block_len = 4096
326 + prefill_worker.block_len_per_layer = [block_len]
327 +
328 + # Override loop to use current test loop
329 + origin_sender_loop = prefill_worker.sender_loop
330 + prefill_worker.sender_loop = asyncio.get_event_loop()
331 +
332 + # A request is finished on Producer and ready to be sent.
333 + transfer_id = "xfer-req-1"
334 + send_meta = SendBlockMeta(
335 + p_req_id="p-req-1",
336 + transfer_id=transfer_id,
337 + local_block_ids=[10, 11],
338 + ready=asyncio.Event(),
339 + )
340 + prefill_worker.reqs_need_send[transfer_id] = send_meta
341 + send_meta.ready.set()
342 +
343 + # Remote consumer request metadata
344 + xfer_meta = MooncakeXferMetadata(
345 + remote_hostname="consumer-host",
346 + remote_port=54321,
347 + remote_tp_size=1,
348 + remote_tp_rank=0,
349 + req_blocks={"d-req-1": (transfer_id, [20, 21])},
350 + kv_caches_base_addr=[0x2000],
351 + block_lens=[block_len],
352 + )
353 +
354 + mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
355 + mock_socket.send_multipart = AsyncMock()
356 + identity = b"consumer-id"
357 +
358 + with patch.object(
359 + prefill_worker, "_send_blocks", return_value=0
360 + ) as mock_send_blocks:
361 + # Normal case: 2 blocks to 2 blocks
362 + # Worker processes the consumer's request
363 + await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
364 + # Verify transfer parameters are correct
365 + src_ptr = 0x1000 + 10 * block_len
366 + dst_ptr = 0x2000 + 20 * block_len
367 + length = 2 * block_len
368 + mock_send_blocks.assert_called_once_with(
369 + "consumer-host:54321", [src_ptr], [dst_ptr], [length]
370 + )
371 + mock_socket.send_multipart.assert_called_once()
372 +
373 + # Verify the response sent back to the consumer
374 + sent_call = mock_socket.send_multipart.call_args[0][0]
375 + sent_identity, sent_payload = sent_call
376 + assert sent_identity == identity
377 + response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
378 + assert response.status == MooncakeXferResponseStatus.FINISH
379 + assert response.ok_reqs == ["d-req-1"]
380 +
381 + # Verify internal state cleanup
382 + assert transfer_id not in prefill_worker.reqs_need_send
383 + assert "p-req-1" in prefill_worker.finished_sending_reqs
384 +
385 + # More cases:
386 + # Consumer only needs 1 block (less than P)
387 + mock_send_blocks.reset_mock()
388 + mock_socket.send_multipart.reset_mock()
389 + prefill_worker.reqs_need_send[transfer_id] = send_meta
390 + send_meta.sent = 0
391 + send_meta.ready.set()
392 + xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20])
393 + # Worker processes the consumer's request
394 + await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
395 + # Verify transfer parameters are correct: 11 to 20
396 + src_ptr = 0x1000 + 11 * block_len
397 + dst_ptr = 0x2000 + 20 * block_len
398 + length = 1 * block_len
399 + mock_send_blocks.assert_called_once_with(
400 + "consumer-host:54321", [src_ptr], [dst_ptr], [length]
401 + )
402 + mock_socket.send_multipart.assert_called_once()
403 +
404 + # Consumer needs 3 blocks (more than P, error case)
405 + mock_send_blocks.reset_mock()
406 + mock_socket.send_multipart.reset_mock()
407 + prefill_worker.reqs_need_send[transfer_id] = send_meta
408 + send_meta.sent = 0
409 + send_meta.ready.set()
410 + xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21, 22])
411 + # Worker processes the consumer's request
412 + await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
413 + # This should not be called because error.
414 + mock_send_blocks.assert_not_called()
415 + mock_socket.send_multipart.assert_called_once()
416 + _, sent_payload = mock_socket.send_multipart.call_args[0][0]
417 + response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
418 + assert response.err_msg == "P num blocks less than D"
419 + assert response.err_reqs == ["d-req-1"]
420 +
421 + # Timeout
422 + mock_send_blocks.reset_mock()
423 + mock_socket.send_multipart.reset_mock()
424 + prefill_worker.reqs_need_send[transfer_id] = send_meta
425 + send_meta.sent = 0
426 + send_meta.ready.clear()
427 + xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21])
428 + # Worker processes the consumer's request
429 + await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
430 + # This should not be called because timeout.
431 + mock_send_blocks.assert_not_called()
432 + mock_socket.send_multipart.assert_called_once()
433 + _, sent_payload = mock_socket.send_multipart.call_args[0][0]
434 + response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
435 + assert response.err_msg == "Timeout waiting for P side ready."
436 + assert response.err_reqs == ["d-req-1"]
437 +
438 + # Transfer error
439 + with patch.object(
440 + prefill_worker, "_send_blocks", return_value=123
441 + ) as mock_send_blocks:
442 + mock_socket.send_multipart.reset_mock()
443 + prefill_worker.reqs_need_send[transfer_id] = send_meta
444 + send_meta.sent = 0
445 + send_meta.ready.set()
446 + xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21])
447 + # Worker processes the consumer's request
448 + await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
449 + mock_send_blocks.assert_called_once()
450 + mock_socket.send_multipart.assert_called_once()
451 + _, sent_payload = mock_socket.send_multipart.call_args[0][0]
452 + response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
453 + assert response.err_msg == "Mooncake transfer engine returned 123"
454 + assert response.err_reqs == ["d-req-1"]
455 +
456 + # Clean up
457 + prefill_worker.sender_loop = origin_sender_loop
458 + prefill_worker.shutdown()
459 +
460 +
461 +@pytest.mark.asyncio
462 +async def test_kv_consumuer(monkeypatch):
463 + """
464 + Simulates a Consumer Worker (Decoder) initiating a pull from a Producer.
465 +
466 + Verifies that MooncakeXferMetadata is correctly serialized and sent via ZMQ.
467 + """
468 +
469 + vllm_config = create_vllm_config(
470 + kv_connector="MooncakeConnector", kv_role="kv_consumer"
471 + )
472 +
473 + with set_current_vllm_config(vllm_config), patch_worker_dependencies() as mocks:
474 + decode_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
475 + decode_worker = decode_connector.connector_worker
476 + decode_worker.kv_caches_base_addr = [0x1000]
477 + decode_worker.rpc_port = 54321
478 +
479 + # A request to pull data arrives.
480 + pull_metas = {
481 + "d-req-1": PullReqMeta(
482 + d_req_id="d-req-1",
483 + transfer_id="xfer-req-1",
484 + local_block_ids=[100, 101],
485 + remote_engine_id="p-engine",
486 + remote_bootstrap_addr="http://bootstrap:33333",
487 + pull_tasks_count=1,
488 + )
489 + }
490 + decode_worker._remote_agents = {"p-engine": {0: {0: "tcp://producer:1234"}}}
491 + decode_worker._tp_size["p-engine"] = 1
492 +
493 + # Mock the response from the producer.
494 + mock_response = MooncakeXferResponse(
495 + status=MooncakeXferResponseStatus.FINISH, ok_reqs=["d-req-1"]
496 + )
497 + encoded_response = decode_worker._encoder.encode(mock_response)
498 + mocks["mock_socket_object"].recv.return_value = encoded_response
499 +
500 + # Trigger the receive logic.
501 + decode_worker.receive_kv("p-engine", pull_metas)
502 + await asyncio.sleep(1) # Allow async task to run
503 +
504 + # Verify the metadata sent to the producer.
505 + mocks["mock_make_zmq"].assert_called_with(
506 + decode_worker.async_zmq_ctx,
507 + "tcp://producer:1234",
508 + zmq.DEALER,
509 + bind=False,
510 + linger=0,
511 + )
512 + sent_payload = mocks["mock_socket_object"].send.call_args[0][0]
513 + sent_meta = decode_worker._xfer_meta_decoder.decode(sent_payload)
514 +
515 + assert sent_meta.remote_hostname == "127.0.0.1"
516 + assert sent_meta.remote_port == 54321
517 + assert sent_meta.req_blocks["d-req-1"] == ("xfer-req-1", [100, 101])
518 +
519 + # Verify internal state is updated correctly.
520 + assert "d-req-1" in decode_worker.finished_recving_reqs
521 +
522 + # Clean up
523 + decode_worker.shutdown()
524 +
525 +
526 +@pytest.mark.asyncio
527 +async def test_worker_get_finished_timeout(monkeypatch):
528 + """Tests the cleanup mechanism for requests."""
529 +
530 + vllm_config = create_vllm_config(
531 + kv_connector="MooncakeConnector", kv_role="kv_producer"
532 + )
533 + with set_current_vllm_config(vllm_config), patch_worker_dependencies():
534 + prefill_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
535 + prefill_worker = prefill_connector.connector_worker
536 +
537 + # Add an expired request (expire_time is in the past).
538 + prefill_worker.reqs_need_send["tx-expired"] = SendBlockMeta(
539 + p_req_id="p-req-expired",
540 + transfer_id="tx-expired",
541 + local_block_ids=[1, 2],
542 + ready=MagicMock(),
543 + expire_time=time.perf_counter() - 100,
544 + )
545 +
546 + # Add a non-expired request.
547 + prefill_worker.reqs_need_send["tx-active"] = SendBlockMeta(
548 + p_req_id="p-req-active",
549 + transfer_id="tx-active",
550 + local_block_ids=[3, 4],
551 + ready=MagicMock(),
552 + expire_time=time.perf_counter() + 100,
553 + )
554 +
555 + finished_reqs = await prefill_worker.fetch_finished_sending_reqs()
556 +
557 + assert "p-req-expired" in finished_reqs
558 + assert "p-req-active" not in finished_reqs
559 + assert "tx-expired" not in prefill_worker.reqs_need_send
560 + assert "tx-active" in prefill_worker.reqs_need_send
561 +
562 +
563 +def test_register_kv_caches():
564 + """Tests the memory registration logic with the underlying Mooncake engine."""
565 +
566 + vllm_config = create_vllm_config(
567 + kv_connector="MooncakeConnector", kv_role="kv_consumer"
568 + )
569 +
570 + with (
571 + set_current_vllm_config(vllm_config),
572 + patch_worker_dependencies(),
573 + patch(
574 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Event"
575 + ),
576 + patch(
577 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Thread"
578 + ) as mock_thread,
579 + ):
580 + connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
581 + worker = connector.connector_worker
582 + mock_thread.return_value.is_alive.return_value = False
583 +
584 + kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
585 + num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
586 + )
587 + tensor1 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
588 + tensor2 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
589 + kv_caches = {"layer0": tensor1, "layer1": tensor2}
590 +
591 + with patch.object(
592 + worker.engine, "batch_register_memory", return_value=0
593 + ) as mock_batch_register:
594 + connector.register_kv_caches(kv_caches)
595 +
596 + mock_batch_register.assert_called_once()
597 + registered_ptrs, registered_lens = mock_batch_register.call_args[0]
598 + expected_ptrs = {
599 + tensor.data_ptr()
600 + for kv_pair in kv_caches.values()
601 + for tensor in kv_pair
602 + }
603 + assert set(registered_ptrs) == expected_ptrs
604 + assert set(registered_lens) == {tensor1[0].nbytes}
605 +
606 + # Verify block_len_per_layer is set correctly.
607 + assert len(worker.block_len_per_layer) == len(registered_ptrs)
608 + for bl in worker.block_len_per_layer:
609 + assert bl == tensor1[0].nbytes // tensor1.shape[1]
610 +
611 +
612 +@pytest.mark.asyncio
613 +@patch(
614 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
615 + "mooncake_connector.TransferEngine",
616 + FakeMooncakeWrapper,
617 +)
618 +@pytest.mark.parametrize("d_tp_size", [1, 4], ids=["p_tp2_d_tp1", "p_tp2_d_tp4"])
619 +async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
620 + """
621 + Tests heterogeneous TP support in the producer transfer path.
622 +
623 + Verifies correct pointer and offset calculation when producer TP=2
624 + sends to consumer with TP=1 (P>D) or TP=4 (P<D).
625 +
626 + Parametrized cases:
627 + - P TP=2 > D TP=1: one D rank receives; dst_offset based on P rank
628 + - P TP=2 < D TP=4: two D ranks receive; src_offset based on D rank
629 + """
630 +
631 + P_TP_SIZE = 2
632 + P_TP_RANK = 0
633 + LOCAL_BLOCK_LEN = 4096
634 +
635 + local_block_len = LOCAL_BLOCK_LEN
636 + remote_block_len = LOCAL_BLOCK_LEN * P_TP_SIZE // d_tp_size
637 +
638 + monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
639 + vllm_config = create_vllm_config(
640 + kv_connector="MooncakeConnector", kv_role="kv_producer"
641 + )
642 +
643 + with set_current_vllm_config(vllm_config), patch_worker_dependencies():
644 + prefill_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
645 + prefill_worker = prefill_connector.connector_worker
646 +
647 + # Override TP rank/size to simulate P TP=2
648 + prefill_worker.tp_rank = P_TP_RANK
649 + prefill_worker.tp_size = P_TP_SIZE
650 + # Update shared dict so kv_topo sees correct TP size
651 + prefill_worker._tp_size[prefill_worker.engine_id] = P_TP_SIZE
652 + prefill_worker.kv_topo.tp_rank = P_TP_RANK
653 +
654 + prefill_worker.kv_caches_base_addr = [0x1000]
655 + prefill_worker.block_len_per_layer = [local_block_len]
656 +
657 + origin_sender_loop = prefill_worker.sender_loop
658 + prefill_worker.sender_loop = asyncio.get_event_loop()
659 +
660 + transfer_id = "xfer-hetero-1"
661 + local_block_ids = [10, 11]
662 + send_meta = SendBlockMeta(
663 + p_req_id="p-req-h1",
664 + transfer_id=transfer_id,
665 + local_block_ids=local_block_ids,
666 + ready=asyncio.Event(),
667 + )
668 + prefill_worker.reqs_need_send[transfer_id] = send_meta
669 + send_meta.ready.set()
670 +
671 + # Compute target D ranks using the production code path
672 + target_d_ranks = prefill_worker.kv_topo.get_target_remote_ranks(d_tp_size)
673 +
674 + mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
675 + mock_socket.send_multipart = AsyncMock()
676 + identity = b"consumer-hetero"
677 +
678 + # Assign different remote block IDs per D rank
679 + d_rank_remote_blocks = {
680 + rank: [20 + i * 10, 21 + i * 10] for i, rank in enumerate(target_d_ranks)
681 + }
682 +
683 + with patch.object(
684 + prefill_worker, "_send_blocks", return_value=0
685 + ) as mock_send_blocks:
686 + for d_rank in target_d_ranks:
687 + remote_block_ids = d_rank_remote_blocks[d_rank]
688 + xfer_meta = MooncakeXferMetadata(
689 + remote_hostname="consumer-host",
690 + remote_port=54321,
691 + remote_tp_size=d_tp_size,
692 + remote_tp_rank=d_rank,
693 + req_blocks={
694 + f"d-req-h1-r{d_rank}": (
695 + transfer_id,
696 + remote_block_ids,
697 + )
698 + },
699 + kv_caches_base_addr=[0x2000],
700 + block_lens=[remote_block_len],
701 + )
702 +
703 + mock_send_blocks.reset_mock()
704 + mock_socket.send_multipart.reset_mock()
705 +
706 + await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
707 +
708 + # Verify _send_blocks was called
709 + mock_send_blocks.assert_called_once()
710 + call_args = mock_send_blocks.call_args[0]
711 + src_ptrs = call_args[1]
712 + dst_ptrs = call_args[2]
713 + lengths = call_args[3]
714 +
715 + # Heterogeneous TP: blocks cannot be coalesced because
716 + # local and remote block_lens differ
717 + assert len(src_ptrs) == len(local_block_ids)
718 + assert len(dst_ptrs) == len(local_block_ids)
719 + assert len(lengths) == len(local_block_ids)
720 +
721 + # Compute expected offsets based on TP ratio
722 + if d_tp_size <= P_TP_SIZE:
723 + tp_ratio = P_TP_SIZE // d_tp_size
724 + expected_src_off = 0
725 + expected_dst_off = (P_TP_RANK % tp_ratio) * local_block_len
726 + expected_xfer_len = local_block_len
727 + else:
728 + ratio_abs = d_tp_size // P_TP_SIZE
729 + expected_src_off = (d_rank % ratio_abs) * remote_block_len
730 + expected_dst_off = 0
731 + expected_xfer_len = remote_block_len
732 +
733 + for idx, (lblk, rblk) in enumerate(
734 + zip(local_block_ids, remote_block_ids)
735 + ):
736 + assert src_ptrs[idx] == (
737 + 0x1000 + lblk * local_block_len + expected_src_off
738 + )
739 + assert dst_ptrs[idx] == (
740 + 0x2000 + rblk * remote_block_len + expected_dst_off
741 + )
742 + assert lengths[idx] == expected_xfer_len
743 +
744 + # Verify successful response sent back to consumer
745 + mock_socket.send_multipart.assert_called_once()
746 + _, sent_payload = mock_socket.send_multipart.call_args[0][0]
747 + response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
748 + assert response.status == MooncakeXferResponseStatus.FINISH
749 + assert response.ok_reqs == [f"d-req-h1-r{d_rank}"]
750 +
751 + # After serving all D ranks, the request should be complete
752 + assert transfer_id not in prefill_worker.reqs_need_send
753 + assert "p-req-h1" in prefill_worker.finished_sending_reqs
754 +
755 + prefill_worker.sender_loop = origin_sender_loop
756 + prefill_worker.shutdown()
tests/v1/kv_connector/unit/utils.py
+4 -2
@@ -100,6 +100,8 @@ def create_vllm_config(
100 hf_overrides: dict[str, Any] | None = None,
101 attention_backend: str | None = None,
102 kv_load_failure_policy: Literal["recompute", "fail"] = "fail",
103 + kv_connector: str = "NixlConnector",
104 + kv_role: str = "kv_both",
105 ) -> VllmConfig:
106 """Initialize VllmConfig For Testing."""
107 model_config = ModelConfig(
@@ -124,8 +126,8 @@ def create_vllm_config(
126 enable_prefix_caching=True,
127 )
128 kv_transfer_config = KVTransferConfig(
127 - kv_connector="NixlConnector",
128 - kv_role="kv_both",
129 + kv_connector=kv_connector,
130 + kv_role=kv_role,
131 enable_permute_local_kv=enable_permute_local_kv,
132 kv_connector_extra_config=kv_connector_extra_config or {},
133 kv_load_failure_policy=kv_load_failure_policy,
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+13 -9
@@ -47,14 +47,17 @@ from vllm.v1.attention.backends.utils import get_kv_cache_layout
47 from vllm.v1.core.sched.output import SchedulerOutput
48 from vllm.v1.request import RequestStatus
49
50 +logger = init_logger(__name__)
51 +
52 try:
53 from mooncake.engine import TransferEngine
52 -except ImportError as e:
53 - raise ImportError(
54 +except ImportError:
55 + logger.warning(
56 "Please install mooncake by following the instructions at "
57 "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md "
58 "to run VLLM with MooncakeTransferEngine."
57 - ) from e
59 + )
60 + TransferEngine = None
61
62 if TYPE_CHECKING:
63 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
@@ -64,8 +67,6 @@ if TYPE_CHECKING:
67 ReqId = str # Internal scheduler request ID
68 TransferId = str # KV transfer coordination ID (shared by P/D)
69
67 -logger = init_logger(__name__)
68 -
70
71 @dataclass(frozen=True)
72 class TransferRegion:
@@ -638,6 +639,9 @@ class MooncakeConnectorWorker:
639 """Implementation of Worker side methods"""
640
641 def __init__(self, vllm_config: VllmConfig, engine_id: str):
642 + if TransferEngine is None:
643 + logger.error("Mooncake is not available")
644 + raise RuntimeError("Mooncake is not available")
645 logger.info("Initializing Mooncake Transfer Engine worker %s", engine_id)
646
647 self.vllm_config = vllm_config
@@ -721,9 +725,7 @@ class MooncakeConnectorWorker:
725 # Start bootstrap server on global rank 0.
726 if should_launch_bootstrap_server(vllm_config):
727 _, port = get_mooncake_bootstrap_addr(vllm_config)
724 - self.bootstrap_server = MooncakeBootstrapServer(
725 - vllm_config, "0.0.0.0", port
726 - )
728 + self.bootstrap_server = MooncakeBootstrapServer("0.0.0.0", port)
729 self.bootstrap_server.start()
730
731 if not self.is_kv_producer:
@@ -778,7 +780,9 @@ class MooncakeConnectorWorker:
780 if self.sender_loop.is_running():
781 self.sender_loop.call_soon_threadsafe(self.sender_loop.stop)
782 self._sender_listener_t.join()
781 - if should_launch_bootstrap_server(self.vllm_config):
783 + if should_launch_bootstrap_server(self.vllm_config) and hasattr(
784 + self, "bootstrap_server"
785 + ):
786 self.bootstrap_server.shutdown()
787 if not self.is_kv_producer and self.receiver_loop.is_running():
788 self.receiver_loop.call_soon_threadsafe(self.receiver_loop.stop)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py
+1 -2
@@ -8,7 +8,6 @@ import uvicorn
8 from fastapi import FastAPI, HTTPException
9 from pydantic import BaseModel
10
11 -from vllm.config import VllmConfig
11 from vllm.distributed.kv_transfer.kv_connector.utils import EngineId
12 from vllm.logger import init_logger
13
@@ -38,7 +37,7 @@ class MooncakeBootstrapServer:
37 Prefiller workers register their connection info (IP, port, ranks) here.
38 """
39
41 - def __init__(self, vllm_config: VllmConfig, host: str, port: int):
40 + def __init__(self, host: str, port: int):
41 self.workers: dict[int, EngineEntry] = {}
42
43 self.host = host