fix infer_state init

Author: shihaobai

lightllm/common/basemodel/basemodel.py

@@ -480,7 +480,7 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod
             model_input0.max_len_in_batch,
             infer_state0.mem_index,
         )
-        infer_state0.init_some_extra_state(self, input_ids0)
+        infer_state0.init_some_extra_state(self, model_input0)
 
         infer_state1 = self._create_inferstate(model_input1, 1)
         init_req_to_token_indexes(
@@ -491,7 +491,7 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod
             model_input1.max_len_in_batch,
             infer_state1.mem_index,
         )
-        infer_state1.init_some_extra_state(self, input_ids1)
+        infer_state1.init_some_extra_state(self, model_input1)
 
         model_output0, model_output1 = self._overlap_tpsp_context_forward(
             input_ids0, infer_state0, input_ids1=input_ids1, infer_state1=infer_state1
@@ -576,15 +576,15 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
                 infer_state0.b_seq_len,
                 infer_state0.mem_index,
             )
-            infer_state0.init_some_extra_state(self, model_input0.input_ids)
+            infer_state0.init_some_extra_state(self, model_input0)
             infer_state1 = self._create_inferstate(model_input1, 1)
             copy_kv_index_to_req(
                 self.req_manager.req_to_token_indexs,
                 infer_state1.b_req_idx,
                 infer_state1.b_seq_len,
                 infer_state1.mem_index,
             )
-            infer_state1.init_some_extra_state(self, model_input1.input_ids)
+            infer_state1.init_some_extra_state(self, model_input1)
 
             model_output0, model_output1 = self._overlap_tpsp_token_forward(
                 model_input0.input_ids, infer_state0, input_ids1=model_input1.input_ids, infer_state1=infer_state1