1

Signed-off-by: MengqingCao <[email protected]>

Author: MengqingCao

vllm_ascend/worker/model_runner_v1.py

@@ -2774,6 +2774,7 @@ def _allocate_kv_cache_tensors(
             # TODO: REFACTOR ME to sharing hybrid cache
             for idx in range(len(kv_cache_tensor.shared_by)):
                 layer_name = kv_cache_tensor.shared_by[idx]
+                print(30*"-", f"layer_name: {layer_name}")
                 if "linear_attn" in layer_name:
                     # for mamba linear attention
                     if self.vllm_config.kv_transfer_config is None:
@@ -2788,7 +2789,8 @@ def _allocate_kv_cache_tensors(
                         tensor = self._align_memory(
                             tensor, alignment)[:kv_cache_tensor.size]
                     kv_cache_raw_tensors[layer_name] = tensor
-                elif "attn" in layer_name:
+                elif "attn" in layer_name and layer_name not in kv_cache_raw_tensors.keys():
+                    print(30*"/", f"layer_name: {layer_name}")
                     # NOTE: We need to init k cache tensor (nope cache tensor in mla) and
                     # v cache tensor (rope cache tensor in mla) separately to support llmdatadist,
                     # as it only support the 0-dim of kv_cache is `num_blocks`.
@@ -2862,14 +2864,14 @@ def _allocate_kv_cache_tensors(
                                 and "linear_attn" not in layer_name_inner):
                             kv_cache_raw_tensors[layer_name_inner] = (k_tensor, v_tensor) if \
                                 not self.use_sparse else (k_tensor, v_tensor, k_cache_tensor)
-                            break
 
         layer_names = set()
         for group in kv_cache_config.kv_cache_groups:
             for layer_name in group.layer_names:
                 if layer_name in self.runner_only_attn_layers:
                     continue
                 layer_names.add(layer_name)
+                print(30*"=", f"layer_name: {layer_name}: kv_cache_raw_tensors[layer_name]: {id(kv_cache_raw_tensors[layer_name])}")
         assert layer_names == set(kv_cache_raw_tensors.keys(
         )), "Some layers are not correctly initialized"