@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2743dc4f92

[Feature][Multimodal] Add cross-encoder caching to Mooncake P2P (#56242)

Signed-off-by: Zhou ziheng <jiaranran2@gmail.com> Co-authored-by: Tianyu Guo <guoty@inferact.ai> Co-authored-by: jiangkuaixue123 <jiangxiaozhou111@163.com> Co-authored-by: Codex <noreply@openai.com>

Zhou ziheng committed Sep 17, 2026 at 13:16 UTC 2743dc4f92275128700888be273610d7d145fd70
15 files changed +1863 -154
docs/features/cross_encoder_cache.md new
+109
@@ -0,0 +1,109 @@
1 +# Cross-Encoder Output Reuse
2 +
3 +Cross-Encoder output reuse lets an Encoder load image embeddings computed by
4 +another Encoder from a shared Mooncake Store. Local cache misses check the Store
5 +before encoding; newly computed outputs are published asynchronously. Both paths
6 +use the existing `ECMooncakeConnector` P2P delivery to Prefill.
7 +
8 +## Requirements
9 +
10 +- A working [disaggregated Encoder](disagg_encoder.md) deployment using
11 + `ECMooncakeConnector`. See the [Mooncake integration example](../../tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh)
12 + for the base E/PD configuration.
13 +- Model Runner V2 (`VLLM_USE_V2_MODEL_RUNNER=1`), Encoder TP=1, and no dynamic LoRA.
14 +- `--mm-processor-cache-gb` greater than zero to preserve content identifiers.
15 +- Mooncake 0.3.12 or later with a running RAM Store. Disk offloading
16 + (`enable_offload: true`) is not supported.
17 +
18 +Shared reuse supports images with contiguous 2D FP16, BF16 or FP32 encoder outputs.
19 +Other modalities use normal encoding; non-contiguous outputs are not published.
20 +
21 +## Usage
22 +
23 +On each Encoder, add `cross_encoder_cache: true` to the existing producer
24 +`--ec-transfer-config`, preserving its P2P settings:
25 +
26 +```json
27 +{
28 + "ec_connector": "ECMooncakeConnector",
29 + "ec_role": "ec_producer",
30 + "ec_connector_extra_config": {
31 + "cross_encoder_cache": true
32 + }
33 +}
34 +```
35 +
36 +Prefill keeps its existing `ECMooncakeConnector` consumer configuration and does
37 +not need a Store client.
38 +
39 +Create a JSON file pointing to an existing Store. For an independently managed
40 +RAM pool, use `standalone-store` with `global_segment_size` set to zero:
41 +
42 +```json
43 +{
44 + "metadata_server": "http://STORE_HOST:2379/metadata",
45 + "master_server_address": "STORE_HOST:50051",
46 + "protocol": "tcp",
47 + "device_name": "",
48 + "mode": "standalone-store",
49 + "global_segment_size": 0,
50 + "local_buffer_size": "4GB"
51 +}
52 +```
53 +
54 +Replace the addresses with your Store endpoints and set the path on each Encoder
55 +before starting vLLM:
56 +
57 +```bash
58 +export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json
59 +```
60 +
61 +This connects to an existing Store; it does not start the storage services.
62 +For Store setup and tenant configuration, see the
63 +[Mooncake Store guide](mooncake_store_connector_usage.md#prerequisites).
64 +
65 +## Configuration
66 +
67 +These options belong in the producer's `ec_connector_extra_config`:
68 +
69 +| Option | Default | Description |
70 +| --- | --- | --- |
71 +| `cross_encoder_cache` | `false` | Enable shared output reuse. |
72 +| `embedding_cache_prefix` | `""` | Namespace prefix for shared embeddings. |
73 +| `embedding_model_identity` | Configured model path | Override the model field in Store keys. Matching multimodal identifiers are still required. |
74 +| `store_max_pending_items` | `32` | Maximum pending publications per Encoder. Must be positive. |
75 +| `store_max_pending_bytes` | `2147483648` (2 GiB) | Maximum retained tensor storage for pending publications per Encoder. Must be positive. |
76 +| `store_read_buffer_bytes` | `134217728` (128 MiB) | Maximum reusable CPU staging buffer per Encoder; pinned for CUDA outputs. Must be positive. |
77 +
78 +Encoders sharing outputs must use the same immutable weights, compatible
79 +preprocessing and matching multimodal identifiers. Use the same configured model
80 +path for automatic identifiers: `embedding_model_identity` alone does not enable
81 +reuse across different paths. Change `embedding_cache_prefix` when replacing
82 +weights in place or changing output-affecting settings. Caller-provided UUIDs
83 +must consistently identify the same input; see [cached inputs](multimodal_inputs.md#cached-inputs).
84 +
85 +## Limitations
86 +
87 +- Store reads are synchronous. Reuse skips Encoder computation, but not
88 + preprocessing, P2P delivery or the scheduler's Encoder budget reservation.
89 +- Hits are read in batches through one lazily allocated, registered CPU staging
90 + buffer. Each chunk is copied to independent output tensors before the buffer
91 + is reused. Reads are split by `store_read_buffer_bytes`; an individual object
92 + larger than this capacity (including its 24-byte header) falls back to encoding.
93 + Staging memory is additional to the Store's `local_buffer_size` and is released
94 + on healthy shutdown. The first load includes allocation and registration costs.
95 +- Publication is best-effort. A request can finish before its outputs reach the
96 + Store, so concurrent cold requests may still encode the same image.
97 +- Publication budgets count each retained backing storage once across pending
98 + views. Its charge is released after the last view is safely reclaimed.
99 + Exceeding a budget skips the write. Store client buffers and the P2P pool
100 + consume additional memory.
101 +- Cache misses and recoverable read errors fall back to encoding. Recoverable
102 + publication errors skip the write. Unexpected native errors or unconfirmed
103 + I/O completion or buffer release can fail the worker.
104 +- Incompatible objects are rejected without replacement; they can trigger
105 + repeated fallback until evicted.
106 +- An independent Store can retain embeddings across Encoder restarts. Capacity,
107 + eviction and resilience to Store failures depend on the Store deployment.
108 +- The `protocol:v3` namespace isolates the compact embedding format from older
109 + cache objects. Encoders using different format versions do not share hits.
docs/features/disagg_encoder.md
+3
@@ -45,6 +45,9 @@ Below ready-to-run scripts shows the workflow:
45
46 ## 3 Test Script
47
48 +For optional cross-Encoder output reuse while retaining Mooncake P2P delivery,
49 +see [Cross-encoder output reuse](cross_encoder_cache.md).
50 +
51 Please refer to the directories `tests/v1/ec_connector`
52
53 ## 4 Development
tests/v1/ec_connector/unit/test_cross_encoder_cache.py new
+726
@@ -0,0 +1,726 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import ctypes
5 +import gc
6 +import struct
7 +import threading
8 +import weakref
9 +from concurrent.futures import Future, ThreadPoolExecutor
10 +from dataclasses import replace
11 +from types import SimpleNamespace
12 +from unittest.mock import MagicMock, patch
13 +
14 +import pytest
15 +import torch
16 +
17 +from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
18 +from vllm.config import MultiModalConfig
19 +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
20 +from vllm.distributed.ec_transfer.ec_connector.mooncake.config import MooncakeECConfig
21 +from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import (
22 + ECMooncakeConnectorMetadata,
23 + ECMooncakePushSpec,
24 +)
25 +from vllm.distributed.ec_transfer.ec_connector.mooncake_ec_connector import (
26 + ECMooncakeConnector,
27 +)
28 +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import (
29 + backend as store_backend,
30 +)
31 +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import (
32 + store_client,
33 +)
34 +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
35 + TensorSpec,
36 + build_embedding_namespace,
37 + make_embedding_key,
38 +)
39 +
40 +pytestmark = pytest.mark.cpu_test
41 +KEY = (
42 + "test@embedding@model:model@revision:revision"
43 + "@encoder:encoder@dtype:torch.float32@protocol:v3"
44 +)
45 +SPEC = TensorSpec((2, 2), "torch.float32", 16)
46 +
47 +
48 +def test_embedding_key_combines_namespace_and_identifier():
49 + config = create_ec_vllm_config(ec_role="ec_producer")
50 + config.ec_transfer_config.ec_connector_extra_config.update(
51 + embedding_cache_prefix="team@cache", embedding_model_identity="shared/model"
52 + )
53 + config.model_config.revision = "rev:1"
54 + config.model_config.dtype = torch.float32
55 + config.model_config.multimodal_config = None
56 + assert make_embedding_key(build_embedding_namespace(config), "image@1") == (
57 + "team@cache@embedding@model:shared/model@revision:rev:1"
58 + "@encoder:encoder:default@dtype:torch.float32@protocol:v3@id:image@1"
59 + )
60 +
61 +
62 +def test_shared_reuse_requires_content_stable_identifiers():
63 + """Encoder-only cache-off rendering uses process-local counter identifiers."""
64 + config = create_ec_vllm_config(ec_role="ec_producer")
65 + config.ec_transfer_config.ec_connector = "ECMooncakeConnector"
66 + config.ec_transfer_config.ec_connector_extra_config["cross_encoder_cache"] = True
67 + config.model_config.multimodal_config = MultiModalConfig(mm_processor_cache_gb=0)
68 + config.use_v2_model_runner = True
69 + config.lora_config = None
70 + with pytest.raises(ValueError, match="mm_processor_cache_gb"):
71 + MooncakeECConfig.from_vllm_config(config)
72 +
73 +
74 +class MemoryStore:
75 + """Only the native boundary is faked; codec, resolver and publisher are real."""
76 +
77 + def __init__(self):
78 + self.objects: dict[str, bytes] = {}
79 + self.registered: dict[int, int] = {}
80 +
81 + def close(self):
82 + return 0
83 +
84 + def batch_is_exist(self, keys):
85 + return [int(k in self.objects) for k in keys]
86 +
87 + def register_buffer(self, addr, size):
88 + self.registered[addr] = size
89 + return 0
90 +
91 + def unregister_buffer(self, addr):
92 + del self.registered[addr]
93 + return 0
94 +
95 + def batch_put_from_multi_buffers(self, keys, pointers, sizes, config):
96 + [key], [ptrs], [lengths] = keys, pointers, sizes
97 + self.objects[key] = b"".join(
98 + ctypes.string_at(p, n) for p, n in zip(ptrs, lengths, strict=True)
99 + )
100 + return [0]
101 +
102 + def batch_get_into(self, keys, pointers, sizes):
103 + results = []
104 + for key, ptr, size in zip(keys, pointers, sizes, strict=True):
105 + assert any(
106 + base <= ptr and ptr + size <= base + capacity
107 + for base, capacity in self.registered.items()
108 + )
109 + obj = self.objects.get(key)
110 + if obj is None:
111 + results.append(-704)
112 + elif len(obj) > size:
113 + results.append(-600)
114 + else:
115 + ctypes.memmove(ptr, obj, len(obj))
116 + results.append(len(obj))
117 + return results
118 +
119 +
120 +@pytest.fixture
121 +def backend():
122 + instance = store_backend.MooncakeEmbeddingStoreBackend(
123 + store_client.MooncakeEmbeddingStoreClient(
124 + MemoryStore(), read_buffer_bytes=1024
125 + ),
126 + KEY,
127 + max_pending_items=32,
128 + max_pending_bytes=2 * 1024**3,
129 + )
130 + yield instance
131 + # Some tests deliberately poison I/O; always stop the Python executor.
132 + instance._executor.shutdown(wait=True)
133 + if not instance.store_client._poisoned and (
134 + instance.store_client._read_buffer is not None
135 + or instance.store_client._put_header is not None
136 + ):
137 + instance.store_client.close()
138 +
139 +
140 +def test_tensor_round_trip(backend):
141 + client, native = backend.store_client, backend.store_client.store
142 + native.register_buffer = MagicMock(wraps=native.register_buffer)
143 + header_ptr = None
144 + for index, dtype in enumerate((torch.float16, torch.bfloat16, torch.float32)):
145 + tensor = torch.arange(6, dtype=dtype).reshape(index + 1, -1)
146 + key = make_embedding_key(KEY, str(dtype))
147 + client.put_tensor(key, tensor)
148 + pointer = ctypes.addressof(client._put_header)
149 + header_ptr = header_ptr or pointer
150 + assert pointer == header_ptr
151 + expected = TensorSpec(tuple(tensor.shape), str(dtype), tensor.nbytes)
152 + with torch.device("meta"):
153 + loaded = client.load_tensors({key: expected}, "cpu")[key]
154 + assert client._read_buffer.device.type == "cpu"
155 + assert loaded.dtype == dtype and torch.equal(loaded, tensor)
156 + assert len(native.objects[key]) == 24 + tensor.nbytes
157 + assert (
158 + sum(c.args[0] == header_ptr for c in native.register_buffer.call_args_list) == 1
159 + )
160 + assert len(native.registered) == 2
161 +
162 +
163 +@pytest.mark.parametrize("capacity", [80, 120])
164 +def test_batch_reads_reuse_registration_without_aliasing_outputs(backend, capacity):
165 + client, native = backend.store_client, backend.store_client.store
166 + client._read_buffer_bytes = capacity
167 + tensors = {key: torch.full((2, 2), float(i)) for i, key in enumerate("abc")}
168 + for key, tensor in tensors.items():
169 + client.put_tensor(key, tensor)
170 + native.register_buffer = MagicMock(wraps=native.register_buffer)
171 + native.unregister_buffer = MagicMock(wraps=native.unregister_buffer)
172 + native.batch_get_into = MagicMock(wraps=native.batch_get_into)
173 + expected = dict.fromkeys(tensors, SPEC)
174 + first = client.load_tensors(expected, "cpu")
175 + for key, tensor in tensors.items():
176 + native.objects[key] = native.objects[key][:24] + struct.pack(
177 + "<4f", *[tensor[0, 0].item() + 10] * 4
178 + )
179 + second = client.load_tensors(expected, "cpu")
180 + assert all(torch.equal(first[key], tensor) for key, tensor in tensors.items())
181 + assert all(torch.equal(second[key], tensor + 10) for key, tensor in tensors.items())
182 + assert native.batch_get_into.call_count == (4 if capacity == 80 else 2)
183 + native.register_buffer.assert_called_once()
184 + native.unregister_buffer.assert_not_called()
185 + client.close()
186 + assert native.unregister_buffer.call_count == 2
187 + assert not native.registered
188 +
189 +
190 +@pytest.mark.parametrize("mode", ["miss", "oversized"])
191 +def test_unloadable_inputs_do_not_allocate_staging(backend, mode):
192 + client, native = backend.store_client, backend.store_client.store
193 + client._read_buffer_bytes = 39
194 + if mode == "oversized":
195 + client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2))
196 + native.register_buffer = MagicMock(wraps=native.register_buffer)
197 + assert not backend.resolve_inputs({"a": SPEC}, {}, "cpu")
198 + assert client._read_buffer is None
199 + if mode == "miss":
200 + assert client._put_header is None
201 + client.close()
202 + native.register_buffer.assert_not_called()
203 +
204 +
205 +@pytest.mark.parametrize(
206 + "results", [[40, -800], [40], [40, True], [[40], [40]], None, [40, 41]]
207 +)
208 +def test_unsafe_batch_result_retains_whole_staging(backend, results):
209 + client, native = backend.store_client, backend.store_client.store
210 + native.batch_get_into = MagicMock(return_value=results)
211 + native.unregister_buffer = MagicMock(wraps=native.unregister_buffer)
212 + for _ in range(2):
213 + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
214 + client.load_tensors(dict.fromkeys("ab", SPEC), "cpu")
215 + native.batch_get_into.assert_called_once()
216 + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
217 + client.close()
218 + native.unregister_buffer.assert_not_called()
219 + assert client._read_buffer is client._unsafe_owners[0]
220 +
221 +
222 +@pytest.mark.parametrize("failure", ["allocation", "copy", "record", "completion"])
223 +def test_copy_failure_preserves_staging_lifetime(backend, failure):
224 + """Emulate CUDA allocation/event boundaries without requiring a GPU."""
225 + client = backend.store_client
226 + for key in "ab":
227 + client.put_tensor(key, torch.ones(2, 2))
228 + allocate, copy = torch.empty, torch.Tensor.copy_
229 + targets: list[torch.Tensor] = []
230 + copies: list[torch.Tensor] = []
231 + ready = MagicMock()
232 + if failure == "record":
233 + ready.record.side_effect = RuntimeError("record failure")
234 + elif failure == "completion":
235 + ready.synchronize.side_effect = RuntimeError("copy failure")
236 +
237 + def allocate_on_cpu(*args, **kwargs):
238 + is_target = kwargs.get("device") == torch.device("cuda")
239 + if is_target and targets and failure == "allocation":
240 + raise torch.OutOfMemoryError("allocation failure")
241 + kwargs.update(device="cpu", pin_memory=False)
242 + tensor = allocate(*args, **kwargs)
243 + if is_target:
244 + targets.append(tensor)
245 + return tensor
246 +
247 + def record_copy(target, source, **kwargs):
248 + if copies and failure == "copy":
249 + raise RuntimeError("copy failure")
250 + copies.append(target)
251 + return copy(target, source, **kwargs)
252 +
253 + with (
254 + patch.object(torch, "empty", side_effect=allocate_on_cpu),
255 + patch.object(torch, "Event", return_value=ready),
256 + patch.object(torch.accelerator, "current_stream"),
257 + patch.object(torch.Tensor, "copy_", new=record_copy),
258 + pytest.raises(RuntimeError, match="allocation failure|unconfirmed"),
259 + ):
260 + client.load_tensors(dict.fromkeys("ab", SPEC), "cuda")
261 + if failure == "allocation":
262 + assert not copies and not client._poisoned
263 + else:
264 + assert len(copies) == (1 if failure == "copy" else 2)
265 + assert client._poisoned
266 + assert any(owner is client._read_buffer for owner in client._unsafe_owners)
267 + assert any(owner is ready for owner in client._unsafe_owners)
268 + assert all(
269 + any(owner is target for owner in client._unsafe_owners)
270 + for target in targets
271 + )
272 + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
273 + client.close()
274 +
275 +
276 +def test_non_matrix_output_is_rejected_before_registration(backend):
277 + client = backend.store_client
278 + client.store.register_buffer = MagicMock()
279 + with pytest.raises(store_client.EmbeddingStoreOperationError, match="2D"):
280 + client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2, 2))
281 + client.store.register_buffer.assert_not_called()
282 +
283 +
284 +def make_worker(backend):
285 + from vllm.distributed.ec_transfer.ec_connector.mooncake.worker import (
286 + ECMooncakeWorker,
287 + )
288 +
289 + worker = object.__new__(ECMooncakeWorker)
290 + worker.is_producer = True
291 + worker.is_consumer = False
292 + worker._output_store = backend
293 + worker._buffer_device = "cpu"
294 + worker._producer_pushes = MagicMock()
295 + worker._producer_pushes.reserve.return_value = (MagicMock(), True)
296 + worker._producer_pushes.poll.return_value = []
297 + worker._producer_pushes.cancel_requests.return_value = []
298 + worker._producer_pushes.pending = False
299 + worker._control_executor = MagicMock()
300 + worker._consumer_memory = MagicMock()
301 + worker._consumer_memory.drain_reclaimed.return_value = set()
302 + worker._completed_loads = set()
303 + worker._failed_loads = set()
304 + worker._failed_saves = set()
305 + worker._push_ready = threading.Event()
306 + worker._flush_pending_pushes = MagicMock()
307 + worker._bind_push_source = MagicMock()
308 + connector = object.__new__(ECMooncakeConnector)
309 + connector._worker = worker
310 + connector._role = ECConnectorRole.WORKER
311 + connector._is_producer = True
312 + connector._is_consumer = False
313 + connector._connector_metadata = None
314 + return connector
315 +
316 +
317 +@pytest.mark.parametrize(
318 + "mode",
319 + [
320 + "miss",
321 + "legacy",
322 + "hit",
323 + "partial",
324 + "get-failure",
325 + "short-data",
326 + "lease-expired",
327 + "registration-failure",
328 + ],
329 +)
330 +def test_reuse_pipeline(backend, mode):
331 + from vllm.multimodal.inputs import (
332 + MultiModalFeatureSpec,
333 + MultiModalKwargsItem,
334 + PlaceholderRange,
335 + )
336 + from vllm.v1.worker.gpu.ec_connector import ActiveECConnector
337 + from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
338 + from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
339 + from vllm.v1.worker.gpu.model_states.interface import ModelState
340 +
341 + native = backend.store_client.store
342 + values = {"a": 1.0, "b": 2.0, "c": 3.0}
343 + if mode == "legacy":
344 + old_namespace = KEY.replace("@protocol:v3", "@protocol:v2")
345 + for key in "abc":
346 + native.objects[make_embedding_key(old_namespace, key)] = b"old format"
347 + if mode == "partial":
348 + hits = {"a", "c"}
349 + elif mode not in ("miss", "legacy"):
350 + hits = set("abc")
351 + else:
352 + hits = set()
353 + for key in hits:
354 + backend.store_client.put_tensor(
355 + make_embedding_key(KEY, key), torch.full((2, 2), values[key])
356 + )
357 + if mode == "short-data":
358 + key = make_embedding_key(KEY, "b")
359 + native.objects[key] = native.objects[key][:-1]
360 + hits.remove("b")
361 + if mode in ("get-failure", "lease-expired"):
362 + read = native.batch_get_into
363 +
364 + def failed_read(keys, pointers, sizes):
365 + results = read(keys, pointers, sizes)
366 + results[keys.index(make_embedding_key(KEY, "b"))] = (
367 + -707 if mode == "lease-expired" else -704
368 + )
369 + return results
370 +
371 + native.batch_get_into = failed_read
372 + hits.remove("b")
373 + if mode == "registration-failure":
374 + native.register_buffer = MagicMock(return_value=-500)
375 + hits.clear()
376 + connector = make_worker(backend)
377 + metadata = ECMooncakeConnectorMetadata(
378 + store_candidates=dict.fromkeys("abc", SPEC),
379 + pushes=[
380 + ECMooncakePushSpec(
381 + mm_hash=k,
382 + shape=(2, 2),
383 + dtype="float32",
384 + nbytes=16,
385 + consumer_zmq="ipc:///tmp/test",
386 + transfer_id=k,
387 + request_id="req",
388 + )
389 + for k in "abc"
390 + ],
391 + )
392 + cache = EncoderCache()
393 + cache.mm_features["req"] = [
394 + MultiModalFeatureSpec(
395 + data=MultiModalKwargsItem({}),
396 + modality="image",
397 + identifier=k,
398 + mm_position=PlaceholderRange(offset=i * 2, length=2),
399 + )
400 + for i, k in enumerate("abc")
401 + ]
402 + runner = EncoderRunner(
403 + model=None,
404 + max_num_tokens=8,
405 + hidden_size=2,
406 + encoder_cache=cache,
407 + dtype=torch.float32,
408 + device=torch.device("cpu"),
409 + )
410 + runner.enable_timing = False
411 + missing = [k for k in "abc" if k not in hits]
412 + runner.execute_mm_encoder = MagicMock(
413 + return_value=[torch.full((2, 2), values[k] + 10) for k in missing]
414 + )
415 + state = SimpleNamespace(encoder_runner=runner, encoder_cache=cache)
416 + with patch(
417 + "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=connector
418 + ):
419 + adapter = ActiveECConnector(SimpleNamespace(), cache.encoder_outputs)
420 + with adapter.maybe_get_output(
421 + SimpleNamespace(ec_connector_metadata=metadata, finished_req_ids=set())
422 + ) as output:
423 + assert runner.prepare_mm_inputs({"req": [0, 1, 2]})[0] == missing
424 + ModelState.execute_mm_encoder(state, {"req": [0, 1, 2]})
425 + assert runner.execute_mm_encoder.call_count == bool(missing)
426 + assert output.ec_connector_worker_meta.pending_saves is False
427 + assert all(
428 + torch.equal(
429 + cache.encoder_outputs[k],
430 + torch.full((2, 2), values[k] + (0 if k in hits else 10)),
431 + )
432 + for k in "abc"
433 + )
434 + assert {
435 + c.args[1] for c in connector._worker._bind_push_source.call_args_list
436 + } == set("abc")
437 + backend.shutdown()
438 + assert not native.registered
439 + if mode in ("miss", "legacy"):
440 + other = store_backend.MooncakeEmbeddingStoreBackend(
441 + store_client.MooncakeEmbeddingStoreClient(native, read_buffer_bytes=1024),
442 + KEY,
443 + max_pending_items=32,
444 + max_pending_bytes=2 * 1024**3,
445 + )
446 + try:
447 + loaded: dict[str, torch.Tensor] = {}
448 + assert other.resolve_inputs(
449 + dict.fromkeys("abc", SPEC),
450 + loaded,
451 + "cpu",
452 + ) == set("abc")
453 + assert all(torch.equal(loaded[k], cache.encoder_outputs[k]) for k in "abc")
454 + finally:
455 + other.shutdown()
456 +
457 +
458 +@pytest.mark.parametrize(
459 + "mismatch",
460 + ["revision", "shape", "dtype", "magic", "short-header", "short-data", "long-data"],
461 +)
462 +def test_incompatible_output_is_a_miss(backend, mismatch):
463 + key = make_embedding_key(KEY, "a")
464 + dtype = torch.float16 if mismatch == "dtype" else torch.float32
465 + backend.store_client.put_tensor(key, torch.ones(2, 2, dtype=dtype))
466 + if mismatch == "revision":
467 + backend.namespace = KEY.replace("@revision:revision", "@revision:other")
468 + objects = backend.store_client.store.objects
469 + if mismatch == "magic":
470 + objects[key] = b"\0" * 4 + objects[key][4:]
471 + elif mismatch == "short-header":
472 + objects[key] = objects[key][:23]
473 + elif mismatch == "short-data":
474 + objects[key] = objects[key][:-1]
475 + elif mismatch == "long-data":
476 + objects[key] += b"extra"
477 + expected = replace(SPEC, shape=(4, 1)) if mismatch == "shape" else SPEC
478 + outputs: dict[str, torch.Tensor] = {}
479 + assert not backend.resolve_inputs({"a": expected}, outputs, "cpu")
480 + assert not outputs
481 + tensor = torch.ones(2, 2)
482 + backend.store_client.put_tensor(
483 + make_embedding_key(backend.namespace, "valid"), tensor
484 + )
485 + assert backend.resolve_inputs({"valid": SPEC}, outputs, "cpu") == {"valid"}
486 + assert torch.equal(outputs["valid"], tensor)
487 + assert len(backend.store_client.store.registered) == 2
488 +
489 +
490 +def test_partial_registration_rejection_releases_publication(backend):
491 + """A rejected header registration must not prevent a later publication."""
492 + native = backend.store_client.store
493 + register = native.register_buffer
494 + put = native.batch_put_from_multi_buffers
495 + native.register_buffer = lambda addr, size: (
496 + -500 if size == 24 else register(addr, size)
497 + )
498 + native.batch_put_from_multi_buffers = MagicMock()
499 + tensor = torch.ones(2, 2)
500 + owner = weakref.ref(tensor)
501 + assert backend._enqueue_save("a", tensor)
502 + storage_owner = weakref.ref(tensor.untyped_storage())
503 + future = backend._pending["a"][0]
504 + with pytest.raises(store_client.EmbeddingStoreOperationError, match="register"):
505 + future.result(timeout=5)
506 + del tensor
507 + gc.collect()
508 + assert owner() is None and not native.registered
509 + native.batch_put_from_multi_buffers.assert_not_called()
510 + backend.resolve_inputs({}, {}, "cpu")
511 + assert backend._pending_bytes == 0
512 + gc.collect()
513 + assert storage_owner() is None # The failed Future is still alive.
514 + assert backend.store_client._put_header is None
515 + native.register_buffer = register
516 + native.batch_put_from_multi_buffers = put
517 + backend.store_client.put_tensor("retry", torch.ones(2, 2))
518 + assert len(native.registered) == 1
519 +
520 +
521 +@pytest.mark.parametrize(
522 + "operation",
523 + ["read-unregister", "header-unregister", "put", "rejected-put", "unregister"],
524 +)
525 +def test_native_io_owners(backend, operation):
526 + """Uncertain I/O retains owners; completed rejection frees them even via Future."""
527 + native, client = backend.store_client.store, backend.store_client
528 + put = native.batch_put_from_multi_buffers
529 + backend.max_pending_bytes = SPEC.nbytes
530 + tensor = torch.ones(2, 2)
531 + owner = weakref.ref(tensor)
532 + rejected = operation == "rejected-put"
533 + if operation in ("read-unregister", "header-unregister"):
534 + key = make_embedding_key(KEY, "a")
535 + client.put_tensor(key, tensor)
536 + if operation == "read-unregister":
537 + client.load_tensors({key: SPEC}, "cpu")
538 + native.unregister_buffer = MagicMock(return_value=-500)
539 + with pytest.raises(store_client.EmbeddingStoreError, match="unregister"):
540 + client.close()
541 + owner = weakref.ref(
542 + client._read_buffer
543 + if operation == "read-unregister"
544 + else client._put_header
545 + )
546 + else:
547 + if operation == "unregister":
548 + native.unregister_buffer = MagicMock(return_value=-500)
549 + else:
550 + native.batch_put_from_multi_buffers = MagicMock(
551 + return_value=[-1700 if rejected else -800]
552 + )
553 + assert backend._enqueue_save("a", tensor)
554 + future = backend._pending["a"][0]
555 + with pytest.raises(store_client.EmbeddingStoreError):
556 + future.result(timeout=5)
557 + del tensor
558 + gc.collect()
559 + if rejected:
560 + assert owner() is None # A live Future must not retain the tensor.
561 + tensor = torch.ones(4, 4) # Too large for admission after reaping.
562 + caller_owner = weakref.ref(tensor)
563 + backend._step_candidates = {"next"}
564 + backend.save_output("next", tensor)
565 + del tensor
566 + gc.collect()
567 + assert caller_owner() is None
568 + assert backend._pending_bytes == 0 and len(native.registered) == 1
569 + header = client._put_header
570 + native.batch_put_from_multi_buffers = put
571 + client.put_tensor("retry", torch.ones(1, 4))
572 + assert client._put_header is header
573 + else:
574 + assert owner() is not None and native.registered
575 + header = bytes(client._put_header)
576 + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
577 + client.put_tensor("retry", torch.ones(1, 4))
578 + assert bytes(client._put_header) == header
579 + with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"):
580 + client.close()
581 +
582 +
583 +def test_publication_budget_charges_and_releases_backing_storage(backend):
584 + tensor = torch.arange(16, dtype=torch.float32)[:4].view(2, 2)
585 + owner = weakref.ref(tensor)
586 + future: Future[None] = Future()
587 + backend.max_pending_items = 1
588 + backend.max_pending_bytes = tensor.nbytes
589 + with patch.object(backend._executor, "submit", return_value=future):
590 + assert not backend._enqueue_save("a", tensor)
591 + backend.max_pending_bytes = 64
592 + assert backend._enqueue_save("a", tensor)
593 + storage_owner = weakref.ref(tensor.untyped_storage())
594 + assert not backend._enqueue_save("b", tensor)
595 + del tensor
596 + gc.collect()
597 + backend.reap() # The final step-end poll precedes publication completion.
598 + assert owner() is not None
599 + backend._save(make_embedding_key(KEY, "a"), backend._pending["a"][1])
600 + future.set_result(None)
601 + assert owner() is None and storage_owner() is not None
602 +
603 + def load_after_reaping(*args):
604 + assert not backend._pending and backend._pending_bytes == 0
605 + assert storage_owner() is None
606 + return {}
607 +
608 + with patch.object(backend.store_client, "load_tensors", load_after_reaping):
609 + assert not backend.resolve_inputs({"b": SPEC}, {}, "cpu")
610 +
611 +
612 +@pytest.mark.parametrize("rejected", [False, True])
613 +def test_shared_storage_budget_is_released_after_last_publication(backend, rejected):
614 + batch = torch.arange(16, dtype=torch.float32).view(4, 4)
615 + first: Future[None] = Future()
616 + last: Future[None] = Future()
617 + separate: Future[None] = Future()
618 + backend.max_pending_items = 2
619 + backend.max_pending_bytes = batch.nbytes
620 + with patch.object(backend._executor, "submit", side_effect=[first, last, separate]):
621 + assert backend._enqueue_save("a", batch[:2])
622 + assert backend._enqueue_save("b", batch[2:])
623 + assert backend._pending_bytes == batch.nbytes
624 + assert not backend._enqueue_save("a", batch[:2])
625 + assert not backend._enqueue_save("extra", batch[:1])
626 + if rejected:
627 + first.set_exception(store_client.EmbeddingStoreOperationError("rejected"))
628 + else:
629 + first.set_result(None)
630 + backend.reap()
631 + assert backend._pending_bytes == batch.nbytes
632 + assert not backend._enqueue_save("c", torch.ones(2, 2))
633 + last.set_result(None)
634 + backend.reap()
635 + assert backend._pending_bytes == 0
636 + assert backend._enqueue_save("c", torch.ones(2, 2))
637 + assert backend._pending_bytes == 16
638 + separate.set_result(None)
639 + backend.reap()
640 + assert backend._pending_bytes == 0
641 +
642 +
643 +@pytest.mark.parametrize("shared_storage", [False, True])
644 +def test_completion_race_preserves_failure_and_reclaims_other_items(
645 + backend, shared_storage
646 +):
647 + class RacingFuture(Future[None]):
648 + def done(self):
649 + observed = super().done()
650 + if not observed:
651 + self.set_exception(store_client.EmbeddingStoreError("uncertain"))
652 + return observed
653 +
654 + racing = RacingFuture()
655 + other: Future[None] = Future()
656 + tensor = torch.ones(2, 2)
657 + other_tensor = tensor[:1] if shared_storage else torch.ones(2, 2)
658 + with patch.object(backend._executor, "submit", side_effect=[racing, other]):
659 + assert backend._enqueue_save("race", tensor)
660 + assert backend._enqueue_save("other", other_tensor)
661 + backend.reap()
662 + assert backend._pending_bytes == (16 if shared_storage else 32)
663 + other.set_result(None)
664 + with pytest.raises(store_client.EmbeddingStoreError):
665 + backend.reap()
666 + assert backend._pending_bytes == 16
667 + with pytest.raises(store_client.EmbeddingStoreError):
668 + backend._enqueue_save("later", torch.ones(2, 2))
669 +
670 +
671 +def test_readiness_and_shutdown_do_not_block_p2p_completion(backend):
672 + ready, waiting, entered = [threading.Event() for _ in range(3)]
673 +
674 + original = backend._executor.shutdown
675 +
676 + def closing(*args, **kwargs):
677 + entered.set()
678 + return original(*args, **kwargs)
679 +
680 + native = backend.store_client.store
681 +
682 + def close_store():
683 + assert ready.is_set() and not native.registered
684 + assert backend._pending_bytes == 0
685 + return 0
686 +
687 + native.close = MagicMock(side_effect=close_store)
688 + put = native.batch_put_from_multi_buffers
689 +
690 + def waiting_put(*args):
691 + waiting.set()
692 + assert ready.wait(5)
693 + assert backend.store_client._put_header is not None and native.registered
694 + return put(*args)
695 +
696 + native.batch_put_from_multi_buffers = waiting_put
697 + connector = make_worker(backend)
698 + connector.bind_connector_metadata(
699 + ECMooncakeConnectorMetadata(store_candidates={"a": SPEC})
700 + )
701 + cache: dict[str, torch.Tensor] = {}
702 + with (
703 + patch.object(
704 + store_backend, "_record_tensor_ready_event", return_value=MagicMock()
705 + ),
706 + patch.object(backend._executor, "shutdown", side_effect=closing),
707 + ThreadPoolExecutor(max_workers=1) as closer,
708 + ):
709 + connector.start_save_caches(encoder_cache=cache)
710 + cache["a"] = torch.arange(4, dtype=torch.float32).view(2, 2)
711 + connector.save_caches(cache, "a")
712 + assert connector.build_connector_worker_meta().pending_saves is False
713 + assert waiting.wait(5) and not backend.store_client.store.objects
714 + close = closer.submit(backend.shutdown)
715 + try:
716 + assert entered.wait(5) and not close.done()
717 + assert not backend._enqueue_save("late", torch.ones(2, 2))
718 + finally:
719 + ready.set()
720 + close.result(timeout=5)
721 + native.close.assert_called_once_with()
722 + assert not backend.store_client.store.registered and backend._pending_bytes == 0
723 +
724 + assert backend.store_client.store.objects[make_embedding_key(KEY, "a")][
725 + 24:
726 + ] == struct.pack("<4f", 0, 1, 2, 3)
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+27 -10
@@ -17,6 +17,7 @@ import pytest
17 import torch
18
19 from tests.v1.attention.utils import dense_kv_cache_views
20 +from vllm.distributed import mooncake_store
21 from vllm.distributed.kv_events import KVEventAggregator
22 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
23 rdma_utils,
@@ -402,6 +403,7 @@ def _install_fake_mooncake(monkeypatch, store_instance: MagicMock):
403 fake_store_module = types.ModuleType("mooncake.store")
404 fake_store_module.MooncakeDistributedStore = lambda: store_instance # type: ignore[attr-defined]
405 fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined]
406 + fake_store_module.ObjectDataType = SimpleNamespace(TENSOR=1) # type: ignore[attr-defined]
407 fake_mooncake_module = types.ModuleType("mooncake")
408 fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined]
409 monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module)
@@ -590,7 +592,7 @@ def test_pool_key_cache_prefix_namespaces_and_disambiguates():
592 def test_default_local_buffer_size_matches_pr40900():
593 """PR-40900 shipped a 4 GiB default for local_buffer_size; the dual-mode
594 patch preserves it (and the JSON key) so unchanged PR-40900 configs work."""
593 - assert worker.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3
595 + assert mooncake_store.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3
596
597
598 def test_get_requester_local_hostname_prefers_override(monkeypatch):
@@ -2957,6 +2959,7 @@ def test_requester_worker_group_semantics_string_true_enables(
2959 fake_store_module = types.ModuleType("mooncake.store")
2960 fake_store_module.MooncakeDistributedStore = lambda: store # type: ignore[attr-defined]
2961 fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined]
2962 + fake_store_module.ObjectDataType = SimpleNamespace(TENSOR=1) # type: ignore[attr-defined]
2963 fake_mooncake_module = types.ModuleType("mooncake")
2964 fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined]
2965 monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module)
@@ -4377,10 +4380,10 @@ def test_config_defaults_to_embedded():
4380 """A JSON without explicit mode parses as embedded with 4 GiB segment."""
4381 cfg = _make_config()
4382 assert cfg.mode == "embedded"
4380 - assert cfg.global_segment_size == worker.DEFAULT_GLOBAL_SEGMENT_SIZE
4381 - assert cfg.local_buffer_size == worker.DEFAULT_LOCAL_BUFFER_SIZE
4383 + assert cfg.global_segment_size == mooncake_store.DEFAULT_GLOBAL_SEGMENT_SIZE
4384 + assert cfg.local_buffer_size == mooncake_store.DEFAULT_LOCAL_BUFFER_SIZE
4385 assert cfg.enable_offload is False
4383 - assert cfg.tenant_id == worker.DEFAULT_TENANT_ID
4386 + assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID
4387
4388
4389 def test_config_pr40900_unchanged(tmp_path):
@@ -4402,7 +4405,7 @@ def test_config_pr40900_unchanged(tmp_path):
4405 assert cfg.global_segment_size == 4 * 1024**3
4406 assert cfg.local_buffer_size == 4 * 1024**3
4407 assert cfg.enable_offload is False
4405 - assert cfg.tenant_id == worker.DEFAULT_TENANT_ID
4408 + assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID
4409
4410
4411 def test_config_from_file_normalizes_tenant_id(tmp_path):
@@ -4432,7 +4435,7 @@ def test_config_from_file_normalizes_empty_tenant_id_to_default(tmp_path):
4435
4436 cfg = worker.MooncakeStoreConfig.from_file(config_path)
4437
4435 - assert cfg.tenant_id == worker.DEFAULT_TENANT_ID
4438 + assert cfg.tenant_id == mooncake_store.DEFAULT_TENANT_ID
4439
4440
4441 def test_config_from_file_rejects_non_string_tenant_id(tmp_path):
@@ -4584,7 +4587,20 @@ def test_topology_embedded_cpu_only(tmp_path, monkeypatch):
4587 assert w.disk_offload_buffer_budget_bytes is None
4588
4589
4587 -def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch):
4590 +def _create_store_client_for_tenant_test(consumer):
4591 + if consumer == "ec":
4592 + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.store_client import ( # noqa: E501
4593 + create_mooncake_embedding_store_client,
4594 + )
4595 +
4596 + return create_mooncake_embedding_store_client()
4597 + return worker.MooncakeStoreWorker(
4598 + _make_vllm_config(rank=1), _make_kv_cache_config()
4599 + )
4600 +
4601 +
4602 +@pytest.mark.parametrize("consumer", ["kv", "ec"])
4603 +def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch, consumer):
4604 store = MagicMock()
4605 store.setup.return_value = 0
4606 _install_fake_mooncake(monkeypatch, store)
@@ -4605,12 +4621,13 @@ def test_topology_forwards_non_default_tenant_id(tmp_path, monkeypatch):
4621 ),
4622 )
4623
4608 - worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config())
4624 + _create_store_client_for_tenant_test(consumer)
4625
4626 assert store.setup.call_args.kwargs == {"tenant_id": "tenant-a"}
4627
4628
4613 -def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch):
4629 +@pytest.mark.parametrize("consumer", ["kv", "ec"])
4630 +def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch, consumer):
4631 store = MagicMock()
4632 setup_error = TypeError(
4633 "setup(): incompatible function arguments; "
@@ -4636,7 +4653,7 @@ def test_non_default_tenant_preserves_setup_type_error(tmp_path, monkeypatch):
4653 )
4654
4655 with pytest.raises(TypeError) as exc_info:
4639 - worker.MooncakeStoreWorker(_make_vllm_config(rank=1), _make_kv_cache_config())
4656 + _create_store_client_for_tenant_test(consumer)
4657
4658 assert exc_info.value is setup_error
4659
vllm/config/ec_transfer.py
+2
@@ -23,6 +23,8 @@ class ECTransferConfig:
23 safetensors) and ``ECMooncakeConnector`` (Mooncake TransferEngine RDMA;
24 requires ``mooncake-transfer-engine`` and matching producer/consumer
25 ``ec_connector_extra_config``; see ``mooncake_ec_connector`` module docstring).
26 + Set ``cross_encoder_cache`` in Mooncake extra config to reuse shared
27 + Encoder outputs from Store before encoding, retaining P2P delivery.
28 """
29
30 engine_id: str | None = None
vllm/distributed/ec_transfer/ec_connector/mooncake/config.py
+32
@@ -69,6 +69,10 @@ class MooncakeECConfig:
69 control_timeout_ms: int
70 push_wait_timeout_s: float
71 pool_size: int
72 + cross_encoder_cache: bool = False
73 + store_max_pending_items: int = 32
74 + store_max_pending_bytes: int = 2 * 1024**3
75 + store_read_buffer_bytes: int = 128 * 1024**2
76
77 @classmethod
78 def from_vllm_config(cls, vllm_config: VllmConfig) -> MooncakeECConfig:
@@ -94,6 +98,24 @@ class MooncakeECConfig:
98 "ec_buffer_size", ec_config.ec_buffer_size
99 )
100 get = ec_config.get_from_extra_config
101 + shared_reuse = get("cross_encoder_cache", False)
102 + if not isinstance(shared_reuse, bool):
103 + raise ValueError("cross_encoder_cache must be a boolean")
104 + if shared_reuse:
105 + if not ec_config.is_encode_only:
106 + raise ValueError(
107 + "cross_encoder_cache requires an Encoder-only producer"
108 + )
109 + if not vllm_config.use_v2_model_runner:
110 + raise ValueError("cross_encoder_cache requires Model Runner V2")
111 + if vllm_config.lora_config is not None:
112 + raise ValueError("cross_encoder_cache does not support LoRA")
113 + mm_config = vllm_config.model_config.multimodal_config
114 + if mm_config is not None and mm_config.mm_processor_cache_gb == 0:
115 + raise ValueError(
116 + "cross_encoder_cache requires mm_processor_cache_gb > 0 "
117 + "to preserve content identifiers across Encoders and restarts"
118 + )
119 control_port = int(ec_config.ec_port) + (
120 parallel_config.data_parallel_index * parallel_config.tensor_parallel_size
121 )
@@ -120,4 +142,14 @@ class MooncakeECConfig:
142 "push_wait_timeout_s", get("push_wait_timeout_s", 60)
143 ),
144 pool_size=registered_buffer_size,
145 + cross_encoder_cache=shared_reuse,
146 + store_max_pending_items=_positive_int(
147 + "store_max_pending_items", get("store_max_pending_items", 32)
148 + ),
149 + store_max_pending_bytes=_positive_int(
150 + "store_max_pending_bytes", get("store_max_pending_bytes", 2 * 1024**3)
151 + ),
152 + store_read_buffer_bytes=_positive_int(
153 + "store_read_buffer_bytes", get("store_read_buffer_bytes", 128 * 1024**2)
154 + ),
155 )
vllm/distributed/ec_transfer/ec_connector/mooncake/metadata.py
+4
@@ -10,6 +10,9 @@ from vllm.distributed.ec_transfer.ec_connector.base import (
10 ECConnectorMetadata,
11 ECConnectorWorkerMetadata,
12 )
13 +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
14 + TensorSpec,
15 +)
16
17
18 @dataclass
@@ -46,6 +49,7 @@ class ECMooncakeConnectorMetadata(ECConnectorMetadata):
49 loads: list[ECMooncakeLoadSpec] = field(default_factory=list)
50 pushes: list[ECMooncakePushSpec] = field(default_factory=list)
51 freed: list[str] | None = None
52 + store_candidates: dict[str, TensorSpec] = field(default_factory=dict)
53
54
55 @dataclass
vllm/distributed/ec_transfer/ec_connector/mooncake/scheduler.py
+34 -10
@@ -45,6 +45,9 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake.state import (
45 from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import (
46 ensure_mooncake_available,
47 )
48 +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
49 + TensorSpec,
50 +)
51 from vllm.distributed.ec_transfer.ec_connector.utils import (
52 PlaceholderMetadataResolver,
53 collect_ec_item_metadata,
@@ -76,6 +79,8 @@ class ECMooncakeScheduler:
79 ensure_mooncake_available()
80 config = MooncakeECConfig.from_vllm_config(vllm_config)
81
82 + self._store_candidates: list[tuple[str, str, TensorSpec]] = []
83 + self._cross_encoder_cache = config.cross_encoder_cache
84 self._is_producer = config.is_producer
85 self._is_consumer = config.is_consumer
86 self._control_addr = config.control_addr
@@ -427,28 +432,40 @@ class ECMooncakeScheduler:
432 transfer_id = f"{request.request_id}:{index}"
433 if not consumer_zmq or transfer_id in self._prepared_push_transfer_ids:
434 return
430 - num_tokens = request.get_num_encoder_embeds(index)
431 - dtype = self._model_config.dtype
432 - assert isinstance(dtype, torch.dtype)
433 - assert self._encoder_cache_hidden_dim is not None
434 - dtype_name = str(dtype).split(".")[-1]
435 - shape = (num_tokens, self._encoder_cache_hidden_dim)
436 - nbytes = math.prod(shape) * dtype.itemsize
435 + output = self._encoder_output_spec(request, index)
436 self._pushes_to_prepare[transfer_id] = ECMooncakePushSpec(
437 mm_hash=mm_hash,
439 - nbytes=nbytes,
440 - shape=shape,
441 - dtype=dtype_name,
438 + nbytes=output.nbytes,
439 + shape=output.shape,
440 + dtype=output.dtype.removeprefix("torch."),
441 consumer_zmq=str(consumer_zmq),
442 transfer_id=transfer_id,
443 request_id=request.request_id,
444 )
445 self._prepared_push_transfer_ids.add(transfer_id)
446
447 + def _encoder_output_spec(self, request: Any, index: int) -> TensorSpec:
448 + dtype = self._model_config.dtype
449 + assert isinstance(dtype, torch.dtype)
450 + assert self._encoder_cache_hidden_dim is not None
451 + shape = (request.get_num_encoder_embeds(index), self._encoder_cache_hidden_dim)
452 + return TensorSpec(shape, str(dtype), math.prod(shape) * dtype.itemsize)
453 +
454 def update_state_after_alloc(self, request: Any, index: int) -> None:
455 self._local_cache.add(request.mm_features[index].identifier)
456 if self._is_producer:
457 self._prepare_push_spec(request, index)
458 + if (
459 + self._cross_encoder_cache
460 + and request.mm_features[index].modality == "image"
461 + ):
462 + self._store_candidates.append(
463 + (
464 + request.request_id,
465 + request.mm_features[index].identifier,
466 + self._encoder_output_spec(request, index),
467 + )
468 + )
469
470 def update_state_after_free(self, request: Any, index: int) -> None:
471 if not self._is_consumer:
@@ -474,6 +491,13 @@ class ECMooncakeScheduler:
491 meta = ECMooncakeConnectorMetadata(
492 freed=scheduler_output.free_encoder_mm_hashes
493 )
494 + preempted = scheduler_output.preempted_req_ids or set()
495 + meta.store_candidates = {
496 + identifier: output
497 + for request_id, identifier, output in self._store_candidates
498 + if request_id not in preempted
499 + }
500 + self._store_candidates.clear()
501 for push_spec in self._pushes_to_prepare.values():
502 meta.pushes.append(push_spec)
503 self._pushes_to_prepare.clear()
vllm/distributed/ec_transfer/ec_connector/mooncake/worker.py
+48
@@ -49,6 +49,9 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import (
49 MooncakeTransfer,
50 ensure_mooncake_available,
51 )
52 +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
53 + build_embedding_namespace,
54 +)
55 from vllm.logger import init_logger
56 from vllm.utils.network_utils import get_ip
57
@@ -59,6 +62,10 @@ _T = TypeVar("_T")
62 if TYPE_CHECKING:
63 from vllm.config import VllmConfig
64
65 + from ..mooncake_store_embedding.backend import (
66 + MooncakeEmbeddingStoreBackend,
67 + )
68 +
69 _RESERVATION_REFRESH_SECONDS = _RESERVATION_TTL_SECONDS / 2
70 _MAX_CANCELLED_TRANSFER_IDS = 1 << 16
71 _CANCEL_ATTEMPTS = 2
@@ -87,6 +94,17 @@ class ECMooncakeWorker:
94 def __init__(self, vllm_config: VllmConfig) -> None:
95 ensure_mooncake_available()
96 config = MooncakeECConfig.from_vllm_config(vllm_config)
97 + self._store_config = (
98 + (
99 + build_embedding_namespace(vllm_config),
100 + config.store_max_pending_items,
101 + config.store_max_pending_bytes,
102 + config.store_read_buffer_bytes,
103 + )
104 + if config.cross_encoder_cache
105 + else None
106 + )
107 + self._output_store: MooncakeEmbeddingStoreBackend | None = None
108 self.is_producer = config.is_producer
109 self.is_consumer = config.is_consumer
110 self._buffer_device = config.buffer_device
@@ -173,6 +191,21 @@ class ECMooncakeWorker:
191 self._is_receiving_rank = True
192
193 def start_services(self) -> None:
194 + if self._store_config is not None and self._output_store is None:
195 + from ..mooncake_store_embedding.backend import (
196 + MooncakeEmbeddingStoreBackend,
197 + )
198 + from ..mooncake_store_embedding.store_client import (
199 + create_mooncake_embedding_store_client,
200 + )
201 +
202 + namespace, max_items, max_bytes, read_bytes = self._store_config
203 + self._output_store = MooncakeEmbeddingStoreBackend(
204 + create_mooncake_embedding_store_client(read_buffer_bytes=read_bytes),
205 + namespace,
206 + max_pending_items=max_items,
207 + max_pending_bytes=max_bytes,
208 + )
209 if not self.is_consumer or self._control_server is not None:
210 return
211 self._resolve_consumer_rank()
@@ -470,6 +503,14 @@ class ECMooncakeWorker:
503 tensor = encoder_cache.get(mm_hash)
504 if tensor is not None:
505 self._bind_push_source(tensor, mm_hash)
506 + if self._output_store is not None:
507 + loaded = self._output_store.resolve_inputs(
508 + metadata.store_candidates, encoder_cache, self._buffer_device
509 + )
510 + # Loaded entries precede the runner's new-output snapshot, so they
511 + # need explicit P2P source binding here.
512 + for mm_hash in loaded:
513 + self._bind_push_source(encoder_cache[mm_hash], mm_hash)
514
515 def _reserve_batch(self, records: list[ProducerPushRecord]) -> None:
516 """Resolve independent item futures with one reserve RPC per shard."""
@@ -851,6 +892,8 @@ class ECMooncakeWorker:
892 if not self.is_producer:
893 return None, None
894
895 + if self._output_store is not None:
896 + self._output_store.reap()
897 for record in self._producer_pushes.cancel_requests(finished_req_ids):
898 self._producer_pushes.submit_cancel(
899 record,
@@ -866,6 +909,8 @@ class ECMooncakeWorker:
909 return
910 tensor = encoder_cache[mm_hash]
911 self._bind_push_source(tensor, mm_hash)
912 + if self._output_store is not None:
913 + self._output_store.save_output(mm_hash, tensor)
914
915 def build_connector_worker_meta(self) -> ECMooncakeWorkerMetadata | None:
916 if self.is_consumer and not self._is_receiving_rank:
@@ -938,3 +983,6 @@ class ECMooncakeWorker:
983 )
984 self._producer_memory.close()
985 self._transfer.close()
986 + if self._output_store is not None:
987 + self._output_store.shutdown()
988 + self._output_store = None
vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/__init__.py new
+3
@@ -0,0 +1,3 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Mooncake Store support for cross-Encoder output reuse."""
vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/backend.py new
+214
@@ -0,0 +1,214 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Encoder-side Mooncake Store resolver and publisher."""
4 +
5 +from __future__ import annotations
6 +
7 +import threading
8 +import traceback
9 +from collections.abc import Mapping
10 +from concurrent.futures import Future, ThreadPoolExecutor
11 +from dataclasses import dataclass
12 +
13 +import torch
14 +
15 +from vllm.logger import init_logger
16 +
17 +from .data import TensorSpec, make_embedding_key
18 +from .store_client import (
19 + EmbeddingStoreError,
20 + EmbeddingStoreOperationError,
21 + MooncakeEmbeddingStoreClient,
22 +)
23 +
24 +logger = init_logger(__name__)
25 +
26 +
27 +@dataclass
28 +class _StoreSave:
29 + tensor: torch.Tensor | None
30 + ready_event: torch.Event | None
31 + storage_ptr: int
32 +
33 + def release(self) -> None:
34 + self.tensor = None
35 + self.ready_event = None
36 +
37 +
38 +class MooncakeEmbeddingStoreBackend:
39 + """Resolve immutable encoder outputs and publish misses asynchronously."""
40 +
41 + def __init__(
42 + self,
43 + store_client: MooncakeEmbeddingStoreClient,
44 + namespace: str,
45 + *,
46 + max_pending_items: int,
47 + max_pending_bytes: int,
48 + ) -> None:
49 + self.store_client = store_client
50 + self.namespace = namespace
51 + self.max_pending_items = max_pending_items
52 + self.max_pending_bytes = max_pending_bytes
53 + self._executor = ThreadPoolExecutor(
54 + max_workers=1,
55 + thread_name_prefix="ec-mooncake-store",
56 + )
57 + self._pending: dict[str, tuple[Future[None], _StoreSave]] = {}
58 + self._pending_storages: dict[int, tuple[torch.UntypedStorage, int]] = {}
59 + self._pending_lock = threading.Lock()
60 + self._pending_bytes = 0
61 + self._step_candidates: set[str] = set()
62 + self._closed = False
63 + self._fatal_error: BaseException | None = None
64 +
65 + def save_output(self, identifier: str, tensor: torch.Tensor) -> None:
66 + if identifier in self._step_candidates:
67 + self.reap()
68 + self._enqueue_save(identifier, tensor)
69 +
70 + def resolve_inputs(
71 + self,
72 + expected_tensors: Mapping[str, TensorSpec],
73 + encoder_cache: dict[str, torch.Tensor],
74 + device: torch.device | str,
75 + ) -> set[str]:
76 + """Load outputs before encoding; unresolved items retain normal computation."""
77 + self.reap()
78 + self._step_candidates = set(expected_tensors)
79 + candidates = [key for key in expected_tensors if key not in encoder_cache]
80 + if not candidates:
81 + return set()
82 +
83 + pool_keys = [
84 + make_embedding_key(self.namespace, identifier) for identifier in candidates
85 + ]
86 + try:
87 + exists = self.store_client.batch_exists(pool_keys)
88 + except (EmbeddingStoreOperationError, OSError):
89 + logger.warning(
90 + "Mooncake embedding Store lookup failed; falling back to encoder",
91 + exc_info=True,
92 + )
93 + return set()
94 +
95 + hits = {
96 + pool_key: expected_tensors[identifier]
97 + for identifier, pool_key, hit in zip(
98 + candidates, pool_keys, exists, strict=True
99 + )
100 + if hit
101 + }
102 + try:
103 + tensors = self.store_client.load_tensors(hits, device)
104 + except (EmbeddingStoreOperationError, OSError):
105 + logger.warning(
106 + "Mooncake embedding Store GET failed; falling back to encoder",
107 + exc_info=True,
108 + )
109 + return set()
110 + loaded: set[str] = set()
111 + for identifier, pool_key in zip(candidates, pool_keys, strict=True):
112 + if pool_key in tensors:
113 + encoder_cache[identifier] = tensors[pool_key]
114 + loaded.add(identifier)
115 + return loaded
116 +
117 + def _enqueue_save(self, identifier: str, tensor: torch.Tensor) -> bool:
118 + """Admit contiguous outputs, charging each retained storage once."""
119 + if not tensor.is_contiguous():
120 + return False
121 + storage = tensor.untyped_storage()
122 + storage_ptr = storage.data_ptr()
123 + with self._pending_lock:
124 + if self._fatal_error is not None:
125 + raise EmbeddingStoreError(
126 + "Store publisher has failed"
127 + ) from self._fatal_error
128 + if self._closed or identifier in self._pending:
129 + return False
130 + _, ref_count = self._pending_storages.get(storage_ptr, (storage, 0))
131 + budget_bytes = 0 if ref_count else storage.nbytes()
132 + if (
133 + len(self._pending) >= self.max_pending_items
134 + or self._pending_bytes + budget_bytes > self.max_pending_bytes
135 + ):
136 + return False
137 + # Only admitted work creates an event. Runtime errors propagate.
138 + save = _StoreSave(
139 + tensor,
140 + _record_tensor_ready_event(tensor),
141 + storage_ptr,
142 + )
143 + future = self._executor.submit(
144 + self._save, make_embedding_key(self.namespace, identifier), save
145 + )
146 + self._pending[identifier] = (future, save)
147 + # Keep the storage alive until reaping so its address cannot be reused.
148 + self._pending_storages[storage_ptr] = (storage, ref_count + 1)
149 + self._pending_bytes += budget_bytes
150 + return True
151 +
152 + def _save(self, pool_key: str, save: _StoreSave) -> None:
153 + try:
154 + if self.store_client.exists(pool_key):
155 + return
156 + tensor = save.tensor
157 + assert tensor is not None
158 + if save.ready_event is not None:
159 + save.ready_event.synchronize()
160 + self.store_client.put_tensor(pool_key, tensor)
161 + except BaseException as error:
162 + # Returned client frames can otherwise keep tensors via a Future's
163 + # traceback. Uncertain I/O owners remain retained by the client.
164 + traceback.clear_frames(error.__traceback__)
165 + raise
166 + finally:
167 + tensor = None
168 + save.release()
169 +
170 + def reap(self) -> None:
171 + with self._pending_lock:
172 + for identifier, (future, save) in list(self._pending.items()):
173 + # Observe once; a racing completion is handled on the next poll.
174 + if not future.done():
175 + continue
176 + try:
177 + future.result() # Already done: never waits on Store I/O.
178 + except (EmbeddingStoreOperationError, OSError) as error:
179 + logger.warning("Store PUT failed for %s: %s", identifier, error)
180 + # This failure is consumed; do not retain publisher/caller frames.
181 + error.__traceback__ = None
182 + except BaseException as error:
183 + # The client may still own unsafe native-I/O buffers. Keep
184 + # their charge and fatal result until worker termination.
185 + if self._fatal_error is None:
186 + self._fatal_error = error
187 + continue
188 + del self._pending[identifier]
189 + storage, ref_count = self._pending_storages[save.storage_ptr]
190 + if ref_count == 1:
191 + del self._pending_storages[save.storage_ptr]
192 + self._pending_bytes -= storage.nbytes()
193 + else:
194 + self._pending_storages[save.storage_ptr] = (storage, ref_count - 1)
195 + fatal = self._fatal_error
196 + if fatal is not None:
197 + raise fatal
198 +
199 + def shutdown(self) -> None:
200 + with self._pending_lock:
201 + if self._closed:
202 + return
203 + self._closed = True
204 + self._executor.shutdown(wait=True, cancel_futures=False)
205 + self.reap()
206 + self.store_client.close()
207 +
208 +
209 +def _record_tensor_ready_event(tensor: torch.Tensor) -> torch.Event | None:
210 + if tensor.device.type != "cuda":
211 + return None
212 + event = torch.Event()
213 + event.record(torch.accelerator.current_stream(tensor.device))
214 + return event
vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/data.py new
+51
@@ -0,0 +1,51 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Content identities and output contracts for the shared Encoder Store."""
4 +
5 +from __future__ import annotations
6 +
7 +from dataclasses import dataclass
8 +from typing import TYPE_CHECKING
9 +
10 +if TYPE_CHECKING:
11 + from vllm.config import VllmConfig
12 +
13 +EMBEDDING_CACHE_KEY_VERSION = "v3"
14 +
15 +
16 +def make_embedding_key(namespace: str, identifier: str) -> str:
17 + return f"{namespace}@id:{identifier}"
18 +
19 +
20 +@dataclass(frozen=True)
21 +class TensorSpec:
22 + """Contiguous output contract used by the Encoder and Store codec."""
23 +
24 + shape: tuple[int, ...]
25 + dtype: str
26 + nbytes: int
27 +
28 +
29 +def build_embedding_namespace(vllm_config: VllmConfig) -> str:
30 + model_config = vllm_config.model_config
31 + ec_config = vllm_config.ec_transfer_config
32 + assert ec_config is not None
33 + extra_config = ec_config.ec_connector_extra_config
34 +
35 + multimodal_config = model_config.multimodal_config
36 + encoder_hash = (
37 + multimodal_config.compute_hash()
38 + if multimodal_config is not None
39 + else "encoder:default"
40 + )
41 + cache_prefix = str(extra_config.get("embedding_cache_prefix", ""))
42 + model_name = str(extra_config.get("embedding_model_identity", model_config.model))
43 + prefix = f"{cache_prefix}@" if cache_prefix else ""
44 + return (
45 + f"{prefix}embedding"
46 + f"@model:{model_name}"
47 + f"@revision:{model_config.revision or 'default'}"
48 + f"@encoder:{encoder_hash}"
49 + f"@dtype:{model_config.dtype}"
50 + f"@protocol:{EMBEDDING_CACHE_KEY_VERSION}"
51 + )
vllm/distributed/ec_transfer/ec_connector/mooncake_store_embedding/store_client.py new
+463
@@ -0,0 +1,463 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Thin Mooncake Store client for embedding objects.
4 +
5 +Originally adapted from vLLM PR #47302 (0d1f71f5d36c). The v3 key namespace
6 +uses a compact header for complete 2D embeddings, not that PR's tensor format.
7 +"""
8 +
9 +from __future__ import annotations
10 +
11 +import ctypes
12 +import struct
13 +import threading
14 +from collections.abc import Callable, Iterator, Mapping
15 +from contextlib import contextmanager
16 +from enum import IntEnum
17 +from typing import Any
18 +
19 +import torch
20 +
21 +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import (
22 + TensorSpec,
23 +)
24 +from vllm.distributed.mooncake_store import MooncakeStoreConfig, setup_mooncake_store
25 +from vllm.logger import init_logger
26 +from vllm.utils.network_utils import get_ip
27 +
28 +logger = init_logger(__name__)
29 +
30 +# A failed transfer can still use registered memory after the binding returns.
31 +# Poisoned clients retain their owners until process exit, including if normal
32 +# worker teardown drops the connector. New operations are rejected.
33 +_UNSAFE_CLIENTS: list[Any] = []
34 +
35 +
36 +# Values from Mooncake v0.3.12 / v0.3.12.post1 types.h (ErrorCode).
37 +# Python exposes integer returns, not this enum.
38 +class _MooncakeErrorCode(IntEnum):
39 + NO_AVAILABLE_HANDLE = -200
40 + INVALID_PARAMS = -600
41 + INVALID_REPLICA = -702
42 + REPLICA_IS_NOT_READY = -703
43 + OBJECT_NOT_FOUND = -704
44 + OBJECT_ALREADY_EXISTS = -705
45 + LEASE_EXPIRED = -707
46 + RPC_FAIL = -900
47 + RPC_TIMEOUT = -901
48 + TENANT_QUOTA_EXCEEDED = -1700
49 +
50 +
51 +# For RAM batch reads and single-replica writes, these rejections occur
52 +# before submission or after completed I/O. Transfer timeouts can leave I/O
53 +# in flight; all other failures retain registered buffers and poison the client.
54 +_SAFE_IO_REJECTIONS = frozenset(
55 + {
56 + _MooncakeErrorCode.NO_AVAILABLE_HANDLE,
57 + _MooncakeErrorCode.REPLICA_IS_NOT_READY,
58 + _MooncakeErrorCode.OBJECT_NOT_FOUND,
59 + _MooncakeErrorCode.OBJECT_ALREADY_EXISTS,
60 + _MooncakeErrorCode.RPC_FAIL,
61 + _MooncakeErrorCode.RPC_TIMEOUT,
62 + }
63 +)
64 +# Batch GET validates capacity/replicas before I/O and checks leases afterwards.
65 +# PUT rejects tenant quotas before allocating replicas or submitting transfers.
66 +_SAFE_GET_REJECTIONS = _SAFE_IO_REJECTIONS | {
67 + _MooncakeErrorCode.INVALID_PARAMS,
68 + _MooncakeErrorCode.INVALID_REPLICA,
69 + _MooncakeErrorCode.LEASE_EXPIRED,
70 +}
71 +_SAFE_PUT_REJECTIONS = _SAFE_IO_REJECTIONS | {_MooncakeErrorCode.TENANT_QUOTA_EXCEEDED}
72 +
73 +
74 +_MOONCAKE_TENSOR_OBJECT_MAGIC = 0x4D4F4F4E
75 +# Magic, dtype, token count, hidden dimension. Format versioning is in the key.
76 +_MOONCAKE_TENSOR_HEADER = struct.Struct("<IIQQ")
77 +_MOONCAKE_DTYPE_TO_TORCH_DTYPE = {
78 + 0: "torch.float32",
79 + 11: "torch.float16",
80 + 12: "torch.bfloat16",
81 +}
82 +_TORCH_DTYPE_TO_MOONCAKE_DTYPE = {
83 + value: key for key, value in _MOONCAKE_DTYPE_TO_TORCH_DTYPE.items()
84 +}
85 +
86 +
87 +def create_mooncake_embedding_store_client(
88 + read_buffer_bytes: int = 128 * 1024**2,
89 +) -> MooncakeEmbeddingStoreClient:
90 + try:
91 + from mooncake.store import ( # type: ignore
92 + MooncakeDistributedStore,
93 + ObjectDataType,
94 + ReplicateConfig,
95 + )
96 + except ImportError as e:
97 + raise ImportError(
98 + "Install mooncake with Store support to enable cross_encoder_cache."
99 + ) from e
100 +
101 + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils
102 +
103 + config = MooncakeStoreConfig.load_from_config()
104 + if config.enable_offload:
105 + raise ValueError(
106 + "cross_encoder_cache supports a RAM Store; disable enable_offload"
107 + )
108 + store = MooncakeDistributedStore()
109 + local_ip = get_ip()
110 + local_hostname = rdma_utils.get_requester_local_hostname(local_ip)
111 + setup_mooncake_store(store, config, local_hostname)
112 +
113 + logger.info(
114 + "Initialized embedding Mooncake store mode=%s global_segment_size=%d "
115 + "local_buffer_size=%d",
116 + config.mode,
117 + config.global_segment_size,
118 + config.local_buffer_size,
119 + )
120 + replicate_config = ReplicateConfig()
121 + replicate_config.data_type = ObjectDataType.TENSOR
122 + return MooncakeEmbeddingStoreClient(
123 + store, replicate_config=replicate_config, read_buffer_bytes=read_buffer_bytes
124 + )
125 +
126 +
127 +class EmbeddingStoreError(RuntimeError):
128 + """A Store failure that must propagate to the worker."""
129 +
130 +
131 +class EmbeddingStoreOperationError(EmbeddingStoreError):
132 + """A rejected or completed operation that permits normal Encoder fallback."""
133 +
134 +
135 +def _batch_get_completed_safely(results: Any, capacities: list[int]) -> bool:
136 + return (
137 + isinstance(results, list)
138 + and len(results) == len(capacities)
139 + and all(
140 + type(result) is int
141 + and (0 <= result <= capacity or result in _SAFE_GET_REJECTIONS)
142 + for result, capacity in zip(results, capacities, strict=True)
143 + )
144 + )
145 +
146 +
147 +class MooncakeEmbeddingStoreClient:
148 + """Wraps Mooncake object and buffer APIs used by embedding transfer."""
149 +
150 + def __init__(
151 + self,
152 + store: Any,
153 + replicate_config: Any | None = None,
154 + *,
155 + read_buffer_bytes: int = 128 * 1024**2,
156 + ):
157 + self._lifetime_lock = threading.Lock()
158 + self._unsafe_owners: list[Any] = []
159 + self._poisoned = False
160 + self.store = store
161 + self.replicate_config = replicate_config
162 + self._read_buffer_bytes = read_buffer_bytes
163 + self._read_buffer: torch.Tensor | None = None
164 + self._put_header: ctypes.Array[ctypes.c_char] | None = None
165 +
166 + def _check_healthy(self) -> None:
167 + if self._poisoned:
168 + raise EmbeddingStoreError(
169 + "Mooncake I/O completion is unconfirmed; worker must exit"
170 + )
171 +
172 + def _poison(self, owners: list[Any]) -> None:
173 + with self._lifetime_lock:
174 + self._unsafe_owners.extend(owners)
175 + if not self._poisoned:
176 + self._poisoned = True
177 + _UNSAFE_CLIENTS.append(self)
178 +
179 + @contextmanager
180 + def _registered_io(
181 + self,
182 + buffers: list[tuple[Any, int, int]],
183 + ) -> Iterator[Callable[[int], None]]:
184 + self._check_healthy()
185 + registered = []
186 + submitted = False
187 + confirmed = False
188 +
189 + def confirm(result: int) -> None:
190 + nonlocal confirmed
191 + if type(result) is not int or (
192 + result < 0 and result not in _SAFE_PUT_REJECTIONS
193 + ):
194 + raise EmbeddingStoreError(
195 + "Mooncake I/O completion is unconfirmed; retaining buffers"
196 + )
197 + confirmed = True
198 +
199 + try:
200 + for owner, addr, size in buffers:
201 + ret = self.store.register_buffer(addr, size)
202 + if ret != 0:
203 + raise EmbeddingStoreOperationError(
204 + f"Failed to register embedding buffer: {ret}"
205 + )
206 + registered.append(addr)
207 + submitted = True
208 + yield confirm
209 + except BaseException as error:
210 + if submitted and not confirmed:
211 + self._poison([owner for owner, _, _ in buffers])
212 + raise EmbeddingStoreError(
213 + "Mooncake I/O completion is unconfirmed; retaining buffers"
214 + ) from error
215 + raise
216 + finally:
217 + if not submitted or confirmed:
218 + try:
219 + for addr in reversed(registered):
220 + self._unregister_buffer(addr)
221 + except BaseException:
222 + self._poison([owner for owner, _, _ in buffers])
223 + raise
224 +
225 + def close(self) -> None:
226 + """Close the supported Mooncake binding after all I/O is drained."""
227 + self._check_healthy()
228 + if self._read_buffer is not None:
229 + try:
230 + self._unregister_buffer(self._read_buffer.data_ptr())
231 + except BaseException:
232 + self._poison([self._read_buffer])
233 + raise
234 + self._read_buffer = None
235 + if self._put_header is not None:
236 + try:
237 + self._unregister_buffer(ctypes.addressof(self._put_header))
238 + except BaseException:
239 + self._poison([self._put_header])
240 + raise
241 + self._put_header = None
242 + ret = self.store.close()
243 + if ret != 0:
244 + raise EmbeddingStoreError(
245 + f"failed to close embedding Mooncake Store: {ret}"
246 + )
247 +
248 + def exists(self, pool_key: str) -> bool:
249 + return self.batch_exists([pool_key])[0]
250 +
251 + def batch_exists(self, keys: list[str]) -> list[bool]:
252 + self._check_healthy()
253 + if not keys:
254 + return []
255 +
256 + states = self.store.batch_is_exist(keys)
257 + if len(states) != len(keys):
258 + raise RuntimeError(
259 + "Mooncake Store returned an unexpected number of lookup results"
260 + )
261 + return [state == 1 for state in states]
262 +
263 + def load_tensors(
264 + self, expected: Mapping[str, TensorSpec], device: torch.device | str
265 + ) -> dict[str, torch.Tensor]:
266 + """Synchronously load runner inputs into independently owned tensors."""
267 + self._check_healthy()
268 + device = torch.device(device)
269 + loaded: dict[str, torch.Tensor] = {}
270 + chunk: dict[str, TensorSpec] = {}
271 + chunk_bytes = 0
272 + for key, spec in expected.items():
273 + size = _MOONCAKE_TENSOR_HEADER.size + spec.nbytes
274 + if size > self._read_buffer_bytes:
275 + logger.warning(
276 + "Skipping Store GET for %s: object bytes=%d exceed read buffer=%d",
277 + key,
278 + size,
279 + self._read_buffer_bytes,
280 + )
281 + continue
282 + if chunk_bytes + size > self._read_buffer_bytes:
283 + loaded.update(self._load_chunk(chunk, device))
284 + chunk = {}
285 + chunk_bytes = 0
286 + chunk[key] = spec
287 + chunk_bytes += size
288 + if chunk:
289 + loaded.update(self._load_chunk(chunk, device))
290 + return loaded
291 +
292 + def _load_chunk(
293 + self, expected: dict[str, TensorSpec], device: torch.device
294 + ) -> dict[str, torch.Tensor]:
295 + self._check_healthy()
296 + if self._read_buffer is None:
297 + buffer = torch.empty(
298 + self._read_buffer_bytes,
299 + dtype=torch.uint8,
300 + device="cpu",
301 + pin_memory=device.type == "cuda",
302 + )
303 + ret = self.store.register_buffer(buffer.data_ptr(), buffer.nbytes)
304 + if ret != 0:
305 + raise EmbeddingStoreOperationError(
306 + f"Failed to register embedding read buffer: {ret}"
307 + )
308 + self._read_buffer = buffer
309 + buffer = self._read_buffer
310 + base = buffer.data_ptr()
311 + sizes = [
312 + _MOONCAKE_TENSOR_HEADER.size + spec.nbytes for spec in expected.values()
313 + ]
314 + offsets: list[int] = []
315 + end = 0
316 + for size in sizes:
317 + offsets.append(end)
318 + end += size
319 + try:
320 + results = self.store.batch_get_into(
321 + list(expected), [base + offset for offset in offsets], sizes
322 + )
323 + if not _batch_get_completed_safely(results, sizes):
324 + raise EmbeddingStoreError("Unsafe batch GET results")
325 + except BaseException as error:
326 + self._poison([buffer])
327 + raise EmbeddingStoreError(
328 + "Mooncake I/O completion is unconfirmed; retaining read buffer"
329 + ) from error
330 +
331 + valid: list[tuple[str, TensorSpec, int]] = []
332 + for (key, spec), offset, size, result in zip(
333 + expected.items(), offsets, sizes, results, strict=True
334 + ):
335 + if result != size:
336 + logger.warning(
337 + "Store GET skipped for %s: result=%d, expected bytes=%d",
338 + key,
339 + result,
340 + size,
341 + )
342 + continue
343 + try:
344 + _validate_mooncake_tensor_metadata(
345 + key,
346 + ctypes.string_at(base + offset, _MOONCAKE_TENSOR_HEADER.size),
347 + spec,
348 + )
349 + except EmbeddingStoreOperationError as error:
350 + logger.warning("Store GET skipped for %s: %s", key, error)
351 + continue
352 + valid.append((key, spec, offset + _MOONCAKE_TENSOR_HEADER.size))
353 + if not valid:
354 + return {}
355 + # Allocate every destination before submitting any asynchronous copy.
356 + targets = {
357 + key: torch.empty(
358 + spec.shape,
359 + dtype=getattr(torch, spec.dtype.removeprefix("torch.")),
360 + device=device,
361 + )
362 + for key, spec, _ in valid
363 + }
364 + ready = torch.Event() if device.type == "cuda" else None
365 + try:
366 + for key, spec, offset in valid:
367 + targets[key].view(torch.uint8).view(-1).copy_(
368 + buffer[offset : offset + spec.nbytes],
369 + non_blocking=ready is not None,
370 + )
371 + if ready is not None:
372 + ready.record(torch.accelerator.current_stream(device))
373 + ready.synchronize()
374 + except BaseException as error:
375 + if ready is not None:
376 + self._poison([buffer, *targets.values(), ready])
377 + raise EmbeddingStoreError(
378 + "Embedding copy completion is unconfirmed; retaining buffers"
379 + ) from error
380 + raise
381 + return targets
382 +
383 + def put_tensor(self, pool_key: str, tensor: torch.Tensor) -> None:
384 + self._check_healthy()
385 + if not tensor.is_contiguous():
386 + raise EmbeddingStoreOperationError("embedding tensor must be contiguous")
387 + metadata = _encode_mooncake_tensor_metadata(tensor)
388 + # The single publisher reuses this header only after PUT completion.
389 + if self._put_header is None:
390 + header = ctypes.create_string_buffer(len(metadata))
391 + ret = self.store.register_buffer(ctypes.addressof(header), len(metadata))
392 + if ret != 0:
393 + raise EmbeddingStoreOperationError(
394 + f"Failed to register embedding header: {ret}"
395 + )
396 + self._put_header = header
397 + header_ptr = ctypes.addressof(self._put_header)
398 + ctypes.memmove(header_ptr, metadata, len(metadata))
399 + data_size = tensor.numel() * tensor.element_size()
400 + buffers = [
401 + (tensor, tensor.data_ptr(), data_size),
402 + ]
403 + with self._registered_io(buffers) as confirm:
404 + [result] = self.store.batch_put_from_multi_buffers(
405 + [pool_key],
406 + [[header_ptr, tensor.data_ptr()]],
407 + [[len(metadata), data_size]],
408 + self.replicate_config,
409 + )
410 + confirm(result)
411 + if result < 0:
412 + raise EmbeddingStoreOperationError(
413 + f"failed to put embedding tensor for {pool_key}"
414 + )
415 +
416 + def _unregister_buffer(self, addr: int) -> None:
417 + try:
418 + ret = self.store.unregister_buffer(addr)
419 + except Exception as error:
420 + raise EmbeddingStoreError(
421 + f"could not confirm buffer unregistration addr={addr:#x}"
422 + ) from error
423 + if ret != 0:
424 + raise EmbeddingStoreError(
425 + f"failed to unregister embedding buffer addr={addr:#x}: {ret}"
426 + )
427 +
428 +
429 +def _validate_mooncake_tensor_metadata(
430 + pool_key: str,
431 + metadata: bytes,
432 + expected: TensorSpec,
433 +) -> None:
434 + magic, dtype, num_tokens, hidden_dim = _MOONCAKE_TENSOR_HEADER.unpack(metadata)
435 + if magic != _MOONCAKE_TENSOR_OBJECT_MAGIC:
436 + raise EmbeddingStoreOperationError(
437 + f"invalid Mooncake tensor metadata header for {pool_key}"
438 + )
439 + if _MOONCAKE_DTYPE_TO_TORCH_DTYPE.get(dtype) != expected.dtype:
440 + raise EmbeddingStoreOperationError(
441 + f"unexpected Mooncake tensor dtype for {pool_key}: {dtype}"
442 + )
443 + if (num_tokens, hidden_dim) != expected.shape:
444 + raise EmbeddingStoreOperationError(
445 + f"embedding tensor shape mismatch for {pool_key}: expected={expected.shape}"
446 + )
447 +
448 +
449 +def _encode_mooncake_tensor_metadata(tensor: torch.Tensor) -> bytes:
450 + dtype = str(tensor.dtype)
451 + if dtype not in _TORCH_DTYPE_TO_MOONCAKE_DTYPE:
452 + raise EmbeddingStoreOperationError(
453 + f"unsupported embedding tensor dtype: {dtype}"
454 + )
455 + if tensor.ndim != 2:
456 + raise EmbeddingStoreOperationError(
457 + f"embedding tensor must be 2D: shape={tuple(tensor.shape)}"
458 + )
459 + return _MOONCAKE_TENSOR_HEADER.pack(
460 + _MOONCAKE_TENSOR_OBJECT_MAGIC,
461 + _TORCH_DTYPE_TO_MOONCAKE_DTYPE[dtype],
462 + *tensor.shape,
463 + )
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+3 -134
@@ -11,7 +11,6 @@ and MooncakeDistributedStore integration.
11 """
12
13 import dataclasses
14 -import json
14 import math
15 import os
16 import queue
@@ -20,10 +19,8 @@ import threading
19 import time
20 from collections.abc import Callable, Sequence
21 from concurrent.futures import Future, ThreadPoolExecutor
23 -from dataclasses import dataclass
24 -from typing import Any, Literal, TypeVar
22 +from typing import Any, TypeVar
23
26 -import regex as re
24 import torch
25 import zmq
26
@@ -67,6 +64,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import
64 decode_lookup_response,
65 encode_lookup_response,
66 )
67 +from vllm.distributed.mooncake_store import MooncakeStoreConfig, setup_mooncake_store
68 from vllm.logger import init_logger
69 from vllm.utils.math_utils import cdiv
70 from vllm.utils.network_utils import get_ip, make_zmq_socket
@@ -94,10 +92,6 @@ from .metrics import MooncakeStoreConnectorStats
92
93 logger = init_logger(__name__)
94
97 -DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
98 -DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
99 -DEFAULT_TENANT_ID = "default"
100 -
95 MOONCAKE_NO_AVAILABLE_HANDLE = -200
96 _T = TypeVar("_T")
97
@@ -149,115 +143,6 @@ _DIRECT_IO_ALIGNMENT = 4096
143 _DIRECT_IO_PADDING_BYTES = 2 * _DIRECT_IO_ALIGNMENT
144
145
152 -MooncakeMode = Literal["embedded", "standalone-store"]
153 -
154 -
155 -@dataclass
156 -class MooncakeStoreConfig:
157 - """Configuration for MooncakeDistributedStore.
158 -
159 - ``mode`` selects the topology: ``embedded`` (each rank contributes
160 - ``global_segment_size`` in-process) or ``standalone-store`` (rank
161 - contributes 0; an external ``mooncake_client`` process owns the pool
162 - and the SSD tier).
163 - """
164 -
165 - metadata_server: str
166 - master_server_address: str
167 - protocol: str
168 - device_name: str
169 - mode: MooncakeMode = "embedded"
170 - global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE
171 - local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE
172 - enable_offload: bool = False
173 - tenant_id: str = DEFAULT_TENANT_ID
174 -
175 - def __post_init__(self) -> None:
176 - if self.mode not in ("embedded", "standalone-store"):
177 - raise ValueError(f"unknown Mooncake mode: {self.mode!r}")
178 - if self.local_buffer_size <= 0:
179 - raise ValueError("local_buffer_size must be > 0")
180 - if self.mode == "embedded" and self.global_segment_size == 0:
181 - raise ValueError("embedded mode requires global_segment_size > 0")
182 - if self.mode == "standalone-store" and self.global_segment_size != 0:
183 - raise ValueError("standalone-store mode requires global_segment_size == 0")
184 -
185 - @staticmethod
186 - def from_file(file_path: str) -> "MooncakeStoreConfig":
187 - with open(file_path) as file:
188 - config = json.load(file)
189 - return MooncakeStoreConfig(
190 - metadata_server=config.get("metadata_server", ""),
191 - master_server_address=config.get("master_server_address", ""),
192 - protocol=config.get("protocol", "rdma"),
193 - device_name=config.get("device_name", ""),
194 - mode=config.get("mode", "embedded"),
195 - global_segment_size=_parse_size(
196 - config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
197 - ),
198 - local_buffer_size=_parse_size(
199 - config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)
200 - ),
201 - enable_offload=bool(config.get("enable_offload", False)),
202 - tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)),
203 - )
204 -
205 - @staticmethod
206 - def load_from_config() -> "MooncakeStoreConfig":
207 - config_path = os.getenv("MOONCAKE_CONFIG_PATH")
208 - if not config_path:
209 - raise ValueError(
210 - "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
211 - )
212 - return MooncakeStoreConfig.from_file(config_path)
213 -
214 -
215 -def _normalize_tenant_id(value: Any) -> str:
216 - if value is None:
217 - return DEFAULT_TENANT_ID
218 - if not isinstance(value, str):
219 - raise TypeError(
220 - f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}"
221 - )
222 - tenant_id = value.strip()
223 - return tenant_id if tenant_id else DEFAULT_TENANT_ID
224 -
225 -
226 -def _parse_size(value: Any) -> int:
227 - """Parse storage size strings with units: GB, MB, KB, B."""
228 - if isinstance(value, int):
229 - return value
230 - if not isinstance(value, str):
231 - try:
232 - return int(value)
233 - except (TypeError, ValueError) as e:
234 - raise TypeError(f"Unsupported type for size: {type(value)}") from e
235 -
236 - cleaned = value.strip().lower()
237 - if not cleaned:
238 - raise ValueError("Size cannot be empty.")
239 -
240 - unit_multipliers = {
241 - "gb": 1024**3,
242 - "mb": 1024**2,
243 - "kb": 1024,
244 - "b": 1,
245 - }
246 - match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned)
247 - if not match:
248 - raise ValueError(f"Invalid format: '{value}'")
249 -
250 - number_str = match.group(1)
251 - unit = match.group(2) or "b"
252 - multiplier = unit_multipliers[unit]
253 -
254 - try:
255 - numeric_value = float(number_str)
256 - except ValueError as exc:
257 - raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc
258 - return int(numeric_value * multiplier)
259 -
260 -
146 def _align_up(value: int, alignment: int) -> int:
147 return ((value + alignment - 1) // alignment) * alignment
148
@@ -1595,23 +1480,7 @@ class MooncakeStoreWorker:
1480 self.store = MooncakeDistributedStore()
1481 local_ip = get_ip()
1482 local_hostname = rdma_utils.get_requester_local_hostname(local_ip)
1598 - setup_kwargs: dict[str, str] = {}
1599 - if store_config.tenant_id != DEFAULT_TENANT_ID:
1600 - setup_kwargs["tenant_id"] = store_config.tenant_id
1601 - ret = self.store.setup(
1602 - local_hostname,
1603 - store_config.metadata_server,
1604 - store_config.global_segment_size,
1605 - store_config.local_buffer_size,
1606 - store_config.protocol,
1607 - store_config.device_name,
1608 - store_config.master_server_address,
1609 - **setup_kwargs,
1610 - )
1611 - if ret != 0:
1612 - msg = "Initialize MooncakeDistributedStore failed."
1613 - logger.error(msg)
1614 - raise RuntimeError(msg)
1483 + setup_mooncake_store(self.store, store_config, local_hostname)
1484
1485 preferred_segment = rdma_utils.get_configured_preferred_segment(extra_config)
1486 self.preferred_segment = preferred_segment
vllm/distributed/mooncake_store.py new
+144
@@ -0,0 +1,144 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Shared configuration and setup for EC and KV Mooncake Store clients."""
4 +
5 +import json
6 +import os
7 +from dataclasses import dataclass
8 +from typing import Any, Literal
9 +
10 +import regex as re
11 +
12 +DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
13 +DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
14 +DEFAULT_TENANT_ID = "default"
15 +
16 +
17 +MooncakeMode = Literal["embedded", "standalone-store"]
18 +
19 +
20 +@dataclass
21 +class MooncakeStoreConfig:
22 + """Configuration for MooncakeDistributedStore.
23 +
24 + ``mode`` selects the topology: ``embedded`` (each rank contributes
25 + ``global_segment_size`` in-process) or ``standalone-store`` (rank
26 + contributes 0; an external ``mooncake_client`` process owns the pool
27 + and the SSD tier).
28 + """
29 +
30 + metadata_server: str
31 + master_server_address: str
32 + protocol: str
33 + device_name: str
34 + mode: MooncakeMode = "embedded"
35 + global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE
36 + local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE
37 + enable_offload: bool = False
38 + tenant_id: str = DEFAULT_TENANT_ID
39 +
40 + def __post_init__(self) -> None:
41 + if self.mode not in ("embedded", "standalone-store"):
42 + raise ValueError(f"unknown Mooncake mode: {self.mode!r}")
43 + if self.local_buffer_size <= 0:
44 + raise ValueError("local_buffer_size must be > 0")
45 + if self.mode == "embedded" and self.global_segment_size == 0:
46 + raise ValueError("embedded mode requires global_segment_size > 0")
47 + if self.mode == "standalone-store" and self.global_segment_size != 0:
48 + raise ValueError("standalone-store mode requires global_segment_size == 0")
49 +
50 + @staticmethod
51 + def from_file(file_path: str) -> "MooncakeStoreConfig":
52 + with open(file_path) as file:
53 + config = json.load(file)
54 + return MooncakeStoreConfig(
55 + metadata_server=config.get("metadata_server", ""),
56 + master_server_address=config.get("master_server_address", ""),
57 + protocol=config.get("protocol", "rdma"),
58 + device_name=config.get("device_name", ""),
59 + mode=config.get("mode", "embedded"),
60 + global_segment_size=_parse_size(
61 + config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
62 + ),
63 + local_buffer_size=_parse_size(
64 + config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)
65 + ),
66 + enable_offload=bool(config.get("enable_offload", False)),
67 + tenant_id=_normalize_tenant_id(config.get("tenant_id", DEFAULT_TENANT_ID)),
68 + )
69 +
70 + @staticmethod
71 + def load_from_config() -> "MooncakeStoreConfig":
72 + config_path = os.getenv("MOONCAKE_CONFIG_PATH")
73 + if not config_path:
74 + raise ValueError(
75 + "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
76 + )
77 + return MooncakeStoreConfig.from_file(config_path)
78 +
79 +
80 +def _normalize_tenant_id(value: Any) -> str:
81 + if value is None:
82 + return DEFAULT_TENANT_ID
83 + if not isinstance(value, str):
84 + raise TypeError(
85 + f"tenant_id must be a string or null, got {type(value).__name__}: {value!r}"
86 + )
87 + tenant_id = value.strip()
88 + return tenant_id if tenant_id else DEFAULT_TENANT_ID
89 +
90 +
91 +def _parse_size(value: Any) -> int:
92 + """Parse storage size strings with units: GB, MB, KB, B."""
93 + if isinstance(value, int):
94 + return value
95 + if not isinstance(value, str):
96 + try:
97 + return int(value)
98 + except (TypeError, ValueError) as e:
99 + raise TypeError(f"Unsupported type for size: {type(value)}") from e
100 +
101 + cleaned = value.strip().lower()
102 + if not cleaned:
103 + raise ValueError("Size cannot be empty.")
104 +
105 + unit_multipliers = {
106 + "gb": 1024**3,
107 + "mb": 1024**2,
108 + "kb": 1024,
109 + "b": 1,
110 + }
111 + match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned)
112 + if not match:
113 + raise ValueError(f"Invalid format: '{value}'")
114 +
115 + number_str = match.group(1)
116 + unit = match.group(2) or "b"
117 + multiplier = unit_multipliers[unit]
118 +
119 + try:
120 + numeric_value = float(number_str)
121 + except ValueError as exc:
122 + raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc
123 + return int(numeric_value * multiplier)
124 +
125 +
126 +def setup_mooncake_store(
127 + store: Any, config: MooncakeStoreConfig, local_hostname: str
128 +) -> None:
129 + """Initialize a native Store, preserving explicit tenant configuration."""
130 + setup_kwargs: dict[str, str] = {}
131 + if config.tenant_id != DEFAULT_TENANT_ID:
132 + setup_kwargs["tenant_id"] = config.tenant_id
133 + ret = store.setup(
134 + local_hostname,
135 + config.metadata_server,
136 + config.global_segment_size,
137 + config.local_buffer_size,
138 + config.protocol,
139 + config.device_name,
140 + config.master_server_address,
141 + **setup_kwargs,
142 + )
143 + if ret != 0:
144 + raise RuntimeError("Initialize MooncakeDistributedStore failed.")