paged transfer.

Author: Weichao Luo

lightllm/common/basemodel/basemodel.py

@@ -166,6 +166,10 @@ def _init_kv_move_buffer(self):
         # p d 分离的推理模式下才需要做这一步初始化
         if self.run_mode in ["prefill", "decode"]:
             self.mem_manager.alloc_kv_move_buffer(self.mem_manager.size)
+        elif self.run_mode in ["nixl_prefill", "nixl_decode"]:
+            page_num = int(os.getenv("PD_NIXL_MOVE_PAGE_NUM", 32))
+            page_size = int(os.getenv("PD_NIXL_MOVE_PAGE_SIZE", 1024))
+            self.mem_manager.alloc_paged_kv_move_buffer(page_num, page_size)
 
     def _check_mem_size(self):
         self.max_total_token_num = self.mem_manager.size

lightllm/common/deepseek2_mem_manager.py

@@ -36,6 +36,12 @@ def alloc_kv_move_buffer(self, max_req_total_len):
         self.token_dim_size = self.kv_move_buffer.shape[-1] * self.kv_move_buffer.shape[-2]
         return
 
+    def alloc_paged_kv_move_buffer(self, page_num, page_size):
+        self.kv_move_buffer = torch.empty(
+            (page_num, page_size, self.layer_num, self.head_num, self.head_dim), dtype=self.dtype, device="cuda"
+        )
+        return
+
     def send_to_decode_node(
         self,
         move_tasks: List[KVMoveTask],

lightllm/common/mem_manager.py

@@ -92,6 +92,14 @@ def alloc_kv_move_buffer(self, max_req_total_len):
         self.token_dim_size = self.kv_move_buffer.shape[-2] * self.kv_move_buffer.shape[-1]
         return
 
+    def alloc_paged_kv_move_buffer(self, page_num, page_size):
+        if isinstance(self, MemoryManager) and type(self) != MemoryManager:
+            raise NotImplementedError("subclass need reimpl this method")
+        self.kv_move_buffer = torch.empty(
+            (page_num, page_size, self.layer_num, 2 * self.head_num, self.head_dim), dtype=self.dtype, device="cuda"
+        )
+        return
+
     def send_to_decode_node(
         self,
         move_tasks: List[KVMoveTask],

lightllm/server/router/model_infer/infer_batch.py

@@ -264,8 +264,6 @@ def init_all(self):
             self.shm_req = g_infer_context.shm_req_manager.get_req_obj_by_index(self.shm_index)
             self.shm_req.link_prompt_ids_shm_array()
             self.shm_req.link_logprobs_shm_array()
-            if isinstance(self.shm_req, PDNIXLChunkedPrefillReq):
-                self.in_prefill_or_transfer = False
 
             self.sampling_param: InferSamplingParams = InferSamplingParams(self.shm_req, self.vocab_size)
             if self.sampling_param.shm_param.input_penalty:

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

@@ -291,7 +291,7 @@ def _post_handle(
             req_obj.update_finish_status(self.eos_id)
 
             if extra_post_req_handle_chunk_func is not None:
-                extra_post_req_handle_chunk_func(req_obj)
+                extra_post_req_handle_chunk_func(req_obj, next_token_id, next_token_logprob)
 
             if extra_post_req_handle_func is not None:
                 extra_post_req_handle_func(req_obj, next_token_id, next_token_logprob)

lightllm/server/router/model_infer/mode_backend/pd_nixl/impl_for_pd_base.py

@@ -2,6 +2,7 @@
 import torch
 import torch.multiprocessing as mp
 import threading
+from concurrent.futures import ThreadPoolExecutor
 from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend
 from typing import List, Tuple, Dict
 from lightllm.server.router.model_infer.infer_batch import g_infer_context, InferReq
@@ -11,7 +12,7 @@
 from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample
 from lightllm.server.multimodal_params import MultimodalParams
 
-from .pd_remote_prefill_obj import RemotePrefillTask, RemotePrefillServerInfo, RemotePrefillRequest
+from .pd_remote_prefill_obj import RemotePrefillTask, RemotePrefillServerInfo, RemotePrefillRequest, RemoteTransferStatusType
 
 from .impl_for_pd_base import PDNIXLBackendBase
 
@@ -24,7 +25,10 @@ def __init__(self, prefill_task_queue: mp.Queue, prefill_done_queue: mp.Queue, n
 
     def init_custom(self):
         super().init_custom()
-        self.wait_prefill_thread = threading.Thread(target=self._prefill_wait_loop, daemon=True)
+        self.wait_prefill_thread = threading.Thread(target=self._start_async_loop,
+                                                    args=(self._prefill_wait_loop_async,),
+                                                    daemon=True)
+        self.wait_move_page_pool = ThreadPoolExecutor(max_workers=4)
         self.wait_prefill_thread.start()
         return
 
@@ -44,9 +48,15 @@ def _build_remote_prefill_task(self, index: int, kwargs: Dict, req: InferReq):
             multimodal_params=MultimodalParams.from_dict(req.multimodal_params),
             local_cached_len=req.cur_kv_len,
             token_ids=mem_indexes[b_start_loc[index] : b_start_loc[index + 1]],
+            page_ids=self.page_scheduer.borrow() # get page ids for this request, blocking when not enough pages
         )
         return RemotePrefillTask(server_info=prefill_node_info, prefill_request=prefill_request)
 
+    def _trigger_remote_prefill(self, req_id: int, index: int, kwargs: Dict, req: InferReq):
+        remote_prefill_task = self._build_remote_prefill_task(index, kwargs, req)
+        self.request_to_page_ids[req_id] = remote_prefill_task.prefill_request.page_ids
+        self.to_remote_queue.put(remote_prefill_task)
+
     def prefill(self, reqs: List[Tuple]):
         self._init_reqs(reqs, init_req_obj=False)
         return
@@ -74,9 +84,11 @@ def decode(self):
                 # since the token index are the same across TPs, we only need to trigger prefill on master
                 if self.is_master_in_dp:
                     run_req.remote_prefill_start = time.time()
-                    self.to_remote_queue.put(self._build_remote_prefill_task(idx, kwargs, run_req))
+                    # since this function may blocking the calling thread, so we do it in a thread pool
+                    self.wait_move_page_pool.submit(self._trigger_remote_prefill,
+                                                    shm_req.group_req_id, idx, kwargs, run_req)
 
-                shm_req.set_pd_req_rank_state(self.rank_in_dp, 0)  # set in progress state
+                shm_req.set_pd_req_rank_state(self.rank_in_dp, RemoteTransferStatusType.IN_PROGRESS.value)  # set in progress state
                 run_req.in_prefill_or_transfer = True
                 self.remote_prefilled_reqs[shm_req.group_req_id] = run_req
 

lightllm/server/router/model_infer/mode_backend/pd_nixl/impl_for_pd_decode.py

@@ -10,7 +10,7 @@
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.server.router.model_infer.mode_backend.dp_backend.pre_process import padded_prepare_decode_inputs
 
-from .impl_for_pd_decode import PDNIXLBackendForDecodeNode
+from .impl_for_pd_decode import PDNIXLBackendForDecodeNode, RemoteTransferStatusType
 
 logger = init_logger(__name__)
 
@@ -55,9 +55,11 @@ def decode(self):
                 # since the token index are the same across TPs, we only need to trigger prefill on master
                 if self.is_master_in_dp:
                     run_req.remote_prefill_start = time.time()
-                    self.to_remote_queue.put(self._build_remote_prefill_task(idx, kwargs, run_req))
+                    # since this function may blocking the calling thread, so we do it in a thread pool
+                    self.wait_move_page_pool.submit(self._trigger_remote_prefill,
+                                                    shm_req.group_req_id, idx, kwargs, run_req)
 
-                shm_req.set_pd_req_rank_state(self.rank_in_dp, 0)  # set in progress state
+                shm_req.set_pd_req_rank_state(self.rank_in_dp, RemoteTransferStatusType.IN_PROGRESS.value)  # set in progress state
                 run_req.in_prefill_or_transfer = True
                 self.remote_prefilled_reqs[shm_req.group_req_id] = run_req
 

lightllm/server/router/model_infer/mode_backend/pd_nixl/impl_for_pd_decode_dp.py

@@ -18,9 +18,15 @@ def __init__(self, transfer_task_queue: mp.Queue, transfer_done_queue: mp.Queue,
 
     def init_custom(self):
         super().init_custom()
-        self.handle_prefill_loop_thread = threading.Thread(target=self._handle_prefill_loop, daemon=True)
-        self.wait_transfer_loop_thread = threading.Thread(target=self._wait_transfer_loop, daemon=True)
-        self.handle_transfer_loop_thread = threading.Thread(target=self._handle_transfer_loop, daemon=True)
+        self.handle_prefill_loop_thread = threading.Thread(target=self._start_async_loop,
+                                                           args=(self._handle_prefill_loop,),
+                                                           daemon=True)
+        self.wait_transfer_loop_thread = threading.Thread(target=self._start_async_loop,
+                                                          args=(self._wait_page_transfer_loop,),
+                                                          daemon=True)
+        self.handle_transfer_loop_thread = threading.Thread(target=self._start_async_loop,
+                                                            args=(self._handle_transfer_loop,),
+                                                            daemon=True)
 
         self.handle_prefill_loop_thread.start()
         self.handle_transfer_loop_thread.start()

lightllm/server/router/model_infer/mode_backend/pd_nixl/impl_for_pd_prefill.py

@@ -41,12 +41,6 @@ def decode(self):
         self._prefill_abort_remote(aborted_reqs)
         self._filter_reqs(aborted_reqs + ok_finished_reqs)
 
-        # if ok_finished_reqs:
-        #     for req in ok_finished_reqs:
-        #         self._transfer_kv_to_remote(req)
-        #     self._filter_reqs(ok_finished_reqs)
-        #     ok_finished_reqs.clear()
-
         current_dp_prefill_num = len(prefill_reqs)
         self.reduce_tensor.fill_(current_dp_prefill_num)
         dist.all_reduce(self.reduce_tensor, op=dist.ReduceOp.MAX, group=None, async_op=False)