[HiCache] Split the host-memory budget across co-located ranks (#35540)

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -9,7 +9,6 @@ if TYPE_CHECKING:
     from sglang.srt.mem_cache.hicache_storage import PoolName
     from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
 
-import psutil
 import torch
 
 from sglang.kernels.ops.kvcache.hicache import (
@@ -46,7 +45,7 @@ logger = logging.getLogger(__name__)
 from sglang.srt.mem_cache.pool_host import HostKVCache
 from sglang.srt.mem_cache.pool_host.base import (
     _WRITE_BACK_STAGING_PAGE_CHUNK,
-    HICACHE_HOST_MEMORY_RESERVE_BYTES,
+    host_memory_budget_bytes,
     synchronized,
 )
 from sglang.srt.mem_cache.pool_host.common import (
@@ -216,8 +215,7 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache):
         self.gpu_device = device_buffers[0].device if device_buffers else device
 
         requested_bytes = self.layer_num * num_host_pages * self.item_bytes
-        host_mem = psutil.virtual_memory()
-        available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
+        available_bytes = host_memory_budget_bytes()
         if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory for V4 paged pool {pool_name}. "
@@ -620,8 +618,7 @@ class DeepSeekV4StateHostPool(HostKVCache):
         self.size_per_token = self.state_page_bytes
 
         requested_bytes = self.layer_num * num_host_pages * self.state_page_bytes
-        host_mem = psutil.virtual_memory()
-        available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
+        available_bytes = host_memory_budget_bytes()
         if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory for V4 state pool {pool_name}. "
@@ -1132,8 +1129,7 @@ class DSAIndexerPoolHost(HostKVCache):
 
         buf_elem_size = self.page_num * self.layer_num * self.indexer_page_stride_size
         requested_bytes = buf_elem_size * self.indexer_dtype.itemsize
-        host_mem = psutil.virtual_memory()
-        available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
+        available_bytes = host_memory_budget_bytes()
         if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory for DSA indexer hierarchical cache. "

python/sglang/srt/mem_cache/pool_host/base.py

@@ -9,11 +9,13 @@ from typing import Optional
 import psutil
 import torch
 
+from sglang.srt.distributed.parallel_state import get_world_group
 from sglang.srt.mem_cache.memory_pool import KVCache
 from sglang.srt.mem_cache.pool_host.common import (
     _cuda_host_unregister,
     get_allocator_from_storage,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import is_cuda, is_hip
 
 logger = logging.getLogger(__name__)
@@ -27,6 +29,36 @@ HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3)
 _WRITE_BACK_STAGING_PAGE_CHUNK = 64
 
 
+def ranks_per_host() -> int:
+    """Number of ranks of this job running on the same machine as this one.
+
+    Derived as world_size // nnodes: the launcher slices ranks uniformly
+    across nodes (resolution asserts divisibility), so no hostname collective
+    is needed — a collective here would have to be issued the same number of
+    times on every rank, and ranks build different numbers of host pools.
+    """
+    if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+        return 1
+    try:
+        world_group = get_world_group()
+    except AssertionError:
+        return 1
+    if world_group.world_size == 1:
+        return 1
+    return max(world_group.world_size // get_parallel().nnodes, 1)
+
+
+def host_memory_budget_bytes() -> int:
+    """Host RAM this rank may claim for a HiCache pool.
+
+    psutil reports the whole machine, so co-located ranks each see the same free
+    memory; without the split every rank sizes its pool against all of it and
+    the host is oversubscribed by the number of ranks it holds.
+    """
+    free = psutil.virtual_memory().available - HICACHE_HOST_MEMORY_RESERVE_BYTES
+    return free // ranks_per_host()
+
+
 def sync_fixed_hicache_size(size: int, host_size: int) -> int:
     """Sync fixed-size HiCache token capacity across PP ranks.
 
@@ -139,9 +171,8 @@ class HostKVCache(abc.ABC):
             )
 
         # Verify there is enough available host memory.
-        host_mem = psutil.virtual_memory()
         requested_bytes = self.size * self.size_per_token
-        available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
+        available_bytes = host_memory_budget_bytes()
         if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory available. Requesting "

python/sglang/srt/mem_cache/pool_host/mamba.py

@@ -5,13 +5,12 @@ import threading
 from typing import Optional
 
 import numpy as np
-import psutil
 import torch
 
 from sglang.srt.mem_cache.memory_pool import MambaPool
 from sglang.srt.mem_cache.pool_host.base import (
-    HICACHE_HOST_MEMORY_RESERVE_BYTES,
     HostKVCache,
+    host_memory_budget_bytes,
     sync_fixed_hicache_size,
     synchronized,
 )
@@ -96,9 +95,8 @@ class MambaPoolHost(HostKVCache):
                 device_pool.size,
             )
 
-        host_mem = psutil.virtual_memory()
         requested_bytes = self.size * self.size_per_token
-        available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
+        available_bytes = host_memory_budget_bytes()
         if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory available. Requesting "

python/sglang/srt/mem_cache/pool_host/mha.py

@@ -4,7 +4,6 @@ import logging
 import threading
 from typing import Sequence
 
-import psutil
 import torch
 
 from sglang.kernels.ops.kvcache.hicache import (
@@ -32,8 +31,8 @@ from sglang.kernels.ops.kvcache.hicache import (
 from sglang.srt.mem_cache.memory_pool import MHATokenToKOnlyPool, MHATokenToKVPool
 from sglang.srt.mem_cache.pool_host.base import (
     _WRITE_BACK_STAGING_PAGE_CHUNK,
-    HICACHE_HOST_MEMORY_RESERVE_BYTES,
     HostKVCache,
+    host_memory_budget_bytes,
 )
 from sglang.srt.mem_cache.pool_host.common import (
     ALLOC_MEMORY_FUNCS,
@@ -724,9 +723,8 @@ class MHATokenToKOnlyPoolHost(HostKVCache):
         self.page_num = anchor_host.page_num
         self.size_per_token = self.get_size_per_token()
 
-        host_mem = psutil.virtual_memory()
         requested_bytes = self.size * self.size_per_token
-        available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
+        available_bytes = host_memory_budget_bytes()
         if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory for MiniMax index-K hierarchical cache. "

test/registered/unit/mem_cache/test_mem_pool_host.py

@@ -2,6 +2,7 @@
 
 import threading
 import unittest
+import unittest.mock
 
 import torch
 
@@ -10,8 +11,10 @@ from sglang.srt.mem_cache.memory_pool_host import (
     DeepSeekV4PagedHostPool,
     LogicalHostPool,
 )
+from sglang.srt.mem_cache.pool_host import base
 from sglang.srt.mem_cache.pool_host.mamba import MambaPoolHost
 from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
+from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -198,5 +201,42 @@ class TestLazyHostPoolRelease(CustomTestCase):
             pool.free(torch.tensor([0]))
 
 
+class TestHostMemoryBudget(CustomTestCase):
+    # Pinned so the two budget reads below see identical free memory; the real
+    # psutil value drifts between calls and would flake the equality checks.
+    _AVAILABLE = base.HICACHE_HOST_MEMORY_RESERVE_BYTES + 64 * (1024**3)
+
+    def _budget_with_ranks(self, ranks):
+        # Deliberate single-accessor stub: isolates the budget math from the
+        # topology derivation, which the ranks_per_host case below covers.
+        fake_mem = unittest.mock.Mock(available=self._AVAILABLE)
+        with unittest.mock.patch.object(
+            base, "ranks_per_host", return_value=ranks
+        ), unittest.mock.patch.object(
+            base.psutil, "virtual_memory", return_value=fake_mem
+        ):
+            return base.host_memory_budget_bytes()
+
+    def test_budget_is_split_across_co_located_ranks(self):
+        solo = self._budget_with_ranks(1)
+        self.assertEqual(self._budget_with_ranks(4), solo // 4)
+
+    def test_reserve_is_taken_before_the_split(self):
+        # Each rank must not get its own copy of the reserve.
+        budget = self._budget_with_ranks(8)
+        self.assertLessEqual(
+            budget * 8, self._AVAILABLE - base.HICACHE_HOST_MEMORY_RESERVE_BYTES
+        )
+
+    def test_ranks_per_host_divides_world_size_by_nodes(self):
+        # The launcher slices ranks uniformly across nodes, so the co-located
+        # rank count is world_size // nnodes — no hostname collective.
+        fake_group = unittest.mock.Mock(world_size=16)
+        with get_context().override_server_args(nnodes=2), unittest.mock.patch.object(
+            torch.distributed, "is_initialized", return_value=True
+        ), unittest.mock.patch.object(base, "get_world_group", return_value=fake_group):
+            self.assertEqual(base.ranks_per_host(), 8)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_minimax_sparse_pool_host_unit.py

@@ -9,9 +9,7 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
     HybridCacheController,
 )
 from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
-from sglang.srt.mem_cache.memory_pool_host import (
-    HICACHE_HOST_MEMORY_RESERVE_BYTES,
-)
+from sglang.srt.mem_cache.pool_host.base import HICACHE_HOST_MEMORY_RESERVE_BYTES
 from sglang.srt.mem_cache.pool_host.common import (
     ALLOC_MEMORY_FUNCS,
     alloc_with_pin_memory,