opt: reduce gpu memeory alloc for req_param_cache

Author: niushengxiao

lightllm/common/basemodel/triton_kernel/apply_penalty_cache.py

@@ -7,7 +7,7 @@
 
 
 @triton.jit
-def _fwd_kernel_apply_penalty_cache(
+def _kernel_apply_penalty_cache(
     Logits,
     req_idxs,
     presence_penalty,
@@ -27,7 +27,7 @@ def _fwd_kernel_apply_penalty_cache(
 
     batch_ids = BLOCK_P * block_idx + tl.arange(0, BLOCK_P)
     batch_ids_count = tl.load(
-        p_token_vocabs + token_idx * stride_p_token_vocabs_b + batch_ids,
+        p_token_vocabs + cur_batch * stride_p_token_vocabs_b + batch_ids,
         mask=batch_ids < stride_p_token_vocabs_b,
         other=0,
     )
@@ -43,7 +43,7 @@ def _fwd_kernel_apply_penalty_cache(
 
 
 @triton.jit
-def _eos_penalty(
+def _kernel_eos_penalty(
     Logits,
     req_idxs,
     p_token_lens,
@@ -71,26 +71,58 @@ def _eos_penalty(
     return
 
 
+@triton.jit
+def _kernel_bincount(
+    req_idxs,
+    input,
+    output,
+    input_lens,
+    stride_input_b,
+    stride_output_b,
+    BLOCK_SIZE: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    req_idx = tl.load(req_idxs + cur_batch)
+    block_idx = tl.program_id(1)
+    input_ptr = input + req_idx * stride_input_b + block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    input_len = tl.load(input_lens + req_idx)
+    mask = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < input_len
+    token_id = tl.load(input_ptr, mask=mask, other=0)
+    tl.atomic_add(output + cur_batch * stride_output_b + token_id, 1, mask=mask)
+    return
+
+
 @torch.no_grad()
 def apply_penalty_cache(
     Logits,
     req_idxs,
     presence_penalty,
     freqency_penalty,
     repetition_penalty,
-    p_token_vocabs,
+    p_token_ids,
     p_token_lens,
     exponential_decay_length_penalties,
     length_penalty_idx,
     eos_ids,
     mask_eos_reqs,
-    is_eos_penalty=False,
+    vocab_size: tl.constexpr,
+    is_eos_penalty: tl.constexpr = False,
 ):
     assert Logits.is_contiguous()
     BLOCK_P = 1024
-    num_warps = 8
-    vocab_size = p_token_vocabs.shape[1]
-    _fwd_kernel_apply_penalty_cache[(Logits.shape[0], triton.cdiv(vocab_size, BLOCK_P))](
+    num_warps = 4
+    p_token_vocabs = torch.zeros((Logits.shape[0], vocab_size), dtype=torch.int32, device="cuda")
+    _kernel_bincount[(Logits.shape[0], triton.cdiv(p_token_ids.stride(0), BLOCK_P))](
+        req_idxs,
+        p_token_ids,
+        p_token_vocabs,
+        p_token_lens,
+        p_token_ids.stride(0),
+        p_token_vocabs.stride(0),
+        num_warps=num_warps,
+        BLOCK_SIZE=BLOCK_P,
+    )
+    _kernel_apply_penalty_cache[(Logits.shape[0], triton.cdiv(vocab_size, BLOCK_P))](
         Logits,
         req_idxs,
         presence_penalty,
@@ -103,8 +135,7 @@ def apply_penalty_cache(
         BLOCK_P=BLOCK_P,
     )
     if is_eos_penalty:
-        p_token_lens = p_token_vocabs[req_idxs].sum(dim=1).cuda() if p_token_lens is None else p_token_lens
-        _eos_penalty[(Logits.shape[0],)](
+        _kernel_eos_penalty[(Logits.shape[0],)](
             Logits,
             req_idxs,
             p_token_lens,
@@ -121,11 +152,13 @@ def apply_penalty_cache(
 
 if __name__ == "__main__":
     from .apply_penalty import apply_penalty
+    from torch.nn.utils.rnn import pad_sequence
 
     bs = 200
     vocab_size = 150000
-    p_tokens = 2000
-    repseats = 1000
+    p_tokens = 3000
+    max_token_len = 16384
+    repseats = max_token_len // p_tokens
     req_idxs = torch.arange(bs).cuda()
     logits = torch.randn((bs, vocab_size), dtype=torch.float32).cuda()
     logits2 = logits.clone()
@@ -144,7 +177,7 @@ def apply_penalty_cache(
         i += s_l
     p_token_counts = torch.randint(1, repseats, (p_seq_len.sum(),)).cuda()
     p_cumsum_seq_len = p_seq_len.cumsum(dim=0).cuda()
-    p_token_vocabs = torch.zeros((bs, vocab_size), dtype=torch.int16).cuda()
+    p_token_vocabs = torch.zeros((bs, vocab_size), dtype=torch.int32).cuda()
     i = 0
     b = 0
     for token_id, count in zip(p_token_ids, p_token_counts):
@@ -154,7 +187,12 @@ def apply_penalty_cache(
             b += 1
             i = 0
 
-    p_token_lens = p_token_vocabs.sum(dim=1).cuda()
+    p_token_lens = p_token_vocabs.cuda().sum(dim=1)
+    assert p_token_lens.max() < max_token_len
+    token_idx = torch.arange(vocab_size).cuda()
+    sequences = [token_idx.repeat_interleave(p_token_vocabs[b]) for b in range(bs)]  # shape = [sum(rep_nums)]
+    p_token_mat = pad_sequence(sequences, batch_first=True, padding_value=0).to(torch.int32).cuda()
+
     length_penalty_idx = torch.randint(0, p_tokens, (bs,)).cuda()
     len_idx = torch.tensor([max(p_token_lens[i] - length_penalty_idx[i], 0) for i in range(bs)]).cuda()
     mask_eos_reqs = torch.randint(1, p_tokens, (bs,)).cuda()
@@ -180,12 +218,13 @@ def apply_penalty_cache(
         presence_penalty,
         freqency_penalty,
         repetition_penalty,
-        p_token_vocabs,
+        p_token_mat,
         p_token_lens,
         exponential_decay_length_penalties,
         length_penalty_idx,
         eos_ids,
         mask_eos_reqs,
+        vocab_size,
     )
     fn1()
     fn2()

lightllm/common/req_manager.py

@@ -60,6 +60,7 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana
         self.mem_manager = mem_manager
         self.req_sample_parms_manager = None
         self.max_request_num = max_request_num
+        self.max_sequence_length = max_sequence_length
         self.HOLD_REQUEST_ID = max_request_num
 
     def alloc(self):
@@ -69,7 +70,7 @@ 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
+            self.req_sample_parms_manager.p_token_lens[free_req_indexes] = 0
 
         if self.req_list.is_all_free():
             logger.debug(f"freed all request size {self.req_list.can_alloc_size}")

lightllm/server/router/model_infer/infer_batch.py

@@ -22,7 +22,7 @@
 
 
 class ReqSampleParmsManager:
-    def __init__(self, max_request_num, vocab_size):
+    def __init__(self, max_request_num, max_seq_len, vocab_size):
         self.presence_penalties = torch.empty(max_request_num, dtype=torch.float, device="cpu", pin_memory=True).cuda(
             non_blocking=True
         )
@@ -50,9 +50,13 @@ def __init__(self, max_request_num, vocab_size):
         self.mask_eos_reqs = torch.empty(max_request_num, dtype=torch.int32, device="cpu", pin_memory=True).cuda(
             non_blocking=True
         )
-        self.p_token_vocabs = torch.zeros(
-            (max_request_num, vocab_size), dtype=torch.int16, device="cpu", pin_memory=True
+        self.p_token_ids = torch.zeros(
+            (max_request_num, max_seq_len), dtype=torch.int32, device="cpu", pin_memory=True
         ).cuda(non_blocking=True)
+        self.p_token_lens = torch.zeros((max_request_num,), dtype=torch.int32, device="cpu", pin_memory=True).cuda(
+            non_blocking=True
+        )
+        self.vocab_size = vocab_size
 
     def set_sample_params(self, req_idx, sampling_param):
         self.presence_penalties[req_idx] = sampling_param.shm_param.presence_penalty
@@ -83,7 +87,9 @@ def register(
         self, req_manager: ReqManager, radix_cache: RadixCache, shm_req_manager: ShmReqManager, vocab_size: int
     ):
         if os.getenv("ENABLE_REQ_PARAM_CACHE", False):
-            req_manager.req_sample_parms_manager = ReqSampleParmsManager(req_manager.max_request_num, vocab_size)
+            req_manager.req_sample_parms_manager = ReqSampleParmsManager(
+                req_manager.max_request_num, req_manager.max_sequence_length, vocab_size
+            )
         self.req_manager = req_manager
         self.radix_cache = radix_cache
         self.shm_req_manager = shm_req_manager
@@ -317,9 +323,11 @@ def init_all(self):
                     self.req_idx, self.sampling_param
                 )
                 if self.sampling_param.shm_param.input_penalty:
-                    dct = collections.Counter(self.shm_req.get_prompt_ids())
-                    for idx, count in dct.items():
-                        g_infer_context.req_manager.req_sample_parms_manager.p_token_vocabs[self.req_idx][idx] = count
+                    ids_len = len(self.shm_req.get_prompt_ids())
+                    g_infer_context.req_manager.req_sample_parms_manager.p_token_ids[self.req_idx][
+                        :ids_len
+                    ] = self.shm_req.get_prompt_ids()
+                    g_infer_context.req_manager.req_sample_parms_manager.p_token_lens[self.req_idx] = ids_len
 
             if self.sampling_param.shm_param.input_penalty:
                 self.out_token_id_count = collections.Counter(self.shm_req.get_prompt_ids())

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

@@ -161,12 +161,13 @@ def sample_in_cache(logits, reqs, eos_id: List[int] = [2]):
         params.presence_penalties,
         params.frequency_penalties,
         params.repetition_penalties,
-        params.p_token_vocabs,
-        None,
+        params.p_token_ids,
+        params.p_token_lens,
         params.exponential_decay_length_penalties,
         params.length_penalty_idx,
         eos_ids,
         params.mask_eos_reqs,
+        params.vocab_size,
     )
 
     logits.div_(params.temperatures[req_idxs].view((-1, 1)))
@@ -209,4 +210,6 @@ def _get_req_idxs(reqs: List[InferReq]):
 
 
 def _update_repeatition_tokens(req_idxs, token_ids):
-    g_infer_context.req_manager.req_sample_parms_manager.p_token_vocabs[req_idxs, token_ids] += 1
+    token_idxs = g_infer_context.req_manager.req_sample_parms_manager.p_token_lens[req_idxs]
+    g_infer_context.req_manager.req_sample_parms_manager.p_token_ids[req_idxs, token_idxs] = token_ids
+    g_infer_context.req_manager.req_sample_parms_manager.p_token_lens[req_idxs] += 1