From 36e74c9ea4feb5ade38ffa1ea96f24dd73316e02 Mon Sep 17 00:00:00 2001 From: Zhewen Li Date: Sat, 16 May 2026 13:34:15 -0700 Subject: [PATCH] [KV Connector] Support disk offloading in MooncakeStoreConnector (#42689) Signed-off-by: Zhewen Li Co-authored-by: Zhewen Li Co-authored-by: Claude Opus 4.7 (1M context) --- .../mooncake_store_connector_usage.md | 65 +- .../unit/test_mooncake_store_worker.py | 835 ++++++++++++++++++ .../kv_connector/v1/mooncake/rdma_utils.py | 115 +++ .../kv_connector/v1/mooncake/store/worker.py | 407 +++++++-- vllm/envs.py | 18 + 5 files changed, 1385 insertions(+), 55 deletions(-) create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/mooncake/rdma_utils.py diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md index 3520cd6e872..7fcbe1cefa7 100644 --- a/docs/features/mooncake_store_connector_usage.md +++ b/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 diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index b808e148045..c8b9437616a 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/rdma_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/rdma_utils.py new file mode 100644 index 00000000000..9ee0d2cc542 --- /dev/null +++ b/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 "" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 487542c5917..3e832042971 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/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 "", + 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}" + ) diff --git a/vllm/envs.py b/vllm/envs.py index 9d3542c1bee..bdf7298188a 100755 --- a/vllm/envs.py +++ b/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: