fix

Author: hiworldwzj

lightllm/common/basemodel/batch_objs.py

@@ -2,6 +2,7 @@
 from dataclasses import dataclass, field
 from typing import Optional
 from typing import List
+from lightllm.utils.envs_utils import get_env_start_args
 
 
 @dataclass
@@ -21,6 +22,15 @@ class ModelInput:
     b_req_idx: torch.Tensor = None
     b_mtp_index: torch.Tensor = None
     b_seq_len: torch.Tensor = None
+    # 只会在 diverse_mode 下的 decode 阶段真正被使用的参数, 用于记录共享的radix cache中的长度
+    b_shared_seq_len: torch.Tensor = None
+    # 只会在 diverse_mode 下的 decode 阶段真正被使用的参数, 用于记录请求间的共享关系。
+    # 举列说明:
+    # b_shared_seq_len : [10, 10, 10, 11, 11, 11, 11]
+    # b_mark_shared_group: [0, 0, 3, 0, 0, 0, 4]
+    # b_mark_shared_group 中每一个不为0的位置都代表其与前面多少个请求形成一个共享前缀组。属于
+    # 同一个共享前缀组的请求, 其在对应的 b_shared_seq_len 中的内容必然相同。
+    b_mark_shared_group: torch.Tensor = None
     mem_indexes: torch.Tensor = None
     is_prefill: bool = False
     b_ready_cache_len: torch.Tensor = None
@@ -52,6 +62,16 @@ def to_cuda(self):
             self.b_ready_cache_len = self.b_ready_cache_len.cuda(non_blocking=True)
         if self.b_prefill_start_loc is not None:
             self.b_prefill_start_loc = self.b_prefill_start_loc.cuda(non_blocking=True)
+        if not self.is_prefill and get_env_start_args().diverse_mode:
+            batch_size = len(self.b_req_idx)
+            if self.b_mark_shared_group is None:
+                self.b_mark_shared_group = torch.ones(size=(batch_size,), dtype=torch.int32, device="cuda")
+            else:
+                self.b_mark_shared_group = self.b_mark_shared_group.cuda(non_blocking=True)
+            if self.b_shared_seq_len is None:
+                self.b_shared_seq_len = torch.zeros(size=(batch_size,), dtype=torch.int32, device="cuda")
+            else:
+                self.b_shared_seq_len = self.b_shared_seq_len.cuda(non_blocking=True)
 
 
 @dataclass

lightllm/server/router/model_infer/infer_batch.py

@@ -434,6 +434,9 @@ def remove_master_req(self):
         else:
             logger.warning(f"try to remove master req, but related_master_req is None, req id {self.req_id}")
 
+    def get_radix_cache_shared_len(self):
+        return 0 if self.shared_kv_node is None else self.shared_kv_node.node_prefix_total_len
+
     def get_output_len(self):
         return self.cur_output_len
 

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

@@ -4,6 +4,7 @@
 from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context
 from lightllm.common.basemodel.infer_lock import g_infer_state_lock
 from lightllm.common.basemodel.batch_objs import ModelInput
+from lightllm.utils.envs_utils import get_env_start_args, get_diverse_max_batch_shared_group_size
 
 
 def prepare_prefill_inputs(
@@ -99,12 +100,16 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
     b_mtp_index = []
     b_seq_len = []
     b_q_seq_len = []
+    b_shared_seq_len = []
+    max_batch_shared_group_size = get_diverse_max_batch_shared_group_size()
     for req in req_objs:
+        _radix_shared_len = req.get_radix_cache_shared_len()
         run_reqs.append(req)
         b_req_idx.append(req.req_idx)
         seq_len = req.get_cur_total_len()
         assert req.cur_kv_len == seq_len - 1, f"{req.cur_kv_len} {seq_len}"
         b_seq_len.append(seq_len)
+        b_shared_seq_len.append(_radix_shared_len)
         total_token_num += seq_len
         max_len_in_batch = max(max_len_in_batch, seq_len)
         b_mtp_index.append(0)
@@ -114,6 +119,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
             b_req_idx.append(req.req_idx)
             seq_len += 1
             b_seq_len.append(seq_len)
+            b_shared_seq_len.append(_radix_shared_len)
             total_token_num += seq_len
             max_len_in_batch = max(max_len_in_batch, seq_len)
             b_mtp_index.append(step + 1)
@@ -124,7 +130,36 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
 
     b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu")
     b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu")
+    b_shared_seq_len = torch.tensor(b_shared_seq_len, dtype=torch.int32, device="cpu")
     b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu")
+    if get_env_start_args().diverse_mode:
+        b_mark_shared_group = []
+        shared_nodes = [req.shared_kv_node for req in run_reqs]
+        _current_group = []
+        for node in shared_nodes:
+            if not _current_group:
+                _current_group.append(node)
+            elif node == _current_group[-1]:
+                _current_group.append(node)
+            else:
+                b_mark_shared_group.extend([0 for _ in range(len(_current_group))])
+                b_mark_shared_group[-1] = len(_current_group)
+                _current_group.clear()
+                _current_group.append(node)
+
+            if len(_current_group) == max_batch_shared_group_size:
+                b_mark_shared_group.extend([0 for _ in range(len(_current_group))])
+                b_mark_shared_group[-1] = len(_current_group)
+                _current_group.clear()
+        if _current_group:
+            b_mark_shared_group.extend([0 for _ in range(len(_current_group))])
+            b_mark_shared_group[-1] = len(_current_group)
+            _current_group.clear()
+
+        assert len(b_mark_shared_group) == len(run_reqs)
+        b_mark_shared_group = torch.tensor(b_mark_shared_group, dtype=torch.int32, device="cpu")
+    else:
+        b_mark_shared_group = None
 
     # dynamic prompt cache 准备 token
     g_infer_state_lock.acquire()
@@ -144,6 +179,8 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
         b_req_idx=b_req_idx,
         b_mtp_index=b_mtp_index,
         b_seq_len=b_seq_len,
+        b_shared_seq_len=b_shared_seq_len,
+        b_mark_shared_group=b_mark_shared_group,
         is_prefill=False,
     )
     return model_input, run_reqs

lightllm/utils/envs_utils.py

@@ -193,4 +193,9 @@ def enable_radix_tree_timer_merge() -> bool:
 
 @lru_cache(maxsize=None)
 def get_radix_tree_merge_update_delta() -> int:
-    return int(os.getenv("LIGHTLMM_RADIX_TREE_MERGE_DELTA", 6000))
+    return int(os.getenv("LIGHTLLM_RADIX_TREE_MERGE_DELTA", 6000))
+
+
+@lru_cache(maxsize=None)
+def get_diverse_max_batch_shared_group_size() -> int:
+    return int(os.getenv("LIGHTLLM_MAX_BATCH_SHARED_GROUP_SIZE", 4))