From d0058e64da72b9cb9438f80ce32d733f4f460433 Mon Sep 17 00:00:00 2001 From: Venkat Raman Date: Mon, 19 Oct 2026 16:21:27 +0200 Subject: [PATCH] [HiCache] Size each rank's host pool from the full host headroom, not headroom // ranks (upstream #38157) The kimi-k3 v0.5.22 prefill/decode line: v0.5.22 plus this one layer. It is 331f64442c (the v0.5.21 kimi-k3 line's layer 5b, fork PR #13) rebased onto v0.5.22, with the two conflicts against v0.5.22's pool_host changes resolved. host_memory_budget_bytes() outside a budget scope returns free // ranks_per_host(), where free is the host headroom left when the rank asks. On an 8-GPU pod the ranks arrive at different times, so a rank that checks after its peers have pinned their pools divides a smaller headroom by 8 again and is refused although every pool fits (seen 2026-10-01 on Kimi-K3: "Requesting 24.77 GB but only have 22.74 GB free"). v0.5.22 still does this with an explicit --hicache-size (hicache_auto_size.py opens no budget scope; pool_host/base.py and pool_host/mamba.py check against free // ranks_per_host()). The guard now checks a pool against the whole remaining headroom; the two callers that size a pool (HiCache auto-sizing, the unified page-envelope pool when --hicache-size and --hicache-ratio are 0) keep this rank's share through host_memory_budget_bytes_per_rank(). The unit tests that pinned the split are updated. Not carried from the v0.5.21 line, for the prefill/decode copy (prefill at DCP 1 with HiCache and Mooncake L3, decode without HiCache): #41692 to #41695 (file L3 and DCP > 1 L3 only) and #42039 (in v0.5.22 as #42420 and #42788). The aggregated copy (Mooncake L3 at DCP 8) stays on verda/kimi-k3-2026-10-03. Validated on HGX B300 with Kimi-K3 (19 Oct 2026, v0.5.22 image + this layer as a Python overlay): same-time A/B against the v0.5.21 line on a 1 prefill + 2 decode copy (weka median 24 twice, AgentX 48), and a 2 prefill + 2 decode copy (AgentX 56 twice, 64): level or better, no cross-request leaks, no broken answers. Upstream: https://github.com/sgl-project/sglang/pull/38157 Issue: https://github.com/sgl-project/sglang/issues/38156 Co-Authored-By: Claude Opus 5.5 --- .../sglang/srt/mem_cache/hicache_auto_size.py | 7 ++++-- python/sglang/srt/mem_cache/pool_host/base.py | 24 +++++++++++++++---- .../sglang/srt/mem_cache/pool_host/unified.py | 9 +++++-- .../unit/mem_cache/test_mem_pool_host.py | 14 ++++++----- 4 files changed, 40 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/mem_cache/hicache_auto_size.py b/python/sglang/srt/mem_cache/hicache_auto_size.py index f97bc226fc66..6757ed933ef8 100644 --- a/python/sglang/srt/mem_cache/hicache_auto_size.py +++ b/python/sglang/srt/mem_cache/hicache_auto_size.py @@ -15,7 +15,7 @@ MLATokenToKVPool, ) from sglang.srt.mem_cache.pool_host.base import ( - host_memory_budget_bytes, + host_memory_budget_bytes_per_rank, host_memory_budget_scope, ranks_per_host, ) @@ -93,7 +93,10 @@ def auto_size_hicache( return requested = get_memory().hicache_ratio device_bytes = _estimate_hicache_bytes(params, draft_plan) - budget = int(host_memory_budget_bytes(auto_size=True) * fraction) + # Auto-sizing keeps the per-rank split: host_memory_budget_bytes() returns + # the whole remaining headroom, since dividing it inside the guard refused + # late ranks whose explicitly sized pools fit (upstream #38156). + budget = int(host_memory_budget_bytes_per_rank(auto_size=True) * fraction) ratio = min(requested, budget * (1 - _ALLOCATION_SLACK_FRACTION) / device_bytes) # One collective before any pool is built: PP stages own different pool # counts, so a per-pool collective could deadlock. diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index 630e1f856466..b113dd99f6e4 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/python/sglang/srt/mem_cache/pool_host/base.py @@ -86,12 +86,19 @@ def host_memory_budget_bytes( ) -> int: """Host RAM this rank may claim for a HiCache pool. - Bound machine availability by the visible cgroup limits before splitting - among local ranks. Independent engines with separate container budgets - therefore size against their own remaining allowance. + Bound machine availability by the visible cgroup limits. Independent + engines with separate container budgets therefore size against their own + remaining allowance. Auto-sizing requires successful cgroup discovery; explicit pool sizing can fall back to checking host availability when discovery fails. + The allowance is not divided by ranks_per_host(): outside a budget scope + each rank reads the host's remaining headroom at a different moment, so a + rank that checks after its peers have pinned their pools would be refused + even though all pools fit (upstream #38156). Explicitly sized pools are + checked one at a time against the whole headroom; HiCache auto-sizing + divides by ranks_per_host() itself before booking a budget scope. + Inside host_memory_budget_scope, requested_bytes is booked against the snapshot when it fits; the allowance before booking is returned. """ @@ -105,7 +112,16 @@ def host_memory_budget_bytes( available_host_memory_bytes(allow_cgroup_fallback=not auto_size) - HICACHE_HOST_MEMORY_RESERVE_BYTES ) - return free // ranks_per_host() + return max(0, free) + + +def host_memory_budget_bytes_per_rank(*, auto_size: bool = False) -> int: + """Host RAM one rank may plan for when every local rank sizes itself the + same way (HiCache auto-sizing): the remaining headroom split evenly. + + Kept out of host_memory_budget_bytes() so the explicit-size guard checks a + pool against the whole headroom (upstream #38156).""" + return host_memory_budget_bytes(auto_size=auto_size) // ranks_per_host() def sync_fixed_hicache_size(size: int, host_size: int) -> int: diff --git a/python/sglang/srt/mem_cache/pool_host/unified.py b/python/sglang/srt/mem_cache/pool_host/unified.py index 4ffd5e8843e5..557db1b2509f 100644 --- a/python/sglang/srt/mem_cache/pool_host/unified.py +++ b/python/sglang/srt/mem_cache/pool_host/unified.py @@ -6,7 +6,11 @@ import torch -from sglang.srt.mem_cache.pool_host.base import HostKVCache, host_memory_budget_bytes +from sglang.srt.mem_cache.pool_host.base import ( + HostKVCache, + host_memory_budget_bytes, + host_memory_budget_bytes_per_rank, +) from sglang.srt.mem_cache.pool_host.common import ( ALLOC_MEMORY_FUNCS, _cuda_host_unregister, @@ -111,7 +115,8 @@ def __init__( ] if host_size == 0 and host_to_device_ratio == 0: - total_bytes = host_memory_budget_bytes() + # Sizing, not checking: keep this rank's share (see base.py). + total_bytes = host_memory_budget_bytes_per_rank() elif host_size > 0: total_bytes = max(nominal_view_bytes) else: diff --git a/test/registered/unit/mem_cache/test_mem_pool_host.py b/test/registered/unit/mem_cache/test_mem_pool_host.py index 1c8f3702a119..42ac1baf84a6 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/test/registered/unit/mem_cache/test_mem_pool_host.py @@ -449,15 +449,17 @@ def _budget_with_ranks(self, ranks): ): return base.host_memory_budget_bytes() - def test_budget_is_split_across_co_located_ranks(self): + def test_budget_is_not_split_across_co_located_ranks(self): + # Outside a budget scope every rank sees the whole remaining headroom; + # splitting by rank count refused late ranks whose pools fit (#38156). solo = self._budget_with_ranks(1) - self.assertEqual(self._budget_with_ranks(4), solo // 4) + self.assertEqual(self._budget_with_ranks(4), solo) + self.assertEqual(self._budget_with_ranks(8), solo) - def test_reserve_is_taken_before_the_split(self): - # Each rank must not get its own copy of the reserve. + def test_reserve_is_taken_from_the_budget(self): budget = self._budget_with_ranks(8) - self.assertLessEqual( - budget * 8, self._AVAILABLE - base.HICACHE_HOST_MEMORY_RESERVE_BYTES + self.assertEqual( + budget, self._AVAILABLE - base.HICACHE_HOST_MEMORY_RESERVE_BYTES ) def test_ranks_per_host_divides_world_size_by_nodes(self):