Update Dummy Run To Suppport Mutil-Batch Execution (PaddlePaddle#6123)

Author: iosmers

fastdeploy/worker/xpu_model_runner.py

@@ -1298,16 +1298,34 @@ def initialize_attn_backend(self) -> None:
             )
         self.attn_backends.append(attn_backend)
 
-    def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int):
-        """Set dummy prefill inputs to share_inputs"""
-        full_length = min(num_tokens // batch_size, self.model_config.max_model_len - 10)
-        input_length = int(full_length - 512)
+    def get_input_length_list(self, num_tokens: int, batch_size: int, expected_decode_len: int):
+        """
+        Args:
+            num_tokens (int): The total number of tokens across all sequences.
+            batch_size (int): The number of sequences (requests) in the batch.
+            expected_decode_len (int): The expected number of tokens every sequence should be generated by the model.
+        Returns:
+            List[int]: A list of integers representing the sequence length for each request.
+                    This list is crafted to maximize the total number of blocks.
+        """
+        max_dec_len = expected_decode_len + 1
+        input_length = min(num_tokens // batch_size, self.model_config.max_model_len - max_dec_len)
         block_num = (
             input_length + self.cache_config.block_size - 1
         ) // self.cache_config.block_size + self.cache_config.enc_dec_block_num
+        input_length_list = [input_length] * batch_size
+        len_of_input_length_list = len(input_length_list)
+        max_dec_len_list = [max_dec_len] * len_of_input_length_list
+        return input_length_list, max_dec_len_list, block_num
+
+    def _dummy_prefill_inputs(self, input_length_list: List[int], max_dec_len_list: List[int], block_num: int):
+        """Set dummy prefill inputs to share_inputs"""
+        batch_size = len(input_length_list)
 
         for i in range(batch_size):
             idx = i
+            input_length = input_length_list[idx]
+            max_dec_len = max_dec_len_list[idx]
             self.share_inputs["input_ids"][idx : idx + 1, :input_length] = np.array([5] * input_length)
             self.share_inputs["prompt_ids"][idx : idx + 1, :input_length] = np.array([5] * input_length)
             self.share_inputs["eos_token_id"][:] = np.array([2], dtype="int64").reshape(-1, 1)
@@ -1317,7 +1335,7 @@ def _dummy_prefill_inputs(self, num_tokens: int, batch_size: int):
             self.share_inputs["seq_lens_encoder"][idx : idx + 1] = input_length
             self.share_inputs["seq_lens_decoder"][idx : idx + 1] = 0
             self.share_inputs["step_idx"][idx : idx + 1] = 0
-            self.share_inputs["max_dec_len"][idx : idx + 1] = 10
+            self.share_inputs["max_dec_len"][idx : idx + 1] = max_dec_len
             self.share_inputs["stop_flags"][idx : idx + 1] = False
 
             self.share_inputs["first_token_ids"][idx : idx + 1] = self.share_inputs["input_ids"][idx : idx + 1, :1]
@@ -1333,14 +1351,27 @@ def _dummy_run(
         self,
         num_tokens: paddle.Tensor,
         batch_size: paddle.Tensor,
+        expected_decode_len: int = 1,
         in_capturing: bool = False,
     ) -> paddle.Tensor:
         """
         Use dummy inputs to run before formal execution.
         Args:
-            num_tokens: Expected number of tokens generated
+            num_tokens: Number of the input tokens
+            batch_size: Batch size
+            expected_decode_len: Expected decode length
+            in_capturing: Is cuda graph in capturing state
         """
-        self._dummy_prefill_inputs(num_tokens, batch_size)
+        input_length_list, max_dec_len_list, block_num = self.get_input_length_list(
+            num_tokens=num_tokens,
+            batch_size=batch_size,
+            expected_decode_len=expected_decode_len,
+        )
+        self._dummy_prefill_inputs(
+            input_length_list=input_length_list,
+            max_dec_len_list=max_dec_len_list,
+            block_num=block_num,
+        )
 
         if self.speculative_method in ["mtp"]:
             self.proposer.dummy_prefill_inputs(
@@ -1350,7 +1381,7 @@ def _dummy_run(
             )
 
         while True:
-            self.execute_model(is_dummy_run=True)
+            self.execute_model(is_dummy_run=True, in_capturing=in_capturing)
 
             if int((self.share_inputs["seq_lens_this_time"] > 0).sum()) == 0:
                 break
@@ -1417,6 +1448,7 @@ def execute_model(
         model_forward_batch: Optional[List[Request]] = None,
         num_running_requests: int = None,
         is_dummy_run: bool = False,
+        in_capturing: bool = False,
     ) -> Optional[ModelRunnerOutput]:
         """
         The Entrance of model execute.