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):