@tduncan / Nvidia-Optimization-Stack-VLLM / commits / bf2866f8bf

[KV Connector] Add decode offloading to Mooncake Store consumers (#52466)

Signed-off-by: z-zanez <zhouzh93@mail2.sysu.edu.cn> Signed-off-by: Guanyi Chen <939416532@qq.com> Co-authored-by: z-zanez <zhouzh93@mail2.sysu.edu.cn>

Guanyi Chen committed Aug 20, 2026 at 08:40 UTC bf2866f8bf5bb20628e2b93835be3c281a9b4ca4
7 files changed +518 -72
docs/features/mooncake_store_connector_usage.md
+23
@@ -136,6 +136,25 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \
136 }'
137 ```
138
139 +To also offload newly completed decode KV blocks, add the following extra
140 +configuration to the decoder's `MooncakeStoreConnector` entry. This currently
141 +requires the prefill and decode instances to use the same tensor-parallel size
142 +and compatible KV-cache topology.
143 +
144 +When decode processing starts, the consumer checks the block-aligned prompt
145 +prefix and fills any blocks missing from the Store. Subsequent saves append
146 +newly completed decode blocks. This keeps a complete, reusable prefix in the
147 +Store and also covers deployments where prompt KV is delivered directly by
148 +`MooncakeConnector` instead of through the Store.
149 +
150 +```json
151 +{
152 + "kv_connector_extra_config": {
153 + "save_decode_cache": true
154 + }
155 +}
156 +```
157 +
158 **Proxy:**
159
160 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.
@@ -230,6 +249,10 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants
249 - `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
250 - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
251 - `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).
252 +- `save_decode_cache` (bool): Enable offloading decode tokens' KV cache. A `kv_consumer` does not save during prefill; when decode starts, it fills any missing block-aligned prompt prefix before appending completed decode blocks. Default: `false`.
253 +
254 +Decode offloading uses the existing TP-rank key namespace. Cross-TP sharing is
255 +not yet supported.
256
257 ## Notes
258
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+187 -2
@@ -3,6 +3,8 @@
3
4 from types import SimpleNamespace
5
6 +import pytest
7 +
8 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
9 LoadSpec,
10 MooncakeStoreWorkerMetadata,
@@ -16,10 +18,16 @@ from vllm.v1.core.block_pool import BlockPool
18
19
20 def _make_bare_scheduler(
19 - *, hash_block_size: int = 16, enable_partial_hash_hits: bool = False
21 + *,
22 + hash_block_size: int = 16,
23 + enable_partial_hash_hits: bool = False,
24 + kv_role: str = "kv_both",
25 + save_decode_cache: bool = False,
26 ) -> MooncakeStoreScheduler:
27 scheduler = object.__new__(MooncakeStoreScheduler)
22 - scheduler.kv_role = "kv_both"
28 + scheduler.kv_role = kv_role
29 + scheduler.save_decode_cache = save_decode_cache
30 + scheduler.enable_kv_events = False
31 scheduler.lookup_async = False
32 scheduler.enable_lookup = True
33 scheduler._block_size = 16
@@ -56,6 +64,74 @@ def _make_scheduler_output(*, scheduled_spec_tokens: list[int] | None):
64 )
65
66
67 +def _make_decode_scheduler_output(
68 + *, num_computed_tokens: int, num_scheduled_tokens: int = 1
69 +) -> SimpleNamespace:
70 + return SimpleNamespace(
71 + finished_req_ids=set(),
72 + preempted_req_ids=set(),
73 + scheduled_new_reqs=[],
74 + scheduled_cached_reqs=SimpleNamespace(
75 + req_ids=["req-0"],
76 + # The block that becomes full was allocated on an earlier step.
77 + new_block_ids=[()],
78 + num_computed_tokens=[num_computed_tokens],
79 + resumed_req_ids=set(),
80 + ),
81 + num_scheduled_tokens={"req-0": num_scheduled_tokens},
82 + scheduled_spec_decode_tokens={},
83 + )
84 +
85 +
86 +def _make_new_scheduler_output() -> SimpleNamespace:
87 + request = SimpleNamespace(
88 + req_id="req-0",
89 + num_computed_tokens=0,
90 + prompt_token_ids=list(range(32)),
91 + prefill_token_ids=None,
92 + block_ids=([0, 1],),
93 + block_hashes=[b"h0", b"h1"],
94 + )
95 + return SimpleNamespace(
96 + request=request,
97 + finished_req_ids=set(),
98 + preempted_req_ids=set(),
99 + scheduled_new_reqs=[request],
100 + scheduled_cached_reqs=SimpleNamespace(
101 + req_ids=[],
102 + new_block_ids=[],
103 + num_computed_tokens=[],
104 + resumed_req_ids=set(),
105 + ),
106 + num_scheduled_tokens={"req-0": 32},
107 + scheduled_spec_decode_tokens={},
108 + )
109 +
110 +
111 +def test_scheduler_only_tracks_token_ids_for_kv_events():
112 + for enable_kv_events in (False, True):
113 + scheduler = _make_bare_scheduler()
114 + scheduler.enable_kv_events = enable_kv_events
115 + scheduler._unfinished_requests["req-0"] = (
116 + _make_new_scheduler_output().request,
117 + ([0, 1],),
118 + )
119 +
120 + meta = scheduler.build_connector_meta(_make_new_scheduler_output())
121 +
122 + tracker = scheduler._request_trackers["req-0"]
123 + req_meta = meta.requests[0]
124 + if enable_kv_events:
125 + assert tracker.token_ids == list(range(32))
126 + assert req_meta.token_ids == list(range(32))
127 + assert req_meta.token_ids_start == 0
128 + tracker.token_ids.append(99)
129 + assert req_meta.token_ids == list(range(32))
130 + else:
131 + assert tracker.token_ids is None
132 + assert req_meta.token_ids is None
133 +
134 +
135 def _make_preemption_scheduler_output():
136 return SimpleNamespace(
137 finished_req_ids=set(),
@@ -103,6 +179,29 @@ def _add_unfinished_request(
179 )
180
181
182 +def _setup_decode_request(
183 + *,
184 + kv_role: str = "kv_consumer",
185 + save_decode_cache: bool = False,
186 + token_len: int = 47,
187 +) -> tuple[MooncakeStoreScheduler, RequestTracker]:
188 + scheduler = _make_bare_scheduler(
189 + kv_role=kv_role, save_decode_cache=save_decode_cache
190 + )
191 + token_ids = list(range(token_len + 1))
192 + _add_unfinished_request(
193 + scheduler,
194 + token_ids=token_ids,
195 + block_hashes=[b"h0", b"h1", b"h2"],
196 + prefill_end_tokens=32,
197 + )
198 + tracker = scheduler._request_trackers["req-0"]
199 + tracker.token_len = token_len
200 + tracker.allocated_block_ids = ([0, 1, 2],)
201 + tracker.token_ids = token_ids[:token_len]
202 + return scheduler, tracker
203 +
204 +
205 def test_cached_request_with_spec_decode_does_not_save_scheduled_drafts():
206 # Drafts in scheduled_spec_decode_tokens are not appended to all_token_ids
207 # yet, so the tracker's token_len does not advance and num_tokens_to_save
@@ -149,6 +248,92 @@ def test_cached_request_without_spec_decode_keeps_current_step_save_overlap():
248 assert tracker.num_saved_tokens == 48
249
250
251 +@pytest.mark.parametrize("kv_role", ["kv_consumer", "kv_both"])
252 +def test_decode_tracking_is_skipped_by_default(kv_role):
253 + scheduler, tracker = _setup_decode_request(kv_role=kv_role)
254 +
255 + meta = scheduler.build_connector_meta(
256 + _make_decode_scheduler_output(num_computed_tokens=47)
257 + )
258 +
259 + assert meta.requests == []
260 + assert tracker.token_len == 47
261 + assert tracker.num_saved_tokens == 32
262 + assert tracker.allocated_block_ids == ([0, 1, 2],)
263 + assert tracker.token_ids == list(range(47))
264 +
265 +
266 +def test_fresh_consumer_first_decode_save_can_backfill_missing_prompt():
267 + scheduler = _make_bare_scheduler(
268 + kv_role="kv_consumer",
269 + save_decode_cache=True,
270 + )
271 + scheduler.enable_kv_events = True
272 + new_output = _make_new_scheduler_output()
273 + request = new_output.request
274 + request.block_ids = ([0, 1, 2],)
275 + request.block_hashes = [b"h0", b"h1", b"h2"]
276 + request.all_token_ids = list(range(48))
277 + request.num_output_placeholders = 0
278 + scheduler._unfinished_requests["req-0"] = (request, request.block_ids)
279 +
280 + # The consumer does not save during prefill, so its first decode save must
281 + # still cover the prompt. The worker deduplicates prompt blocks already in
282 + # the Store and fills any that are missing, preserving a complete prefix.
283 + assert scheduler.build_connector_meta(new_output).requests == []
284 + tracker = scheduler._request_trackers["req-0"]
285 + assert tracker.num_saved_tokens == 0
286 +
287 + # The first decode step checks/saves the block-aligned prompt. The worker's
288 + # Store dedup turns this into a no-op when the producer already saved it.
289 + [prompt_meta] = scheduler.build_connector_meta(
290 + _make_decode_scheduler_output(num_computed_tokens=32)
291 + ).requests
292 + assert prompt_meta.token_len_chunk == 32
293 + assert prompt_meta.token_ids_start == 0
294 + assert prompt_meta.token_ids == list(range(32))
295 +
296 + for num_computed_tokens in range(33, 48):
297 + meta = scheduler.build_connector_meta(
298 + _make_decode_scheduler_output(
299 + num_computed_tokens=num_computed_tokens,
300 + )
301 + )
302 + if num_computed_tokens < 47:
303 + assert meta.requests == []
304 +
305 + [req_meta] = meta.requests
306 + assert req_meta.token_len_chunk == 48
307 + assert req_meta.token_ids_start == 32
308 + assert req_meta.token_ids == list(range(32, 48))
309 + assert tracker.num_saved_tokens == 48
310 +
311 +
312 +@pytest.mark.parametrize("token_len, saved_tokens", [(46, 32), (47, 48)])
313 +def test_consumer_saves_only_full_decode_blocks(token_len, saved_tokens):
314 + scheduler, tracker = _setup_decode_request(
315 + save_decode_cache=True, token_len=token_len
316 + )
317 +
318 + meta = scheduler.build_connector_meta(
319 + _make_decode_scheduler_output(num_computed_tokens=token_len)
320 + )
321 +
322 + assert tracker.token_len == token_len + 1
323 + assert tracker.num_saved_tokens == saved_tokens
324 + if saved_tokens == 32:
325 + assert meta.requests == []
326 + else:
327 + [req_meta] = meta.requests
328 + assert req_meta.can_save is True
329 + assert req_meta.token_len_chunk == 48
330 + assert req_meta.block_ids == ([0, 1, 2],)
331 + assert req_meta.token_ids == list(range(32, 48))
332 + assert req_meta.token_ids_start == 32
333 + tracker.token_ids.append(999)
334 + assert req_meta.token_ids == list(range(32, 48))
335 +
336 +
337 def test_preemption_resets_tracker():
338 scheduler = _make_bare_scheduler()
339 _add_unfinished_request(
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+157 -3
@@ -36,7 +36,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
36 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
37 MooncakeStoreConnectorStats,
38 )
39 -from vllm.v1.core.kv_cache_utils import BlockHash
39 +from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
40
41
42 class _RecordingBlockHashes:
@@ -230,6 +230,7 @@ def _make_vllm_config(
230 extra_config: dict[str, object] | None = None,
231 rank: int = 0,
232 decode_context_parallel_size: int = 1,
233 + kv_role: str = "kv_both",
234 ) -> SimpleNamespace:
235 return SimpleNamespace(
236 model_config=_FakeModelConfig(),
@@ -240,7 +241,9 @@ def _make_vllm_config(
241 decode_context_parallel_size=decode_context_parallel_size,
242 prefill_context_parallel_size=1,
243 ),
243 - kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config),
244 + kv_transfer_config=_FakeKVTransferConfig(
245 + kv_role=kv_role, extra_config=extra_config
246 + ),
247 cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=10),
248 kv_events_config=SimpleNamespace(enable_kv_cache_events=False),
249 speculative_config=None,
@@ -943,15 +946,20 @@ def test_stale_store_job_cannot_touch_a_reused_request_id():
946 live = _make_store_req("req-a", [b"a0", b"a1"])
947 live.store_job_id = 2
948 thread.add_request(live)
949 + thread._retry_token_ids["req-a"] = (32, list(range(32, 64)))
950
951 thread.finish_store_job(stale)
952 thread._record_saved(stale, 64)
953 thread._mark_request_skipped_for_pressure(stale)
954 + thread._update_retry_token_ids(stale, False, 0, list(range(32)))
955 + thread._update_retry_token_ids(stale, True, 0, None)
956
957 assert thread.is_live_store_job(live)
958 assert not thread.is_live_store_job(stale)
959 assert thread._saved_offset.get("req-a") is None
960 assert "req-a" not in thread._skip_store_requests
961 + assert thread._get_retry_token_ids(stale) is None
962 + assert thread._get_retry_token_ids(live) == (32, list(range(32, 64)))
963
964
965 def test_store_recving_thread_reports_failed_block_ids():
@@ -1655,6 +1663,39 @@ def test_requester_worker_init_skips_disk_budget_when_offload_disabled(
1663 assert w.disk_offload_buffer_budget_bytes is None
1664
1665
1666 +def test_save_decode_cache_keeps_transfer_path_enabled(tmp_path, monkeypatch):
1667 + store = MagicMock()
1668 + store.setup.return_value = 0
1669 + _install_fake_mooncake(monkeypatch, store)
1670 + _patch_worker_runtime(monkeypatch)
1671 + monkeypatch.setenv(
1672 + "MOONCAKE_CONFIG_PATH",
1673 + _write_mooncake_config(
1674 + tmp_path,
1675 + {
1676 + "metadata_server": "http://metadata/endpoint",
1677 + "protocol": "tcp",
1678 + "device_name": "",
1679 + "master_server_address": "10.0.0.7:50051",
1680 + },
1681 + ),
1682 + )
1683 +
1684 + w = worker.MooncakeStoreWorker(
1685 + _make_vllm_config(
1686 + kv_role="kv_consumer",
1687 + extra_config={
1688 + "enable_lookup": False,
1689 + "save_decode_cache": True,
1690 + },
1691 + ),
1692 + _make_kv_cache_config(),
1693 + )
1694 +
1695 + assert w.can_put is True
1696 + assert w._capacity_only is False
1697 +
1698 +
1699 def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
1700 tmp_path,
1701 monkeypatch,
@@ -1962,7 +2003,7 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():
2003 store = MagicMock()
2004 store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
2005 store.batch_put_from_multi_buffers.side_effect = (
1965 - lambda keys, addrs, sizes, *_args: [256] * len(keys)
2006 + lambda keys, addrs, sizes, *_args: ([256] * len(keys))
2007 )
2008
2009 full_spec = FullAttentionSpec(
@@ -2176,6 +2217,85 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata():
2217 assert swa_event.block_hashes == [maybe_convert_block_hash(BlockHash(hs[3]))]
2218
2219
2220 +def _make_event_store_req(token_len: int, token_ids_start: int = 0) -> ReqMeta:
2221 + num_blocks = token_len // 16
2222 + return ReqMeta(
2223 + req_id="r0",
2224 + token_len_chunk=token_len,
2225 + block_ids=(list(range(num_blocks)),),
2226 + block_hashes=[f"a{i}".encode() for i in range(num_blocks)],
2227 + can_save=True,
2228 + token_ids=list(range(token_ids_start, token_len)),
2229 + token_ids_start=token_ids_start,
2230 + )
2231 +
2232 +
2233 +@pytest.mark.parametrize(
2234 + ("saved_offset", "put_step", "exists", "stored_indices", "parent_indices"),
2235 + [
2236 + pytest.param(16, 1, [0, 0, 0], [1, 2, 3], [0, 1, 2], id="suffix"),
2237 + pytest.param(0, 1, [1, 0, 1, 0], [1, 3], [0, 2], id="dedup-holes"),
2238 + pytest.param(0, 2, [0, 0], [0, 2], [None, 1], id="tp-stride"),
2239 + ],
2240 +)
2241 +def test_store_sending_thread_kv_events_use_request_chain_parents(
2242 + saved_offset, put_step, exists, stored_indices, parent_indices
2243 +):
2244 + store = MagicMock()
2245 + store.batch_is_exist.return_value = exists
2246 + store.batch_put_from_multi_buffers.return_value = [256] * len(stored_indices)
2247 + thread = _make_store_sending_thread(store, put_step=put_step)
2248 + thread.enable_kv_event = True
2249 +
2250 + thread._saved_offset["r0"] = saved_offset
2251 + _run_store_req(thread, _make_event_store_req(64, saved_offset))
2252 +
2253 + events = thread.get_kv_events()
2254 + assert [event.block_hashes for event in events] == [
2255 + [maybe_convert_block_hash(BlockHash(f"a{i}".encode()))] for i in stored_indices
2256 + ]
2257 + assert [event.parent_block_hash for event in events] == [
2258 + (None if i is None else maybe_convert_block_hash(BlockHash(f"a{i}".encode())))
2259 + for i in parent_indices
2260 + ]
2261 + assert thread._retry_token_ids == {}
2262 +
2263 +
2264 +def test_store_sending_thread_kv_events_retry_without_covered_tokens():
2265 + store = MagicMock()
2266 + store.batch_is_exist.return_value = [0, 0]
2267 + store.batch_put_from_multi_buffers.return_value = [256, 256]
2268 + thread = _make_store_sending_thread(store)
2269 + thread.enable_kv_event = True
2270 +
2271 + _run_store_req(thread, _make_event_store_req(32, 16))
2272 +
2273 + retry_event, suffix_event = thread.get_kv_events()
2274 + assert retry_event.token_ids == []
2275 + assert suffix_event.token_ids == list(range(16, 32))
2276 +
2277 +
2278 +def test_store_sending_thread_kv_events_recover_suffix_after_put_failure():
2279 + store = MagicMock()
2280 + store.batch_is_exist.side_effect = ([0], [0, 0])
2281 + store.batch_put_from_multi_buffers.side_effect = ([-1], [256, 256])
2282 + thread = _make_store_sending_thread(store)
2283 + thread.enable_kv_event = True
2284 +
2285 + _run_store_req(thread, _make_event_store_req(16))
2286 +
2287 + assert thread.get_kv_events() == []
2288 + assert thread._retry_token_ids["r0"] == (0, list(range(16)))
2289 +
2290 + _run_store_req(thread, _make_event_store_req(32, 16))
2291 +
2292 + retry_event, suffix_event = thread.get_kv_events()
2293 + assert retry_event.token_ids == list(range(16))
2294 + assert suffix_event.token_ids == list(range(16, 32))
2295 + assert thread._retry_token_ids == {}
2296 + assert thread._saved_offset["r0"] == 32
2297 +
2298 +
2299 def _auto_set_ready_event(*args, **kwargs):
2300 """Side effect for mocked thread constructors that auto-sets ready_event."""
2301 for arg in args:
@@ -2214,6 +2334,7 @@ def _make_bare_worker(
2334 num_gpu_blocks: int = 10,
2335 block_size: int = 16,
2336 kv_role: str = "kv_both",
2337 + save_decode_cache: bool = False,
2338 ) -> mooncake_store_worker.MooncakeStoreWorker:
2339 """Construct a MooncakeStoreWorker via __new__, bypassing __init__.
2340
@@ -2227,14 +2348,17 @@ def _make_bare_worker(
2348 worker.store = MagicMock()
2349 worker.store.register_buffer.return_value = 0
2350 worker.kv_role = kv_role
2351 + worker.can_put = kv_role in ("kv_producer", "kv_both") or save_decode_cache
2352 worker._capacity_only = False
2353 worker.block_size = block_size
2354 worker.tp_rank = 0
2355 worker.enable_kv_events = False
2356 + worker.load_async = True
2357 worker.kv_send_thread = None
2358 worker.kv_recv_threads = []
2359 worker.num_recv_threads = 1
2360 worker.recv_request_queue = queue.Queue()
2361 + worker.finished_store_req = set()
2362 worker.tp_size = 1
2363 worker.num_kv_head = 1
2364 worker.pp_size = 1
@@ -2743,6 +2867,36 @@ def test_lookup_applies_swa_mask_before_accessing_hashes():
2867 # ---------------------------------------------------------------------------
2868
2869
2870 +@pytest.mark.parametrize("save_decode_cache", [False, True])
2871 +def test_consumer_starts_send_thread_only_when_put_is_enabled(save_decode_cache):
2872 + num_blocks = 10
2873 + tensor = torch.zeros(num_blocks, 64, dtype=torch.float16)
2874 +
2875 + worker = _make_bare_worker(
2876 + kv_role="kv_consumer", save_decode_cache=save_decode_cache
2877 + )
2878 + _register_with_mocked_threads(worker, {"layer0": tensor})
2879 + assert (worker.kv_send_thread is not None) == save_decode_cache
2880 +
2881 +
2882 +def test_putting_consumer_queues_decode_save():
2883 + w = _make_bare_worker(kv_role="kv_consumer", save_decode_cache=True)
2884 + send_thread = MagicMock()
2885 + w.kv_send_thread = send_thread
2886 + req = _make_store_req("decode-req", [b"h0", b"h1"])
2887 + req.store_job_id = 1
2888 + meta = mooncake_store_worker.MooncakeStoreConnectorMetadata(set(), set())
2889 + meta.add_request(req)
2890 +
2891 + with patch.object(torch.cuda, "Event") as event_cls:
2892 + event = event_cls.return_value
2893 + w.get_finished(set(), meta)
2894 +
2895 + event.record.assert_called_once_with()
2896 + send_thread.add_request.assert_called_once_with(req)
2897 + assert req.current_event is event
2898 +
2899 +
2900 def test_register_kv_caches_blocks_first_single_segment():
2901 """Blocks-first layout (FlashInfer/MLA): one segment per layer."""
2902 num_blocks = 10
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+6 -4
@@ -141,12 +141,14 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
141 assert vllm_config.kv_transfer_config is not None
142 assert kv_cache_config is not None, "kv_cache_config is required"
143 self.kv_role = vllm_config.kv_transfer_config.kv_role
144 + extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
145 + save_decode_cache = extra_config.get("save_decode_cache", False)
146 # Capacity-only: contributes its segment to the store pool but transfers
147 # no KV, so the KV-cache-shape invariants below cannot be reached.
146 - self._capacity_only = self.kv_role == "kv_consumer" and not (
147 - vllm_config.kv_transfer_config.kv_connector_extra_config.get(
148 - "enable_lookup", True
149 - )
148 + self._capacity_only = (
149 + self.kv_role == "kv_consumer"
150 + and not extra_config.get("enable_lookup", True)
151 + and not save_decode_cache
152 )
153 if not self._capacity_only:
154 self._validate_kv_cache_config(vllm_config, kv_cache_config)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+9 -6
@@ -362,10 +362,11 @@ class ReqMeta:
362
363 can_save: bool | None = None
364 load_spec: LoadSpec | None = None
365 - is_last_chunk: bool | None = None
365 current_event: torch.cuda.Event | None = None
366
367 token_ids: list[int] | None = None
368 + # Absolute request offset represented by token_ids[0].
369 + token_ids_start: int = 0
370 num_prompt_tokens: int | None = None
371 # Identifies this store job for the engine's lifetime. A request id cannot
372 # serve that purpose: it is reused once a preempted request resumes, so it
@@ -384,14 +385,14 @@ class ReqMeta:
385 load_spec: LoadSpec | None = None,
386 skip_save: bool | None = False,
387 block_hashes: list[BlockHash] | None = None,
387 - is_last_chunk: bool | None = None,
388 ) -> "ReqMeta | None":
389 """Create ReqMeta from a RequestTracker."""
390 if block_hashes is None:
391 block_hashes = []
392 input_token_len = tracker.token_len
393
394 - chunk_boundary = cdiv(tracker.num_saved_tokens + 1, block_size) * block_size
394 + token_ids_start = tracker.num_saved_tokens
395 + chunk_boundary = cdiv(token_ids_start + 1, block_size) * block_size
396 num_tokens_to_save = input_token_len // block_size * block_size
397
398 skip_save = skip_save or num_tokens_to_save < chunk_boundary
@@ -408,8 +409,10 @@ class ReqMeta:
409 tracker.num_saved_tokens = num_tokens_to_save
410
411 token_ids = None
411 - if tracker.token_ids:
412 - token_ids = tracker.token_ids
412 + if tracker.token_ids and not skip_save:
413 + # Scheduler tracking continues while this job is handled by an
414 + # asynchronous worker, so metadata must own a stable snapshot.
415 + token_ids = tracker.token_ids[token_ids_start:num_tokens_to_save]
416
417 if load_spec is not None and load_spec.can_load:
418 logger.debug(
@@ -433,8 +436,8 @@ class ReqMeta:
436 can_save=not skip_save,
437 load_spec=load_spec,
438 block_hashes=block_hashes,
436 - is_last_chunk=is_last_chunk,
439 token_ids=token_ids,
440 + token_ids_start=token_ids_start,
441 num_prompt_tokens=tracker.prefill_end_tokens,
442 )
443
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+42 -33
@@ -63,6 +63,11 @@ class MooncakeStoreScheduler:
63 self.lookup_async = kvc_extra_config.get("lookup_async", False)
64 # Skips lookup CPU cost on instances that never load KV from the store.
65 self.enable_lookup = kvc_extra_config.get("enable_lookup", True)
66 + self.save_decode_cache = kvc_extra_config.get("save_decode_cache", False)
67 + kv_event_config = vllm_config.kv_events_config
68 + self.enable_kv_events = bool(
69 + kv_event_config and kv_event_config.enable_kv_cache_events
70 + )
71 self.client = LookupKeyClient(vllm_config)
72
73 # Align with the engine's own scheduler_block_size and hash_block_size.
@@ -182,7 +187,8 @@ class MooncakeStoreScheduler:
187 self, scheduler_output: SchedulerOutput
188 ) -> KVConnectorMetadata:
189 """Build connector metadata for this scheduler step."""
185 - force_skip_save = self.kv_role == "kv_consumer"
190 + is_consumer = self.kv_role == "kv_consumer"
191 + can_process_cached = not is_consumer or self.save_decode_cache
192
193 for finished_req_id in scheduler_output.finished_req_ids:
194 self.client.discard(finished_req_id)
@@ -227,37 +233,38 @@ class MooncakeStoreScheduler:
233 token_len=num_tokens_to_compute,
234 allocated_block_ids=unfolded_block_ids,
235 num_saved_tokens=0,
230 - token_ids=prefill_tokens[:num_tokens_to_compute],
236 + token_ids=(
237 + prefill_tokens[:num_tokens_to_compute]
238 + if self.enable_kv_events
239 + else None
240 + ),
241 prefill_end_tokens=len(prefill_tokens),
242 )
243 self._request_trackers[request.req_id] = request_tracker
244
235 - last_chunk_tokens_num = (
236 - len(prefill_tokens) // self._block_size * self._block_size
237 - )
238 -
245 req_meta = ReqMeta.from_request_tracker(
246 request_tracker,
247 self._block_size,
248 load_spec=load_spec,
243 - skip_save=force_skip_save,
249 + # A consumer may write decode KV without becoming a prefill
250 + # producer. Loads are still carried by the same metadata.
251 + skip_save=is_consumer,
252 block_hashes=request_real.block_hashes,
245 - is_last_chunk=(request_tracker.token_len >= last_chunk_tokens_num),
253 )
254 if req_meta is not None:
255 meta.add_request(req_meta)
256
257 # Handle cached (running, or MRV1 resumed-from-preemption) requests
258 cached_reqs = scheduler_output.scheduled_cached_reqs
252 - if not force_skip_save:
259 + if can_process_cached:
260 for i, req_id in enumerate(cached_reqs.req_ids):
261 new_block_ids = cached_reqs.new_block_ids[i]
255 - if not new_block_ids:
256 - continue
262
263 req_meta = None
264 if req_id in cached_reqs.resumed_req_ids:
265 # Resumed after preemption
266 + if not new_block_ids:
267 + continue
268 if isinstance(new_block_ids, tuple):
269 new_block_ids = tuple(b.copy() for b in new_block_ids)
270 else:
@@ -277,27 +284,33 @@ class MooncakeStoreScheduler:
284 token_len=num_tokens_to_compute,
285 allocated_block_ids=new_block_ids,
286 num_saved_tokens=0,
280 - token_ids=prefill_tokens[:num_tokens_to_compute].copy(),
287 + token_ids=(
288 + prefill_tokens[:num_tokens_to_compute]
289 + if self.enable_kv_events
290 + else None
291 + ),
292 prefill_end_tokens=len(prefill_tokens),
293 )
294 self._request_trackers[req_id] = request_tracker
295
285 - last_chunk_tokens_num = (
286 - len(prefill_tokens) // self._block_size * self._block_size
287 - )
296 req_meta = ReqMeta.from_request_tracker(
297 request_tracker,
298 self._block_size,
299 load_spec=load_spec,
292 - skip_save=force_skip_save,
300 + skip_save=is_consumer,
301 block_hashes=request_real.block_hashes,
294 - is_last_chunk=(
295 - request_tracker.token_len >= last_chunk_tokens_num
296 - ),
302 )
303 else:
304 # Decode/chunked request
305 request_tracker = self._request_trackers[req_id]
306 + num_computed_token = cached_reqs.num_computed_tokens[i]
307 + # Use the tracker's snapshot of the prefill range so resumed
308 + # requests keep saving past the original prompt boundary.
309 + prefill_end = request_tracker.prefill_end_tokens
310 + is_decode = num_computed_token >= prefill_end
311 + if is_decode and not self.save_decode_cache:
312 + continue
313 +
314 num_new_tokens = scheduler_output.num_scheduled_tokens[req_id]
315 req_tuple = self._unfinished_requests.get(req_id)
316 if req_tuple:
@@ -307,30 +320,26 @@ class MooncakeStoreScheduler:
320 num_current_tokens : num_current_tokens + num_new_tokens
321 ]
322 request_tracker.token_len += len(new_token_ids)
323 + if request_tracker.token_ids is not None:
324 + request_tracker.token_ids.extend(new_token_ids)
325 else:
326 raise ValueError(
327 f"Request {req_id} is not in _unfinished_requests"
328 )
314 - num_computed_token = cached_reqs.num_computed_tokens[i]
315 - # Use the tracker's snapshot of the prefill range so resumed
316 - # requests keep saving past the original prompt boundary.
317 - prefill_end = request_tracker.prefill_end_tokens
318 - if num_computed_token >= prefill_end:
329 + # A block is usually allocated before the step that fills
330 + # it, so reaching a save boundary does not imply that this
331 + # step has new block ids.
332 + if new_block_ids:
333 + request_tracker.update(new_block_ids)
334 + if is_consumer and not is_decode:
335 continue
320 - request_tracker.update(new_block_ids)
336
322 - last_chunk_tokens_num = (
323 - prefill_end // self._block_size * self._block_size
324 - )
337 req_meta = ReqMeta.from_request_tracker(
338 request_tracker,
339 self._block_size,
340 load_spec=None,
329 - skip_save=force_skip_save,
341 + skip_save=False,
342 block_hashes=unfinished_req.block_hashes,
331 - is_last_chunk=(
332 - request_tracker.token_len >= last_chunk_tokens_num
333 - ),
343 )
344
345 if req_meta is not None:
@@ -370,7 +379,7 @@ class MooncakeStoreScheduler:
379 # emit an offload-only ReqMeta (token_len_chunk=0 skips the normal
380 # save; can_save=True takes the normal enqueue path).
381 step_partial_tails = getattr(scheduler_output, "partial_tail_offloads", None)
373 - if step_partial_tails and not force_skip_save:
382 + if step_partial_tails and not is_consumer:
383 pending = dict(step_partial_tails)
384 for req_meta in meta.requests:
385 if req_meta.can_save:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+94 -24
@@ -530,6 +530,9 @@ class KVCacheStoreSendingThread(KVTransferThread):
530 # Per-request high-water mark of tokens actually persisted; the next
531 # batch resumes here, so pressure-skipped or failed ranges are retried.
532 self._saved_offset: dict[str, int] = {}
533 + # Retained only after a failed store so retry events can recover the
534 + # token suffix without full snapshots on the normal path.
535 + self._retry_token_ids: dict[str, tuple[int, list[int]]] = {}
536
537 def add_request(self, request: ReqMeta) -> None:
538 # Register before enqueueing so a job is never picked up unledgered.
@@ -552,6 +555,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
555 del self.stored_requests[req_id]
556 self._skip_store_requests.discard(req_id)
557 self._saved_offset.pop(req_id, None)
558 + self._retry_token_ids.pop(req_id, None)
559
560 def finish_store_job(self, req_meta: ReqMeta) -> None:
561 """Retire a job from the ledger and report its blocks as no longer read.
@@ -585,6 +589,36 @@ class KVCacheStoreSendingThread(KVTransferThread):
589 if req_meta.store_job_id in self.stored_requests.get(req_meta.req_id, ()):
590 self._saved_offset[req_meta.req_id] = token_len
591
592 + def _get_retry_token_ids(self, req_meta: ReqMeta) -> tuple[int, list[int]] | None:
593 + """Return retry state only if this store job is still live."""
594 + with self.done_task_lock:
595 + if req_meta.store_job_id not in self.stored_requests.get(
596 + req_meta.req_id, ()
597 + ):
598 + return None
599 + return self._retry_token_ids.get(req_meta.req_id)
600 +
601 + def _update_retry_token_ids(
602 + self,
603 + req_meta: ReqMeta,
604 + save_completed: bool,
605 + token_ids_start: int,
606 + event_token_ids: list[int] | None,
607 + ) -> None:
608 + """Update retry state without letting a stale job touch a reused ID."""
609 + with self.done_task_lock:
610 + if req_meta.store_job_id not in self.stored_requests.get(
611 + req_meta.req_id, ()
612 + ):
613 + return
614 + if save_completed:
615 + self._retry_token_ids.pop(req_meta.req_id, None)
616 + elif event_token_ids is not None:
617 + self._retry_token_ids[req_meta.req_id] = (
618 + token_ids_start,
619 + event_token_ids,
620 + )
621 +
622 def _should_skip_request(self, req_id: str) -> bool:
623 with self.done_task_lock:
624 return self._store_pressure_active and req_id in self._skip_store_requests
@@ -785,18 +819,30 @@ class KVCacheStoreSendingThread(KVTransferThread):
819 def _handle_request(self, req_meta: ReqMeta):
820 # The single `finally` is the only way out, so the scheduler releases
821 # this job's GPU block references however the job ends.
822 + save_completed = False
823 + token_len = 0
824 + req_id = req_meta.req_id
825 + event_token_ids = req_meta.token_ids
826 + token_ids_start = req_meta.token_ids_start
827 try:
828 # Cache hits are always a multiple of ``lcm_block_size`` tokens,
829 # which is also ``store_mask``'s precondition.
830 lcm_block_size = self.coord.lcm_block_size
831 token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size
832 block_ids_per_group = req_meta.block_ids
794 - req_id = req_meta.req_id
833 current_event = req_meta.current_event
834
835 if not self.is_live_store_job(req_meta):
836 return
837
838 + if self.enable_kv_event:
839 + retry_token_ids = self._get_retry_token_ids(req_meta)
840 + if retry_token_ids is not None and event_token_ids is not None:
841 + retry_start, retry_ids = retry_token_ids
842 + if retry_start + len(retry_ids) == token_ids_start:
843 + event_token_ids = retry_ids + event_token_ids
844 + token_ids_start = retry_start
845 +
846 if self._should_skip_request(req_id):
847 logger.debug(
848 "Skipping Mooncake store for request %s while CPU/disk "
@@ -852,6 +898,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
898
899 if not keys:
900 self._record_saved(req_meta, token_len)
901 + save_completed = True
902 return
903
904 # Check which blocks already exist (dedup)
@@ -878,6 +925,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
925
926 if not missing_indices:
927 self._record_saved(req_meta, token_len)
928 + save_completed = True
929 return
930
931 if len(missing_indices) != len(keys):
@@ -927,12 +975,11 @@ class KVCacheStoreSendingThread(KVTransferThread):
975 addrs.extend(group_addrs)
976 sizes.extend(group_sizes)
977
930 - # parent_block_hash chains live within a group, not across.
978 if self.enable_kv_event:
932 - prev_key_per_group: dict[int, Any] = {}
979 new_block_hashes = [
980 maybe_convert_block_hash(bh) for bh in kv_event_block_hashes
981 ]
982 + token_ids_end = token_ids_start + len(event_token_ids or ())
983
984 for idx, (s, e, g_idx) in enumerate(
985 zip(starts, ends, group_indices, strict=True)
@@ -940,13 +987,24 @@ class KVCacheStoreSendingThread(KVTransferThread):
987 db = self.token_databases[g_idx]
988 if self.enable_kv_event:
989 token_ids = (
943 - req_meta.token_ids[s:e]
944 - if req_meta.token_ids is not None
945 - else None
990 + event_token_ids[s - token_ids_start : e - token_ids_start]
991 + if event_token_ids is not None
992 + and token_ids_start <= s
993 + and e <= token_ids_end
994 + else []
995 )
996 stored_event = BlockStored(
997 block_hashes=[new_block_hashes[idx]],
949 - parent_block_hash=prev_key_per_group.get(g_idx),
998 + # Derive the direct predecessor from the unfiltered
999 + # request chain. Adjacent PUTs need not be adjacent in
1000 + # that chain after Store dedup, masks, or TP striding.
1001 + parent_block_hash=(
1002 + maybe_convert_block_hash(
1003 + req_meta.block_hashes[s // db.hash_block_size - 1]
1004 + )
1005 + if s > 0
1006 + else None
1007 + ),
1008 token_ids=token_ids,
1009 block_size=db.block_size,
1010 lora_id=None,
@@ -955,7 +1013,6 @@ class KVCacheStoreSendingThread(KVTransferThread):
1013 group_idx=g_idx,
1014 )
1015 stored_events.append(stored_event)
958 - prev_key_per_group[g_idx] = new_block_hashes[idx]
1016
1017 if current_event is not None:
1018 current_event.synchronize()
@@ -984,6 +1041,13 @@ class KVCacheStoreSendingThread(KVTransferThread):
1041 )
1042 if failed:
1043 failed_codes = set(res[i] for i in failed)
1044 + if self.enable_kv_event:
1045 + failed_indices = set(failed)
1046 + stored_events = [
1047 + event
1048 + for i, event in enumerate(stored_events)
1049 + if i not in failed_indices
1050 + ]
1051 logger.warning(
1052 "batch_put failed: %d/%d keys failed "
1053 "(codes=%s, batch_bytes=%d, num_keys=%d), "
@@ -1008,6 +1072,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
1072 )
1073 else:
1074 self._record_saved(req_meta, token_len)
1075 + save_completed = True
1076 if self._clear_store_pressure():
1077 logger.info(
1078 "Mooncake CPU/disk offloading pressure cleared "
@@ -1023,10 +1088,18 @@ class KVCacheStoreSendingThread(KVTransferThread):
1088 num_failed_keys=len(keys),
1089 )
1090 logger.error("Failed to put key %s, error: %s", keys, e)
1091 + stored_events.clear()
1092
1093 if self.enable_kv_event and stored_events:
1094 self.update_kv_event(stored_events)
1095 finally:
1096 + if self.enable_kv_event and token_len:
1097 + self._update_retry_token_ids(
1098 + req_meta,
1099 + save_completed,
1100 + token_ids_start,
1101 + event_token_ids,
1102 + )
1103 self.finish_store_job(req_meta)
1104 self.request_queue.task_done()
1105
@@ -1274,15 +1347,19 @@ class MooncakeStoreWorker:
1347 self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_size > 1 else 0
1348
1349 assert vllm_config.kv_transfer_config is not None
1277 - self.kv_role = vllm_config.kv_transfer_config.kv_role
1278 - self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
1279 - "load_async", True
1350 + kv_role = vllm_config.kv_transfer_config.kv_role
1351 + assert kv_role is not None
1352 + self.kv_role = kv_role
1353 + extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
1354 + self.can_put = self.kv_role in ("kv_producer", "kv_both") or (
1355 + extra_config.get("save_decode_cache", False)
1356 )
1357 + self.load_async = extra_config.get("load_async", True)
1358 # Mirrors MooncakeStoreConnector._capacity_only.
1282 - self._capacity_only = self.kv_role == "kv_consumer" and not (
1283 - vllm_config.kv_transfer_config.kv_connector_extra_config.get(
1284 - "enable_lookup", True
1285 - )
1359 + self._capacity_only = (
1360 + self.kv_role == "kv_consumer"
1361 + and not extra_config.get("enable_lookup", True)
1362 + and not self.can_put
1363 )
1364 self.cache_config = vllm_config.cache_config
1365 self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes(
@@ -1294,11 +1371,6 @@ class MooncakeStoreWorker:
1371
1372 # Initialize MooncakeDistributedStore with its own TransferEngine
1373 store_config = MooncakeStoreConfig.load_from_config()
1297 - extra_config = (
1298 - vllm_config.kv_transfer_config.kv_connector_extra_config
1299 - if vllm_config.kv_transfer_config
1300 - else {}
1301 - )
1374 self.store = MooncakeDistributedStore()
1375 local_ip = get_ip()
1376 local_hostname = rdma_utils.get_requester_local_hostname(local_ip)
@@ -1617,7 +1689,7 @@ class MooncakeStoreWorker:
1689 db.set_block_len(block_lens)
1690
1691 # Start transfer threads
1620 - if self.kv_role in ["kv_producer", "kv_both"]:
1692 + if self.can_put:
1693 ready_event_sending = threading.Event()
1694 self.kv_send_thread = KVCacheStoreSendingThread(
1695 self.store,
@@ -1700,7 +1772,7 @@ class MooncakeStoreWorker:
1772
1773 assert self.load_async, "load_async must be True for better performance."
1774 # Issue stores with CUDA event synchronization.
1703 - if self.kv_role in ["kv_producer", "kv_both"]:
1775 + if self.can_put:
1776 current_event = None
1777 for request in meta.requests:
1778 if request.can_save:
@@ -1714,8 +1786,6 @@ class MooncakeStoreWorker:
1786 request.current_event = current_event
1787 assert self.kv_send_thread is not None
1788 self.kv_send_thread.add_request(request)
1717 -
1718 - if self.kv_role in ["kv_producer", "kv_both"]:
1789 self._close_ended_store_requests(finished_req_ids, meta)
1790
1791 # Blocks read by a store job are released by the scheduler when the job