[#15178][fix] Account for reclaimable memory in V2 KV cache sizing on integrated GPUs

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -23,8 +23,9 @@
 import tensorrt_llm
 import tensorrt_llm.bindings.executor as trtllm
 from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version,
-                                 is_sm_100f, prefer_pinned,
-                                 str_dtype_to_binding, torch_dtype_to_str)
+                                 is_device_integrated, is_sm_100f,
+                                 prefer_pinned, str_dtype_to_binding,
+                                 torch_dtype_to_str)
 from tensorrt_llm.inputs.multimodal import MultimodalParams
 
 # isort: off
@@ -73,6 +74,7 @@
                                            MambaHybridCacheManagerV2,
                                            MixedMambaHybridCacheManager,
                                            use_py_mamba_cache_manager)
+from .kv_cache.memory_utils import get_integrated_gpu_free_memory
 from .llm_request import ExecutorResponse, LlmRequestState
 from .model_engine import ModelEngine
 from .py_executor import PyExecutor, PyExecutorResources
@@ -1364,6 +1366,28 @@ def scaled(cost: CacheCost) -> CacheCost:
                         include_request_cost=include_request_cost))
         return None
 
+    def _get_device_free_memory(self) -> tuple[int, int]:
+        """Return ``(free, total)`` device bytes for KV cache sizing.
+
+        Equivalent to ``torch.cuda.mem_get_info()`` except for the V2 KV cache
+        manager on integrated GPUs, where free memory also counts reclaimable
+        page cache (see ``get_integrated_gpu_free_memory``). The V2 manager
+        fails at construction when its quota cannot hold the live recurrent
+        states, so counting page cache as consumed device memory makes hybrid
+        models fail to start on unified-memory systems after a cold image pull.
+        Discrete GPUs and the V1 manager are unchanged.
+        """
+        free, total = torch.cuda.mem_get_info()
+        if not (self._is_kv_cache_manager_v2 and is_device_integrated()):
+            return free, total
+        adjusted = get_integrated_gpu_free_memory(free, total)
+        if adjusted != free:
+            logger.info(
+                "Integrated GPU: sizing the V2 KV cache from "
+                f"{adjusted / GB:.2f} GiB of free memory including reclaimable "
+                f"page cache (CUDA reports {free / GB:.2f} GiB free).")
+        return adjusted, total
+
     def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
                         allocated_bytes: int) -> int:
         """
@@ -1747,7 +1771,7 @@ def try_prepare_estimation(self) -> bool:
             ) if self._kv_cache_config.max_tokens is not None else estimate_max_tokens
             self._kv_cache_config.avg_seq_len = self._max_seq_len
             if self._is_kv_cache_manager_v2:
-                free_mem, _ = torch.cuda.mem_get_info()
+                free_mem, _ = self._get_device_free_memory()
                 max_gpu_total_bytes = int(
                     self._kv_cache_config.free_gpu_memory_fraction * free_mem)
                 if (self._max_gpu_total_bytes_in is not None
@@ -1831,7 +1855,7 @@ def configure_kv_cache_capacity(self,
 
         torch.cuda.empty_cache()
         torch.cuda.reset_peak_memory_stats()
-        end, total_gpu_memory = torch.cuda.mem_get_info()
+        end, total_gpu_memory = self._get_device_free_memory()
         total_used_bytes = total_gpu_memory - end
         model_bytes = torch.cuda.memory_stats()["allocated_bytes.all.current"]
         logger.info(
@@ -1880,7 +1904,7 @@ def configure_kv_cache_capacity(self,
 
                 # Clear the caching allocator before measuring the current memory usage
                 torch.cuda.empty_cache()
-                end, total_gpu_memory = torch.cuda.mem_get_info()
+                end, total_gpu_memory = self._get_device_free_memory()
                 torch_used_bytes = torch.cuda.memory_stats(
                 )["allocated_bytes.all.current"]
             finally:
@@ -4831,7 +4855,12 @@ def _try_infer_num_experts(model_config: ModelConfig) -> int:
     return num_experts
 
 
-def _adjust_torch_mem_fraction():
+def _adjust_torch_mem_fraction(*, use_reclaimable_memory: bool = False) -> None:
+    """Cap PyTorch allocations, counting shared page cache for V2 integrated GPUs.
+
+    Args:
+        use_reclaimable_memory: Whether the executor uses V2 memory accounting.
+    """
     # If true, adjust PyTorch CUDA memory fraction to correspond to the
     # total GPU memory minus the statically allocated engine memory.
     # If false, set the PyTorch CUDA memory fraction to 1.0.
@@ -4869,8 +4898,11 @@ def _adjust_torch_mem_fraction():
         return
     mem_reserved = torch.cuda.memory_reserved()
     mem_free, mem_total = torch.cuda.mem_get_info()
+    if use_reclaimable_memory and is_device_integrated():
+        mem_free = get_integrated_gpu_free_memory(mem_free, mem_total)
     safety_margin = 32 * 1024**2
-    mem_torch_max = mem_free + mem_reserved - safety_margin
+    mem_torch_max = min(mem_total,
+                        max(0, mem_free + mem_reserved - safety_margin))
     mem_torch_fraction = mem_torch_max / mem_total
     logger.info(
         f"Setting PyTorch memory fraction to {mem_torch_fraction} ({mem_torch_max / 1024**3} GiB)"

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -39,6 +39,7 @@
 from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout
 from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp
 from tensorrt_llm._torch.model_config import _get_kv_cache_layer_specs
+from tensorrt_llm._torch.pyexecutor.kv_cache.memory_utils import get_integrated_gpu_free_memory
 from tensorrt_llm._torch.utils import helix_local_len, maybe_compile
 from tensorrt_llm._utils import (
     TensorWrapper,
@@ -499,6 +500,21 @@ def _estimate_full_attn_size_per_token(
     )
 
 
+def _get_available_memory_for_auto_host_tier(shares_device_memory: bool) -> float:
+    """Read available memory before allocating either cache tier.
+
+    Args:
+        shares_device_memory: Whether host and device allocations share RAM.
+    """
+    if shares_device_memory:
+        cuda_free, cuda_total = torch.cuda.mem_get_info()
+        return get_integrated_gpu_free_memory(cuda_free, cuda_total)
+    try:
+        return os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_AVPHYS_PAGES")
+    except (ValueError, OSError):
+        return float("inf")
+
+
 def _compute_auto_host_tier_quota(
     quota: int,
     local_ranks: int,
@@ -514,10 +530,10 @@ def _compute_auto_host_tier_quota(
     pinnable-memory (RLIMIT_MEMLOCK) limit.
 
     On an integrated GPU the device pool and the host tier are the same
-    physical memory, so ``mem_available`` counts bytes the device quota is
-    about to consume. Charging the device quota against the budget keeps the
-    two tiers from double-counting; when nothing is left the tier is dropped
-    rather than falling back to the device quota.
+    physical memory, so ``mem_available`` counts bytes the local device pools
+    are about to consume. Charge a device quota for every local rank before
+    sharing the remaining budget between host tiers; when nothing is left the
+    tier is dropped rather than falling back to the device quota.
 
     Args:
         quota: Device (GPU) KV cache quota in bytes; must be positive.
@@ -527,15 +543,15 @@ def _compute_auto_host_tier_quota(
         memlock_limit: RLIMIT_MEMLOCK soft limit in bytes, or
             ``float("inf")`` if unlimited or unknown.
         shares_device_memory: Whether host and device memory are one physical
-            pool (integrated GPU). Charges ``quota`` against ``mem_available``
-            and allows a zero result.
+            pool (integrated GPU). Charges ``quota * local_ranks`` against
+            ``mem_available`` and allows a zero result.
 
     Returns:
         Host cache tier quota in bytes. Positive, except on integrated GPUs
         with no headroom left, where 0 means "provision no host tier".
     """
     if shares_device_memory and mem_available != float("inf"):
-        mem_available = max(0.0, mem_available - quota)
+        mem_available = max(0.0, mem_available - quota * local_ranks)
     candidates = [quota]
     if mem_available != float("inf"):
         candidates.append(int(mem_available / local_ranks * 0.5))
@@ -549,8 +565,9 @@ def _compute_auto_host_tier_quota(
             # OOM-killed.
             logger.warning(
                 "KV cache manager v2 auto host tier sizing has no headroom "
-                "left on this integrated GPU after reserving the "
-                f"{quota / (1 << 30):.2f}GiB device quota; provisioning no "
+                "left on this integrated GPU after reserving "
+                f"{local_ranks} local device quota(s) of "
+                f"{quota / (1 << 30):.2f}GiB each; provisioning no "
                 "host tier. Lower kv_cache_config.max_tokens or "
                 "free_gpu_memory_fraction, or set an explicit "
                 "kv_cache_config.host_cache_size."
@@ -2294,10 +2311,8 @@ def append_to_kv_heads_per_layer(
             # ~170GiB device quota each on a 975GiB node).
             local_ranks = max(1, Distributed.get(mapping).local_world_size)
 
-            try:
-                mem_available = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_AVPHYS_PAGES")
-            except (ValueError, OSError):
-                mem_available = float("inf")
+            shares_device_memory = is_device_integrated()
+            mem_available = _get_available_memory_for_auto_host_tier(shares_device_memory)
             try:
                 _soft, _hard = resource.getrlimit(resource.RLIMIT_MEMLOCK)
                 memlock_limit = _soft if _soft != resource.RLIM_INFINITY else float("inf")
@@ -2308,7 +2323,7 @@ def append_to_kv_heads_per_layer(
                 local_ranks,
                 mem_available,
                 memlock_limit,
-                shares_device_memory=is_device_integrated(),
+                shares_device_memory=shares_device_memory,
             )
             logger.info(
                 f"KV cache manager v2 auto host tier sizing: "

tensorrt_llm/_torch/pyexecutor/kv_cache/memory_utils.py

@@ -0,0 +1,111 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+
+def _read_meminfo_available_bytes(meminfo_path: str = "/proc/meminfo") -> int | None:
+    """Read the kernel's MemAvailable estimate in bytes from a procfs file."""
+    try:
+        for line in Path(meminfo_path).read_text().splitlines():
+            if line.startswith("MemAvailable:"):
+                return max(0, int(line.split()[1]) * 1024)
+    except (OSError, UnicodeError, ValueError, IndexError):
+        return None
+    return None
+
+
+def _read_cgroup_headroom(cgroup: Path) -> int | None:
+    """Read one cgroup v2 limit, crediting only readily reclaimable file cache."""
+    try:
+        limit_text = (cgroup / "memory.max").read_text().strip()
+        if limit_text == "max":
+            return None
+        limit = int(limit_text)
+        current = int((cgroup / "memory.current").read_text())
+    except (OSError, UnicodeError, ValueError):
+        return None
+
+    reclaimable = 0
+    try:
+        stats = dict(line.split() for line in (cgroup / "memory.stat").read_text().splitlines())
+        # Exclude shmem, pages awaiting writeback, and unevictable pages.
+        # These counters can overlap, and unevictable includes anonymous
+        # pages, so subtracting all is deliberately conservative. Missing
+        # counters get no file-cache credit.
+        unavailable = sum(
+            max(0, int(stats[key]))
+            for key in ("shmem", "file_dirty", "file_writeback", "unevictable")
+        )
+        reclaimable = max(0, int(stats["file"]) - unavailable)
+    except (OSError, UnicodeError, ValueError, KeyError):
+        pass
+    return max(0, limit - max(0, current - reclaimable))
+
+
+def _read_cgroup_memory_headroom_bytes(
+    cgroup_dir: str = "/sys/fs/cgroup", proc_cgroup_path: str = "/proc/self/cgroup"
+) -> int | None:
+    """Return the tightest finite headroom from the process cgroup and ancestors.
+
+    Args:
+        cgroup_dir: Root of the cgroup v2 mount.
+        proc_cgroup_path: Procfs file identifying the process's cgroup.
+    """
+    root = Path(cgroup_dir)
+    leaf = root
+    try:
+        for line in Path(proc_cgroup_path).read_text().splitlines():
+            hierarchy, controllers, relative = line.split(":", 2)
+            if hierarchy == "0" and not controllers:
+                parts = Path(relative.lstrip("/")).parts
+                if ".." not in parts:
+                    candidate = root.joinpath(*parts)
+                    if candidate.is_dir():
+                        leaf = candidate
+                break
+    except (OSError, UnicodeError, ValueError):
+        pass
+
+    # A private namespace or a subtree mount can expose the process cgroup at
+    # the mount root, even if the reported host-relative path is not visible.
+    headrooms = []
+    while True:
+        headroom = _read_cgroup_headroom(leaf)
+        if headroom is not None:
+            headrooms.append(headroom)
+        if leaf == root:
+            break
+        leaf = leaf.parent
+    return min(headrooms) if headrooms else None
+
+
+def get_integrated_gpu_free_memory(cuda_free: int, cuda_total: int) -> int:
+    """Estimate allocatable shared memory, including reclaimable host page cache.
+
+    MemAvailable can include file cache that CUDA does not report as free on
+    integrated GPUs. Finite cgroup headroom is a hard upper bound, including
+    when it is below CUDA free or MemAvailable cannot be read.
+
+    Args:
+        cuda_free: Free bytes reported by CUDA; used if MemAvailable is absent.
+        cuda_total: Device capacity in bytes, an upper bound on the estimate.
+    """
+    available = _read_meminfo_available_bytes()
+    available = cuda_free if available is None else max(cuda_free, available)
+    available = min(available, cuda_total)
+    cgroup_headroom = _read_cgroup_memory_headroom_bytes()
+    if cgroup_headroom is not None:
+        available = min(available, cgroup_headroom)
+    return max(0, available)

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -1370,7 +1370,9 @@ def allocation_scope(current_stage: ExecutorMemoryType):
             py_executor.shutdown()
             raise
 
-    _adjust_torch_mem_fraction()
+    _adjust_torch_mem_fraction(
+        use_reclaimable_memory=kv_cache_creator is not None
+        and kv_cache_creator._is_kv_cache_manager_v2)
 
     if mapping.rank == 0:
         logger.info(f"LLM Args:\n{llm_args}")