[Feat] Support load async (#166)

* [Feat] support load_async

* [Feat] Support asynchronous loading of KV cache

Signed-off-by: y00945504 <yuhui87@huawei.com>

* [Fix bug] adapt unittest and continue to load next req when finish dealing with load_async req

Signed-off-by: flesher0813 <1208954694@qq.com>

* [Fix bug] func check return value not match

Signed-off-by: flesher0813 <1208954694@qq.com>

---------

Signed-off-by: y00945504 <yuhui87@huawei.com>
Signed-off-by: flesher0813 <1208954694@qq.com>
Co-authored-by: y00945504 <yuhui87@huawei.com>

Author: flesher0813

test/test_uc_connector.py

@@ -109,6 +109,7 @@ def init_uc(
             ucconnector.layerwise_load_tasks: dict[
                 str, dict[str, tuple[Task, Task]]
             ] = {}
+            ucconnector._need_load_reqs: dict[str, Union[list[int], list[Task]]] = {}
         return ucconnector
 
     def test_get_num_new_matched_tokens_hit(self):

ucm/integration/vllm/uc_connector.py

@@ -26,7 +26,7 @@
 import hashlib
 import pickle
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Generator, List, Optional
+from typing import TYPE_CHECKING, Any, Generator, List, Optional, Union
 
 import torch
 from vllm.config import VllmConfig
@@ -88,6 +88,8 @@ class ReqMeta:
     load_paras: Optional[LoadPara] = None
     # Save information
     save_paras: Optional[SavePara] = None
+    # Mark request which need load async
+    load_async: bool = False
 
 
 @dataclass
@@ -103,13 +105,15 @@ def add_request(
         vllm_block_ids: list[int],
         load_paras: Optional[LoadPara] = None,
         save_paras: Optional[SavePara] = None,
+        load_async: bool = False,
     ) -> None:
         self.requests.append(
             ReqMeta(
                 request_id=request_id,
                 vllm_block_ids=vllm_block_ids,
                 load_paras=load_paras,
                 save_paras=save_paras,
+                load_async=load_async,
             )
         )
 
@@ -136,6 +140,7 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         )
         self.element_size = vllm_config.model_config.dtype.itemsize
         self.kv_role = vllm_config.kv_transfer_config.kv_role
+        self._need_load_reqs: dict[str, Union[list[int], list[Task]]] = {}
         if (
             self._vllm_config.kv_transfer_config is not None
             and "ucm_connector_name"
@@ -326,35 +331,41 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
                 tensors, offsets = self.get_tensor_and_offset_layerwise(
                     fetch_block_ids, kv_layer, layer_name
                 )
-                if not self.use_layerwise:
-                    task = self.connector.load(
-                        fetch_block_hashes, offsets[:blocks_len], tensors[:blocks_len]
+                k_task_id = self.connector.load(
+                    fetch_block_hashes, offsets[:blocks_len], tensors[:blocks_len]
+                )
+                v_task_id = None
+                if not self.is_mla:
+                    v_task_id = self.connector.load(
+                        fetch_block_hashes,
+                        offsets[blocks_len:],
+                        tensors[blocks_len:],
                     )
-                    assert self.connector.wait(task) == 0
+                if request.request_id not in self.layerwise_load_tasks:
+                    self.layerwise_load_tasks[request.request_id] = {}
+                self.layerwise_load_tasks[request.request_id][layer_name] = (
+                    k_task_id,
+                    v_task_id,
+                )
+
+            if request.load_async:
+                for _, (k_task, v_task) in self.layerwise_load_tasks[
+                    request.request_id
+                ].items():
+                    if request.request_id not in self._need_load_reqs:
+                        self._need_load_reqs[request.request_id] = []
+                    self._need_load_reqs[request.request_id].append(k_task)
                     if not self.is_mla:
-                        task = self.connector.load(
-                            fetch_block_hashes,
-                            offsets[blocks_len:],
-                            tensors[blocks_len:],
-                        )
-                        assert self.connector.wait(task) == 0
-                else:
-                    k_task_id = self.connector.load(
-                        fetch_block_hashes, offsets[:blocks_len], tensors[:blocks_len]
-                    )
-                    v_task_id = None
+                        self._need_load_reqs[request.request_id].append(v_task)
+                continue
+
+            if not self.use_layerwise:
+                for _, (k_task, v_task) in self.layerwise_load_tasks[
+                    request.request_id
+                ].items():
+                    assert self.connector.wait(k_task) == 0
                     if not self.is_mla:
-                        v_task_id = self.connector.load(
-                            fetch_block_hashes,
-                            offsets[blocks_len:],
-                            tensors[blocks_len:],
-                        )
-                    if request.request_id not in self.layerwise_load_tasks:
-                        self.layerwise_load_tasks[request.request_id] = {}
-                    self.layerwise_load_tasks[request.request_id][layer_name] = (
-                        k_task_id,
-                        v_task_id,
-                    )
+                        assert self.connector.wait(v_task) == 0
 
     def wait_for_layer_load(self, layer_name: str) -> None:
         """
@@ -415,7 +426,7 @@ def save_kv_layer(
         assert attn_metadata is not None, "The attn_metadata should not be None."
 
         for request in metadata.requests:
-            if request.save_paras is None:
+            if request.save_paras is None or request.load_async:
                 continue
 
             save_param = request.save_paras
@@ -535,6 +546,30 @@ def wait_for_tasks():
         self.dump_tasks.clear()
         return success_dumped_blocks if success_dumped_blocks else None
 
+    def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
+        """Get the finished recving and sending requests."""
+        done_recving: set[str] = set()
+        for req_id, tasks in self._need_load_reqs.items():
+            unfinished_tasks = []
+            for task in tasks:
+                ret = self.connector.check(task)
+                if ret == 0:
+                    # Remove this assertion after support recompute load failed reqs
+                    assert self.connector.wait(task) == 0
+                    continue
+                if ret != -1:
+                    raise ValueError(f"Task {task.get_id()} Not Found")
+                unfinished_tasks.append(task)
+            if not unfinished_tasks:
+                done_recving.add(req_id)
+            self._need_load_reqs[req_id] = unfinished_tasks
+
+        # remove the finished requests
+        for req_id in list(done_recving):
+            self._need_load_reqs.pop(req_id, None)
+
+        return None, done_recving
+
     # ==============================
     # Scheduler-side methods
     # ==============================
@@ -588,6 +623,12 @@ def md5(input) -> int:
             can_load=False,
         )
 
+        need_load_tokens = max(num_external_computed_tokens - num_computed_tokens, 0)
+        if hasattr(self, "kv_role") and self.kv_role == "kv_consumer":
+            if need_load_tokens > 0:
+                self._need_load_reqs[request.request_id] = []
+                return need_load_tokens, True
+
         num_max_cached_tokens = max(num_external_computed_tokens, num_computed_tokens)
         num_blocks_need_save = (
             len(request.all_token_ids) - num_max_cached_tokens
@@ -608,7 +649,7 @@ def md5(input) -> int:
             f"num_computed_tokens = {num_computed_tokens}.\n"
         )
 
-        return max(num_external_computed_tokens - num_computed_tokens, 0), False
+        return need_load_tokens, False
 
     def update_state_after_alloc(
         self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
@@ -623,6 +664,12 @@ def update_state_after_alloc(
         if num_external_tokens > 0:
             self.load_paras[request.request_id].can_load = True
 
+        if request.request_id in self._need_load_reqs:
+            local_block_ids = (
+                blocks.get_unhashed_block_ids() if num_external_tokens > 0 else []
+            )
+            self._need_load_reqs[request.request_id] = local_block_ids
+
     def build_connector_meta(
         self, scheduler_output: SchedulerOutput
     ) -> KVConnectorMetadata:
@@ -636,6 +683,16 @@ def build_connector_meta(
             scheduler_output (SchedulerOutput): the scheduler output object.
         """
         meta = UCConnectorV1Metadata()
+
+        for req_id, block_ids in self._need_load_reqs.items():
+            meta.add_request(
+                req_id,
+                vllm_block_ids=block_ids,
+                load_paras=self.load_paras[req_id],
+                load_async=True,
+            )
+        self._need_load_reqs.clear()
+
         for new_req in scheduler_output.scheduled_new_reqs:
             # Load kv is only supported for new reqs
             new_scheduled_blocks = (

ucm/store/base.py

@@ -137,3 +137,16 @@ def commit(self, block_ids: List[str], is_success: bool = True) -> None:
             is_success(bool): if False, we need release block
         """
         pass
+
+    @abstractmethod
+    def check(self, task: Task) -> int:
+        """
+        check if kv transfer task finished.
+
+        Args:
+            task (Task): transfer engine task.
+        Returns:
+            0 - finished
+            others - in process.
+        """
+        pass

ucm/store/ucm_dram.py

@@ -187,3 +187,15 @@ def commit(self, block_ids: List[str], is_success: bool = True) -> None:
         """
         if is_success:
             self.cached_blocks.update(block_ids)
+
+    def check(self, task: Task) -> int:
+        """
+        check if kv transfer task finished.
+
+        Args:
+            task (Task): transfer engine task.
+        Returns:
+            0 - finished
+            others - in process.
+        """
+        pass

ucm/store/ucm_mooncake.py

@@ -331,3 +331,15 @@ def shutdown(self):
                 self.loop.close()
 
         self.store.close()
+
+    def check(self, task: Task) -> int:
+        """
+        check if kv transfer task finished.
+
+        Args:
+            task (Task): transfer engine task.
+        Returns:
+            0 - finished
+            others - in process.
+        """
+        pass

ucm/store/ucm_nfs_store.py

@@ -185,3 +185,22 @@ def commit(self, block_ids: List[str], is_success: bool = True) -> None:
             logger.warning(f"commit {block_ids} to {is_success}")
         ucmnfsstore.CommitBatch(block_ids, is_success)
         logger.debug("Succeed in committing kv cache.")
+
+    def check(self, task: Task) -> int:
+        """
+        check if kv transfer task finished.
+
+        Args:
+            task (Task): transfer engine task.
+        Returns:
+            0 - finished
+            -1 - in process
+            others - errors.
+        """
+        ret, finish = ucmnfsstore.Check(task.get_id())
+        if ret == 0 and finish == True:
+            return 0
+        elif ret == 0:
+            return -1
+        logger.error(f"check {task.get_id()} return {ret}")
+        return ret