From 8f2a3ad6d7d68c58ae65b61a75bb2115449addca Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 4 Aug 2026 04:21:36 -0700 Subject: [PATCH] [mem_cache] Label HiCache host pools and clarify post-capture KV sizing logs (#33445) --- .../hybrid_cache/hybrid_pool_assembler.py | 8 ++++++- .../sglang/srt/mem_cache/kv_cache_builder.py | 1 + python/sglang/srt/mem_cache/pool_host/base.py | 11 ++++++++-- python/sglang/srt/mem_cache/pool_host/mha.py | 3 +++ python/sglang/srt/mem_cache/pool_host/mla.py | 3 +++ .../kv_pool_runtime.py | 21 ++++++++++++++++++- .../mem_cache/test_post_capture_kv_sizing.py | 7 ++++--- 7 files changed, 47 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 8faec5d77..1096cd958 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/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", ) diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 3e6ac0639..50f797101 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index 19770824e..57dc8c1a2 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 1af4bb972..150c37b72 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index 7066b6348..e31152072 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py b/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py index 96756e067..c0a16bc56 100644 --- a/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py +++ b/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( diff --git a/test/registered/mem_cache/test_post_capture_kv_sizing.py b/test/registered/mem_cache/test_post_capture_kv_sizing.py index e4638297d..a641e12c9 100644 --- a/test/registered/mem_cache/test_post_capture_kv_sizing.py +++ b/test/registered/mem_cache/test_post_capture_kv_sizing.py @@ -5,7 +5,7 @@ CUDA graphs, then sizes and physically backs the pool from measured free memory. This test launches a server with the feature enabled and asserts that: 1. the post-capture sizing path actually ran (log line present, not a silent no-op skip via post_capture_kv_sizing_planned), - 2. the pool was sized to a positive max_total_num_tokens, and + 2. the pool was backed to a positive KV size, and 3. gsm8k accuracy is unchanged vs. the default sizing path. """ @@ -69,14 +69,15 @@ class TestPostCaptureKVSizing(CustomTestCase): def test_post_capture_sizing_ran(self): """The post-capture path must actually execute, not silently skip.""" m = re.search( - r"Post-capture KV sizing: max_total_num_tokens=(\d+)", self._server_logs() + r"Post-capture KV sizing: KV cache allocated\..*?KV size: ([\d.]+) GB", + self._server_logs(), ) self.assertIsNotNone( m, "Post-capture KV sizing log line not found; the feature was gated off " "or the resize path did not run.", ) - self.assertGreater(int(m.group(1)), 0) + self.assertGreater(float(m.group(1)), 0) def test_server_info_pool_sized(self): info = requests.get(f"{self.base_url}/server_info").json()