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 3, 19:00 UTC
[KV Transfer] Add MooncakeStoreConnector for KV cache offloading via Mooncake distributed store (#40900)
Signed-off-by: leichao.lc <leichao.lc@antgroup.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: leichao.lc <leichao.lc@antgroup.com> Co-authored-by: ivanium <yifanqiao@inferact.ai> Co-authored-by: aoshen524 <aoshen@inferact.ai> Co-authored-by: Dao007forever <daole@inferact.ai> Co-authored-by: Teng Ma <sima.mt@alibaba-inc.com> Co-authored-by: Pz1116 <zpbzpb123123@gmail.com> Co-authored-by: foraxe <1055696449@qq.com> Co-authored-by: Skywalker-EP <173423846@qq.com> Co-authored-by: fems14 <1804143737@qq.com> Co-authored-by: jianzs <zheng.shoujian@outlook.com> Co-authored-by: baxingpiaochong <771405853@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Chao Lei committed
May 13, 2026 at 07:09 UTC
ebeb09d82261e0f841ea10c109adbdaefd3ce459
10 files changed
+2600
docs/features/mooncake_store_connector_usage.md
new
+161
@@ -0,0 +1,161 @@
1
+# MooncakeStoreConnector Usage Guide
2
+
3
+MooncakeStoreConnector is a KV cache connector that uses [MooncakeDistributedStore](https://github.com/kvcache-ai/Mooncake) as a shared KV cache pool. Unlike `MooncakeConnector` which does direct point-to-point KV transfer between prefiller and decoder, MooncakeStoreConnector enables KV cache offloading to an external distributed store, supporting:
4
+
5
+- **CPU offloading**: Extend effective KV cache capacity by offloading to CPU memory via Mooncake's transfer engine.
6
+- **Prefix caching across instances**: Hash-based deduplication allows multiple vLLM instances to share cached KV blocks through the store.
7
+- **Single-node and multi-node deployment**: Works both as a standalone KV cache extension and in disaggregated prefill-decode setups.
8
+
9
+## Prerequisites
10
+
11
+### Install Mooncake
12
+
13
+Install mooncake through pip:
14
+
15
+```bash
16
+uv pip install mooncake-transfer-engine
17
+```
18
+
19
+Refer to the [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions and building from source.
20
+
21
+### Start the Mooncake Master Server
22
+
23
+The Mooncake master manages metadata and coordinates the distributed store. Start it before launching vLLM:
24
+
25
+```bash
26
+mooncake_master --port 50051
27
+```
28
+
29
+Default ports:
30
+
31
+- RPC: 50051
32
+
33
+Multiple vLLM instances can share the same master server.
34
+
35
+### Configure Mooncake
36
+
37
+Create a JSON configuration file (e.g., `mooncake_config.json`):
38
+
39
+```json
40
+{
41
+ "metadata_server": "P2PHANDSHAKE",
42
+ "master_server_address": "127.0.0.1:50051",
43
+ "global_segment_size": "80GB",
44
+ "local_buffer_size": "4GB",
45
+ "protocol": "rdma",
46
+ "device_name": ""
47
+}
48
+```
49
+
50
+- `protocol`: Use `"rdma"` for best performance. `"tcp"` works as a fallback.
51
+- `global_segment_size`: CPU memory contributed to the distributed pool (per GPU).
52
+- `local_buffer_size`: Private buffer for this node's own operations (per GPU).
53
+
54
+Set the config path via environment variable:
55
+
56
+```bash
57
+export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json
58
+```
59
+
60
+## Usage
61
+
62
+### Single-Node KV Cache Offloading
63
+
64
+Use MooncakeStoreConnector to offload KV cache to CPU memory, extending the effective cache size:
65
+
66
+```bash
67
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
68
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
69
+ --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}'
70
+```
71
+
72
+### Disaggregated Prefill-Decode (XpYd)
73
+
74
+In disaggregated prefill-decode mode, use `MultiConnector` to combine `MooncakeConnector` (point-to-point KV transfer) with `MooncakeStoreConnector` (shared KV cache pool). This enables both direct P2P transfer between prefiller and decoder, and cross-instance prefix cache sharing via the distributed store.
75
+**Prefiller Node:**
76
+
77
+```bash
78
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
79
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50052 \
80
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
81
+ --port 8100 \
82
+ --kv-transfer-config '{
83
+ "kv_connector": "MultiConnector",
84
+ "kv_role": "kv_producer",
85
+ "kv_connector_extra_config": {
86
+ "connectors": [
87
+ {
88
+ "kv_connector": "MooncakeConnector",
89
+ "kv_role": "kv_producer"
90
+ },
91
+ {
92
+ "kv_connector": "MooncakeStoreConnector",
93
+ "kv_role": "kv_producer"
94
+ }
95
+ ]
96
+ }
97
+ }'
98
+```
99
+
100
+**Decoder Node:**
101
+
102
+```bash
103
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
104
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50053 \
105
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
106
+ --port 8200 \
107
+ --kv-transfer-config '{
108
+ "kv_connector": "MultiConnector",
109
+ "kv_role": "kv_consumer",
110
+ "kv_connector_extra_config": {
111
+ "connectors": [
112
+ {
113
+ "kv_connector": "MooncakeConnector",
114
+ "kv_role": "kv_consumer"
115
+ },
116
+ {
117
+ "kv_connector": "MooncakeStoreConnector",
118
+ "kv_role": "kv_consumer"
119
+ }
120
+ ]
121
+ }
122
+ }'
123
+```
124
+
125
+**Proxy:**
126
+
127
+A disaggregation proxy is required to route requests between prefiller and decoder nodes. The proxy assigns `do_remote_prefill=True` / `do_remote_decode=True` to coordinate P2P transfer via `MooncakeConnector`. Refer to the [MooncakeConnector usage guide](mooncake_connector_usage.md) for proxy setup details.
128
+
129
+## Environment Variables
130
+
131
+| Variable | Description | Default |
132
+| --- | --- | --- |
133
+| `MOONCAKE_CONFIG_PATH` | Path to Mooncake JSON config file | (required) |
134
+| `VLLM_MOONCAKE_BOOTSTRAP_PORT` | Bootstrap port for MooncakeConnector P2P transfer (disagg mode only) | 8998 |
135
+
136
+## KV Transfer Config
137
+
138
+### KV Role Options
139
+
140
+- **kv_producer**: For prefiller instances that store KV caches to the pool.
141
+- **kv_consumer**: For decoder instances that load KV caches from the pool.
142
+- **kv_both**: The instance both stores and loads KV caches. Use this for single-node CPU offloading.
143
+
144
+### kv_connector_extra_config
145
+
146
+- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
147
+- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
148
+- `discard_partial_chunks` (bool): Discard partial block chunks during store. Default: `true`.
149
+- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
150
+
151
+## Notes
152
+
153
+### Cross-DP Prefix Cache Hits
154
+
155
+When running with data parallelism, set a fixed `PYTHONHASHSEED` so that block hashes are consistent across DP ranks:
156
+
157
+```bash
158
+PYTHONHASHSEED=0 vllm serve ...
159
+```
160
+
161
+Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits.
tests/v1/kv_connector/unit/test_mooncake_store_connector.py
new
+258
@@ -0,0 +1,258 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from unittest.mock import MagicMock, patch
5
+
6
+from vllm.config import set_current_vllm_config
7
+from vllm.distributed.kv_events import BlockStored
8
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
9
+ KVConnectorRole,
10
+)
11
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
12
+ connector,
13
+ worker,
14
+)
15
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
16
+ MooncakeStoreConnectorMetadata,
17
+)
18
+from vllm.v1.outputs import KVConnectorOutput
19
+
20
+from .utils import create_vllm_config
21
+
22
+
23
+def _make_vllm_config():
24
+ return create_vllm_config(
25
+ kv_connector="MooncakeStoreConnector",
26
+ kv_role="kv_both",
27
+ )
28
+
29
+
30
+def _make_block_stored() -> BlockStored:
31
+ return BlockStored(
32
+ block_hashes=[b"hash"],
33
+ parent_block_hash=None,
34
+ token_ids=[1, 2, 3],
35
+ block_size=16,
36
+ lora_id=None,
37
+ medium="cpu",
38
+ lora_name=None,
39
+ )
40
+
41
+
42
+def test_scheduler_role_initializes_store_scheduler_only():
43
+ vllm_config = _make_vllm_config()
44
+
45
+ with (
46
+ set_current_vllm_config(vllm_config),
47
+ patch(
48
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
49
+ "connector.MooncakeStoreScheduler"
50
+ ) as mock_scheduler,
51
+ patch(
52
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
53
+ "connector.MooncakeStoreWorker"
54
+ ) as mock_worker,
55
+ ):
56
+ conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
57
+
58
+ mock_scheduler.assert_called_once_with(vllm_config)
59
+ mock_worker.assert_not_called()
60
+ assert conn.connector_scheduler is mock_scheduler.return_value
61
+ assert conn.connector_worker is None
62
+
63
+
64
+def test_worker_role_initializes_store_worker_on_rank0():
65
+ vllm_config = _make_vllm_config()
66
+
67
+ with (
68
+ set_current_vllm_config(vllm_config),
69
+ patch(
70
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
71
+ "connector.MooncakeStoreScheduler"
72
+ ) as mock_scheduler,
73
+ patch(
74
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
75
+ "connector.MooncakeStoreWorker"
76
+ ) as mock_worker,
77
+ ):
78
+ conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
79
+
80
+ mock_scheduler.assert_not_called()
81
+ mock_worker.assert_called_once_with(vllm_config)
82
+ assert conn.connector_scheduler is None
83
+ assert conn.connector_worker is mock_worker.return_value
84
+
85
+
86
+def test_worker_role_initializes_on_nonzero_rank():
87
+ vllm_config = _make_vllm_config()
88
+ vllm_config.parallel_config.rank = 1
89
+
90
+ with (
91
+ set_current_vllm_config(vllm_config),
92
+ patch(
93
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
94
+ "connector.MooncakeStoreWorker"
95
+ ) as mock_worker,
96
+ ):
97
+ connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
98
+
99
+ mock_worker.assert_called_once_with(vllm_config)
100
+
101
+
102
+def test_lookup_rpc_path_uses_data_parallel_index_in_dense_dp():
103
+ vllm_config = _make_vllm_config()
104
+ vllm_config.parallel_config.data_parallel_rank = 0
105
+ vllm_config.parallel_config.data_parallel_index = 3
106
+
107
+ path = worker.get_zmq_rpc_path_lookup(vllm_config)
108
+
109
+ assert path.endswith("_dp_rank3")
110
+
111
+
112
+def test_lookup_rpc_path_uses_local_rank_when_local_engines_only():
113
+ vllm_config = _make_vllm_config()
114
+ vllm_config.parallel_config.data_parallel_index = 7
115
+ vllm_config.parallel_config.data_parallel_rank_local = 1
116
+ vllm_config.parallel_config.data_parallel_hybrid_lb = True
117
+
118
+ path = worker.get_zmq_rpc_path_lookup(vllm_config)
119
+
120
+ assert path.endswith("_dp_rank1")
121
+
122
+
123
+def test_worker_methods_delegate_to_store_worker():
124
+ vllm_config = _make_vllm_config()
125
+ kv_caches = {"layer0": MagicMock()}
126
+ metadata = MooncakeStoreConnectorMetadata(set(), set())
127
+ finished_req_ids = {"req-1"}
128
+
129
+ with (
130
+ set_current_vllm_config(vllm_config),
131
+ patch(
132
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
133
+ "connector.MooncakeStoreWorker"
134
+ ) as mock_worker_cls,
135
+ ):
136
+ conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
137
+
138
+ worker_inst = mock_worker_cls.return_value
139
+ worker_inst.get_finished.return_value = ({"req-1"}, {"req-2"})
140
+ conn.bind_connector_metadata(metadata)
141
+
142
+ conn.register_kv_caches(kv_caches)
143
+ result = conn.get_finished(finished_req_ids)
144
+
145
+ worker_inst.register_kv_caches.assert_called_once_with(kv_caches)
146
+ worker_inst.get_finished.assert_called_once_with(finished_req_ids, metadata)
147
+ assert result == ({"req-1"}, {"req-2"})
148
+
149
+
150
+def test_get_kv_connector_kv_cache_events_returns_none_when_empty():
151
+ vllm_config = _make_vllm_config()
152
+
153
+ with (
154
+ set_current_vllm_config(vllm_config),
155
+ patch(
156
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
157
+ "connector.MooncakeStoreWorker"
158
+ ) as mock_worker_cls,
159
+ ):
160
+ conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
161
+
162
+ mock_worker_cls.return_value.get_kv_events.return_value = []
163
+ assert conn.get_kv_connector_kv_cache_events() is None
164
+
165
+
166
+def test_get_kv_connector_kv_cache_events_wraps_worker_events():
167
+ vllm_config = _make_vllm_config()
168
+ event = _make_block_stored()
169
+
170
+ with (
171
+ set_current_vllm_config(vllm_config),
172
+ patch(
173
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
174
+ "connector.MooncakeStoreWorker"
175
+ ) as mock_worker_cls,
176
+ ):
177
+ conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
178
+
179
+ mock_worker_cls.return_value.get_kv_events.return_value = [event]
180
+ kv_events = conn.get_kv_connector_kv_cache_events()
181
+
182
+ assert isinstance(kv_events, connector.MooncakeStoreKVEvents)
183
+ assert kv_events.get_number_of_workers() == 1
184
+ assert kv_events.get_all_events() == [event]
185
+
186
+
187
+def test_prefer_cross_layer_blocks_from_config():
188
+ # Default: disabled
189
+ vllm_config = _make_vllm_config()
190
+ with (
191
+ set_current_vllm_config(vllm_config),
192
+ patch(
193
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
194
+ "connector.MooncakeStoreScheduler"
195
+ ),
196
+ ):
197
+ conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
198
+ assert conn.prefer_cross_layer_blocks is False
199
+
200
+ # Enabled via config
201
+ vllm_config_enabled = create_vllm_config(
202
+ kv_connector="MooncakeStoreConnector",
203
+ kv_role="kv_both",
204
+ kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
205
+ )
206
+ with (
207
+ set_current_vllm_config(vllm_config_enabled),
208
+ patch(
209
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
210
+ "connector.MooncakeStoreScheduler"
211
+ ),
212
+ ):
213
+ conn_enabled = connector.MooncakeStoreConnector(
214
+ vllm_config_enabled, KVConnectorRole.SCHEDULER
215
+ )
216
+ assert conn_enabled.prefer_cross_layer_blocks is True
217
+
218
+
219
+def test_register_cross_layers_kv_cache_delegates_to_worker():
220
+ vllm_config = _make_vllm_config()
221
+
222
+ with (
223
+ set_current_vllm_config(vllm_config),
224
+ patch(
225
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
226
+ "connector.MooncakeStoreWorker"
227
+ ) as mock_worker_cls,
228
+ ):
229
+ conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
230
+
231
+ fake_tensor = MagicMock()
232
+ fake_backend = MagicMock()
233
+ conn.register_cross_layers_kv_cache(fake_tensor, fake_backend)
234
+
235
+ worker_inst = mock_worker_cls.return_value
236
+ worker_inst.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor)
237
+
238
+
239
+def test_update_connector_output_and_take_events():
240
+ vllm_config = _make_vllm_config()
241
+ event = _make_block_stored()
242
+
243
+ with (
244
+ set_current_vllm_config(vllm_config),
245
+ patch(
246
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
247
+ "connector.MooncakeStoreScheduler"
248
+ ),
249
+ ):
250
+ conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
251
+
252
+ kv_events = connector.MooncakeStoreKVEvents(num_workers=1)
253
+ kv_events.add_events([event])
254
+ conn.update_connector_output(KVConnectorOutput(kv_cache_events=kv_events))
255
+
256
+ assert conn._kv_cache_events is kv_events
257
+ assert list(conn.take_events()) == [event]
258
+ assert conn._kv_cache_events is None
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
new
+300
@@ -0,0 +1,300 @@
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, patch
6
+
7
+import torch
8
+
9
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
10
+ worker,
11
+)
12
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
13
+ ChunkedTokenDatabase,
14
+ KeyMetadata,
15
+ ReqMeta,
16
+)
17
+
18
+
19
+def _make_store_sending_thread(
20
+ store: MagicMock,
21
+) -> worker.KVCacheStoreSendingThread:
22
+ token_database = ChunkedTokenDatabase(
23
+ KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
24
+ )
25
+ token_database.set_kv_caches_base_addr([0x1000])
26
+ token_database.set_block_len([256])
27
+ thread = worker.KVCacheStoreSendingThread(
28
+ store=store,
29
+ token_database=token_database,
30
+ block_size=16,
31
+ tp_rank=0,
32
+ put_step=1,
33
+ kv_role="kv_producer",
34
+ ready_event=threading.Event(),
35
+ )
36
+ thread.request_queue.task_done = MagicMock()
37
+ return thread
38
+
39
+
40
+def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
41
+ return ReqMeta(
42
+ req_id=req_id,
43
+ token_len_chunk=32,
44
+ block_ids=[0, 1],
45
+ block_hashes=block_hashes,
46
+ can_save=True,
47
+ original_block_size=16,
48
+ )
49
+
50
+
51
+def test_store_sending_thread_skips_request_during_cpu_pressure():
52
+ store = MagicMock()
53
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
54
+ store.batch_put_from_multi_buffers.side_effect = [
55
+ [-200, -200],
56
+ [256, 256],
57
+ [256, 256],
58
+ ]
59
+ thread = _make_store_sending_thread(store)
60
+
61
+ thread.add_stored_request("req-a")
62
+ thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
63
+
64
+ assert thread._store_pressure_active is True
65
+ assert "req-a" in thread._skip_store_requests
66
+ assert store.batch_put_from_multi_buffers.call_count == 1
67
+
68
+ thread.add_stored_request("req-a")
69
+ thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))
70
+
71
+ assert store.batch_put_from_multi_buffers.call_count == 1
72
+
73
+ thread.add_stored_request("req-b")
74
+ thread._handle_request(_make_store_req("req-b", [b"b0", b"b1"]))
75
+
76
+ assert thread._store_pressure_active is False
77
+ assert "req-a" not in thread._skip_store_requests
78
+ assert store.batch_put_from_multi_buffers.call_count == 2
79
+
80
+ thread.add_stored_request("req-a")
81
+ thread._handle_request(_make_store_req("req-a", [b"a4", b"a5"]))
82
+
83
+ assert store.batch_put_from_multi_buffers.call_count == 3
84
+
85
+
86
+def test_store_sending_thread_only_skips_on_no_available_handle():
87
+ store = MagicMock()
88
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
89
+ store.batch_put_from_multi_buffers.side_effect = [
90
+ [-500, -500],
91
+ [256, 256],
92
+ ]
93
+ thread = _make_store_sending_thread(store)
94
+
95
+ thread.add_stored_request("req-a")
96
+ thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
97
+
98
+ assert thread._store_pressure_active is False
99
+ assert "req-a" not in thread._skip_store_requests
100
+ assert store.batch_put_from_multi_buffers.call_count == 1
101
+
102
+ thread.add_stored_request("req-a")
103
+ thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))
104
+
105
+ assert store.batch_put_from_multi_buffers.call_count == 2
106
+
107
+
108
+# ---------------------------------------------------------------------------
109
+# Helpers for register_kv_caches tests
110
+# ---------------------------------------------------------------------------
111
+
112
+
113
+def _auto_set_ready_event(*args, **kwargs):
114
+ """Side effect for mocked thread constructors that auto-sets ready_event."""
115
+ for arg in args:
116
+ if isinstance(arg, threading.Event):
117
+ arg.set()
118
+ for val in kwargs.values():
119
+ if isinstance(val, threading.Event):
120
+ val.set()
121
+ return MagicMock()
122
+
123
+
124
+def _make_bare_worker(
125
+ *,
126
+ num_gpu_blocks: int = 10,
127
+ block_size: int = 16,
128
+ kv_role: str = "kv_both",
129
+) -> worker.MooncakeStoreWorker:
130
+ """Construct a MooncakeStoreWorker via __new__, bypassing __init__.
131
+
132
+ Sets only the attributes that register_kv_caches() reads so we can
133
+ test the stride-based layout detection without a real
134
+ MooncakeDistributedStore.
135
+ """
136
+ w = object.__new__(worker.MooncakeStoreWorker)
137
+ w.cache_config = MagicMock()
138
+ w.cache_config.num_gpu_blocks = num_gpu_blocks
139
+ w.store = MagicMock()
140
+ w.store.register_buffer.return_value = 0
141
+ w.use_mla = False
142
+ w.token_database = ChunkedTokenDatabase(
143
+ KeyMetadata("test-model", 0, 0, 0, 0), block_size=block_size
144
+ )
145
+ w.kv_role = kv_role
146
+ w.block_size = block_size
147
+ w.tp_rank = 0
148
+ w.put_step = 1
149
+ w.enable_kv_events = False
150
+ w.kv_send_thread = None
151
+ w.kv_recv_thread = None
152
+ return w
153
+
154
+
155
+# ---------------------------------------------------------------------------
156
+# register_kv_caches tests
157
+# ---------------------------------------------------------------------------
158
+
159
+
160
+def test_register_kv_caches_blocks_first_single_segment():
161
+ """Blocks-first layout (FlashInfer/MLA): one segment per layer."""
162
+ num_blocks = 10
163
+ page_size_elements = 64 # elements per block
164
+ w = _make_bare_worker(num_gpu_blocks=num_blocks)
165
+
166
+ # Shape: (num_blocks, page_size_elements) — blocks outermost, no outer_dims
167
+ tensor = torch.zeros(num_blocks, page_size_elements, dtype=torch.float16)
168
+
169
+ with (
170
+ patch(
171
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
172
+ "worker.KVCacheStoreSendingThread",
173
+ side_effect=_auto_set_ready_event,
174
+ ),
175
+ patch(
176
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
177
+ "worker.KVCacheStoreRecvingThread",
178
+ side_effect=_auto_set_ready_event,
179
+ ),
180
+ ):
181
+ w.register_kv_caches({"layer0": tensor})
182
+
183
+ assert len(w.kv_caches_base_addr) == 1
184
+ assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
185
+
186
+ expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
187
+ assert len(w.block_len) == 1
188
+ assert w.block_len[0] == expected_block_len
189
+
190
+ w.store.register_buffer.assert_called_once_with(
191
+ tensor.untyped_storage().data_ptr(),
192
+ tensor.untyped_storage().nbytes(),
193
+ )
194
+
195
+
196
+def test_register_kv_caches_kv_first_two_segments():
197
+ """K/V-first layout (FlashAttn): two segments (K, V) per layer."""
198
+ num_blocks = 10
199
+ block_size_tokens = 16
200
+ num_kv_heads = 4
201
+ head_size = 8
202
+
203
+ w = _make_bare_worker(num_gpu_blocks=num_blocks)
204
+
205
+ # Shape: (2, num_blocks, block_size, num_kv_heads, head_size) — K/V outermost
206
+ tensor = torch.zeros(
207
+ 2,
208
+ num_blocks,
209
+ block_size_tokens,
210
+ num_kv_heads,
211
+ head_size,
212
+ dtype=torch.float16,
213
+ )
214
+
215
+ with (
216
+ patch(
217
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
218
+ "worker.KVCacheStoreSendingThread",
219
+ side_effect=_auto_set_ready_event,
220
+ ),
221
+ patch(
222
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
223
+ "worker.KVCacheStoreRecvingThread",
224
+ side_effect=_auto_set_ready_event,
225
+ ),
226
+ ):
227
+ w.register_kv_caches({"layer0": tensor})
228
+
229
+ # K/V-first: dim 0 has stride > page_size, so 2 segments
230
+ assert len(w.kv_caches_base_addr) == 2
231
+ assert len(w.block_len) == 2
232
+
233
+ el = tensor.element_size()
234
+ seg_stride = tensor.stride(0) * el # stride of the K/V dim in bytes
235
+ base = tensor.untyped_storage().data_ptr()
236
+ assert w.kv_caches_base_addr[0] == base
237
+ assert w.kv_caches_base_addr[1] == base + seg_stride
238
+ assert w.block_len[0] == seg_stride // num_blocks
239
+ assert w.block_len[1] == seg_stride // num_blocks
240
+
241
+
242
+def test_register_kv_caches_cross_layer_single_segment():
243
+ """Cross-layer tensor: single segment with block_len = page_size * num_layers."""
244
+ num_blocks = 10
245
+ num_layers = 4
246
+ per_layer_page_elements = 64 # elements per layer per block
247
+
248
+ w = _make_bare_worker(num_gpu_blocks=num_blocks)
249
+
250
+ # Cross-layer blocks-first tensor: all layers packed into a single
251
+ # contiguous block. Shape (num_blocks, num_layers * per_layer_page)
252
+ # mimics the physical layout after stride reordering.
253
+ total_page_elements = num_layers * per_layer_page_elements
254
+ tensor = torch.zeros(num_blocks, total_page_elements, dtype=torch.float16)
255
+
256
+ with (
257
+ patch(
258
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
259
+ "worker.KVCacheStoreSendingThread",
260
+ side_effect=_auto_set_ready_event,
261
+ ),
262
+ patch(
263
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
264
+ "worker.KVCacheStoreRecvingThread",
265
+ side_effect=_auto_set_ready_event,
266
+ ),
267
+ ):
268
+ # Use the cross-layer wrapper key, same as register_cross_layers_kv_caches
269
+ w.register_kv_caches({"__cross_layer__": tensor})
270
+
271
+ assert len(w.kv_caches_base_addr) == 1
272
+ assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
273
+
274
+ expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
275
+ # block_len should be per_layer_page_size * num_layers
276
+ assert (
277
+ expected_block_len
278
+ == num_layers * per_layer_page_elements * tensor.element_size()
279
+ )
280
+ assert len(w.block_len) == 1
281
+ assert w.block_len[0] == expected_block_len
282
+
283
+ # Also verify via register_cross_layers_kv_caches wrapper
284
+ w2 = _make_bare_worker(num_gpu_blocks=num_blocks)
285
+ with (
286
+ patch(
287
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
288
+ "worker.KVCacheStoreSendingThread",
289
+ side_effect=_auto_set_ready_event,
290
+ ),
291
+ patch(
292
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
293
+ "worker.KVCacheStoreRecvingThread",
294
+ side_effect=_auto_set_ready_event,
295
+ ),
296
+ ):
297
+ w2.register_cross_layers_kv_caches(tensor)
298
+
299
+ assert w2.kv_caches_base_addr == w.kv_caches_base_addr
300
+ assert w2.block_len == w.block_len
vllm/distributed/kv_transfer/kv_connector/factory.py
+5
@@ -197,6 +197,11 @@ KVConnectorFactory.register_connector(
197
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
198
"MooncakeConnector",
199
)
200
+KVConnectorFactory.register_connector(
201
+ "MooncakeStoreConnector",
202
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.connector",
203
+ "MooncakeStoreConnector",
204
+)
205
KVConnectorFactory.register_connector(
206
"FlexKVConnectorV1",
207
"vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector",
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py
+10
@@ -8,6 +8,7 @@ import uvicorn
8
from fastapi import FastAPI, HTTPException
9
from pydantic import BaseModel
10
11
+from vllm.config import ParallelConfig
12
from vllm.distributed.kv_transfer.kv_connector.utils import EngineId
13
from vllm.logger import init_logger
14
@@ -16,6 +17,15 @@ WorkerAddr = str
17
logger = init_logger(__name__)
18
19
20
+def get_mooncake_dp_engine_index(parallel_config: ParallelConfig) -> int:
21
+ """Return the per-engine DP index used for Mooncake side channels."""
22
+ if parallel_config.local_engines_only:
23
+ assert parallel_config.data_parallel_rank_local is not None
24
+ return parallel_config.data_parallel_rank_local
25
+
26
+ return parallel_config.data_parallel_index
27
+
28
+
29
class RegisterWorkerPayload(BaseModel):
30
engine_id: EngineId
31
dp_rank: int
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
new
+229
@@ -0,0 +1,229 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#
4
+# Adapted from vllm-project/vllm-ascend
5
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
6
+"""MooncakeStoreConnector - KV cache connector using MooncakeDistributedStore.
7
+
8
+Unlike MooncakeConnector which does direct P2P transfer, this connector
9
+uses MooncakeDistributedStore as a shared KV cache pool. Both producer
10
+and consumer instances read/write KV to/from the store independently,
11
+enabling prefix caching via hash-based deduplication.
12
+"""
13
+
14
+from collections.abc import Iterable
15
+from typing import Any
16
+
17
+import torch
18
+
19
+from vllm.config import VllmConfig
20
+from vllm.distributed.kv_events import (
21
+ KVCacheEvent,
22
+ KVConnectorKVEvents,
23
+ KVEventAggregator,
24
+)
25
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
26
+ KVConnectorBase_V1,
27
+ KVConnectorMetadata,
28
+ KVConnectorRole,
29
+)
30
+from vllm.forward_context import ForwardContext
31
+from vllm.logger import init_logger
32
+from vllm.v1.attention.backend import AttentionMetadata
33
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
34
+from vllm.v1.core.sched.output import SchedulerOutput
35
+from vllm.v1.kv_cache_interface import KVCacheConfig
36
+from vllm.v1.outputs import KVConnectorOutput
37
+from vllm.v1.request import Request
38
+
39
+from .data import MooncakeStoreConnectorMetadata
40
+from .scheduler import MooncakeStoreScheduler
41
+from .worker import MooncakeStoreWorker
42
+
43
+logger = init_logger(__name__)
44
+
45
+
46
+class MooncakeStoreKVEvents(KVConnectorKVEvents):
47
+ """KV event aggregation for MooncakeStoreConnector."""
48
+
49
+ def __init__(self, num_workers: int) -> None:
50
+ self._aggregator = KVEventAggregator(num_workers)
51
+
52
+ def add_events(self, events: list[KVCacheEvent]) -> None:
53
+ self._aggregator.add_events(events)
54
+
55
+ def aggregate(self) -> "MooncakeStoreKVEvents":
56
+ common_events = self._aggregator.get_common_events()
57
+ self._aggregator.clear_events()
58
+ self._aggregator.add_events(common_events)
59
+ self._aggregator.reset_workers()
60
+ return self
61
+
62
+ def increment_workers(self, count: int = 1) -> None:
63
+ self._aggregator.increment_workers(count)
64
+
65
+ def get_all_events(self) -> list[KVCacheEvent]:
66
+ return self._aggregator.get_all_events()
67
+
68
+ def get_number_of_workers(self) -> int:
69
+ return self._aggregator.get_number_of_workers()
70
+
71
+ def clear_events(self) -> None:
72
+ self._aggregator.clear_events()
73
+ self._aggregator.reset_workers()
74
+
75
+ def __repr__(self) -> str:
76
+ return f"<MooncakeStoreKVEvents events={self.get_all_events()}>"
77
+
78
+
79
+class MooncakeStoreConnector(KVConnectorBase_V1):
80
+ """KV connector using MooncakeDistributedStore as shared KV pool."""
81
+
82
+ @property
83
+ def prefer_cross_layer_blocks(self) -> bool:
84
+ extra_config = self._kv_transfer_config.kv_connector_extra_config
85
+ return (
86
+ str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
87
+ == "true"
88
+ )
89
+
90
+ def __init__(
91
+ self,
92
+ vllm_config: VllmConfig,
93
+ role: KVConnectorRole,
94
+ kv_cache_config: KVCacheConfig | None = None,
95
+ ):
96
+ super().__init__(
97
+ vllm_config=vllm_config,
98
+ role=role,
99
+ kv_cache_config=kv_cache_config, # type: ignore[arg-type]
100
+ )
101
+ assert vllm_config.kv_transfer_config is not None
102
+ self.kv_role = vllm_config.kv_transfer_config.kv_role
103
+ self._kv_cache_events: MooncakeStoreKVEvents | None = None
104
+
105
+ self.connector_scheduler: MooncakeStoreScheduler | None = None
106
+ self.connector_worker: MooncakeStoreWorker | None = None
107
+
108
+ if role == KVConnectorRole.SCHEDULER:
109
+ self.connector_scheduler = MooncakeStoreScheduler(vllm_config)
110
+ else:
111
+ self.connector_worker = MooncakeStoreWorker(vllm_config)
112
+
113
+ # ============================================================
114
+ # Scheduler-side methods
115
+ # ============================================================
116
+
117
+ def get_num_new_matched_tokens(
118
+ self,
119
+ request: Request,
120
+ num_computed_tokens: int,
121
+ ) -> tuple[int, bool]:
122
+ assert self.connector_scheduler is not None
123
+ return self.connector_scheduler.get_num_new_matched_tokens(
124
+ request, num_computed_tokens
125
+ )
126
+
127
+ def update_state_after_alloc(
128
+ self,
129
+ request: Request,
130
+ blocks: KVCacheBlocks,
131
+ num_external_tokens: int,
132
+ ):
133
+ assert self.connector_scheduler is not None
134
+ return self.connector_scheduler.update_state_after_alloc(
135
+ request, blocks, num_external_tokens
136
+ )
137
+
138
+ def build_connector_meta(
139
+ self,
140
+ scheduler_output: SchedulerOutput,
141
+ ) -> KVConnectorMetadata:
142
+ assert self.connector_scheduler is not None
143
+ return self.connector_scheduler.build_connector_meta(scheduler_output)
144
+
145
+ def request_finished(
146
+ self,
147
+ request: Request,
148
+ block_ids: list[int],
149
+ ) -> tuple[bool, dict[str, Any] | None]:
150
+ assert self.connector_scheduler is not None
151
+ return self.connector_scheduler.request_finished(request, block_ids)
152
+
153
+ def update_connector_output(self, connector_output: KVConnectorOutput):
154
+ kv_cache_events = connector_output.kv_cache_events
155
+ if not kv_cache_events or not isinstance(
156
+ kv_cache_events, MooncakeStoreKVEvents
157
+ ):
158
+ return
159
+
160
+ if self._kv_cache_events is None:
161
+ self._kv_cache_events = kv_cache_events
162
+ else:
163
+ self._kv_cache_events.add_events(kv_cache_events.get_all_events())
164
+ self._kv_cache_events.increment_workers(
165
+ kv_cache_events.get_number_of_workers()
166
+ )
167
+
168
+ def take_events(self) -> Iterable[KVCacheEvent]:
169
+ if self._kv_cache_events is not None:
170
+ self._kv_cache_events.aggregate()
171
+ yield from self._kv_cache_events.get_all_events()
172
+ self._kv_cache_events.clear_events()
173
+ self._kv_cache_events = None
174
+
175
+ # ============================================================
176
+ # Worker-side methods
177
+ # ============================================================
178
+
179
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
180
+ assert self.connector_worker is not None
181
+ self.connector_worker.register_kv_caches(kv_caches)
182
+
183
+ def register_cross_layers_kv_cache(
184
+ self, kv_cache: torch.Tensor, attn_backend: type
185
+ ):
186
+ assert self.connector_worker is not None
187
+ self.connector_worker.register_cross_layers_kv_caches(kv_cache)
188
+
189
+ def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
190
+ # No-op: loads are issued in get_finished() for compute overlap.
191
+ pass
192
+
193
+ def wait_for_layer_load(self, layer_name: str) -> None:
194
+ # No layerwise support - no-op
195
+ return
196
+
197
+ def save_kv_layer(
198
+ self,
199
+ layer_name: str,
200
+ kv_layer: torch.Tensor,
201
+ attn_metadata: AttentionMetadata,
202
+ **kwargs: Any,
203
+ ) -> None:
204
+ # No layerwise support - no-op
205
+ return
206
+
207
+ def wait_for_save(self):
208
+ # No-op: stores are issued in get_finished() for compute overlap.
209
+ pass
210
+
211
+ def get_finished(
212
+ self, finished_req_ids: set[str]
213
+ ) -> tuple[set[str] | None, set[str] | None]:
214
+ assert self.connector_worker is not None
215
+ metadata = self._get_connector_metadata()
216
+ assert isinstance(metadata, MooncakeStoreConnectorMetadata)
217
+ return self.connector_worker.get_finished(finished_req_ids, metadata)
218
+
219
+ def get_kv_connector_kv_cache_events(
220
+ self,
221
+ ) -> MooncakeStoreKVEvents | None:
222
+ assert self.connector_worker is not None
223
+ events = self.connector_worker.get_kv_events()
224
+ if not events:
225
+ return None
226
+
227
+ kv_events = MooncakeStoreKVEvents(num_workers=1)
228
+ kv_events.add_events(events)
229
+ return kv_events
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
new
+276
@@ -0,0 +1,276 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#
4
+# Adapted from vllm-project/vllm-ascend
5
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
6
+"""Data classes for MooncakeStoreConnector."""
7
+
8
+from collections.abc import Iterable
9
+from dataclasses import dataclass
10
+
11
+import torch
12
+
13
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
14
+ KVConnectorMetadata,
15
+)
16
+from vllm.logger import init_logger
17
+from vllm.utils.math_utils import cdiv
18
+from vllm.v1.core.kv_cache_utils import BlockHash
19
+
20
+logger = init_logger(__name__)
21
+
22
+
23
+@dataclass
24
+class KeyMetadata:
25
+ """Metadata for constructing pool keys."""
26
+
27
+ model_name: str
28
+ tp_rank: int
29
+ pcp_rank: int
30
+ dcp_rank: int
31
+ pp_rank: int
32
+
33
+
34
+@dataclass(order=True)
35
+class PoolKey:
36
+ """Key for addressing KV cache blocks in the distributed store."""
37
+
38
+ key_metadata: KeyMetadata
39
+ chunk_hash: str
40
+
41
+ def __hash__(self):
42
+ return hash(
43
+ (
44
+ self.key_metadata.model_name,
45
+ self.key_metadata.tp_rank,
46
+ self.key_metadata.pcp_rank,
47
+ self.key_metadata.dcp_rank,
48
+ self.key_metadata.pp_rank,
49
+ self.chunk_hash,
50
+ )
51
+ )
52
+
53
+ def to_string(self) -> str:
54
+ return (
55
+ f"{self.key_metadata.model_name}"
56
+ f"@tp_rank:{self.key_metadata.tp_rank}"
57
+ f"@pcp{self.key_metadata.pcp_rank}"
58
+ f"@dcp{self.key_metadata.dcp_rank}"
59
+ f"@pp_rank:{self.key_metadata.pp_rank}"
60
+ f"@{self.chunk_hash}"
61
+ )
62
+
63
+
64
+class ChunkedTokenDatabase:
65
+ """Maps token positions to store keys and GPU memory addresses."""
66
+
67
+ def __init__(self, metadata: KeyMetadata, block_size: int):
68
+ self.metadata = metadata
69
+ self.block_size = block_size
70
+ self.kv_caches_base_addr: list[int] = []
71
+ self.block_len: list[int] = []
72
+
73
+ def _make_key_by_hash(self, chunk_hash: str) -> PoolKey:
74
+ return PoolKey(self.metadata, chunk_hash)
75
+
76
+ def set_kv_caches_base_addr(self, kv_caches_base_addr: list[int]):
77
+ self.kv_caches_base_addr = kv_caches_base_addr
78
+
79
+ def set_block_len(self, block_len: list[int]):
80
+ for length in block_len:
81
+ if length % self.block_size != 0:
82
+ raise ValueError(f"block_len {length} % {self.block_size} != 0")
83
+ self.block_len = block_len
84
+
85
+ def prepare_value(
86
+ self, start: int, end: int, block_ids: list[int]
87
+ ) -> tuple[list[int], list[int], int]:
88
+ """Compute memory addresses and sizes for a token range.
89
+
90
+ Returns:
91
+ (addr_list, size_list, block_id)
92
+ """
93
+ addr_list = []
94
+ size_list = []
95
+ block_id = block_ids[start // self.block_size]
96
+ length = len(self.block_len)
97
+ for index, base_addr in enumerate(self.kv_caches_base_addr):
98
+ addr = base_addr + block_id * self.block_len[index % length]
99
+ size = self.block_len[index % length] // self.block_size * (end - start)
100
+ addr_list.append(addr)
101
+ size_list.append(size)
102
+ return addr_list, size_list, block_id
103
+
104
+ def process_tokens(
105
+ self,
106
+ token_len: int,
107
+ block_hashes: list[BlockHash] | list[str],
108
+ mask_num: int = 0,
109
+ ) -> Iterable[tuple[int, int, PoolKey]]:
110
+ """Process tokens and yield (start_idx, end_idx, pool_key) tuples.
111
+
112
+ Args:
113
+ token_len: Total number of tokens.
114
+ block_hashes: Block hashes for each block.
115
+ mask_num: Number of tokens to skip from the beginning.
116
+ """
117
+ if not block_hashes:
118
+ return
119
+ if not isinstance(block_hashes[0], str):
120
+ block_hashes = [
121
+ h.hex() # type: ignore[union-attr]
122
+ for h in block_hashes
123
+ ]
124
+ for chunk_id, hash_val in enumerate(block_hashes):
125
+ start_idx = chunk_id * self.block_size
126
+ if start_idx >= token_len:
127
+ break
128
+ end_idx = min(start_idx + self.block_size, token_len)
129
+ if start_idx < mask_num:
130
+ continue
131
+ else:
132
+ yield (
133
+ start_idx,
134
+ end_idx,
135
+ self._make_key_by_hash(
136
+ hash_val # type: ignore[arg-type]
137
+ ),
138
+ )
139
+
140
+
141
+@dataclass
142
+class LoadSpec:
143
+ """Specification for loading KV cache from external store."""
144
+
145
+ vllm_cached_tokens: int
146
+ kvpool_cached_tokens: int
147
+ can_load: bool
148
+ token_len: int = 0
149
+
150
+
151
+@dataclass
152
+class RequestTracker:
153
+ """Tracks per-request state across scheduler ticks."""
154
+
155
+ req_id: str
156
+ token_len: int
157
+ allocated_block_ids: list[int]
158
+ num_saved_tokens: int = 0
159
+ token_ids: list[int] | None = None
160
+ # Snapshot of the prefill range length at tracker creation time.
161
+ # For a fresh request this is len(prompt). For a resumed-from-preemption
162
+ # request it includes previously-generated tokens, which are re-prefilled.
163
+ prefill_end_tokens: int = 0
164
+
165
+ def update(
166
+ self,
167
+ new_block_ids: tuple[list[int], ...] | list[int],
168
+ ) -> None:
169
+ if len(new_block_ids) == 0:
170
+ new_block_ids = []
171
+ elif isinstance(new_block_ids, tuple):
172
+ new_block_ids = new_block_ids[0]
173
+ elif isinstance(new_block_ids, list):
174
+ pass
175
+ else:
176
+ raise ValueError(f"Unsupported new_block_ids type {type(new_block_ids)}")
177
+ self.allocated_block_ids.extend(new_block_ids)
178
+
179
+
180
+@dataclass
181
+class ReqMeta:
182
+ """Per-request metadata for store put/get operations."""
183
+
184
+ req_id: str
185
+ token_len_chunk: int
186
+ block_ids: list[int]
187
+ block_hashes: list[BlockHash]
188
+
189
+ can_save: bool | None = None
190
+ load_spec: LoadSpec | None = None
191
+ is_last_chunk: bool | None = None
192
+ current_event: torch.cuda.Event | None = None
193
+
194
+ token_ids: list[int] | None = None
195
+ original_block_size: int | None = None
196
+
197
+ @staticmethod
198
+ def from_request_tracker(
199
+ tracker: RequestTracker,
200
+ block_size: int,
201
+ load_spec: LoadSpec | None = None,
202
+ skip_save: bool | None = False,
203
+ block_hashes: list[BlockHash] | None = None,
204
+ is_last_chunk: bool | None = None,
205
+ discard_partial_chunks: bool = True,
206
+ original_block_size: int | None = None,
207
+ ) -> "ReqMeta | None":
208
+ """Create ReqMeta from a RequestTracker."""
209
+ if block_hashes is None:
210
+ block_hashes = []
211
+ input_token_len = tracker.token_len
212
+
213
+ chunk_boundary = (
214
+ cdiv(tracker.num_saved_tokens + 1, block_size) * block_size
215
+ if discard_partial_chunks
216
+ else 0
217
+ )
218
+ num_tokens_to_save = (
219
+ (input_token_len // block_size * block_size)
220
+ if discard_partial_chunks
221
+ else input_token_len
222
+ )
223
+
224
+ skip_save = skip_save or num_tokens_to_save < chunk_boundary
225
+ if skip_save and load_spec is None:
226
+ return None
227
+
228
+ if not skip_save:
229
+ tracker.num_saved_tokens = num_tokens_to_save
230
+
231
+ token_ids = None
232
+ if tracker.token_ids:
233
+ token_ids = tracker.token_ids
234
+
235
+ if load_spec is not None and load_spec.can_load:
236
+ logger.debug(
237
+ "Scheduled to load %d tokens for request %s",
238
+ load_spec.kvpool_cached_tokens,
239
+ tracker.req_id,
240
+ )
241
+ else:
242
+ load_spec = None
243
+
244
+ logger.debug(
245
+ "request:%s, meta save spec:%s, meta load spec:%s",
246
+ tracker.req_id,
247
+ not skip_save,
248
+ load_spec,
249
+ )
250
+ return ReqMeta(
251
+ req_id=tracker.req_id,
252
+ token_len_chunk=num_tokens_to_save,
253
+ block_ids=tracker.allocated_block_ids,
254
+ can_save=not skip_save,
255
+ load_spec=load_spec,
256
+ block_hashes=block_hashes,
257
+ is_last_chunk=is_last_chunk,
258
+ token_ids=token_ids,
259
+ original_block_size=original_block_size,
260
+ )
261
+
262
+
263
+class MooncakeStoreConnectorMetadata(KVConnectorMetadata):
264
+ """Metadata passed from scheduler to worker."""
265
+
266
+ def __init__(
267
+ self,
268
+ unfinished_request_ids: set[str],
269
+ preempted_req_ids: set[str],
270
+ ):
271
+ self.requests: list[ReqMeta] = []
272
+ self.unfinished_request_ids = unfinished_request_ids
273
+ self.preempted_req_ids = preempted_req_ids
274
+
275
+ def add_request(self, req_meta: ReqMeta) -> None:
276
+ self.requests.append(req_meta)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
new
+380
@@ -0,0 +1,380 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#
4
+# Adapted from vllm-project/vllm-ascend
5
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
6
+"""Scheduler-side logic for MooncakeStoreConnector."""
7
+
8
+from typing import Any
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
12
+ KVConnectorMetadata,
13
+)
14
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
15
+ LoadSpec,
16
+ MooncakeStoreConnectorMetadata,
17
+ ReqMeta,
18
+ RequestTracker,
19
+)
20
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import ( # noqa: E501
21
+ LookupKeyClient,
22
+)
23
+from vllm.logger import init_logger
24
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
25
+from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput
26
+from vllm.v1.request import Request
27
+
28
+logger = init_logger(__name__)
29
+
30
+
31
+def _new_req_prefill_tokens(request: NewRequestData) -> list[int]:
32
+ """Tokens this prefill will compute KV for.
33
+
34
+ Under the v2 model runner, resumed-from-preemption requests appear in
35
+ ``scheduled_new_reqs`` with ``prefill_token_ids`` set to the request's full
36
+ token list (prompt + previously-generated). For all other cases this falls
37
+ back to the original prompt.
38
+ """
39
+ if request.prefill_token_ids is not None:
40
+ return request.prefill_token_ids
41
+ assert request.prompt_token_ids is not None
42
+ return request.prompt_token_ids
43
+
44
+
45
+class MooncakeStoreScheduler:
46
+ """Scheduler-side component for MooncakeStoreConnector."""
47
+
48
+ def __init__(self, vllm_config: VllmConfig):
49
+ assert vllm_config.kv_transfer_config is not None
50
+ self.kv_role = vllm_config.kv_transfer_config.kv_role
51
+ self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
52
+ "load_async", True
53
+ )
54
+ self.client = LookupKeyClient(vllm_config)
55
+
56
+ self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
57
+ self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
58
+ self.original_block_size = vllm_config.cache_config.block_size
59
+ self._block_size = vllm_config.cache_config.block_size
60
+ if self.pcp_size > 1:
61
+ self._block_size *= self.pcp_size
62
+ if self.dcp_size > 1:
63
+ self._block_size *= self.dcp_size
64
+
65
+ self._discard_partial_chunks = (
66
+ vllm_config.kv_transfer_config.get_from_extra_config(
67
+ "discard_partial_chunks", True
68
+ )
69
+ )
70
+
71
+ # Per-request state
72
+ self.load_specs: dict[str, LoadSpec] = {} # to be loaded
73
+ self._request_trackers: dict[str, RequestTracker] = {} # scheduled new requests
74
+ self._preempted_req_ids: set[str] = set() # preempted requests
75
+ self._unfinished_requests: dict[str, tuple[Request, list[int]]] = {}
76
+ self._unfinished_request_ids: set[str] = set()
77
+
78
+ def get_num_new_matched_tokens(
79
+ self,
80
+ request: Request,
81
+ num_computed_tokens: int,
82
+ ) -> tuple[int, bool]:
83
+ """Check for external KV cache hit."""
84
+ # Look up against the full prefill range, not just the prompt.
85
+ if self._discard_partial_chunks:
86
+ token_len = request.num_tokens // self._block_size * self._block_size
87
+ else:
88
+ token_len = request.num_tokens
89
+
90
+ if token_len < self._block_size:
91
+ return 0, False
92
+
93
+ num_external_hit_tokens = self.client.lookup(token_len, request.block_hashes)
94
+
95
+ if num_external_hit_tokens == request.num_tokens:
96
+ num_external_hit_tokens -= 1
97
+
98
+ if num_external_hit_tokens < num_computed_tokens:
99
+ need_to_allocate = 0
100
+ else:
101
+ need_to_allocate = num_external_hit_tokens - num_computed_tokens
102
+
103
+ logger.debug(
104
+ "Reqid: %s, Total tokens %d, kvpool hit tokens: %d, need to load: %d",
105
+ request.request_id,
106
+ request.num_tokens,
107
+ num_external_hit_tokens,
108
+ need_to_allocate,
109
+ )
110
+
111
+ if need_to_allocate <= 0:
112
+ return 0, False
113
+
114
+ self.load_specs[request.request_id] = LoadSpec(
115
+ vllm_cached_tokens=num_computed_tokens,
116
+ kvpool_cached_tokens=num_external_hit_tokens,
117
+ can_load=False,
118
+ )
119
+
120
+ return need_to_allocate, self.load_async
121
+
122
+ def update_state_after_alloc(
123
+ self,
124
+ request: Request,
125
+ blocks: KVCacheBlocks,
126
+ num_external_tokens: int,
127
+ ):
128
+ """Update state after block allocation."""
129
+ local_block_ids: list[int] = []
130
+ if num_external_tokens > 0:
131
+ local_block_ids = blocks.get_block_ids()[0]
132
+
133
+ self._unfinished_requests[request.request_id] = (request, local_block_ids)
134
+ self._unfinished_request_ids.add(request.request_id)
135
+
136
+ if request.request_id not in self.load_specs:
137
+ return
138
+
139
+ if num_external_tokens == 0:
140
+ self.load_specs[request.request_id].can_load = False
141
+ return
142
+
143
+ assert (
144
+ num_external_tokens > 0
145
+ and num_external_tokens
146
+ == self.load_specs[request.request_id].kvpool_cached_tokens
147
+ - self.load_specs[request.request_id].vllm_cached_tokens
148
+ ), (
149
+ f"Mismatch in number of tokens: {num_external_tokens} vs "
150
+ f"{self.load_specs[request.request_id].kvpool_cached_tokens} - "
151
+ f"{self.load_specs[request.request_id].vllm_cached_tokens}"
152
+ f" for request {request.request_id}"
153
+ )
154
+
155
+ self.load_specs[request.request_id].can_load = True
156
+
157
+ def build_connector_meta(
158
+ self, scheduler_output: SchedulerOutput
159
+ ) -> KVConnectorMetadata:
160
+ """Build connector metadata for this scheduler step."""
161
+ force_skip_save = self.kv_role == "kv_consumer"
162
+
163
+ for finished_req_id in scheduler_output.finished_req_ids:
164
+ self.load_specs.pop(finished_req_id, None)
165
+ self._request_trackers.pop(finished_req_id, None)
166
+ self._unfinished_requests.pop(finished_req_id, None)
167
+ self._unfinished_request_ids.discard(finished_req_id)
168
+ self._preempted_req_ids.discard(finished_req_id)
169
+
170
+ preempted_ids = scheduler_output.preempted_req_ids or set()
171
+ self._preempted_req_ids.update(preempted_ids)
172
+ for req_id in preempted_ids:
173
+ self._request_trackers.pop(req_id, None)
174
+ self._unfinished_requests.pop(req_id, None)
175
+
176
+ meta = MooncakeStoreConnectorMetadata(
177
+ self._unfinished_request_ids,
178
+ preempted_ids,
179
+ )
180
+
181
+ # Handle new requests
182
+ for request in scheduler_output.scheduled_new_reqs:
183
+ load_spec = self.load_specs.pop(request.req_id, None)
184
+ num_tokens_to_compute = (
185
+ request.num_computed_tokens
186
+ + scheduler_output.num_scheduled_tokens[request.req_id]
187
+ )
188
+ assert request.req_id in self._unfinished_requests
189
+ request_tuple = self._unfinished_requests.get(request.req_id)
190
+ request_real = request_tuple[0] # type: ignore[index]
191
+
192
+ if not isinstance(request.block_ids[0], list):
193
+ unfolded_block_ids = request.block_ids.copy()
194
+ else:
195
+ # TODO: support HMA
196
+ unfolded_block_ids = request.block_ids[0].copy()
197
+
198
+ prefill_tokens = _new_req_prefill_tokens(request)
199
+ request_tracker = RequestTracker(
200
+ req_id=request.req_id,
201
+ token_len=num_tokens_to_compute,
202
+ allocated_block_ids=unfolded_block_ids,
203
+ num_saved_tokens=0,
204
+ token_ids=prefill_tokens[:num_tokens_to_compute],
205
+ prefill_end_tokens=len(prefill_tokens),
206
+ )
207
+ self._request_trackers[request.req_id] = request_tracker
208
+
209
+ last_chunk_tokens_num = (
210
+ (len(prefill_tokens) // self._block_size * self._block_size)
211
+ if self._discard_partial_chunks
212
+ else len(prefill_tokens)
213
+ )
214
+
215
+ req_meta = ReqMeta.from_request_tracker(
216
+ request_tracker,
217
+ self._block_size,
218
+ load_spec=load_spec,
219
+ skip_save=force_skip_save,
220
+ block_hashes=request_real.block_hashes,
221
+ is_last_chunk=(request_tracker.token_len >= last_chunk_tokens_num),
222
+ discard_partial_chunks=self._discard_partial_chunks,
223
+ original_block_size=self.original_block_size,
224
+ )
225
+ if req_meta is not None:
226
+ meta.add_request(req_meta)
227
+
228
+ # Handle cached (running, or MRV1 resumed-from-preemption) requests
229
+ cached_reqs = scheduler_output.scheduled_cached_reqs
230
+ if not force_skip_save:
231
+ for i, req_id in enumerate(cached_reqs.req_ids):
232
+ new_block_ids = cached_reqs.new_block_ids[i]
233
+ if not new_block_ids:
234
+ continue
235
+
236
+ req_meta = None
237
+ if req_id in self._preempted_req_ids:
238
+ # Resumed after preemption
239
+ if isinstance(new_block_ids, tuple):
240
+ block_ids_list = new_block_ids[0].copy()
241
+ else:
242
+ block_ids_list = new_block_ids.copy()
243
+ self._preempted_req_ids.discard(req_id)
244
+ load_spec = self.load_specs.pop(req_id, None)
245
+ request_tuple = self._unfinished_requests.get(req_id)
246
+ request_real = request_tuple[0] # type: ignore[index]
247
+ num_tokens_to_compute = (
248
+ request_real.num_computed_tokens
249
+ + scheduler_output.num_scheduled_tokens[req_id]
250
+ )
251
+ # On resume, the request re-prefills prompt + previously
252
+ # generated tokens (all_token_ids).
253
+ prefill_tokens = list(request_real.all_token_ids)
254
+ request_tracker = RequestTracker(
255
+ req_id=req_id,
256
+ token_len=num_tokens_to_compute,
257
+ allocated_block_ids=block_ids_list,
258
+ num_saved_tokens=0,
259
+ token_ids=prefill_tokens[:num_tokens_to_compute].copy(),
260
+ prefill_end_tokens=len(prefill_tokens),
261
+ )
262
+ self._request_trackers[req_id] = request_tracker
263
+
264
+ last_chunk_tokens_num = (
265
+ (len(prefill_tokens) // self._block_size * self._block_size)
266
+ if self._discard_partial_chunks
267
+ else len(prefill_tokens)
268
+ )
269
+ req_meta = ReqMeta.from_request_tracker(
270
+ request_tracker,
271
+ self._block_size,
272
+ load_spec=load_spec,
273
+ skip_save=force_skip_save,
274
+ block_hashes=request_real.block_hashes,
275
+ is_last_chunk=(
276
+ request_tracker.token_len >= last_chunk_tokens_num
277
+ ),
278
+ discard_partial_chunks=self._discard_partial_chunks,
279
+ original_block_size=self.original_block_size,
280
+ )
281
+ else:
282
+ # Decode/chunked request
283
+ request_tracker = self._request_trackers[req_id]
284
+ num_new_tokens = scheduler_output.num_scheduled_tokens[req_id]
285
+ req_tuple = self._unfinished_requests.get(req_id)
286
+ if req_tuple:
287
+ unfinished_req = req_tuple[0]
288
+ num_current_tokens = request_tracker.token_len
289
+ new_token_ids = unfinished_req.all_token_ids[
290
+ num_current_tokens : num_current_tokens + num_new_tokens
291
+ ]
292
+ request_tracker.token_len += len(new_token_ids)
293
+ else:
294
+ raise ValueError(
295
+ f"Request {req_id} is not in _unfinished_requests"
296
+ )
297
+ num_computed_token = cached_reqs.num_computed_tokens[i]
298
+ # Use the tracker's snapshot of the prefill range so resumed
299
+ # requests keep saving past the original prompt boundary.
300
+ prefill_end = request_tracker.prefill_end_tokens
301
+ if num_computed_token >= prefill_end:
302
+ continue
303
+ request_tracker.update(new_block_ids)
304
+
305
+ last_chunk_tokens_num = (
306
+ (prefill_end // self._block_size * self._block_size)
307
+ if self._discard_partial_chunks
308
+ else prefill_end
309
+ )
310
+ req_meta = ReqMeta.from_request_tracker(
311
+ request_tracker,
312
+ self._block_size,
313
+ load_spec=None,
314
+ skip_save=force_skip_save,
315
+ block_hashes=unfinished_req.block_hashes,
316
+ is_last_chunk=(
317
+ request_tracker.token_len >= last_chunk_tokens_num
318
+ ),
319
+ discard_partial_chunks=self._discard_partial_chunks,
320
+ original_block_size=self.original_block_size,
321
+ )
322
+
323
+ if req_meta is not None:
324
+ meta.add_request(req_meta)
325
+
326
+ # Handle requests with pending load specs not yet scheduled
327
+ request_ids = [req.req_id for req in scheduler_output.scheduled_new_reqs]
328
+ for request_id, (
329
+ unfinished_req,
330
+ block_ids,
331
+ ) in self._unfinished_requests.items():
332
+ if request_id not in request_ids and request_id not in cached_reqs.req_ids:
333
+ load_spec = self.load_specs.pop(request_id, None)
334
+ if not load_spec:
335
+ continue
336
+ num_tokens_to_compute = load_spec.kvpool_cached_tokens
337
+ if (num_tokens_to_compute % self._block_size != 0) and (
338
+ num_tokens_to_compute == unfinished_req.num_tokens - 1
339
+ ):
340
+ num_tokens_to_compute = num_tokens_to_compute + 1
341
+ request_tracker = RequestTracker(
342
+ req_id=request_id,
343
+ token_len=num_tokens_to_compute,
344
+ allocated_block_ids=block_ids,
345
+ num_saved_tokens=0,
346
+ )
347
+ self._request_trackers[request_id] = request_tracker
348
+ req_meta = ReqMeta.from_request_tracker(
349
+ request_tracker,
350
+ self._block_size,
351
+ load_spec=load_spec,
352
+ skip_save=None,
353
+ block_hashes=unfinished_req.block_hashes,
354
+ discard_partial_chunks=self._discard_partial_chunks,
355
+ )
356
+ if req_meta is not None:
357
+ meta.add_request(req_meta)
358
+
359
+ return meta
360
+
361
+ def request_finished(
362
+ self,
363
+ request: Request,
364
+ block_ids: list[int],
365
+ ) -> tuple[bool, dict[str, Any] | None]:
366
+ """Determine whether to delay freeing blocks for async save."""
367
+ if self.kv_role == "kv_consumer":
368
+ return False, None
369
+ tracker = self._request_trackers.get(request.request_id)
370
+ assert tracker is not None
371
+ if tracker.num_saved_tokens <= 0:
372
+ return False, None
373
+ delay_free_blocks = len(block_ids) > 0
374
+ if delay_free_blocks:
375
+ logger.debug(
376
+ "Delaying free of %d blocks for request %s",
377
+ len(block_ids),
378
+ request.request_id,
379
+ )
380
+ return delay_free_blocks, None
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
new
+979
@@ -0,0 +1,979 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#
4
+# The transfer-thread scaffolding (KVTransferThread, KVCacheStoreSendingThread,
5
+# KVCacheStoreRecvingThread) is adapted from vllm-project/vllm-ascend
6
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
7
+"""Worker-side logic for MooncakeStoreConnector.
8
+
9
+Includes the store worker, transfer threads, lookup server,
10
+and MooncakeDistributedStore integration.
11
+"""
12
+
13
+import json
14
+import os
15
+import queue
16
+import threading
17
+from collections import defaultdict
18
+from dataclasses import dataclass
19
+from typing import Any
20
+
21
+import regex as re
22
+import torch
23
+import zmq
24
+
25
+import vllm.envs as envs
26
+from vllm.config import VllmConfig
27
+from vllm.distributed import (
28
+ get_dcp_group,
29
+ get_pcp_group,
30
+ get_tensor_model_parallel_rank,
31
+ get_tensor_model_parallel_world_size,
32
+)
33
+from vllm.distributed.kv_events import BlockStored
34
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
35
+ get_mooncake_dp_engine_index,
36
+)
37
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
38
+ ChunkedTokenDatabase,
39
+ KeyMetadata,
40
+ MooncakeStoreConnectorMetadata,
41
+ ReqMeta,
42
+)
43
+from vllm.logger import init_logger
44
+from vllm.utils.network_utils import get_ip, make_zmq_socket
45
+from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
46
+from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
47
+
48
+logger = init_logger(__name__)
49
+
50
+DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
51
+DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
52
+MOONCAKE_NO_AVAILABLE_HANDLE = -200
53
+
54
+
55
+@dataclass
56
+class MooncakeStoreConfig:
57
+ """Configuration for MooncakeDistributedStore."""
58
+
59
+ metadata_server: str
60
+ global_segment_size: int
61
+ local_buffer_size: int
62
+ protocol: str
63
+ device_name: str
64
+ master_server_address: str
65
+
66
+ @staticmethod
67
+ def from_file(file_path: str) -> "MooncakeStoreConfig":
68
+ with open(file_path) as file:
69
+ config = json.load(file)
70
+ return MooncakeStoreConfig(
71
+ metadata_server=config.get("metadata_server", ""),
72
+ global_segment_size=_parse_size(
73
+ config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
74
+ ),
75
+ local_buffer_size=_parse_size(
76
+ config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)
77
+ ),
78
+ protocol=config.get("protocol", "rdma"),
79
+ device_name=config.get("device_name", ""),
80
+ master_server_address=config.get("master_server_address", ""),
81
+ )
82
+
83
+ @staticmethod
84
+ def load_from_env() -> "MooncakeStoreConfig":
85
+ config_path = os.getenv("MOONCAKE_CONFIG_PATH")
86
+ if not config_path:
87
+ raise ValueError(
88
+ "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
89
+ )
90
+ return MooncakeStoreConfig.from_file(config_path)
91
+
92
+
93
+def _parse_size(value: Any) -> int:
94
+ """Parse storage size strings with units: GB, MB, KB, B."""
95
+ if isinstance(value, int):
96
+ return value
97
+ if not isinstance(value, str):
98
+ try:
99
+ return int(value)
100
+ except (TypeError, ValueError) as e:
101
+ raise TypeError(f"Unsupported type for size: {type(value)}") from e
102
+
103
+ cleaned = value.strip().lower()
104
+ if not cleaned:
105
+ raise ValueError("Size cannot be empty.")
106
+
107
+ unit_multipliers = {
108
+ "gb": 1024**3,
109
+ "mb": 1024**2,
110
+ "kb": 1024,
111
+ "b": 1,
112
+ }
113
+ match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned)
114
+ if not match:
115
+ raise ValueError(f"Invalid format: '{value}'")
116
+
117
+ number_str = match.group(1)
118
+ unit = match.group(2) or "b"
119
+ multiplier = unit_multipliers[unit]
120
+
121
+ try:
122
+ numeric_value = float(number_str)
123
+ except ValueError as exc:
124
+ raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc
125
+ return int(numeric_value * multiplier)
126
+
127
+
128
+# ============================================================
129
+# Transfer Threads
130
+# ============================================================
131
+
132
+
133
+class KVTransferThread(threading.Thread):
134
+ """Base class for async KV cache transfer threads."""
135
+
136
+ def __init__(
137
+ self,
138
+ store: Any,
139
+ token_database: ChunkedTokenDatabase,
140
+ block_size: int,
141
+ tp_rank: int,
142
+ ready_event: threading.Event,
143
+ name: str,
144
+ ):
145
+ super().__init__(daemon=True, name=name)
146
+ self.store = store
147
+ self.ready_event = ready_event
148
+ self.block_size = block_size
149
+ self.tp_rank = tp_rank
150
+ self.token_database = token_database
151
+ self.done_task_lock = threading.Lock()
152
+ self.request_queue: queue.Queue[Any] = queue.Queue()
153
+ self.finished_requests: set[str] = set()
154
+ self.kv_event_lock = threading.Lock()
155
+ self.kv_events: list[BlockStored] = []
156
+
157
+ def add_request(self, request: ReqMeta) -> None:
158
+ self.request_queue.put(request)
159
+
160
+ def get_and_clear_finished_requests(self) -> set[str]:
161
+ with self.done_task_lock:
162
+ finished = self.finished_requests.copy()
163
+ self.finished_requests.clear()
164
+ return finished
165
+
166
+ def set_finished_request(self, req_id: str):
167
+ with self.done_task_lock:
168
+ self.finished_requests.add(req_id)
169
+
170
+ def run(self):
171
+ self.ready_event.set()
172
+ while True:
173
+ try:
174
+ request_data = self.request_queue.get()
175
+ if request_data is None:
176
+ logger.warning("Received a None request!")
177
+ self.request_queue.task_done()
178
+ continue
179
+ self._handle_request(request_data)
180
+ except Exception as e:
181
+ logger.error("Error in %s: %s", self.name, e)
182
+
183
+ def _handle_request(self, req_meta: Any):
184
+ pass
185
+
186
+ def update_kv_event(self, events: list[BlockStored]):
187
+ with self.kv_event_lock:
188
+ self.kv_events.extend(events)
189
+
190
+ def get_kv_events(self) -> list[BlockStored]:
191
+ with self.kv_event_lock:
192
+ events = self.kv_events.copy()
193
+ self.kv_events.clear()
194
+ return events
195
+
196
+
197
+class KVCacheStoreSendingThread(KVTransferThread):
198
+ """Background thread for storing KV cache blocks to the store."""
199
+
200
+ def __init__(
201
+ self,
202
+ store: Any,
203
+ token_database: ChunkedTokenDatabase,
204
+ block_size: int,
205
+ tp_rank: int,
206
+ put_step: int,
207
+ kv_role: str,
208
+ ready_event: threading.Event,
209
+ enable_kv_event: bool = False,
210
+ ):
211
+ super().__init__(
212
+ store,
213
+ token_database,
214
+ block_size,
215
+ tp_rank,
216
+ ready_event,
217
+ name="KVCacheStoreSendingThread",
218
+ )
219
+ self.put_step = put_step
220
+ self.kv_role = kv_role
221
+ self.stored_requests: defaultdict[str, int] = defaultdict(int)
222
+ self.enable_kv_event = enable_kv_event
223
+
224
+ # Pause store requests when CPU offloading is under pressure.
225
+ self._store_pressure_active = False
226
+ self._skip_store_requests: set[str] = set()
227
+
228
+ def add_stored_request(self, req_id: str):
229
+ with self.done_task_lock:
230
+ self.stored_requests[req_id] += 1
231
+
232
+ def dec_stored_request(self, req_id: str):
233
+ with self.done_task_lock:
234
+ if req_id in self.stored_requests:
235
+ self.stored_requests[req_id] -= 1
236
+
237
+ def delete_finished_stored_request(self, req_id: str):
238
+ with self.done_task_lock:
239
+ if req_id in self.stored_requests:
240
+ del self.stored_requests[req_id]
241
+ self._skip_store_requests.discard(req_id)
242
+
243
+ def _should_skip_request(self, req_id: str) -> bool:
244
+ with self.done_task_lock:
245
+ return self._store_pressure_active and req_id in self._skip_store_requests
246
+
247
+ def _mark_request_skipped_for_pressure(self, req_id: str) -> bool:
248
+ with self.done_task_lock:
249
+ already_skipped = req_id in self._skip_store_requests
250
+ self._store_pressure_active = True
251
+ self._skip_store_requests.add(req_id)
252
+ return already_skipped
253
+
254
+ def _clear_store_pressure(self) -> bool:
255
+ with self.done_task_lock:
256
+ if not self._store_pressure_active and not self._skip_store_requests:
257
+ return False
258
+ self._store_pressure_active = False
259
+ self._skip_store_requests.clear()
260
+ return True
261
+
262
+ def _handle_request(self, req_meta: ReqMeta):
263
+ token_len = req_meta.token_len_chunk
264
+ block_ids = req_meta.block_ids
265
+ req_id = req_meta.req_id
266
+ current_event = req_meta.current_event
267
+
268
+ if req_id not in self.stored_requests:
269
+ self.request_queue.task_done()
270
+ return
271
+ if self._should_skip_request(req_id):
272
+ logger.debug(
273
+ "Skipping Mooncake store for request %s while CPU offloading "
274
+ "is under pressure",
275
+ req_id,
276
+ )
277
+ self.dec_stored_request(req_id)
278
+ self.request_queue.task_done()
279
+ return
280
+
281
+ starts = []
282
+ ends = []
283
+ keys = []
284
+ block_hashes: list[BlockHash] = []
285
+ for index, (start, end, key) in enumerate(
286
+ self.token_database.process_tokens(token_len, req_meta.block_hashes)
287
+ ):
288
+ starts.append(start)
289
+ ends.append(end)
290
+ keys.append(key.to_string())
291
+ block_hashes.append(req_meta.block_hashes[index])
292
+
293
+ # Apply put_step striding for TP
294
+ starts = starts[self.tp_rank % self.put_step :: self.put_step]
295
+ ends = ends[self.tp_rank % self.put_step :: self.put_step]
296
+ keys = keys[self.tp_rank % self.put_step :: self.put_step]
297
+ block_hashes = block_hashes[self.tp_rank % self.put_step :: self.put_step]
298
+
299
+ if not keys:
300
+ self.dec_stored_request(req_id)
301
+ return
302
+
303
+ # Check which blocks already exist (dedup)
304
+ exists_states = self.store.batch_is_exist(keys)
305
+ missing_indices = [i for i, exists in enumerate(exists_states) if exists != 1]
306
+
307
+ if not missing_indices:
308
+ self.dec_stored_request(req_id)
309
+ return
310
+
311
+ starts = [starts[i] for i in missing_indices]
312
+ ends = [ends[i] for i in missing_indices]
313
+ keys = [keys[i] for i in missing_indices]
314
+ block_hashes = [block_hashes[i] for i in missing_indices]
315
+
316
+ logger.debug(
317
+ "Storing KV cache for %d out of %d blocks "
318
+ "(missing_count=%d) for request %s",
319
+ len(keys),
320
+ token_len // self.block_size,
321
+ len(missing_indices),
322
+ req_id,
323
+ )
324
+
325
+ addrs = []
326
+ sizes = []
327
+ stored_events: list[BlockStored] = []
328
+ prev_key = None
329
+ new_block_hashes = [maybe_convert_block_hash(bh) for bh in block_hashes]
330
+
331
+ for index, start in enumerate(starts):
332
+ addr, size, _ = self.token_database.prepare_value(
333
+ start, ends[index], block_ids
334
+ )
335
+ addrs.append(addr)
336
+ sizes.append(size)
337
+
338
+ if self.enable_kv_event:
339
+ token_ids = (
340
+ req_meta.token_ids[start : ends[index]]
341
+ if req_meta.token_ids is not None
342
+ else None
343
+ )
344
+ stored_event = BlockStored(
345
+ block_hashes=[new_block_hashes[index]],
346
+ parent_block_hash=prev_key,
347
+ token_ids=token_ids,
348
+ block_size=req_meta.original_block_size,
349
+ lora_id=None,
350
+ medium="cpu",
351
+ lora_name=None,
352
+ )
353
+ stored_events.append(stored_event)
354
+ prev_key = new_block_hashes[index]
355
+
356
+ if current_event is not None:
357
+ current_event.synchronize()
358
+
359
+ try:
360
+ res = self.store.batch_put_from_multi_buffers(keys, addrs, sizes)
361
+ failed = [i for i, v in enumerate(res) if v < 0]
362
+ if failed:
363
+ # Compute total bytes attempted for this batch
364
+ total_bytes = sum(sum(s) if isinstance(s, list) else s for s in sizes)
365
+ failed_codes = set(res[i] for i in failed)
366
+ logger.warning(
367
+ "batch_put failed: %d/%d keys failed "
368
+ "(codes=%s, batch_bytes=%d, num_keys=%d), "
369
+ "first_key=%s",
370
+ len(failed),
371
+ len(keys),
372
+ failed_codes,
373
+ total_bytes,
374
+ len(keys),
375
+ keys[0] if keys else "N/A",
376
+ )
377
+ if (
378
+ MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes
379
+ and not self._mark_request_skipped_for_pressure(req_id)
380
+ ):
381
+ logger.warning(
382
+ "Detected Mooncake CPU offloading pressure "
383
+ "(NO_AVAILABLE_HANDLE); skipping future store "
384
+ "batches for request %s until a later store "
385
+ "batch succeeds",
386
+ req_id,
387
+ )
388
+ elif self._clear_store_pressure():
389
+ logger.info(
390
+ "Mooncake CPU offloading pressure cleared after a "
391
+ "successful store batch"
392
+ )
393
+ except Exception as e:
394
+ logger.error("Failed to put key %s, error: %s", keys, e)
395
+
396
+ if self.enable_kv_event and stored_events:
397
+ self.update_kv_event(stored_events)
398
+
399
+ self.dec_stored_request(req_id)
400
+ self.request_queue.task_done()
401
+
402
+
403
+class KVCacheStoreRecvingThread(KVTransferThread):
404
+ """Background thread for loading KV cache blocks from the store."""
405
+
406
+ def __init__(
407
+ self,
408
+ store: Any,
409
+ token_database: ChunkedTokenDatabase,
410
+ block_size: int,
411
+ tp_rank: int,
412
+ ready_event: threading.Event,
413
+ ):
414
+ super().__init__(
415
+ store,
416
+ token_database,
417
+ block_size,
418
+ tp_rank,
419
+ ready_event,
420
+ name="KVCacheStoreRecvingThread",
421
+ )
422
+
423
+ def _handle_request(self, req_meta: ReqMeta):
424
+ token_len = req_meta.load_spec.token_len # type: ignore[union-attr]
425
+ req_id = req_meta.req_id
426
+ mask_num = (
427
+ req_meta.load_spec.vllm_cached_tokens # type: ignore[union-attr]
428
+ // self.block_size
429
+ * self.block_size
430
+ )
431
+
432
+ addr_list = []
433
+ size_list = []
434
+ key_list = []
435
+ for start, end, key in self.token_database.process_tokens(
436
+ token_len, req_meta.block_hashes, mask_num
437
+ ):
438
+ addr, size, _ = self.token_database.prepare_value(
439
+ start, end, req_meta.block_ids
440
+ )
441
+ key_list.append(key.to_string())
442
+ addr_list.append(addr)
443
+ size_list.append(size)
444
+
445
+ # Rotate lists by tp_rank for load balancing
446
+ key_list_c = (
447
+ key_list[self.tp_rank % len(key_list) :]
448
+ + key_list[: self.tp_rank % len(key_list)]
449
+ )
450
+ addr_list_c = (
451
+ addr_list[self.tp_rank % len(addr_list) :]
452
+ + addr_list[: self.tp_rank % len(addr_list)]
453
+ )
454
+ size_list_c = (
455
+ size_list[self.tp_rank % len(size_list) :]
456
+ + size_list[: self.tp_rank % len(size_list)]
457
+ )
458
+
459
+ try:
460
+ res = self.store.batch_get_into_multi_buffers(
461
+ key_list_c, addr_list_c, size_list_c
462
+ )
463
+ failed = [
464
+ (key, value)
465
+ for key, value in zip(key_list_c, res, strict=True)
466
+ if value < 0
467
+ ]
468
+ if failed:
469
+ logger.warning(
470
+ "Failed to get %d Mooncake keys (batch_keys=%d, first_failures=%s)",
471
+ len(failed),
472
+ len(key_list_c),
473
+ failed[:3],
474
+ )
475
+ except Exception as e:
476
+ logger.warning(
477
+ "Failed to get Mooncake batch %s, error: %s",
478
+ key_list_c[:3],
479
+ e,
480
+ )
481
+
482
+ self.set_finished_request(req_id)
483
+ self.request_queue.task_done()
484
+
485
+
486
+# ============================================================
487
+# Store Worker
488
+# ============================================================
489
+
490
+
491
+class MooncakeStoreWorker:
492
+ """Worker-side component for MooncakeStoreConnector."""
493
+
494
+ def __init__(self, vllm_config: VllmConfig):
495
+ try:
496
+ from mooncake.store import MooncakeDistributedStore # type: ignore
497
+ except ImportError as e:
498
+ raise ImportError(
499
+ "Please install mooncake by following the instructions at "
500
+ "https://github.com/kvcache-ai/Mooncake/blob/main/doc/"
501
+ "en/build.md to run vLLM with MooncakeStoreConnector."
502
+ ) from e
503
+
504
+ model_config = vllm_config.model_config
505
+ parallel_config = vllm_config.parallel_config
506
+
507
+ self.dp_rank = get_mooncake_dp_engine_index(parallel_config)
508
+ self.tp_rank = get_tensor_model_parallel_rank()
509
+ self.tp_size = get_tensor_model_parallel_world_size()
510
+ self.pp_size = parallel_config.pipeline_parallel_size
511
+ self.pp_rank = (parallel_config.rank // self.tp_size) % self.pp_size
512
+
513
+ self.pcp_size = get_pcp_group().world_size
514
+ self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0
515
+ self.dcp_size = get_dcp_group().world_size
516
+ self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_size > 1 else 0
517
+
518
+ assert vllm_config.kv_transfer_config is not None
519
+ self.kv_role = vllm_config.kv_transfer_config.kv_role
520
+ self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
521
+ "load_async", True
522
+ )
523
+ self.cache_config = vllm_config.cache_config
524
+ self.original_block_size = self.cache_config.block_size
525
+ self.block_size = self.cache_config.block_size
526
+ if self.pcp_size > 1:
527
+ self.block_size *= self.pcp_size
528
+ if self.dcp_size > 1:
529
+ self.block_size *= self.dcp_size
530
+ self.num_layers = model_config.get_num_layers(parallel_config)
531
+
532
+ self.use_mla = False
533
+ if (
534
+ hasattr(model_config, "use_mla")
535
+ and isinstance(model_config.use_mla, bool)
536
+ and model_config.use_mla
537
+ ):
538
+ self.use_mla = True
539
+
540
+ if self.use_mla:
541
+ self.num_kv_head = 1
542
+ else:
543
+ self.num_kv_head = model_config.get_total_num_kv_heads()
544
+
545
+ if self.num_kv_head < self.tp_size:
546
+ self.put_step = self.tp_size // self.num_kv_head
547
+ self.head_or_tp_rank = self.tp_rank // self.put_step
548
+ else:
549
+ self.head_or_tp_rank = self.tp_rank
550
+ self.put_step = 1
551
+
552
+ self.metadata = KeyMetadata(
553
+ model_name=model_config.model.rstrip("/").split("/")[-1],
554
+ tp_rank=self.head_or_tp_rank,
555
+ pcp_rank=self.pcp_rank,
556
+ dcp_rank=self.dcp_rank,
557
+ pp_rank=self.pp_rank,
558
+ )
559
+
560
+ self.token_database = ChunkedTokenDatabase(self.metadata, self.block_size)
561
+
562
+ # Initialize MooncakeDistributedStore with its own TransferEngine
563
+ store_config = MooncakeStoreConfig.load_from_env()
564
+ self.store = MooncakeDistributedStore()
565
+
566
+ local_seg = get_ip()
567
+ config_dict = {
568
+ "local_hostname": local_seg,
569
+ "metadata_server": store_config.metadata_server,
570
+ "global_segment_size": str(store_config.global_segment_size),
571
+ "local_buffer_size": str(store_config.local_buffer_size),
572
+ "protocol": store_config.protocol,
573
+ "rdma_devices": store_config.device_name,
574
+ "master_server_addr": store_config.master_server_address,
575
+ }
576
+ ret = self.store.setup(config_dict)
577
+ if ret != 0:
578
+ msg = "Initialize MooncakeDistributedStore failed."
579
+ logger.error(msg)
580
+ raise RuntimeError(msg)
581
+
582
+ kv_event_config = vllm_config.kv_events_config
583
+ self.enable_kv_events = False
584
+ if kv_event_config and kv_event_config.enable_kv_cache_events:
585
+ self.enable_kv_events = True
586
+
587
+ self.kv_send_thread: KVCacheStoreSendingThread | None = None
588
+ self.kv_recv_thread: KVCacheStoreRecvingThread | None = None
589
+ self.finished_store_req: set[str] = set()
590
+
591
+ # Start lookup server on rank 0 for scheduler-side prefix queries
592
+ self.lookup_server: LookupKeyServer | None = None
593
+ if vllm_config.parallel_config.rank == 0:
594
+ self.lookup_server = LookupKeyServer(self, vllm_config)
595
+
596
+ def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
597
+ """Register a cross-layers KV cache tensor.
598
+
599
+ Wraps the unified tensor in a single-entry dict so that the
600
+ existing stride-based logic in register_kv_caches() produces
601
+ the correct single-segment result (block_len = page_size * num_layers).
602
+ """
603
+ self.register_kv_caches({"__cross_layer__": kv_cache})
604
+
605
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
606
+ """Register KV cache tensors and start transfer threads."""
607
+ # TODO(yifan): we haven't supported HMA yet.
608
+ first_kv_cache = next(iter(kv_caches.values()))
609
+
610
+ # num_blocks from cache_config is authoritative (set after
611
+ # profiling, before KV cache allocation).
612
+ assert self.cache_config.num_gpu_blocks is not None
613
+ self.num_blocks = self.cache_config.num_gpu_blocks
614
+
615
+ # Detect the KV cache memory layout using the stride-based
616
+ # approach from simple_kv_offload/worker.py.
617
+ #
618
+ # The physical layout varies across attention backends:
619
+ # FlashAttn/ROCm : (2, num_blocks, ...) → K/V outermost
620
+ # FlashInfer/MLA : (num_blocks, ...) → blocks outermost
621
+ #
622
+ # We derive page_size_bytes = storage.nbytes() // num_blocks,
623
+ # then classify dims: any dim whose byte-stride exceeds
624
+ # page_size_bytes must be an outer segment dim (e.g. the K/V
625
+ # dim of size 2). For those backends we register each segment
626
+ # (K, V) as a separate base-address so that the per-block
627
+ # offset arithmetic in prepare_value() stays correct.
628
+ storage = first_kv_cache.untyped_storage()
629
+ el = first_kv_cache.element_size()
630
+ page_size_bytes = storage.nbytes() // self.num_blocks
631
+ outer_dims = [
632
+ d
633
+ for d in range(first_kv_cache.ndim)
634
+ if first_kv_cache.stride(d) * el > page_size_bytes
635
+ ]
636
+
637
+ # Register buffers with the store (deduplicate shared storages)
638
+ # and record per-segment base addresses for every layer.
639
+ seen_ptrs: set[int] = set()
640
+ self.kv_caches_base_addr: list[int] = []
641
+ self.block_len: list[int] = []
642
+
643
+ for cache in kv_caches.values():
644
+ cache_storage = cache.untyped_storage()
645
+ base_addr = cache_storage.data_ptr()
646
+ region_len = cache_storage.nbytes()
647
+
648
+ if base_addr not in seen_ptrs:
649
+ seen_ptrs.add(base_addr)
650
+ ret = self.store.register_buffer(base_addr, region_len)
651
+ if ret != 0:
652
+ logger.error(
653
+ "register_buffer failed for addr %#x len %d: %d",
654
+ base_addr,
655
+ region_len,
656
+ ret,
657
+ )
658
+
659
+ if not outer_dims:
660
+ # Blocks-first layout (FlashInfer / MLA): one segment.
661
+ self.kv_caches_base_addr.append(base_addr)
662
+ self.block_len.append(page_size_bytes)
663
+ else:
664
+ # K/V-first layout (FlashAttn / ROCm): split segments.
665
+ seg_stride = cache.stride(outer_dims[0]) * el
666
+ for idx in range(cache.shape[outer_dims[0]]):
667
+ self.kv_caches_base_addr.append(base_addr + idx * seg_stride)
668
+ self.block_len.append(seg_stride // self.num_blocks)
669
+
670
+ logger.info(
671
+ "Registering KV_Caches. use_mla: %s, shape %s, "
672
+ "num_blocks: %d, block_len: %s, "
673
+ "per_key_bytes: %d, "
674
+ "num_segments: %d",
675
+ self.use_mla,
676
+ first_kv_cache.shape,
677
+ self.num_blocks,
678
+ list(set(self.block_len)),
679
+ sum(self.block_len),
680
+ len(self.kv_caches_base_addr),
681
+ )
682
+
683
+ self.token_database.set_kv_caches_base_addr(self.kv_caches_base_addr)
684
+ self.token_database.set_block_len(self.block_len)
685
+
686
+ # Start transfer threads
687
+ if self.kv_role in ["kv_producer", "kv_both"]:
688
+ ready_event_sending = threading.Event()
689
+ self.kv_send_thread = KVCacheStoreSendingThread(
690
+ self.store,
691
+ self.token_database,
692
+ self.block_size,
693
+ self.tp_rank,
694
+ self.put_step,
695
+ self.kv_role,
696
+ ready_event_sending,
697
+ self.enable_kv_events,
698
+ )
699
+ self.kv_send_thread.start()
700
+
701
+ ready_event_recving = threading.Event()
702
+ self.kv_recv_thread = KVCacheStoreRecvingThread(
703
+ self.store,
704
+ self.token_database,
705
+ self.block_size,
706
+ self.tp_rank,
707
+ ready_event_recving,
708
+ )
709
+ self.kv_recv_thread.start()
710
+ ready_event_recving.wait()
711
+
712
+ def start_load_kv(
713
+ self,
714
+ metadata: MooncakeStoreConnectorMetadata,
715
+ ):
716
+ """No-op: loads are issued in get_finished() for overlap."""
717
+ pass
718
+
719
+ def wait_for_save(
720
+ self,
721
+ metadata: MooncakeStoreConnectorMetadata,
722
+ ):
723
+ """No-op: stores are issued in get_finished() for overlap."""
724
+ pass
725
+
726
+ def get_finished(
727
+ self,
728
+ finished_req_ids: set[str],
729
+ meta: MooncakeStoreConnectorMetadata,
730
+ ) -> tuple[set[str], set[str]]:
731
+ """Issue all I/O and get completed send/recv request IDs.
732
+
733
+ All load and store I/O requests are issued here (after model
734
+ compute is launched on the compute stream) for better
735
+ compute-I/O overlap.
736
+ """
737
+ # Issue async loads
738
+ for request in meta.requests:
739
+ load_spec = request.load_spec
740
+ if load_spec is None or not load_spec.can_load:
741
+ continue
742
+
743
+ token_len = request.token_len_chunk
744
+ if (load_spec.kvpool_cached_tokens % self.block_size != 0) and (
745
+ load_spec.kvpool_cached_tokens == token_len - 1
746
+ ):
747
+ token_len = load_spec.kvpool_cached_tokens + 1
748
+ else:
749
+ token_len = load_spec.kvpool_cached_tokens
750
+ load_spec.token_len = token_len
751
+
752
+ assert self.kv_recv_thread is not None
753
+ self.kv_recv_thread.add_request(request)
754
+
755
+ assert self.load_async, "load_async must be True for better performance."
756
+ # Issue stores with CUDA event synchronization
757
+ if self.kv_role in ["kv_producer", "kv_both"]:
758
+ current_event = None
759
+ for request in meta.requests:
760
+ if request.can_save:
761
+ current_event = torch.cuda.Event()
762
+ current_event.record()
763
+ break
764
+
765
+ for request in meta.requests:
766
+ if not request.can_save:
767
+ continue
768
+ request.current_event = current_event
769
+ assert self.kv_send_thread is not None
770
+ self.kv_send_thread.add_stored_request(request.req_id)
771
+ self.kv_send_thread.add_request(request)
772
+
773
+ # Check completion of previously queued transfers
774
+ done_sending = (
775
+ self._get_and_clear_finished_sending(finished_req_ids, meta)
776
+ if self.kv_role in ["kv_producer", "kv_both"]
777
+ else set()
778
+ )
779
+
780
+ done_recving = (
781
+ self.kv_recv_thread.get_and_clear_finished_requests()
782
+ if self.load_async and self.kv_recv_thread is not None
783
+ else set()
784
+ )
785
+
786
+ logger.debug(
787
+ "Completed send: %d, recv: %d, tp_rank: %d",
788
+ len(done_sending),
789
+ len(done_recving),
790
+ self.tp_rank,
791
+ )
792
+ return done_sending, done_recving
793
+
794
+ def _get_and_clear_finished_sending(
795
+ self,
796
+ finished_req_ids: set[str],
797
+ meta: MooncakeStoreConnectorMetadata,
798
+ ) -> set[str]:
799
+ assert self.kv_send_thread is not None
800
+ finished_sending: set[str] = set()
801
+
802
+ for req_id in meta.preempted_req_ids:
803
+ self.kv_send_thread.delete_finished_stored_request(req_id)
804
+
805
+ for req_id in self.kv_send_thread.stored_requests.copy():
806
+ if (
807
+ self.kv_send_thread.stored_requests[req_id] == 0
808
+ and req_id in self.finished_store_req
809
+ ):
810
+ self.finished_store_req.remove(req_id)
811
+ finished_sending.add(req_id)
812
+ self.kv_send_thread.delete_finished_stored_request(req_id)
813
+
814
+ for req_id in finished_req_ids:
815
+ req_remain_jobs = self.kv_send_thread.stored_requests.get(req_id)
816
+ if req_remain_jobs == 0:
817
+ finished_sending.add(req_id)
818
+ self.kv_send_thread.delete_finished_stored_request(req_id)
819
+ elif req_remain_jobs is not None:
820
+ self.finished_store_req.add(req_id)
821
+
822
+ return finished_sending
823
+
824
+ def lookup(
825
+ self,
826
+ token_len: int,
827
+ block_hashes: list[BlockHash],
828
+ ) -> int:
829
+ """Check how many prefix tokens exist in the store.
830
+
831
+ Checks across all TP ranks and PP ranks.
832
+ """
833
+ end = 0
834
+ keys: list[str] = []
835
+ try:
836
+ starts: list[int] = []
837
+ for start, end, key in self.token_database.process_tokens(
838
+ token_len, block_hashes
839
+ ):
840
+ keys.append(key.to_string())
841
+ starts.append(start)
842
+
843
+ # Expand keys for all TP ranks
844
+ multi_tp_keys = keys[:]
845
+ for i in range(1, min(self.tp_size, self.num_kv_head)):
846
+ for item in keys:
847
+ new_str = item.replace("@tp_rank:0", f"@tp_rank:{i}", 1)
848
+ multi_tp_keys.append(new_str)
849
+
850
+ # Expand keys for all PP ranks
851
+ pp_base_keys = multi_tp_keys.copy()
852
+ for i in range(1, self.pp_size):
853
+ for item in pp_base_keys:
854
+ new_str = item.replace("@pp_rank:0", f"@pp_rank:{i}", 1)
855
+ multi_tp_keys.append(new_str)
856
+
857
+ res = self.store.batch_is_exist(multi_tp_keys)
858
+
859
+ num_block = len(keys)
860
+ multi_tp_values = [
861
+ res[i * num_block : (i + 1) * num_block]
862
+ for i in range(min(self.tp_size, self.num_kv_head) * self.pp_size)
863
+ ]
864
+ index = self._find_min_first_non_one_index(multi_tp_values)
865
+ if index != -1:
866
+ return starts[index]
867
+ except Exception as e:
868
+ logger.error("Remote connection failed in lookup: %s", e)
869
+ return 0
870
+ return end
871
+
872
+ @staticmethod
873
+ def _find_min_first_non_one_index(
874
+ arr: list[list[int]],
875
+ ) -> int:
876
+ try:
877
+ return min(idx for row in arr for idx, val in enumerate(row) if val != 1)
878
+ except ValueError:
879
+ return -1
880
+
881
+ def get_kv_events(self) -> list[BlockStored]:
882
+ if self.enable_kv_events and self.kv_send_thread is not None:
883
+ return self.kv_send_thread.get_kv_events()
884
+ return []
885
+
886
+
887
+# ============================================================
888
+# Lookup Key Server
889
+# ============================================================
890
+
891
+
892
+class LookupKeyServer:
893
+ """ZMQ server on worker rank 0 for handling prefix lookup queries."""
894
+
895
+ def __init__(
896
+ self,
897
+ store_worker: MooncakeStoreWorker,
898
+ vllm_config: VllmConfig,
899
+ ):
900
+ self.decoder = MsgpackDecoder()
901
+ self.ctx = zmq.Context() # type: ignore[attr-defined]
902
+ socket_path = get_zmq_rpc_path_lookup(vllm_config)
903
+ self._ipc_path = socket_path.removeprefix("ipc://")
904
+ if os.path.exists(self._ipc_path):
905
+ os.unlink(self._ipc_path)
906
+ self.socket = make_zmq_socket(
907
+ self.ctx,
908
+ socket_path,
909
+ zmq.REP, # type: ignore[attr-defined]
910
+ bind=True,
911
+ )
912
+
913
+ self.store_worker = store_worker
914
+ self.running = True
915
+
916
+ def process_request():
917
+ while self.running:
918
+ all_frames = self.socket.recv_multipart(copy=False)
919
+ token_len = int.from_bytes(all_frames[0], byteorder="big")
920
+ hash_frames = all_frames[1:]
921
+ hashes_str = self.decoder.decode(hash_frames)
922
+ result = self.store_worker.lookup(token_len, hashes_str)
923
+ response = result.to_bytes(4, "big")
924
+ self.socket.send(response)
925
+
926
+ self.thread = threading.Thread(target=process_request, daemon=True)
927
+ self.thread.start()
928
+
929
+ def close(self):
930
+ self.socket.close(linger=0)
931
+ if os.path.exists(self._ipc_path):
932
+ os.unlink(self._ipc_path)
933
+
934
+
935
+# ============================================================
936
+# Lookup Key Client
937
+# ============================================================
938
+
939
+
940
+class LookupKeyClient:
941
+ """ZMQ client for querying prefix cache hits from worker."""
942
+
943
+ def __init__(self, vllm_config: VllmConfig):
944
+ self.encoder = MsgpackEncoder()
945
+ self.ctx = zmq.Context() # type: ignore[attr-defined]
946
+ socket_path = get_zmq_rpc_path_lookup(vllm_config)
947
+ self.socket = make_zmq_socket(
948
+ self.ctx,
949
+ socket_path,
950
+ zmq.REQ, # type: ignore[attr-defined]
951
+ bind=False,
952
+ )
953
+
954
+ def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
955
+ hash_strs = [h.hex() for h in block_hashes]
956
+ hash_frames = self.encoder.encode(hash_strs)
957
+ token_len_bytes = token_len.to_bytes(4, byteorder="big")
958
+ all_frames = [token_len_bytes] + list(hash_frames)
959
+ self.socket.send_multipart(all_frames, copy=False)
960
+ resp = self.socket.recv()
961
+ result = int.from_bytes(resp, "big")
962
+ return result
963
+
964
+ def close(self):
965
+ self.socket.close(linger=0)
966
+
967
+
968
+def get_zmq_rpc_path_lookup(vllm_config: VllmConfig) -> str:
969
+ """Construct IPC path for ZMQ lookup socket."""
970
+ dp_rank = get_mooncake_dp_engine_index(vllm_config.parallel_config)
971
+ base_url = envs.VLLM_RPC_BASE_PATH
972
+ rpc_port = 0
973
+ assert vllm_config.kv_transfer_config is not None
974
+ extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
975
+ if "lookup_rpc_port" in extra_config:
976
+ rpc_port = extra_config["lookup_rpc_port"]
977
+ uid = os.getuid()
978
+ logger.debug("Base URL: %s, RPC Port: %s, UID: %s", base_url, rpc_port, uid)
979
+ return f"ipc://{base_url}/lookup_rpc_port_{rpc_port}_uid{uid}_dp_rank{dp_rank}"