feat: add b_last_mem_indx in the InferReq

Author: niushengxiao

lightllm/common/basemodel/basemodel.py

@@ -687,9 +687,7 @@ def _check_max_len_infer(self):
             b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda")
             b_seq_len[:] = self.batch_max_tokens
             b_ready_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda")
-            mem_indexes = self.req_manager.alloc_mem_indices(
-                len(dummy_input_ids), b_req_idx, b_seq_len, b_ready_cache_len
-            ).cuda()
+            mem_indexes = self.req_manager.alloc_mem_indices(len(dummy_input_ids), b_seq_len, b_ready_cache_len).cuda()
             total_token_num = self.batch_max_tokens
             b_mtp_index = torch.zeros(1, dtype=torch.int32, device="cuda")
             model_input = ModelInput(
@@ -765,7 +763,7 @@ def _autotune_warmup(self):
                 total_token_num = input_len
                 b_mtp_index = torch.zeros(1, dtype=torch.int32, device="cuda")
                 mem_indexes = self.req_manager.alloc_mem_indices(
-                    len(dummy_input_ids), b_req_idx, b_seq_len, b_ready_cache_len
+                    len(dummy_input_ids), b_seq_len, b_ready_cache_len
                 ).cuda()
                 model_input = ModelInput(
                     batch_size=1,

lightllm/common/basemodel/cuda_graph.py

@@ -201,8 +201,11 @@ def warmup(self, model):
             )
             b_seq_len = torch.empty(batch_size, dtype=torch.int32, device="cuda")
             b_seq_len.fill_(seq_len)
+            b_last_mem_index = torch.zeros_like(b_seq_len)
             b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
-            mem_indexes = model.req_manager.alloc_mem_indices(len(input_ids), b_req_idx, b_seq_len).cuda()
+            mem_indexes = model.req_manager.alloc_mem_indices(
+                len(input_ids), b_seq_len, b_last_mem_index=b_last_mem_index
+            ).cuda()
 
             model_input = ModelInput(
                 batch_size=batch_size,
@@ -257,8 +260,11 @@ def warmup_overlap(self, model):
                 )
                 b_seq_len = torch.empty(batch_size, dtype=torch.int32, device="cuda")
                 b_seq_len.fill_(seq_len)
+                b_last_mem_index = torch.zeros_like(b_seq_len)
                 b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cuda")
-                mem_indexes = model.req_manager.alloc_mem_indices(len(input_ids), b_req_idx, b_seq_len).cuda()
+                mem_indexes = model.req_manager.alloc_mem_indices(
+                    len(input_ids), b_seq_len, b_last_mem_index=b_last_mem_index
+                ).cuda()
 
                 micro_batch = ModelInput(
                     is_prefill=False,

lightllm/common/req_manager.py

@@ -71,10 +71,21 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana
     def calc_real_need_token_num(self, need_token_num, b_seq_len, b_ready_cache_len=None):
         return max(need_token_num, self._get_need_paged_token_num(b_seq_len, b_ready_cache_len))
 
-    def alloc_mem_indices(self, need_size, b_req_idx=None, b_seq_len=None, b_ready_cache_len=None) -> torch.Tensor:
+    def calc_last_mem_index_in_prefill(self, mem_indices, b_seq_len, b_ready_cache_len=None):
+        b_token_len = b_seq_len
+        if b_ready_cache_len is not None:
+            b_token_len = b_seq_len - b_ready_cache_len
+        b_token_len_cumsum = torch.cumsum(b_token_len, dim=0)
+        b_last_mem_index = mem_indices[b_token_len_cumsum - 1]
+        return b_last_mem_index
+
+    # b_ready_cache_len为None时才需要b_last_mem_index
+    def alloc_mem_indices(
+        self, need_size, b_seq_len=None, b_ready_cache_len=None, b_last_mem_index=None
+    ) -> torch.Tensor:
         page_size = get_page_size()
-        if page_size > 1 and b_req_idx is not None and b_seq_len is not None:
-            return self._alloc_paged_mem_indices(b_req_idx, page_size, b_seq_len, b_ready_cache_len)
+        if page_size > 1 and b_seq_len is not None:
+            return self._alloc_paged_mem_indices(page_size, b_seq_len, b_ready_cache_len, b_last_mem_index)
         else:
             return self.mem_manager.alloc(need_size)
 
@@ -114,12 +125,11 @@ def _expand_by_page_size(self, b_token_len, page_size):
         p_token_len[last_page_positions] = remainders
         return need_pages_num, p_token_len
 
-    def _alloc_paged_mem_indices(self, b_req_idx, page_size, b_seq_len, b_ready_cache_len):
+    def _alloc_paged_mem_indices(self, page_size, b_seq_len, b_ready_cache_len, b_last_mem_index):
+        b_seq_len = b_seq_len.cpu()
         if b_ready_cache_len is not None:
             # prefill
-            b_seq_len = b_seq_len.cpu()
             b_ready_cache_len = b_ready_cache_len.cpu()
-
             b_token_len = b_seq_len - b_ready_cache_len
             total_pages_needed, p_token_len = self._expand_by_page_size(b_token_len, page_size)
             paged_token_idxs = self.mem_manager.alloc(total_pages_needed * page_size)
@@ -128,19 +138,17 @@ def _alloc_paged_mem_indices(self, b_req_idx, page_size, b_seq_len, b_ready_cach
             return pages[mask]
         else:
             # decode
-            b_seq_len = b_seq_len.cuda()
-            b_req_idx = b_req_idx.cuda()
+            assert b_last_mem_index is not None
+            b_last_mem_index = b_last_mem_index.cpu()
             need_new_page_mask = (b_seq_len - 1) % page_size == 0
-            new_pages_num = need_new_page_mask.sum().cpu()
+            new_pages_num = need_new_page_mask.sum()
             token_idxs = torch.zeros_like(b_seq_len, device=b_seq_len.device)
             if new_pages_num > 0:
-                new_pages_tokens = self.mem_manager.alloc(new_pages_num * page_size).cuda()
+                new_pages_tokens = self.mem_manager.alloc(new_pages_num * page_size)
                 token_idxs[need_new_page_mask] = new_pages_tokens[::page_size]
-
             mask = ~need_new_page_mask
             if mask.any():
-                seq_lens = b_seq_len[mask]
-                token_idxs[mask] = self.req_to_token_indexs[b_req_idx[mask], seq_lens - 2] + 1
+                token_idxs[mask] = b_last_mem_index[mask] + 1
         return token_idxs
 
     def _get_need_paged_token_num(self, b_seq_len, b_ready_cache_len=None):

lightllm/server/router/model_infer/infer_batch.py

@@ -288,6 +288,7 @@ def __init__(
         self.shm_index = shm_index
         self.multimodal_params = multimodal_params
         self.vocab_size = vocab_size
+        self.last_kv_mem_index = -1
 
         # 请求需要被暂停
         self.wait_pause = False

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

@@ -82,8 +82,13 @@ def padded_prepare_prefill_inputs(
         )
         g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(token_num)
     mem_indexes = g_infer_context.req_manager.alloc_mem_indices(
-        input_ids.shape[0] - padded_req_num, b_req_idx, b_seq_len, b_ready_cache_len
+        input_ids.shape[0] - padded_req_num, b_seq_len, b_ready_cache_len
     )
+    b_last_mem_index = g_infer_context.req_manager.calc_last_mem_index_in_prefill(
+        mem_indexes, b_seq_len, b_ready_cache_len
+    )
+    for i, req in enumerate(req_objs):
+        req.last_kv_mem_index = b_last_mem_index[i].item()
 
     g_infer_state_lock.release()
 
@@ -123,6 +128,7 @@ def padded_prepare_decode_inputs(
     b_req_idx = []
     b_mtp_index = []
     b_seq_len = []
+    b_last_mem_index = []
     for req in req_objs:
         run_reqs.append(req)
         b_req_idx.append(req.req_idx)
@@ -132,6 +138,7 @@ def padded_prepare_decode_inputs(
         total_token_num += seq_len
         max_len_in_batch = max(max_len_in_batch, seq_len)
         b_mtp_index.append(0)
+        b_last_mem_index.append(req.last_kv_mem_index)
         # process the draft tokens.
         for step in range(req.mtp_step):
             run_reqs.append(req)
@@ -164,15 +171,18 @@ def padded_prepare_decode_inputs(
     b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu")
     b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu")
     b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu")
+    b_last_mem_index = torch.tensor(b_last_mem_index, dtype=torch.int32, device="cpu")
 
     # dynamic prompt cache 准备 token
     g_infer_state_lock.acquire()
     if g_infer_context.radix_cache is not None:
         token_num = g_infer_context.req_manager.calc_real_need_token_num(b_seq_len.shape[0] - padded_req_num, b_seq_len)
         g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(token_num)
     mem_indexes = g_infer_context.req_manager.alloc_mem_indices(
-        b_seq_len.shape[0] - padded_req_num, b_req_idx, b_seq_len
+        b_seq_len.shape[0] - padded_req_num, b_seq_len, b_last_mem_index=b_last_mem_index
     )
+    for i, req in enumerate(req_objs):
+        req.last_kv_mem_index = mem_indexes[i]
     g_infer_state_lock.release()
 
     if padded_req_num > 0:

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

@@ -59,9 +59,12 @@ def prepare_prefill_inputs(
             input_ids.shape[0], b_seq_len, b_ready_cache_len
         )
         g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(token_num)
-    mem_indexes = g_infer_context.req_manager.alloc_mem_indices(
-        input_ids.shape[0], b_req_idx, b_seq_len, b_ready_cache_len
+    mem_indexes = g_infer_context.req_manager.alloc_mem_indices(input_ids.shape[0], b_seq_len, b_ready_cache_len)
+    b_last_mem_index = g_infer_context.req_manager.calc_last_mem_index_in_prefill(
+        mem_indexes, b_seq_len, b_ready_cache_len
     )
+    for i, req in enumerate(req_objs):
+        req.last_kv_mem_index = b_last_mem_index[i].item()
     g_infer_state_lock.release()
 
     model_input = ModelInput(
@@ -90,6 +93,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
     b_req_idx = []
     b_mtp_index = []
     b_seq_len = []
+    b_last_mem_index = []
     for req in req_objs:
         run_reqs.append(req)
         b_req_idx.append(req.req_idx)
@@ -99,6 +103,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
         total_token_num += seq_len
         max_len_in_batch = max(max_len_in_batch, seq_len)
         b_mtp_index.append(0)
+        b_last_mem_index.append(req.last_kv_mem_index)
         # process the draft tokens.
         for step in range(req.mtp_step):
             run_reqs.append(req)
@@ -112,13 +117,18 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
     b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu")
     b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu")
     b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu")
+    b_last_mem_index = torch.tensor(b_last_mem_index, dtype=torch.int32, device="cpu")
 
     # dynamic prompt cache 准备 token
     g_infer_state_lock.acquire()
     if g_infer_context.radix_cache is not None:
         token_num = g_infer_context.req_manager.calc_real_need_token_num(b_seq_len.shape[0], b_seq_len)
         g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(token_num)
-    mem_indexes = g_infer_context.req_manager.alloc_mem_indices(b_seq_len.shape[0], b_req_idx, b_seq_len)
+    mem_indexes = g_infer_context.req_manager.alloc_mem_indices(
+        b_seq_len.shape[0], b_seq_len, b_last_mem_index=b_last_mem_index
+    )
+    for i, req in enumerate(req_objs):
+        req.last_kv_mem_index = mem_indexes[i]
     g_infer_state_lock.release()
 
     model_input = ModelInput(

test/benchmark/static_inference/model_infer.py

@@ -258,7 +258,8 @@ def run_forward_once(
         b_seq_len[i] = input_len
 
     total_token_num = batch_size * input_len
-    mem_indexes = model_part.req_manager.alloc_mem_indices(test_data.shape[0], b_req_idx, b_seq_len, b_ready_cache_len)
+    mem_indexes = model_part.req_manager.alloc_mem_indices(test_data.shape[0], b_seq_len, b_ready_cache_len)
+    b_last_mem_index = model_part.req_manager.calc_last_mem_index_in_prefill(mem_indexes, b_seq_len, b_ready_cache_len)
     b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cpu")
     rank_id = model_kvargs["rank_id"]
 
@@ -321,7 +322,10 @@ def run_forward_once(
         step_start = time.time()
         total_token_num += batch_size
         b_seq_len += 1
-        mem_indexes = model_part.req_manager.alloc_mem_indices(predict_ids.shape[0], b_req_idx, b_seq_len)
+        mem_indexes = model_part.req_manager.alloc_mem_indices(
+            predict_ids.shape[0], b_seq_len, b_last_mem_index=b_last_mem_index
+        )
+        b_last_mem_index = mem_indexes
         max_len_in_batch = input_len + i + 1
         logits = decode_fn(
             model_part,

test/benchmark/static_inference/model_infer_mtp.py

@@ -124,9 +124,8 @@ def run_forward_once(args, input_len, output_len, batch_size, main_model, draft_
         b_seq_len[i] = input_len
 
     total_token_num = input_len * batch_size
-    mem_indexes = main_model.req_manager.alloc_mem_indices(
-        test_data.shape[0], b_req_idx, b_seq_len, b_ready_cache_len
-    ).cuda()
+    mem_indexes = main_model.req_manager.alloc_mem_indices(test_data.shape[0], b_seq_len, b_ready_cache_len).cuda()
+    b_last_mem_index = main_model.req_manager.calc_last_mem_index_in_prefill(mem_indexes, b_seq_len, b_ready_cache_len)
     # Main model Prefill
     model_input = ModelInput(
         batch_size=batch_size,
@@ -194,7 +193,7 @@ def run_forward_once(args, input_len, output_len, batch_size, main_model, draft_
     nopad_b_seq_idx = torch.tensor(nopad_b_seq_idx, dtype=torch.int32, device="cuda")
     nopad_b_seq_len = torch.tensor(nopad_b_seq_len, dtype=torch.int32, device="cuda")
     mem_indexes = main_model.req_manager.alloc_mem_indices(
-        batch_size * (len(draft_models) + 1), nopad_b_seq_idx, nopad_b_seq_len
+        batch_size * (len(draft_models) + 1), nopad_b_seq_len, b_last_mem_index=b_last_mem_index
     ).cuda()
 
     model_input = ModelInput(