mooncake store connector

Co-authored-by: fems14 <1804143737@qq.com>
Co-authored-by: Dreamerleader <2270923832@qq.com>
Co-authored-by: Pz1116 <zpbzpb123123@gmail.com>
Co-authored-by: lizy124 <1950471827@qq.com>
Co-authored-by: zouyida2052 <zouyida2002@gmail.com>
Signed-off-by: LCAIZJ <leichao139636@163.com>

Author: 6 people

vllm_ascend/attention/mla_v1.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Type, TypeVar
+from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Type, TypeVar, List
 
 import torch
 import torch_npu
@@ -12,6 +12,10 @@
 from vllm.model_executor.layers.linear import (LinearBase,
                                                UnquantizedLinearMethod)
 from vllm.utils import cdiv, round_down
+from vllm.distributed.kv_transfer import (get_kv_transfer_group,
+                                          has_kv_transfer_group,
+                                          is_v1_kv_transfer_group)
+from vllm.forward_context import ForwardContext, get_forward_context
 
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.attention.attention_v1 import AscendAttentionState
@@ -976,6 +980,7 @@ def forward(
         assert attn_metadata.num_decodes is not None and \
         attn_metadata.num_prefills is not None and \
         attn_metadata.num_decode_tokens is not None
+        self.wait_for_kv_layer_from_connector(layer.layer_name)
         num_decode_tokens = attn_metadata.num_decode_tokens
         # Inputs and outputs may be padded for CUDA graphs
         output_padded = output
@@ -1046,4 +1051,36 @@ def forward(
                     is_force_scatter=self.enable_shared_expert_dp)[0]
                 current_ms_metadata.after_comm_event.record()
         del o_proj_input
+        self.maybe_save_kv_layer_to_connector(layer_name=layer.layer_name, kv_cache_layer=kv_cache)
         return output_padded
+
+    def wait_for_kv_layer_from_connector(self, layer_name: str):
+        if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
+            return
+
+        connector = get_kv_transfer_group()
+
+        forward_context: ForwardContext = get_forward_context()
+        attn_metadata = forward_context.attn_metadata
+        if attn_metadata is None:
+            return
+        assert isinstance(attn_metadata, AscendMLAMetadata)
+        connector.wait_for_layer_load(layer_name)
+
+    def maybe_save_kv_layer_to_connector(
+        self,
+        layer_name: str,
+        kv_cache_layer: List[torch.Tensor],
+    ):
+        if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
+            return
+
+        connector = get_kv_transfer_group()
+
+        forward_context: ForwardContext = get_forward_context()
+        attn_metadata = forward_context.attn_metadata
+        if attn_metadata is None:
+            return
+        assert isinstance(attn_metadata, AscendMLAMetadata)
+        connector.save_kv_layer(layer_name, kv_cache_layer,
+                                attn_metadata)

vllm_ascend/distributed/__init__.py

@@ -26,3 +26,8 @@
 KVConnectorFactory.register_connector(
     "MooncakeConnectorV1", "vllm_ascend.distributed.mooncake_connector",
     "MooncakeConnector")
+
+KVConnectorFactory.register_connector(
+    "MooncakeConnectorStoreV1",
+    "vllm_ascend.distributed.mooncake.mooncake_store_connector_v1",
+    "MooncakeConnectorV1")