fix(kv-cache): log capacity after block resolution (#1764)

Signed-off-by: AlpinDale <alpindale@gmail.com>

Author: AlpinDale

aphrodite/v1/core/kv_cache_utils.py

@@ -1752,6 +1752,20 @@ def get_kv_cache_capacity(aphrodite_config: AphroditeConfig, kv_cache_config: KV
     return int(max_concurrency * max_model_len), max_concurrency
 
 
+def update_kv_cache_capacity(aphrodite_config: AphroditeConfig, kv_cache_config: KVCacheConfig) -> None:
+    """Store and log the resolved KV cache capacity."""
+    num_tokens, max_concurrency = get_kv_cache_capacity(aphrodite_config, kv_cache_config)
+    aphrodite_config.cache_config.kv_cache_size_tokens = num_tokens
+    aphrodite_config.cache_config.kv_cache_max_concurrency = max_concurrency
+    max_model_len = aphrodite_config.model_config.max_model_len
+    logger.info_once(
+        "GPU KV cache size: %s tokens, Maximum concurrency for %s tokens per request: %.2fx",
+        f"{num_tokens:,}",
+        f"{max_model_len:,}",
+        max_concurrency,
+    )
+
+
 def _max_memory_usage_bytes_from_groups(
     aphrodite_config: AphroditeConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
@@ -2055,21 +2069,6 @@ def get_kv_cache_configs(
             assert tensor.size % num_blocks_old == 0
             tensor.size = tensor.size // num_blocks_old * min_num_blocks
 
-        if len(kv_cache_config.kv_cache_groups) > 0:
-            max_model_len = aphrodite_config.model_config.max_model_len
-            # GPU KV cache size in tokens = max_concurrency * max_model_len:
-            # the total tokens of context the pool can hold at peak
-            # utilization. Sourcing this from the concurrency calculation
-            # handles hybrid layouts correctly.
-            num_tokens, max_concurrency = get_kv_cache_capacity(aphrodite_config, kv_cache_config)
-
-            logger.info_once("GPU KV cache size: %s tokens", f"{num_tokens:,}")
-            logger.info_once(
-                "Maximum concurrency for %s tokens per request: %.2fx",
-                f"{max_model_len:,}",
-                max_concurrency,
-            )
-
     return kv_cache_configs
 
 

aphrodite/v1/engine/core.py

@@ -46,11 +46,11 @@
 from aphrodite.v1.core.kv_cache_utils import (
     BlockHash,
     generate_scheduler_kv_cache_config,
-    get_kv_cache_capacity,
     get_kv_cache_configs,
     get_request_block_hasher,
     init_none_hash,
     resolve_kv_cache_block_sizes,
+    update_kv_cache_capacity,
 )
 from aphrodite.v1.core.sched.interface import PauseState, SchedulerInterface
 from aphrodite.v1.core.sched.output import SchedulerOutput
@@ -287,9 +287,7 @@ def _initialize_kv_caches(self, aphrodite_config: AphroditeConfig) -> KVCacheCon
         kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
         if kv_cache_groups:
             aphrodite_config.cache_config.block_size = min(g.kv_cache_spec.block_size for g in kv_cache_groups)
-            num_tokens, max_concurrency = get_kv_cache_capacity(aphrodite_config, scheduler_kv_cache_config)
-            aphrodite_config.cache_config.kv_cache_size_tokens = num_tokens
-            aphrodite_config.cache_config.kv_cache_max_concurrency = max_concurrency
+            update_kv_cache_capacity(aphrodite_config, scheduler_kv_cache_config)
 
         aphrodite_config.validate_block_size()