fix

Author: wangzaijun

lightllm/server/router/model_infer/mode_backend/pd_nixl/decode_node_impl/decode_impl_for_dp.py

@@ -3,7 +3,7 @@
 from lightllm.utils.log_utils import init_logger
 from typing import List, Tuple
 from lightllm.server.router.model_infer.mode_backend.dp_backend.impl import DPChunkedPrefillBackend
-from .decode_impl import NIXLDecodeNode
+from .decode_impl import NIXLDecodeNode, NIXLChunckedTransTaskGroup
 
 logger = init_logger(__name__)
 
@@ -32,13 +32,18 @@ def _decode_node_gen_trans_tasks(self, req_obj: InferReq):
         return NIXLDecodeNode._decode_node_gen_trans_tasks(self, req_obj=req_obj)
 
     def _create_nixl_trans_task(
-        self, req_obj: InferReq, mem_indexes: List[int], kv_start_index: int, kv_end_index: int, is_first_task: bool
+        self,
+        req_obj: InferReq,
+        mem_indexes: List[int],
+        kv_start_index: int,
+        kv_end_index: int,
+        group: NIXLChunckedTransTaskGroup,
     ):
-        return NIXLDecodeNode._decode_node_gen_trans_tasks(
+        return NIXLDecodeNode._create_nixl_trans_task(
             self,
             req_obj=req_obj,
             mem_indexes=mem_indexes,
             kv_start_index=kv_start_index,
             kv_end_index=kv_end_index,
-            is_first_task=is_first_task,
+            group=group,
         )

lightllm/server/router/model_infer/mode_backend/pd_nixl/prefill_node_impl/prefill_impl_for_dp.py

@@ -27,12 +27,11 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]:
     def _prefill_chuncked_handle_func(
         self, req_obj: InferReq, next_token_id: int, next_token_prob: float, output_len: int
     ):
-        NIXLChunckedPrefillForPrefillNode._prefill_chuncked_handle_func(
+        return NIXLChunckedPrefillForPrefillNode._prefill_chuncked_handle_func(
             self, req_obj=req_obj, next_token_id=next_token_id, next_token_prob=next_token_prob, output_len=output_len
         )
-        return
 
     def _create_nixl_trans_task(self, req_obj: InferReq, kv_start_index: int, kv_end_index: int):
-        NIXLChunckedPrefillForPrefillNode._create_nixl_trans_task(
+        return NIXLChunckedPrefillForPrefillNode._create_nixl_trans_task(
             self, req_obj=req_obj, kv_start_index=kv_start_index, kv_end_index=kv_end_index
         )