Support only one GPU store partially with bug

Author: jayfeather9

lightllm/server/router/dynamic_prompt/hiradix_cache.py

@@ -17,6 +17,9 @@ def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager, max_
         logger.info("Initializing HiRadixCache")
         self.rank_in_node = rank_in_node
         try:
+            # TODO: determine by model type && dp, tp
+            store_once = True  # Deepseek -> True, Llama -> False
+            self.do_store = store_once and self.rank_in_node == 0
             self.is_hi_radix_cache = True
             all_buffers = self.mem_manager.kv_buffer
             all_buffers = all_buffers.view(all_buffers.shape[0], all_buffers.shape[1], -1)
@@ -37,83 +40,111 @@ def __init__(self, unique_name, total_token_num, rank_in_node, mem_manager, max_
     # then when the decode finishes, do syncronize to see whether this can be free
     # no buffer, parallel insert inputs
     def insert_disk(self, req_id, key, value):
+        if not self.do_store:
+            return
         if req_id in self.working_tasks:
-            self.wait_till_finish(req_id)
+            self.abort_req_store_task(req_id)
         self.working_tasks[req_id] = self.py_cache_service.create(tokens=key, kv_page_indexer=value, mode="w")
         logger.info(f"Created store task for req {req_id}.")
 
-    def wait_till_finish(self, req_id):
-        if req_id not in self.working_tasks:
+    def abort_req_store_task(self, req_id):
+        if not self.do_store:
+            return
+        if self.working_tasks[req_id].ready():
+            logger.info(f"Calling abort for req {req_id}, but is finished.")
             return
-        starting_time = time.time()
-        while not self.working_tasks[req_id].ready():
-            time.sleep(0.01)
-        logger.info(f"Waited {time.time() - starting_time}s for req {req_id}.")
-
-    # def insert(self, key, value=None):
-    #     if value is None:
-    #         value = key
-
-    #     assert len(key) == len(value)  # and len(key) >= 1
-    #     if len(key) == 0:
-    #         return 0
-
-    #     # current implement is serial, TODO: make it parallel
-    #     # if no hi_cache_buffer, work with normal radix cache
-    #     if self.hi_cache_kv_buffer is not None:
-    #         do_copy = False
-    #         # and if is moving, ignore this insert request
-    #         with self.moving_lock:
-    #             if (not self.start_store_task) and self.write_task is not None:
-    #                 if self.write_task.ready():
-    #                     logger.info(f"HiCache of [{self.rank_in_node}]: stored len = {self.hi_cache_buffer_len}")
-    #                     self.start_store_task = True # ensure ready => start new only one kvcache stores
-    #                     do_copy = True
-    #             elif self.write_task is None and self.starting:
-    #                 self.starting = False
-    #                 self.start_store_task = True
-    #                 do_copy = True
-
-    #         if do_copy:
-    #             # copy the key and value to the hi_cache_buffer
-    #             self.hi_cache_key_buffer[:len(key)].copy_(key)
-    #             self.hi_cache_buffer_len = len(key)
-    #             for buffer_index, index in enumerate(value):
-    #                 kv_data = self.mem_manager.get_index_kv_buffer(index)
-    #                 self.mem_manager.load_index_kv_buffer(self.hi_cache_kv_buffer[buffer_index], kv_data)
-    #             # create a new thread to store the buffer
-    #             self._store_buffer()
-
-    #     return self._insert_helper(self.root_node, key, value)
-
-    # def _store_buffer(self):
-    #     logger.info(f"Storing buffer size = {self.hi_cache_buffer_len}")
-    #     assert self.hi_cache_buffer_len > 0
-    #     assert self.hi_cache_kv_buffer is not None
-    #     key = self.hi_cache_key_buffer[:self.hi_cache_buffer_len].tolist()
-    #     self.write_task = self.py_cache_service.create(
-    #         tokens=key, kv_page_indexer=self.hi_cache_kv_buffer[:self.hi_cache_buffer_len], mode="w")
-    #     with self.moving_lock:
-    #         self.start_store_task = False
+        logger.info(f"Aborting req {req_id} unfinished.")
+        self.py_cache_service.az5(self.working_tasks[req_id])
+
+    # TODO: finish this function to only update new ones
+    def _reinsert_helper(self, node: TreeNode, key, value, ans_value_list: list, update_refs=False):
+        if node.is_leaf():
+            self.evict_tree_set.discard(node)
+
+        if update_refs:
+            node.ref_counter += 1
+            # from 0 to 1 need update refs token num
+            if node.ref_counter == 1:
+                self.refed_tokens_num.arr[0] += len(node.token_mem_index_value)
+
+        try:
+            if len(key) == 0:
+                return node
+
+            first_key_id = key[0].item()
+            if first_key_id in node.children.keys():
+                child: TreeNode = node.children[first_key_id]
+                prefix_len = match(key, child.token_id_key)
+                if prefix_len == len(key):
+                    if child.is_leaf():
+                        self.evict_tree_set.discard(child)
+                    child.update_time()
+                    ans_value_list.append(child.token_mem_index_value)
+                    if child.is_leaf():
+                        self.evict_tree_set.add(child)
+                    return prefix_len
+
+                elif prefix_len < len(key) and prefix_len < len(child.token_id_key):
+                    if child.is_leaf():
+                        self.evict_tree_set.discard(child)
+
+                    key = key[prefix_len:]
+                    value = value[prefix_len:]
+                    split_parent_node = child.split_node(prefix_len)
+                    new_node = split_parent_node.add_and_return_new_child(key, value)
+                    # update total token num
+                    self.tree_total_tokens_num.arr[0] += len(new_node.token_mem_index_value)
+
+                    if split_parent_node.is_leaf():
+                        self.evict_tree_set.add(split_parent_node)
+                    if new_node.is_leaf():
+                        self.evict_tree_set.add(new_node)
+
+                    if child.is_leaf():
+                        self.evict_tree_set.add(child)
+                    return prefix_len
+                elif prefix_len < len(key) and prefix_len == len(child.token_id_key):
+                    return prefix_len + self._insert_helper(child, key[prefix_len:], value[prefix_len:])
+                else:
+                    assert False, "can not run to here"
+
+            else:
+                new_node = node.add_and_return_new_child(key, value)
+                # update total token num
+                self.tree_total_tokens_num.arr[0] += len(new_node.token_mem_index_value)
+                ans_value_list.append(new_node.token_mem_index_value)
+                if update_refs:
+                    new_node.ref_counter += 1
+                    if new_node.ref_counter == 1:
+                        self.refed_tokens_num.arr[0] += len(new_node.token_mem_index_value)
+                if new_node.is_leaf():
+                    self.evict_tree_set.add(new_node)
+                return new_node
+        finally:
+            node.update_time()
+            if node.is_leaf():
+                self.evict_tree_set.add(node)
 
     def match_prefix(self, key, update_refs=False):
         st_time = time.time()
         assert len(key) != 0
         ans_value_list = []
-        tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=update_refs)
+        tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=False)
         # add a parameter if get long enough (>50%)
         first_query_time = time.time()
         logger.info(f"HiCache of [{self.rank_in_node}]: No.1 First GPU query took {first_query_time - st_time}")
         max_len = self._query_hi_cache(key)  # x64
         hi_cache_query_time = time.time()
         logger.info(f"HiCache of [{self.rank_in_node}]: No.2 Disk query took {hi_cache_query_time - first_query_time}")
-        logger.info(f"Matched {len(ans_value_list)} from gpu and {max_len} from disk.")
+        logger.info(f"Matched {sum(len(s) for s in ans_value_list)} from gpu and {max_len} from disk.")
         pull_hi_cache = False
-        if max_len > len(ans_value_list):
+        if max_len > sum(len(s) for s in ans_value_list):
             pull_hi_cache = True
             try:
                 self.free_radix_cache_to_get_enough_token(max_len)
             except:
+                if update_refs:
+                    tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=update_refs)
                 pull_hi_cache = False
         if pull_hi_cache:
             buffers = self.mem_manager.alloc(max_len)
@@ -133,7 +164,10 @@ def match_prefix(self, key, update_refs=False):
             logger.info(f"HiCache of [{self.rank_in_node}]: No.4 Reinsert took {insert_time - hicache_pull_time}")
             ans_value_list = []
             tree_node = self._match_prefix_helper(self.root_node, key, ans_value_list, update_refs=update_refs)
-            logger.info(f"HiCache of [{self.rank_in_node}]: No.5 Re match prefix took {time.time() - insert_time}")
+            logger.info(
+                f"HiCache of [{self.rank_in_node}]: No.5 Re match prefix took {time.time() - insert_time}"
+                + f" matched {sum(len(s) for s in ans_value_list)} tokens"
+            )
         if tree_node != self.root_node:
             if len(ans_value_list) != 0:
                 value = torch.concat(ans_value_list)

lightllm/server/router/model_infer/infer_batch.py

@@ -110,7 +110,7 @@ def free_a_req_mem(self, free_token_index: List, req: "InferReq", is_group_finis
                     req.shared_kv_node = None
 
             if self.radix_cache.is_hi_radix_cache:
-                self.radix_cache.wait_till_finish(req.req_id)
+                self.radix_cache.abort_req_store_task(req.req_id)
 
     def _save_promptcache_kvbuffer(self):
         """

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -165,13 +165,6 @@ def decode(self):
         """This method can be overridden in subclasses."""
         raise NotImplementedError()
 
-    def store_hicache_after_prefill(self, run_reqs):
-        if self.use_hi_dynamic_prompt_cache and self.radix_cache is not None:
-            for req in run_reqs:
-                key = torch.tensor(req.get_input_token_ids()[0 : req.cur_kv_len], dtype=torch.int64, device="cpu")
-                value = self.model.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].detach().cpu()
-                self.radix_cache.insert_disk(req.req_id, key, value)
-
     def pause_reqs(self, req_ids):
         if self.dp_size_in_node != 1:
             req_ids = [req_id for req_id in req_ids if req_id in g_infer_context.requests_mapping]
@@ -350,6 +343,23 @@ def _overlap_req_init_and_filter(
 
         return
 
+    def _overlap_store_prefill_reqs(self, run_reqs: List[InferReq]):
+        if run_reqs:
+            with torch.cuda.stream(g_infer_context.get_overlap_stream()):
+                if self.use_hi_dynamic_prompt_cache and self.radix_cache is not None:
+                    for req in run_reqs:
+                        if req.cur_output_len > 1:
+                            continue
+                        key = torch.tensor(
+                            req.get_input_token_ids()[0 : req.cur_kv_len], dtype=torch.int64, device="cpu"
+                        )
+                        value = self.model.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].detach().cpu()
+                        self.radix_cache.insert_disk(req.req_id, key, value)
+
+            torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream())
+
+        return
+
     # 一些可以复用的通用功能函数
     def _post_init_reqs(self, uninit_reqs: List[InferReq]):
         """

lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py

@@ -43,6 +43,7 @@ def decode(self):
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
             )
+            self._overlap_store_prefill_reqs(run_reqs=run_reqs)
             next_token_ids, next_token_probs = sample(logits, run_reqs, self.eos_id)
             next_token_ids = next_token_ids.detach().cpu().numpy()
             next_token_logprobs = torch.log(next_token_probs).detach().cpu().numpy()
@@ -59,7 +60,6 @@ def decode(self):
                     prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
                 )
                 logits = self.model.forward(**kwargs)
-                self.store_hicache_after_prefill(run_reqs)
                 self._overlap_req_init_and_filter(
                     uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
                 )

lightllm/server/router/model_infer/mode_backend/continues_batch/impl.py

@@ -35,8 +35,6 @@ def decode(self):
             )
             logits = self.model.forward(**kwargs)
 
-            self.store_hicache_after_prefill(run_reqs)
-
             self._overlap_req_init_and_filter(
                 uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
             )