Fix Comment

Signed-off-by: yizhang-nv <187001205+yizhang-nv@users.noreply.github.com>

Author: yizhang-nv

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -675,11 +675,9 @@ def _general_warmup(self,
             self.kv_cache_manager_key)
         token_num_upper_bound = min(self.max_num_tokens,
                                     self.batch_size * (self.max_seq_len - 1))
-        curr_max_num_tokens = min(
-            kv_cache_manager.get_num_available_tokens(
-                token_num_upper_bound=token_num_upper_bound,
-                max_num_draft_tokens=self.original_max_draft_len),
-            token_num_upper_bound)
+        curr_max_num_tokens = kv_cache_manager.get_num_available_tokens(
+            token_num_upper_bound=token_num_upper_bound,
+            max_num_draft_tokens=self.original_max_draft_len)
         max_batch_size = min(
             self.batch_size,
             curr_max_num_tokens // (1 + self.runtime_draft_len))
@@ -730,11 +728,9 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager):
             self.kv_cache_manager_key)
         token_num_upper_bound = min(self.max_num_tokens,
                                     self.batch_size * (self.max_seq_len - 1))
-        curr_max_num_tokens = min(
-            kv_cache_manager.get_num_available_tokens(
-                token_num_upper_bound=token_num_upper_bound,
-                max_num_draft_tokens=self.original_max_draft_len),
-            token_num_upper_bound)
+        curr_max_num_tokens = kv_cache_manager.get_num_available_tokens(
+            token_num_upper_bound=token_num_upper_bound,
+            max_num_draft_tokens=self.original_max_draft_len)
 
         cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None)
         with self.no_cuda_graph(), autotune(cache_path=cache_path):

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -1002,10 +1002,13 @@ def get_num_kv_blocks(self, num_tokens: int) -> int:
         return (num_tokens + self.tokens_per_block - 1) // self.tokens_per_block
 
     def get_num_available_tokens(self,
+                                 token_num_upper_bound: int,
                                  max_num_draft_tokens: int = 0,
                                  **kwargs) -> int:
-        return (self.get_num_free_blocks() * self.tokens_per_block -
-                self.num_extra_kv_tokens - max_num_draft_tokens)
+        return min(
+            token_num_upper_bound,
+            self.get_num_free_blocks() * self.tokens_per_block -
+            self.num_extra_kv_tokens - max_num_draft_tokens)
 
     def get_buffers(self,
                     layer_idx: int,