[mooncake] support CUSTOM_MEM_POOL in vllm (#49300)

Signed-off-by: bruce.xu <bruce.xb@alibaba-inc.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: bruce.xu <bruce.xb@alibaba-inc.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>

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."""

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",

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.
 

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)

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

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)

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