[KV Connector] Support sleep mode with MooncakeConnector over RDMA

docs/features/nixl_connector_usage.md

@@ -398,7 +398,7 @@ The `kv_load_failure_policy` setting controls how the system handles failures wh
 
 ### For NVIDIA GB-series GPUs
 
-GB-series GPUs support multi-node NVLink. NIXL supports this capability, but KVCache must be registered as VMM during KVCache registration. To enable this feature, you need to set `--enable-cumem-allocator` or `--enable-sleep-mode` flags, and set `UCX_CUDA_IPC_ENABLE_MNNVL: 'y'` env. Otherwise, NIXL can only use RDMA/TCP for cross-node KVCache transfers.
+GB-series GPUs support multi-node NVLink. NIXL supports this capability, but KVCache must be registered as VMM during KVCache registration. To enable this feature, you need to set the `--enable-cumem-allocator` flag, and set `UCX_CUDA_IPC_ENABLE_MNNVL: 'y'` env. Otherwise, NIXL can only use RDMA/TCP for cross-node KVCache transfers.
 
 ## Experimental Feature
 

tests/v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh

@@ -13,6 +13,7 @@ python3 -m pytest -s -x "${IMPORT_CANARY}"
 tp_configs=(
   "GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=1"
   "GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2"
+  "GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=1 SLEEP_WAKE=1"
 )
 
 configs=("${tp_configs[@]}")

tests/v1/kv_connector/mooncake_integration/run_accuracy_test.sh

@@ -59,6 +59,12 @@ GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2}
 BLOCK_SIZE=${BLOCK_SIZE:-128}
 # Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
 VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
+# Set to 1 to put every instance to sleep and wake it up before the eval.
+SLEEP_WAKE=${SLEEP_WAKE:-0}
+if [[ "$SLEEP_WAKE" == "1" ]]; then
+  export VLLM_SERVER_DEV_MODE=1
+  VLLM_SERVE_EXTRA_ARGS="${VLLM_SERVE_EXTRA_ARGS:+${VLLM_SERVE_EXTRA_ARGS},}--enable-sleep-mode"
+fi
 
 # Resolve the repository root from the script location instead of `.git`.
 SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
@@ -226,9 +232,17 @@ run_tests_for_model() {
   # Wait for the proxy to start
   sleep 5
 
+  SLEEP_WAKE_URLS=""
+  if [[ "$SLEEP_WAKE" == "1" ]]; then
+    for PORT in "${PREFILL_PORTS[@]}" "${DECODE_PORTS[@]}"; do
+      SLEEP_WAKE_URLS="${SLEEP_WAKE_URLS:+${SLEEP_WAKE_URLS},}http://localhost:${PORT}"
+    done
+  fi
+
   # Run lm eval for this model
   echo "Running tests for $model_name"
-  TEST_MODEL=$model_name PROXY_PORT=$PROXY_PORT python3 -m pytest -s -x \
+  TEST_MODEL=$model_name PROXY_PORT=$PROXY_PORT SLEEP_WAKE_URLS=$SLEEP_WAKE_URLS \
+    python3 -m pytest -s -x \
     "${GIT_ROOT}/tests/v1/kv_connector/mooncake_integration/test_accuracy.py"
 
   # Clean up before running next model

tests/v1/kv_connector/mooncake_integration/test_accuracy.py

@@ -4,8 +4,13 @@
 
 import lm_eval
 import openai
+import requests
 
 BASE_URL = f"http://localhost:{os.environ.get('PROXY_PORT', 8192)}/v1"
+# Comma-separated server URLs to sleep and wake up before the eval.
+SLEEP_WAKE_URLS = [
+    url for url in os.environ.get("SLEEP_WAKE_URLS", "").split(",") if url
+]
 NUM_CONCURRENT = 100
 TASK = "gsm8k"
 FILTER = "exact_match,strict-match"
@@ -35,9 +40,18 @@ def run_simple_prompt():
     print("-" * 50)
 
 
+def sleep_and_wake_up(urls: list[str]):
+    for url in urls:
+        requests.post(f"{url}/sleep", params={"level": 1}).raise_for_status()
+    for url in urls:
+        requests.post(f"{url}/wake_up").raise_for_status()
+
+
 def test_accuracy():
     """Run the end to end accuracy test."""
     run_simple_prompt()
+    if SLEEP_WAKE_URLS:
+        sleep_and_wake_up(SLEEP_WAKE_URLS)
 
     model_args = (
         f"model={MODEL_NAME},"

tests/v1/kv_connector/unit/test_config.py

@@ -101,6 +101,7 @@ def test_kv_connector(
 def _build_config(
     *,
     kv_connector: str | None,
+    kv_connector_extra_config: dict | None = None,
     enable_sleep_mode: bool = False,
     enable_cumem_allocator: bool = False,
 ) -> VllmConfig:
@@ -109,7 +110,11 @@ def _build_config(
     from types import SimpleNamespace
 
     kv_transfer_config = (
-        KVTransferConfig(kv_connector=kv_connector, kv_role="kv_both")
+        KVTransferConfig(
+            kv_connector=kv_connector,
+            kv_role="kv_both",
+            kv_connector_extra_config=kv_connector_extra_config or {},
+        )
         if kv_connector is not None
         else None
     )
@@ -141,7 +146,7 @@ def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch):
     """Sleep mode routes KV allocations through CuMemAllocator's pool, which
     auto-disables expandable_segments (see #40812)."""
     monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-    _build_config(kv_connector="NixlConnector", enable_sleep_mode=True)
+    _build_config(kv_connector="MooncakeConnector", enable_sleep_mode=True)
 
 
 def test_kv_connector_allows_expandable_segments_with_cumem_allocator(
@@ -152,6 +157,27 @@ def test_kv_connector_allows_expandable_segments_with_cumem_allocator(
     _build_config(kv_connector="NixlConnector", enable_cumem_allocator=True)
 
 
+@pytest.mark.parametrize(
+    ("kv_connector", "extra_config", "env"),
+    [
+        ("NixlConnector", {}, None),
+        ("MooncakeConnector", {"mooncake_protocol": "tcp"}, None),
+        ("MooncakeConnector", {}, "MC_USE_TENT"),
+    ],
+)
+def test_sleep_mode_refuses_unsupported_kv_connector(
+    monkeypatch, kv_connector, extra_config, env
+):
+    if env:
+        monkeypatch.setenv(env, "1")
+    with pytest.raises(ValueError, match=f"^{kv_connector} does not support sleep"):
+        _build_config(
+            kv_connector=kv_connector,
+            kv_connector_extra_config=extra_config,
+            enable_sleep_mode=True,
+        )
+
+
 def test_kv_connector_allows_other_alloc_conf(monkeypatch):
     """Other PYTORCH_CUDA_ALLOC_CONF values must not be rejected."""
     monkeypatch.setenv(

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -232,6 +232,9 @@ def batch_transfer_sync_write(
     def batch_register_memory(self, buffer_addresses, capacities) -> int:
         return 0
 
+    def batch_unregister_memory(self, buffer_addresses) -> int:
+        return 0
+
 
 def test_align_transfer_regions_uses_layer_name_occurrences():
     """Repeated layer names should align by occurrence order."""
@@ -2574,3 +2577,29 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
 
         prefill_worker.sender_loop = origin_sender_loop
         prefill_worker.shutdown()
+
+
+def test_sleep_unregisters_and_restores_kv_caches():
+    spec = FullAttentionSpec(
+        block_size=maybe_update_block_size(16),
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+    )
+    layer_names = ["model.layers.0.self_attn", "model.layers.1.self_attn"]
+    with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, reg):
+        for layer_name in layer_names:
+            worker._layer_specs[layer_name] = spec
+        raw = torch.zeros(2 * 2 * spec.page_size_bytes, dtype=torch.int8)
+        views = dense_kv_cache_views(raw, spec, 2, 2, KVCacheLayout.LBHNC)
+        worker.register_kv_caches(dict(zip(layer_names, views)))
+        reg.reset_mock()
+        with patch.object(
+            worker.engine, "batch_unregister_memory", return_value=0
+        ) as unreg:
+            worker.unregister_kv_caches()
+            worker.unregister_kv_caches()
+        unreg.assert_called_once_with([raw.data_ptr()])
+        worker.restore_kv_caches()
+        worker.restore_kv_caches()
+        reg.assert_called_once_with([raw.data_ptr()], [raw.nbytes])

vllm/config/vllm.py

@@ -1199,6 +1199,21 @@ def _verify_kv_transfer_compat(self) -> None:
         ):
             return
 
+        if self.model_config is not None and self.model_config.enable_sleep_mode:
+            # Lazy import to avoid circular dependencies
+            from vllm.distributed.kv_transfer.kv_connector.factory import (
+                KVConnectorFactory,
+            )
+
+            connector_cls = KVConnectorFactory.get_connector_class(
+                self.kv_transfer_config
+            )
+            if not connector_cls.supports_sleep_mode(self.kv_transfer_config):
+                raise ValueError(
+                    f"{self.kv_transfer_config.kv_connector} does not support "
+                    "sleep mode."
+                )
+
         # PyTorch's expandable_segments allocator uses CUDA VMM, which can
         # remap a virtual address range to different physical pages over the
         # engine's lifetime. KV connectors that pin KV cache memory (e.g.

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -56,6 +56,7 @@
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
+    from vllm.config.kv_transfer import KVTransferConfig
     from vllm.distributed.kv_events import KVCacheEvent, KVConnectorKVEvents
     from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
         KVConnectorPromMetrics,
@@ -205,6 +206,11 @@ def requires_kv_delivery(self) -> bool:
         """
         return self._kv_transfer_config.is_kv_producer
 
+    @classmethod
+    def supports_sleep_mode(cls, kv_transfer_config: "KVTransferConfig") -> bool:
+        """Whether the connector can follow the KV cache remap of sleep mode."""
+        return True
+
     def __init__(
         self,
         vllm_config: "VllmConfig",
@@ -293,6 +299,15 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """
         return
 
+    def unregister_kv_caches(self) -> None:
+        """Drop KV cache registrations before sleep. Transfers are already
+        drained. Idempotent."""
+        return
+
+    def restore_kv_caches(self) -> None:
+        """Redo `unregister_kv_caches` after wake-up. Idempotent."""
+        return
+
     def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
         """Set the xPU-specific ops for copying KV between host and device.
         Needed when host buffer is used for kv transfer (e.g., in NixlConnector)

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

@@ -3,6 +3,7 @@
 import asyncio
 import itertools
 import logging
+import os
 import queue
 import threading
 import time
@@ -21,7 +22,7 @@
 from huggingface_hub.utils import httpx
 
 from vllm import envs
-from vllm.config import VllmConfig
+from vllm.config import KVTransferConfig, VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     EngineId,
     TransferTopology,
@@ -800,10 +801,34 @@ def request_finished_all_groups(
     ############################################################
     # Worker Side Methods
     ############################################################
+    @classmethod
+    def supports_sleep_mode(cls, kv_transfer_config: KVTransferConfig) -> bool:
+        # Sleep mode is only validated over RDMA, not NVLink or TENT.
+        if any(
+            v in os.environ
+            for v in (
+                "MC_FORCE_MNNVL",
+                "MC_INTRANODE_NVLINK",
+                "MC_USE_TENT",
+                "MC_USE_TEV1",
+            )
+        ):
+            return False
+        extra_config = kv_transfer_config.kv_connector_extra_config
+        return extra_config.get("mooncake_protocol", "rdma") == "rdma"
+
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         assert self.connector_worker is not None
         self.connector_worker.register_kv_caches(kv_caches)
 
+    def unregister_kv_caches(self) -> None:
+        assert self.connector_worker is not None
+        self.connector_worker.unregister_kv_caches()
+
+    def restore_kv_caches(self) -> None:
+        assert self.connector_worker is not None
+        self.connector_worker.restore_kv_caches()
+
     def get_finished(
         self, finished_req_ids: set[str]
     ) -> tuple[set[str] | None, set[str] | None]:
@@ -1239,6 +1264,8 @@ def __init__(
         self.region_row_offsets: list[int] = []
         self.opaque_packed_storages: set[int] = set()
         self.seen_base_addresses: list[int] = []
+        self._kv_data_lens: list[int] = []
+        self._kv_unregistered = False
         # Aligned regions depend only on the peer's registered layout.
         # The third item is an error string when alignment cannot proceed.
         self._prepared_transfer_regions: dict[
@@ -2156,6 +2183,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
 
         self.kv_caches_base_addr = region_base_addresses
         self.seen_base_addresses = kv_data_ptrs
+        self._kv_data_lens = kv_data_lens
 
         if not kv_data_ptrs:
             raise RuntimeError("No KV cache tensors were registered with Mooncake.")
@@ -2211,6 +2239,24 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                     f"{ready_timeout:.0f}s."
                 )
 
+    def unregister_kv_caches(self) -> None:
+        if self._kv_unregistered or not self.seen_base_addresses:
+            return
+        ret_value = self.engine.batch_unregister_memory(self.seen_base_addresses)
+        if ret_value != 0:
+            raise RuntimeError("Mooncake batch memory unregistration failed.")
+        self._kv_unregistered = True
+
+    def restore_kv_caches(self) -> None:
+        if not self._kv_unregistered:
+            return
+        ret_value = self.engine.batch_register_memory(
+            self.seen_base_addresses, self._kv_data_lens
+        )
+        if ret_value != 0:
+            raise RuntimeError("Mooncake batch memory registration failed.")
+        self._kv_unregistered = False
+
     async def fetch_finished_recving_reqs(self) -> set[ReqId]:
         finished_recving_reqs = self.finished_recving_reqs
         self.finished_recving_reqs = set()

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

@@ -18,7 +18,7 @@
 
 import torch
 
-from vllm.config import VllmConfig
+from vllm.config import KVTransferConfig, VllmConfig
 from vllm.distributed.kv_events import (
     BlockStored,
     KVCacheEvent,
@@ -127,6 +127,10 @@ def __repr__(self) -> str:
 class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
     """KV connector using MooncakeDistributedStore as shared KV pool."""
 
+    @classmethod
+    def supports_sleep_mode(cls, kv_transfer_config: KVTransferConfig) -> bool:
+        return False
+
     @staticmethod
     def _validate_kv_cache_config(
         vllm_config: VllmConfig, kv_cache_config: KVCacheConfig

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

@@ -16,7 +16,7 @@
 import torch
 import zmq
 
-from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.config import CUDAGraphMode, KVTransferConfig, VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     BlockIds,
     clip_ssm_state_blocks,
@@ -229,6 +229,10 @@ def resolve_moriio_transfer_ack(
 class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
     _cache_hit_source = CacheHitSource.P2P
 
+    @classmethod
+    def supports_sleep_mode(cls, kv_transfer_config: KVTransferConfig) -> bool:
+        return False
+
     @property
     def supports_divergent_local_hybrid_hits(self) -> bool:
         # The READ path always transfers the recurrent-state slot, including