diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 77448760f5e9..4fd07a230738 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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)" diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 9101c06bec7e..23a8b21d4138 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/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: " diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/memory_utils.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/memory_utils.py new file mode 100644 index 000000000000..234997da8983 --- /dev/null +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index b9ed5fc73501..774101244735 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/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}") diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_integrated_gpu_memory.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_integrated_gpu_memory.py new file mode 100644 index 000000000000..d652ec9c4962 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_integrated_gpu_memory.py @@ -0,0 +1,256 @@ +# 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. + +import importlib.util +from functools import partial +from pathlib import Path +from unittest.mock import patch + +import pytest + +# This module is pure Python and can also be tested without a built TRT-LLM wheel. +_MODULE_PATH = ( + Path(__file__).resolve().parents[5] / "tensorrt_llm/_torch/pyexecutor/kv_cache/memory_utils.py" +) +_SPEC = importlib.util.spec_from_file_location("kv_memory_utils", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +memory_utils = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(memory_utils) + +pytestmark = pytest.mark.cpu_only +_GiB = 1 << 30 + + +def _write_cgroup( + path: Path, + limit: int | str, + current: int, + file_bytes: int = 0, + shmem: int = 0, + file_dirty: int = 0, + file_writeback: int = 0, + unevictable: int = 0, +) -> None: + """Populate a synthetic cgroup v2 directory.""" + path.mkdir(parents=True, exist_ok=True) + (path / "memory.max").write_text(f"{limit}\n") + (path / "memory.current").write_text(f"{current}\n") + (path / "memory.stat").write_text( + f"anon 123\nfile {file_bytes}\nshmem {shmem}\nkernel 0\n" + f"file_dirty {file_dirty}\nfile_writeback {file_writeback}\nunevictable {unevictable}\n" + ) + + +@pytest.mark.parametrize("value,expected", [("1000", 1024000), ("0", 0), ("-1", 0), ("bad", None)]) +def test_meminfo_available(tmp_path: Path, value: str, expected: int | None) -> None: + """Parse MemAvailable and reject malformed values.""" + path = tmp_path / "meminfo" + path.write_text(f"MemFree: 1 kB\nMemAvailable: {value} kB\n") + assert memory_utils._read_meminfo_available_bytes(str(path)) == expected + + +@pytest.mark.parametrize("content", [None, "MemTotal: 1 kB\n", "MemAvailable:\n"]) +def test_meminfo_unavailable(tmp_path: Path, content: str | None) -> None: + """Missing files or fields return no estimate.""" + path = tmp_path / "meminfo" + if content is not None: + path.write_text(content) + assert memory_utils._read_meminfo_available_bytes(str(path)) is None + + +@pytest.mark.parametrize( + "limit,current,file_bytes,shmem,expected", + [ + (100, 90, 60, 0, 70), + (100, 90, 80, 70, 20), + (100, 90, 30, 40, 10), + (100, 120, 5, 0, 0), + (100, 90, 200, 0, 100), + ("max", 90, 80, 0, None), + ], +) +def test_cgroup_reclaimable_memory( + tmp_path: Path, + limit: int | str, + current: int, + file_bytes: int, + shmem: int, + expected: int | None, +) -> None: + """Shared memory never increases the reclaimable-file estimate.""" + _write_cgroup(tmp_path, limit, current, file_bytes, shmem) + proc = tmp_path / "proc_cgroup" + proc.write_text("0::/\n") + assert memory_utils._read_cgroup_memory_headroom_bytes(str(tmp_path), str(proc)) == expected + + +@pytest.mark.parametrize("stat", [None, "file 80\n", "file 80\nshmem bad\n", "malformed\n"]) +def test_cgroup_missing_stats_preserve_limit(tmp_path: Path, stat: str | None) -> None: + """Unreadable reclaimable stats cannot discard a finite memory limit.""" + _write_cgroup(tmp_path, 100, 90) + path = tmp_path / "memory.stat" + path.unlink() + if stat is not None: + path.write_text(stat) + assert ( + memory_utils._read_cgroup_memory_headroom_bytes(str(tmp_path), str(tmp_path / "absent")) + == 10 + ) + + +@pytest.mark.parametrize( + "shmem,file_dirty,file_writeback,unevictable,expected", + [ + (0, 60, 0, 0, 30), + (0, 0, 60, 0, 30), + (0, 0, 0, 60, 30), + (0, 30, 20, 10, 30), + (10, 20, 15, 5, 40), + (20, 40, 30, 20, 10), + ], +) +def test_cgroup_non_reclaimable_file_pages( + tmp_path: Path, + shmem: int, + file_dirty: int, + file_writeback: int, + unevictable: int, + expected: int, +) -> None: + """Exclude pages requiring writeback or unavailable to reclaim. + + The counters can overlap, and unevictable can include anonymous pages. + Conservatively subtracting them may exhaust file credit, never free bytes. + """ + _write_cgroup(tmp_path, 100, 90, 80, shmem, file_dirty, file_writeback, unevictable) + assert memory_utils._read_cgroup_headroom(tmp_path) == expected + + +@pytest.mark.parametrize("counter", ["file_dirty", "file_writeback", "unevictable"]) +@pytest.mark.parametrize("value", [None, "bad"]) +def test_cgroup_incomplete_file_stats_get_no_credit( + tmp_path: Path, counter: str, value: str | None +) -> None: + """Unknown non-reclaimable bytes cannot grant credit or remove the hard cap.""" + _write_cgroup(tmp_path, 100, 90, 80) + path = tmp_path / "memory.stat" + replacement = "" if value is None else f"{counter} {value}\n" + path.write_text(path.read_text().replace(f"{counter} 0\n", replacement)) + assert memory_utils._read_cgroup_headroom(tmp_path) == 10 + + +def test_integrated_gpu_excludes_unavailable_cgroup_cache(tmp_path: Path) -> None: + """Real procfs/cgroup parsing must cap even a larger CUDA-free estimate.""" + root = tmp_path / "cgroup" + _write_cgroup(root, 100 * _GiB, 90 * _GiB, 80 * _GiB, 10 * _GiB, 20 * _GiB, 15 * _GiB, 5 * _GiB) + proc = tmp_path / "proc_cgroup" + proc.write_text("0::/\n") + meminfo = tmp_path / "meminfo" + meminfo.write_text(f"MemAvailable: {80 * _GiB // 1024} kB\n") + # Redirect only file locations; execute the real parsers and budget logic. + with ( + patch.object( + memory_utils, + "_read_meminfo_available_bytes", + side_effect=partial(memory_utils._read_meminfo_available_bytes, str(meminfo)), + ), + patch.object( + memory_utils, + "_read_cgroup_memory_headroom_bytes", + side_effect=partial( + memory_utils._read_cgroup_memory_headroom_bytes, str(root), str(proc) + ), + ), + ): + # Only 30 GiB of the 80 GiB file counter is credited. Together with + # 10 GiB of unused cgroup allowance, the allocatable budget is 40 GiB. + assert memory_utils.get_integrated_gpu_free_memory(70 * _GiB, 120 * _GiB) == 40 * _GiB + + +@pytest.mark.parametrize( + "limits,expected", [(("max", "max", 40), 20), ((200, 50, 100), 30), ((30, 100, 200), 10)] +) +def test_cgroup_process_and_ancestors( + tmp_path: Path, limits: tuple[int | str, ...], expected: int +) -> None: + """The tightest headroom can be at the process, parent, or mount root.""" + root = tmp_path / "cgroup" + for path, limit in zip([root, root / "user.slice", root / "user.slice/session.scope"], limits): + _write_cgroup(path, limit, 30, 15, 5) + proc = tmp_path / "proc_cgroup" + proc.write_text("0::/user.slice/session.scope\n") + assert memory_utils._read_cgroup_memory_headroom_bytes(str(root), str(proc)) == expected + + +@pytest.mark.parametrize("relative", ["/", "/host/container", "/../../outside"]) +def test_cgroup_namespaced_root(tmp_path: Path, relative: str) -> None: + """A private namespace still enforces the limit exposed at its mount root.""" + root = tmp_path / "cgroup" + _write_cgroup(root, 100, 90, 80, 70) + proc = tmp_path / "proc_cgroup" + proc.write_text(f"0::{relative}\n") + assert memory_utils._read_cgroup_memory_headroom_bytes(str(root), str(proc)) == 20 + + +def test_cgroup_parent_with_sibling_usage(tmp_path: Path) -> None: + """Each ancestor uses its own usage and cache, including sibling charges.""" + root = tmp_path / "cgroup" + _write_cgroup(root, 100, 95, 10, 5) + _write_cgroup(root / "child", 50, 10, 0, 0) + proc = tmp_path / "proc_cgroup" + proc.write_text("0::/child\n") + assert memory_utils._read_cgroup_memory_headroom_bytes(str(root), str(proc)) == 10 + + +def test_cgroup_unavailable(tmp_path: Path) -> None: + """No readable cgroup limit leaves the host estimate uncapped.""" + assert ( + memory_utils._read_cgroup_memory_headroom_bytes(str(tmp_path), str(tmp_path / "absent")) + is None + ) + + +@pytest.mark.parametrize( + "cuda_free,cuda_total,available,cgroup,expected", + [ + (4, 120, 60, None, 60), + (4, 120, 200, None, 120), + (4, 120, 60, 20, 20), + (8, 120, 6, None, 8), + (8, 120, 60, 2, 2), + (8, 120, 60, 0, 0), + (4, 120, None, None, 4), + (4, 120, None, 2, 2), + ], +) +def test_integrated_gpu_bounds( + cuda_free: int, cuda_total: int, available: int | None, cgroup: int | None, expected: int +) -> None: + """Device capacity and finite cgroup headroom cap even the CUDA fallback.""" + with ( + patch.object( + memory_utils, + "_read_meminfo_available_bytes", + return_value=None if available is None else available * _GiB, + ), + patch.object( + memory_utils, + "_read_cgroup_memory_headroom_bytes", + return_value=None if cgroup is None else cgroup * _GiB, + ), + ): + assert ( + memory_utils.get_integrated_gpu_free_memory(cuda_free * _GiB, cuda_total * _GiB) + == expected * _GiB + ) diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py index 5c06f0766e20..dac58d67137a 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py @@ -26,6 +26,7 @@ from tensorrt_llm._torch.pyexecutor._util import ( CacheCost, KvCacheCreator, + _adjust_torch_mem_fraction, _create_kv_cache_manager, _derive_v2_layer_type_attention_windows, _get_num_pool_groups_for_estimation, @@ -37,6 +38,8 @@ KVCacheLayerSpec, KVCacheManagerV2, Role, + _compute_auto_host_tier_quota, + _get_available_memory_for_auto_host_tier, _get_static_cache_size_layer_components, ) from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine @@ -2063,6 +2066,7 @@ def test_create_kv_cache_manager_layer_specs_preserve_window_guards(guard: str) def test_profiling_failure_releases_resources_and_restores_response_state(): creator = object.__new__(KvCacheCreator) + creator._is_kv_cache_manager_v2 = False creator._mapping = SimpleNamespace(cp_config={}) creator._kv_cache_config = SimpleNamespace(free_gpu_memory_fraction=0.9) creator._skip_est = False @@ -2096,3 +2100,253 @@ def test_profiling_failure_releases_resources_and_restores_response_state(): warmup_scope.__exit__.assert_called_once() assert executor.enable_iter_perf_stats is True assert executor.set_gather_responses.call_args.args == (False,) + + +# --------------------------------------------------------------------------- +# Integrated-GPU (unified memory) free-memory accounting for the V2 manager +# --------------------------------------------------------------------------- + +_UTIL = "tensorrt_llm._torch.pyexecutor._util" +_MEMORY = "tensorrt_llm._torch.pyexecutor.kv_cache.memory_utils" +_MANAGER = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" +_GiB = 1 << 30 + + +@pytest.mark.parametrize( + "is_v2,integrated,expected_free", + [ + (True, True, 60 * _GiB), # the only case that changes + (True, False, 4 * _GiB), # discrete GPU: CUDA value + (False, True, 4 * _GiB), # V1 manager: CUDA value + (False, False, 4 * _GiB), + ], +) +def test_device_free_memory_is_gated_on_v2_and_integrated( + is_v2: bool, integrated: bool, expected_free: int +) -> None: + """Only V2 integrated GPUs use reclaimable memory for KV capacity.""" + creator = object.__new__(KvCacheCreator) + creator._is_kv_cache_manager_v2 = is_v2 + with ( + patch("torch.cuda.mem_get_info", return_value=(4 * _GiB, 120 * _GiB)), + patch(f"{_UTIL}.is_device_integrated", return_value=integrated), + patch(f"{_MEMORY}._read_meminfo_available_bytes", return_value=60 * _GiB), + patch(f"{_MEMORY}._read_cgroup_memory_headroom_bytes", return_value=None), + ): + assert creator._get_device_free_memory() == (expected_free, 120 * _GiB) + + +def _make_v2_estimation_creator( + max_gpu_total_bytes: int | None = None, +) -> tuple[KvCacheCreator, KvCacheConfig]: + """Construct a minimal V2 creator for memory budgeting tests.""" + quota_kwargs = {} + if max_gpu_total_bytes is not None: + quota_kwargs["max_gpu_total_bytes"] = max_gpu_total_bytes + kv_cache_config = KvCacheConfig( + free_gpu_memory_fraction=0.9, + enable_block_reuse=False, + **quota_kwargs, + ) + model_engine = Mock() + model_engine.model.model_config.attn_backend = "TRTLLM" + model_engine.model.model_config.is_encoder_decoder = False + model_engine.mm_encoder_output_budget_bytes = None + model_engine.mm_encoder_cache = None + llm_args = Mock(cache_transceiver_config=None, enable_chunked_prefill=True) + with patch.object( + KvCacheCreator, "_get_model_kv_cache_manager_cls", return_value=KVCacheManagerV2 + ): + creator = KvCacheCreator( + model_engine=model_engine, + mapping=Mock(cp_config={}), + net_max_seq_len=3000, + kv_connector_manager=None, + max_num_tokens=8192, + max_beam_width=1, + tokens_per_block=32, + max_seq_len=3000, + max_batch_size=512, + kv_cache_config=kv_cache_config, + llm_args=llm_args, + speculative_config=None, + sparse_attention_config=None, + profiling_stage_data=None, + is_disagg=False, + ) + return creator, kv_cache_config + + +@pytest.mark.parametrize( + "integrated,user_cap,expected", + [ + # Integrated: 0.9 x MemAvailable (60 GiB) instead of 0.9 x MemFree (4 GiB). + (True, None, int(0.9 * 60 * _GiB)), + # An explicit max_gpu_total_bytes still caps the provisional quota. + (True, 10 * _GiB, 10 * _GiB), + # Discrete GPU: unchanged. + (False, None, int(0.9 * 4 * _GiB)), + (False, 1 * _GiB, 1 * _GiB), + ], +) +def test_v2_provisional_quota_counts_reclaimable_memory_on_integrated_gpus( + integrated: bool, user_cap: int | None, expected: int +) -> None: + """Reclaimable memory increases capacity without bypassing user caps.""" + creator, kv_cache_config = _make_v2_estimation_creator(max_gpu_total_bytes=user_cap) + with ( + patch.object(creator, "_get_token_num_for_estimation", return_value=16384), + patch("torch.cuda.mem_get_info", return_value=(4 * _GiB, 120 * _GiB)), + patch(f"{_UTIL}.is_device_integrated", return_value=integrated), + patch(f"{_MEMORY}._read_meminfo_available_bytes", return_value=60 * _GiB), + patch(f"{_MEMORY}._read_cgroup_memory_headroom_bytes", return_value=None), + ): + assert creator.try_prepare_estimation() + assert kv_cache_config.max_gpu_total_bytes == expected + + +@pytest.mark.parametrize("integrated", [False, True]) +def test_v2_final_budget_excludes_reclaimable_memory_from_peak(integrated: bool) -> None: + """The final V2 budget subtracts only memory the process actually holds. + + The profiling measurements use the same free value, so on an integrated GPU + reclaimable page cache no longer shows up as non-torch peak usage. The + budget formula itself is unchanged. + """ + creator, kv_cache_config = _make_v2_estimation_creator() + py_executor = Mock() + py_executor.warmup_scope.return_value = nullcontext() + py_executor.dist.mapping.rank = 0 + py_executor.dist.broadcast.side_effect = lambda value, root: value + py_executor.enqueue_requests.return_value = [1] + py_executor.await_responses.return_value = [] + kv_manager = Mock() + kv_manager.get_kv_cache_stats.return_value = SimpleNamespace(allocated_bytes=10 * _GiB) + py_executor.resource_manager.resource_managers.get.side_effect = [kv_manager, None] + captured = {} + + def _capture( + peak_memory: int, total_gpu_memory: int, fraction: float, allocated_bytes: int + ) -> int: + """Capture the final budgeting inputs without allocating a cache.""" + captured.update(peak=peak_memory, total=total_gpu_memory, allocated=allocated_bytes) + return 1 * _GiB + + torch_bytes = {"allocated_bytes.all.current": 60 * _GiB, "allocated_bytes.all.peak": 61 * _GiB} + with ( + patch.object(creator, "_get_token_num_for_estimation", return_value=16384), + # CUDA: 4 GiB free (MemFree); kernel: 60 GiB available incl. page cache. + patch("torch.cuda.mem_get_info", return_value=(4 * _GiB, 120 * _GiB)), + patch(f"{_UTIL}.is_device_integrated", return_value=integrated), + patch(f"{_MEMORY}._read_meminfo_available_bytes", return_value=60 * _GiB), + patch(f"{_MEMORY}._read_cgroup_memory_headroom_bytes", return_value=None), + patch("torch.cuda.memory_stats", return_value=torch_bytes), + patch("torch.cuda.empty_cache"), + patch("torch.cuda.reset_peak_memory_stats"), + patch.object(creator, "_encode_dummy_inputs", return_value=None), + patch.object(creator, "_cal_max_memory", side_effect=_capture), + patch.object(creator, "_get_kv_size_per_token", return_value=CacheCost(slope=1)), + patch(f"{_UTIL}.get_attention_workspace_bytes_per_token", return_value=0), + patch(f"{_UTIL}.get_attention_workspace_fixed_bytes", return_value=0), + ): + assert creator.try_prepare_estimation() + creator.configure_kv_cache_capacity(py_executor) + + used_outside_torch = (120 - (60 if integrated else 4)) * _GiB - 60 * _GiB + assert captured["total"] == 120 * _GiB + assert captured["allocated"] == 10 * _GiB + assert captured["peak"] == 61 * _GiB + max(used_outside_torch, 0) + assert kv_cache_config.max_gpu_total_bytes == 1 * _GiB + + +@pytest.mark.parametrize( + "is_v2,integrated", [(True, True), (True, False), (False, True), (False, False)] +) +@pytest.mark.parametrize("available,cgroup", [(20, None), (20, 5), (20, 0), (None, 1), (200, None)]) +def test_torch_mem_fraction_leaves_headroom_for_lazy_allocations( + is_v2: bool, integrated: bool, available: int | None, cgroup: int | None +) -> None: + """The allocator can consume reclaimable memory after the KV pool is built.""" + reserved, raw_free, total = 64 * _GiB, 2 * _GiB, 128 * _GiB + with ( + patch("torch.cuda.set_per_process_memory_fraction") as set_fraction, + patch("torch.cuda.get_allocator_backend", return_value="cudaMallocAsync"), + patch("torch.cuda.memory_reserved", return_value=reserved), + patch("torch.cuda.mem_get_info", return_value=(raw_free, total)), + patch(f"{_UTIL}.is_device_integrated", return_value=integrated), + patch( + f"{_MEMORY}._read_meminfo_available_bytes", + return_value=None if available is None else available * _GiB, + ), + patch( + f"{_MEMORY}._read_cgroup_memory_headroom_bytes", + return_value=None if cgroup is None else cgroup * _GiB, + ), + ): + _adjust_torch_mem_fraction(use_reclaimable_memory=is_v2) + free = raw_free + if is_v2 and integrated: + free = raw_free if available is None else min(total, max(raw_free, available * _GiB)) + if cgroup is not None: + free = min(free, cgroup * _GiB) + expected = min(total, max(0, reserved + free - 32 * 1024**2)) / total + assert set_fraction.call_args_list[0].args == (1.0,) + assert set_fraction.call_args_list[1].args == (expected,) + assert 0 <= expected <= 1 + if is_v2 and integrated and available == 20 and cgroup is None: + # A lazy 4 GiB workspace would exceed the old raw-free cap but fits now. + assert reserved + 4 * _GiB > reserved + raw_free - 32 * 1024**2 + assert reserved + 4 * _GiB < expected * total + + +@pytest.mark.parametrize( + "integrated,available,cgroup,quota,local_ranks,memlock,expected_gib", + [ + (True, 90, None, 81, 1, float("inf"), 4.5), + (True, 90, 85, 81, 1, float("inf"), 2), + (True, 90, 10, 81, 1, float("inf"), 0), + (True, None, None, 27, 1, float("inf"), 1.5), + (True, 90, None, 81, 1, 5 * _GiB, 4), + (False, 90, None, 81, 1, float("inf"), 15), + (True, 90, None, 40, 2, float("inf"), 2.5), + (True, 90, None, 45, 2, float("inf"), 0), + (True, 90, None, 20, 4, float("inf"), 1.25), + (True, 90, None, 40, 2, 1.25 * _GiB, 1), + (False, 90, None, 40, 2, float("inf"), 7.5), + ], +) +def test_auto_host_tier_uses_shared_memory_budget( + integrated: bool, + available: int | None, + cgroup: int | None, + quota: int, + local_ranks: int, + memlock: float, + expected_gib: float, +) -> None: + """Host sizing accounts for every local device pool in the shared budget.""" + with ( + patch("torch.cuda.mem_get_info", return_value=(30 * _GiB, 120 * _GiB)), + patch(f"{_MANAGER}.os.sysconf", side_effect=[4096, 30 * _GiB // 4096]), + patch( + f"{_MEMORY}._read_meminfo_available_bytes", + return_value=None if available is None else available * _GiB, + ), + patch( + f"{_MEMORY}._read_cgroup_memory_headroom_bytes", + return_value=None if cgroup is None else cgroup * _GiB, + ), + patch(f"{_MANAGER}.logger.warning") as warn, + ): + available_bytes = _get_available_memory_for_auto_host_tier(integrated) + host_quota = _compute_auto_host_tier_quota( + quota * _GiB, local_ranks, available_bytes, memlock, shares_device_memory=integrated + ) + assert host_quota == expected_gib * _GiB + assert warn.called == (expected_gib == 0) + if integrated: + device_bytes = quota * _GiB * local_ranks + if device_bytes <= available_bytes: + assert device_bytes + host_quota * local_ranks <= available_bytes + else: + assert host_quota == 0 diff --git a/tests/unittest/_torch/pyexecutor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/pyexecutor/test_py_executor_creator_mla_cache_reuse_sync.py index 4487ec7d356a..996b598d5fdf 100644 --- a/tests/unittest/_torch/pyexecutor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/pyexecutor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -16,6 +16,7 @@ import sys from contextlib import contextmanager from types import SimpleNamespace +from unittest.mock import Mock import pytest @@ -128,6 +129,8 @@ def start_worker(self): class _DummyKvCacheCreator: """Mock KV cache creator that builds a dummy KV cache manager with reuse settings.""" + _is_kv_cache_manager_v2 = False + def __init__(self, **kwargs): """Initialize with KV cache configuration. @@ -335,7 +338,10 @@ def _run_create_py_executor( py_executor_creator, "get_spec_resource_manager", lambda *args, **kwargs: None ) monkeypatch.setattr(py_executor_creator, "get_spec_drafter", lambda *args, **kwargs: None) - monkeypatch.setattr(py_executor_creator, "_adjust_torch_mem_fraction", lambda: None) + adjust_torch_mem_fraction = Mock() + monkeypatch.setattr( + py_executor_creator, "_adjust_torch_mem_fraction", adjust_torch_mem_fraction + ) monkeypatch.setattr(py_executor_creator, "log_memory_usage", lambda *args, **kwargs: None) monkeypatch.setattr(py_executor_creator, "is_mla", lambda _: is_mla_model) monkeypatch.setattr( @@ -399,6 +405,9 @@ def _create_py_executor_resources(**kwargs): checkpoint_dir=None, ) assert created_executors == [py_executor] + adjust_torch_mem_fraction.assert_called_once_with( + use_reclaimable_memory=kv_cache_creator_cls._is_kv_cache_manager_v2 + ) kv_cache_manager = py_executor.resource_manager.get_resource_manager( ResourceManagerType.KV_CACHE_MANAGER ) @@ -410,6 +419,19 @@ def _create_py_executor_resources(**kwargs): ) +@pytest.mark.parametrize("is_v2", [False, True], ids=["v1", "v2"]) +def test_allocator_option_follows_resolved_kv_cache_manager( + monkeypatch: pytest.MonkeyPatch, is_v2: bool +) -> None: + """Executor creation forwards the resolved manager's allocator option.""" + monkeypatch.setattr(_DummyKvCacheCreator, "_is_kv_cache_manager_v2", is_v2) + _run_create_py_executor( + monkeypatch, + sm_version=90, + kv_cache_quant_algo=QuantAlgo.NO_QUANT, + ) + + @pytest.mark.parametrize("requires_encoder_features", [False, True]) @pytest.mark.parametrize("requested_reuse", [False, True]) @pytest.mark.parametrize("wrapped_model", [False, True])