feat: replace page idxs with token idxs in paged_mem_manager

Author: niushengxiao

lightllm/common/paged_mem_manager.py

@@ -18,12 +18,6 @@ def cdiv(a, b):
 class PagedMemoryManager(MemoryManager):
     def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9):
         super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction)
-        page_size = get_page_size()
-        self.mem_page_state = torch.arange(
-            0, cdiv(self.size, page_size), dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True
-        )
-        self.mark_page_start = 0
-        self.can_use_page_size = cdiv(self.size, page_size)
 
     def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
         self.kv_buffer = torch.empty(
@@ -53,42 +47,23 @@ def check_cache_page_valid(self, values: torch.Tensor):
         return True
 
     def alloc(self, need_size) -> torch.Tensor:
-        if self.can_use_page_size < need_size:
-            raise RuntimeError(
-                f"No available pages for alloc. remaining: {self.can_use_page_size}, needed: {need_size}"
-            )
-        new_pages = self.mem_page_state[self.mark_page_start : self.mark_page_start + need_size].cuda()
-        self.mark_page_start += need_size
-        self.can_use_page_size -= need_size
-        self.can_use_mem_size -= need_size * get_page_size()
-        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
-        return new_pages
+        assert need_size % get_page_size() == 0, "Need size must be a multiple of page size"
+        return super().alloc(need_size)
 
     def free(self, free_index: Union[torch.Tensor, List[int]]):
-        self.can_use_mem_size += len(free_index)
-        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
-
         page_size = get_page_size()
-        if isinstance(free_index, list):
-            free_index = torch.tensor(free_index, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True)
-
-        if len(free_index) == 0:
-            return
+        if page_size == 1:
+            return super().free(free_index)
 
+        if isinstance(free_index, list):
+            free_index = torch.tensor(free_index)
         base_free_index = free_index[free_index % page_size == 0]
-        page_indices = base_free_index // page_size
-        for page_idx in sorted(page_indices, reverse=True):  # 逆序放回，保持池的相对顺序
-            self.mark_page_start -= 1
-            self.mem_page_state[self.mark_page_start] = page_idx
-            self.can_use_page_size += 1
-
+        if len(base_free_index) == 0:
+            return
+        token_idxs = base_free_index[:, None] + torch.arange(page_size, device=free_index.device)
+        token_idxs = token_idxs.flatten()
+        super().free(token_idxs)
         return
 
     def free_all(self):
         super().free_all()
-        page_size = get_page_size()
-        self.mark_page_start = 0
-        self.can_use_page_size = cdiv(self.size, page_size)
-        self.mem_page_state = torch.arange(
-            0, cdiv(self.size, page_size), dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True
-        )

lightllm/common/req_manager.py

@@ -71,7 +71,7 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana
         self.max_request_num = max_request_num
         self.HOLD_REQUEST_ID = max_request_num
 
-    def expand_by_page_size(self, b_token_len, page_size):
+    def _expand_by_page_size(self, b_token_len, page_size):
         # 将seq_len按page整数倍展开，例如seq_len = [9,9,9] -> page_len = [4,4,1,4,4,1,4,4,1], page_size = 4
         b_page_len = cdiv(b_token_len, page_size)
         need_pages_num = b_page_len.sum()
@@ -82,36 +82,28 @@ def expand_by_page_size(self, b_token_len, page_size):
         p_token_len[last_page_positions] = remainders
         return need_pages_num, b_page_len, p_token_len
 
-    def alloc_paged_token_indices(self, b_req_idx, page_size, b_seq_len, b_ready_cache_len):
+    def _alloc_paged_token_indices(self, b_req_idx, page_size, b_seq_len, b_ready_cache_len):
         if b_ready_cache_len is not None:
             # prefill
-            b_req_idx = b_req_idx.cuda()
-            b_seq_len = b_seq_len.cuda()
-            b_ready_cache_len = b_ready_cache_len.cuda()
+            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, b_page_len, p_token_len = self.expand_by_page_size(b_token_len, page_size)
-            allocated_pages = self.mem_manager.alloc(total_pages_needed)
-
-            def get_offsets_by_length(b_len, max_len):
-                # 例：b_len = [3,4,5] -> [0,1,2,0,1,2,3,0,1,2,3,4]
-                offsets = torch.arange(max_len, dtype=b_len.dtype, device=b_len.device)
-                offset_mask = offsets.unsqueeze(0) < b_len.unsqueeze(1)
-                return torch.masked_select(offsets, offset_mask)
-
-            token_offsets = get_offsets_by_length(p_token_len, page_size)
-            page_bases = allocated_pages * page_size
-            return torch.repeat_interleave(page_bases, p_token_len) + token_offsets
+            total_pages_needed, b_page_len, 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)
+            pages = paged_token_idxs.view(-1, page_size)
+            mask = torch.arange(page_size, device=p_token_len.device) < p_token_len.unsqueeze(1)
+            return pages[mask]
         else:
             # decode
             b_seq_len = b_seq_len.cuda()
             b_req_idx = b_req_idx.cuda()
             need_new_page_mask = (b_seq_len - 1) % page_size == 0
-            new_pages_num = need_new_page_mask.sum()
+            new_pages_num = need_new_page_mask.sum().cpu()
             token_idxs = torch.zeros_like(b_seq_len, device=b_seq_len.device)
             if new_pages_num > 0:
-                new_pages = self.mem_manager.alloc(new_pages_num)
-                token_idxs[need_new_page_mask] = new_pages * page_size
+                new_pages_tokens = self.mem_manager.alloc(new_pages_num * page_size).cuda()
+                token_idxs[need_new_page_mask] = new_pages_tokens[::page_size]
 
             mask = ~need_new_page_mask
             if mask.any():
@@ -122,10 +114,28 @@ def get_offsets_by_length(b_len, max_len):
                 )
         return token_idxs
 
+    def _get_need_paged_token_num(self, b_seq_len, b_ready_cache_len=None):
+        page_size = get_page_size()
+        if page_size == 1:
+            return 0
+
+        need_new_pages = 0
+        if b_ready_cache_len is not None:
+            need_tokens_array = b_seq_len - b_ready_cache_len
+            need_pages_array = (need_tokens_array + page_size - 1) // page_size
+            need_new_pages = need_pages_array.sum()
+        else:
+            mask = (b_seq_len - 1) % page_size == 0
+            need_new_pages = mask.sum()
+        return need_new_pages * page_size
+
+    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_token_indices(self, need_size, b_req_idx, b_seq_len, b_ready_cache_len=None) -> torch.Tensor:
         page_size = get_page_size()
         if page_size > 1:
-            return self.alloc_paged_token_indices(b_req_idx, page_size, b_seq_len, b_ready_cache_len)
+            return self._alloc_paged_token_indices(b_req_idx, page_size, b_seq_len, b_ready_cache_len)
         else:
             return self.mem_manager.alloc(need_size)
 

lightllm/server/router/dynamic_prompt/paged_radix_cache.py

@@ -391,39 +391,17 @@ def _print_helper(self, node: TreeNode, indent):
             self._print_helper(child, indent=indent + 2)
         return
 
-    def free_radix_cache_to_get_enough_token(self, need_token_num=None, b_seq_len=None, b_ready_cache_len=None):
+    def free_radix_cache_to_get_enough_token(self, need_token_num):
         assert self.mem_manager is not None
-        need_pages = 0
-        can_use_pages = 0
-        if hasattr(self.mem_manager, "can_use_page_size") and self.page_size > 1 and b_seq_len is not None:
-
-            def get_need_page_size(page_size, b_seq_len, b_ready_cache_len=None):
-                need_new_pages = 0
-                if b_ready_cache_len is not None:
-                    need_tokens_array = b_seq_len - b_ready_cache_len
-                    need_pages_array = (need_tokens_array + page_size - 1) // page_size
-                    need_new_pages = need_pages_array.sum()
-                else:
-                    mask = (b_seq_len - 1) % page_size == 0
-                    need_new_pages = mask.sum()
-                return need_new_pages
-
-            need_pages = get_need_page_size(self.page_size, b_seq_len, b_ready_cache_len)
-            can_use_pages = self.mem_manager.can_use_page_size
-        if need_token_num > self.mem_manager.can_use_mem_size or need_pages > can_use_pages:
-            need_evict_single_token_num = need_token_num - self.mem_manager.can_use_mem_size
-            need_evict_page_token_num = (need_pages - can_use_pages) * self.page_size
-            need_evict_token_num = max(need_evict_single_token_num, need_evict_page_token_num)
-            remaining_tokens = self.get_tree_total_tokens_num() - self.get_refed_tokens_num()
-            need_evict_token_num = min(need_evict_token_num, remaining_tokens)
+        if need_token_num > self.mem_manager.can_use_mem_size:
+            need_evict_token_num = need_token_num - self.mem_manager.can_use_mem_size
             release_mems = []
 
             def release_mem(mem_index):
                 release_mems.append(mem_index)
                 return
 
             self.evict(need_evict_token_num, release_mem)
-            if release_mems:
-                mem_index = torch.concat(release_mems)
-                self.mem_manager.free(mem_index)
+            mem_index = torch.concat(release_mems)
+            self.mem_manager.free(mem_index)
         return

lightllm/server/router/dynamic_prompt/radix_cache.py

@@ -333,7 +333,7 @@ def _print_helper(self, node: TreeNode, indent):
             self._print_helper(child, indent=indent + 2)
         return
 
-    def free_radix_cache_to_get_enough_token(self, need_token_num=None, b_seq_len=None, b_ready_cache_len=None):
+    def free_radix_cache_to_get_enough_token(self, need_token_num):
         assert self.mem_manager is not None
         if need_token_num > self.mem_manager.can_use_mem_size:
             need_evict_token_num = need_token_num - self.mem_manager.can_use_mem_size

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

@@ -77,9 +77,10 @@ def padded_prepare_prefill_inputs(
     # dynamic prompt cache 准备 token
     g_infer_state_lock.acquire()
     if g_infer_context.radix_cache is not None:
-        g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(
+        token_num = g_infer_context.req_manager.calc_real_need_token_num(
             input_ids.shape[0] - padded_req_num, 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_token_indices(
         input_ids.shape[0] - padded_req_num, b_req_idx, b_seq_len, b_ready_cache_len
     )
@@ -167,7 +168,8 @@ def padded_prepare_decode_inputs(
     # dynamic prompt cache 准备 token
     g_infer_state_lock.acquire()
     if g_infer_context.radix_cache is not None:
-        g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(b_seq_len.shape[0] - padded_req_num, b_seq_len)
+        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_token_indices(
         b_seq_len.shape[0] - padded_req_num, b_req_idx, b_seq_len
     )

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

@@ -55,9 +55,10 @@ def prepare_prefill_inputs(
     # dynamic prompt cache 准备 token
     g_infer_state_lock.acquire()
     if g_infer_context.radix_cache is not None:
-        g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(
+        token_num = g_infer_context.req_manager.calc_real_need_token_num(
             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_token_indices(
         input_ids.shape[0], b_req_idx, b_seq_len, b_ready_cache_len
     )
@@ -115,7 +116,8 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
     # dynamic prompt cache 准备 token
     g_infer_state_lock.acquire()
     if g_infer_context.radix_cache is not None:
-        g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(b_seq_len.shape[0], b_seq_len)
+        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_token_indices(b_seq_len.shape[0], b_req_idx, b_seq_len)
     g_infer_state_lock.release()