From e4aecc244ab8016e471f9c2b48a83d4fd52ce894 Mon Sep 17 00:00:00 2001 From: Venkat Raman Date: Tue, 27 Oct 2026 12:17:40 +0100 Subject: [PATCH] [HiCache] Size each rank's host pool from the full host headroom, not headroom // ranks (upstream #38157) Layer 2 of the Kimi-K3 nightly line on upstream main 5884e4fec6 (the lmsysorg/sglang:nightly-dev-cu13-20261027-5884e4fe image). host_memory_budget_bytes() outside a budget scope returned free // ranks_per_host(), so on an 8-GPU host a rank that asks after its peers have pinned their pools divides a smaller headroom by 8 again and is refused although every pool fits. The guard now checks a pool against the whole remaining headroom; the callers that size a pool keep this rank's share through host_memory_budget_bytes_per_rank(). Source: 894d2ed8b5 on the kimi-k3-nightly-3c560089 line (the change of fork PR #14 on the v0.5.22 line), cherry-picked unchanged; the test hunk applies at a 4-line offset. Upstream #38157 does the same job a different way and is still open. Callers added upstream since 3c560089e1 were checked: the MiniMax-M3.1 NVFP4 host pool (pool_host/minimax_m3_nvfp4.py) and its HiCache guard in hybrid_pool_assembler.py check an explicitly sized pool, so they take the whole headroom like the other host pools; no new caller sizes a pool. Tests (CPU): test_mem_pool_host.py with this commit's version on 5884e4fec6: 2 failed, the two cases that pinned the split (the budget came back as the headroom divided by 4 and by 8 ranks), 28 passed; with this commit 30 passed. Fork PR (v0.5.22 line): https://github.com/datacrunch-research/sglang/pull/14 Upstream: https://github.com/sgl-project/sglang/pull/38157 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 65c9297af5e6..d7d4489d9fd3 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 3572f6f60811..f9a8caa264a9 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/test/registered/unit/mem_cache/test_mem_pool_host.py @@ -687,15 +687,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):