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 4, 06:15 UTC
[P/D][Mooncake] Add KVConnectorStats for transfer observability (#40414)
Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai>
Zhewen Li committed
May 5, 2026 at 02:17 UTC
bee126165fb39e770030e9c3b4a43aee4acd3175
3 files changed
+477
-3
tests/v1/kv_connector/unit/test_mooncake_stats.py
new
+281
@@ -0,0 +1,281 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import threading
5
+from unittest.mock import MagicMock
6
+
7
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
8
+ MooncakeConnector,
9
+ MooncakeConnectorWorker,
10
+ SendBlockMeta,
11
+)
12
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.stats import (
13
+ MooncakeKVConnectorStats,
14
+)
15
+
16
+
17
+def test_is_empty_on_fresh_stats():
18
+ stats = MooncakeKVConnectorStats()
19
+ assert stats.is_empty()
20
+ assert stats.num_successful_transfers == 0
21
+
22
+
23
+def test_record_transfer_and_reduce():
24
+ stats = MooncakeKVConnectorStats()
25
+ # 1 MB transfer in 1 ms -> 1000 MB/s throughput
26
+ stats.record_transfer(duration_s=0.001, total_bytes=1 * 2**20, num_descs=4)
27
+ # 2 MB transfer in 2 ms
28
+ stats.record_transfer(duration_s=0.002, total_bytes=2 * 2**20, num_descs=6)
29
+ assert not stats.is_empty()
30
+ assert stats.num_successful_transfers == 2
31
+
32
+ reduced = stats.reduce()
33
+ assert reduced["Num successful transfers"] == 2
34
+ # avg = (1 + 2) / 2 = 1.5 ms
35
+ assert reduced["Avg xfer time (ms)"] == 1.5
36
+ assert reduced["Avg MB per transfer"] == 1.5
37
+ # 3 MB total / 3 ms total = 1000 MB/s
38
+ assert reduced["Throughput (MB/s)"] == 1000.0
39
+ assert reduced["Avg number of descriptors"] == 5.0
40
+ assert reduced["Num failed transfers"] == 0
41
+ assert reduced["Num failed recvs"] == 0
42
+ assert reduced["Num KV expired reqs"] == 0
43
+
44
+
45
+def test_record_failures_keeps_stats_non_empty():
46
+ stats = MooncakeKVConnectorStats()
47
+ stats.record_failed_transfer()
48
+ stats.record_failed_recv()
49
+ stats.record_kv_expired_req()
50
+ assert not stats.is_empty()
51
+
52
+ reduced = stats.reduce()
53
+ # No successful transfers -> latency/throughput all zero, but failure
54
+ # counters still surface.
55
+ assert reduced["Num successful transfers"] == 0
56
+ assert reduced["Num failed transfers"] == 1
57
+ assert reduced["Num failed recvs"] == 1
58
+ assert reduced["Num KV expired reqs"] == 1
59
+
60
+
61
+def test_aggregate_sums_observations():
62
+ a = MooncakeKVConnectorStats()
63
+ b = MooncakeKVConnectorStats()
64
+ a.record_transfer(duration_s=0.001, total_bytes=1 * 2**20, num_descs=1)
65
+ b.record_transfer(duration_s=0.002, total_bytes=2 * 2**20, num_descs=2)
66
+ b.record_failed_transfer()
67
+
68
+ a.aggregate(b)
69
+
70
+ assert a.num_successful_transfers == 2
71
+ reduced = a.reduce()
72
+ assert reduced["Num successful transfers"] == 2
73
+ assert reduced["Num failed transfers"] == 1
74
+
75
+
76
+def test_aggregate_with_empty_other_is_noop():
77
+ a = MooncakeKVConnectorStats()
78
+ a.record_transfer(duration_s=0.001, total_bytes=1, num_descs=1)
79
+ b = MooncakeKVConnectorStats()
80
+
81
+ a.aggregate(b)
82
+
83
+ assert a.num_successful_transfers == 1
84
+
85
+
86
+def test_getstate_drops_lock_and_setstate_recreates_it():
87
+ # KVConnectorStats subclasses must be picklable (worker→scheduler IPC),
88
+ # but threading.Lock isn't — so __getstate__ strips it and __setstate__
89
+ # rebuilds a fresh per-process lock.
90
+ original = MooncakeKVConnectorStats()
91
+ original.record_transfer(duration_s=0.01, total_bytes=2048, num_descs=3)
92
+
93
+ state = original.__getstate__()
94
+ assert "_lock" not in state
95
+
96
+ rebuilt = MooncakeKVConnectorStats.__new__(MooncakeKVConnectorStats)
97
+ rebuilt.__setstate__(state)
98
+ assert rebuilt.data == original.data
99
+ # Lock works on the receiver side.
100
+ rebuilt.record_transfer(duration_s=0.02, total_bytes=4096, num_descs=5)
101
+ assert rebuilt.num_successful_transfers == 2
102
+
103
+
104
+def test_concurrent_writers_keep_row_lengths_aligned():
105
+ # Multiple writers + a snapshot reader must never produce a snapshot
106
+ # with mismatched column lengths — reduce()'s
107
+ # len(descs) == num_successful_transfers assertion would fire.
108
+ stats = MooncakeKVConnectorStats()
109
+ stop = threading.Event()
110
+ writer_count = 4
111
+ snapshots: list[MooncakeKVConnectorStats] = []
112
+
113
+ def writer():
114
+ i = 0
115
+ while not stop.is_set():
116
+ stats.record_transfer(
117
+ duration_s=0.001 + i * 1e-9,
118
+ total_bytes=1024 + i,
119
+ num_descs=1 + (i % 8),
120
+ )
121
+ i += 1
122
+
123
+ def snapper():
124
+ while not stop.is_set():
125
+ snap = stats.clone_and_reset()
126
+ if not snap.is_empty():
127
+ # Force the same path the logger walks; reduce() will
128
+ # blow up on torn rows via its internal assert.
129
+ snap.reduce()
130
+ snapshots.append(snap)
131
+
132
+ threads = [threading.Thread(target=writer) for _ in range(writer_count)]
133
+ snapshotter = threading.Thread(target=snapper)
134
+ for t in threads:
135
+ t.start()
136
+ snapshotter.start()
137
+ # Short fixed window — long enough to interleave thousands of ops.
138
+ threading.Event().wait(0.2)
139
+ stop.set()
140
+ for t in threads:
141
+ t.join()
142
+ snapshotter.join()
143
+
144
+ # Final drain so we don't lose the in-flight tail.
145
+ final = stats.clone_and_reset()
146
+ if not final.is_empty():
147
+ final.reduce()
148
+ snapshots.append(final)
149
+
150
+ # Every snapshot's columns must have identical lengths (the invariant
151
+ # the lock protects), and the union must contain at least one row.
152
+ total_rows = 0
153
+ for snap in snapshots:
154
+ n = len(snap.data["transfer_duration"])
155
+ assert len(snap.data["bytes_transferred"]) == n
156
+ assert len(snap.data["num_descriptors"]) == n
157
+ total_rows += n
158
+ assert total_rows > 0
159
+
160
+
161
+def test_clone_and_reset_hands_off_old_data():
162
+ stats = MooncakeKVConnectorStats()
163
+ stats.record_transfer(duration_s=0.001, total_bytes=1, num_descs=1)
164
+ stats.record_failed_recv()
165
+
166
+ snapshot = stats.clone_and_reset()
167
+
168
+ assert snapshot.num_successful_transfers == 1
169
+ assert not snapshot.is_empty()
170
+ # Original is now empty.
171
+ assert stats.is_empty()
172
+ assert stats.num_successful_transfers == 0
173
+ # Recording on the original does not mutate the snapshot.
174
+ stats.record_transfer(duration_s=0.005, total_bytes=2, num_descs=2)
175
+ assert snapshot.num_successful_transfers == 1
176
+
177
+
178
+def test_build_kv_connector_stats_none_returns_empty_instance():
179
+ out = MooncakeConnector.build_kv_connector_stats()
180
+ assert isinstance(out, MooncakeKVConnectorStats)
181
+ assert out.is_empty()
182
+
183
+
184
+def test_build_kv_connector_stats_with_data_round_trips():
185
+ original = MooncakeKVConnectorStats()
186
+ original.record_transfer(duration_s=0.01, total_bytes=1024, num_descs=3)
187
+ original.record_failed_transfer()
188
+
189
+ # Serialized form is the .data dict; build should reconstruct an instance
190
+ # that behaves the same.
191
+ rebuilt = MooncakeConnector.build_kv_connector_stats(data=original.data)
192
+
193
+ assert isinstance(rebuilt, MooncakeKVConnectorStats)
194
+ assert rebuilt.num_successful_transfers == 1
195
+ assert rebuilt.reduce()["Num failed transfers"] == 1
196
+
197
+
198
+def _bare_worker() -> MooncakeConnectorWorker:
199
+ """Construct a MooncakeConnectorWorker skipping __init__ (full init requires
200
+ a live TransferEngine). Only the attributes touched by the methods under
201
+ test are populated; role flags and async_zmq_ctx keep __del__'s shutdown
202
+ path a no-op."""
203
+ worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker)
204
+ worker.xfer_stats = MooncakeKVConnectorStats()
205
+ worker.engine = MagicMock()
206
+ worker.async_zmq_ctx = MagicMock()
207
+ worker.is_kv_consumer = True
208
+ worker.is_kv_producer = True
209
+ return worker
210
+
211
+
212
+def test_send_blocks_records_success():
213
+ worker = _bare_worker()
214
+ worker.engine.batch_transfer_sync_write.return_value = 0
215
+
216
+ ret = worker._send_blocks(
217
+ "host:1234",
218
+ src_ptrs=[0x1000, 0x2000],
219
+ dst_ptrs=[0x3000, 0x4000],
220
+ lengths=[1024, 2048],
221
+ )
222
+
223
+ assert ret == 0
224
+ assert worker.xfer_stats.num_successful_transfers == 1
225
+ data = worker.xfer_stats.data
226
+ assert data["bytes_transferred"] == [1024 + 2048]
227
+ assert data["num_descriptors"] == [2]
228
+ assert data["num_failed_transfers"] == []
229
+
230
+
231
+def test_send_blocks_records_failure():
232
+ worker = _bare_worker()
233
+ worker.engine.batch_transfer_sync_write.return_value = 1 # non-zero = fail
234
+
235
+ ret = worker._send_blocks("host:1234", [0x1000], [0x2000], [4096])
236
+
237
+ assert ret == 1
238
+ assert worker.xfer_stats.num_successful_transfers == 0
239
+ assert worker.xfer_stats.data["num_failed_transfers"] == [1]
240
+
241
+
242
+def test_get_kv_connector_stats_returns_none_when_empty():
243
+ worker = _bare_worker()
244
+
245
+ assert worker.get_kv_connector_stats() is None
246
+
247
+
248
+def test_get_kv_connector_stats_returns_and_resets():
249
+ worker = _bare_worker()
250
+ worker.engine.batch_transfer_sync_write.return_value = 0
251
+ worker._send_blocks("host:1234", [0x1000], [0x2000], [4096])
252
+
253
+ snapshot = worker.get_kv_connector_stats()
254
+ assert isinstance(snapshot, MooncakeKVConnectorStats)
255
+ assert snapshot.num_successful_transfers == 1
256
+
257
+ # Second call returns None because the worker's stats were reset.
258
+ assert worker.get_kv_connector_stats() is None
259
+
260
+
261
+def test_expired_request_bumps_counter():
262
+ import asyncio
263
+
264
+ worker = _bare_worker()
265
+ worker.reqs_need_send = {
266
+ "tid1": SendBlockMeta(
267
+ p_req_id="req1",
268
+ transfer_id="tid1",
269
+ local_block_ids=[0, 1],
270
+ ready=asyncio.Event(),
271
+ expire_time=-1.0, # Already expired.
272
+ sending=0,
273
+ ),
274
+ }
275
+ worker.finished_sending_reqs = set()
276
+
277
+ asyncio.run(worker.fetch_finished_sending_reqs())
278
+
279
+ assert worker.xfer_stats.data["num_kv_expired_reqs"] == [1]
280
+ # Expired transfer also cleaned out of reqs_need_send.
281
+ assert "tid1" not in worker.reqs_need_send
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+50
-3
@@ -31,10 +31,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
31
KVConnectorRole,
32
SupportsHMA,
33
)
34
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
35
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
36
MooncakeBootstrapServer,
37
RegisterWorkerPayload,
38
)
39
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.stats import (
40
+ MooncakeKVConnectorStats,
41
+)
42
from vllm.distributed.parallel_state import (
43
get_pp_group,
44
get_tensor_model_parallel_rank,
@@ -457,6 +461,25 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
461
def wait_for_save(self):
462
pass
463
464
+ def get_kv_connector_stats(self) -> KVConnectorStats | None:
465
+ """Return worker-local transfer stats since the last call.
466
+
467
+ Note the P/D asymmetry: because Mooncake is P-push (P calls
468
+ batch_transfer_sync_write), P records successful transfer latency,
469
+ bytes, and descriptor counts, while D only records failures
470
+ (recv/ZMQ errors). Aggregated NIXL-style dashboards will find
471
+ successful-transfer metrics on the P worker, not D.
472
+ """
473
+ if self.connector_worker is None:
474
+ return None
475
+ return self.connector_worker.get_kv_connector_stats()
476
+
477
+ @classmethod
478
+ def build_kv_connector_stats(
479
+ cls, data: dict[str, Any] | None = None
480
+ ) -> KVConnectorStats | None:
481
+ return MooncakeKVConnectorStats(data=data or {})
482
+
483
484
class MooncakeConnectorScheduler:
485
"""Implementation of Scheduler side methods"""
@@ -816,6 +839,8 @@ class MooncakeConnectorWorker:
839
self.finished_sending_reqs: set[ReqId] = set()
840
self.finished_recving_reqs: set[ReqId] = set()
841
842
+ self.xfer_stats = MooncakeKVConnectorStats()
843
+
844
self.block_size = vllm_config.cache_config.block_size
845
self.model_config = vllm_config.model_config
846
self.cache_config = vllm_config.cache_config
@@ -1340,11 +1365,23 @@ class MooncakeConnectorWorker:
1365
ret_value = self.engine.batch_transfer_sync_write(
1366
remote_session, src_ptrs, dst_ptrs, lengths
1367
)
1368
+ duration = time.perf_counter() - start_time
1369
if ret_value == 0:
1344
- logger.debug(
1345
- "Sending to %s done, took %s",
1370
+ self.xfer_stats.record_transfer(
1371
+ duration_s=duration,
1372
+ total_bytes=sum(lengths),
1373
+ num_descs=len(src_ptrs),
1374
+ )
1375
+ logger.debug("Sending to %s done, took %s", remote_session, duration)
1376
+ else:
1377
+ self.xfer_stats.record_failed_transfer()
1378
+ logger.warning(
1379
+ "Sending to %s failed (ret=%s) after %s (%d descriptors, %d bytes)",
1380
remote_session,
1347
- time.perf_counter() - start_time,
1381
+ ret_value,
1382
+ duration,
1383
+ len(src_ptrs),
1384
+ sum(lengths),
1385
)
1386
return ret_value
1387
@@ -1445,6 +1482,7 @@ class MooncakeConnectorWorker:
1482
send_meta.p_req_id,
1483
envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT,
1484
)
1485
+ self.xfer_stats.record_kv_expired_req()
1486
finished_sending_reqs.add(send_meta.p_req_id)
1487
expired_transfer_id.append(transfer_id)
1488
@@ -1485,6 +1523,13 @@ class MooncakeConnectorWorker:
1523
1524
return finished_sending_reqs or None, finished_recving_reqs or None
1525
1526
+ def get_kv_connector_stats(self) -> KVConnectorStats | None:
1527
+ """Return transfer stats collected since the last call, or None
1528
+ if nothing has been recorded in this interval."""
1529
+ if self.xfer_stats.is_empty():
1530
+ return None
1531
+ return self.xfer_stats.clone_and_reset()
1532
+
1533
async def receive_kv_from_single_worker(
1534
self,
1535
worker_addr: str,
@@ -1531,6 +1576,7 @@ class MooncakeConnectorWorker:
1576
req_ids,
1577
response.err_msg,
1578
)
1579
+ self.xfer_stats.record_failed_recv()
1580
return
1581
self.process_pulling_result(response, pull_metas)
1582
if response.status == MooncakeXferResponseStatus.FINISH:
@@ -1539,6 +1585,7 @@ class MooncakeConnectorWorker:
1585
logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.")
1586
except Exception as e:
1587
logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e)
1588
+ self.xfer_stats.record_failed_recv()
1589
return
1590
1591
def process_pulling_result(
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py
new
+146
@@ -0,0 +1,146 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Stats container for the Mooncake connector."""
4
+
5
+import threading
6
+from dataclasses import dataclass
7
+from typing import Any
8
+
9
+import numpy as np
10
+
11
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
12
+ KVConnectorStats,
13
+)
14
+
15
+# TODO(mooncake-stats): add MooncakePromMetrics (mirror NixlPromMetrics)
16
+# and wire it via MooncakeConnector.build_prom_metrics in a follow-up PR.
17
+
18
+
19
+@dataclass
20
+class MooncakeKVConnectorStats(KVConnectorStats):
21
+ """Container for Mooncake KV transfer performance metrics.
22
+
23
+ `_lock` serializes record_* against clone_and_reset so each row's
24
+ appends are atomic and column lengths stay aligned. Writers run on
25
+ the sender pool / receiver loop / sender loop; reader runs on the
26
+ main worker thread.
27
+ """
28
+
29
+ def __post_init__(self):
30
+ self._lock = threading.Lock()
31
+ if not self.data:
32
+ self.reset()
33
+
34
+ # threading.Lock is not picklable; strip it from the wire form and
35
+ # rebuild a fresh per-process lock on the receiver side.
36
+ def __getstate__(self) -> dict[str, Any]:
37
+ state = self.__dict__.copy()
38
+ state.pop("_lock", None)
39
+ return state
40
+
41
+ def __setstate__(self, state: dict[str, Any]) -> None:
42
+ self.__dict__.update(state)
43
+ self._lock = threading.Lock()
44
+
45
+ def reset(self):
46
+ self.data: dict[str, list[float | int]] = {
47
+ "transfer_duration": [],
48
+ "bytes_transferred": [],
49
+ "num_descriptors": [],
50
+ "num_failed_transfers": [],
51
+ "num_failed_recvs": [],
52
+ "num_kv_expired_reqs": [],
53
+ }
54
+
55
+ def record_transfer(self, duration_s: float, total_bytes: int, num_descs: int):
56
+ with self._lock:
57
+ self.data["transfer_duration"].append(duration_s)
58
+ self.data["bytes_transferred"].append(total_bytes)
59
+ self.data["num_descriptors"].append(num_descs)
60
+
61
+ # Failure counters store a list of 1s so a future Prom counter can iterate
62
+ # with .inc(list_item), mirroring NIXL's NixlPromMetrics.observe.
63
+ def record_failed_transfer(self):
64
+ with self._lock:
65
+ self.data["num_failed_transfers"].append(1)
66
+
67
+ def record_failed_recv(self):
68
+ with self._lock:
69
+ self.data["num_failed_recvs"].append(1)
70
+
71
+ def record_kv_expired_req(self):
72
+ with self._lock:
73
+ self.data["num_kv_expired_reqs"].append(1)
74
+
75
+ def clone_and_reset(self) -> "MooncakeKVConnectorStats":
76
+ # Copy lists under the lock for length alignment; return a fresh
77
+ # instance so the snapshot has its own _lock.
78
+ with self._lock:
79
+ snapshot_data: dict[str, list[float | int]] = {
80
+ k: list(v) for k, v in self.data.items()
81
+ }
82
+ self.reset()
83
+ return MooncakeKVConnectorStats(data=snapshot_data)
84
+
85
+ def is_empty(self) -> bool:
86
+ return (
87
+ self.num_successful_transfers == 0
88
+ and len(self.data["num_failed_transfers"]) == 0
89
+ and len(self.data["num_failed_recvs"]) == 0
90
+ and len(self.data["num_kv_expired_reqs"]) == 0
91
+ )
92
+
93
+ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
94
+ if not other.is_empty():
95
+ for k, v in other.data.items():
96
+ accumulator = self.data[k]
97
+ assert isinstance(accumulator, list)
98
+ accumulator.extend(v)
99
+ return self
100
+
101
+ def reduce(self) -> dict[str, int | float]:
102
+ num_failed_transfers = len(self.data["num_failed_transfers"])
103
+ num_failed_recvs = len(self.data["num_failed_recvs"])
104
+ num_kv_expired_reqs = len(self.data["num_kv_expired_reqs"])
105
+
106
+ if self.num_successful_transfers == 0:
107
+ return {
108
+ "Num successful transfers": 0,
109
+ "Avg xfer time (ms)": 0,
110
+ "P90 xfer time (ms)": 0,
111
+ "Avg MB per transfer": 0,
112
+ "Throughput (MB/s)": 0,
113
+ "Avg number of descriptors": 0,
114
+ "Num failed transfers": num_failed_transfers,
115
+ "Num failed recvs": num_failed_recvs,
116
+ "Num KV expired reqs": num_kv_expired_reqs,
117
+ }
118
+
119
+ xfer_time = np.asarray(self.data["transfer_duration"])
120
+ mb = np.asarray(self.data["bytes_transferred"]) / 2**20
121
+ descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32)
122
+ n = len(descs)
123
+ assert n == self.num_successful_transfers
124
+
125
+ total_mb = mb.sum()
126
+ avg_mb = total_mb / n
127
+ total_time_seconds = xfer_time.sum()
128
+ throughput_mb_s = (
129
+ total_mb / total_time_seconds if total_time_seconds > 0 else 0.0
130
+ )
131
+
132
+ return {
133
+ "Num successful transfers": n,
134
+ "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),
135
+ "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),
136
+ "Avg MB per transfer": round(avg_mb, 3),
137
+ "Throughput (MB/s)": round(throughput_mb_s, 3),
138
+ "Avg number of descriptors": round(descs.mean(), 1),
139
+ "Num failed transfers": num_failed_transfers,
140
+ "Num failed recvs": num_failed_recvs,
141
+ "Num KV expired reqs": num_kv_expired_reqs,
142
+ }
143
+
144
+ @property
145
+ def num_successful_transfers(self) -> int:
146
+ return len(self.data["transfer_duration"])