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 12, 22:15 UTC
[DeepSeek + LMCache Multiprocess] handle MLA for deepseek model + LMCache Multiprocess connector (#29039)
Signed-off-by: KuntaiDu <kuntai@uchicago.edu>
Kuntai Du committed
Nov 20, 2025 at 09:40 UTC
05c2dee7e9f485f1e76eee084849e07c1c12a68b
1 file changed
+39
-8
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py
+39
-8
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, cast
7
8
import torch
9
import zmq
10
+from lmcache.integration.vllm.utils import mla_enabled
11
from lmcache.utils import init_logger as lmcache_init_logger
12
13
from vllm.config import VllmConfig
@@ -60,17 +61,44 @@ def reformat_block_ids(block_ids: tuple[list[int], ...] | None) -> list[int]:
61
return block_ids[0]
62
63
64
+def extract_world_size_and_kv_rank(
65
+ world_size: int,
66
+ rank: int,
67
+ vllm_config: VllmConfig,
68
+) -> tuple[int, int]:
69
+ """
70
+ Convert the rank for the MLA.
71
+ """
72
+ use_mla = mla_enabled(vllm_config.model_config)
73
+ if not use_mla:
74
+ return world_size, rank
75
+ else:
76
+ # Tensor parallel does not change the KV caches for MLA models.
77
+ # So we need to "exclude" the effect of TP on rank and world size
78
+ tp_size = vllm_config.parallel_config.tensor_parallel_size
79
+ # vLLM constructs TP groups first, and then construct other
80
+ # parallel groups on top of TP groups.
81
+ # for example, TP=4, PP=2,
82
+ # TP group: [0, 1, 2, 3], [4, 5, 6, 7]
83
+ # PP group: [0, 4], [1, 5], [2, 6], [3, 7]
84
+ # So we can "exclude" the effect of TP by rank // tp_size.
85
+ return world_size // tp_size, rank // tp_size
86
+
87
+
88
def create_scheduler_adapter(
89
server_url: str, zmq_context: zmq.Context, vllm_config: VllmConfig
90
) -> LMCacheMPSchedulerAdapter:
66
- # TODO: have a helper function to calculate the correct rank and
67
- # world size for the MLA and other models
91
+ world_size, kv_rank = extract_world_size_and_kv_rank(
92
+ vllm_config.parallel_config.world_size,
93
+ vllm_config.parallel_config.rank,
94
+ vllm_config,
95
+ )
96
return LMCacheMPSchedulerAdapter(
97
server_url,
98
zmq_context,
99
vllm_config.model_config.model,
72
- vllm_config.parallel_config.world_size,
73
- vllm_config.parallel_config.rank,
100
+ world_size,
101
+ kv_rank,
102
vllm_config.cache_config.block_size,
103
)
104
@@ -78,14 +106,17 @@ def create_scheduler_adapter(
106
def create_worker_adapter(
107
server_url: str, zmq_context: zmq.Context, vllm_config: VllmConfig
108
) -> LMCacheMPWorkerAdapter:
81
- # TODO: have a helper function to calculate the correct rank and
82
- # world size for the MLA and other models
109
+ world_size, kv_rank = extract_world_size_and_kv_rank(
110
+ vllm_config.parallel_config.world_size,
111
+ vllm_config.parallel_config.rank,
112
+ vllm_config,
113
+ )
114
return LMCacheMPWorkerAdapter(
115
server_url,
116
zmq_context,
117
vllm_config.model_config.model,
87
- vllm_config.parallel_config.world_size,
88
- vllm_config.parallel_config.rank,
118
+ world_size,
119
+ kv_rank,
120
vllm_config.cache_config.block_size,
121
)
122