From f0b13d46b14f26d7e4107e148ea339ef665d28fa Mon Sep 17 00:00:00 2001 From: ShangmingCai Date: Tue, 10 Dec 2024 16:05:23 +0800 Subject: [PATCH 1/3] Fix value unpack error of simple connector. Signed-off-by: ShangmingCai --- .../kv_transfer/kv_connector/simple_connector.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/simple_connector.py b/vllm/distributed/kv_transfer/kv_connector/simple_connector.py index 5870070a54c75..fa758f3a15d2b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/simple_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/simple_connector.py @@ -118,6 +118,12 @@ def send_kv_caches_and_hidden_states( start_layer = model_executable.model.start_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_hidden_layers = model_config.num_attention_heads + head_size = int(hidden_size / num_hidden_layers) + # query_lens contains new KV caches that are added to vLLM. # so we will send them to decode instance # FIXME(Kuntai): This assume that all requests are prefill. @@ -131,8 +137,6 @@ def send_kv_caches_and_hidden_states( for layer_id in range(start_layer, end_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) value_cache = kv_cache[1].reshape(-1, num_heads, head_size) From ecfb7f060687096a2ff8dce977f355fbf246baad Mon Sep 17 00:00:00 2001 From: ShangmingCai Date: Tue, 10 Dec 2024 18:26:39 +0800 Subject: [PATCH 2/3] retrigger ci Signed-off-by: ShangmingCai From b8fc9aa530957ae4d222ca9e8b2112de821788db Mon Sep 17 00:00:00 2001 From: ShangmingCai Date: Wed, 11 Dec 2024 11:01:31 +0800 Subject: [PATCH 3/3] Fix typo. Signed-off-by: ShangmingCai --- vllm/distributed/kv_transfer/kv_connector/simple_connector.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/simple_connector.py b/vllm/distributed/kv_transfer/kv_connector/simple_connector.py index fa758f3a15d2b..bf4f40ca94e29 100644 --- a/vllm/distributed/kv_transfer/kv_connector/simple_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/simple_connector.py @@ -121,8 +121,8 @@ def send_kv_caches_and_hidden_states( model_config = model_executable.model.config num_heads = model_config.num_key_value_heads hidden_size = model_config.hidden_size - num_hidden_layers = model_config.num_attention_heads - head_size = int(hidden_size / num_hidden_layers) + 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. # so we will send them to decode instance