@tduncan / Nvidia-Optimization-Stack-VLLM / commits / eb0f2ca37f

[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>

void committed Sep 27, 2026 at 08:26 UTC eb0f2ca37f65e496e8c67a7497e818c8a92bae46

7 files changed +193 -4

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+60

@@ -290,6 +290,8 @@ class _FakeKVTransferConfig:

290 class _FakeModelConfig:

291 model = "test-model"

292 use_mla = False

293 + enable_sleep_mode = False

294 + enable_cumem_allocator = False

295

296 def get_num_layers(self, parallel_config) -> int:

297 return 1

@@ -432,6 +434,64 @@ def _patch_worker_runtime(

434 monkeypatch.setattr(worker, "LookupKeyServer", MagicMock())

435

436

437 +def test_create_mem_pool_returns_none_when_unconfigured():

438 + assert worker.MooncakeStoreWorker._create_mem_pool({}, _FakeModelConfig()) is None

439 +

440 +

441 +def test_create_mem_pool_rejects_unknown_pool():

442 + with pytest.raises(

443 + ValueError,

444 + match="Unsupported custom_mem_pool='UNKNOWN'",

445 + ):

446 + worker.MooncakeStoreWorker._create_mem_pool(

447 + {"custom_mem_pool": "unknown"}, _FakeModelConfig()

448 + )

449 +

450 +

451 +@pytest.mark.parametrize(

452 + "conflicting_option", ["enable_sleep_mode", "enable_cumem_allocator"]

453 +)

454 +def test_create_mem_pool_rejects_cumem_conflicts(conflicting_option):

455 + model_config = _FakeModelConfig()

456 + setattr(model_config, conflicting_option, True)

457 +

458 + with pytest.raises(ValueError, match="custom_mem_pool is incompatible"):

459 + worker.MooncakeStoreWorker._create_mem_pool(

460 + {"custom_mem_pool": "NVLINK"}, model_config

461 + )

462 +

463 +

464 +@pytest.mark.parametrize(

465 + ("pool_type", "allocator_name"),

466 + [("nvlink", "NVLinkAllocator"), ("barex", "BarexAllocator")],

467 +)

468 +def test_create_mem_pool_uses_mooncake_allocator(

469 + monkeypatch, pool_type, allocator_name

470 +):

471 + allocator_cls = MagicMock()

472 + allocator = allocator_cls.get_allocator.return_value

473 + allocator_module = types.ModuleType("mooncake.allocator")

474 + setattr(allocator_module, allocator_name, allocator_cls)

475 + mooncake_module = types.ModuleType("mooncake")

476 + mooncake_module.allocator = allocator_module # type: ignore[attr-defined]

477 + monkeypatch.setitem(sys.modules, "mooncake", mooncake_module)

478 + monkeypatch.setitem(sys.modules, "mooncake.allocator", allocator_module)

479 +

480 + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 2)

481 + mem_pool = MagicMock()

482 + mem_pool_ctor = MagicMock(return_value=mem_pool)

483 + monkeypatch.setattr(torch.cuda, "MemPool", mem_pool_ctor)

484 +

485 + result = worker.MooncakeStoreWorker._create_mem_pool(

486 + {"custom_mem_pool": pool_type}, _FakeModelConfig()

487 + )

488 +

489 + assert result is mem_pool

490 + allocator_cls.get_allocator.assert_called_once_with(torch.device("cuda", 2))

491 + allocator.allocator.assert_called_once_with()

492 + mem_pool_ctor.assert_called_once_with(allocator.allocator.return_value)

493 +

494 +

495 def test_pool_key_to_string_without_prefix_is_unchanged():

496 """Default (empty) cache_prefix keeps keys byte-identical to the

497 historical unprefixed format so existing deployments keep their hits."""

tests/v1/kv_connector/unit/test_multi_connector.py

+32

@@ -68,6 +68,7 @@ class MockConnector(KVConnectorBase_V1):

68 # Override just build_kv_connector_stats

69 mock.build_kv_connector_stats = cls.build_kv_connector_stats

70 mock.get_kv_connector_stats.return_value = None

71 + mock.get_mem_pool_context.return_value = None

72 return mock

74 @classmethod

@@ -201,6 +202,36 @@ def mc() -> MultiConnector:

202 return mc

203

204

205 +def test_multi_connector_mem_pool_context_none(mc: MultiConnector):

206 + assert mc.get_mem_pool_context() is None

207 +

208 +

209 +def test_multi_connector_forwards_mem_pool_context(mc: MultiConnector):

210 + context = MagicMock()

211 + provider = MagicMock(spec_set=KVConnectorBase_V1)

212 + provider.get_mem_pool_context.return_value = context

213 + mc._connectors = [mc._connectors[0], provider]

214 +

215 + assert mc.get_mem_pool_context() is context

216 + provider.get_mem_pool_context.assert_called_once_with()

217 +

218 +

219 +def test_multi_connector_rejects_multiple_mem_pool_contexts(mc: MultiConnector):

220 + providers = [

221 + MagicMock(spec_set=KVConnectorBase_V1),

222 + MagicMock(spec_set=KVConnectorBase_V1),

223 + ]

224 + for provider in providers:

225 + provider.get_mem_pool_context.return_value = MagicMock()

226 + mc._connectors = providers

227 +

228 + with pytest.raises(

229 + ValueError,

230 + match="Multiple connectors provide a KV cache memory pool",

231 + ):

232 + mc.get_mem_pool_context()

233 +

234 +

235 # Helper function to compare directories recursively

236 def _compare_directories(dir1: Path, dir2: Path) -> bool:

237 """Compares two directories recursively for identical content."""

@@ -320,6 +351,7 @@ def test_multi_example_connector_consistency():

351 # First three events are from initialization. Layer hooks run before the

352 # deferred load starts after the forward pass.

353 expected_worker_prefix = [

354 + "get_mem_pool_context",

355 "register_kv_caches",

356 "set_host_xfer_buffer_ops",

357 "get_handshake_metadata",

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -42,6 +42,7 @@ The class provides the following primitives:

42 import enum

43 from abc import ABC, abstractmethod

44 from collections.abc import Callable, Iterable

45 +from contextlib import AbstractContextManager

46 from dataclasses import dataclass, field

47 from typing import TYPE_CHECKING, Any, Literal

@@ -227,6 +228,13 @@ class KVConnectorBase_V1(ABC):

228 # Worker-side methods

229 # ==============================

230

231 + def get_mem_pool_context(self) -> AbstractContextManager | None:

232 + """Return a custom KV cache allocation context, if configured.

233 +

234 + Returning None uses the engine's default memory pool.

235 + """

236 + return None

237 +

238 def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:

239 """Set the connector metadata from the scheduler.

240

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

+12

@@ -13,6 +13,7 @@ enabling prefix caching via hash-based deduplication.

14 from collections import Counter

15 from collections.abc import Iterable, Sequence

16 +from contextlib import AbstractContextManager

17 from typing import Any

19 import torch

@@ -322,6 +323,17 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

323 # Worker-side methods

324 # ============================================================

325

326 + def get_mem_pool_context(self) -> AbstractContextManager | None:

327 + """Return a context manager for the custom MemPool, or None.

328 +

329 + Called by the Worker before ``initialize_kv_cache`` so that KV

330 + cache is allocated from the Mooncake-managed pool when

331 + ``custom_mem_pool`` is set in ``kv_connector_extra_config``.

332 + """

333 + if self.connector_worker is None:

334 + return None

335 + return self.connector_worker.get_mem_pool_context()

336 +

337 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

338 assert self.connector_worker is not None

339 self.connector_worker.register_kv_caches(kv_caches)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+48 -1

@@ -19,13 +19,14 @@ import threading

19 import time

20 from collections.abc import Callable, Sequence

21 from concurrent.futures import Future, ThreadPoolExecutor

22 +from contextlib import AbstractContextManager

23 from typing import Any, TypeVar

25 import torch

26 import zmq

28 import vllm.envs as envs

28 -from vllm.config import VllmConfig

29 +from vllm.config import ModelConfig, VllmConfig

30 from vllm.distributed import (

31 get_dcp_group,

32 get_pcp_group,

@@ -1414,6 +1415,41 @@ class KVCacheStoreRecvingThread(KVTransferThread):

1415 class MooncakeStoreWorker:

1416 """Worker-side component for MooncakeStoreConnector."""

1417

1418 + @staticmethod

1419 + def _create_mem_pool(

1420 + extra_config: dict[str, Any], model_config: ModelConfig

1421 + ) -> torch.cuda.MemPool | None:

1422 + pool_type = str(extra_config.get("custom_mem_pool") or "").upper()

1423 + if not pool_type:

1424 + return None

1425 + if pool_type not in ("NVLINK", "BAREX"):

1426 + raise ValueError(

1427 + f"Unsupported custom_mem_pool={pool_type!r}, "

1428 + "expected 'NVLINK' or 'BAREX'"

1429 + )

1430 + if model_config.enable_sleep_mode or model_config.enable_cumem_allocator:

1431 + raise ValueError(

1432 + "custom_mem_pool is incompatible with enable_sleep_mode "

1433 + "or enable_cumem_allocator; CuMemAllocator cannot manage "

1434 + "allocations from the custom pool."

1435 + )

1436 + try:

1437 + if pool_type == "NVLINK":

1438 + from mooncake.allocator import NVLinkAllocator as allocator_cls

1439 + else: # pool_type == "BAREX"

1440 + from mooncake.allocator import BarexAllocator as allocator_cls

1441 + except ImportError as e:

1442 + raise ImportError(

1443 + f"custom_mem_pool={pool_type!r} requires "

1444 + "mooncake-transfer-engine>=0.3.8. Please upgrade Mooncake."

1445 + ) from e

1446 +

1447 + device = torch.device("cuda", torch.accelerator.current_device_index())

1448 + allocator = allocator_cls.get_allocator(device)

1449 + mem_pool = torch.cuda.MemPool(allocator.allocator())

1450 + logger.info("Using Mooncake custom memory pool: %s", pool_type)

1451 + return mem_pool

1452 +

1453 def __init__(

1454 self,

1455 vllm_config: VllmConfig,

@@ -1471,6 +1507,10 @@ class MooncakeStoreWorker:

1507

1508 # Initialize MooncakeDistributedStore with its own TransferEngine

1509 store_config = MooncakeStoreConfig.load_from_config()

1510 + # Validate and create the custom memory pool before initializing the

1511 + # store so failures cannot leak a partially initialized store handle.

1512 + self._mem_pool = self._create_mem_pool(extra_config, model_config)

1513 +

1514 self.store = MooncakeDistributedStore()

1515 local_ip = get_ip()

1516 local_hostname = rdma_utils.get_requester_local_hostname(local_ip)

@@ -1741,6 +1781,13 @@ class MooncakeStoreWorker:

1781 )

1782 return token_dbs

1783

1784 + def get_mem_pool_context(self) -> AbstractContextManager | None:

1785 + """Return a context manager for the custom MemPool, or None if

1786 + no custom pool is configured."""

1787 + if self._mem_pool is None:

1788 + return None

1789 + return torch.cuda.use_mem_pool(self._mem_pool)

1790 +

1791 def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int:

1792 if self.dcp_size > 1:

1793 return 1

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

+21

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import copy

4 from collections.abc import Callable, Iterable

5 +from contextlib import AbstractContextManager

6 from dataclasses import dataclass

7 from typing import TYPE_CHECKING, Any, cast

@@ -294,6 +295,26 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):

295 # ==============================

296 # Worker-side methods

297 # ==============================

298 + def get_mem_pool_context(self) -> AbstractContextManager | None:

299 + """Forward a custom KV cache memory pool from a child connector.

300 +

301 + KV cache is allocated once and can only live in one pool, so at

302 + most one child may provide a context.

303 + """

304 + found: list[tuple[str, AbstractContextManager]] = []

305 + for connector in self._connectors:

306 + if (ctx := connector.get_mem_pool_context()) is not None:

307 + found.append((type(connector).__name__, ctx))

308 +

309 + if len(found) > 1:

310 + names = [name for name, _ in found]

311 + raise ValueError(

312 + f"Multiple connectors provide a KV cache memory pool {names}; "

313 + "KV cache is allocated once and can only live in one pool. "

314 + "Configure custom_mem_pool on at most one connector."

315 + )

316 + return found[0][1] if found else None

317 +

318 def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:

319 for c in self._connectors:

320 c.start_load_kv(forward_context, **kwargs)

vllm/v1/worker/gpu_worker.py

+12 -3

@@ -801,11 +801,20 @@ class Worker(WorkerBase):

801 # related to kv cache connector (e.g. kv cache sharing layers).

802 ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)

803

804 + # If the connector provides a custom memory pool (e.g. Mooncake

805 + # NVLink/BAREX), use it for KV cache allocation; otherwise fall

806 + # back to the standard CuMem pool.

807 + mem_pool_context = (

808 + get_kv_transfer_group().get_mem_pool_context()

809 + if has_kv_transfer_group()

810 + else None

811 + )

812 + if mem_pool_context is None:

813 + mem_pool_context = self._maybe_get_memory_pool_context(tag="kv_cache")

814 +

815 self.model_runner.initialize_kv_cache(

816 kv_cache_config,

806 - kv_cache_allocation_context=self._maybe_get_memory_pool_context(

807 - tag="kv_cache"

808 - ),

817 + kv_cache_allocation_context=mem_pool_context,

818 )

819

820 # Build KV-zero metadata outside the CuMem pool so the bookkeeping