improve postprocess

Author: hiworldwzj

lightllm/common/req_manager.py

@@ -58,7 +58,7 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana
             (max_request_num + 1, max_sequence_length), dtype=torch.int32, device="cuda"
         )
         self.mem_manager = mem_manager
-        self.req_sample_parms_manager = None
+        self.req_sampling_params_manager = ReqSamplingParamsManager(max_request_num)
         self.max_request_num = max_request_num
         self.HOLD_REQUEST_ID = max_request_num
 
@@ -68,8 +68,6 @@ def alloc(self):
     def free(self, free_req_indexes: List[int], free_token_index):
         for req_index in free_req_indexes:
             self.req_list.free(req_index)
-        if self.req_sample_parms_manager is not None:
-            self.req_sample_parms_manager.p_token_vocabs[free_req_indexes] = 0
 
         if self.req_list.is_all_free():
             logger.debug(f"freed all request size {self.req_list.can_alloc_size}")
@@ -88,3 +86,38 @@ def free_token(self, free_token_index):
     def free_all(self):
         self.req_list = _ReqLinkedList(self.max_request_num)
         return
+
+
+class ReqSamplingParamsManager:
+    """
+    ReqSamplingParamsManager 将输出采样参数中，确定比较固定的部分，纳入到 gpu buffer中进行管理，这样可以更快捷的
+    利用cuda kernel 将采样参数提取为以batch 为单位的采样参数，对于哪些比较动态，或者存在特殊处理的后处理参数，
+    则保留从 InferSamplingParams 中进行动态读取和动态组batch， 具体使用可以参考
+    lightllm/server/router/model_infer/mode_backend/generic_post_process.py 文件中的使用方式。
+    """
+
+    def __init__(self, max_request_num):
+        self.req_to_presence_penalty = torch.zeros(max_request_num + 1, dtype=torch.float32, device="cuda")
+        self.req_to_frequency_penalty = torch.zeros(max_request_num + 1, dtype=torch.float32, device="cuda")
+        self.req_to_repetition_penalty = torch.zeros(max_request_num + 1, dtype=torch.float32, device="cuda")
+        self.req_to_temperature = torch.zeros(max_request_num + 1, dtype=torch.float32, device="cuda")
+        self.req_to_exponential_decay_length_penalty = torch.zeros(
+            max_request_num + 1, dtype=torch.float32, device="cuda"
+        )
+
+    def init_req_sampling_params(self, req):
+        # fix cycle loop import
+        from lightllm.server.router.model_infer.infer_batch import InferReq
+
+        req: InferReq = req
+
+        shm_param = req.sampling_param.shm_param
+        self.req_to_presence_penalty[req.req_idx].fill_(shm_param.presence_penalty)
+        self.req_to_frequency_penalty[req.req_idx].fill_(shm_param.frequency_penalty)
+        self.req_to_repetition_penalty[req.req_idx].fill_(shm_param.repetition_penalty)
+        self.req_to_temperature[req.req_idx].fill_(shm_param.temperature)
+        exponential_decay_length_penalty = shm_param.exponential_decay_length_penalty.to_tuple()
+        self.req_to_exponential_decay_length_penalty[req.req_id].fill_(exponential_decay_length_penalty[1])
+
+    def get_sampling_batch_params(self, req_idx_list: List[int]):
+        pass

lightllm/server/router/model_infer/infer_batch.py

@@ -22,11 +22,6 @@
 logger = init_logger(__name__)
 
 
-class ReqSampleParmsManager:
-    def __init__(self, max_request_num, vocab_size):
-        self.p_token_vocabs = torch.zeros((max_request_num, vocab_size), dtype=torch.int16, device="cuda")
-
-
 @dataclass
 class InferenceContext:
     req_manager: ReqManager = None  # gpu 请求管理
@@ -42,8 +37,6 @@ class InferenceContext:
     def register(
         self, req_manager: ReqManager, radix_cache: RadixCache, shm_req_manager: ShmReqManager, vocab_size: int
     ):
-        if enable_env_vars("ENABLE_REQ_PARAM_CACHE"):
-            req_manager.req_sample_parms_manager = ReqSampleParmsManager(req_manager.max_request_num, vocab_size)
         self.req_manager = req_manager
         self.radix_cache = radix_cache
         self.shm_req_manager = shm_req_manager
@@ -272,18 +265,11 @@ def init_all(self):
             self.shm_req.link_logprobs_shm_array()
             self.sampling_param: InferSamplingParams = InferSamplingParams(self.shm_req, self.vocab_size)
 
-            if enable_env_vars("ENABLE_REQ_PARAM_CACHE"):
-                if self.sampling_param.shm_param.input_penalty:
-                    idxs = torch.bincount(self.shm_req.get_prompt_ids())
-                    g_infer_context.req_manager.req_sample_parms_manager.p_token_vocabs[self.req_idx][
-                        : len(idxs)
-                    ] = idxs
-                self.out_token_id_count = None
+            g_infer_context.req_manager.req_sampling_params_manager.init_req_sampling_params(self)
+            if self.sampling_param.shm_param.input_penalty:
+                self.out_token_id_count = collections.Counter(self.shm_req.get_prompt_ids())
             else:
-                if self.sampling_param.shm_param.input_penalty:
-                    self.out_token_id_count = collections.Counter(self.shm_req.get_prompt_ids())
-                else:
-                    self.out_token_id_count = collections.defaultdict(int)
+                self.out_token_id_count = collections.defaultdict(int)
 
             self.stop_sequences = self.sampling_param.shm_param.stop_sequences.to_list()
             # token healing mode 才被使用的管理对象