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

Signed-off-by: Zhewen Li <[email protected]>
Co-authored-by: Zhewen Li <[email protected]>
Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>

docs/features/mooncake_store_connector_usage.md

@@ -2,7 +2,7 @@
 
 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:
 
-- **CPU offloading**: Extend effective KV cache capacity by offloading to CPU memory via Mooncake's transfer engine.
+- **CPU/disk offloading**: Extend effective KV cache capacity by offloading to CPU memory or disk via Mooncake's transfer engine.
 - **Prefix caching across instances**: Hash-based deduplication allows multiple vLLM instances to share cached KV blocks through the store.
 - **Single-node and multi-node deployment**: Works both as a standalone KV cache extension and in disaggregated prefill-decode setups.
 
@@ -38,18 +38,29 @@ Create a JSON configuration file (e.g., `mooncake_config.json`):
 
 ```json
 {
+  "mode": "embedded",
   "metadata_server": "P2PHANDSHAKE",
   "master_server_address": "127.0.0.1:50051",
   "global_segment_size": "80GB",
   "local_buffer_size": "4GB",
   "protocol": "rdma",
-  "device_name": ""
+  "device_name": "",
+  "enable_offload": false
 }
 ```
 
+- `mode`: Topology selection. `"embedded"` (default, PR-40900 baseline) has each
+  vLLM rank contribute `global_segment_size` to the pool in-process.
+  `"standalone-store"` makes ranks pure requesters — an external
+  `mooncake_client` process owns the CPU pool and (optionally) the SSD tier.
 - `protocol`: Use `"rdma"` for best performance. `"tcp"` works as a fallback.
-- `global_segment_size`: CPU memory contributed to the distributed pool (per GPU).
+- `global_segment_size`: CPU memory contributed to the distributed pool (per
+  GPU). Must be `> 0` in `embedded` mode and `0` in `standalone-store` mode.
 - `local_buffer_size`: Private buffer for this node's own operations (per GPU).
+- `enable_offload`: When `true`, vLLM allocates a DirectIO staging buffer so
+  large prefills do not exceed the owner's SSD-write budget. Set this together
+  with the matching `--enable_offload=true` flag on `mooncake_master` and on
+  the external `mooncake_client` (if any).
 
 Set the config path via environment variable:
 
@@ -126,12 +137,60 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \
 
 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.
 
+### Disk Offloading
+
+Disk offloading is most commonly run in `standalone-store` mode: an external
+`mooncake_client` process owns the CPU pool and the SSD tier, and each vLLM
+rank is a pure requester. This avoids per-rank duplication of the SSD pool
+and keeps DirectIO budget tracking on a single process.
+
+Three things need to be aligned for end-to-end disk offloading:
+
+1. **`mooncake_master`** is started with `--enable_offload=true`.
+2. **`mooncake_client`** (the owner) is started with `--enable_offload=true`
+   plus an SSD path via `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`.
+3. **vLLM-side** sets `"enable_offload": true` in the JSON config file (this is
+   read by the connector and is **not** an environment variable).
+
+Example `mooncake_config.json` for the vLLM side:
+
+```json
+{
+  "mode": "standalone-store",
+  "metadata_server": "P2PHANDSHAKE",
+  "master_server_address": "127.0.0.1:50051",
+  "global_segment_size": 0,
+  "local_buffer_size": "4GB",
+  "protocol": "rdma",
+  "device_name": "mlx5_0",
+  "enable_offload": true
+}
+```
+
+Steer this rank to the local owner segment with:
+
+```bash
+export MOONCAKE_PREFERRED_SEGMENT=127.0.0.1:50053
+```
+
+The owner's SSD directory, on-disk eviction policy, and the DirectIO staging
+buffer size are controlled on the `mooncake_client` side via the standard
+Mooncake environment variables (`MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`,
+`MOONCAKE_BUCKET_EVICTION_POLICY`, `MOONCAKE_USE_URING`,
+`MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES`,
+`MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES`, etc.). Those are independent of
+the vLLM JSON config.
+
 ## Environment Variables
 
 | Variable | Description | Default |
 | --- | --- | --- |
 | `MOONCAKE_CONFIG_PATH` | Path to Mooncake JSON config file | (required) |
 | `VLLM_MOONCAKE_BOOTSTRAP_PORT` | Bootstrap port for MooncakeConnector P2P transfer (disagg mode only) | 8998 |
+| `MOONCAKE_PREFERRED_SEGMENT` | Pin this rank's replicas to a specific owner segment (`host:port`); used in `standalone-store` mode | — |
+| `MOONCAKE_REQUESTER_LOCAL_HOSTNAME` | Override the hostname the vLLM rank registers with Mooncake as a requester. Defaults to the rank's resolved IP. | — |
+| `VLLM_MOONCAKE_STORE_TIER_LOG` | When `1`, logs a per-batch tier summary (memory vs disk hits) for observability | disabled |
+| `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 |
 
 ## KV Transfer Config
 

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

@@ -1,23 +1,36 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
+import logging
+import math
+import sys
 import threading
+import types
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
+import pytest
 import torch
 
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
+    rdma_utils,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
     worker,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
     ChunkedTokenDatabase,
     KeyMetadata,
+    LoadSpec,
     ReqMeta,
 )
 
 
 def _make_store_sending_thread(
     store: MagicMock,
+    *,
+    replicate_config: object | None = None,
 ) -> worker.KVCacheStoreSendingThread:
     token_database = ChunkedTokenDatabase(
         KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
@@ -32,11 +45,55 @@ def _make_store_sending_thread(
         put_step=1,
         kv_role="kv_producer",
         ready_event=threading.Event(),
+        replicate_config=replicate_config,
     )
     thread.request_queue.task_done = MagicMock()
     return thread
 
 
+def _make_store_recving_thread(
+    store: MagicMock,
+    *,
+    disk_offload_buffer_budget_bytes: int | None = None,
+) -> worker.KVCacheStoreRecvingThread:
+    token_database = ChunkedTokenDatabase(
+        KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
+    )
+    token_database.set_kv_caches_base_addr([0x1000])
+    token_database.set_block_len([256])
+    thread = worker.KVCacheStoreRecvingThread(
+        store=store,
+        token_database=token_database,
+        block_size=16,
+        tp_rank=0,
+        ready_event=threading.Event(),
+        disk_offload_buffer_budget_bytes=disk_offload_buffer_budget_bytes,
+    )
+    thread.request_queue.task_done = MagicMock()
+    return thread
+
+
+def _make_load_req(
+    req_id: str,
+    block_hashes: list[bytes],
+    *,
+    token_len: int,
+    vllm_cached_tokens: int = 0,
+) -> ReqMeta:
+    return ReqMeta(
+        req_id=req_id,
+        token_len_chunk=token_len,
+        block_ids=list(range(len(block_hashes))),
+        block_hashes=block_hashes,
+        load_spec=LoadSpec(
+            vllm_cached_tokens=vllm_cached_tokens,
+            kvpool_cached_tokens=token_len,
+            can_load=True,
+            token_len=token_len,
+        ),
+    )
+
+
 def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
     return ReqMeta(
         req_id=req_id,
@@ -48,6 +105,233 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
     )
 
 
+_DISK_OFFLOAD_SINGLE_KEY_BYTES = worker._estimate_disk_offload_staging_bytes([256])
+_DISK_OFFLOAD_USABLE_BUDGET_RATIO = 0.9
+_DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS = 4 * _DISK_OFFLOAD_SINGLE_KEY_BYTES
+_DISK_OFFLOAD_BUDGET_FOR_SPLIT = math.ceil(
+    2 * _DISK_OFFLOAD_SINGLE_KEY_BYTES / _DISK_OFFLOAD_USABLE_BUDGET_RATIO
+)  # Allows two 256-byte chunks but not the third.
+_DISK_OFFLOAD_BUDGET_TOO_SMALL = (
+    _DISK_OFFLOAD_SINGLE_KEY_BYTES - 1
+)  # Smaller than a single 256-byte chunk.
+
+
+class _FakeKVTransferConfig:
+    def __init__(
+        self,
+        *,
+        kv_role: str = "kv_both",
+        extra_config: dict[str, object] | None = None,
+    ) -> None:
+        self.kv_role = kv_role
+        self.kv_connector_extra_config = extra_config or {}
+
+    def get_from_extra_config(self, key: str, default: object) -> object:
+        return self.kv_connector_extra_config.get(key, default)
+
+
+class _FakeModelConfig:
+    model = "test-model"
+    use_mla = False
+
+    def get_num_layers(self, parallel_config) -> int:
+        return 1
+
+    def get_total_num_kv_heads(self) -> int:
+        return 1
+
+
+def _make_vllm_config(
+    *, extra_config: dict[str, object] | None = None
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        model_config=_FakeModelConfig(),
+        parallel_config=SimpleNamespace(
+            pipeline_parallel_size=1,
+            rank=0,
+        ),
+        kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config),
+        cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=10),
+        kv_events_config=SimpleNamespace(enable_kv_cache_events=False),
+    )
+
+
+def _write_mooncake_config(tmp_path, config: dict[str, object]) -> str:
+    config_path = tmp_path / "mooncake_config.json"
+    config_path.write_text(json.dumps(config), encoding="utf-8")
+    return str(config_path)
+
+
+def _install_fake_mooncake(monkeypatch, store_instance: MagicMock):
+    class FakeReplicateConfig:
+        def __init__(self) -> None:
+            self.preferred_segment = ""
+
+    fake_store_module = types.ModuleType("mooncake.store")
+    fake_store_module.MooncakeDistributedStore = lambda: store_instance  # type: ignore[attr-defined]
+    fake_store_module.ReplicateConfig = FakeReplicateConfig  # type: ignore[attr-defined]
+    fake_mooncake_module = types.ModuleType("mooncake")
+    fake_mooncake_module.store = fake_store_module  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module)
+    monkeypatch.setitem(sys.modules, "mooncake.store", fake_store_module)
+    return FakeReplicateConfig
+
+
+def _patch_worker_runtime(monkeypatch, *, local_ip: str = "10.0.0.7") -> None:
+    single_rank_group = SimpleNamespace(world_size=1, rank_in_group=0)
+    monkeypatch.setattr(worker, "get_mooncake_dp_engine_index", lambda _: 0)
+    monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
+    monkeypatch.setattr(worker, "get_dcp_group", lambda: single_rank_group)
+    monkeypatch.setattr(worker, "get_ip", lambda: local_ip)
+
+
+def test_default_local_buffer_size_matches_pr40900():
+    """PR-40900 shipped a 4 GiB default for local_buffer_size; the dual-mode
+    patch preserves it (and the JSON key) so unchanged PR-40900 configs work."""
+    assert worker.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3
+
+
+def test_get_requester_local_hostname_prefers_override(monkeypatch):
+    monkeypatch.setenv("MOONCAKE_REQUESTER_LOCAL_HOSTNAME", "worker-a:50053")
+
+    assert rdma_utils.get_requester_local_hostname("10.0.0.7") == "worker-a:50053"
+
+
+def test_get_configured_preferred_segment_returns_explicit_override():
+    assert (
+        rdma_utils.get_configured_preferred_segment(
+            {"preferred_segment": "10.0.0.7:50053"}
+        )
+        == "10.0.0.7:50053"
+    )
+
+
+def test_get_configured_preferred_segment_prefers_explicit_over_env(monkeypatch):
+    monkeypatch.setenv("MOONCAKE_PREFERRED_SEGMENT", "10.0.0.8:50053")
+
+    assert (
+        rdma_utils.get_configured_preferred_segment(
+            {"preferred_segment": "10.0.0.7:50053"}
+        )
+        == "10.0.0.7:50053"
+    )
+
+
+def test_get_configured_preferred_segment_returns_env_override(monkeypatch):
+    monkeypatch.setenv("MOONCAKE_PREFERRED_SEGMENT", "10.0.0.8:50053")
+
+    assert rdma_utils.get_configured_preferred_segment({}) == "10.0.0.8:50053"
+
+
+def test_get_configured_preferred_segment_rejects_empty_override():
+    with pytest.raises(ValueError, match="preferred_segment"):
+        rdma_utils.get_configured_preferred_segment({"preferred_segment": "  "})
+
+
+def test_get_configured_worker_rnic_prefers_explicit_device_name(monkeypatch):
+    store_config = worker.MooncakeStoreConfig(
+        metadata_server="",
+        local_buffer_size=1,
+        protocol="rdma",
+        device_name="rocep139s0",
+        master_server_address="",
+    )
+
+    assert (
+        rdma_utils.get_configured_worker_rnic(
+            protocol=store_config.protocol,
+            configured_device=store_config.device_name,
+        )
+        == "rocep139s0"
+    )
+
+
+def test_get_configured_worker_rnic_selects_device_from_explicit_csv(monkeypatch):
+    monkeypatch.setattr(
+        rdma_utils,
+        "get_current_physical_gpu_index",
+        lambda: 1,
+    )
+    store_config = worker.MooncakeStoreConfig(
+        metadata_server="",
+        local_buffer_size=1,
+        protocol="rdma",
+        device_name="rocep139s0,rocep140s0",
+        master_server_address="",
+    )
+
+    assert (
+        rdma_utils.get_configured_worker_rnic(
+            protocol=store_config.protocol,
+            configured_device=store_config.device_name,
+        )
+        == "rocep140s0"
+    )
+
+
+def test_get_configured_worker_rnic_warns_and_returns_empty_for_rdma_with_no_device(
+    caplog, monkeypatch
+):
+    """No device configured + protocol=rdma → emit a clear warning and return ""
+    so the C++ side handles auto-selection. There is no Python-side fallback."""
+    monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
+    with caplog.at_level(logging.WARNING):
+        result = rdma_utils.get_configured_worker_rnic(
+            protocol="rdma",
+            configured_device="",
+        )
+    assert result == ""
+    warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
+    assert any("No RDMA devices specified" in r.message for r in warnings), (
+        f"expected fallback warning, got {[r.message for r in warnings]}"
+    )
+
+
+def test_get_configured_worker_rnic_silent_for_tcp_with_no_device(caplog, monkeypatch):
+    """protocol=tcp + no device → return "" silently (no RDMA, no warning)."""
+    monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
+    with caplog.at_level(logging.WARNING):
+        result = rdma_utils.get_configured_worker_rnic(
+            protocol="tcp",
+            configured_device="",
+        )
+    assert result == ""
+    warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
+    assert not any("RDMA" in r.message for r in warnings), (
+        "did not expect RDMA warning for tcp protocol, got "
+        f"{[r.message for r in warnings]}"
+    )
+
+
+def test_get_configured_worker_rnic_rejects_short_explicit_csv(monkeypatch):
+    monkeypatch.setattr(
+        rdma_utils,
+        "get_current_physical_gpu_index",
+        lambda: 2,
+    )
+    with pytest.raises(ValueError, match="does not cover local GPU 2"):
+        rdma_utils.get_configured_worker_rnic(
+            protocol="rdma",
+            configured_device="rocep139s0,rocep140s0",
+        )
+
+
+class _ReplicaDesc:
+    def __init__(self, tier: str):
+        self.tier = tier
+
+    def is_memory_replica(self) -> bool:
+        return self.tier == "memory"
+
+    def is_disk_replica(self) -> bool:
+        return self.tier == "disk"
+
+    def is_local_disk_replica(self) -> bool:
+        return self.tier == "disk"
+
+
 def test_store_sending_thread_skips_request_during_cpu_pressure():
     store = MagicMock()
     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():
     assert store.batch_put_from_multi_buffers.call_count == 2
 
 
+def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set():
+    store = MagicMock()
+    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
+    store.batch_put_from_multi_buffers.return_value = [256, 256]
+    replicate_config = SimpleNamespace(preferred_segment="10.0.0.7:50053")
+    thread = _make_store_sending_thread(store, replicate_config=replicate_config)
+
+    thread.add_stored_request("req-a")
+    thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
+
+    assert store.batch_put_from_multi_buffers.call_count == 1
+    call_args = store.batch_put_from_multi_buffers.call_args.args
+    assert len(call_args) == 4
+    assert call_args[3] is replicate_config
+
+
+def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_segment():  # noqa: E501
+    """Without a preferred_segment the SendingThread still forwards a
+    (default-constructed) ReplicateConfig so the C++ side always sees a
+    well-defined config object."""
+    store = MagicMock()
+    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
+    store.batch_put_from_multi_buffers.return_value = [256, 256]
+    replicate_config = SimpleNamespace()
+    thread = _make_store_sending_thread(store, replicate_config=replicate_config)
+
+    thread.add_stored_request("req-a")
+    thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
+
+    assert store.batch_put_from_multi_buffers.call_count == 1
+    call_args = store.batch_put_from_multi_buffers.call_args.args
+    assert len(call_args) == 4
+    assert call_args[3] is replicate_config
+
+
+def test_estimate_disk_offload_staging_bytes_sums_multi_segment_sizes():
+    assert worker._estimate_disk_offload_staging_bytes([256, 512]) == 12288
+
+
+def test_recv_thread_uses_single_batch_when_no_disk_offload_budget(monkeypatch):
+    monkeypatch.delenv("VLLM_MOONCAKE_STORE_TIER_LOG", raising=False)
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [256, 256, 256]
+    thread = _make_store_recving_thread(store, disk_offload_buffer_budget_bytes=None)
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 1
+    keys, addrs, sizes = store.batch_get_into_multi_buffers.call_args.args
+    assert keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
+    ]
+    assert sizes == [[256], [256], [256]]
+    store.batch_get_replica_desc.assert_not_called()
+
+
+def test_recv_thread_logs_tier_summary_when_enabled(monkeypatch, caplog_vllm):
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_TIER_LOG", "1")
+    caplog_vllm.set_level(logging.INFO, logger=worker.logger.name)
+
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [256, 256, -10]
+    thread = _make_store_recving_thread(store, disk_offload_buffer_budget_bytes=None)
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+    expected_keys = [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
+    ]
+    store.batch_get_replica_desc.return_value = {
+        expected_keys[0]: [_ReplicaDesc("memory")],
+        expected_keys[1]: [_ReplicaDesc("disk")],
+        expected_keys[2]: [],
+    }
+
+    thread._handle_request(req)
+
+    assert store.batch_get_replica_desc.call_args.args == (expected_keys,)
+    assert store.method_calls[0][0] == "batch_get_replica_desc"
+    assert store.method_calls[1][0] == "batch_get_into_multi_buffers"
+
+    messages = [record.getMessage() for record in caplog_vllm.records]
+    assert any(
+        "Mooncake load tier summary" in message
+        and "req_id=req-a" in message
+        and "batch_keys=3" in message
+        and "memory_keys=1" in message
+        and "disk_keys=1" in message
+        and "unknown_keys=1" in message
+        and "success_keys=2" in message
+        and "failed_keys=1" in message
+        and "bytes_by_tier={'memory': 256, 'disk': 256, 'unknown': 0}" in message
+        for message in messages
+    )
+
+
+def test_recv_thread_uses_ratio_scaled_budget_for_first_pass_split():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.side_effect = [
+        [256],
+        [256],
+    ]
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=2 * _DISK_OFFLOAD_SINGLE_KEY_BYTES,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1"],
+        token_len=32,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 2
+    first_keys = store.batch_get_into_multi_buffers.call_args_list[0].args[0]
+    second_keys = store.batch_get_into_multi_buffers.call_args_list[1].args[0]
+    assert first_keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+    ]
+    assert second_keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+    ]
+
+
+def test_recv_thread_splits_disk_offload_loads_by_budget():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.side_effect = [
+        [256, 256],
+        [256],
+    ]
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_FOR_SPLIT,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 2
+
+    first_keys = store.batch_get_into_multi_buffers.call_args_list[0].args[0]
+    second_keys = store.batch_get_into_multi_buffers.call_args_list[1].args[0]
+    first_addrs = store.batch_get_into_multi_buffers.call_args_list[0].args[1]
+    second_addrs = store.batch_get_into_multi_buffers.call_args_list[1].args[1]
+    first_sizes = store.batch_get_into_multi_buffers.call_args_list[0].args[2]
+    second_sizes = store.batch_get_into_multi_buffers.call_args_list[1].args[2]
+    assert first_keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+    ]
+    assert second_keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
+    ]
+    base_addr = thread.token_database.kv_caches_base_addr[0]
+    block_len = thread.token_database.block_len[0]
+    assert first_addrs == [[base_addr], [base_addr + block_len]]
+    assert second_addrs == [[base_addr + 2 * block_len]]
+    expected_size = block_len
+    assert first_sizes == [[expected_size], [expected_size]]
+    assert second_sizes == [[expected_size]]
+
+
+def test_recv_thread_stops_after_first_failing_disk_offload_sub_batch():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [-10, -10]
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_FOR_SPLIT,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 1
+
+
+def test_recv_thread_skips_split_when_budget_holds_all_keys():
+    """PR-36 removed the count-based split trigger; with budget for 3 keys,
+    all three should be requested in a single call."""
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [256, 256, 256]
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 1
+    assert store.batch_get_into_multi_buffers.call_args_list[0].args[0] == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
+    ]
+
+
+def test_recv_thread_reports_unsplittable_key_larger_than_budget():
+    store = MagicMock()
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_TOO_SMALL,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0"],
+        token_len=16,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 0
+
+
+def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "global_segment_size": "4gb",
+                "local_buffer_size": "64mb",
+                "protocol": "rdma",
+                "device_name": "mlx5_0",
+                "master_server_address": "10.0.0.7:50051",
+                "enable_offload": True,
+            },
+        ),
+    )
+    w = worker.MooncakeStoreWorker(_make_vllm_config())
+
+    assert not hasattr(w, "_isolate_offload_resources")
+    assert store.setup.call_args.args == (
+        "10.0.0.7",
+        "http://metadata/endpoint",
+        4 * 1024 * 1024 * 1024,  # global_segment_size: "4gb" honored
+        64 * 1024 * 1024,
+        "rdma",
+        "mlx5_0",
+        "10.0.0.7:50051",
+    )
+
+
+def test_requester_worker_init_prefers_local_hostname_override(
+    tmp_path,
+    monkeypatch,
+):
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv("MOONCAKE_REQUESTER_LOCAL_HOSTNAME", "worker-a:50053")
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "local_buffer_size": "64mb",
+                "protocol": "tcp",
+                "device_name": "",
+                "master_server_address": "10.0.0.7:50051",
+            },
+        ),
+    )
+    worker.MooncakeStoreWorker(_make_vllm_config())
+
+    assert store.setup.call_args.args[0] == "worker-a:50053"
+
+
+def test_requester_worker_init_skips_disk_budget_when_offload_disabled(
+    tmp_path,
+    monkeypatch,
+):
+    """enable_offload=False zeroes out the disk budget so we don't generate
+    redundant owner GET-RPCs."""
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "protocol": "tcp",
+                "device_name": "",
+                "master_server_address": "10.0.0.7:50051",
+                "enable_offload": False,
+            },
+        ),
+    )
+    w = worker.MooncakeStoreWorker(_make_vllm_config())
+
+    assert w.disk_offload_buffer_budget_bytes is None
+
+
+def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
+    tmp_path,
+    monkeypatch,
+):
+    store = MagicMock()
+    store.setup.return_value = 0
+    fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "protocol": "tcp",
+                "device_name": "",
+                "master_server_address": "10.0.0.7:50051",
+            },
+        ),
+    )
+    w = worker.MooncakeStoreWorker(
+        _make_vllm_config(
+            extra_config={
+                "preferred_segment": "10.0.0.7:50053",
+            }
+        )
+    )
+
+    assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
+    assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
+
+
 # ---------------------------------------------------------------------------
 # Helpers for register_kv_caches tests
 # ---------------------------------------------------------------------------
@@ -147,8 +797,10 @@ def _make_bare_worker(
     w.tp_rank = 0
     w.put_step = 1
     w.enable_kv_events = False
+    w.disk_offload_buffer_budget_bytes = None
     w.kv_send_thread = None
     w.kv_recv_thread = None
+    w.store_replicate_config = SimpleNamespace()
     return w
 
 
@@ -298,3 +950,186 @@ def test_register_kv_caches_cross_layer_single_segment():
 
     assert w2.kv_caches_base_addr == w.kv_caches_base_addr
     assert w2.block_len == w.block_len
+
+
+# ---------------------------------------------------------------------------
+# Dual-mode (embedded / standalone-store) config validation tests
+# ---------------------------------------------------------------------------
+
+
+def _make_config(**overrides):
+    """Build a MooncakeStoreConfig with sensible defaults for validation tests.
+
+    Required dataclass fields are populated; callers override only the field
+    under test.
+    """
+    base = dict(
+        metadata_server="http://metadata/endpoint",
+        master_server_address="10.0.0.7:50051",
+        protocol="rdma",
+        device_name="mlx5_0",
+    )
+    base.update(overrides)
+    return worker.MooncakeStoreConfig(**base)
+
+
+def test_config_defaults_to_embedded():
+    """A JSON without explicit mode parses as embedded with 4 GiB segment."""
+    cfg = _make_config()
+    assert cfg.mode == "embedded"
+    assert cfg.global_segment_size == worker.DEFAULT_GLOBAL_SEGMENT_SIZE
+    assert cfg.local_buffer_size == worker.DEFAULT_LOCAL_BUFFER_SIZE
+    assert cfg.enable_offload is False
+
+
+def test_config_pr40900_unchanged(tmp_path):
+    """A literal PR-40900 config (no mode, no enable_offload, no preferred_segment)
+    parses without raising and resolves to embedded mode."""
+    config_path = _write_mooncake_config(
+        tmp_path,
+        {
+            "metadata_server": "http://metadata/endpoint",
+            "global_segment_size": "4GB",
+            "local_buffer_size": "4GB",
+            "protocol": "rdma",
+            "device_name": "mlx5_0",
+            "master_server_address": "10.0.0.7:50051",
+        },
+    )
+    cfg = worker.MooncakeStoreConfig.from_file(config_path)
+    assert cfg.mode == "embedded"
+    assert cfg.global_segment_size == 4 * 1024**3
+    assert cfg.local_buffer_size == 4 * 1024**3
+    assert cfg.enable_offload is False
+
+
+def test_config_embedded_rejects_zero_segment():
+    with pytest.raises(
+        ValueError, match=r"embedded mode requires global_segment_size > 0"
+    ):
+        _make_config(mode="embedded", global_segment_size=0)
+
+
+def test_config_standalone_store_rejects_nonzero_segment():
+    with pytest.raises(
+        ValueError,
+        match=r"standalone-store mode requires global_segment_size == 0",
+    ):
+        _make_config(mode="standalone-store", global_segment_size=4 * 1024**3)
+
+
+def test_config_standalone_store_accepts_zero_segment():
+    cfg = _make_config(mode="standalone-store", global_segment_size=0)
+    assert cfg.mode == "standalone-store"
+    assert cfg.global_segment_size == 0
+
+
+def test_config_unknown_mode():
+    with pytest.raises(ValueError, match=r"unknown Mooncake mode"):
+        _make_config(mode="something-else")
+
+
+def test_config_zero_local_buffer():
+    with pytest.raises(ValueError, match=r"local_buffer_size must be > 0"):
+        _make_config(local_buffer_size=0)
+
+
+# ---------------------------------------------------------------------------
+# End-to-end topology tests
+# Covers the two supported recipes:
+#   (A) standalone-store mode + disk offload (mode="standalone-store",
+#       segment=0, enable_offload=true, preferred_segment set)
+#   (B) embedded mode + CPU only      (mode default, segment>0,
+#       enable_offload=false, no preferred_segment)
+# ---------------------------------------------------------------------------
+
+
+def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch):
+    """standalone-store + disk: global_segment_size=0, enable_offload=True,
+    preferred_segment set. Assert setup() positional args, ReplicateConfig
+    wiring, and that the disk-offload buffer budget is allocated."""
+    store = MagicMock()
+    store.setup.return_value = 0
+    fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "mode": "standalone-store",
+                "metadata_server": "http://metadata/endpoint",
+                "global_segment_size": 0,
+                "local_buffer_size": "1GB",
+                "protocol": "rdma",
+                "device_name": "mlx5_0",
+                "master_server_address": "10.0.0.7:50051",
+                "enable_offload": True,
+            },
+        ),
+    )
+
+    w = worker.MooncakeStoreWorker(
+        _make_vllm_config(extra_config={"preferred_segment": "10.0.0.7:50053"})
+    )
+
+    # setup() receives global_segment_size=0 and the configured local buffer.
+    assert store.setup.call_args.args == (
+        "10.0.0.7",
+        "http://metadata/endpoint",
+        0,
+        1024 * 1024 * 1024,
+        "rdma",
+        "mlx5_0",
+        "10.0.0.7:50051",
+    )
+    # ReplicateConfig is built and carries the preferred_segment.
+    assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
+    assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
+    # Disk-offload staging budget is allocated (enable_offload=True).
+    assert w.disk_offload_buffer_budget_bytes is not None
+    assert w.disk_offload_buffer_budget_bytes > 0
+
+
+def test_topology_embedded_cpu_only(tmp_path, monkeypatch):
+    """embedded + CPU-only: no mode key (defaults to embedded),
+    global_segment_size>0, enable_offload absent, no preferred_segment.
+    This is the PR-40900 baseline recipe."""
+    store = MagicMock()
+    store.setup.return_value = 0
+    fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "global_segment_size": "4GB",
+                "local_buffer_size": "4GB",
+                "protocol": "rdma",
+                "device_name": "mlx5_0",
+                "master_server_address": "10.0.0.7:50051",
+            },
+        ),
+    )
+
+    w = worker.MooncakeStoreWorker(_make_vllm_config())
+
+    # setup() receives global_segment_size=4 GiB (rank contributes a segment).
+    assert store.setup.call_args.args == (
+        "10.0.0.7",
+        "http://metadata/endpoint",
+        4 * 1024 * 1024 * 1024,
+        4 * 1024 * 1024 * 1024,
+        "rdma",
+        "mlx5_0",
+        "10.0.0.7:50051",
+    )
+    # No preferred_segment — ReplicateConfig is default-constructed (so the
+    # preferred_segment field keeps its default value).
+    assert w.preferred_segment is None
+    assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
+    assert w.store_replicate_config.preferred_segment == ""
+    # No disk budget — enable_offload was absent (defaults to False).
+    assert w.disk_offload_buffer_budget_bytes is None

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/rdma_utils.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Mooncake requester config helpers."""
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def normalize_string_override(value: Any) -> str | None:
+    if not isinstance(value, str):
+        return None
+    normalized = value.strip()
+    return normalized or None
+
+
+def get_current_physical_gpu_index() -> int | None:
+    try:
+        from vllm.platforms import current_platform
+    except ImportError:
+        return None
+
+    try:
+        device_index = torch.accelerator.current_device_index()
+        physical_device_id = current_platform.device_id_to_physical_device_id(
+            device_index
+        )
+        return int(physical_device_id)
+    except Exception:
+        return None
+
+
+def get_requester_local_hostname(local_ip: str) -> str:
+    override = normalize_string_override(envs.MOONCAKE_REQUESTER_LOCAL_HOSTNAME)
+    if override is not None:
+        return override
+    return local_ip
+
+
+def get_configured_preferred_segment(
+    extra_config: Mapping[str, Any],
+) -> str | None:
+    preferred_segment = normalize_string_override(extra_config.get("preferred_segment"))
+    if preferred_segment is not None:
+        return preferred_segment
+    if extra_config.get("preferred_segment") is not None:
+        raise ValueError(
+            "Mooncake preferred_segment override must be a non-empty string"
+        )
+
+    env_value = normalize_string_override(envs.MOONCAKE_PREFERRED_SEGMENT)
+    if env_value is not None:
+        logger.info(
+            "Mooncake preferred_segment from MOONCAKE_PREFERRED_SEGMENT: %s",
+            env_value,
+        )
+        return env_value
+    return None
+
+
+def _get_explicit_worker_rnic(device_list: str) -> str:
+    entries = [entry.strip() for entry in device_list.split(",")]
+    if any(not entry for entry in entries):
+        raise ValueError(
+            "Mooncake worker device_name contains an empty RDMA device entry"
+        )
+    if len(entries) == 1:
+        return entries[0]
+
+    gpu_index = get_current_physical_gpu_index()
+    if gpu_index is None:
+        raise RuntimeError(
+            "Mooncake RDMA requester could not determine the local physical GPU index"
+        )
+    if gpu_index >= len(entries):
+        raise ValueError(
+            "Mooncake worker device list does not cover local GPU "
+            f"{gpu_index}: {device_list}"
+        )
+    device_name = entries[gpu_index]
+    logger.info(
+        "Mooncake selected worker RNIC %s from explicit device list for local GPU %s",
+        device_name,
+        gpu_index,
+    )
+    return device_name
+
+
+def get_configured_worker_rnic(
+    *,
+    protocol: str,
+    configured_device: str,
+) -> str:
+    normalized_device = normalize_string_override(configured_device)
+    if normalized_device is not None:
+        return _get_explicit_worker_rnic(normalized_device)
+
+    if protocol not in {"rdma", "efa"}:
+        return ""
+
+    logger.warning(
+        "No RDMA devices specified for Mooncake backend (protocol=%s). "
+        "Set 'device_name' in mooncake_config.json to a single RNIC name "
+        "or a comma-separated CSV indexed by physical GPU; falling back to "
+        "Mooncake's built-in auto-selection, which may converge on the same "
+        "NIC across all DP ranks and saturate bandwidth.",
+        protocol,
+    )
+    return ""

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

@@ -13,10 +13,11 @@ and MooncakeDistributedStore integration.
 import json
 import os
 import queue
+import socket
 import threading
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 import regex as re
 import torch
@@ -31,10 +32,11 @@ from vllm.distributed import (
     get_tensor_model_parallel_world_size,
 )
 from vllm.distributed.kv_events import BlockStored
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
     get_mooncake_dp_engine_index,
 )
-from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
     ChunkedTokenDatabase,
     KeyMetadata,
     MooncakeStoreConnectorMetadata,
@@ -49,19 +51,48 @@ logger = init_logger(__name__)
 
 DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
 DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
+
 MOONCAKE_NO_AVAILABLE_HANDLE = -200
 
+# Mirrors FileStorageConfig::local_buffer_size in Mooncake C++.
+DEFAULT_MOONCAKE_DISK_STAGING_BUFFER_BYTES = 1280 * 1024 * 1024
+
+# Mirrors DirectIO alignment in Mooncake's AllocateBatch.
+_DIRECT_IO_ALIGNMENT = 4096
+_DIRECT_IO_PADDING_BYTES = 2 * _DIRECT_IO_ALIGNMENT
+
+
+MooncakeMode = Literal["embedded", "standalone-store"]
+
 
 @dataclass
 class MooncakeStoreConfig:
-    """Configuration for MooncakeDistributedStore."""
+    """Configuration for MooncakeDistributedStore.
+
+    ``mode`` selects the topology: ``embedded`` (each rank contributes
+    ``global_segment_size`` in-process) or ``standalone-store`` (rank
+    contributes 0; an external ``mooncake_client`` process owns the pool
+    and the SSD tier).
+    """
 
     metadata_server: str
-    global_segment_size: int
-    local_buffer_size: int
+    master_server_address: str
     protocol: str
     device_name: str
-    master_server_address: str
+    mode: MooncakeMode = "embedded"
+    global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE
+    local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE
+    enable_offload: bool = False
+
+    def __post_init__(self) -> None:
+        if self.mode not in ("embedded", "standalone-store"):
+            raise ValueError(f"unknown Mooncake mode: {self.mode!r}")
+        if self.local_buffer_size <= 0:
+            raise ValueError("local_buffer_size must be > 0")
+        if self.mode == "embedded" and self.global_segment_size == 0:
+            raise ValueError("embedded mode requires global_segment_size > 0")
+        if self.mode == "standalone-store" and self.global_segment_size != 0:
+            raise ValueError("standalone-store mode requires global_segment_size == 0")
 
     @staticmethod
     def from_file(file_path: str) -> "MooncakeStoreConfig":
@@ -69,15 +100,17 @@ class MooncakeStoreConfig:
             config = json.load(file)
         return MooncakeStoreConfig(
             metadata_server=config.get("metadata_server", ""),
+            master_server_address=config.get("master_server_address", ""),
+            protocol=config.get("protocol", "rdma"),
+            device_name=config.get("device_name", ""),
+            mode=config.get("mode", "embedded"),
             global_segment_size=_parse_size(
                 config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
             ),
             local_buffer_size=_parse_size(
                 config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)
             ),
-            protocol=config.get("protocol", "rdma"),
-            device_name=config.get("device_name", ""),
-            master_server_address=config.get("master_server_address", ""),
+            enable_offload=bool(config.get("enable_offload", False)),
         )
 
     @staticmethod
@@ -125,6 +158,158 @@ def _parse_size(value: Any) -> int:
     return int(numeric_value * multiplier)
 
 
+def _align_up(value: int, alignment: int) -> int:
+    return ((value + alignment - 1) // alignment) * alignment
+
+
+def _estimate_disk_offload_staging_bytes(size_list: list[int]) -> int:
+    data_size = sum(size_list)
+    return _align_up(data_size, _DIRECT_IO_ALIGNMENT) + _DIRECT_IO_PADDING_BYTES
+
+
+def _get_usable_disk_offload_buffer_budget_bytes(raw_budget_bytes: int) -> int:
+    return max(1, int(raw_budget_bytes * envs.VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO))
+
+
+def _split_disk_offload_load_batches(
+    keys: list[str],
+    addrs: list[list[int]],
+    sizes: list[list[int]],
+    usable_budget_bytes: int,
+    raw_budget_bytes: int,
+) -> tuple[list[tuple[list[str], list[list[int]], list[list[int]]]], str | None]:
+    """Split a GET into sub-batches that fit the owner's staging buffer.
+
+    ``addrs[i]`` / ``sizes[i]`` are scatter-gather lists (K/V or multi-layer
+    segments) for key ``i``. ``usable_budget_bytes`` caps a multi-key batch;
+    ``raw_budget_bytes`` is the hard per-key cap.
+
+    Returns ``(batches, oversize_key)``. Aborts with ``([], key)`` if any
+    single key exceeds ``raw_budget_bytes``; otherwise ``oversize_key`` is
+    ``None``.
+    """
+    batches: list[tuple[list[str], list[list[int]], list[list[int]]]] = []
+    batch_keys: list[str] = []
+    batch_addrs: list[list[int]] = []
+    batch_sizes: list[list[int]] = []
+    batch_bytes = 0
+
+    for key, addr, size in zip(keys, addrs, sizes, strict=True):
+        key_bytes = _estimate_disk_offload_staging_bytes(size)
+        if key_bytes > raw_budget_bytes:
+            return [], key
+        if key_bytes > usable_budget_bytes:
+            if batch_keys:
+                batches.append((batch_keys, batch_addrs, batch_sizes))
+                batch_keys, batch_addrs, batch_sizes = [], [], []
+                batch_bytes = 0
+            batches.append(([key], [addr], [size]))
+            continue
+        if batch_keys and batch_bytes + key_bytes > usable_budget_bytes:
+            batches.append((batch_keys, batch_addrs, batch_sizes))
+            batch_keys, batch_addrs, batch_sizes = [], [], []
+            batch_bytes = 0
+        batch_keys.append(key)
+        batch_addrs.append(addr)
+        batch_sizes.append(size)
+        batch_bytes += key_bytes
+
+    if batch_keys:
+        batches.append((batch_keys, batch_addrs, batch_sizes))
+    return batches, None
+
+
+def _call_replica_predicate(replica_desc: Any, method_name: str) -> bool:
+    method = getattr(replica_desc, method_name, None)
+    if method is None:
+        return False
+    try:
+        return bool(method())
+    except Exception:
+        return False
+
+
+def _classify_replica_tier(replica_descs: Any) -> str:
+    if not replica_descs:
+        return "unknown"
+    try:
+        replica_desc = replica_descs[0]
+    except (IndexError, KeyError, TypeError):
+        return "unknown"
+
+    if _call_replica_predicate(replica_desc, "is_memory_replica"):
+        return "memory"
+    if _call_replica_predicate(
+        replica_desc, "is_disk_replica"
+    ) or _call_replica_predicate(replica_desc, "is_local_disk_replica"):
+        return "disk"
+    return "unknown"
+
+
+def _get_replica_tiers_by_key(store: Any, keys: list[str]) -> dict[str, str]:
+    tiers_by_key = {key: "unknown" for key in keys}
+    try:
+        replica_descs_by_key = store.batch_get_replica_desc(keys)
+    except Exception as e:
+        logger.warning(
+            "Failed to get Mooncake replica descriptors for tier logging "
+            "(batch_keys=%d, error=%s); marking tiers unknown",
+            len(keys),
+            e,
+        )
+        return tiers_by_key
+
+    for key in keys:
+        if hasattr(replica_descs_by_key, "get"):
+            replica_descs = replica_descs_by_key.get(key)
+        else:
+            try:
+                replica_descs = replica_descs_by_key[key]
+            except (KeyError, TypeError):
+                replica_descs = None
+        tiers_by_key[key] = _classify_replica_tier(replica_descs)
+    return tiers_by_key
+
+
+def _log_mooncake_load_tier_summary(
+    req_id: str,
+    batch_keys: list[str],
+    load_results: list[int],
+    tiers_by_key: dict[str, str],
+) -> None:
+    tier_counts = {"memory": 0, "disk": 0, "unknown": 0}
+    bytes_by_tier = {"memory": 0, "disk": 0, "unknown": 0}
+    success_keys = 0
+    failed_keys = 0
+
+    for index, key in enumerate(batch_keys):
+        tier = tiers_by_key.get(key, "unknown")
+        if tier not in tier_counts:
+            tier = "unknown"
+        tier_counts[tier] += 1
+
+        value = load_results[index] if index < len(load_results) else -1
+        if value >= 0:
+            success_keys += 1
+            bytes_by_tier[tier] += int(value)
+        else:
+            failed_keys += 1
+
+    logger.info(
+        "Mooncake load tier summary: req_id=%s batch_keys=%d "
+        "memory_keys=%d disk_keys=%d unknown_keys=%d "
+        "success_keys=%d failed_keys=%d bytes_by_tier=%s",
+        req_id,
+        len(batch_keys),
+        tier_counts["memory"],
+        tier_counts["disk"],
+        tier_counts["unknown"],
+        success_keys,
+        failed_keys,
+        bytes_by_tier,
+    )
+
+
 # ============================================================
 # Transfer Threads
 # ============================================================
@@ -207,6 +392,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
         kv_role: str,
         ready_event: threading.Event,
         enable_kv_event: bool = False,
+        replicate_config: Any = None,
     ):
         super().__init__(
             store,
@@ -220,8 +406,11 @@ class KVCacheStoreSendingThread(KVTransferThread):
         self.kv_role = kv_role
         self.stored_requests: defaultdict[str, int] = defaultdict(int)
         self.enable_kv_event = enable_kv_event
+        # Caller always passes a non-None ReplicateConfig — see
+        # MooncakeStoreWorker.__init__ where store_replicate_config is built.
+        self.replicate_config = replicate_config
 
-        # Pause store requests when CPU offloading is under pressure.
+        # Pause store requests when CPU/disk offloading is under pressure.
         self._store_pressure_active = False
         self._skip_store_requests: set[str] = set()
 
@@ -270,7 +459,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
             return
         if self._should_skip_request(req_id):
             logger.debug(
-                "Skipping Mooncake store for request %s while CPU offloading "
+                "Skipping Mooncake store for request %s while CPU/disk offloading "
                 "is under pressure",
                 req_id,
             )
@@ -357,7 +546,12 @@ class KVCacheStoreSendingThread(KVTransferThread):
             current_event.synchronize()
 
         try:
-            res = self.store.batch_put_from_multi_buffers(keys, addrs, sizes)
+            res = self.store.batch_put_from_multi_buffers(
+                keys,
+                addrs,
+                sizes,
+                self.replicate_config,
+            )
             failed = [i for i, v in enumerate(res) if v < 0]
             if failed:
                 # Compute total bytes attempted for this batch
@@ -379,7 +573,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
                     and not self._mark_request_skipped_for_pressure(req_id)
                 ):
                     logger.warning(
-                        "Detected Mooncake CPU offloading pressure "
+                        "Detected Mooncake CPU/disk offloading pressure "
                         "(NO_AVAILABLE_HANDLE); skipping future store "
                         "batches for request %s until a later store "
                         "batch succeeds",
@@ -387,7 +581,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
                     )
             elif self._clear_store_pressure():
                 logger.info(
-                    "Mooncake CPU offloading pressure cleared after a "
+                    "Mooncake CPU/disk offloading pressure cleared after a "
                     "successful store batch"
                 )
         except Exception as e:
@@ -410,6 +604,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
         block_size: int,
         tp_rank: int,
         ready_event: threading.Event,
+        disk_offload_buffer_budget_bytes: int | None = None,
     ):
         super().__init__(
             store,
@@ -419,6 +614,14 @@ class KVCacheStoreRecvingThread(KVTransferThread):
             ready_event,
             name="KVCacheStoreRecvingThread",
         )
+        self.disk_offload_buffer_budget_bytes = disk_offload_buffer_budget_bytes
+        self.usable_disk_offload_buffer_budget_bytes = (
+            None
+            if disk_offload_buffer_budget_bytes is None
+            else _get_usable_disk_offload_buffer_budget_bytes(
+                disk_offload_buffer_budget_bytes
+            )
+        )
 
     def _handle_request(self, req_meta: ReqMeta):
         token_len = req_meta.load_spec.token_len  # type: ignore[union-attr]
@@ -456,26 +659,69 @@ class KVCacheStoreRecvingThread(KVTransferThread):
             + size_list[: self.tp_rank % len(size_list)]
         )
 
-        try:
-            res = self.store.batch_get_into_multi_buffers(
-                key_list_c, addr_list_c, size_list_c
+        load_batches = [(key_list_c, addr_list_c, size_list_c)]
+        if self.usable_disk_offload_buffer_budget_bytes is not None:
+            total_staging_bytes = sum(
+                _estimate_disk_offload_staging_bytes(size) for size in size_list_c
             )
-            failed = [
-                (key, value)
-                for key, value in zip(key_list_c, res, strict=True)
-                if value < 0
-            ]
-            if failed:
-                logger.warning(
-                    "Failed to get %d Mooncake keys (batch_keys=%d, first_failures=%s)",
-                    len(failed),
-                    len(key_list_c),
-                    failed[:3],
+            if total_staging_bytes > self.usable_disk_offload_buffer_budget_bytes:
+                assert self.disk_offload_buffer_budget_bytes is not None
+                load_batches, oversized_key = _split_disk_offload_load_batches(
+                    key_list_c,
+                    addr_list_c,
+                    size_list_c,
+                    self.usable_disk_offload_buffer_budget_bytes,
+                    self.disk_offload_buffer_budget_bytes,
                 )
+                if oversized_key is not None:
+                    oversized_key_index = key_list_c.index(oversized_key)
+                    oversized_key_bytes = _estimate_disk_offload_staging_bytes(
+                        size_list_c[oversized_key_index]
+                    )
+                    logger.warning(
+                        "Skipping Mooncake load for request %s because key %s "
+                        "requires %d staging bytes, exceeding budget %d",
+                        req_id,
+                        oversized_key,
+                        oversized_key_bytes,
+                        self.disk_offload_buffer_budget_bytes,
+                    )
+                    self.set_finished_request(req_id)
+                    self.request_queue.task_done()
+                    return
+
+        current_batch_keys: list[str] = key_list_c
+        try:
+            for batch_keys, batch_addrs, batch_sizes in load_batches:
+                current_batch_keys = batch_keys
+                tiers_by_key: dict[str, str] | None = None
+                if envs.VLLM_MOONCAKE_STORE_TIER_LOG:
+                    tiers_by_key = _get_replica_tiers_by_key(self.store, batch_keys)
+                res = self.store.batch_get_into_multi_buffers(
+                    batch_keys, batch_addrs, batch_sizes
+                )
+                if tiers_by_key is not None:
+                    _log_mooncake_load_tier_summary(
+                        req_id, batch_keys, res, tiers_by_key
+                    )
+                failed = [
+                    (key, value)
+                    for key, value in zip(batch_keys, res, strict=True)
+                    if value < 0
+                ]
+                if failed:
+                    logger.warning(
+                        "Failed to get %d Mooncake keys from sub-batch "
+                        "(batch_keys=%d, first_failures=%s)",
+                        len(failed),
+                        len(batch_keys),
+                        failed[:3],
+                    )
+                    break
         except Exception as e:
             logger.warning(
-                "Failed to get Mooncake batch %s, error: %s",
-                key_list_c[:3],
+                "Failed to get Mooncake sub-batch %s, error: %s",
+                current_batch_keys[:3],
                 e,
             )
 
@@ -493,7 +739,10 @@ class MooncakeStoreWorker:
 
     def __init__(self, vllm_config: VllmConfig):
         try:
-            from mooncake.store import MooncakeDistributedStore  # type: ignore
+            from mooncake.store import (  # type: ignore
+                MooncakeDistributedStore,
+                ReplicateConfig,
+            )
         except ImportError as e:
             raise ImportError(
                 "Please install mooncake by following the instructions at "
@@ -561,24 +810,79 @@ class MooncakeStoreWorker:
 
         # Initialize MooncakeDistributedStore with its own TransferEngine
         store_config = MooncakeStoreConfig.load_from_env()
+        extra_config = (
+            vllm_config.kv_transfer_config.kv_connector_extra_config
+            if vllm_config.kv_transfer_config
+            else {}
+        )
+        store_config.device_name = rdma_utils.get_configured_worker_rnic(
+            protocol=store_config.protocol,
+            configured_device=store_config.device_name,
+        )
         self.store = MooncakeDistributedStore()
-
-        local_seg = get_ip()
-        config_dict = {
-            "local_hostname": local_seg,
-            "metadata_server": store_config.metadata_server,
-            "global_segment_size": str(store_config.global_segment_size),
-            "local_buffer_size": str(store_config.local_buffer_size),
-            "protocol": store_config.protocol,
-            "rdma_devices": store_config.device_name,
-            "master_server_addr": store_config.master_server_address,
-        }
-        ret = self.store.setup(config_dict)
+        local_ip = get_ip()
+        local_hostname = rdma_utils.get_requester_local_hostname(local_ip)
+        ret = self.store.setup(
+            local_hostname,
+            store_config.metadata_server,
+            store_config.global_segment_size,
+            store_config.local_buffer_size,
+            store_config.protocol,
+            store_config.device_name,
+            store_config.master_server_address,
+        )
         if ret != 0:
             msg = "Initialize MooncakeDistributedStore failed."
             logger.error(msg)
             raise RuntimeError(msg)
 
+        preferred_segment = rdma_utils.get_configured_preferred_segment(extra_config)
+        self.preferred_segment = preferred_segment
+        self.store_replicate_config = ReplicateConfig()
+        if preferred_segment is not None:
+            self.store_replicate_config.preferred_segment = preferred_segment
+
+        logger.info(
+            "Mooncake mode=%s (global_segment_size=%d, local_buffer_size=%d, "
+            "preferred_segment=%s, enable_offload=%s)",
+            store_config.mode,
+            store_config.global_segment_size,
+            store_config.local_buffer_size,
+            preferred_segment or "<none>",
+            store_config.enable_offload,
+        )
+        if store_config.mode == "embedded":
+            if store_config.enable_offload and preferred_segment is None:
+                logger.warning(
+                    "enable_offload is set in embedded mode without "
+                    "preferred_segment; SSD tier will only see puts that "
+                    "happen to land on the owner segment."
+                )
+            if preferred_segment is not None:
+                logger.warning(
+                    "preferred_segment=%s with mode=embedded: rank-"
+                    "contributed segments will be idle.",
+                    preferred_segment,
+                )
+        elif (
+            store_config.mode == "standalone-store" and not store_config.enable_offload
+        ):
+            logger.warning(
+                "standalone-store mode without enable_offload: large prefills "
+                "may exceed the owner DirectIO budget."
+            )
+
+        self.disk_offload_buffer_budget_bytes = (
+            DEFAULT_MOONCAKE_DISK_STAGING_BUFFER_BYTES
+            if store_config.enable_offload
+            else None
+        )
+
+        # Start lookup server on rank 0 for scheduler-side prefix queries
+        self.lookup_server: LookupKeyServer | None = None
+        if vllm_config.parallel_config.rank == 0:
+            self.lookup_server = LookupKeyServer(self, vllm_config)
+
         kv_event_config = vllm_config.kv_events_config
         self.enable_kv_events = False
         if kv_event_config and kv_event_config.enable_kv_cache_events:
@@ -588,11 +892,6 @@ class MooncakeStoreWorker:
         self.kv_recv_thread: KVCacheStoreRecvingThread | None = None
         self.finished_store_req: set[str] = set()
 
-        # Start lookup server on rank 0 for scheduler-side prefix queries
-        self.lookup_server: LookupKeyServer | None = None
-        if vllm_config.parallel_config.rank == 0:
-            self.lookup_server = LookupKeyServer(self, vllm_config)
-
     def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
         """Register a cross-layers KV cache tensor.
 
@@ -695,6 +994,7 @@ class MooncakeStoreWorker:
                 self.kv_role,
                 ready_event_sending,
                 self.enable_kv_events,
+                self.store_replicate_config,
             )
             self.kv_send_thread.start()
 
@@ -705,6 +1005,7 @@ class MooncakeStoreWorker:
             self.block_size,
             self.tp_rank,
             ready_event_recving,
+            disk_offload_buffer_budget_bytes=self.disk_offload_buffer_budget_bytes,
         )
         self.kv_recv_thread.start()
         ready_event_recving.wait()
@@ -967,13 +1268,15 @@ class LookupKeyClient:
 
 def get_zmq_rpc_path_lookup(vllm_config: VllmConfig) -> str:
     """Construct IPC path for ZMQ lookup socket."""
+    assert vllm_config.kv_transfer_config is not None
     dp_rank = get_mooncake_dp_engine_index(vllm_config.parallel_config)
     base_url = envs.VLLM_RPC_BASE_PATH
     rpc_port = 0
-    assert vllm_config.kv_transfer_config is not None
+    hostname = socket.gethostname()
     extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
     if "lookup_rpc_port" in extra_config:
         rpc_port = extra_config["lookup_rpc_port"]
-    uid = os.getuid()
-    logger.debug("Base URL: %s, RPC Port: %s, UID: %s", base_url, rpc_port, uid)
-    return f"ipc://{base_url}/lookup_rpc_port_{rpc_port}_uid{uid}_dp_rank{dp_rank}"
+    logger.debug("Base URL: %s, RPC Port: %s", base_url, rpc_port)
+    return (
+        f"ipc://{base_url}/lookup_rpc_port_{rpc_port}_host_{hostname}_dp_rank{dp_rank}"
+    )

vllm/envs.py

@@ -191,6 +191,10 @@ if TYPE_CHECKING:
     VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
     VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600
     VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998
+    VLLM_MOONCAKE_STORE_TIER_LOG: bool = False
+    VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO: float = 0.9
+    MOONCAKE_PREFERRED_SEGMENT: str | None = None
+    MOONCAKE_REQUESTER_LOCAL_HOSTNAME: str | None = None
     VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840
     VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1
     VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
@@ -1389,6 +1393,20 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_MOONCAKE_BOOTSTRAP_PORT": lambda: int(
         os.getenv("VLLM_MOONCAKE_BOOTSTRAP_PORT", "8998")
     ),
+    # Log per-batch memory/disk tier breakdown on external GETs.
+    "VLLM_MOONCAKE_STORE_TIER_LOG": lambda: (
+        os.getenv("VLLM_MOONCAKE_STORE_TIER_LOG", "False").lower() in ("true", "1")
+    ),
+    # Fraction of the owner's DirectIO staging buffer to fill per GET batch.
+    "VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO": lambda: float(
+        os.getenv("VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO", "0.9")
+    ),
+    # Pin this rank to a specific owner segment ("host:port").
+    "MOONCAKE_PREFERRED_SEGMENT": lambda: os.getenv("MOONCAKE_PREFERRED_SEGMENT"),
+    # Override the hostname the rank registers as a Mooncake requester.
+    "MOONCAKE_REQUESTER_LOCAL_HOSTNAME": lambda: os.getenv(
+        "MOONCAKE_REQUESTER_LOCAL_HOSTNAME"
+    ),
     # Flashinfer MoE backend for vLLM's fused Mixture-of-Experts support.
     # Both require compute capability 10.0 or above.
     # Available options: