fix

Author: hiworldwzj

lightllm/server/router/model_infer/infer_batch.py

@@ -78,6 +78,16 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache:
             req_objs.append(r_obj)
 
         self.infer_req_ids.extend(request_ids)
+
+        # 多输出模式下需要将请求添加到各自的组对象 InferReqGroup 中
+        if get_env_start_args().diverse_mode:
+            for r_id in request_ids:
+                req: InferReq = g_infer_context.requests_mapping[r_id]
+                group_req_id = req.shm_req.group_req_id
+                if group_req_id not in g_infer_context.group_mapping:
+                    g_infer_context.group_mapping[group_req_id] = InferReqGroup(group_req_id=group_req_id)
+                g_infer_context.group_mapping[group_req_id].add_req(r_id)
+
         return req_objs
 
     def free_a_req_mem(self, free_token_index: List, req: "InferReq", is_group_finished: bool):

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

@@ -4,33 +4,20 @@
     g_infer_context,
     InferReq,
     InferReqGroup,
-    InferSamplingParams,
 )
 from typing import List, Tuple
-from lightllm.utils.log_utils import init_logger
-from lightllm.server.tokenizer import get_tokenizer
 from lightllm.server.req_id_generator import convert_sub_id_to_group_id
 from lightllm.server.router.model_infer.mode_backend.pre import (
     prepare_prefill_inputs,
-    prepare_decode_inputs,
 )
 from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample
+from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack
 
 
 class DiversehBackend(ModeBackend):
     def __init__(self) -> None:
         super().__init__()
-
-    def init_custom(self):
-        pass
-
-    def build_group(self, req_ids: List[int]):
-        for r_id in req_ids:
-            req: InferReq = g_infer_context.requests_mapping[r_id]
-            group_req_id = req.shm_req.group_req_id
-            if group_req_id not in g_infer_context.group_mapping:
-                g_infer_context.group_mapping[group_req_id] = InferReqGroup(group_req_id=group_req_id)
-            g_infer_context.group_mapping[group_req_id].add_req(r_id)
+        self.prefill = self.beam_prefill
 
     def diverse_copy(self, groups: List[InferReqGroup]):
         batch_idx = []
@@ -46,64 +33,36 @@ def diverse_copy(self, groups: List[InferReqGroup]):
             run_reqs.extend(req_group.get_all_reqs())
         return batch_idx, run_reqs
 
-    def decode(self):
-        uninit_reqs, aborted_reqs, ok_finished_reqs, prefill_reqs, decode_reqs = self._get_classed_reqs(
-            g_infer_context.infer_req_ids,
-            strict_prefill=True,
+    def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq]):
+        group_reqs = [
+            g_infer_context.requests_mapping[req.req_id]
+            for req in prefill_reqs
+            if convert_sub_id_to_group_id(req.req_id) == req.req_id
+        ]
+        groups = [
+            g_infer_context.group_mapping[req.req_id]
+            for req in prefill_reqs
+            if convert_sub_id_to_group_id(req.req_id) == req.req_id
+        ]
+        model_input, group_run_reqs = prepare_prefill_inputs(
+            group_reqs, is_chuncked_mode=not self.disable_chunked_prefill, is_multimodal=self.is_multimodal
         )
+        model_output = self.model.forward(model_input)
+        logits = model_output.logits
 
-        if aborted_reqs:
-            g_infer_context.filter_reqs(aborted_reqs)
-        if prefill_reqs:
-            group_reqs = [
-                g_infer_context.requests_mapping[req.req_id]
-                for req in prefill_reqs
-                if convert_sub_id_to_group_id(req.req_id) == req.req_id
-            ]
-            groups = [
-                g_infer_context.group_mapping[req.req_id]
-                for req in prefill_reqs
-                if convert_sub_id_to_group_id(req.req_id) == req.req_id
-            ]
-            model_input, group_run_reqs = prepare_prefill_inputs(
-                group_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
-            )
-            model_output = self.model.forward(model_input)
-            logits = model_output.logits
-
-            uninit_req_ids = [req.req_id for req in uninit_reqs]
-            self._overlap_req_init_and_filter(
-                uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
-            )
-            self.build_group(uninit_req_ids)
-            batch_idx, run_reqs = self.diverse_copy(groups)
-            logits = logits[batch_idx]
-            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()
+        batch_idx, run_reqs = self.diverse_copy(groups)
+        logits = logits[batch_idx]
 
-            self._post_handle(
-                run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=True, do_filter_finished_reqs=False
-            )
+        next_token_ids_gpu, next_token_probs_gpu = sample(model_output.logits, run_reqs, self.eos_id)
+        next_token_ids_cpu = next_token_ids_gpu.detach().cpu().numpy()
+        next_token_logprobs_cpu = torch.log(next_token_probs_gpu).detach().cpu().numpy()
 
-        if decode_reqs:
-            model_input, run_reqs = prepare_decode_inputs(decode_reqs)
-            model_output = self.model.forward(model_input)
-            logits = model_output.logits
-            uninit_req_ids = [req.req_id for req in uninit_reqs]
-            self._overlap_req_init_and_filter(
-                uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
-            )
-            self.build_group(uninit_req_ids)
-
-            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()
-
-            self._post_handle(
-                run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=False, do_filter_finished_reqs=False
-            )
-        uninit_req_ids = [req.req_id for req in uninit_reqs]
-        self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
-        self.build_group(uninit_req_ids)
+        update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill)
+        self._post_handle(
+            run_reqs=run_reqs,
+            next_token_ids=next_token_ids_cpu,
+            next_token_logprobs=next_token_logprobs_cpu,
+            run_reqs_update_packs=update_packs,
+            extra_post_req_handle_func=self.extra_post_req_handle_func,
+        )
         return

lightllm/server/router/model_infer/pin_mem_manager.py

@@ -0,0 +1,44 @@
+import torch
+import threading
+import collections
+from typing import List, Dict
+
+
+class PinMemTensorManager:
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.key_to_tensor_list: Dict[str, List[torch.Tensor]] = collections.defaultdict(list)
+        self.key_to_alloc_index: Dict[str, int] = {}
+
+    def alloc_pin_tensor(self, key: str, size: int, dtype: torch.dtype):
+        """
+        利用4 buffer的 pin mem的cache，加速对pin mem的申请和释放操作。
+        """
+        with self.lock:
+            if key not in self.key_to_tensor_list:
+                self.key_to_tensor_list[key].append(
+                    torch.empty(size=(size,), dtype=dtype, device="cpu", pin_memory=True)
+                )
+                self.key_to_tensor_list[key].append(
+                    torch.empty(size=(size,), dtype=dtype, device="cpu", pin_memory=True)
+                )
+                self.key_to_tensor_list[key].append(
+                    torch.empty(size=(size,), dtype=dtype, device="cpu", pin_memory=True)
+                )
+                self.key_to_tensor_list[key].append(
+                    torch.empty(size=(size,), dtype=dtype, device="cpu", pin_memory=True)
+                )
+                self.key_to_alloc_index[key] = 0
+
+            alloc_index = self.key_to_alloc_index[key]
+            buff_tensor = self.key_to_tensor_list[key][alloc_index]
+            if buff_tensor.numel() < size:
+                self.key_to_tensor_list[key][alloc_index] = torch.empty(
+                    size=(size,), dtype=dtype, device="cpu", pin_memory=True
+                )
+                buff_tensor = self.key_to_tensor_list[key][alloc_index]
+            self.key_to_alloc_index[key] = (alloc_index + 1) % 4
+            return buff_tensor[0:size]
+
+
+g_pin_mem_manager = PinMemTensorManager()