[None][fix] Fix DSACacheManager and RocketCacheManager KV cache estimation ignoring num_layers for draft models

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -1931,7 +1931,9 @@ def shutdown(self):
         super().shutdown()
 
     @staticmethod
-    def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping,
+    def get_cache_size_per_token(model_config: ModelConfig,
+                                 mapping: Mapping,
+                                 num_layers: Optional[int] = None,
                                  **kwargs):
         config = model_config.pretrained_config
         sparse_attn_config = model_config.sparse_attention_config
@@ -1948,9 +1950,8 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping,
         # get head dim
         head_dim = config.kv_lora_rank + config.qk_rope_head_dim
 
-        # provide at least 1 layer to prevent division by zero cache size
-        num_attention_layers = max(
-            len(mapping.pp_layers(model_config.get_num_attention_layers())), 1)
+        num_attention_layers = KVCacheManager._resolve_num_attention_layers(
+            model_config, mapping, num_layers)
         mem_per_token *= num_attention_layers * head_dim
 
         # 1 for K, others for indexer K cache

tensorrt_llm/_torch/attention_backend/sparse/rocket.py

@@ -1037,7 +1037,9 @@ def free_resources(self, request):
         self.kt_cache_manager.free_resources(request)
 
     @staticmethod
-    def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping,
+    def get_cache_size_per_token(model_config: ModelConfig,
+                                 mapping: Mapping,
+                                 num_layers: Optional[int] = None,
                                  **kwargs):
         # get kv cache dtype bytes
         mem_per_token = 2
@@ -1061,9 +1063,8 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping,
             head_dim = config.hidden_size // config.num_attention_heads
         head_dim = head_dim * num_key_value_heads // tp_size
 
-        # provide at least 1 layer to prevent division by zero cache size
-        num_attention_layers = max(
-            len(mapping.pp_layers(model_config.get_num_attention_layers())), 1)
+        num_attention_layers = KVCacheManager._resolve_num_attention_layers(
+            model_config, mapping, num_layers)
         mem_per_token *= num_attention_layers * head_dim
 
         # K and V

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -842,6 +842,26 @@ def calculate_scaling_factor_size_bytes(
         return get_size_in_bytes(cache_size // quant_vector_size,
                                  scaling_factor_dtype)
 
+    @staticmethod
+    def _resolve_num_attention_layers(
+        model_config: ModelConfigPython,
+        mapping: Mapping,
+        num_layers: Optional[int] = None,
+    ) -> int:
+        """Compute the effective number of attention layers for cache sizing.
+
+        When *num_layers* is explicitly provided (e.g. for draft models whose
+        HF config layer count differs from runtime), it is used directly
+        without PP distribution.  Otherwise the layer count is derived from
+        the model config and distributed evenly across PP ranks via
+        ``mapping.pp_layers``.
+        """
+        if num_layers is not None:
+            return max(num_layers, 1)
+        # provide at least 1 layer to prevent division by zero cache size
+        return max(
+            len(mapping.pp_layers(model_config.get_num_attention_layers())), 1)
+
     # TODO: refactor get_cache_size_per_token and get_cache_bytes_per_token to use the same logic
     @staticmethod
     def get_cache_size_per_token(model_config: ModelConfigPython,
@@ -870,18 +890,8 @@ def get_cache_size_per_token(model_config: ModelConfigPython,
             head_dim = head_dim * num_key_value_heads // tp_size
             kv_factor = 2
 
-        # When num_layers is explicitly provided (e.g. for draft models
-        # where the HF config layer count differs from runtime), use it
-        # directly without PP distribution.  Draft layers have their own
-        # PP assignment logic (see get_pp_layers) that doesn't match the
-        # standard uniform split, so pp_layers() would give wrong results.
-        if num_layers is not None:
-            num_attention_layers = max(num_layers, 1)
-        else:
-            # provide at least 1 layer to prevent division by zero cache size
-            num_attention_layers = max(
-                len(mapping.pp_layers(model_config.get_num_attention_layers())),
-                1)
+        num_attention_layers = KVCacheManager._resolve_num_attention_layers(
+            model_config, mapping, num_layers)
         # K and V
         mem_per_token = kv_factor * num_attention_layers * head_dim
         # The data type bytes.
@@ -2521,18 +2531,8 @@ def get_cache_size_per_token(model_config: ModelConfigPython,
             head_dim = head_dim * num_key_value_heads // tp_size
             kv_factor = 2
 
-        # When num_layers is explicitly provided (e.g. for draft models
-        # where the HF config layer count differs from runtime), use it
-        # directly without PP distribution.  Draft layers have their own
-        # PP assignment logic (see get_pp_layers) that doesn't match the
-        # standard uniform split, so pp_layers() would give wrong results.
-        if num_layers is not None:
-            num_attention_layers = max(num_layers, 1)
-        else:
-            # provide at least 1 layer to prevent division by zero cache size
-            num_attention_layers = max(
-                len(mapping.pp_layers(model_config.get_num_attention_layers())),
-                1)
+        num_attention_layers = KVCacheManager._resolve_num_attention_layers(
+            model_config, mapping, num_layers)
         mem_per_token *= num_attention_layers * head_dim
 
         # K and V