[kimi-k3] HiCache: size each rank's host pool from the full host headroom (upstream #38157)

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.

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:

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:

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