[KV Connector] Support sleep mode with the NIXL pull connector

.buildkite/test_areas/disaggregated.yaml

@@ -14,6 +14,7 @@ steps:
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+    - GPU_MEMORY_UTILIZATION=0.6 SLEEP_WAKE=1 bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
   mirror:
     amd:
       label: ":amd: (MI355) Distributed NixlConnector PD accuracy"

docs/features/nixl_connector_usage.md

@@ -400,6 +400,8 @@ The `kv_load_failure_policy` setting controls how the system handles failures wh
 
 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.
 
+With `--enable-sleep-mode` (pull mode and the UCX backend only), the worker sets `UCX_RCACHE_ENABLE=n` and `UCX_TLS=^cuda_ipc` unless they are already set, since UCX's registration cache and a peer's cuda_ipc import would keep the slept KV cache resident. P/D on one node then transfers through the NIC instead of NVLink. Routers must forward `remote_registration_epoch` in `kv_transfer_params`.
+
 ## Experimental Feature
 
 ### Heterogeneous KV Layout support

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

@@ -85,6 +85,12 @@ DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
 ENFORCE_EAGER=${ENFORCE_EAGER:-1}
 # 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
 # Pin concurrent prefiller and non-DP decoder engines to separate internal
 # port windows. DP decoder ranks retain their existing internal port selection.
 PREFILLER_INTERNAL_PORT_BASE=${PREFILLER_INTERNAL_PORT_BASE:-20000}
@@ -313,9 +319,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 python3 -m pytest -s -x "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_accuracy.py
+  TEST_MODEL=$model_name SLEEP_WAKE_URLS=$SLEEP_WAKE_URLS \
+    python3 -m pytest -s -x "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_accuracy.py
 
   # Clean up before running next model
   cleanup_instances

tests/v1/kv_connector/nixl_integration/test_accuracy.py

@@ -4,8 +4,13 @@
 
 import lm_eval
 import openai
+import requests
 
 BASE_URL = "http://localhost: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 = int(os.getenv("NUM_CONCURRENT", "100"))
 TASK = "gsm8k"
 FILTER = "exact_match,strict-match"
@@ -61,9 +66,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)
 
     if "gemma-4" in MODEL_NAME:
         # Gemma4 is quite sensible to having a chat template applied, so we evaluate

tests/v1/kv_connector/unit/test_config.py

@@ -160,7 +160,8 @@ def test_kv_connector_allows_expandable_segments_with_cumem_allocator(
 @pytest.mark.parametrize(
     ("kv_connector", "extra_config", "env"),
     [
-        ("NixlConnector", {}, None),
+        ("NixlConnector", {"backends": ["LIBFABRIC"]}, None),
+        ("NixlPushConnector", {}, None),
         ("MooncakeConnector", {"mooncake_protocol": "tcp"}, None),
         ("MooncakeConnector", {}, "MC_USE_TENT"),
     ],

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -2192,6 +2192,7 @@ def recv_worker():
     worker._engine_by_address = {}
     worker._failed_remote_engines = set()
     worker._invalid_remote_engines = set()
+    worker._remote_registration_epochs = {}
     worker._replicated_pcp_done_sending = set()
     worker._invalid_block_ids = queue.Queue()
     worker._pending_recv_notifs = {}
@@ -2729,6 +2730,52 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
         mock_dereg.assert_any_call("desc2")
 
 
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
+    FakeNixlWrapper,
+)
+def test_sleep_unregisters_and_restores_kv_caches(dist_init):
+    """Sleep drops everything built on the KV cache registration; wake-up
+    registers the same memory again and publishes the new agent metadata."""
+    vllm_config = create_vllm_config()
+    vllm_config.cache_config.kv_cache_layout = "BLHNC"
+    vllm_config.kv_transfer_config.kv_buffer_device = "cpu"
+    spec = FullAttentionSpec(
+        block_size=16, num_kv_heads=4, head_size=16, dtype=torch.float16
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=2,
+        kv_cache_tensors=[],
+        kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)],
+    )
+    raw = torch.zeros(spec.page_size_bytes * 2, dtype=torch.int8)
+    (cache,) = dense_kv_cache_views(
+        raw, spec, 2, num_layers=1, layout=KVCacheLayout.BLHNC
+    )
+    with set_current_vllm_config(vllm_config):
+        connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
+        connector.register_kv_caches({"layer0": cache})
+    worker = connector.connector_worker
+    nixl = worker.nixl_wrapper = MagicMock(wraps=worker.nixl_wrapper)
+    worker._remote_agents["remote"] = {(0, 0): "agent"}
+    descs = list(worker._registered_descs)
+
+    connector.unregister_kv_caches()
+    connector.unregister_kv_caches()  # a discard, then a sleep
+    assert nixl.deregister_memory.call_args_list == [call(d) for d in descs]
+    assert not worker._remote_agents and not worker.src_xfer_handles_by_block_size
+
+    nixl.get_agent_metadata.return_value = b"re-registered"
+    connector.restore_kv_caches()
+    connector.restore_kv_caches()
+    assert nixl.register_memory.call_count == len(descs)
+    assert worker.block_size in worker.src_xfer_handles_by_block_size
+    payload = worker.xfer_handshake_metadata
+    agent = msgspec.msgpack.decode(payload.agent_metadata_bytes, type=NixlAgentMetadata)
+    assert agent.agent_metadata == b"re-registered"
+    connector.shutdown()
+
+
 # ── TTL-based remote engine eviction tests ──────────────────────────
 
 
@@ -3270,6 +3317,18 @@ def test_eviction_clears_pending_recovery(self):
         assert not self.worker._failed_remote_engines
         assert not self.worker._invalid_remote_engines
 
+    @pytest.mark.parametrize("loaded", [True, False])
+    def test_rehandshakes_re_registered_peer(self, loaded):
+        if not loaded:  # our own sleep already dropped the agent
+            self.worker._cleanup_remote_engine("old", log_eviction=False)
+            self.transport.remove_remote_agent.reset_mock()
+        metadata = self._request("old")
+        metadata.reqs_to_recv["old-req"].remote.registration_epoch = 1
+        self.worker.start_load_kv(metadata)
+        self.transport.make_prepped_xfer.assert_not_called()
+        assert self.transport.remove_remote_agent.called is loaded
+        assert "old" in self.worker._handshake_futures
+
 
 def test_transfer_topology_unregister():
     """TransferTopology.unregister_remote_engine removes the engine."""

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

@@ -113,6 +113,9 @@ def __init__(
         # Background thread for handling new handshake requests.
         self._nixl_handshake_listener_t: threading.Thread | None = None
         self._stop_event = threading.Event()
+        # Payloads the listener serves; a republish after a wake-up updates them.
+        self._handshake_data: dict[tuple[int, int], bytes] = {}
+        self._registration_epoch = 0
 
         # Requests that need to start recv/send.
         # New requests are added by update_state_after_alloc in
@@ -299,20 +302,22 @@ def set_xfer_handshake_metadata(
             metadata (dict): the handshake metadata to set.
 
         """
-        encoded_data: dict[tuple[int, int], bytes] = {}
+        if self._handshake_data:
+            # A wake-up re-registered the KV cache: peers must handshake again.
+            self._registration_epoch += 1
         encoder = msgspec.msgpack.Encoder()
         for (pp_rank, tp_rank), rank_metadata in metadata.items():
             if not isinstance(rank_metadata, NixlHandshakePayload):
                 raise ValueError(
                     "NixlConnectorScheduler expects NixlHandshakePayload for "
                     "handshake metadata."
                 )
-            encoded_data[(pp_rank, tp_rank)] = encoder.encode(rank_metadata)
+            self._handshake_data[(pp_rank, tp_rank)] = encoder.encode(rank_metadata)
             logger.debug(
                 "PP rank %d, TP rank %d: encoded NixlHandshakePayload size: %s bytes",
                 pp_rank,
                 tp_rank,
-                str(len(encoded_data[(pp_rank, tp_rank)])),
+                str(len(self._handshake_data[(pp_rank, tp_rank)])),
             )
 
         # Only start the listener when we have metadata to serve.
@@ -321,7 +326,7 @@ def set_xfer_handshake_metadata(
             self._nixl_handshake_listener_t = threading.Thread(
                 target=self._nixl_handshake_listener,
                 args=(
-                    encoded_data,
+                    self._handshake_data,
                     ready_event,
                     self._stop_event,
                     self.side_channel_host,

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -641,6 +641,11 @@ def __init__(
                 else nixl_agent_config(num_threads=num_threads, capture_telemetry=True)
             )
 
+        if vllm_config.model_config.enable_sleep_mode:
+            # UCX's registration cache and peers' cuda_ipc imports would keep the
+            # slept KV cache resident. UCX reads both when the agent is created.
+            os.environ.setdefault("UCX_RCACHE_ENABLE", "n")
+            os.environ.setdefault("UCX_TLS", "^cuda_ipc")
         self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config)
         # Map of engine_id -> {(pp_rank, tp_rank): agent_name, ...}.
         # non-PP remote uses pp_rank 0, i.e. (0, tp_rank).
@@ -800,6 +805,7 @@ def __init__(
         self.dst_uses_region_group_mapping: dict[EngineId, bool] = {}
         self.dst_region_mem_types: dict[EngineId, list[str]] = {}
         self._registered_descs: list[Any] = []
+        self._kv_unregistered = False
 
         # In progress transfers.
         # [req_id -> list[handle]]
@@ -1825,6 +1831,35 @@ def _ple_page_len(self) -> int:
         )
         return self._ple_block_len
 
+    def unregister_kv_caches(self) -> None:
+        if self._kv_unregistered or not self._registered_descs:
+            return
+        self._release_local_xfer_handles()
+        # Remote engines hold local handles too; they handshake again on demand.
+        for engine_id in list(self._remote_agents):
+            self._cleanup_remote_engine(engine_id, log_eviction=False)
+        for descs in self._registered_descs:
+            self.nixl_wrapper.deregister_memory(descs)
+        self._kv_unregistered = True
+
+    def restore_kv_caches(self) -> None:
+        if not self._kv_unregistered:
+            return
+        for descs in self._registered_descs:
+            self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
+        self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = (
+            self.register_local_xfer_handler(self.block_size)
+        )
+        # Peers load the new registration from the agent metadata.
+        payload = self.xfer_handshake_metadata
+        assert payload is not None
+        agent_metadata = msgspec.msgpack.decode(
+            payload.agent_metadata_bytes, type=NixlAgentMetadata
+        )
+        agent_metadata.agent_metadata = self.nixl_wrapper.get_agent_metadata()
+        payload.agent_metadata_bytes = msgspec.msgpack.encode(agent_metadata)
+        self._kv_unregistered = False
+
     def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray:
         """Build desc regions (conv sub-projections + ssm) per layer for
         local mamba blocks with DS conv layout, as an Nx3 uint64 array.
@@ -3729,8 +3764,7 @@ def __del__(self):
         with contextlib.suppress(Exception):
             self.shutdown()
 
-    def _finish_shutdown(self) -> None:
-        self._recving_transfers.clear()
+    def _release_local_xfer_handles(self) -> None:
         try:
             for handle in self.src_xfer_handles_by_block_size.values():
                 self.nixl_wrapper.release_dlist_handle(handle)
@@ -3744,19 +3778,25 @@ def _finish_shutdown(self) -> None:
             for dram_handle in self._dram_src_handles_by_block_size.values():
                 if dram_handle is not None:
                     self.nixl_wrapper.release_dlist_handle(dram_handle)
+        finally:
+            self.src_xfer_handles_by_block_size.clear()
+            self.src_xfer_handles_by_tp_ratio.clear()
+            self._dram_src_handles_by_tp_ratio.clear()
+            self._dram_src_handles_by_block_size.clear()
+
+    def _finish_shutdown(self) -> None:
+        self._recving_transfers.clear()
+        try:
+            self._release_local_xfer_handles()
         except Exception:
             logger.exception("NIXL dlist-handle release failed at shutdown.")
-        self.src_xfer_handles_by_block_size.clear()
-        self.src_xfer_handles_by_tp_ratio.clear()
-        self._dram_src_handles_by_tp_ratio.clear()
-        self._dram_src_handles_by_block_size.clear()
         try:
             for engine_id in list(self._remote_agents):
                 self._cleanup_remote_engine(engine_id, log_eviction=False)
         except Exception:
             logger.exception("NIXL remote-engine cleanup failed at shutdown.")
         try:
-            for desc in self._registered_descs:
+            for desc in [] if self._kv_unregistered else self._registered_descs:
                 self.nixl_wrapper.deregister_memory(desc)
         finally:
             self._registered_descs.clear()

vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py

@@ -142,6 +142,7 @@ def __init__(
     ############################################################
     @classmethod
     def supports_sleep_mode(cls, kv_transfer_config: KVTransferConfig) -> bool:
+        # Push producers keep the decoder's KV cache registration across its sleep.
         return False
 
     @classmethod
@@ -367,6 +368,19 @@ def __init__(
                 vllm_config, self.engine_id, kv_cache_config
             )
 
+    @classmethod
+    def supports_sleep_mode(cls, kv_transfer_config: KVTransferConfig) -> bool:
+        # Sleep mode is only validated with the UCX backend.
+        return kv_transfer_config.get_from_extra_config("backends", ["UCX"]) == ["UCX"]
+
+    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 start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
         assert self.connector_worker is not None
         assert isinstance(self.connector_worker, NixlPullConnectorWorker)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py

@@ -50,8 +50,9 @@
 #  12: Add per-region member names for PP push
 #  13: Add packed-member layouts and order-independent packed-push backend hashes
 #  14: Add Qwen3.8-Flash-Next PLE page length - #59997
+#  15: Add remote_registration_epoch to kv_transfer_params
 #
-NIXL_CONNECTOR_VERSION: int = 14
+NIXL_CONNECTOR_VERSION: int = 15
 
 
 @dataclass
@@ -236,6 +237,8 @@ class RemoteMeta:
     request_id: str
     blocks_expiry_time: float | None = None
     num_tokens: int | None = None
+    # Bumped each time the remote re-registers its KV cache after a sleep.
+    registration_epoch: int = 0
 
 
 @dataclass
@@ -336,5 +339,6 @@ def add_new_req_to_recv(
             port=kv_transfer_params["remote_port"],
             blocks_expiry_time=kv_transfer_params.get("remote_blocks_expiry_time"),
             num_tokens=kv_transfer_params.get("remote_num_tokens"),
+            registration_epoch=kv_transfer_params.get("remote_registration_epoch", 0),
         )
         self.reqs_to_recv[request_id] = req

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py

@@ -288,5 +288,6 @@ def request_finished(
             pp_size=self.vllm_config.parallel_config.pipeline_parallel_size,
             remote_num_tokens=remote_num_tokens,
             remote_blocks_expiry_time=blocks_expiry_time,
+            remote_registration_epoch=self._registration_epoch,
             transfer_mode=self._TRANSFER_MODE,
         )