@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 36e74c9ea4
Fork 0

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

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

Synced Nov 3, 13:00 UTC

[KV Connector] Support disk offloading in MooncakeStoreConnector (#42689)

Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Zhewen Li committed May 16, 2026 at 13:34 UTC 36e74c9ea4feb5ade38ffa1ea96f24dd73316e02
5 files changed +1385 -55
docs/features/mooncake_store_connector_usage.md
+62 -3
@@ -2,7 +2,7 @@
2
3 MooncakeStoreConnector is a KV cache connector that uses [MooncakeDistributedStore](https://github.com/kvcache-ai/Mooncake) as a shared KV cache pool. Unlike `MooncakeConnector` which does direct point-to-point KV transfer between prefiller and decoder, MooncakeStoreConnector enables KV cache offloading to an external distributed store, supporting:
4
5 -- **CPU offloading**: Extend effective KV cache capacity by offloading to CPU memory via Mooncake's transfer engine.
5 +- **CPU/disk offloading**: Extend effective KV cache capacity by offloading to CPU memory or disk via Mooncake's transfer engine.
6 - **Prefix caching across instances**: Hash-based deduplication allows multiple vLLM instances to share cached KV blocks through the store.
7 - **Single-node and multi-node deployment**: Works both as a standalone KV cache extension and in disaggregated prefill-decode setups.
8
@@ -38,18 +38,29 @@ Create a JSON configuration file (e.g., `mooncake_config.json`):
38
39 ```json
40 {
41 + "mode": "embedded",
42 "metadata_server": "P2PHANDSHAKE",
43 "master_server_address": "127.0.0.1:50051",
44 "global_segment_size": "80GB",
45 "local_buffer_size": "4GB",
46 "protocol": "rdma",
46 - "device_name": ""
47 + "device_name": "",
48 + "enable_offload": false
49 }
50 ```
51
52 +- `mode`: Topology selection. `"embedded"` (default, PR-40900 baseline) has each
53 + vLLM rank contribute `global_segment_size` to the pool in-process.
54 + `"standalone-store"` makes ranks pure requesters — an external
55 + `mooncake_client` process owns the CPU pool and (optionally) the SSD tier.
56 - `protocol`: Use `"rdma"` for best performance. `"tcp"` works as a fallback.
51 -- `global_segment_size`: CPU memory contributed to the distributed pool (per GPU).
57 +- `global_segment_size`: CPU memory contributed to the distributed pool (per
58 + GPU). Must be `> 0` in `embedded` mode and `0` in `standalone-store` mode.
59 - `local_buffer_size`: Private buffer for this node's own operations (per GPU).
60 +- `enable_offload`: When `true`, vLLM allocates a DirectIO staging buffer so
61 + large prefills do not exceed the owner's SSD-write budget. Set this together
62 + with the matching `--enable_offload=true` flag on `mooncake_master` and on
63 + the external `mooncake_client` (if any).
64
65 Set the config path via environment variable:
66
@@ -126,12 +137,60 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \
137
138 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.
139
140 +### Disk Offloading
141 +
142 +Disk offloading is most commonly run in `standalone-store` mode: an external
143 +`mooncake_client` process owns the CPU pool and the SSD tier, and each vLLM
144 +rank is a pure requester. This avoids per-rank duplication of the SSD pool
145 +and keeps DirectIO budget tracking on a single process.
146 +
147 +Three things need to be aligned for end-to-end disk offloading:
148 +
149 +1. **`mooncake_master`** is started with `--enable_offload=true`.
150 +2. **`mooncake_client`** (the owner) is started with `--enable_offload=true`
151 + plus an SSD path via `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`.
152 +3. **vLLM-side** sets `"enable_offload": true` in the JSON config file (this is
153 + read by the connector and is **not** an environment variable).
154 +
155 +Example `mooncake_config.json` for the vLLM side:
156 +
157 +```json
158 +{
159 + "mode": "standalone-store",
160 + "metadata_server": "P2PHANDSHAKE",
161 + "master_server_address": "127.0.0.1:50051",
162 + "global_segment_size": 0,
163 + "local_buffer_size": "4GB",
164 + "protocol": "rdma",
165 + "device_name": "mlx5_0",
166 + "enable_offload": true
167 +}
168 +```
169 +
170 +Steer this rank to the local owner segment with:
171 +
172 +```bash
173 +export MOONCAKE_PREFERRED_SEGMENT=127.0.0.1:50053
174 +```
175 +
176 +The owner's SSD directory, on-disk eviction policy, and the DirectIO staging
177 +buffer size are controlled on the `mooncake_client` side via the standard
178 +Mooncake environment variables (`MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`,
179 +`MOONCAKE_BUCKET_EVICTION_POLICY`, `MOONCAKE_USE_URING`,
180 +`MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES`,
181 +`MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES`, etc.). Those are independent of
182 +the vLLM JSON config.
183 +
184 ## Environment Variables
185
186 | Variable | Description | Default |
187 | --- | --- | --- |
188 | `MOONCAKE_CONFIG_PATH` | Path to Mooncake JSON config file | (required) |
189 | `VLLM_MOONCAKE_BOOTSTRAP_PORT` | Bootstrap port for MooncakeConnector P2P transfer (disagg mode only) | 8998 |
190 +| `MOONCAKE_PREFERRED_SEGMENT` | Pin this rank's replicas to a specific owner segment (`host:port`); used in `standalone-store` mode | — |
191 +| `MOONCAKE_REQUESTER_LOCAL_HOSTNAME` | Override the hostname the vLLM rank registers with Mooncake as a requester. Defaults to the rank's resolved IP. | — |
192 +| `VLLM_MOONCAKE_STORE_TIER_LOG` | When `1`, logs a per-batch tier summary (memory vs disk hits) for observability | disabled |
193 +| `VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO` | Fraction of the owner's DirectIO staging buffer that the requester will fill in a single `batch_get_into_multi_buffers` call. Lower → more conservative pre-split, more round trips. | 0.9 |
194
195 ## KV Transfer Config
196
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+835
@@ -1,23 +1,36 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import json
5 +import logging
6 +import math
7 +import sys
8 import threading
9 +import types
10 +from types import SimpleNamespace
11 from unittest.mock import MagicMock, patch
12
13 +import pytest
14 import torch
15
16 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
17 + rdma_utils,
18 +)
19 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
20 worker,
21 )
22 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
23 ChunkedTokenDatabase,
24 KeyMetadata,
25 + LoadSpec,
26 ReqMeta,
27 )
28
29
30 def _make_store_sending_thread(
31 store: MagicMock,
32 + *,
33 + replicate_config: object | None = None,
34 ) -> worker.KVCacheStoreSendingThread:
35 token_database = ChunkedTokenDatabase(
36 KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
@@ -32,11 +45,55 @@ def _make_store_sending_thread(
45 put_step=1,
46 kv_role="kv_producer",
47 ready_event=threading.Event(),
48 + replicate_config=replicate_config,
49 + )
50 + thread.request_queue.task_done = MagicMock()
51 + return thread
52 +
53 +
54 +def _make_store_recving_thread(
55 + store: MagicMock,
56 + *,
57 + disk_offload_buffer_budget_bytes: int | None = None,
58 +) -> worker.KVCacheStoreRecvingThread:
59 + token_database = ChunkedTokenDatabase(
60 + KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
61 + )
62 + token_database.set_kv_caches_base_addr([0x1000])
63 + token_database.set_block_len([256])
64 + thread = worker.KVCacheStoreRecvingThread(
65 + store=store,
66 + token_database=token_database,
67 + block_size=16,
68 + tp_rank=0,
69 + ready_event=threading.Event(),
70 + disk_offload_buffer_budget_bytes=disk_offload_buffer_budget_bytes,
71 )
72 thread.request_queue.task_done = MagicMock()
73 return thread
74
75
76 +def _make_load_req(
77 + req_id: str,
78 + block_hashes: list[bytes],
79 + *,
80 + token_len: int,
81 + vllm_cached_tokens: int = 0,
82 +) -> ReqMeta:
83 + return ReqMeta(
84 + req_id=req_id,
85 + token_len_chunk=token_len,
86 + block_ids=list(range(len(block_hashes))),
87 + block_hashes=block_hashes,
88 + load_spec=LoadSpec(
89 + vllm_cached_tokens=vllm_cached_tokens,
90 + kvpool_cached_tokens=token_len,
91 + can_load=True,
92 + token_len=token_len,
93 + ),
94 + )
95 +
96 +
97 def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
98 return ReqMeta(
99 req_id=req_id,
@@ -48,6 +105,233 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
105 )
106
107
108 +_DISK_OFFLOAD_SINGLE_KEY_BYTES = worker._estimate_disk_offload_staging_bytes([256])
109 +_DISK_OFFLOAD_USABLE_BUDGET_RATIO = 0.9
110 +_DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS = 4 * _DISK_OFFLOAD_SINGLE_KEY_BYTES
111 +_DISK_OFFLOAD_BUDGET_FOR_SPLIT = math.ceil(
112 + 2 * _DISK_OFFLOAD_SINGLE_KEY_BYTES / _DISK_OFFLOAD_USABLE_BUDGET_RATIO
113 +) # Allows two 256-byte chunks but not the third.
114 +_DISK_OFFLOAD_BUDGET_TOO_SMALL = (
115 + _DISK_OFFLOAD_SINGLE_KEY_BYTES - 1
116 +) # Smaller than a single 256-byte chunk.
117 +
118 +
119 +class _FakeKVTransferConfig:
120 + def __init__(
121 + self,
122 + *,
123 + kv_role: str = "kv_both",
124 + extra_config: dict[str, object] | None = None,
125 + ) -> None:
126 + self.kv_role = kv_role
127 + self.kv_connector_extra_config = extra_config or {}
128 +
129 + def get_from_extra_config(self, key: str, default: object) -> object:
130 + return self.kv_connector_extra_config.get(key, default)
131 +
132 +
133 +class _FakeModelConfig:
134 + model = "test-model"
135 + use_mla = False
136 +
137 + def get_num_layers(self, parallel_config) -> int:
138 + return 1
139 +
140 + def get_total_num_kv_heads(self) -> int:
141 + return 1
142 +
143 +
144 +def _make_vllm_config(
145 + *, extra_config: dict[str, object] | None = None
146 +) -> SimpleNamespace:
147 + return SimpleNamespace(
148 + model_config=_FakeModelConfig(),
149 + parallel_config=SimpleNamespace(
150 + pipeline_parallel_size=1,
151 + rank=0,
152 + ),
153 + kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config),
154 + cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=10),
155 + kv_events_config=SimpleNamespace(enable_kv_cache_events=False),
156 + )
157 +
158 +
159 +def _write_mooncake_config(tmp_path, config: dict[str, object]) -> str:
160 + config_path = tmp_path / "mooncake_config.json"
161 + config_path.write_text(json.dumps(config), encoding="utf-8")
162 + return str(config_path)
163 +
164 +
165 +def _install_fake_mooncake(monkeypatch, store_instance: MagicMock):
166 + class FakeReplicateConfig:
167 + def __init__(self) -> None:
168 + self.preferred_segment = ""
169 +
170 + fake_store_module = types.ModuleType("mooncake.store")
171 + fake_store_module.MooncakeDistributedStore = lambda: store_instance # type: ignore[attr-defined]
172 + fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined]
173 + fake_mooncake_module = types.ModuleType("mooncake")
174 + fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined]
175 + monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module)
176 + monkeypatch.setitem(sys.modules, "mooncake.store", fake_store_module)
177 + return FakeReplicateConfig
178 +
179 +
180 +def _patch_worker_runtime(monkeypatch, *, local_ip: str = "10.0.0.7") -> None:
181 + single_rank_group = SimpleNamespace(world_size=1, rank_in_group=0)
182 + monkeypatch.setattr(worker, "get_mooncake_dp_engine_index", lambda _: 0)
183 + monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: 0)
184 + monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: 1)
185 + monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
186 + monkeypatch.setattr(worker, "get_dcp_group", lambda: single_rank_group)
187 + monkeypatch.setattr(worker, "get_ip", lambda: local_ip)
188 +
189 +
190 +def test_default_local_buffer_size_matches_pr40900():
191 + """PR-40900 shipped a 4 GiB default for local_buffer_size; the dual-mode
192 + patch preserves it (and the JSON key) so unchanged PR-40900 configs work."""
193 + assert worker.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3
194 +
195 +
196 +def test_get_requester_local_hostname_prefers_override(monkeypatch):
197 + monkeypatch.setenv("MOONCAKE_REQUESTER_LOCAL_HOSTNAME", "worker-a:50053")
198 +
199 + assert rdma_utils.get_requester_local_hostname("10.0.0.7") == "worker-a:50053"
200 +
201 +
202 +def test_get_configured_preferred_segment_returns_explicit_override():
203 + assert (
204 + rdma_utils.get_configured_preferred_segment(
205 + {"preferred_segment": "10.0.0.7:50053"}
206 + )
207 + == "10.0.0.7:50053"
208 + )
209 +
210 +
211 +def test_get_configured_preferred_segment_prefers_explicit_over_env(monkeypatch):
212 + monkeypatch.setenv("MOONCAKE_PREFERRED_SEGMENT", "10.0.0.8:50053")
213 +
214 + assert (
215 + rdma_utils.get_configured_preferred_segment(
216 + {"preferred_segment": "10.0.0.7:50053"}
217 + )
218 + == "10.0.0.7:50053"
219 + )
220 +
221 +
222 +def test_get_configured_preferred_segment_returns_env_override(monkeypatch):
223 + monkeypatch.setenv("MOONCAKE_PREFERRED_SEGMENT", "10.0.0.8:50053")
224 +
225 + assert rdma_utils.get_configured_preferred_segment({}) == "10.0.0.8:50053"
226 +
227 +
228 +def test_get_configured_preferred_segment_rejects_empty_override():
229 + with pytest.raises(ValueError, match="preferred_segment"):
230 + rdma_utils.get_configured_preferred_segment({"preferred_segment": " "})
231 +
232 +
233 +def test_get_configured_worker_rnic_prefers_explicit_device_name(monkeypatch):
234 + store_config = worker.MooncakeStoreConfig(
235 + metadata_server="",
236 + local_buffer_size=1,
237 + protocol="rdma",
238 + device_name="rocep139s0",
239 + master_server_address="",
240 + )
241 +
242 + assert (
243 + rdma_utils.get_configured_worker_rnic(
244 + protocol=store_config.protocol,
245 + configured_device=store_config.device_name,
246 + )
247 + == "rocep139s0"
248 + )
249 +
250 +
251 +def test_get_configured_worker_rnic_selects_device_from_explicit_csv(monkeypatch):
252 + monkeypatch.setattr(
253 + rdma_utils,
254 + "get_current_physical_gpu_index",
255 + lambda: 1,
256 + )
257 + store_config = worker.MooncakeStoreConfig(
258 + metadata_server="",
259 + local_buffer_size=1,
260 + protocol="rdma",
261 + device_name="rocep139s0,rocep140s0",
262 + master_server_address="",
263 + )
264 +
265 + assert (
266 + rdma_utils.get_configured_worker_rnic(
267 + protocol=store_config.protocol,
268 + configured_device=store_config.device_name,
269 + )
270 + == "rocep140s0"
271 + )
272 +
273 +
274 +def test_get_configured_worker_rnic_warns_and_returns_empty_for_rdma_with_no_device(
275 + caplog, monkeypatch
276 +):
277 + """No device configured + protocol=rdma → emit a clear warning and return ""
278 + so the C++ side handles auto-selection. There is no Python-side fallback."""
279 + monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
280 + with caplog.at_level(logging.WARNING):
281 + result = rdma_utils.get_configured_worker_rnic(
282 + protocol="rdma",
283 + configured_device="",
284 + )
285 + assert result == ""
286 + warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
287 + assert any("No RDMA devices specified" in r.message for r in warnings), (
288 + f"expected fallback warning, got {[r.message for r in warnings]}"
289 + )
290 +
291 +
292 +def test_get_configured_worker_rnic_silent_for_tcp_with_no_device(caplog, monkeypatch):
293 + """protocol=tcp + no device → return "" silently (no RDMA, no warning)."""
294 + monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
295 + with caplog.at_level(logging.WARNING):
296 + result = rdma_utils.get_configured_worker_rnic(
297 + protocol="tcp",
298 + configured_device="",
299 + )
300 + assert result == ""
301 + warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
302 + assert not any("RDMA" in r.message for r in warnings), (
303 + "did not expect RDMA warning for tcp protocol, got "
304 + f"{[r.message for r in warnings]}"
305 + )
306 +
307 +
308 +def test_get_configured_worker_rnic_rejects_short_explicit_csv(monkeypatch):
309 + monkeypatch.setattr(
310 + rdma_utils,
311 + "get_current_physical_gpu_index",
312 + lambda: 2,
313 + )
314 + with pytest.raises(ValueError, match="does not cover local GPU 2"):
315 + rdma_utils.get_configured_worker_rnic(
316 + protocol="rdma",
317 + configured_device="rocep139s0,rocep140s0",
318 + )
319 +
320 +
321 +class _ReplicaDesc:
322 + def __init__(self, tier: str):
323 + self.tier = tier
324 +
325 + def is_memory_replica(self) -> bool:
326 + return self.tier == "memory"
327 +
328 + def is_disk_replica(self) -> bool:
329 + return self.tier == "disk"
330 +
331 + def is_local_disk_replica(self) -> bool:
332 + return self.tier == "disk"
333 +
334 +
335 def test_store_sending_thread_skips_request_during_cpu_pressure():
336 store = MagicMock()
337 store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
@@ -105,6 +389,372 @@ def test_store_sending_thread_only_skips_on_no_available_handle():
389 assert store.batch_put_from_multi_buffers.call_count == 2
390
391
392 +def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set():
393 + store = MagicMock()
394 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
395 + store.batch_put_from_multi_buffers.return_value = [256, 256]
396 + replicate_config = SimpleNamespace(preferred_segment="10.0.0.7:50053")
397 + thread = _make_store_sending_thread(store, replicate_config=replicate_config)
398 +
399 + thread.add_stored_request("req-a")
400 + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
401 +
402 + assert store.batch_put_from_multi_buffers.call_count == 1
403 + call_args = store.batch_put_from_multi_buffers.call_args.args
404 + assert len(call_args) == 4
405 + assert call_args[3] is replicate_config
406 +
407 +
408 +def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_segment(): # noqa: E501
409 + """Without a preferred_segment the SendingThread still forwards a
410 + (default-constructed) ReplicateConfig so the C++ side always sees a
411 + well-defined config object."""
412 + store = MagicMock()
413 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
414 + store.batch_put_from_multi_buffers.return_value = [256, 256]
415 + replicate_config = SimpleNamespace()
416 + thread = _make_store_sending_thread(store, replicate_config=replicate_config)
417 +
418 + thread.add_stored_request("req-a")
419 + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
420 +
421 + assert store.batch_put_from_multi_buffers.call_count == 1
422 + call_args = store.batch_put_from_multi_buffers.call_args.args
423 + assert len(call_args) == 4
424 + assert call_args[3] is replicate_config
425 +
426 +
427 +def test_estimate_disk_offload_staging_bytes_sums_multi_segment_sizes():
428 + assert worker._estimate_disk_offload_staging_bytes([256, 512]) == 12288
429 +
430 +
431 +def test_recv_thread_uses_single_batch_when_no_disk_offload_budget(monkeypatch):
432 + monkeypatch.delenv("VLLM_MOONCAKE_STORE_TIER_LOG", raising=False)
433 + store = MagicMock()
434 + store.batch_get_into_multi_buffers.return_value = [256, 256, 256]
435 + thread = _make_store_recving_thread(store, disk_offload_buffer_budget_bytes=None)
436 +
437 + req = _make_load_req(
438 + "req-a",
439 + [b"a0", b"a1", b"a2"],
440 + token_len=48,
441 + )
442 +
443 + thread._handle_request(req)
444 +
445 + assert store.batch_get_into_multi_buffers.call_count == 1
446 + keys, addrs, sizes = store.batch_get_into_multi_buffers.call_args.args
447 + assert keys == [
448 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
449 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
450 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
451 + ]
452 + assert sizes == [[256], [256], [256]]
453 + store.batch_get_replica_desc.assert_not_called()
454 +
455 +
456 +def test_recv_thread_logs_tier_summary_when_enabled(monkeypatch, caplog_vllm):
457 + monkeypatch.setenv("VLLM_MOONCAKE_STORE_TIER_LOG", "1")
458 + caplog_vllm.set_level(logging.INFO, logger=worker.logger.name)
459 +
460 + store = MagicMock()
461 + store.batch_get_into_multi_buffers.return_value = [256, 256, -10]
462 + thread = _make_store_recving_thread(store, disk_offload_buffer_budget_bytes=None)
463 +
464 + req = _make_load_req(
465 + "req-a",
466 + [b"a0", b"a1", b"a2"],
467 + token_len=48,
468 + )
469 + expected_keys = [
470 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
471 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
472 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
473 + ]
474 + store.batch_get_replica_desc.return_value = {
475 + expected_keys[0]: [_ReplicaDesc("memory")],
476 + expected_keys[1]: [_ReplicaDesc("disk")],
477 + expected_keys[2]: [],
478 + }
479 +
480 + thread._handle_request(req)
481 +
482 + assert store.batch_get_replica_desc.call_args.args == (expected_keys,)
483 + assert store.method_calls[0][0] == "batch_get_replica_desc"
484 + assert store.method_calls[1][0] == "batch_get_into_multi_buffers"
485 +
486 + messages = [record.getMessage() for record in caplog_vllm.records]
487 + assert any(
488 + "Mooncake load tier summary" in message
489 + and "req_id=req-a" in message
490 + and "batch_keys=3" in message
491 + and "memory_keys=1" in message
492 + and "disk_keys=1" in message
493 + and "unknown_keys=1" in message
494 + and "success_keys=2" in message
495 + and "failed_keys=1" in message
496 + and "bytes_by_tier={'memory': 256, 'disk': 256, 'unknown': 0}" in message
497 + for message in messages
498 + )
499 +
500 +
501 +def test_recv_thread_uses_ratio_scaled_budget_for_first_pass_split():
502 + store = MagicMock()
503 + store.batch_get_into_multi_buffers.side_effect = [
504 + [256],
505 + [256],
506 + ]
507 + thread = _make_store_recving_thread(
508 + store,
509 + disk_offload_buffer_budget_bytes=2 * _DISK_OFFLOAD_SINGLE_KEY_BYTES,
510 + )
511 +
512 + req = _make_load_req(
513 + "req-a",
514 + [b"a0", b"a1"],
515 + token_len=32,
516 + )
517 +
518 + thread._handle_request(req)
519 +
520 + assert store.batch_get_into_multi_buffers.call_count == 2
521 + first_keys = store.batch_get_into_multi_buffers.call_args_list[0].args[0]
522 + second_keys = store.batch_get_into_multi_buffers.call_args_list[1].args[0]
523 + assert first_keys == [
524 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
525 + ]
526 + assert second_keys == [
527 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
528 + ]
529 +
530 +
531 +def test_recv_thread_splits_disk_offload_loads_by_budget():
532 + store = MagicMock()
533 + store.batch_get_into_multi_buffers.side_effect = [
534 + [256, 256],
535 + [256],
536 + ]
537 + thread = _make_store_recving_thread(
538 + store,
539 + disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_FOR_SPLIT,
540 + )
541 +
542 + req = _make_load_req(
543 + "req-a",
544 + [b"a0", b"a1", b"a2"],
545 + token_len=48,
546 + )
547 +
548 + thread._handle_request(req)
549 +
550 + assert store.batch_get_into_multi_buffers.call_count == 2
551 +
552 + first_keys = store.batch_get_into_multi_buffers.call_args_list[0].args[0]
553 + second_keys = store.batch_get_into_multi_buffers.call_args_list[1].args[0]
554 + first_addrs = store.batch_get_into_multi_buffers.call_args_list[0].args[1]
555 + second_addrs = store.batch_get_into_multi_buffers.call_args_list[1].args[1]
556 + first_sizes = store.batch_get_into_multi_buffers.call_args_list[0].args[2]
557 + second_sizes = store.batch_get_into_multi_buffers.call_args_list[1].args[2]
558 + assert first_keys == [
559 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
560 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
561 + ]
562 + assert second_keys == [
563 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
564 + ]
565 + base_addr = thread.token_database.kv_caches_base_addr[0]
566 + block_len = thread.token_database.block_len[0]
567 + assert first_addrs == [[base_addr], [base_addr + block_len]]
568 + assert second_addrs == [[base_addr + 2 * block_len]]
569 + expected_size = block_len
570 + assert first_sizes == [[expected_size], [expected_size]]
571 + assert second_sizes == [[expected_size]]
572 +
573 +
574 +def test_recv_thread_stops_after_first_failing_disk_offload_sub_batch():
575 + store = MagicMock()
576 + store.batch_get_into_multi_buffers.return_value = [-10, -10]
577 + thread = _make_store_recving_thread(
578 + store,
579 + disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_FOR_SPLIT,
580 + )
581 +
582 + req = _make_load_req(
583 + "req-a",
584 + [b"a0", b"a1", b"a2"],
585 + token_len=48,
586 + )
587 +
588 + thread._handle_request(req)
589 +
590 + assert store.batch_get_into_multi_buffers.call_count == 1
591 +
592 +
593 +def test_recv_thread_skips_split_when_budget_holds_all_keys():
594 + """PR-36 removed the count-based split trigger; with budget for 3 keys,
595 + all three should be requested in a single call."""
596 + store = MagicMock()
597 + store.batch_get_into_multi_buffers.return_value = [256, 256, 256]
598 + thread = _make_store_recving_thread(
599 + store,
600 + disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS,
601 + )
602 +
603 + req = _make_load_req(
604 + "req-a",
605 + [b"a0", b"a1", b"a2"],
606 + token_len=48,
607 + )
608 +
609 + thread._handle_request(req)
610 +
611 + assert store.batch_get_into_multi_buffers.call_count == 1
612 + assert store.batch_get_into_multi_buffers.call_args_list[0].args[0] == [
613 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
614 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
615 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
616 + ]
617 +
618 +
619 +def test_recv_thread_reports_unsplittable_key_larger_than_budget():
620 + store = MagicMock()
621 + thread = _make_store_recving_thread(
622 + store,
623 + disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_TOO_SMALL,
624 + )
625 +
626 + req = _make_load_req(
627 + "req-a",
628 + [b"a0"],
629 + token_len=16,
630 + )
631 +
632 + thread._handle_request(req)
633 +
634 + assert store.batch_get_into_multi_buffers.call_count == 0
635 +
636 +
637 +def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):
638 + store = MagicMock()
639 + store.setup.return_value = 0
640 + _install_fake_mooncake(monkeypatch, store)
641 + _patch_worker_runtime(monkeypatch)
642 + monkeypatch.setenv(
643 + "MOONCAKE_CONFIG_PATH",
644 + _write_mooncake_config(
645 + tmp_path,
646 + {
647 + "metadata_server": "http://metadata/endpoint",
648 + "global_segment_size": "4gb",
649 + "local_buffer_size": "64mb",
650 + "protocol": "rdma",
651 + "device_name": "mlx5_0",
652 + "master_server_address": "10.0.0.7:50051",
653 + "enable_offload": True,
654 + },
655 + ),
656 + )
657 + w = worker.MooncakeStoreWorker(_make_vllm_config())
658 +
659 + assert not hasattr(w, "_isolate_offload_resources")
660 + assert store.setup.call_args.args == (
661 + "10.0.0.7",
662 + "http://metadata/endpoint",
663 + 4 * 1024 * 1024 * 1024, # global_segment_size: "4gb" honored
664 + 64 * 1024 * 1024,
665 + "rdma",
666 + "mlx5_0",
667 + "10.0.0.7:50051",
668 + )
669 +
670 +
671 +def test_requester_worker_init_prefers_local_hostname_override(
672 + tmp_path,
673 + monkeypatch,
674 +):
675 + store = MagicMock()
676 + store.setup.return_value = 0
677 + _install_fake_mooncake(monkeypatch, store)
678 + _patch_worker_runtime(monkeypatch)
679 + monkeypatch.setenv("MOONCAKE_REQUESTER_LOCAL_HOSTNAME", "worker-a:50053")
680 + monkeypatch.setenv(
681 + "MOONCAKE_CONFIG_PATH",
682 + _write_mooncake_config(
683 + tmp_path,
684 + {
685 + "metadata_server": "http://metadata/endpoint",
686 + "local_buffer_size": "64mb",
687 + "protocol": "tcp",
688 + "device_name": "",
689 + "master_server_address": "10.0.0.7:50051",
690 + },
691 + ),
692 + )
693 + worker.MooncakeStoreWorker(_make_vllm_config())
694 +
695 + assert store.setup.call_args.args[0] == "worker-a:50053"
696 +
697 +
698 +def test_requester_worker_init_skips_disk_budget_when_offload_disabled(
699 + tmp_path,
700 + monkeypatch,
701 +):
702 + """enable_offload=False zeroes out the disk budget so we don't generate
703 + redundant owner GET-RPCs."""
704 + store = MagicMock()
705 + store.setup.return_value = 0
706 + _install_fake_mooncake(monkeypatch, store)
707 + _patch_worker_runtime(monkeypatch)
708 + monkeypatch.setenv(
709 + "MOONCAKE_CONFIG_PATH",
710 + _write_mooncake_config(
711 + tmp_path,
712 + {
713 + "metadata_server": "http://metadata/endpoint",
714 + "protocol": "tcp",
715 + "device_name": "",
716 + "master_server_address": "10.0.0.7:50051",
717 + "enable_offload": False,
718 + },
719 + ),
720 + )
721 + w = worker.MooncakeStoreWorker(_make_vllm_config())
722 +
723 + assert w.disk_offload_buffer_budget_bytes is None
724 +
725 +
726 +def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
727 + tmp_path,
728 + monkeypatch,
729 +):
730 + store = MagicMock()
731 + store.setup.return_value = 0
732 + fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store)
733 + _patch_worker_runtime(monkeypatch)
734 + monkeypatch.setenv(
735 + "MOONCAKE_CONFIG_PATH",
736 + _write_mooncake_config(
737 + tmp_path,
738 + {
739 + "metadata_server": "http://metadata/endpoint",
740 + "protocol": "tcp",
741 + "device_name": "",
742 + "master_server_address": "10.0.0.7:50051",
743 + },
744 + ),
745 + )
746 + w = worker.MooncakeStoreWorker(
747 + _make_vllm_config(
748 + extra_config={
749 + "preferred_segment": "10.0.0.7:50053",
750 + }
751 + )
752 + )
753 +
754 + assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
755 + assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
756 +
757 +
758 # ---------------------------------------------------------------------------
759 # Helpers for register_kv_caches tests
760 # ---------------------------------------------------------------------------
@@ -147,8 +797,10 @@ def _make_bare_worker(
797 w.tp_rank = 0
798 w.put_step = 1
799 w.enable_kv_events = False
800 + w.disk_offload_buffer_budget_bytes = None
801 w.kv_send_thread = None
802 w.kv_recv_thread = None
803 + w.store_replicate_config = SimpleNamespace()
804 return w
805
806
@@ -298,3 +950,186 @@ def test_register_kv_caches_cross_layer_single_segment():
950
951 assert w2.kv_caches_base_addr == w.kv_caches_base_addr
952 assert w2.block_len == w.block_len
953 +
954 +
955 +# ---------------------------------------------------------------------------
956 +# Dual-mode (embedded / standalone-store) config validation tests
957 +# ---------------------------------------------------------------------------
958 +
959 +
960 +def _make_config(**overrides):
961 + """Build a MooncakeStoreConfig with sensible defaults for validation tests.
962 +
963 + Required dataclass fields are populated; callers override only the field
964 + under test.
965 + """
966 + base = dict(
967 + metadata_server="http://metadata/endpoint",
968 + master_server_address="10.0.0.7:50051",
969 + protocol="rdma",
970 + device_name="mlx5_0",
971 + )
972 + base.update(overrides)
973 + return worker.MooncakeStoreConfig(**base)
974 +
975 +
976 +def test_config_defaults_to_embedded():
977 + """A JSON without explicit mode parses as embedded with 4 GiB segment."""
978 + cfg = _make_config()
979 + assert cfg.mode == "embedded"
980 + assert cfg.global_segment_size == worker.DEFAULT_GLOBAL_SEGMENT_SIZE
981 + assert cfg.local_buffer_size == worker.DEFAULT_LOCAL_BUFFER_SIZE
982 + assert cfg.enable_offload is False
983 +
984 +
985 +def test_config_pr40900_unchanged(tmp_path):
986 + """A literal PR-40900 config (no mode, no enable_offload, no preferred_segment)
987 + parses without raising and resolves to embedded mode."""
988 + config_path = _write_mooncake_config(
989 + tmp_path,
990 + {
991 + "metadata_server": "http://metadata/endpoint",
992 + "global_segment_size": "4GB",
993 + "local_buffer_size": "4GB",
994 + "protocol": "rdma",
995 + "device_name": "mlx5_0",
996 + "master_server_address": "10.0.0.7:50051",
997 + },
998 + )
999 + cfg = worker.MooncakeStoreConfig.from_file(config_path)
1000 + assert cfg.mode == "embedded"
1001 + assert cfg.global_segment_size == 4 * 1024**3
1002 + assert cfg.local_buffer_size == 4 * 1024**3
1003 + assert cfg.enable_offload is False
1004 +
1005 +
1006 +def test_config_embedded_rejects_zero_segment():
1007 + with pytest.raises(
1008 + ValueError, match=r"embedded mode requires global_segment_size > 0"
1009 + ):
1010 + _make_config(mode="embedded", global_segment_size=0)
1011 +
1012 +
1013 +def test_config_standalone_store_rejects_nonzero_segment():
1014 + with pytest.raises(
1015 + ValueError,
1016 + match=r"standalone-store mode requires global_segment_size == 0",
1017 + ):
1018 + _make_config(mode="standalone-store", global_segment_size=4 * 1024**3)
1019 +
1020 +
1021 +def test_config_standalone_store_accepts_zero_segment():
1022 + cfg = _make_config(mode="standalone-store", global_segment_size=0)
1023 + assert cfg.mode == "standalone-store"
1024 + assert cfg.global_segment_size == 0
1025 +
1026 +
1027 +def test_config_unknown_mode():
1028 + with pytest.raises(ValueError, match=r"unknown Mooncake mode"):
1029 + _make_config(mode="something-else")
1030 +
1031 +
1032 +def test_config_zero_local_buffer():
1033 + with pytest.raises(ValueError, match=r"local_buffer_size must be > 0"):
1034 + _make_config(local_buffer_size=0)
1035 +
1036 +
1037 +# ---------------------------------------------------------------------------
1038 +# End-to-end topology tests
1039 +# Covers the two supported recipes:
1040 +# (A) standalone-store mode + disk offload (mode="standalone-store",
1041 +# segment=0, enable_offload=true, preferred_segment set)
1042 +# (B) embedded mode + CPU only (mode default, segment>0,
1043 +# enable_offload=false, no preferred_segment)
1044 +# ---------------------------------------------------------------------------
1045 +
1046 +
1047 +def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch):
1048 + """standalone-store + disk: global_segment_size=0, enable_offload=True,
1049 + preferred_segment set. Assert setup() positional args, ReplicateConfig
1050 + wiring, and that the disk-offload buffer budget is allocated."""
1051 + store = MagicMock()
1052 + store.setup.return_value = 0
1053 + fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store)
1054 + _patch_worker_runtime(monkeypatch)
1055 + monkeypatch.setenv(
1056 + "MOONCAKE_CONFIG_PATH",
1057 + _write_mooncake_config(
1058 + tmp_path,
1059 + {
1060 + "mode": "standalone-store",
1061 + "metadata_server": "http://metadata/endpoint",
1062 + "global_segment_size": 0,
1063 + "local_buffer_size": "1GB",
1064 + "protocol": "rdma",
1065 + "device_name": "mlx5_0",
1066 + "master_server_address": "10.0.0.7:50051",
1067 + "enable_offload": True,
1068 + },
1069 + ),
1070 + )
1071 +
1072 + w = worker.MooncakeStoreWorker(
1073 + _make_vllm_config(extra_config={"preferred_segment": "10.0.0.7:50053"})
1074 + )
1075 +
1076 + # setup() receives global_segment_size=0 and the configured local buffer.
1077 + assert store.setup.call_args.args == (
1078 + "10.0.0.7",
1079 + "http://metadata/endpoint",
1080 + 0,
1081 + 1024 * 1024 * 1024,
1082 + "rdma",
1083 + "mlx5_0",
1084 + "10.0.0.7:50051",
1085 + )
1086 + # ReplicateConfig is built and carries the preferred_segment.
1087 + assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
1088 + assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
1089 + # Disk-offload staging budget is allocated (enable_offload=True).
1090 + assert w.disk_offload_buffer_budget_bytes is not None
1091 + assert w.disk_offload_buffer_budget_bytes > 0
1092 +
1093 +
1094 +def test_topology_embedded_cpu_only(tmp_path, monkeypatch):
1095 + """embedded + CPU-only: no mode key (defaults to embedded),
1096 + global_segment_size>0, enable_offload absent, no preferred_segment.
1097 + This is the PR-40900 baseline recipe."""
1098 + store = MagicMock()
1099 + store.setup.return_value = 0
1100 + fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store)
1101 + _patch_worker_runtime(monkeypatch)
1102 + monkeypatch.setenv(
1103 + "MOONCAKE_CONFIG_PATH",
1104 + _write_mooncake_config(
1105 + tmp_path,
1106 + {
1107 + "metadata_server": "http://metadata/endpoint",
1108 + "global_segment_size": "4GB",
1109 + "local_buffer_size": "4GB",
1110 + "protocol": "rdma",
1111 + "device_name": "mlx5_0",
1112 + "master_server_address": "10.0.0.7:50051",
1113 + },
1114 + ),
1115 + )
1116 +
1117 + w = worker.MooncakeStoreWorker(_make_vllm_config())
1118 +
1119 + # setup() receives global_segment_size=4 GiB (rank contributes a segment).
1120 + assert store.setup.call_args.args == (
1121 + "10.0.0.7",
1122 + "http://metadata/endpoint",
1123 + 4 * 1024 * 1024 * 1024,
1124 + 4 * 1024 * 1024 * 1024,
1125 + "rdma",
1126 + "mlx5_0",
1127 + "10.0.0.7:50051",
1128 + )
1129 + # No preferred_segment — ReplicateConfig is default-constructed (so the
1130 + # preferred_segment field keeps its default value).
1131 + assert w.preferred_segment is None
1132 + assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
1133 + assert w.store_replicate_config.preferred_segment == ""
1134 + # No disk budget — enable_offload was absent (defaults to False).
1135 + assert w.disk_offload_buffer_budget_bytes is None
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/rdma_utils.py new
+115
@@ -0,0 +1,115 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Mooncake requester config helpers."""
4 +
5 +from collections.abc import Mapping
6 +from typing import Any
7 +
8 +import torch
9 +
10 +import vllm.envs as envs
11 +from vllm.logger import init_logger
12 +
13 +logger = init_logger(__name__)
14 +
15 +
16 +def normalize_string_override(value: Any) -> str | None:
17 + if not isinstance(value, str):
18 + return None
19 + normalized = value.strip()
20 + return normalized or None
21 +
22 +
23 +def get_current_physical_gpu_index() -> int | None:
24 + try:
25 + from vllm.platforms import current_platform
26 + except ImportError:
27 + return None
28 +
29 + try:
30 + device_index = torch.accelerator.current_device_index()
31 + physical_device_id = current_platform.device_id_to_physical_device_id(
32 + device_index
33 + )
34 + return int(physical_device_id)
35 + except Exception:
36 + return None
37 +
38 +
39 +def get_requester_local_hostname(local_ip: str) -> str:
40 + override = normalize_string_override(envs.MOONCAKE_REQUESTER_LOCAL_HOSTNAME)
41 + if override is not None:
42 + return override
43 + return local_ip
44 +
45 +
46 +def get_configured_preferred_segment(
47 + extra_config: Mapping[str, Any],
48 +) -> str | None:
49 + preferred_segment = normalize_string_override(extra_config.get("preferred_segment"))
50 + if preferred_segment is not None:
51 + return preferred_segment
52 + if extra_config.get("preferred_segment") is not None:
53 + raise ValueError(
54 + "Mooncake preferred_segment override must be a non-empty string"
55 + )
56 +
57 + env_value = normalize_string_override(envs.MOONCAKE_PREFERRED_SEGMENT)
58 + if env_value is not None:
59 + logger.info(
60 + "Mooncake preferred_segment from MOONCAKE_PREFERRED_SEGMENT: %s",
61 + env_value,
62 + )
63 + return env_value
64 + return None
65 +
66 +
67 +def _get_explicit_worker_rnic(device_list: str) -> str:
68 + entries = [entry.strip() for entry in device_list.split(",")]
69 + if any(not entry for entry in entries):
70 + raise ValueError(
71 + "Mooncake worker device_name contains an empty RDMA device entry"
72 + )
73 + if len(entries) == 1:
74 + return entries[0]
75 +
76 + gpu_index = get_current_physical_gpu_index()
77 + if gpu_index is None:
78 + raise RuntimeError(
79 + "Mooncake RDMA requester could not determine the local physical GPU index"
80 + )
81 + if gpu_index >= len(entries):
82 + raise ValueError(
83 + "Mooncake worker device list does not cover local GPU "
84 + f"{gpu_index}: {device_list}"
85 + )
86 + device_name = entries[gpu_index]
87 + logger.info(
88 + "Mooncake selected worker RNIC %s from explicit device list for local GPU %s",
89 + device_name,
90 + gpu_index,
91 + )
92 + return device_name
93 +
94 +
95 +def get_configured_worker_rnic(
96 + *,
97 + protocol: str,
98 + configured_device: str,
99 +) -> str:
100 + normalized_device = normalize_string_override(configured_device)
101 + if normalized_device is not None:
102 + return _get_explicit_worker_rnic(normalized_device)
103 +
104 + if protocol not in {"rdma", "efa"}:
105 + return ""
106 +
107 + logger.warning(
108 + "No RDMA devices specified for Mooncake backend (protocol=%s). "
109 + "Set 'device_name' in mooncake_config.json to a single RNIC name "
110 + "or a comma-separated CSV indexed by physical GPU; falling back to "
111 + "Mooncake's built-in auto-selection, which may converge on the same "
112 + "NIC across all DP ranks and saturate bandwidth.",
113 + protocol,
114 + )
115 + return ""
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+355 -52
@@ -13,10 +13,11 @@ and MooncakeDistributedStore integration.
13 import json
14 import os
15 import queue
16 +import socket
17 import threading
18 from collections import defaultdict
19 from dataclasses import dataclass
19 -from typing import Any
20 +from typing import Any, Literal
21
22 import regex as re
23 import torch
@@ -31,10 +32,11 @@ from vllm.distributed import (
32 get_tensor_model_parallel_world_size,
33 )
34 from vllm.distributed.kv_events import BlockStored
35 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils
36 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
37 get_mooncake_dp_engine_index,
38 )
37 -from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
39 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
40 ChunkedTokenDatabase,
41 KeyMetadata,
42 MooncakeStoreConnectorMetadata,
@@ -49,19 +51,48 @@ logger = init_logger(__name__)
51
52 DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
53 DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
54 +
55 MOONCAKE_NO_AVAILABLE_HANDLE = -200
56
57 +# Mirrors FileStorageConfig::local_buffer_size in Mooncake C++.
58 +DEFAULT_MOONCAKE_DISK_STAGING_BUFFER_BYTES = 1280 * 1024 * 1024
59 +
60 +# Mirrors DirectIO alignment in Mooncake's AllocateBatch.
61 +_DIRECT_IO_ALIGNMENT = 4096
62 +_DIRECT_IO_PADDING_BYTES = 2 * _DIRECT_IO_ALIGNMENT
63 +
64 +
65 +MooncakeMode = Literal["embedded", "standalone-store"]
66 +
67
68 @dataclass
69 class MooncakeStoreConfig:
57 - """Configuration for MooncakeDistributedStore."""
70 + """Configuration for MooncakeDistributedStore.
71 +
72 + ``mode`` selects the topology: ``embedded`` (each rank contributes
73 + ``global_segment_size`` in-process) or ``standalone-store`` (rank
74 + contributes 0; an external ``mooncake_client`` process owns the pool
75 + and the SSD tier).
76 + """
77
78 metadata_server: str
60 - global_segment_size: int
61 - local_buffer_size: int
79 + master_server_address: str
80 protocol: str
81 device_name: str
64 - master_server_address: str
82 + mode: MooncakeMode = "embedded"
83 + global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE
84 + local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE
85 + enable_offload: bool = False
86 +
87 + def __post_init__(self) -> None:
88 + if self.mode not in ("embedded", "standalone-store"):
89 + raise ValueError(f"unknown Mooncake mode: {self.mode!r}")
90 + if self.local_buffer_size <= 0:
91 + raise ValueError("local_buffer_size must be > 0")
92 + if self.mode == "embedded" and self.global_segment_size == 0:
93 + raise ValueError("embedded mode requires global_segment_size > 0")
94 + if self.mode == "standalone-store" and self.global_segment_size != 0:
95 + raise ValueError("standalone-store mode requires global_segment_size == 0")
96
97 @staticmethod
98 def from_file(file_path: str) -> "MooncakeStoreConfig":
@@ -69,15 +100,17 @@ class MooncakeStoreConfig:
100 config = json.load(file)
101 return MooncakeStoreConfig(
102 metadata_server=config.get("metadata_server", ""),
103 + master_server_address=config.get("master_server_address", ""),
104 + protocol=config.get("protocol", "rdma"),
105 + device_name=config.get("device_name", ""),
106 + mode=config.get("mode", "embedded"),
107 global_segment_size=_parse_size(
108 config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
109 ),
110 local_buffer_size=_parse_size(
111 config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)
112 ),
78 - protocol=config.get("protocol", "rdma"),
79 - device_name=config.get("device_name", ""),
80 - master_server_address=config.get("master_server_address", ""),
113 + enable_offload=bool(config.get("enable_offload", False)),
114 )
115
116 @staticmethod
@@ -125,6 +158,158 @@ def _parse_size(value: Any) -> int:
158 return int(numeric_value * multiplier)
159
160
161 +def _align_up(value: int, alignment: int) -> int:
162 + return ((value + alignment - 1) // alignment) * alignment
163 +
164 +
165 +def _estimate_disk_offload_staging_bytes(size_list: list[int]) -> int:
166 + data_size = sum(size_list)
167 + return _align_up(data_size, _DIRECT_IO_ALIGNMENT) + _DIRECT_IO_PADDING_BYTES
168 +
169 +
170 +def _get_usable_disk_offload_buffer_budget_bytes(raw_budget_bytes: int) -> int:
171 + return max(1, int(raw_budget_bytes * envs.VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO))
172 +
173 +
174 +def _split_disk_offload_load_batches(
175 + keys: list[str],
176 + addrs: list[list[int]],
177 + sizes: list[list[int]],
178 + usable_budget_bytes: int,
179 + raw_budget_bytes: int,
180 +) -> tuple[list[tuple[list[str], list[list[int]], list[list[int]]]], str | None]:
181 + """Split a GET into sub-batches that fit the owner's staging buffer.
182 +
183 + ``addrs[i]`` / ``sizes[i]`` are scatter-gather lists (K/V or multi-layer
184 + segments) for key ``i``. ``usable_budget_bytes`` caps a multi-key batch;
185 + ``raw_budget_bytes`` is the hard per-key cap.
186 +
187 + Returns ``(batches, oversize_key)``. Aborts with ``([], key)`` if any
188 + single key exceeds ``raw_budget_bytes``; otherwise ``oversize_key`` is
189 + ``None``.
190 + """
191 + batches: list[tuple[list[str], list[list[int]], list[list[int]]]] = []
192 + batch_keys: list[str] = []
193 + batch_addrs: list[list[int]] = []
194 + batch_sizes: list[list[int]] = []
195 + batch_bytes = 0
196 +
197 + for key, addr, size in zip(keys, addrs, sizes, strict=True):
198 + key_bytes = _estimate_disk_offload_staging_bytes(size)
199 + if key_bytes > raw_budget_bytes:
200 + return [], key
201 + if key_bytes > usable_budget_bytes:
202 + if batch_keys:
203 + batches.append((batch_keys, batch_addrs, batch_sizes))
204 + batch_keys, batch_addrs, batch_sizes = [], [], []
205 + batch_bytes = 0
206 + batches.append(([key], [addr], [size]))
207 + continue
208 + if batch_keys and batch_bytes + key_bytes > usable_budget_bytes:
209 + batches.append((batch_keys, batch_addrs, batch_sizes))
210 + batch_keys, batch_addrs, batch_sizes = [], [], []
211 + batch_bytes = 0
212 + batch_keys.append(key)
213 + batch_addrs.append(addr)
214 + batch_sizes.append(size)
215 + batch_bytes += key_bytes
216 +
217 + if batch_keys:
218 + batches.append((batch_keys, batch_addrs, batch_sizes))
219 + return batches, None
220 +
221 +
222 +def _call_replica_predicate(replica_desc: Any, method_name: str) -> bool:
223 + method = getattr(replica_desc, method_name, None)
224 + if method is None:
225 + return False
226 + try:
227 + return bool(method())
228 + except Exception:
229 + return False
230 +
231 +
232 +def _classify_replica_tier(replica_descs: Any) -> str:
233 + if not replica_descs:
234 + return "unknown"
235 + try:
236 + replica_desc = replica_descs[0]
237 + except (IndexError, KeyError, TypeError):
238 + return "unknown"
239 +
240 + if _call_replica_predicate(replica_desc, "is_memory_replica"):
241 + return "memory"
242 + if _call_replica_predicate(
243 + replica_desc, "is_disk_replica"
244 + ) or _call_replica_predicate(replica_desc, "is_local_disk_replica"):
245 + return "disk"
246 + return "unknown"
247 +
248 +
249 +def _get_replica_tiers_by_key(store: Any, keys: list[str]) -> dict[str, str]:
250 + tiers_by_key = {key: "unknown" for key in keys}
251 + try:
252 + replica_descs_by_key = store.batch_get_replica_desc(keys)
253 + except Exception as e:
254 + logger.warning(
255 + "Failed to get Mooncake replica descriptors for tier logging "
256 + "(batch_keys=%d, error=%s); marking tiers unknown",
257 + len(keys),
258 + e,
259 + )
260 + return tiers_by_key
261 +
262 + for key in keys:
263 + if hasattr(replica_descs_by_key, "get"):
264 + replica_descs = replica_descs_by_key.get(key)
265 + else:
266 + try:
267 + replica_descs = replica_descs_by_key[key]
268 + except (KeyError, TypeError):
269 + replica_descs = None
270 + tiers_by_key[key] = _classify_replica_tier(replica_descs)
271 + return tiers_by_key
272 +
273 +
274 +def _log_mooncake_load_tier_summary(
275 + req_id: str,
276 + batch_keys: list[str],
277 + load_results: list[int],
278 + tiers_by_key: dict[str, str],
279 +) -> None:
280 + tier_counts = {"memory": 0, "disk": 0, "unknown": 0}
281 + bytes_by_tier = {"memory": 0, "disk": 0, "unknown": 0}
282 + success_keys = 0
283 + failed_keys = 0
284 +
285 + for index, key in enumerate(batch_keys):
286 + tier = tiers_by_key.get(key, "unknown")
287 + if tier not in tier_counts:
288 + tier = "unknown"
289 + tier_counts[tier] += 1
290 +
291 + value = load_results[index] if index < len(load_results) else -1
292 + if value >= 0:
293 + success_keys += 1
294 + bytes_by_tier[tier] += int(value)
295 + else:
296 + failed_keys += 1
297 +
298 + logger.info(
299 + "Mooncake load tier summary: req_id=%s batch_keys=%d "
300 + "memory_keys=%d disk_keys=%d unknown_keys=%d "
301 + "success_keys=%d failed_keys=%d bytes_by_tier=%s",
302 + req_id,
303 + len(batch_keys),
304 + tier_counts["memory"],
305 + tier_counts["disk"],
306 + tier_counts["unknown"],
307 + success_keys,
308 + failed_keys,
309 + bytes_by_tier,
310 + )
311 +
312 +
313 # ============================================================
314 # Transfer Threads
315 # ============================================================
@@ -207,6 +392,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
392 kv_role: str,
393 ready_event: threading.Event,
394 enable_kv_event: bool = False,
395 + replicate_config: Any = None,
396 ):
397 super().__init__(
398 store,
@@ -220,8 +406,11 @@ class KVCacheStoreSendingThread(KVTransferThread):
406 self.kv_role = kv_role
407 self.stored_requests: defaultdict[str, int] = defaultdict(int)
408 self.enable_kv_event = enable_kv_event
409 + # Caller always passes a non-None ReplicateConfig — see
410 + # MooncakeStoreWorker.__init__ where store_replicate_config is built.
411 + self.replicate_config = replicate_config
412
224 - # Pause store requests when CPU offloading is under pressure.
413 + # Pause store requests when CPU/disk offloading is under pressure.
414 self._store_pressure_active = False
415 self._skip_store_requests: set[str] = set()
416
@@ -270,7 +459,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
459 return
460 if self._should_skip_request(req_id):
461 logger.debug(
273 - "Skipping Mooncake store for request %s while CPU offloading "
462 + "Skipping Mooncake store for request %s while CPU/disk offloading "
463 "is under pressure",
464 req_id,
465 )
@@ -357,7 +546,12 @@ class KVCacheStoreSendingThread(KVTransferThread):
546 current_event.synchronize()
547
548 try:
360 - res = self.store.batch_put_from_multi_buffers(keys, addrs, sizes)
549 + res = self.store.batch_put_from_multi_buffers(
550 + keys,
551 + addrs,
552 + sizes,
553 + self.replicate_config,
554 + )
555 failed = [i for i, v in enumerate(res) if v < 0]
556 if failed:
557 # Compute total bytes attempted for this batch
@@ -379,7 +573,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
573 and not self._mark_request_skipped_for_pressure(req_id)
574 ):
575 logger.warning(
382 - "Detected Mooncake CPU offloading pressure "
576 + "Detected Mooncake CPU/disk offloading pressure "
577 "(NO_AVAILABLE_HANDLE); skipping future store "
578 "batches for request %s until a later store "
579 "batch succeeds",
@@ -387,7 +581,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
581 )
582 elif self._clear_store_pressure():
583 logger.info(
390 - "Mooncake CPU offloading pressure cleared after a "
584 + "Mooncake CPU/disk offloading pressure cleared after a "
585 "successful store batch"
586 )
587 except Exception as e:
@@ -410,6 +604,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
604 block_size: int,
605 tp_rank: int,
606 ready_event: threading.Event,
607 + disk_offload_buffer_budget_bytes: int | None = None,
608 ):
609 super().__init__(
610 store,
@@ -419,6 +614,14 @@ class KVCacheStoreRecvingThread(KVTransferThread):
614 ready_event,
615 name="KVCacheStoreRecvingThread",
616 )
617 + self.disk_offload_buffer_budget_bytes = disk_offload_buffer_budget_bytes
618 + self.usable_disk_offload_buffer_budget_bytes = (
619 + None
620 + if disk_offload_buffer_budget_bytes is None
621 + else _get_usable_disk_offload_buffer_budget_bytes(
622 + disk_offload_buffer_budget_bytes
623 + )
624 + )
625
626 def _handle_request(self, req_meta: ReqMeta):
627 token_len = req_meta.load_spec.token_len # type: ignore[union-attr]
@@ -456,26 +659,69 @@ class KVCacheStoreRecvingThread(KVTransferThread):
659 + size_list[: self.tp_rank % len(size_list)]
660 )
661
459 - try:
460 - res = self.store.batch_get_into_multi_buffers(
461 - key_list_c, addr_list_c, size_list_c
662 + load_batches = [(key_list_c, addr_list_c, size_list_c)]
663 + if self.usable_disk_offload_buffer_budget_bytes is not None:
664 + total_staging_bytes = sum(
665 + _estimate_disk_offload_staging_bytes(size) for size in size_list_c
666 )
463 - failed = [
464 - (key, value)
465 - for key, value in zip(key_list_c, res, strict=True)
466 - if value < 0
467 - ]
468 - if failed:
469 - logger.warning(
470 - "Failed to get %d Mooncake keys (batch_keys=%d, first_failures=%s)",
471 - len(failed),
472 - len(key_list_c),
473 - failed[:3],
667 + if total_staging_bytes > self.usable_disk_offload_buffer_budget_bytes:
668 + assert self.disk_offload_buffer_budget_bytes is not None
669 + load_batches, oversized_key = _split_disk_offload_load_batches(
670 + key_list_c,
671 + addr_list_c,
672 + size_list_c,
673 + self.usable_disk_offload_buffer_budget_bytes,
674 + self.disk_offload_buffer_budget_bytes,
675 + )
676 + if oversized_key is not None:
677 + oversized_key_index = key_list_c.index(oversized_key)
678 + oversized_key_bytes = _estimate_disk_offload_staging_bytes(
679 + size_list_c[oversized_key_index]
680 + )
681 + logger.warning(
682 + "Skipping Mooncake load for request %s because key %s "
683 + "requires %d staging bytes, exceeding budget %d",
684 + req_id,
685 + oversized_key,
686 + oversized_key_bytes,
687 + self.disk_offload_buffer_budget_bytes,
688 + )
689 + self.set_finished_request(req_id)
690 + self.request_queue.task_done()
691 + return
692 +
693 + current_batch_keys: list[str] = key_list_c
694 + try:
695 + for batch_keys, batch_addrs, batch_sizes in load_batches:
696 + current_batch_keys = batch_keys
697 + tiers_by_key: dict[str, str] | None = None
698 + if envs.VLLM_MOONCAKE_STORE_TIER_LOG:
699 + tiers_by_key = _get_replica_tiers_by_key(self.store, batch_keys)
700 + res = self.store.batch_get_into_multi_buffers(
701 + batch_keys, batch_addrs, batch_sizes
702 )
703 + if tiers_by_key is not None:
704 + _log_mooncake_load_tier_summary(
705 + req_id, batch_keys, res, tiers_by_key
706 + )
707 + failed = [
708 + (key, value)
709 + for key, value in zip(batch_keys, res, strict=True)
710 + if value < 0
711 + ]
712 + if failed:
713 + logger.warning(
714 + "Failed to get %d Mooncake keys from sub-batch "
715 + "(batch_keys=%d, first_failures=%s)",
716 + len(failed),
717 + len(batch_keys),
718 + failed[:3],
719 + )
720 + break
721 except Exception as e:
722 logger.warning(
477 - "Failed to get Mooncake batch %s, error: %s",
478 - key_list_c[:3],
723 + "Failed to get Mooncake sub-batch %s, error: %s",
724 + current_batch_keys[:3],
725 e,
726 )
727
@@ -493,7 +739,10 @@ class MooncakeStoreWorker:
739
740 def __init__(self, vllm_config: VllmConfig):
741 try:
496 - from mooncake.store import MooncakeDistributedStore # type: ignore
742 + from mooncake.store import ( # type: ignore
743 + MooncakeDistributedStore,
744 + ReplicateConfig,
745 + )
746 except ImportError as e:
747 raise ImportError(
748 "Please install mooncake by following the instructions at "
@@ -561,24 +810,79 @@ class MooncakeStoreWorker:
810
811 # Initialize MooncakeDistributedStore with its own TransferEngine
812 store_config = MooncakeStoreConfig.load_from_env()
813 + extra_config = (
814 + vllm_config.kv_transfer_config.kv_connector_extra_config
815 + if vllm_config.kv_transfer_config
816 + else {}
817 + )
818 + store_config.device_name = rdma_utils.get_configured_worker_rnic(
819 + protocol=store_config.protocol,
820 + configured_device=store_config.device_name,
821 + )
822 self.store = MooncakeDistributedStore()
565 -
566 - local_seg = get_ip()
567 - config_dict = {
568 - "local_hostname": local_seg,
569 - "metadata_server": store_config.metadata_server,
570 - "global_segment_size": str(store_config.global_segment_size),
571 - "local_buffer_size": str(store_config.local_buffer_size),
572 - "protocol": store_config.protocol,
573 - "rdma_devices": store_config.device_name,
574 - "master_server_addr": store_config.master_server_address,
575 - }
576 - ret = self.store.setup(config_dict)
823 + local_ip = get_ip()
824 + local_hostname = rdma_utils.get_requester_local_hostname(local_ip)
825 + ret = self.store.setup(
826 + local_hostname,
827 + store_config.metadata_server,
828 + store_config.global_segment_size,
829 + store_config.local_buffer_size,
830 + store_config.protocol,
831 + store_config.device_name,
832 + store_config.master_server_address,
833 + )
834 if ret != 0:
835 msg = "Initialize MooncakeDistributedStore failed."
836 logger.error(msg)
837 raise RuntimeError(msg)
838
839 + preferred_segment = rdma_utils.get_configured_preferred_segment(extra_config)
840 + self.preferred_segment = preferred_segment
841 + self.store_replicate_config = ReplicateConfig()
842 + if preferred_segment is not None:
843 + self.store_replicate_config.preferred_segment = preferred_segment
844 +
845 + logger.info(
846 + "Mooncake mode=%s (global_segment_size=%d, local_buffer_size=%d, "
847 + "preferred_segment=%s, enable_offload=%s)",
848 + store_config.mode,
849 + store_config.global_segment_size,
850 + store_config.local_buffer_size,
851 + preferred_segment or "<none>",
852 + store_config.enable_offload,
853 + )
854 + if store_config.mode == "embedded":
855 + if store_config.enable_offload and preferred_segment is None:
856 + logger.warning(
857 + "enable_offload is set in embedded mode without "
858 + "preferred_segment; SSD tier will only see puts that "
859 + "happen to land on the owner segment."
860 + )
861 + if preferred_segment is not None:
862 + logger.warning(
863 + "preferred_segment=%s with mode=embedded: rank-"
864 + "contributed segments will be idle.",
865 + preferred_segment,
866 + )
867 + elif (
868 + store_config.mode == "standalone-store" and not store_config.enable_offload
869 + ):
870 + logger.warning(
871 + "standalone-store mode without enable_offload: large prefills "
872 + "may exceed the owner DirectIO budget."
873 + )
874 +
875 + self.disk_offload_buffer_budget_bytes = (
876 + DEFAULT_MOONCAKE_DISK_STAGING_BUFFER_BYTES
877 + if store_config.enable_offload
878 + else None
879 + )
880 +
881 + # Start lookup server on rank 0 for scheduler-side prefix queries
882 + self.lookup_server: LookupKeyServer | None = None
883 + if vllm_config.parallel_config.rank == 0:
884 + self.lookup_server = LookupKeyServer(self, vllm_config)
885 +
886 kv_event_config = vllm_config.kv_events_config
887 self.enable_kv_events = False
888 if kv_event_config and kv_event_config.enable_kv_cache_events:
@@ -588,11 +892,6 @@ class MooncakeStoreWorker:
892 self.kv_recv_thread: KVCacheStoreRecvingThread | None = None
893 self.finished_store_req: set[str] = set()
894
591 - # Start lookup server on rank 0 for scheduler-side prefix queries
592 - self.lookup_server: LookupKeyServer | None = None
593 - if vllm_config.parallel_config.rank == 0:
594 - self.lookup_server = LookupKeyServer(self, vllm_config)
595 -
895 def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
896 """Register a cross-layers KV cache tensor.
897
@@ -695,6 +994,7 @@ class MooncakeStoreWorker:
994 self.kv_role,
995 ready_event_sending,
996 self.enable_kv_events,
997 + self.store_replicate_config,
998 )
999 self.kv_send_thread.start()
1000
@@ -705,6 +1005,7 @@ class MooncakeStoreWorker:
1005 self.block_size,
1006 self.tp_rank,
1007 ready_event_recving,
1008 + disk_offload_buffer_budget_bytes=self.disk_offload_buffer_budget_bytes,
1009 )
1010 self.kv_recv_thread.start()
1011 ready_event_recving.wait()
@@ -967,13 +1268,15 @@ class LookupKeyClient:
1268
1269 def get_zmq_rpc_path_lookup(vllm_config: VllmConfig) -> str:
1270 """Construct IPC path for ZMQ lookup socket."""
1271 + assert vllm_config.kv_transfer_config is not None
1272 dp_rank = get_mooncake_dp_engine_index(vllm_config.parallel_config)
1273 base_url = envs.VLLM_RPC_BASE_PATH
1274 rpc_port = 0
973 - assert vllm_config.kv_transfer_config is not None
1275 + hostname = socket.gethostname()
1276 extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
1277 if "lookup_rpc_port" in extra_config:
1278 rpc_port = extra_config["lookup_rpc_port"]
977 - uid = os.getuid()
978 - logger.debug("Base URL: %s, RPC Port: %s, UID: %s", base_url, rpc_port, uid)
979 - return f"ipc://{base_url}/lookup_rpc_port_{rpc_port}_uid{uid}_dp_rank{dp_rank}"
1279 + logger.debug("Base URL: %s, RPC Port: %s", base_url, rpc_port)
1280 + return (
1281 + f"ipc://{base_url}/lookup_rpc_port_{rpc_port}_host_{hostname}_dp_rank{dp_rank}"
1282 + )
vllm/envs.py
+18
@@ -191,6 +191,10 @@ if TYPE_CHECKING:
191 VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
192 VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600
193 VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998
194 + VLLM_MOONCAKE_STORE_TIER_LOG: bool = False
195 + VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO: float = 0.9
196 + MOONCAKE_PREFERRED_SEGMENT: str | None = None
197 + MOONCAKE_REQUESTER_LOCAL_HOSTNAME: str | None = None
198 VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840
199 VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1
200 VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
@@ -1389,6 +1393,20 @@ environment_variables: dict[str, Callable[[], Any]] = {
1393 "VLLM_MOONCAKE_BOOTSTRAP_PORT": lambda: int(
1394 os.getenv("VLLM_MOONCAKE_BOOTSTRAP_PORT", "8998")
1395 ),
1396 + # Log per-batch memory/disk tier breakdown on external GETs.
1397 + "VLLM_MOONCAKE_STORE_TIER_LOG": lambda: (
1398 + os.getenv("VLLM_MOONCAKE_STORE_TIER_LOG", "False").lower() in ("true", "1")
1399 + ),
1400 + # Fraction of the owner's DirectIO staging buffer to fill per GET batch.
1401 + "VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO": lambda: float(
1402 + os.getenv("VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO", "0.9")
1403 + ),
1404 + # Pin this rank to a specific owner segment ("host:port").
1405 + "MOONCAKE_PREFERRED_SEGMENT": lambda: os.getenv("MOONCAKE_PREFERRED_SEGMENT"),
1406 + # Override the hostname the rank registers as a Mooncake requester.
1407 + "MOONCAKE_REQUESTER_LOCAL_HOSTNAME": lambda: os.getenv(
1408 + "MOONCAKE_REQUESTER_LOCAL_HOSTNAME"
1409 + ),
1410 # Flashinfer MoE backend for vLLM's fused Mixture-of-Experts support.
1411 # Both require compute capability 10.0 or above.
1412 # Available options: