A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 23:45 UTC
[KV Connector] Update lmcache kv_offloading_backend to use LMCacheMPConnector (#42865)
Signed-off-by: baoloongmao <baoloongmao@tencent.com>
maobaolong committed
Jun 4, 2026 at 10:23 UTC
b58e082d95ffad57a6a9aaffa8b76c862b3bbcf3
2 files changed
+47
-18
tests/v1/kv_connector/unit/test_config.py
+41
-8
@@ -6,25 +6,56 @@
6
import pytest
7
8
from vllm.config import CacheConfig, KVTransferConfig, ParallelConfig, VllmConfig
9
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
10
11
pytestmark = pytest.mark.cpu_test
12
13
14
+class _StubLMCacheMPConnector:
15
+ """Stand-in for LMCacheMPConnector used in config-translation tests.
16
+
17
+ The real connector module hard-imports the optional ``lmcache`` package
18
+ at module load time, which is not installed in the cpu_test image. This
19
+ test only asserts on the connector *name* and the ``extra_config`` dict
20
+ produced by ``VllmConfig``, never instantiates the connector, so a bare
21
+ placeholder class is sufficient. Not subclassing ``SupportsHMA`` mirrors
22
+ the real connector's HMA support (it does not support HMA either)."""
23
+
24
+
25
+@pytest.fixture
26
+def stub_lmcache_mp_connector(monkeypatch):
27
+ """Replace the lazy loader so VllmConfig.__post_init__ does not import
28
+ ``lmcache_mp_connector`` (and thus ``lmcache``) during config tests."""
29
+ monkeypatch.setitem(
30
+ KVConnectorFactory._registry,
31
+ "LMCacheMPConnector",
32
+ lambda: _StubLMCacheMPConnector,
33
+ )
34
+
35
+
36
@pytest.mark.parametrize(
37
"kv_offloading_backend,kv_offloading_size,tp,pp,expected_backend,expected_bytes",
38
[
39
("native", 4.0, 1, 1, "OffloadingConnector", 4.0 * (1 << 30)),
40
# bytes per rank: 8.0 GiB / (2 * 2) = 2.0 GiB
41
("native", 8.0, 2, 2, "OffloadingConnector", 8.0 * (1 << 30)),
19
- ("lmcache", 4.0, 1, 1, "LMCacheConnectorV1", 4.0),
20
- # size per rank: 8.0 GiB / (2 * 2) = 2.0 GiB
21
- ("lmcache", 8.0, 2, 2, "LMCacheConnectorV1", 2.0),
42
+ # ``lmcache`` backend now defaults to LMCacheMPConnector. The KV
43
+ # storage capacity is owned by the standalone LMCache server, so
44
+ # ``kv_offloading_size`` is intentionally not propagated.
45
+ ("lmcache", 4.0, 1, 1, "LMCacheMPConnector", None),
46
+ ("lmcache", 8.0, 2, 2, "LMCacheMPConnector", None),
47
# When kv_offloading_size is None, offloading is disabled (backend is ignored)
48
("native", None, 1, 1, None, None),
49
],
50
)
51
def test_kv_connector(
27
- kv_offloading_backend, kv_offloading_size, tp, pp, expected_backend, expected_bytes
52
+ stub_lmcache_mp_connector,
53
+ kv_offloading_backend,
54
+ kv_offloading_size,
55
+ tp,
56
+ pp,
57
+ expected_backend,
58
+ expected_bytes,
59
):
60
kv_transfer_config = (
61
KVTransferConfig(kv_connector_extra_config={"existing_key": "existing_value"})
@@ -59,10 +90,12 @@ def test_kv_connector(
90
# Existing config should be preserved
91
assert kv_connector_extra_config["existing_key"] == "existing_value"
92
elif kv_offloading_backend == "lmcache":
62
- assert kv_connector_extra_config["lmcache.local_cpu"] is True
63
- assert kv_connector_extra_config["lmcache.max_local_cpu_size"] == expected_bytes
64
- # Existing config should be replaced
65
- assert "existing_key" not in kv_connector_extra_config
93
+ # MP mode does not push lmcache.local_cpu / max_local_cpu_size into
94
+ # extra config (the LMCache server owns capacity). Pre-existing
95
+ # extra config entries are preserved as-is.
96
+ assert "lmcache.local_cpu" not in kv_connector_extra_config
97
+ assert "lmcache.max_local_cpu_size" not in kv_connector_extra_config
98
+ assert kv_connector_extra_config["existing_key"] == "existing_value"
99
100
101
def _build_config(
vllm/config/vllm.py
+6
-10
@@ -780,10 +780,6 @@ class VllmConfig:
780
# If no KVTransferConfig is provided, create a default one.
781
if self.kv_transfer_config is None:
782
self.kv_transfer_config = KVTransferConfig()
783
- num_kv_ranks = (
784
- self.parallel_config.tensor_parallel_size
785
- * self.parallel_config.pipeline_parallel_size
786
- )
783
784
if kv_offloading_backend == "native":
785
if envs.VLLM_USE_SIMPLE_KV_OFFLOAD:
@@ -795,12 +791,12 @@ class VllmConfig:
791
{"cpu_bytes_to_use": kv_offloading_size * (1 << 30)}
792
)
793
elif kv_offloading_backend == "lmcache":
798
- self.kv_transfer_config.kv_connector = "LMCacheConnectorV1"
799
- kv_gb_per_rank = kv_offloading_size / num_kv_ranks
800
- self.kv_transfer_config.kv_connector_extra_config = {
801
- "lmcache.local_cpu": True,
802
- "lmcache.max_local_cpu_size": kv_gb_per_rank,
803
- }
794
+ # Default to LMCache multi-process (MP) mode. The actual KV
795
+ # storage capacity is managed by the standalone LMCache server
796
+ # process, so ``kv_offloading_size`` is not propagated here.
797
+ # ``LMCacheMPConnector`` falls back to ``tcp://localhost:5555``
798
+ # when host/port are not provided via extra_config.
799
+ self.kv_transfer_config.kv_connector = "LMCacheMPConnector"
800
801
# This is the same for all backends
802
self.kv_transfer_config.kv_role = "kv_both"