[None][feat] Run extra general warmup to warm up memory pool

Signed-off-by: Jin Li <[email protected]>

Author: liji-nv

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -598,6 +598,9 @@ def warmup(self, resource_manager: ResourceManager) -> None:
         self._run_torch_compile_warmup(resource_manager)
         self._run_autotuner_warmup(resource_manager)
         self._run_cuda_graph_warmup(resource_manager)
+        if not kv_cache_manager.is_estimating_kv_cache:
+            # Run extra general warmup to warmup memory pool before run real requests.
+            self._general_warmup(resource_manager, reverse=True)
 
         # Set the value back to the original value after all warmups are complete
         self.enable_spec_decode = self.is_spec_decode
@@ -612,8 +615,8 @@ def _general_warmup(self,
                 self.original_max_draft_len), self.max_num_tokens,
             self.batch_size * (self.max_seq_len - 1))
         max_batch_size = min(
-            self.batch_size,
-            curr_max_num_tokens // (1 + self.runtime_draft_len))
+            self.batch_size, curr_max_num_tokens //
+            (1 + self.runtime_draft_len) // self.max_beam_width)
 
         warmup_requests_configs = {
             (1, 1),  # Specialize for 1 token.
@@ -936,8 +939,8 @@ def _create_warmup_request(
 
         blocks_to_use = num_full_seqs * math.ceil(
             max_seq_len / kv_cache_manager.tokens_per_block) + math.ceil(
-                num_left_over_tokens /
-                kv_cache_manager.tokens_per_block) + num_gen_requests
+                num_left_over_tokens / kv_cache_manager.tokens_per_block
+            ) + num_gen_requests * self.max_beam_width
 
         if blocks_to_use > available_blocks:
             return None

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -193,6 +193,7 @@ def __init__(
             idx: offset
             for offset, idx in enumerate(self.pp_layers)
         }
+        self.is_estimating_kv_cache = is_estimating_kv_cache
 
         self.kv_connector_manager = kv_connector_manager