@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c687bf226a

Fork 0

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 5, 06:00 UTC

[LMCache][MP] optimize save when mla enabled (#38810)

Signed-off-by: idellzheng <idellzheng@tencent.com> Co-authored-by: Yihua Cheng <yihua98@uchicago.edu>

chunxiaozheng committed Apr 14, 2026 at 08:56 UTC c687bf226a8e445bdeab8d8bb73b0880aa0e140b

3 files changed +120 -42

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/__init__.py

@@ -7,6 +7,7 @@ from .multi_process_adapter import (

7 LMCacheMPSchedulerAdapter,

8 LMCacheMPWorkerAdapter,

9 LoadStoreOp,

10 + ParallelStrategy,

11 )

13 __all__ = [

@@ -15,4 +16,5 @@ __all__ = [

16 "LMCacheMPSchedulerAdapter",

17 "LMCacheMPWorkerAdapter",

18 "LoadStoreOp",

19 + "ParallelStrategy",

20 ]

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py

+78 -14

@@ -79,6 +79,39 @@ def get_lmcache_chunk_size(

79 return chunk_size

82 +@dataclass

83 +class ParallelStrategy:

84 + use_mla: bool

85 + """Whether to use the MLA."""

86 +

87 + kv_world_size: int

88 + """

89 + The kv world size, kv_world_size may not be equal to the actual_world_size,

90 + in the case of mla, it will 'exclude' the effect of TP, the value is

91 + calculated by `extract_world_size_and_kv_rank` in `lmcache_mp_connector.py`.

92 + """

93 +

94 + kv_worker_id: int

95 + """

96 + The kv worker id of the sub-process, kv_worker_id may not be equal to the

97 + actual_worker_id, in the case of mla, it will 'exclude' the effect of TP,

98 + the value is calculated by `extract_world_size_and_kv_rank` in

99 + `lmcache_mp_connector.py`.

100 + """

101 +

102 + actual_world_size: int

103 + """The actual world size."""

104 +

105 + actual_worker_id: int

106 + """The actual worker id of the sub-process."""

107 +

108 + tp_size: int

109 + """The tensor parallel size."""

110 +

111 + pp_size: int

112 + """The pipeline parallel size."""

113 +

114 +

115 @dataclass

116 class LoadStoreOp:

117 block_ids: list[int]

@@ -111,10 +144,8 @@ class LMCacheMPSchedulerAdapter:

144 server_url: str,

145 context: zmq.Context,

146 model_name: str,

114 - world_size: int,

115 - kv_rank: int,

147 vllm_block_size: int,

117 - tp_size: int = 1,

148 + parallel_strategy: ParallelStrategy,

149 ):

150 """

151 Args:

@@ -122,11 +153,10 @@ class LMCacheMPSchedulerAdapter:

153 context: The ZMQ context

154

155 model_name: The model name used for LMCache keys

125 - world_size: The world size used for LMCache keys

126 - kv_rank: The kv rank used for LMCache keys

156 vllm_block_size: The block size used in vLLM

128 - tp_size: Tensor-parallel size for MLA

129 - multi-reader locking (default 1).

157 + parallel_strategy:

158 + The parallel strategy, which includes `use_mla`,

159 + `world_size`, `worker_id` and so on

160 """

161 self.mq_client = MessageQueueClient(server_url, context)

162

@@ -134,9 +164,7 @@ class LMCacheMPSchedulerAdapter:

164 self.lookup_futures: dict[str, MessagingFuture[LookupResult]] = {}

165

166 self.model_name = model_name

137 - self.world_size = world_size

138 - self.worker_id = kv_rank

139 - self.tp_size = tp_size

167 + self.parallel_strategy = parallel_strategy

168

169 # Read chunk size from lmcache

170 self.chunk_size = get_lmcache_chunk_size(self.mq_client)

@@ -145,6 +173,21 @@ class LMCacheMPSchedulerAdapter:

173 )

174 self.blocks_in_chunk = self.chunk_size // vllm_block_size

175

176 + @property

177 + def world_size(self) -> int:

178 + """The world size."""

179 + return self.parallel_strategy.kv_world_size

180 +

181 + @property

182 + def worker_id(self) -> int:

183 + """The worker id."""

184 + return self.parallel_strategy.kv_worker_id

185 +

186 + @property

187 + def tp_size(self) -> int:

188 + """The tensor parallel size."""

189 + return self.parallel_strategy.tp_size

190 +

191 @_lmcache_nvtx_annotate

192 def maybe_submit_lookup_request(

193 self,

@@ -308,9 +351,8 @@ class LMCacheMPWorkerAdapter:

351 server_url: str,

352 context: zmq.Context,

353 model_name: str,

311 - world_size: int,

312 - kv_rank: int,

354 vllm_block_size: int,

355 + parallel_strategy: ParallelStrategy,

356 ):

357 self.mq_client = MessageQueueClient(server_url, context)

358

@@ -336,8 +378,7 @@ class LMCacheMPWorkerAdapter:

378 self.previously_finished: set[str] = set()

379

380 self.model_name = model_name

339 - self.world_size = world_size

340 - self.worker_id = kv_rank

381 + self.parallel_strategy = parallel_strategy

382

383 # Read chunk size from lmcache

384 chunk_size = get_lmcache_chunk_size(self.mq_client)

@@ -346,6 +387,29 @@ class LMCacheMPWorkerAdapter:

387 )

388 self.blocks_in_chunk = chunk_size // vllm_block_size

389

390 + @property

391 + def world_size(self) -> int:

392 + """The world size."""

393 + return self.parallel_strategy.kv_world_size

394 +

395 + @property

396 + def worker_id(self) -> int:

397 + """The worker id."""

398 + return self.parallel_strategy.kv_worker_id

399 +

400 + @property

401 + def use_mla(self) -> bool:

402 + """Whether to use MLA."""

403 + return self.parallel_strategy.use_mla

404 +

405 + @property

406 + def is_first_rank_of_pp_group(self) -> bool:

407 + """Is the first rank of the pipeline parallel group."""

408 + return (

409 + self.parallel_strategy.actual_worker_id % self.parallel_strategy.tp_size

410 + == 0

411 + )

412 +

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

414 """

415 Register the kv caches with LMCache server

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

+40 -28

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 import enum

4 -import inspect

4 from collections.abc import Iterable

5 from dataclasses import dataclass, field

6 from typing import TYPE_CHECKING, Any, Literal

@@ -28,6 +27,7 @@ try:

27 LMCacheMPSchedulerAdapter,

28 LMCacheMPWorkerAdapter,

29 LoadStoreOp,

30 + ParallelStrategy,

31 )

33 try:

@@ -45,6 +45,7 @@ except ImportError:

45 LMCacheMPSchedulerAdapter,

46 LMCacheMPWorkerAdapter,

47 LoadStoreOp,

48 + ParallelStrategy,

49 )

51 if TYPE_CHECKING:

@@ -64,12 +65,6 @@ if TYPE_CHECKING:

65 logger = lmcache_init_logger(__name__)

67 -def _adapter_accepts_tp_size() -> bool:

68 - """Check if the imported adapter accepts tp_size."""

69 - sig = inspect.signature(LMCacheMPSchedulerAdapter.__init__)

70 - return "tp_size" in sig.parameters

71 -

72 -

68 # Helper functions

69 def reformat_block_ids(block_ids: tuple[list[int], ...] | None) -> list[int]:

70 if block_ids is None:

@@ -105,8 +100,8 @@ def extract_world_size_and_kv_rank(

100 # vLLM constructs TP groups first, and then construct other

101 # parallel groups on top of TP groups.

102 # for example, TP=4, PP=2,

108 - # TP group: [0, 1, 2, 3], [4, 5, 6, 7]

109 - # PP group: [0, 4], [1, 5], [2, 6], [3, 7]

103 + # PP group: [0, 1, 2, 3], [4, 5, 6, 7]

104 + # TP group: [0, 4], [1, 5], [2, 6], [3, 7]

105 # So we can "exclude" the effect of TP by rank // tp_size.

106 return world_size // tp_size, rank // tp_size

107

@@ -123,24 +118,24 @@ def create_scheduler_adapter(

118 vllm_config.parallel_config.rank,

119 vllm_config,

120 )

126 - tp_size = vllm_config.parallel_config.tensor_parallel_size

127 -

128 - # Pass tp_size only when the adapter accepts it so that

129 - # a newer vllm can still work with an older LMCache.

130 - kwargs: dict[str, Any] = {}

131 - if _adapter_accepts_tp_size():

132 - kwargs["tp_size"] = tp_size

133 -

134 - return LMCacheMPSchedulerAdapter(

135 - server_url,

136 - zmq_context,

137 - vllm_config.model_config.model,

121 + parallel_strategy = ParallelStrategy(

122 + mla_enabled(vllm_config.model_config),

123 world_size,

124 kv_rank,

140 - vllm_config.cache_config.block_size,

125 + vllm_config.parallel_config.world_size,

126 + vllm_config.parallel_config.rank,

127 + vllm_config.parallel_config.tensor_parallel_size,

128 + vllm_config.parallel_config.pipeline_parallel_size,

129 + )

130 +

131 + return LMCacheMPSchedulerAdapter(

132 + server_url=server_url,

133 + context=zmq_context,

134 + model_name=vllm_config.model_config.model,

135 + vllm_block_size=vllm_config.cache_config.block_size,

136 + parallel_strategy=parallel_strategy,

137 mq_timeout=mq_timeout,

138 heartbeat_interval=heartbeat_interval,

143 - **kwargs,

139 )

140

141

@@ -156,13 +151,22 @@ def create_worker_adapter(

151 vllm_config.parallel_config.rank,

152 vllm_config,

153 )

159 - return LMCacheMPWorkerAdapter(

160 - server_url,

161 - zmq_context,

162 - vllm_config.model_config.model,

154 + parallel_strategy = ParallelStrategy(

155 + mla_enabled(vllm_config.model_config),

156 world_size,

157 kv_rank,

165 - vllm_config.cache_config.block_size,

158 + vllm_config.parallel_config.world_size,

159 + vllm_config.parallel_config.rank,

160 + vllm_config.parallel_config.tensor_parallel_size,

161 + vllm_config.parallel_config.pipeline_parallel_size,

162 + )

163 +

164 + return LMCacheMPWorkerAdapter(

165 + server_url=server_url,

166 + context=zmq_context,

167 + model_name=vllm_config.model_config.model,

168 + vllm_block_size=vllm_config.cache_config.block_size,

169 + parallel_strategy=parallel_strategy,

170 mq_timeout=mq_timeout,

171 heartbeat_interval=heartbeat_interval,

172 )

@@ -612,6 +616,14 @@ class LMCacheMPConnector(KVConnectorBase_V1):

616

617 This prevents overwrites of paged KV buffer before saving done.

618 """

619 + # In MLA scenario, only the first rank of the pipeline group

620 + # needs to save the KV cache.

621 + if (

622 + self.worker_adapter.use_mla

623 + and not self.worker_adapter.is_first_rank_of_pp_group

624 + ):

625 + return

626 +

627 metadata = self._get_connector_metadata()

628 assert isinstance(metadata, LMCacheMPConnectorMetadata)

629