rebase main.

Author: Weichao Luo

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

@@ -261,7 +261,7 @@ def _post_handle(
         is_chuncked_mode: bool,
         do_filter_finished_reqs: bool,
         extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
-        extra_post_req_handle_chunk_func: Optional[Callable[[InferReq], None]] = None,
+        call_post_handle_for_chunk: bool = False ,
     ) -> List[int]:
         """
         extra_post_req_handle_func 用于提供在一个请求确定输出的时候，给出额外的后处理操作，主要是用于
@@ -282,16 +282,12 @@ def _post_handle(
             if self.is_master_in_dp:
                 shm_req.shm_cur_kv_len = req_obj.cur_kv_len
 
-            if extra_post_req_handle_chunk_func is not None:
-                # 如果存在额外的处理函数，则调用这个函数进行处理。
-                extra_post_req_handle_chunk_func(req_obj)
-
             # 对于没有到达需要输出 token 阶段的请求，直接略过, 说明还
             # 处于chuncked prefill kv 填充的阶段。
             if req_obj.cur_kv_len < req_obj.get_cur_total_len():
                 # chunk transfer
-                if extra_post_req_handle_chunk_func is not None:
-                    extra_post_req_handle_chunk_func(req_obj)
+                if call_post_handle_for_chunk and extra_post_req_handle_func:
+                    extra_post_req_handle_func(req_obj, next_token_id, next_token_logprob)
 
                 continue
 
@@ -314,9 +310,6 @@ def _post_handle(
             # 更新判断请求的 finished 状态
             req_obj.update_finish_status(self.eos_id)
 
-            if extra_post_req_handle_chunk_func is not None:
-                extra_post_req_handle_chunk_func(req_obj, next_token_id, next_token_logprob)
-
             if extra_post_req_handle_func is not None:
                 extra_post_req_handle_func(req_obj, next_token_id, next_token_logprob)
 

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

@@ -46,6 +46,7 @@ def normal_prefill_reqs(
         ok_finished_reqs: List[InferReq],
         mask_func: Optional[Callable[[List[InferReq], torch.Tensor], None]] = None,
         extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
+        call_post_handle_for_chunk: bool = False
     ):
         model_input, run_reqs = prepare_prefill_inputs(
             prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill, is_multimodal=self.is_multimodal
@@ -69,6 +70,7 @@ def normal_prefill_reqs(
             is_chuncked_mode=not self.disable_chunked_prefill,
             do_filter_finished_reqs=False,
             extra_post_req_handle_func=extra_post_req_handle_func,
+            call_post_handle_for_chunk=call_post_handle_for_chunk,
         )
         return
 

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

@@ -1,5 +1,5 @@
 import torch
-from typing import List, Tuple
+from typing import List, Tuple, Callable, Optional
 from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend
 from lightllm.common.basemodel.batch_objs import ModelOutput
 
@@ -52,7 +52,9 @@ def decode(self):
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
         return
 
-    def normal_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int, uninit_reqs, ok_finished_reqs):
+    def normal_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int, uninit_reqs, ok_finished_reqs,
+                            extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
+                            call_post_handle_for_chunk: bool = False):
         model_input, run_reqs, padded_req_num = padded_prepare_prefill_inputs(
             prefill_reqs, is_multimodal=self.is_multimodal
         )
@@ -65,7 +67,9 @@ def normal_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int
             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=True, do_filter_finished_reqs=False
+                run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=True, do_filter_finished_reqs=False,
+                extra_post_req_handle_func=extra_post_req_handle_func,
+                call_post_handle_for_chunk=call_post_handle_for_chunk
             )
         return
 
@@ -117,7 +121,9 @@ def overlap_decode(self, decode_reqs: List[InferReq], max_decode_num: int, unini
             )
         return
 
-    def overlap_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int, uninit_reqs, ok_finished_reqs):
+    def overlap_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int, uninit_reqs, ok_finished_reqs,
+                            extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
+                            call_post_handle_for_chunk: bool = False):
         (
             micro_input,
             run_reqs,
@@ -142,6 +148,8 @@ def overlap_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: in
             next_token_ids = next_token_ids.detach().cpu().numpy()
             next_token_logprobs = torch.log(next_token_probs).detach().cpu().numpy()
             self._post_handle(
-                all_run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=True, do_filter_finished_reqs=False
+                all_run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=True, do_filter_finished_reqs=False,
+                extra_post_req_handle_func=extra_post_req_handle_func,
+                call_post_handle_for_chunk=call_post_handle_for_chunk
             )
         return

lightllm/server/router/model_infer/mode_backend/pd_nixl/impl_for_pd_decode.py

@@ -1,15 +1,12 @@
 import time
-import torch
 import torch.multiprocessing as mp
 import threading
 from concurrent.futures import ThreadPoolExecutor
-from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend
+from lightllm.server.router.model_infer.mode_backend.continues_batch.impl import ContinuesBatchBackend
 from typing import List, Tuple, Dict
 from lightllm.server.router.model_infer.infer_batch import g_infer_context, InferReq
 from lightllm.server.core.objs.req import PDNIXLChunkedPrefillReq
 from lightllm.utils.log_utils import init_logger
-from lightllm.server.router.model_infer.mode_backend.generic_pre_process import prepare_decode_inputs
-from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample
 from lightllm.server.multimodal_params import MultimodalParams
 
 from .pd_remote_prefill_obj import RemotePrefillTask, RemotePrefillServerInfo, RemotePrefillRequest, RemoteTransferStatusType
@@ -93,21 +90,8 @@ def decode(self):
                 self.remote_prefilled_reqs[shm_req.group_req_id] = run_req
 
         if decode_reqs:
-            kwargs, run_reqs = prepare_decode_inputs(decode_reqs)
-            logits = self.model.forward(**kwargs)
-
-            self._overlap_req_init_and_filter(
-                uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True
-            )
-
-            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
-            )
+            ContinuesBatchBackend.normal_decode(
+                self, decode_reqs=decode_reqs, uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs)
 
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
-
         return

lightllm/server/router/model_infer/mode_backend/pd_nixl/impl_for_pd_decode_dp.py

@@ -1,14 +1,11 @@
 import time
 import torch
 import torch.multiprocessing as mp
-import torch.distributed as dist
-from typing import List
 from lightllm.server.router.model_infer.infer_batch import g_infer_context, InferReq
 from lightllm.server.core.objs.req import PDNIXLChunkedPrefillReq
 from lightllm.utils.log_utils import init_logger
-from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample
 from lightllm.utils.envs_utils import get_env_start_args
-from lightllm.server.router.model_infer.mode_backend.dp_backend.pre_process import padded_prepare_decode_inputs
+from lightllm.server.router.model_infer.mode_backend.dp_backend.impl import DPChunkedPrefillBackend
 
 from .impl_for_pd_decode import PDNIXLBackendForDecodeNode, RemoteTransferStatusType
 
@@ -24,7 +21,7 @@ def init_custom(self):
         super().init_custom()
 
         self.reduce_tensor = torch.tensor([0], dtype=torch.int32, device="cuda", requires_grad=False)
-        from lightllm.server.router.model_infer.mode_backend.dp_backend.pre_process import padded_prepare_prefill_inputs
+        from lightllm.server.router.model_infer.mode_backend.pre import padded_prepare_prefill_inputs
 
         kwargs, run_reqs, padded_req_num = padded_prepare_prefill_inputs([], 1, is_multimodal=self.is_multimodal)
         self.model.forward(**kwargs)
@@ -63,62 +60,13 @@ def decode(self):
                 run_req.in_prefill_or_transfer = True
                 self.remote_prefilled_reqs[shm_req.group_req_id] = run_req
 
-        self.reduce_tensor.fill_(len(decode_reqs))
-        dist.all_reduce(self.reduce_tensor, op=dist.ReduceOp.MAX)
-        max_decode_num = self.reduce_tensor.item()
+        max_decode_num = self._dp_all_reduce_decode_req_num(decode_reqs=decode_reqs)
         if max_decode_num != 0:
             if not self.enable_decode_microbatch_overlap:
-                self.normal_decode(decode_reqs, max_decode_num, uninit_reqs, ok_finished_reqs)
+                DPChunkedPrefillBackend.normal_decode(self, decode_reqs, max_decode_num, uninit_reqs, ok_finished_reqs)
             else:
-                self.overlap_decode(decode_reqs, max_decode_num, uninit_reqs, ok_finished_reqs)
-        self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
-        return
-
-    def normal_decode(self, decode_reqs: List[InferReq], max_decode_num: int, uninit_reqs, ok_finished_reqs):
+                DPChunkedPrefillBackend.overlap_decode(self, decode_reqs, max_decode_num, uninit_reqs, ok_finished_reqs)
 
-        kwargs, run_reqs, padded_req_num = padded_prepare_decode_inputs(
-            decode_reqs, max_decode_num, is_multimodal=self.is_multimodal
-        )
-        logits = self.model.forward(**kwargs)
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
-        if len(run_reqs) != 0:
-            logits = logits[0 : len(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()
-            self._post_handle(
-                run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=False, do_filter_finished_reqs=False
-            )
         return
 
-    def overlap_decode(self, decode_reqs: List[InferReq], max_decode_num: int, uninit_reqs, ok_finished_reqs):
-        from lightllm.server.router.model_infer.mode_backend.dp_backend.pre_process import (
-            padded_overlap_prepare_decode_inputs,
-        )
-
-        (
-            micro_batch,
-            run_reqs,
-            padded_req_num,
-            micro_batch1,
-            run_reqs1,
-            padded_req_num1,
-        ) = padded_overlap_prepare_decode_inputs(decode_reqs, max_decode_num, is_multimodal=self.is_multimodal)
-
-        logits, logits1 = self.model.microbatch_overlap_decode(micro_batch, micro_batch1)
-        self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
-        req_num, req_num1 = len(run_reqs), len(run_reqs1)
-        all_logits = torch.empty((req_num + req_num1, logits.shape[1]), dtype=logits.dtype, device=logits.device)
-
-        all_logits[0:req_num, :].copy_(logits[0:req_num, :], non_blocking=True)
-        all_logits[req_num : (req_num + req_num1), :].copy_(logits1[0:req_num1, :], non_blocking=True)
-
-        all_run_reqs = run_reqs + run_reqs1
-        if all_run_reqs:
-            next_token_ids, next_token_probs = sample(all_logits, all_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(
-                all_run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=False, do_filter_finished_reqs=False
-            )
-        return

lightllm/server/router/model_infer/mode_backend/pd_nixl/impl_for_pd_prefill.py

@@ -1,12 +1,9 @@
 import threading
-import torch
 import torch.multiprocessing as mp
 from typing import List, Tuple
-from lightllm.utils.infer_utils import calculate_time, mark_start, mark_end
-from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context
+from lightllm.server.router.model_infer.infer_batch import g_infer_context
 from lightllm.utils.log_utils import init_logger
-from lightllm.server.router.model_infer.mode_backend.generic_pre_process import prepare_prefill_inputs
-from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample
+from lightllm.server.router.model_infer.mode_backend.continues_batch.impl import ContinuesBatchBackend
 from .impl_for_pd_base import PDNIXLBackendBase
 
 logger = init_logger(__name__)
@@ -49,25 +46,10 @@ def decode(self):
         assert len(decode_reqs) == 0
 
         self._prefill_abort_remote(aborted_reqs)
-        self._filter_reqs(aborted_reqs + ok_finished_reqs)
+        self._filter_reqs(aborted_reqs)
 
         if prefill_reqs:
-            kwargs, run_reqs = prepare_prefill_inputs(
-                prefill_reqs, is_chuncked_mode=True, is_multimodal=self.is_multimodal
-            )
-
-            logits = self.model.forward(**kwargs)
-            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=True,
-                do_filter_finished_reqs=False,
-                extra_post_req_handle_chunk_func=self._handle_chunked_transfer,
-            )
+            ContinuesBatchBackend.normal_prefill_reqs(
+                self, prefill_reqs=prefill_reqs, uninit_reqs=uinit_reqs, ok_finished_reqs=ok_finished_reqs,
+                extra_post_req_handle_func=self._handle_chunked_transfer, call_post_handle_for_chunk=True)
         return