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, 18:30 UTC
[Frontend][Bugfix] support prefill decode disaggregation on deepseek (#14824)
Signed-off-by: billishyahao <bill.he@amd.com> Co-authored-by: Zhai Feiyue <80079571+ZhaiFeiyue@users.noreply.github.com>
billishyahao committed
Mar 20, 2025 at 15:00 UTC
742369d35a95999176e1fd0391646dde3aaa3ec3
3 files changed
+70
-21
examples/online_serving/disaggregated_prefill.sh
+9
-4
@@ -8,6 +8,9 @@ set -xe
8
echo "🚧🚧 Warning: The usage of disaggregated prefill is experimental and subject to change 🚧🚧"
9
sleep 1
10
11
+# meta-llama/Meta-Llama-3.1-8B-Instruct or deepseek-ai/DeepSeek-V2-Lite
12
+MODEL_NAME=${HF_MODEL_NAME:-meta-llama/Meta-Llama-3.1-8B-Instruct}
13
+
14
# Trap the SIGINT signal (triggered by Ctrl+C)
15
trap 'cleanup' INT
16
@@ -44,18 +47,20 @@ wait_for_server() {
47
# You can also adjust --kv-ip and --kv-port for distributed inference.
48
49
# prefilling instance, which is the KV producer
47
-CUDA_VISIBLE_DEVICES=0 vllm serve meta-llama/Meta-Llama-3.1-8B-Instruct \
50
+CUDA_VISIBLE_DEVICES=0 vllm serve $MODEL_NAME \
51
--port 8100 \
52
--max-model-len 100 \
53
--gpu-memory-utilization 0.8 \
54
+ --trust-remote-code \
55
--kv-transfer-config \
56
'{"kv_connector":"PyNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2}' &
57
58
# decoding instance, which is the KV consumer
55
-CUDA_VISIBLE_DEVICES=1 vllm serve meta-llama/Meta-Llama-3.1-8B-Instruct \
59
+CUDA_VISIBLE_DEVICES=1 vllm serve $MODEL_NAME \
60
--port 8200 \
61
--max-model-len 100 \
62
--gpu-memory-utilization 0.8 \
63
+ --trust-remote-code \
64
--kv-transfer-config \
65
'{"kv_connector":"PyNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2}' &
66
@@ -78,7 +83,7 @@ sleep 1
83
output1=$(curl -X POST -s http://localhost:8000/v1/completions \
84
-H "Content-Type: application/json" \
85
-d '{
81
-"model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
86
+"model": "'"$MODEL_NAME"'",
87
"prompt": "San Francisco is a",
88
"max_tokens": 10,
89
"temperature": 0
@@ -87,7 +92,7 @@ output1=$(curl -X POST -s http://localhost:8000/v1/completions \
92
output2=$(curl -X POST -s http://localhost:8000/v1/completions \
93
-H "Content-Type: application/json" \
94
-d '{
90
-"model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
95
+"model": "'"$MODEL_NAME"'",
96
"prompt": "Santa Clara is a",
97
"max_tokens": 10,
98
"temperature": 0
vllm/distributed/kv_transfer/kv_connector/simple_connector.py
+60
-17
@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, List, Optional, Tuple, Union
12
13
import torch
14
15
+import vllm.envs as envs
16
from vllm import _custom_ops as ops
17
from vllm.config import VllmConfig
18
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
@@ -37,6 +38,8 @@ class SimpleConnector(KVConnectorBase):
38
39
self.config = config.kv_transfer_config
40
self.tp_size = config.parallel_config.tensor_parallel_size
41
+ self.is_deepseek_mla = config.model_config.is_deepseek_mla
42
+ self.use_mla_opt = not envs.VLLM_MLA_DISABLE
43
44
if self.config.kv_connector == "PyNcclConnector":
45
from vllm.distributed.kv_transfer.kv_pipe.pynccl_pipe import (
@@ -167,8 +170,26 @@ class SimpleConnector(KVConnectorBase):
170
num_heads = int(model_config.num_key_value_heads / self.tp_size)
171
hidden_size = model_config.hidden_size
172
num_attention_heads = model_config.num_attention_heads
170
- head_size = getattr(model_config, "head_dim",
171
- int(hidden_size // num_attention_heads))
173
+
174
+ # Deepseek's MLA (Multi-head Latent Attention) uses two different
175
+ # kv_cache shapes based on whether VLLM_MLA_DISABLE is set to 0.
176
+ # When VLLM_MLA_DISABLE=0 (default), forward absorb is applied,
177
+ # resulting in a kv_cache shape of [num_blks, blk_size, 1,
178
+ # kv_lora_rank + qk_rope_head_dim].
179
+ # When VLLM_MLA_DISABLE=1, standard FA is used instead, leading
180
+ # to a kv_cache shape of [2, num_blks, blk_size,
181
+ # num_key_value_heads / tp, qk_nope_head_dim + qk_rope_head_dim].
182
+ # For more details, see vllm/attention/backends/mla/common.py.
183
+ if self.is_deepseek_mla and self.use_mla_opt:
184
+ head_size = model_config.kv_lora_rank + \
185
+ model_config.qk_rope_head_dim
186
+ num_heads = 1
187
+ elif self.is_deepseek_mla and not self.use_mla_opt:
188
+ head_size = model_config.qk_nope_head_dim + \
189
+ model_config.qk_rope_head_dim
190
+ else:
191
+ head_size = getattr(model_config, "head_dim",
192
+ int(hidden_size // num_attention_heads))
193
194
# query_lens contains new KV caches that are added to vLLM.
195
# so we will send them to decode instance
@@ -192,8 +213,12 @@ class SimpleConnector(KVConnectorBase):
213
for layer_id in range(start_layer, end_layer):
214
kv_cache = kv_caches[layer_id - start_layer]
215
195
- key_cache = kv_cache[0].reshape(-1, num_heads, head_size)
196
- value_cache = kv_cache[1].reshape(-1, num_heads, head_size)
216
+ if self.is_deepseek_mla and self.use_mla_opt:
217
+ key_cache = kv_cache.reshape(-1, num_heads, head_size)
218
+ value_cache = kv_cache.reshape(-1, num_heads, head_size)
219
+ else:
220
+ key_cache = kv_cache[0].reshape(-1, num_heads, head_size)
221
+ value_cache = kv_cache[1].reshape(-1, num_heads, head_size)
222
223
current_slot_mapping = slot_mapping_flat[start_pos:end_pos]
224
@@ -223,6 +248,8 @@ class SimpleConnector(KVConnectorBase):
248
# and hidden states.
249
bypass_model_exec = True
250
251
+ model_config = model_executable.model.config
252
+
253
input_tokens_tensor = model_input.input_tokens
254
seq_lens = model_input.attn_metadata.seq_lens
255
num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
@@ -291,19 +318,35 @@ class SimpleConnector(KVConnectorBase):
318
kv_cache = kv_caches[i - model_executable.model.start_layer]
319
layer = model_executable.model.layers[i]
320
294
- key_cache, value_cache = kv_cache[0], kv_cache[1]
295
- ops.reshape_and_cache_flash(
296
- keys[i - model_executable.model.start_layer].to(
297
- key_cache.device),
298
- values[i - model_executable.model.start_layer].to(
299
- value_cache.device),
300
- key_cache,
301
- value_cache,
302
- slot_mapping[start_pos:end_pos],
303
- layer.self_attn.attn.kv_cache_dtype,
304
- layer.self_attn.attn._k_scale,
305
- layer.self_attn.attn._v_scale,
306
- )
321
+ if self.is_deepseek_mla and self.use_mla_opt:
322
+ layer.self_attn.attn = layer.self_attn.mla_attn
323
+ k_c_normed_k_pe = keys[
324
+ i - model_executable.model.start_layer].to(
325
+ kv_cache.device).squeeze(1)
326
+ k_c_normed = k_c_normed_k_pe[:, :model_config.kv_lora_rank]
327
+ k_pe = k_c_normed_k_pe[:, model_config.kv_lora_rank:]
328
+ ops.concat_and_cache_mla(
329
+ k_c_normed,
330
+ k_pe,
331
+ kv_cache,
332
+ slot_mapping[start_pos:end_pos],
333
+ layer.self_attn.attn.kv_cache_dtype,
334
+ layer.self_attn.attn._k_scale,
335
+ )
336
+ else:
337
+ key_cache, value_cache = kv_cache[0], kv_cache[1]
338
+ ops.reshape_and_cache_flash(
339
+ keys[i - model_executable.model.start_layer].to(
340
+ key_cache.device),
341
+ values[i - model_executable.model.start_layer].to(
342
+ value_cache.device),
343
+ key_cache,
344
+ value_cache,
345
+ slot_mapping[start_pos:end_pos],
346
+ layer.self_attn.attn.kv_cache_dtype,
347
+ layer.self_attn.attn._k_scale,
348
+ layer.self_attn.attn._v_scale,
349
+ )
350
351
hidden_or_intermediate_states_for_one_req.append(hidden)
352
vllm/model_executor/models/deepseek_v2.py
+1
@@ -589,6 +589,7 @@ class DeepseekV2Model(nn.Module):
589
model_config = vllm_config.model_config
590
cache_config = vllm_config.cache_config
591
quant_config = vllm_config.quant_config
592
+ self.config = config
593
594
self.vocab_size = config.vocab_size
595