[https://nvbugs/5570599][fix] Set KVCache free_gpu_memory_fraction fo… (#8780)

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

Author: liji-nv

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -58,7 +58,14 @@ def __init__(
         )
 
     def calculate_max_num_blocks(
-        self, kv_cache_config, head_dim, tokens_per_block, mapping, dtype, kv_factor
+        self,
+        kv_cache_config,
+        head_dim,
+        tokens_per_block,
+        mapping,
+        dtype,
+        kv_factor,
+        enforce_memory_limit,
     ) -> Tuple[int, int]:
         """Calculate the maximum number of blocks needed for the cache."""
         # TODO: this is VERY hacky... Ideally, we want to compute the number of blocks

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -417,6 +417,7 @@ def _create_kv_cache_manager(
                 is_draft=model_engine.is_draft_model,
                 kv_connector_manager=self._kv_connector_manager
                 if not estimating_kv_cache else None,
+                enforce_memory_limit=estimating_kv_cache,
             )
         elif is_nemotron_hybrid(config):
             if self._max_beam_width > 1:
@@ -490,6 +491,7 @@ def _create_kv_cache_manager(
                 is_draft=model_engine.is_draft_model,
                 kv_connector_manager=self._kv_connector_manager
                 if not estimating_kv_cache else None,
+                enforce_memory_limit=estimating_kv_cache,
             )
         # KVCacheManager (Non-draft) modifies the max_seq_len field, update it to self
         if model_engine.kv_cache_manager_key == ResourceManagerType.KV_CACHE_MANAGER:

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -164,6 +164,7 @@ def __init__(
         max_beam_width: int = 1,
         is_draft: bool = False,
         kv_connector_manager: Optional[KvCacheConnectorManager] = None,
+        enforce_memory_limit: bool = False,
     ) -> None:
         self.mapping = mapping
         self.dtype = dtype
@@ -283,6 +284,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
                 mapping=mapping,
                 dtype=dtype,
                 kv_factor=self.kv_factor,
+                enforce_memory_limit=enforce_memory_limit,
             )
             blocks_per_window = {
                 self.max_attention_window_vec[0]:
@@ -563,7 +565,8 @@ def calculate_max_num_blocks(self,
                                  tokens_per_block: int,
                                  mapping: Mapping,
                                  dtype: DataType,
-                                 kv_factor: int = 2):
+                                 kv_factor: int = 2,
+                                 enforce_memory_limit: bool = False):
         free_mem_fraction = (kv_cache_config.free_gpu_memory_fraction
                              if kv_cache_config.free_gpu_memory_fraction
                              is not None else 0.9)
@@ -591,7 +594,7 @@ def calculate_max_num_blocks(self,
         # If user specified a number of tokens
         if kv_cache_config.max_tokens is not None:
             # If user also specified a free gpu memory fraction, take the min
-            if kv_cache_config.free_gpu_memory_fraction is not None:
+            if kv_cache_config.free_gpu_memory_fraction is not None or enforce_memory_limit:
                 max_tokens = min(kv_cache_config.max_tokens, max_tokens)
                 logger.warning(
                     f'Both free_gpu_memory_fraction and max_tokens are set (to {free_mem_fraction} and {max_tokens} with free memory {free_mem / (1 << 32)} of total memory {total_mem / (1<<32)}, respectively). The smaller value will be used.'