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 15, 19:45 UTC
[Bugfix][Disaggregated] Add a check in send_kv_caches_and_hidden_states and fix the reshape of the KVCache (#14369)
Signed-off-by: Mathis Felardos <mathis@mistral.ai>
Mathis Felardos committed
Mar 8, 2025 at 07:39 UTC
980385f8c1839f98e04928cb07c643d20b140cf7
1 file changed
+14
-3
vllm/distributed/kv_transfer/kv_connector/simple_connector.py
+14
-3
@@ -2,7 +2,7 @@
2
"""
3
Simple KV Cache Connector for Distributed Machine Learning Inference
4
5
-The SimpleConnector transfers KV caches between prefill vLLM worker (KV cache
5
+The SimpleConnector transfers KV caches between prefill vLLM worker (KV cache
6
producer) and decode vLLM worker (KV cache consumer) using PyNcclPipe or
7
MooncakePipe.
8
@@ -159,6 +159,7 @@ class SimpleConnector(KVConnectorBase):
159
input_tokens_tensor = model_input.input_tokens
160
seq_lens = model_input.attn_metadata.seq_lens
161
slot_mapping_flat = model_input.attn_metadata.slot_mapping.flatten()
162
+ num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
163
start_layer = model_executable.model.start_layer
164
end_layer = model_executable.model.end_layer
165
@@ -166,7 +167,8 @@ class SimpleConnector(KVConnectorBase):
167
num_heads = int(model_config.num_key_value_heads / self.tp_size)
168
hidden_size = model_config.hidden_size
169
num_attention_heads = model_config.num_attention_heads
169
- head_size = int(hidden_size / num_attention_heads)
170
+ head_size = getattr(model_config, "head_dim",
171
+ int(hidden_size // num_attention_heads))
172
173
# query_lens contains new KV caches that are added to vLLM.
174
# so we will send them to decode instance
@@ -174,6 +176,15 @@ class SimpleConnector(KVConnectorBase):
176
for idx, slen in enumerate(seq_lens):
177
start_pos = sum(seq_lens[:idx])
178
end_pos = start_pos + slen
179
+
180
+ if start_pos >= num_prefill_tokens:
181
+ # vllm/worker/model_runner.py::_prepare_model_input_tensors:
182
+ # - input_tokens[:num_prefill_tokens] contains prefill tokens.
183
+ # - input_tokens[num_prefill_tokens:] contains decode tokens.
184
+ logger.warning("You have some decode requests while using "
185
+ "SimpleConnector. Their KVCache won't be sent.")
186
+ break
187
+
188
current_tokens = input_tokens_tensor[start_pos:end_pos]
189
190
keys, values = [], []
@@ -236,7 +247,7 @@ class SimpleConnector(KVConnectorBase):
247
# - input_tokens[num_prefill_tokens:] contains decode tokens.
248
logger.warning("You should set --enable_chunked_prefill=False "
249
"and --max_num_batched_tokens "
239
- "should be equal to max_seq_len_to_capture")
250
+ "should be equal to --max_seq_len_to_capture")
251
bypass_model_exec = False
252
assert start_pos == num_prefill_tokens
253
break