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 6, 20:30 UTC
[LMCache] Pass TP size in lookup for MLA multi-reader locking (#36129)
Signed-off-by: baoloongmao <baoloongmao@tencent.com> Co-authored-by: Yihua Cheng <yihua98@uchicago.edu>
maobaolong committed
Mar 12, 2026 at 04:45 UTC
12001f2ebc606b471476d47edc22a79af6aca66c
2 files changed
+22
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py
+6
@@ -114,6 +114,7 @@ class LMCacheMPSchedulerAdapter:
114
world_size: int,
115
kv_rank: int,
116
vllm_block_size: int,
117
+ tp_size: int = 1,
118
):
119
"""
120
Args:
@@ -124,6 +125,8 @@ class LMCacheMPSchedulerAdapter:
125
world_size: The world size used for LMCache keys
126
kv_rank: The kv rank used for LMCache keys
127
vllm_block_size: The block size used in vLLM
128
+ tp_size: Tensor-parallel size for MLA
129
+ multi-reader locking (default 1).
130
"""
131
self.mq_client = MessageQueueClient(server_url, context)
132
@@ -133,6 +136,7 @@ class LMCacheMPSchedulerAdapter:
136
self.model_name = model_name
137
self.world_size = world_size
138
self.worker_id = kv_rank
139
+ self.tp_size = tp_size
140
141
# Read chunk size from lmcache
142
self.chunk_size = get_lmcache_chunk_size(self.mq_client)
@@ -281,6 +285,7 @@ class LMCacheMPSchedulerAdapter:
285
start=start,
286
end=end,
287
request_id=request_id,
288
+ tp_size=self.tp_size,
289
)
290
291
def _create_hash_key(
@@ -293,6 +298,7 @@ class LMCacheMPSchedulerAdapter:
298
worker_id=None,
299
chunk_hash=chunk_hash,
300
request_id=request_id,
301
+ tp_size=self.tp_size,
302
)
303
304
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+16
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import enum
4
+import inspect
5
from collections.abc import Iterable
6
from dataclasses import dataclass, field
7
from typing import TYPE_CHECKING, Any, Literal
@@ -52,6 +53,12 @@ if TYPE_CHECKING:
53
logger = lmcache_init_logger(__name__)
54
55
56
+def _adapter_accepts_tp_size() -> bool:
57
+ """Check if the imported adapter accepts tp_size."""
58
+ sig = inspect.signature(LMCacheMPSchedulerAdapter.__init__)
59
+ return "tp_size" in sig.parameters
60
+
61
+
62
# Helper functions
63
def reformat_block_ids(block_ids: tuple[list[int], ...] | None) -> list[int]:
64
if block_ids is None:
@@ -101,6 +108,14 @@ def create_scheduler_adapter(
108
vllm_config.parallel_config.rank,
109
vllm_config,
110
)
111
+ tp_size = vllm_config.parallel_config.tensor_parallel_size
112
+
113
+ # Pass tp_size only when the adapter accepts it so that
114
+ # a newer vllm can still work with an older LMCache.
115
+ kwargs: dict[str, Any] = {}
116
+ if _adapter_accepts_tp_size():
117
+ kwargs["tp_size"] = tp_size
118
+
119
return LMCacheMPSchedulerAdapter(
120
server_url,
121
zmq_context,
@@ -108,6 +123,7 @@ def create_scheduler_adapter(
123
world_size,
124
kv_rank,
125
vllm_config.cache_config.block_size,
126
+ **kwargs,
127
)
128
129