[KV Connector] Support sleep mode with MooncakeConnector over RDMA (#59625)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: Dakai An <dakaian108@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Dakai An <dakaian108@gmail.com> Co-authored-by: Dakai An <77474977+andakai@users.noreply.github.com>
2026-10-14 04:18:47 +00:00
parent f4b8d43a68
commit 3ccc0a90b6
14 changed files with 205 additions and 8 deletions
@@ -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

@@ -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[@]}")
@@ -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
@@ -4,8 +4,13 @@ import os

 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},"
@@ -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(
@@ -232,6 +232,9 @@ class FakeMooncakeWrapper:
    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])
@@ -1199,6 +1199,21 @@ class VllmConfig:
        ):
            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.
@@ -56,6 +56,7 @@ from vllm.v1.outputs import KVConnectorOutput

 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 @@ class KVConnectorBase_V1(ABC):
        """
        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 @@ class KVConnectorBase_V1(ABC):
        """
        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)
@@ -3,6 +3,7 @@
 import asyncio
 import itertools
 import logging
+import os
 import queue
 import threading
 import time
@@ -21,7 +22,7 @@ import zmq.asyncio
 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 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
    ############################################################
    # 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 @@ class MooncakeConnectorWorker:
        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 @@ class MooncakeConnectorWorker:

        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 @@ class MooncakeConnectorWorker:
                    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()
@@ -18,7 +18,7 @@ from typing import Any

 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 @@ class MooncakeStoreKVEvents(KVConnectorKVEvents):
 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
@@ -16,7 +16,7 @@ import numpy as np
 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
@@ -222,6 +222,17 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
    def requires_kv_delivery(self) -> bool:
        return any(c.requires_kv_delivery for c in self._connectors)

+    @classmethod
+    def supports_sleep_mode(cls, kv_transfer_config: "KVTransferConfig") -> bool:
+        children = [
+            KVTransferConfig(**{"engine_id": kv_transfer_config.engine_id, **c})
+            for c in kv_transfer_config.kv_connector_extra_config.get("connectors", [])
+        ]
+        return bool(children) and all(
+            KVConnectorFactory.get_connector_class(c).supports_sleep_mode(c)
+            for c in children
+        )
+
    @classmethod
    def _get_connector_classes_and_configs(
        cls, vllm_config: "VllmConfig"
@@ -252,6 +263,14 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
        for c in self._connectors:
            c.register_kv_caches(kv_caches)

+    def unregister_kv_caches(self) -> None:
+        for c in self._connectors:
+            c.unregister_kv_caches()
+
+    def restore_kv_caches(self) -> None:
+        for c in self._connectors:
+            c.restore_kv_caches()
+
    def bind_kv_cache_manager(self, kv_cache_manager: "KVCacheManager") -> None:
        for connector in self._connectors:
            connector.bind_kv_cache_manager(kv_cache_manager)
@@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Any

 import torch

-from vllm.config import VllmConfig
+from vllm.config import KVTransferConfig, VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.utils import (
    EngineId,
 )
@@ -140,6 +140,10 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
    ############################################################
    # Class Methods
    ############################################################
+    @classmethod
+    def supports_sleep_mode(cls, kv_transfer_config: KVTransferConfig) -> bool:
+        return False
+
    @classmethod
    def get_required_kvcache_layout(cls, vllm_config: VllmConfig):
        if vllm_config.model_config is None:
@@ -308,6 +308,8 @@ class Worker(WorkerBase):
                    name: buffer.cpu().clone() for name, buffer in draft.named_buffers()
                }

+        if has_kv_transfer_group():
+            get_kv_transfer_group().unregister_kv_caches()
        self.sleep_mode_backend.suspend(level)
        if self.vllm_config.model_config.enable_nccl_comm_suspend:
            suspend_device_comms()
@@ -331,6 +333,8 @@ class Worker(WorkerBase):

    def wake_up(self, tags: list[str] | None = None) -> None:
        self.sleep_mode_backend.resume(tags)
+        if has_kv_transfer_group() and (tags is None or "kv_cache" in tags):
+            get_kv_transfer_group().restore_kv_caches()
        if self.vllm_config.model_config.enable_nccl_comm_suspend:
            resume_device_comms()

@@ -356,6 +360,8 @@ class Worker(WorkerBase):
        self.synchronize_device()

    def discard(self, tags: tuple[str, ...]) -> None:
+        if has_kv_transfer_group() and "kv_cache" in tags:
+            get_kv_transfer_group().unregister_kv_caches()
        self.sleep_mode_backend.discard(tags)

    def checkpoint_prepare(self) -> None: