xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-13 21:55:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] Fix value unpack error of simple connector for KVCache transfer. (#11058)

Browse Source 
Signed-off-by: ShangmingCai <csmthu@gmail.com>
This commit is contained in:
shangmingc 2024-12-13 05:19:17 +08:00 committed by GitHub
parent 9f3974a319
commit db6c264a1e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/distributed/kv_transfer/kv_connector/simple_connector.py
View File

@ -118,6 +118,12 @@ class SimpleConnector(KVConnectorBase):
start_layer = model_executable.model.start_layer start_layer = model_executable.model.start_layer
end_layer = model_executable.model.end_layer end_layer = model_executable.model.end_layer
model_config = model_executable.model.config
num_heads = model_config.num_key_value_heads
hidden_size = model_config.hidden_size
num_attention_heads = model_config.num_attention_heads
head_size = int(hidden_size / num_attention_heads)
# query_lens contains new KV caches that are added to vLLM. # query_lens contains new KV caches that are added to vLLM.
# so we will send them to decode instance # so we will send them to decode instance
# FIXME(Kuntai): This assume that all requests are prefill. # FIXME(Kuntai): This assume that all requests are prefill.
@ -131,8 +137,6 @@ class SimpleConnector(KVConnectorBase):
for layer_id in range(start_layer, end_layer): for layer_id in range(start_layer, end_layer):
kv_cache = kv_caches[layer_id - start_layer] kv_cache = kv_caches[layer_id - start_layer]
_, _, num_heads, head_size = kv_cache[0].shape
key_cache = kv_cache[0].reshape(-1, num_heads, head_size) key_cache = kv_cache[0].reshape(-1, num_heads, head_size)
value_cache = kv_cache[1].reshape(-1, num_heads, head_size) value_cache = kv_cache[1].reshape(-1, num_heads, head_size)