chunk transfer.

Author: Weichao Luo

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

@@ -19,6 +19,7 @@
     KVMoveRequest,
     RemotePrefillStatus,
     ThreadSafeDict,
+    TransferState,
 )
 
 logger = init_logger(__name__)
@@ -55,15 +56,17 @@ def handle_remote_prefill(req_status: RemotePrefillStatus):
                 status = req_status.status
                 if status != 1:
                     logger.warning(f"remote prefill reqeust: {group_req_id} done with state: {status}")
+
                 if run_req := self.remote_prefilled_reqs.get(group_req_id, None):
-                    shm_req: PDChunkedPrefillReq = run_req.shm_req
-                    shm_req.set_pd_req_rank_state(self.rank_in_dp, status)
-                    self.remote_prefilled_reqs.pop(group_req_id)
-                    if self.is_master_in_dp:
-                        logger.info(
-                            f"remote prefill reqeust: {group_req_id} done with status: {status} "
-                            f"took: {time.time() - run_req.remote_prefill_start} seconds"
-                        )
+                    if req_status.is_last or status != 1:
+                        shm_req: PDChunkedPrefillReq = run_req.shm_req
+                        shm_req.set_pd_req_rank_state(self.rank_in_dp, status)
+                        self.remote_prefilled_reqs.pop(group_req_id)
+                        if self.is_master_in_dp:
+                            logger.info(
+                                f"remote prefill reqeust: {group_req_id} done with status: {status} "
+                                f"took: {time.time() - run_req.remote_prefill_start} seconds"
+                            )
                 else:
                     if self.is_master_in_dp:
                         logger.warning(f"remote prefill reqeust: {group_req_id} not found")
@@ -100,12 +103,15 @@ def _wait_transfer_loop(self):
                 req: InferReq = self.inflght_transfer_requests[req_id]
                 shm_req: PDChunkedPrefillReq = req.shm_req
                 shm_req.set_pd_req_rank_state(self.rank_in_dp, state)
-                del self.inflght_transfer_requests[req_id]
+                transfer_state = self.remote_prefill_requests[req_id].transfer_state
                 if self.is_master_in_dp:
                     logger.info(
                         f"req: {req_id} kv transfer with state: {state} "
-                        f"took: {time.time() - req.kv_transfer_start} seconds"
+                        f"took: {time.time() - transfer_state.start_time} seconds"
                     )
+                del self.remote_prefill_requests[req_id]
+                del self.inflght_transfer_requests[req_id]
+
             time.sleep(PDNIXLBackendBase._THEAD_WAIT_INTERVAL)
 
     def _handle_prefill_loop(self):
@@ -140,20 +146,37 @@ def _transfer_kv_to_remote(self, req: InferReq):
             logger.info(f"remote prefill request {group_req_id} not found")
             return
 
-        # kick off kv transfer
-        if req.finish_status.is_finished():
-            req.kv_transfer_start = time.time()
-            kv_transfer_req = KVMoveRequest(
-                group_req_id=group_req_id,
-                token_ids=self.model.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].tolist(),
+        remote_request: PrefillRequest = self.remote_prefill_requests[group_req_id]
+        if remote_request.transfer_state is None:
+            remote_request.transfer_state = TransferState(
+                start_time=time.time(),
+                current_kv_len=0,
+                current_chunk_id=0,
             )
-            remote_request = self.remote_prefill_requests[group_req_id]
-            self.nixl_agent.write_blocks(kv_transfer_req, remote_request)
+
+        transfer_state = remote_request.transfer_state
+        token_index = self.model.req_manager.req_to_token_indexs[req.req_idx]
+        is_finished = req.finish_status.is_finished()
+
+        kv_transfer_req = KVMoveRequest(
+            group_req_id=group_req_id,
+            token_ids=token_index[ : req.cur_kv_len].tolist(),
+            prev_kv_len=transfer_state.current_kv_len,
+            cur_kv_len=req.cur_kv_len,
+        )
+        # kick off kv transfer
+        self.nixl_agent.write_blocks(kv_transfer_req, remote_request, is_finished)
+
+        if transfer_state.current_chunk_id == 0:
             shm_req: PDChunkedPrefillReq = req.shm_req
             shm_req.set_pd_req_rank_state(self.rank_in_dp, 0)
-            req.kv_transfering = True
+            req.in_prefill_or_transfer = True
             self.inflght_transfer_requests[group_req_id] = req
 
+        transfer_state.current_kv_len = req.cur_kv_len
+        transfer_state.current_chunk_id += 1
+
+
     def _decode_filter_reqs(
         self, prefill_reqs: List[InferReq], aborted_reqs: List[InferReq], decode_reqs: List[InferReq]
     ):

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

@@ -6,7 +6,10 @@
 
 from lightllm.utils.log_utils import init_logger
 
-from .pd_remote_prefill_obj import RemoteAgent, KVMoveRequest, PrefillRequest, RemotePrefillStatus, ThreadSafeDict
+from .pd_remote_prefill_obj import (
+    RemoteAgent, KVMoveRequest, PrefillRequest,
+    RemotePrefillStatus, ThreadSafeDict, KVMoveRequestState
+    )
 
 
 logger = init_logger(__name__)
@@ -108,11 +111,20 @@ def _get_token_desc_ids(self, token_ids: List[int]):
                 descs_ids.append(layer_id * self.num_tokens + token_id)
         return descs_ids
 
-    def write_blocks(self, request: KVMoveRequest, prefill_request: PrefillRequest):
+    def write_blocks(self, request: KVMoveRequest, prefill_request: PrefillRequest, is_finished: bool):
         group_reqeust_id = request.group_req_id
         skip_kv_move_len = prefill_request.data.local_cached_len
-        src_token_ids = request.token_ids[skip_kv_move_len:]
-        dst_token_ids = prefill_request.data.token_ids
+
+        # current kv len is less than remote cached kv len, just skip
+        if request.cur_kv_len <= skip_kv_move_len:
+            return
+
+        kv_move_start = max(skip_kv_move_len, request.prev_kv_len)
+        kv_move_end   = request.cur_kv_len
+
+        src_token_ids = request.token_ids[kv_move_start:]
+        dst_token_ids = prefill_request.data.token_ids[kv_move_start - skip_kv_move_len: kv_move_end]
+
         remote_agent: RemoteAgent = self.remote_agents[prefill_request.decode_id][
             self.tp_idx
         ]  # TODO one-one mapping now
@@ -124,52 +136,85 @@ def write_blocks(self, request: KVMoveRequest, prefill_request: PrefillRequest):
 
             src_handle = self.local_xfer_handles
             dst_handle = remote_agent.kv_xfer_handles
-            notify_status = RemotePrefillStatus(group_req_id=group_reqeust_id, status=1)
+            notify_status = RemotePrefillStatus(
+                group_req_id=group_reqeust_id,
+                status=1,
+                chunk_id=prefill_request.transfer_state.current_chunk_id,
+                is_last=is_finished)
+
             handle = self.nixl_agent.make_prepped_xfer(
                 "WRITE", src_handle, src_token_descs, dst_handle, dst_token_descs, notify_status.serialize()
             )
 
             status = self.nixl_agent.transfer(handle)
             assert status != "ERR"
 
-            self.inflight_transfers[group_reqeust_id] = (handle, remote_agent, False)
+            if group_reqeust_id not in self.inflight_transfers:
+                self.inflight_transfers[group_reqeust_id] = KVMoveRequestState(
+                    handles=[],
+                    done_handles=[],
+                    remote_agent=remote_agent,
+                    abort=False
+                )
+            self.inflight_transfers[group_reqeust_id].handles.append(handle)
 
             return handle
 
         return None
 
     def send_abort_notify(self, remote_id: int, group_reqeust_id):
         remote_agent: RemoteAgent = self.remote_agents[remote_id][self.tp_idx]
-        notify_status = RemotePrefillStatus(group_req_id=group_reqeust_id, status=-1)
+        notify_status = RemotePrefillStatus(group_req_id=group_reqeust_id, status=-1, chunk_id=-1, is_last=True)
         self.nixl_agent.send_notif(remote_agent.name, notify_status.serialize())
 
         if group_reqeust_id in self.inflight_transfers:
-            self.inflight_transfers[group_reqeust_id][2] = True
+            self.inflight_transfers[group_reqeust_id].abort = True
 
     def get_done_tranfers(self):
         done_req_ids = []
 
-        for req_id, (handle, remote_agent, is_abort) in self.inflight_transfers.items():
-            if is_abort:
+        for req_id, kv_move_state in self.inflight_transfers.items():
+            kv_move_state: KVMoveRequestState
+            if kv_move_state.abort:
                 logger.warning(f"{req_id} Transfer aborted")
                 done_req_ids.append((req_id, -1))
                 continue
 
-            remote_agent: RemoteAgent
-            xfer_state = self.nixl_agent.check_xfer_state(handle)
-            if xfer_state == "DONE":
-                done_req_ids.append((req_id, 1))
-            elif xfer_state == "PROC":
-                continue
-            else:
-                logger.warning(f"{req_id} Transfer failed with state {xfer_state}")
+            remote_agent: RemoteAgent = kv_move_state.remote_agent
+
+            left_handles = []
+            failed = False
+            for handle in kv_move_state.handles:
+                if failed:
+                    left_handles.append(handle)
+                    continue
+
+                xfer_state = self.nixl_agent.check_xfer_state(handle)
+
+                if xfer_state == "DONE":
+                    kv_move_state.done_handles.append(handle)
+                elif xfer_state == "PROC":
+                    left_handles.append(handle)
+                else:
+                    logger.warning(f"{req_id} Transfer failed with state {xfer_state}")
+                    failed = True
+                    kv_move_state.done_handles.append(handle)
+                    notify_failed_status = RemotePrefillStatus(group_req_id=req_id, status=-1, chunk_id=-1, is_last=True)
+                    self.nixl_agent.send_notif(remote_agent.name, notify_failed_status.serialize())
+
+            kv_move_state.handles = left_handles
+
+            if failed:
                 done_req_ids.append((req_id, -1))
-                notify_failed_status = RemotePrefillStatus(group_req_id=req_id, status=-1)
-                self.nixl_agent.send_notif(remote_agent.name, notify_failed_status.serialize())
+            elif len(left_handles) == 0:
+                done_req_ids.append((req_id, 1))
 
         for req_id, _ in done_req_ids:
-            # release will abort inflight transfer
-            self.nixl_agent.release_xfer_handle(self.inflight_transfers[req_id][0])
+            kv_move_state: KVMoveRequestState = self.inflight_transfers[req_id]
+            for handle in kv_move_state.handles + kv_move_state.done_handles:
+                # release will abort inflight transfer
+                self.nixl_agent.release_xfer_handle(handle)
+
             del self.inflight_transfers[req_id]
 
         return done_req_ids

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

@@ -95,13 +95,15 @@ def main_loop(self):
                         queue.put(request)
 
                     success = True
-                    for idx in range(self.tp_size):
+                    for idx in range(self.dist_info.node_world_size):
                         ack = self.from_backend_queue.get()
+                        logger.info(f"received ack from backend {idx}: {ack}")
                         if ack != "OK":
                             success = False
                             break
 
                     self.recv_from_decode.send_pyobj_multipart(client_obj, success)
+                    logger.info(f"Sent ack to decode: {success}")
                     if not success:
                         logger.warning(f"Remote connect failed: {request}")
 
@@ -166,6 +168,7 @@ def _send_nixl_agent(self, socket: SockWithPoller):
         )
 
         success = socket.recv_pyobj(timeout=60)
+        logger.info(f"recv remote nixl connect response {success}")
         if success is None:
             logger.warning("timeout to recv remote nixl connect response")
             return False
@@ -203,6 +206,8 @@ def main_loop(self):
                             RemotePrefillStatus(
                                 group_req_id=prefill_tasks.prefill_request.sampling_params.group_request_id,
                                 status=-1,
+                                chunk_id=-1,
+                                is_last=True,
                             )
                         )
             except Exception as e:
@@ -212,7 +217,7 @@ def main_loop(self):
     def remote_prefill(self, server_id: int, prefill_request: RemotePrefillRequest):
         socket, _ = self.remote_prefill_servers[server_id]
         prefill_request.sampling_params.max_new_tokens = 1
-        socket.send_pyobj(PrefillRequest(type=RemoteRequstType.REMOTE_PREFILL, decode_id=self.id, data=prefill_request))
+        socket.send_pyobj(PrefillRequest(type=RemoteRequstType.REMOTE_PREFILL, decode_id=self.id, data=prefill_request, transfer_state=None))
 
 
 def remote_prefill_server_loop(

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

@@ -14,7 +14,7 @@
 logger = init_logger(__name__)
 
 try:
-    from nixl._api import nixlBind, nixl_prepped_dlist_handle
+    from nixl._api import nixlBind, nixl_prepped_dlist_handle, nixl_xfer_handle
 
 except ImportError:
     logger.error("nixl is not installed, which is required for pd disagreggation!!!")
@@ -53,17 +53,26 @@ class ConnectRequest(RemoteRequest):
     agent_metadatas: List[bytes]
     agent_mem_descs: List[bytes]
 
+@dataclass
+class TransferState:
+    start_time: float
+    current_kv_len: int
+    current_chunk_id: int
 
 @dataclass
 class PrefillRequest(RemoteRequest):
     decode_id: int
     data: RemotePrefillRequest
+    # transfer status
+    transfer_state: Optional[TransferState]
 
 
 @dataclass
 class KVMoveRequest:
     group_req_id: int
     token_ids: List[int]
+    prev_kv_len: int
+    cur_kv_len: int
 
 
 @dataclass
@@ -73,6 +82,13 @@ class RemoteAgent:
     kv_mem_desc: nixlBind.nixlRegDList
     kv_xfer_handles: nixl_prepped_dlist_handle
 
+@dataclass
+class KVMoveRequestState:
+    handles: List[nixl_xfer_handle]
+    done_handles: List[nixl_xfer_handle]
+    remote_agent: RemoteAgent
+    abort: bool
+
 
 @dataclass
 class RemotePrefillStatus:
@@ -160,7 +176,7 @@ def recv_pyobj_multipart(self):
         client_id, data = self.sock.recv_multipart()
         return client_id, pickle.loads(data)
 
-    def send_pyobj_multipart(self, client_id: str, data: Any):
+    def send_pyobj_multipart(self, client_id: bytes, data: Any):
         return self.sock.send_multipart([client_id, pickle.dumps(data)])
 
     def bind(self, addr: str):