[HiCache] Auto-size the host pool to fit available host memory (#40135)

python/sglang/srt/arg_groups/fields/memory.py

@@ -119,6 +119,10 @@ class Memory(msgspec.Struct):
         int,
         "The size of host KV cache memory pool in gigabytes. Overrides --hicache-ratio in either host memory mode.",
     ] = 0
+    hicache_host_memory_fraction: A[
+        Optional[float],
+        "Fraction of the available host memory, bounded by visible cgroup memory.max/memory.high or v1 memory limits (after a 10 GiB reserve) that the HiCache host pools of all ranks on this machine may use. Applies only when neither --hicache-ratio nor --hicache-size is set: the default ratio is then reduced until the pools fit. Lower it when several engines share a memory cgroup.",
+    ] = 0.8
     hicache_write_policy: A[
         str,
         Arg(

python/sglang/srt/arg_groups/hicache_hook.py

@@ -69,16 +69,23 @@ def handle_hicache_ratio_default(server_args: Any):
 
     A decode server keeps the ratio unset here: kv_cache_builder resolves
     it against the retraction-backup backend (1.0 for host_pool, else 2.0).
+
+    An explicit --hicache-ratio or --hicache-size is honored as given, so it
+    resolves --hicache-host-memory-fraction to None (auto-sizing off).
     """
     cfg = resolving_view(server_args)
+    fraction = cfg.hicache_host_memory_fraction
+    if fraction is not None and not 0 < fraction <= 1:
+        raise ValueError("--hicache-host-memory-fraction must be in (0, 1].")
+    fields = {}
     if cfg.hicache_ratio is None and cfg.disaggregation_mode != "decode":
-        declare_resolution(
-            server_args,
-            "_handle_hicache_ratio_default",
-            hicache_ratio=(
-                1.2 if cfg.hicache_host_memory_mode == "buffer_only" else 2.0
-            ),
+        fields["hicache_ratio"] = (
+            1.2 if cfg.hicache_host_memory_mode == "buffer_only" else 2.0
         )
+    if cfg.hicache_ratio is not None or cfg.hicache_size > 0:
+        fields["hicache_host_memory_fraction"] = None
+    if fields:
+        declare_resolution(server_args, "_handle_hicache_ratio_default", **fields)
 
 
 def resolve_hicache_dcp_compatibility(server_args: Any):

python/sglang/srt/mem_cache/hicache_auto_size.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
+from sglang.srt.mem_cache.memory_pool import (
+    HybridLinearKVPool,
+    HybridReqToTokenPool,
+    MHATokenToKVPool,
+    MiniMaxSparseKVPool,
+    MLATokenToKVPool,
+)
+from sglang.srt.mem_cache.pool_host.base import (
+    host_memory_budget_bytes,
+    host_memory_budget_scope,
+    ranks_per_host,
+)
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+from sglang.srt.runtime_context import get_context, get_memory, get_parallel
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+    from sglang.srt.speculative.base_spec_worker import HiCacheDraftPlan
+
+logger = logging.getLogger(__name__)
+
+# Page rounding, allocator metadata and staging buffers are outside the device
+# byte counts the ratio is derived from.
+_ALLOCATION_SLACK_FRACTION = 0.05
+
+_SIZEABLE_POOLS = (
+    MHATokenToKVPool,
+    MLATokenToKVPool,
+    SWAKVPool,
+    HybridLinearKVPool,
+    MiniMaxSparseKVPool,
+)
+
+
+def _pool_bytes(pool) -> int:
+    if isinstance(pool, SWAKVPool):
+        return _pool_bytes(pool.full_kv_pool) + _pool_bytes(pool.swa_kv_pool)
+    if isinstance(pool, HybridLinearKVPool):
+        return _pool_bytes(pool.full_kv_pool)
+    sizes = pool.get_kv_size_bytes()
+    return sum(sizes) if isinstance(sizes, tuple) else sizes
+
+
+def _draft_bytes(target, draft) -> int:
+    if isinstance(draft, BaseSWAKVPool):
+        # Match sidecar construction: only SWA drafts follow target SWA slots.
+        target, draft = target.swa_kv_pool, draft.swa_kv_pool
+    # A sidecar has one host slot per target slot, however few slots the draft has.
+    return _pool_bytes(draft) * target.size // draft.size
+
+
+def _estimate_hicache_bytes(
+    params: CacheInitParams, draft_plan: HiCacheDraftPlan | None
+) -> int:
+    """Device bytes whose host mirrors scale with the HiCache ratio."""
+    pool = params.token_to_kv_pool_allocator.get_kvcache()
+    if not isinstance(pool, _SIZEABLE_POOLS):
+        raise ValueError(
+            f"HiCache auto-sizing does not support {type(pool).__name__}; "
+            "set --hicache-ratio or --hicache-size explicitly."
+        )
+    total = _pool_bytes(pool)
+    if isinstance(params.req_to_token_pool, HybridReqToTokenPool):
+        total += _pool_bytes(params.req_to_token_pool.mamba_pool)
+    drafts = params.mtp_draft_device_pools
+    if draft_plan is not None and draft_plan.mode == "sidecar":
+        drafts = draft_plan.device_pools
+    return total + sum(_draft_bytes(pool, draft) for draft in drafts)
+
+
+@contextmanager
+def auto_size_hicache(
+    params: CacheInitParams, draft_plan: HiCacheDraftPlan | None, *, enabled: bool
+):
+    """Reduce the default HiCache ratio until this machine's host pools fit.
+
+    Resolution nulls the fraction for an explicit --hicache-ratio/--hicache-size.
+    """
+    fraction = get_memory().hicache_host_memory_fraction
+    if not enabled or fraction is None:
+        yield
+        return
+    requested = get_memory().hicache_ratio
+    device_bytes = _estimate_hicache_bytes(params, draft_plan)
+    budget = int(host_memory_budget_bytes() * 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.
+    if torch.distributed.is_initialized():
+        value = torch.tensor([ratio], dtype=torch.float64)
+        torch.distributed.all_reduce(
+            value,
+            op=torch.distributed.ReduceOp.MIN,
+            group=get_parallel().world_group.cpu_group,
+        )
+        ratio = value.item()
+    if ratio <= 0:
+        raise ValueError(
+            "No host memory is left for HiCache after the 10 GiB reserve; "
+            "set --hicache-ratio or --hicache-size explicitly."
+        )
+    get_context().override("hicache.auto_size", hicache_ratio=ratio)
+    logger.info(
+        "HiCache auto-sizing: ratio %.3f -> %.3f; %.1f GiB host memory per rank "
+        "(fraction %.2f, %d ranks on this host), host pools %.1f GiB.",
+        requested,
+        ratio,
+        budget / 1024**3,
+        fraction,
+        ranks_per_host(),
+        device_bytes * ratio / 1024**3,
+    )
+    with host_memory_budget_scope(budget):
+        yield

python/sglang/srt/mem_cache/host_memory.py

@@ -0,0 +1,104 @@
+"""Host-memory headroom bounded by the process's visible cgroup hierarchy."""
+
+from __future__ import annotations
+
+import logging
+import re
+from pathlib import Path, PurePosixPath
+
+import psutil
+
+logger = logging.getLogger(__name__)
+
+
+def _unescape_mount_path(value: str) -> str:
+    return re.sub(r"\\([0-7]{3})", lambda m: chr(int(m[1], 8)), value)
+
+
+def _cgroup_memory_headroom(proc_root: Path = Path("/proc")) -> int | None:
+    memberships = {}
+    try:
+        cgroups = (proc_root / "self/cgroup").read_text()
+        mounts = (proc_root / "self/mountinfo").read_text()
+    except FileNotFoundError:
+        # Non-Linux systems need not expose procfs.
+        return None
+    for line in cgroups.splitlines():
+        _, controllers, path = line.split(":", 2)
+        if not controllers:
+            memberships["cgroup2"] = PurePosixPath(path)
+        elif "memory" in controllers.split(","):
+            memberships["cgroup"] = PurePosixPath(path)
+
+    headroom = None
+    resolved = False
+    for line in mounts.splitlines():
+        before, after = line.split(" - ", 1)
+        filesystem, _, options = after.split()[:3]
+        if filesystem not in memberships:
+            continue
+        if filesystem == "cgroup" and "memory" not in options.split(","):
+            continue
+        fields = before.split()
+        root = PurePosixPath(_unescape_mount_path(fields[3]))
+        mount = Path(_unescape_mount_path(fields[4]))
+        membership = memberships[filesystem]
+        if membership.is_relative_to(root):
+            relative = membership.relative_to(root)
+        elif root != PurePosixPath("/"):
+            # A cgroup namespace can expose membership relative to its root,
+            # while mountinfo still identifies the host-side subtree.
+            relative = membership.relative_to("/")
+        else:
+            continue
+        if ".." in relative.parts:
+            raise ValueError(f"Cannot resolve cgroup memory path: {membership}")
+        directory = mount / relative
+        if not directory.is_dir():
+            continue
+        resolved = True
+        limits = (
+            ("memory.max", "memory.high")
+            if filesystem == "cgroup2"
+            else ("memory.limit_in_bytes",)
+        )
+        usage_name = (
+            "memory.current" if filesystem == "cgroup2" else "memory.usage_in_bytes"
+        )
+        while True:
+            for name in limits:
+                try:
+                    value = (directory / name).read_text().strip()
+                except FileNotFoundError:
+                    # The hierarchy root may not have memory controller files.
+                    continue
+                if value == "max":
+                    continue
+                limit = int(value)
+                # Do not silently ignore an unreadable usage file for a known
+                # limit: falling back to host RAM could overrun the container.
+                usage = int((directory / usage_name).read_text())
+                remaining = max(0, limit - usage)
+                headroom = remaining if headroom is None else min(headroom, remaining)
+            if directory == mount:
+                break
+            directory = directory.parent
+    if memberships and not resolved:
+        raise RuntimeError(
+            "Cannot locate the process memory cgroup in mounted cgroup filesystems"
+        )
+    return headroom
+
+
+def available_host_memory_bytes() -> int:
+    """Conservative allocatable RAM; charged file cache is not assumed reclaimable."""
+    available = psutil.virtual_memory().available
+    cgroup_headroom = _cgroup_memory_headroom()
+    if cgroup_headroom is not None:
+        logger.info(
+            "HiCache memory headroom: host %.1f GiB, cgroup %.1f GiB",
+            available / 1024**3,
+            cgroup_headroom / 1024**3,
+        )
+        available = min(available, cgroup_headroom)
+    return available

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -40,6 +40,7 @@ from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.managers.mm_schedule import init_mm_embedding_cache
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+from sglang.srt.mem_cache.hicache_auto_size import auto_size_hicache
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
 from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
@@ -327,32 +328,36 @@ def build_kv_cache(
         mtp_draft_device_pools=mtp_draft_device_pools,
     )
 
-    tree_cache = create_tree_cache(
-        TreeCacheBuildContext(
-            server_args=server_args,
-            params=params,
-            is_hybrid_swa=is_hybrid_swa,
-            full_tokens_per_layer=full_tokens_per_layer,
-            is_hybrid_ssm=is_hybrid_ssm,
-            is_dsa=is_dsa,
-            enable_hierarchical_cache=enable_hierarchical_cache,
-            disable_radix_cache=disable_radix_cache,
-            effective_chunked_prefill_size=effective_chunked_prefill_size,
-            tp_worker=tp_worker,
-            model_config=model_config,
-            tp_size=ps.tp_size,
-            tp_rank=ps.tp_rank,
-            tp_group=tp_group,
-        )
+    tree_context = TreeCacheBuildContext(
+        server_args=server_args,
+        params=params,
+        is_hybrid_swa=is_hybrid_swa,
+        full_tokens_per_layer=full_tokens_per_layer,
+        is_hybrid_ssm=is_hybrid_ssm,
+        is_dsa=is_dsa,
+        enable_hierarchical_cache=enable_hierarchical_cache,
+        disable_radix_cache=disable_radix_cache,
+        effective_chunked_prefill_size=effective_chunked_prefill_size,
+        tp_worker=tp_worker,
+        model_config=model_config,
+        tp_size=ps.tp_size,
+        tp_rank=ps.tp_rank,
+        tp_group=tp_group,
     )
+    with auto_size_hicache(
+        params,
+        hicache_draft_plan,
+        enabled=enable_hierarchical_cache or retraction_backup == "host_pool",
+    ):
+        tree_cache = create_tree_cache(tree_context)
 
-    if (
-        enable_hierarchical_cache or retraction_backup == "host_pool"
-    ) and hicache_draft_plan is not None:
-        maybe_register_hicache_draft(
-            tree_cache=tree_cache,
-            draft_plan=hicache_draft_plan,
-        )
+        if (
+            enable_hierarchical_cache or retraction_backup == "host_pool"
+        ) and hicache_draft_plan is not None:
+            maybe_register_hicache_draft(
+                tree_cache=tree_cache,
+                draft_plan=hicache_draft_plan,
+            )
 
     if retraction_backup == "host_pool":
         if not isinstance(tree_cache, UnifiedRadixCache):

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

@@ -3,12 +3,14 @@ from __future__ import annotations
 import abc
 import logging
 import threading
+from contextlib import contextmanager
+from contextvars import ContextVar
 from functools import wraps
 from typing import Optional
 
-import psutil
 import torch
 
+from sglang.srt.mem_cache.host_memory import available_host_memory_bytes
 from sglang.srt.mem_cache.memory_pool import KVCache
 from sglang.srt.mem_cache.pool_host.common import (
     _cuda_host_unregister,
@@ -28,6 +30,21 @@ HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3)
 _WRITE_BACK_STAGING_PAGE_CHUNK = 64
 
 
+_host_memory_budget: ContextVar[Optional[int]] = ContextVar(
+    "hicache_host_memory_budget", default=None
+)
+
+
+@contextmanager
+def host_memory_budget_scope(budget_bytes: int):
+    """Book every pool built inside against one snapshot, not re-sampled psutil."""
+    token = _host_memory_budget.set(budget_bytes)
+    try:
+        yield
+    finally:
+        _host_memory_budget.reset(token)
+
+
 def ranks_per_host() -> int:
     """Number of ranks of this job running on the same machine as this one.
 
@@ -48,14 +65,23 @@ def ranks_per_host() -> int:
     return max(launch_world_size // get_parallel().nnodes, 1)
 
 
-def host_memory_budget_bytes() -> int:
+def host_memory_budget_bytes(requested_bytes: int = 0) -> 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.
+    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.
+
+    Inside host_memory_budget_scope, requested_bytes is booked against the
+    snapshot when it fits; the allowance before booking is returned.
     """
-    free = psutil.virtual_memory().available - HICACHE_HOST_MEMORY_RESERVE_BYTES
+    available = _host_memory_budget.get()
+    if available is not None:
+        if requested_bytes <= available:
+            _host_memory_budget.set(available - requested_bytes)
+        return available
+
+    free = available_host_memory_bytes() - HICACHE_HOST_MEMORY_RESERVE_BYTES
     return free // ranks_per_host()
 
 
@@ -172,7 +198,7 @@ class HostKVCache(abc.ABC):
 
         # Verify there is enough available host memory.
         requested_bytes = self.size * self.size_per_token
-        available_bytes = host_memory_budget_bytes()
+        available_bytes = host_memory_budget_bytes(requested_bytes)
         if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory available. Requesting "

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

@@ -111,7 +111,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
-        available_bytes = host_memory_budget_bytes()
+        available_bytes = host_memory_budget_bytes(requested_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/mamba.py

@@ -129,7 +129,7 @@ class MambaPoolHost(HostKVCache):
             )
 
         requested_bytes = self.size * self.size_per_token
-        available_bytes = host_memory_budget_bytes()
+        available_bytes = host_memory_budget_bytes(requested_bytes)
         if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory available. Requesting "

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

@@ -765,7 +765,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache):
         self.size_per_token = self.get_size_per_token()
 
         requested_bytes = self.size * self.size_per_token
-        available_bytes = host_memory_budget_bytes()
+        available_bytes = host_memory_budget_bytes(requested_bytes)
         if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory for MiniMax index-K hierarchical cache. "