opt: add token vocab cache in the post processing

Author: niushengxiao

lightllm/common/basemodel/triton_kernel/apply_penalty.py

@@ -19,9 +19,9 @@ def _fwd_kernel_apply_penalty(
     eos_ids,
     mask_eos_reqs,
     stride_logit_b,
-    stride_logit_s,
     BLOCK_P: tl.constexpr,
     EOS_ID_NUM: tl.constexpr,
+    IS_EOS_PENALTY: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     cur_freqency = tl.load(freqency_penalty + cur_batch)
@@ -46,18 +46,19 @@ def _fwd_kernel_apply_penalty(
         output_ptr = Logits + cur_batch * stride_logit_b + batch_ids
         tl.store(output_ptr, pre_logits, mask=cur_batch_id_offset < cur_batch_end_index)
 
-    mask_eos = tl.load(mask_eos_reqs + cur_batch)
-    exponential_decay_length_penalty = tl.load(exponential_decay_length_penalties + cur_batch)
-    length_penalty = tl.load(length_penalty_idx + cur_batch)
-    penalty_scale = tl.exp2(tl.log2(exponential_decay_length_penalty) * length_penalty) - 1
+    if IS_EOS_PENALTY:
+        mask_eos = tl.load(mask_eos_reqs + cur_batch)
+        exponential_decay_length_penalty = tl.load(exponential_decay_length_penalties + cur_batch)
+        length_penalty = tl.load(length_penalty_idx + cur_batch)
+        penalty_scale = tl.exp2(tl.log2(exponential_decay_length_penalty) * length_penalty) - 1
 
-    for eos_index in range(EOS_ID_NUM):
-        eos_id = tl.load(eos_ids + eos_index)
-        cur_eos_logit_ptr = Logits + cur_batch * stride_logit_b + eos_id
-        cur_eos_logit = tl.load(cur_eos_logit_ptr)
-        cur_eos_logit = cur_eos_logit + tl.abs(cur_eos_logit) * penalty_scale
-        cur_eos_logit = tl.where(mask_eos, -10000000.0, cur_eos_logit)
-        tl.store(cur_eos_logit_ptr, cur_eos_logit)
+        for eos_index in range(EOS_ID_NUM):
+            eos_id = tl.load(eos_ids + eos_index)
+            cur_eos_logit_ptr = Logits + cur_batch * stride_logit_b + eos_id
+            cur_eos_logit = tl.load(cur_eos_logit_ptr)
+            cur_eos_logit = cur_eos_logit + tl.abs(cur_eos_logit) * penalty_scale
+            cur_eos_logit = tl.where(mask_eos, -10000000.0, cur_eos_logit)
+            tl.store(cur_eos_logit_ptr, cur_eos_logit)
     return
 
 
@@ -74,6 +75,7 @@ def apply_penalty(
     length_penalty_idx,
     eos_ids,
     mask_eos_reqs,
+    is_eos_penalty=False,
 ):
     assert Logits.is_contiguous()
     BLOCK_P = 1024
@@ -91,9 +93,9 @@ def apply_penalty(
         eos_ids,
         mask_eos_reqs,
         Logits.stride(0),
-        Logits.stride(1),
         num_warps=num_warps,
         BLOCK_P=BLOCK_P,
         EOS_ID_NUM=eos_ids.shape[0],
+        IS_EOS_PENALTY=is_eos_penalty,
     )
     return

lightllm/common/basemodel/triton_kernel/apply_penalty_cache.py

@@ -0,0 +1,195 @@
+import torch
+
+import triton
+import triton.language as tl
+import torch.nn.functional as F
+import numpy as np
+
+
+@triton.jit
+def _fwd_kernel_apply_penalty_cache(
+    Logits,
+    req_idxs,
+    presence_penalty,
+    freqency_penalty,
+    repetition_penalty,
+    p_token_vocabs,
+    stride_logit_b,
+    stride_p_token_vocabs_b,
+    BLOCK_P: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    block_idx = tl.program_id(1)
+    token_idx = tl.load(req_idxs + cur_batch)
+    cur_freqency = tl.load(freqency_penalty + cur_batch)
+    cur_presence = tl.load(presence_penalty + cur_batch)
+    cur_repetition = tl.load(repetition_penalty + cur_batch)
+
+    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,
+        mask=batch_ids < stride_p_token_vocabs_b,
+        other=0,
+    )
+    row_start_ptr = Logits + cur_batch * stride_logit_b
+    cur_offset = row_start_ptr + batch_ids
+    cur_logits = tl.load(cur_offset, mask=batch_ids_count > 0, other=0.0)
+    rep_logits = tl.where(cur_logits > 0, cur_logits / cur_repetition, cur_logits * cur_repetition)
+    freq_logits = rep_logits - batch_ids_count * cur_freqency
+    pre_logits = freq_logits - cur_presence
+    output_ptr = Logits + cur_batch * stride_logit_b + batch_ids
+    tl.store(output_ptr, pre_logits, mask=batch_ids_count > 0)
+    return
+
+
+@triton.jit
+def _eos_penalty(
+    Logits,
+    p_token_lens,
+    exponential_decay_length_penalties,
+    length_penalty_idx,
+    eos_ids,
+    mask_eos_reqs,
+    stride_logit_b,
+    EOS_ID_NUM: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    exponential_decay_length_penalty = tl.load(exponential_decay_length_penalties + cur_batch)
+    token_lens = tl.load(p_token_lens + cur_batch)
+    length_penalty = tl.maximum(token_lens - tl.load(length_penalty_idx + cur_batch), 0)
+    penalty_scale = tl.exp2(tl.log2(exponential_decay_length_penalty) * length_penalty) - 1
+    mask_eos = tl.load(mask_eos_reqs + cur_batch)
+    for eos_index in range(EOS_ID_NUM):
+        eos_id = tl.load(eos_ids + eos_index)
+        cur_eos_logit_ptr = Logits + cur_batch * stride_logit_b + eos_id
+        cur_eos_logit = tl.load(cur_eos_logit_ptr)
+        cur_eos_logit = cur_eos_logit + tl.abs(cur_eos_logit) * penalty_scale
+        cur_eos_logit = tl.where(token_lens < mask_eos, -10000000.0, cur_eos_logit)
+        tl.store(cur_eos_logit_ptr, cur_eos_logit)
+    return
+
+
+@torch.no_grad()
+def apply_penalty_cache(
+    Logits,
+    req_idxs,
+    presence_penalty,
+    freqency_penalty,
+    repetition_penalty,
+    p_token_vocabs,
+    p_token_lens,
+    exponential_decay_length_penalties,
+    length_penalty_idx,
+    eos_ids,
+    mask_eos_reqs,
+    is_eos_penalty=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))](
+        Logits,
+        req_idxs,
+        presence_penalty,
+        freqency_penalty,
+        repetition_penalty,
+        p_token_vocabs,
+        Logits.stride(0),
+        p_token_vocabs.stride(0),
+        num_warps=num_warps,
+        BLOCK_P=BLOCK_P,
+    )
+    if is_eos_penalty:
+        p_token_lens = p_token_vocabs[req_idxs].count_nonzero(dim=1) if p_token_lens is None else p_token_lens
+        _eos_penalty[(Logits.shape[0],)](
+            Logits,
+            p_token_lens,
+            exponential_decay_length_penalties,
+            length_penalty_idx,
+            eos_ids,
+            mask_eos_reqs,
+            Logits.stride(0),
+            num_warps=num_warps,
+            EOS_ID_NUM=eos_ids.shape[0],
+        )
+    return
+
+
+if __name__ == "__main__":
+    from .apply_penalty import apply_penalty
+
+    bs = 200
+    vocab_size = 150000
+    p_tokens = 2000
+    repseats = 1000
+    req_idxs = torch.arange(bs).cuda()
+    logits = torch.randn((bs, vocab_size), dtype=torch.float32).cuda()
+    logits2 = logits.clone()
+
+    presence_penalty = torch.randn((bs,), dtype=torch.float32).cuda() + 1e-5
+    freqency_penalty = torch.randn((bs,), dtype=torch.float32).cuda()
+    repetition_penalty = torch.randn((bs,), dtype=torch.float32).cuda()
+    exponential_decay_length_penalties = torch.rand(bs).cuda()
+    eos_ids = torch.tensor([999]).cuda()
+
+    p_seq_len = torch.cat([torch.tensor([0]), torch.randint(1, p_tokens, (bs,))]).cuda()
+    p_token_ids = torch.randint(0, vocab_size, (p_seq_len.sum(),)).cuda()
+    i = 0
+    for s_l in p_seq_len[1:]:
+        p_token_ids[i : i + s_l] = torch.randperm(vocab_size)[:s_l]
+        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()
+    i = 0
+    b = 0
+    for token_id, count in zip(p_token_ids, p_token_counts):
+        p_token_vocabs[b][token_id] = count
+        i += 1
+        if i == p_seq_len[b + 1]:
+            b += 1
+            i = 0
+
+    p_token_lens = p_token_vocabs.sum(dim=1).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()
+    mask_bool = torch.tensor([p_token_lens[i] < mask_eos_reqs[i] for i in range(bs)]).cuda()
+
+    fn1 = lambda: apply_penalty(
+        logits,
+        presence_penalty,
+        freqency_penalty,
+        repetition_penalty,
+        p_token_ids,
+        p_token_counts,
+        p_cumsum_seq_len,
+        exponential_decay_length_penalties,
+        len_idx,
+        eos_ids,
+        mask_bool,
+    )
+
+    fn2 = lambda: apply_penalty_cache(
+        logits2,
+        req_idxs,
+        presence_penalty,
+        freqency_penalty,
+        repetition_penalty,
+        p_token_vocabs,
+        p_token_lens,
+        exponential_decay_length_penalties,
+        length_penalty_idx,
+        eos_ids,
+        mask_eos_reqs,
+    )
+    fn1()
+    fn2()
+    cos = F.cosine_similarity(logits, logits2).mean()
+    print("cos =", cos)
+    assert torch.allclose(logits, logits2, atol=1e-2, rtol=0)
+
+    ms1 = triton.testing.do_bench(fn1)
+    ms2 = triton.testing.do_bench(fn2)
+    print("ms1 =", ms1, "ms2 =", ms2)

lightllm/common/req_manager.py

@@ -58,6 +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.max_request_num = max_request_num
         self.HOLD_REQUEST_ID = max_request_num
 
@@ -67,6 +68,8 @@ 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}")

lightllm/server/router/model_infer/infer_batch.py

@@ -17,10 +17,16 @@
 from lightllm.common.basemodel.infer_lock import g_infer_state_lock
 from lightllm.server.multimodal_params import MultimodalParams
 from lightllm.utils.custom_kernel_utis import custom_cat
+from lightllm.utils.envs_utils import enable_env_vars
 
 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 请求管理
@@ -36,6 +42,8 @@ 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
@@ -55,7 +63,6 @@ def get_overlap_stream(self) -> torch.cuda.Stream:
     def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_req_obj=True):
         request_ids = []
         for r in requests:
-
             r_id, r_index, multimodal_params, _ = r
             if r_id not in self.requests_mapping.keys():
                 r_obj = InferReq(
@@ -264,10 +271,19 @@ def init_all(self):
             self.shm_req.link_prompt_ids_shm_array()
             self.shm_req.link_logprobs_shm_array()
             self.sampling_param: InferSamplingParams = InferSamplingParams(self.shm_req, self.vocab_size)
-            if self.sampling_param.shm_param.input_penalty:
-                self.out_token_id_count = collections.Counter(self.shm_req.get_prompt_ids())
+
+            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
             else:
-                self.out_token_id_count = collections.defaultdict(int)
+                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.stop_sequences = self.sampling_param.shm_param.stop_sequences.to_list()
             # token healing mode 才被使用的管理对象

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

@@ -282,7 +282,8 @@ def _post_handle(
             req_obj.set_next_gen_token_id(next_token_id, next_token_logprob)
             req_obj.cur_output_len += 1
 
-            req_obj.out_token_id_count[next_token_id] += 1
+            if req_obj.out_token_id_count is not None:
+                req_obj.out_token_id_count[next_token_id] += 1
             req_obj.update_finish_status(self.eos_id)
 
             if extra_post_req_handle_func is not None:

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

@@ -55,7 +55,8 @@ def prefill(self, run_reqs: List[Tuple]):
             for i in range(req_obj.shm_req.input_len - 1):
                 req_obj.shm_req.shm_logprobs.arr[i + 1] = cur_logprobs[i]
 
-            req_obj.out_token_id_count[next_token_id] += 1
+            if req_obj.out_token_id_count is not None:
+                req_obj.out_token_id_count[next_token_id] += 1
             req_obj.update_finish_status(self.eos_id)
 
             if req_obj.finish_status.is_finished() or req_obj.shm_req.router_aborted:

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

@@ -32,7 +32,8 @@ def prefill(self, reqs: List[Tuple]):
             req_obj.set_next_gen_token_id(next_token_id, next_token_logprob)
             req_obj.cur_output_len += 1
 
-            req_obj.out_token_id_count[next_token_id] += 1
+            if req_obj.out_token_id_count is not None:
+                req_obj.out_token_id_count[next_token_id] += 1
             req_obj.update_finish_status(self.eos_id)
 
             if req_obj.finish_status.is_finished() or req_obj.shm_req.router_aborted: