diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index dab0faaf98..4402095b2e 100644 --- a/docs/features/nixl_connector_usage.md +++ b/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 diff --git a/tests/v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh index 48c446e019..2779806dc9 100755 --- a/tests/v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh +++ b/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[@]}") diff --git a/tests/v1/kv_connector/mooncake_integration/run_accuracy_test.sh b/tests/v1/kv_connector/mooncake_integration/run_accuracy_test.sh index 6436ccdb13..8290bd655b 100755 --- a/tests/v1/kv_connector/mooncake_integration/run_accuracy_test.sh +++ b/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 diff --git a/tests/v1/kv_connector/mooncake_integration/test_accuracy.py b/tests/v1/kv_connector/mooncake_integration/test_accuracy.py index aa4c72faf4..0168471bd1 100644 --- a/tests/v1/kv_connector/mooncake_integration/test_accuracy.py +++ b/tests/v1/kv_connector/mooncake_integration/test_accuracy.py @@ -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}," diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index 019b8d1504..ad9bd18567 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/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( diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index f79d7aa316..2cbb9cc670 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -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]) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index d6e834861a..39015d4c32 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -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. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index 71593f59e5..8997571407 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -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) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 41c1337ecd..bc70d7cba8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/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 @@ 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() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index dd64dd671e..a1e28b97d0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 0e6eccd8ea..b81ace6f9f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 1ded62adda..bc59f3cb27 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -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) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 6e43356c0d..5511edc580 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -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: diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 4fcac2d6c1..aaab37a496 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -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: