From eb0f2ca37f65e496e8c67a7497e818c8a92bae46 Mon Sep 17 00:00:00 2001 From: void Date: Sun, 27 Sep 2026 08:26:05 +0800 Subject: [PATCH] [mooncake] support CUSTOM_MEM_POOL in vllm (#49300) Signed-off-by: bruce.xu Signed-off-by: Yifan Qiao Co-authored-by: bruce.xu Co-authored-by: Yifan Qiao Co-authored-by: OpenAI Codex --- .../unit/test_mooncake_store_worker.py | 60 +++++++++++++++++++ .../kv_connector/unit/test_multi_connector.py | 32 ++++++++++ .../kv_transfer/kv_connector/v1/base.py | 8 +++ .../v1/mooncake/store/connector.py | 12 ++++ .../kv_connector/v1/mooncake/store/worker.py | 49 ++++++++++++++- .../kv_connector/v1/multi_connector.py | 21 +++++++ vllm/v1/worker/gpu_worker.py | 15 ++++- 7 files changed, 193 insertions(+), 4 deletions(-) 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 947e9f6a20..f08213444e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -290,6 +290,8 @@ class _FakeKVTransferConfig: class _FakeModelConfig: model = "test-model" use_mla = False + enable_sleep_mode = False + enable_cumem_allocator = False def get_num_layers(self, parallel_config) -> int: return 1 @@ -432,6 +434,64 @@ def _patch_worker_runtime( monkeypatch.setattr(worker, "LookupKeyServer", MagicMock()) +def test_create_mem_pool_returns_none_when_unconfigured(): + assert worker.MooncakeStoreWorker._create_mem_pool({}, _FakeModelConfig()) is None + + +def test_create_mem_pool_rejects_unknown_pool(): + with pytest.raises( + ValueError, + match="Unsupported custom_mem_pool='UNKNOWN'", + ): + worker.MooncakeStoreWorker._create_mem_pool( + {"custom_mem_pool": "unknown"}, _FakeModelConfig() + ) + + +@pytest.mark.parametrize( + "conflicting_option", ["enable_sleep_mode", "enable_cumem_allocator"] +) +def test_create_mem_pool_rejects_cumem_conflicts(conflicting_option): + model_config = _FakeModelConfig() + setattr(model_config, conflicting_option, True) + + with pytest.raises(ValueError, match="custom_mem_pool is incompatible"): + worker.MooncakeStoreWorker._create_mem_pool( + {"custom_mem_pool": "NVLINK"}, model_config + ) + + +@pytest.mark.parametrize( + ("pool_type", "allocator_name"), + [("nvlink", "NVLinkAllocator"), ("barex", "BarexAllocator")], +) +def test_create_mem_pool_uses_mooncake_allocator( + monkeypatch, pool_type, allocator_name +): + allocator_cls = MagicMock() + allocator = allocator_cls.get_allocator.return_value + allocator_module = types.ModuleType("mooncake.allocator") + setattr(allocator_module, allocator_name, allocator_cls) + mooncake_module = types.ModuleType("mooncake") + mooncake_module.allocator = allocator_module # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "mooncake", mooncake_module) + monkeypatch.setitem(sys.modules, "mooncake.allocator", allocator_module) + + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 2) + mem_pool = MagicMock() + mem_pool_ctor = MagicMock(return_value=mem_pool) + monkeypatch.setattr(torch.cuda, "MemPool", mem_pool_ctor) + + result = worker.MooncakeStoreWorker._create_mem_pool( + {"custom_mem_pool": pool_type}, _FakeModelConfig() + ) + + assert result is mem_pool + allocator_cls.get_allocator.assert_called_once_with(torch.device("cuda", 2)) + allocator.allocator.assert_called_once_with() + mem_pool_ctor.assert_called_once_with(allocator.allocator.return_value) + + def test_pool_key_to_string_without_prefix_is_unchanged(): """Default (empty) cache_prefix keeps keys byte-identical to the historical unprefixed format so existing deployments keep their hits.""" diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 1a85865560..9d04a2ea53 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -68,6 +68,7 @@ class MockConnector(KVConnectorBase_V1): # Override just build_kv_connector_stats mock.build_kv_connector_stats = cls.build_kv_connector_stats mock.get_kv_connector_stats.return_value = None + mock.get_mem_pool_context.return_value = None return mock @classmethod @@ -201,6 +202,36 @@ def mc() -> MultiConnector: return mc +def test_multi_connector_mem_pool_context_none(mc: MultiConnector): + assert mc.get_mem_pool_context() is None + + +def test_multi_connector_forwards_mem_pool_context(mc: MultiConnector): + context = MagicMock() + provider = MagicMock(spec_set=KVConnectorBase_V1) + provider.get_mem_pool_context.return_value = context + mc._connectors = [mc._connectors[0], provider] + + assert mc.get_mem_pool_context() is context + provider.get_mem_pool_context.assert_called_once_with() + + +def test_multi_connector_rejects_multiple_mem_pool_contexts(mc: MultiConnector): + providers = [ + MagicMock(spec_set=KVConnectorBase_V1), + MagicMock(spec_set=KVConnectorBase_V1), + ] + for provider in providers: + provider.get_mem_pool_context.return_value = MagicMock() + mc._connectors = providers + + with pytest.raises( + ValueError, + match="Multiple connectors provide a KV cache memory pool", + ): + mc.get_mem_pool_context() + + # Helper function to compare directories recursively def _compare_directories(dir1: Path, dir2: Path) -> bool: """Compares two directories recursively for identical content.""" @@ -320,6 +351,7 @@ def test_multi_example_connector_consistency(): # First three events are from initialization. Layer hooks run before the # deferred load starts after the forward pass. expected_worker_prefix = [ + "get_mem_pool_context", "register_kv_caches", "set_host_xfer_buffer_ops", "get_handshake_metadata", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index bc3c32ba4c..6df9d1cfbe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -42,6 +42,7 @@ The class provides the following primitives: import enum from abc import ABC, abstractmethod from collections.abc import Callable, Iterable +from contextlib import AbstractContextManager from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal @@ -227,6 +228,13 @@ class KVConnectorBase_V1(ABC): # Worker-side methods # ============================== + def get_mem_pool_context(self) -> AbstractContextManager | None: + """Return a custom KV cache allocation context, if configured. + + Returning None uses the engine's default memory pool. + """ + return None + def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None: """Set the connector metadata from the scheduler. 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 eb4a7fafe9..dd64dd671e 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 @@ -13,6 +13,7 @@ enabling prefix caching via hash-based deduplication. from collections import Counter from collections.abc import Iterable, Sequence +from contextlib import AbstractContextManager from typing import Any import torch @@ -322,6 +323,17 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): # Worker-side methods # ============================================================ + def get_mem_pool_context(self) -> AbstractContextManager | None: + """Return a context manager for the custom MemPool, or None. + + Called by the Worker before ``initialize_kv_cache`` so that KV + cache is allocated from the Mooncake-managed pool when + ``custom_mem_pool`` is set in ``kv_connector_extra_config``. + """ + if self.connector_worker is None: + return None + return self.connector_worker.get_mem_pool_context() + 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) 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 baf8967bc8..173aa324e7 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 @@ -19,13 +19,14 @@ import threading import time from collections.abc import Callable, Sequence from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import AbstractContextManager from typing import Any, TypeVar import torch import zmq import vllm.envs as envs -from vllm.config import VllmConfig +from vllm.config import ModelConfig, VllmConfig from vllm.distributed import ( get_dcp_group, get_pcp_group, @@ -1414,6 +1415,41 @@ class KVCacheStoreRecvingThread(KVTransferThread): class MooncakeStoreWorker: """Worker-side component for MooncakeStoreConnector.""" + @staticmethod + def _create_mem_pool( + extra_config: dict[str, Any], model_config: ModelConfig + ) -> torch.cuda.MemPool | None: + pool_type = str(extra_config.get("custom_mem_pool") or "").upper() + if not pool_type: + return None + if pool_type not in ("NVLINK", "BAREX"): + raise ValueError( + f"Unsupported custom_mem_pool={pool_type!r}, " + "expected 'NVLINK' or 'BAREX'" + ) + if model_config.enable_sleep_mode or model_config.enable_cumem_allocator: + raise ValueError( + "custom_mem_pool is incompatible with enable_sleep_mode " + "or enable_cumem_allocator; CuMemAllocator cannot manage " + "allocations from the custom pool." + ) + try: + if pool_type == "NVLINK": + from mooncake.allocator import NVLinkAllocator as allocator_cls + else: # pool_type == "BAREX" + from mooncake.allocator import BarexAllocator as allocator_cls + except ImportError as e: + raise ImportError( + f"custom_mem_pool={pool_type!r} requires " + "mooncake-transfer-engine>=0.3.8. Please upgrade Mooncake." + ) from e + + device = torch.device("cuda", torch.accelerator.current_device_index()) + allocator = allocator_cls.get_allocator(device) + mem_pool = torch.cuda.MemPool(allocator.allocator()) + logger.info("Using Mooncake custom memory pool: %s", pool_type) + return mem_pool + def __init__( self, vllm_config: VllmConfig, @@ -1471,6 +1507,10 @@ class MooncakeStoreWorker: # Initialize MooncakeDistributedStore with its own TransferEngine store_config = MooncakeStoreConfig.load_from_config() + # Validate and create the custom memory pool before initializing the + # store so failures cannot leak a partially initialized store handle. + self._mem_pool = self._create_mem_pool(extra_config, model_config) + self.store = MooncakeDistributedStore() local_ip = get_ip() local_hostname = rdma_utils.get_requester_local_hostname(local_ip) @@ -1741,6 +1781,13 @@ class MooncakeStoreWorker: ) return token_dbs + def get_mem_pool_context(self) -> AbstractContextManager | None: + """Return a context manager for the custom MemPool, or None if + no custom pool is configured.""" + if self._mem_pool is None: + return None + return torch.cuda.use_mem_pool(self._mem_pool) + def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int: if self.dcp_size > 1: return 1 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 a6f0d4cafb..73b6b0ccf6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy from collections.abc import Callable, Iterable +from contextlib import AbstractContextManager from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast @@ -294,6 +295,26 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): # ============================== # Worker-side methods # ============================== + def get_mem_pool_context(self) -> AbstractContextManager | None: + """Forward a custom KV cache memory pool from a child connector. + + KV cache is allocated once and can only live in one pool, so at + most one child may provide a context. + """ + found: list[tuple[str, AbstractContextManager]] = [] + for connector in self._connectors: + if (ctx := connector.get_mem_pool_context()) is not None: + found.append((type(connector).__name__, ctx)) + + if len(found) > 1: + names = [name for name, _ in found] + raise ValueError( + f"Multiple connectors provide a KV cache memory pool {names}; " + "KV cache is allocated once and can only live in one pool. " + "Configure custom_mem_pool on at most one connector." + ) + return found[0][1] if found else None + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: for c in self._connectors: c.start_load_kv(forward_context, **kwargs) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index def92f601d..778d71bdbc 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -801,11 +801,20 @@ class Worker(WorkerBase): # related to kv cache connector (e.g. kv cache sharing layers). ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config) + # If the connector provides a custom memory pool (e.g. Mooncake + # NVLink/BAREX), use it for KV cache allocation; otherwise fall + # back to the standard CuMem pool. + mem_pool_context = ( + get_kv_transfer_group().get_mem_pool_context() + if has_kv_transfer_group() + else None + ) + if mem_pool_context is None: + mem_pool_context = self._maybe_get_memory_pool_context(tag="kv_cache") + self.model_runner.initialize_kv_cache( kv_cache_config, - kv_cache_allocation_context=self._maybe_get_memory_pool_context( - tag="kv_cache" - ), + kv_cache_allocation_context=mem_pool_context, ) # Build KV-zero metadata outside the CuMem pool so the bookkeeping