[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. "