fix

Author: hiworldwzj

lightllm/common/basemodel/basemodel.py

@@ -463,6 +463,9 @@ def _token_forward(self, input_ids, infer_state: InferStateInfo):
 
     @torch.no_grad()
     def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: ModelInput):
+        model_input0.to_cuda()
+        model_input1.to_cuda()
+
         assert model_input0.mem_indexes.is_cuda
         assert model_input1.mem_indexes.is_cuda
         input_ids0, input_ids1 = model_input0.input_ids, model_input1.input_ids
@@ -500,6 +503,22 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod
 
     @torch.no_grad()
     def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: ModelInput):
+        model_input0.to_cuda()
+        model_input1.to_cuda()
+
+        if model_input0.input_ids is None:
+            model_input0.input_ids = gather_token(
+                self.req_manager.req_sampling_params_manager.req_to_next_token_ids,
+                model_input0.b_req_idx,
+                model_input0.b_mtp_index,
+            )
+        if model_input1.input_ids is None:
+            model_input1.input_ids = gather_token(
+                self.req_manager.req_sampling_params_manager.req_to_next_token_ids,
+                model_input1.b_req_idx,
+                model_input1.b_mtp_index,
+            )
+
         assert model_input0.batch_size == model_input1.batch_size
         assert model_input0.mem_indexes.is_cuda
         assert model_input1.mem_indexes.is_cuda

lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py

@@ -235,8 +235,8 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer
             b_has_out_cpu = (
                 micro_input0.b_prefill_has_output_cpu[0:req_num0] + micro_input1.b_prefill_has_output_cpu[0:req_num1]
             )
-            b_mtp_index = torch.cat(micro_input0.b_mtp_index[0:req_num0], micro_input1.b_mtp_index[0:req_num1])
-            b_req_idx = torch.cat(micro_input0.b_req_idx[0:req_num0], micro_input1.b_req_idx[0:req_num1])
+            b_mtp_index = torch.cat((micro_input0.b_mtp_index[0:req_num0], micro_input1.b_mtp_index[0:req_num1]), dim=0)
+            b_req_idx = torch.cat((micro_input0.b_req_idx[0:req_num0], micro_input1.b_req_idx[0:req_num1]), dim=0)
 
             if (req_num0 + req_num1) > 0:
 
@@ -291,7 +291,7 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe
             micro_input1,
             run_reqs1,
             padded_req_num1,
-        ) = padded_overlap_prepare_decode_inputs(decode_reqs, is_multimodal=self.is_multimodal)
+        ) = padded_overlap_prepare_decode_inputs(req_objs=decode_reqs)
         micro_input0: ModelInput = micro_input0
         micro_input1: ModelInput = micro_input1
 
@@ -305,8 +305,8 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe
 
             logits[0:req_num0, :].copy_(logits0[0:req_num0, :], non_blocking=True)
             logits[req_num0 : (req_num0 + req_num1), :].copy_(logits1[0:req_num1, :], non_blocking=True)
-            b_mtp_index = torch.cat(micro_input0.b_mtp_index[0:req_num0], micro_input1.b_mtp_index[0:req_num1])
-            b_req_idx = torch.cat(micro_input0.b_req_idx[0:req_num0], micro_input1.b_req_idx[0:req_num1])
+            b_mtp_index = torch.cat((micro_input0.b_mtp_index[0:req_num0], micro_input1.b_mtp_index[0:req_num1]), dim=0)
+            b_req_idx = torch.cat((micro_input0.b_req_idx[0:req_num0], micro_input1.b_req_idx[0:req_num1]), dim=0)
 
             run_reqs = run_reqs0 + run_reqs1
             if (req_num0 + req_num1) > 0: