Update vllm/distributed/kv_transfer/kv_connector/simple_connector.py

Co-authored-by: Nick Hill <[email protected]>

Author: da-x

vllm/distributed/kv_transfer/kv_connector/simple_connector.py

@@ -215,9 +215,8 @@ def send_kv_caches_and_hidden_states(
             values = torch.empty_like(keys)
 
             # Copy the relevant parts from kvcache
-            for layer_idx, layer_id in enumerate(range(start_layer,
-                                                       end_layer)):
-                kv_cache = kv_caches[layer_id - start_layer]
+            for layer_idx in range(end_layer - start_layer):
+                kv_cache = kv_caches[layer_idx]
 
                 if self.is_deepseek_mla and self.use_mla_opt:
                     key_cache = kv_cache.reshape(-1, num_heads, head_size)