[mem_cache] Label HiCache host pools and clarify post-capture KV sizing logs (#33445)

python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py

@@ -66,6 +66,7 @@ def build_kv_host_pool(
     use_mla: bool,
     override_kv_cache_dim: Optional[int] = None,
     host_size: Optional[float] = None,
+    pool_label: str = "kv",
 ):
     kv_host_pool_cls = (
         MLATokenToKVPoolHost if use_mla else get_mha_host_pool_cls(kv_pool)
@@ -88,6 +89,7 @@ def build_kv_host_pool(
         page_size,
         server_args.hicache_mem_layout,
         allocator_type=_get_allocator_type(server_args),
+        pool_label=pool_label,
         **kwargs,
     )
 
@@ -219,6 +221,7 @@ def build_hybrid_swa_stack(
         server_args=server_args,
         use_mla=use_mla,
         host_size=kv_host_size,
+        pool_label="full",
     )
     swa_host_pool = build_kv_host_pool(
         kv_pool=swa_kv_pool,
@@ -226,6 +229,7 @@ def build_hybrid_swa_stack(
         server_args=server_args,
         use_mla=use_mla,
         host_size=swa_host_size,
+        pool_label="swa",
     )
 
     # For SWA hybrid, the device alloc/free goes through the inner swa_attn_allocator
@@ -665,6 +669,7 @@ def build_hybrid_mamba_swa_stack(
         server_args=server_args,
         use_mla=False,
         host_size=kv_host_size,
+        pool_label="full",
     )
     swa_host_pool = build_kv_host_pool(
         kv_pool=swa_kv_pool,
@@ -672,6 +677,7 @@ def build_hybrid_mamba_swa_stack(
         server_args=server_args,
         use_mla=False,
         host_size=swa_host_size,
+        pool_label="swa",
     )
     mamba_host_pool = MambaPoolHost(
         mamba_pool,
@@ -1038,7 +1044,7 @@ class _SwaStrategy(StackStrategy):
                 ComponentType.SWA: host_pool_group.get_pool(PoolName.SWA),
             },
             transfer_layer_num=len(full_layer_mapping | swa_layer_mapping),
-            pools_desc="KV + SWA",
+            pools_desc="Full + SWA",
         )
 
 

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -112,6 +112,7 @@ def maybe_register_hicache_draft(
         page_size=page_size,
         layout=server_args.hicache_mem_layout,
         allocator_type=server_args.hicache_storage_backend,
+        pool_label="draft",
     )
     if isinstance(pool, MHATokenToKVPool):
         draft_host_pool = get_mha_host_pool_cls(pool)(pool, **kw)

python/sglang/srt/mem_cache/pool_host/base.py

@@ -94,8 +94,11 @@ class HostKVCache(abc.ABC):
         allocator_type: str = "default",
         dcp_size: int = 1,
         dcp_rank: int = 0,
+        *,
+        pool_label: str = "kv",
     ):
         self.device_pool = device_pool
+        self.pool_label = pool_label
         # page_size arrives widened (x dcp_size); size/page_size/page_num are physical.
         self.dcp_size = dcp_size
         self.dcp_rank = dcp_rank
@@ -127,9 +130,10 @@ class HostKVCache(abc.ABC):
 
         if self.size <= device_pool.size:
             logger.warning(
-                "HiCache host KV pool (%d tokens) is smaller than the device pool (%d tokens);"
+                "HiCache %s host pool (%d tokens) is smaller than the device pool (%d tokens);"
                 "L2 cache effectiveness is reduced."
                 "Consider increasing --hicache-ratio (or --hicache-size) for higher L2 cache hit rate.",
+                pool_label,
                 self.size,
                 device_pool.size,
             )
@@ -147,7 +151,10 @@ class HostKVCache(abc.ABC):
             )
         else:
             logger.info(
-                f"Allocating {requested_bytes / 1e9:.2f} GB host memory for hierarchical KV cache."
+                "Allocating %s hierarchical KV host pool: %d tokens, %.2f GB host memory.",
+                pool_label,
+                self.size,
+                requested_bytes / 1e9,
             )
 
         self.kv_buffer = self.init_kv_buffer()

python/sglang/srt/mem_cache/pool_host/mha.py

@@ -79,6 +79,8 @@ class MHATokenToKVPoolHost(HostKVCache):
         pin_memory: bool = True,
         device: str = "cpu",
         allocator_type: str = "default",
+        *,
+        pool_label: str = "kv",
     ):
         super().__init__(
             device_pool,
@@ -89,6 +91,7 @@ class MHATokenToKVPoolHost(HostKVCache):
             pin_memory,
             device,
             allocator_type,
+            pool_label=pool_label,
         )
         self.element_dim = self.device_pool.head_num * self.device_pool.head_dim
         # The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only

python/sglang/srt/mem_cache/pool_host/mla.py

@@ -64,6 +64,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
         override_kv_cache_dim: Optional[int] = None,
         dcp_size: int = 1,
         dcp_rank: int = 0,
+        *,
+        pool_label: str = "kv",
     ):
         self.override_kv_cache_dim = override_kv_cache_dim
         super().__init__(
@@ -77,6 +79,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
             allocator_type,
             dcp_size=dcp_size,
             dcp_rank=dcp_rank,
+            pool_label=pool_label,
         )
         # The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only
         # helpers in hicache.cuh are guarded by USE_ROCM and the staged

python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py

@@ -102,9 +102,28 @@ def compute_post_capture_kv_resize(
                 capped_reqs,
             )
             capped_max_running_requests = capped_reqs
+    # Two-line summary mirroring the pre-capture pool logs: the resized pool
+    # shape (cf. "Use sliding window memory pool") then its backed KV footprint
+    # and post-resize free memory (cf. "KV Cache is allocated" + "Memory pool end").
+    # Non-hybrid pools leave the full/swa splits unset, so fall back to the single
+    # total for the full-layer count.
+    full_layer_tokens = (
+        config.full_max_total_num_tokens
+        if config.full_max_total_num_tokens is not None
+        else config.max_total_num_tokens
+    )
+    swa_layer_tokens = config.swa_max_total_num_tokens or 0
     logger.info(
-        "Post-capture KV sizing: max_total_num_tokens=%d, free memory=%.2f GB",
+        "Post-capture KV sizing: full_layer_tokens=%d, swa_layer_tokens=%d",
+        full_layer_tokens,
+        swa_layer_tokens,
+    )
+    logger.info(
+        "Post-capture KV sizing: KV cache allocated. dtype: %s, #tokens: %d, "
+        "KV size: %.2f GB, avail mem=%.2f GB",
+        pool.dtype,
         config.max_total_num_tokens,
+        pool.post_capture_backed_bytes / (1 << 30),
         get_available_gpu_memory(model_runner.device, model_runner.gpu_id),
     )
     return PostCaptureKVResize(