diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 5fdbae704172..0828af3d12d9 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1920,7 +1920,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 @@ -1937,9 +1939,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 diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket.py index a76ba4d2e569..fa7e8e4af659 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/rocket.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 43a7503655b6..f2f74381660c 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/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