From 1c82955861e06777638a99824980930ee001c004 Mon Sep 17 00:00:00 2001 From: cctry Date: Wed, 19 Aug 2026 13:01:58 -0700 Subject: [PATCH] [HiCache] Split the host-memory budget across co-located ranks (#35540) --- .../sglang/srt/mem_cache/memory_pool_host.py | 12 ++---- python/sglang/srt/mem_cache/pool_host/base.py | 35 +++++++++++++++- .../sglang/srt/mem_cache/pool_host/mamba.py | 6 +-- python/sglang/srt/mem_cache/pool_host/mha.py | 6 +-- .../unit/mem_cache/test_mem_pool_host.py | 40 +++++++++++++++++++ .../test_minimax_sparse_pool_host_unit.py | 4 +- 6 files changed, 82 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 1ed1645b3..2bb22431b 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/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. " diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index e229b418e..1c258a771 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/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 " diff --git a/python/sglang/srt/mem_cache/pool_host/mamba.py b/python/sglang/srt/mem_cache/pool_host/mamba.py index 8560ec8a5..7e7cda04b 100644 --- a/python/sglang/srt/mem_cache/pool_host/mamba.py +++ b/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 " diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 4646fb255..ea694e924 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/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. " 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 02dc0e531..be68761cc 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_minimax_sparse_pool_host_unit.py b/test/registered/unit/mem_cache/test_minimax_sparse_pool_host_unit.py index 3bdf9df0f..94c98edc2 100644 --- a/test/registered/unit/mem_cache/test_minimax_sparse_pool_host_unit.py +++ b/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,