From b08a6b6e9bf1fe0b178d4aa8275aff744115622e Mon Sep 17 00:00:00 2001 From: jbernloehr Date: Sat, 24 Oct 2026 23:39:48 +0200 Subject: [PATCH] [https://nvbugs/6306663][fix] Account for reclaimable memory in V2 KV cache sizing Use MemAvailable for provisional and profiled V2 KV-cache sizing on integrated GPUs. Preserve the CUDA baseline, explicit user quotas, and V1/discrete behavior. Add CPU regression coverage for accounting, gating, and quota propagation. Signed-off-by: jbernloehr --- tensorrt_llm/_torch/pyexecutor/_util.py | 97 +++++++- .../kv_cache/test_kv_cache_estimation.py | 217 ++++++++++++++++++ 2 files changed, 309 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7593f2ec6f4c..f2164213b718 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 @@ -90,6 +91,70 @@ GB = 1 << 30 + +def _read_meminfo_available_bytes( + meminfo_path: str = "/proc/meminfo") -> Optional[int]: + """Return the kernel's ``MemAvailable`` estimate in bytes, or None.""" + try: + with open(meminfo_path) as f: + for line in f: + if line.startswith("MemAvailable:"): + fields = line.split() + # The value is reported in kB. + return int(fields[1]) * 1024 + except (OSError, ValueError, IndexError): + return None + return None + + +def _read_cgroup_memory_headroom_bytes( + cgroup_dir: str = "/sys/fs/cgroup") -> Optional[int]: + """Return the cgroup v2 memory headroom in bytes, or None if unlimited. + + Page cache charged to the cgroup is reclaimable within it, so it does not + count against the headroom: ``memory.max - (memory.current - file)``. + """ + try: + with open(os.path.join(cgroup_dir, "memory.max")) as f: + limit = f.read().strip() + if limit == "max": + return None + with open(os.path.join(cgroup_dir, "memory.current")) as f: + current = int(f.read().strip()) + file_bytes = 0 + with open(os.path.join(cgroup_dir, "memory.stat")) as f: + for line in f: + key, _, value = line.partition(" ") + if key == "file": + file_bytes = int(value) + break + return max(0, int(limit) - max(0, current - file_bytes)) + except (OSError, ValueError): + return None + + +def get_integrated_gpu_free_memory(cuda_free: int, cuda_total: int) -> int: + """Free device memory on an integrated GPU, counting reclaimable page cache. + + On integrated GPUs (for example GB10 / DGX Spark) device allocations come + from the same physical pool as host memory, and ``cudaMemGetInfo`` reports + only memory that is free right now. Page cache left by container image + pulls and by reading the checkpoint is reclaimed by the kernel when a CUDA + allocation needs it, so it is available to the KV cache even though it is + not reported as free. Use the kernel's ``MemAvailable`` estimate instead, + bounded by the device total and by any cgroup limit, and never below the + CUDA value. Fall back to the CUDA value when the estimate is unavailable. + """ + available = _read_meminfo_available_bytes() + if available is None: + return cuda_free + 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(cuda_free, available) + + # Page size of a standalone drafter's separate KV pool (DFlash, standalone # DSpark). Deliberately not kv_cache_config.tokens_per_block: that describes # the target pool, and target attention backends and deployment configs set it @@ -1307,6 +1372,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: """ @@ -1690,7 +1777,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 @@ -1764,7 +1851,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( @@ -1813,7 +1900,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: diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 7fd42932ef73..07ef7b9609d3 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -32,6 +32,9 @@ _create_kv_cache_manager, _derive_v2_layer_type_attention_windows, _get_num_pool_groups_for_estimation, + _read_cgroup_memory_headroom_bytes, + _read_meminfo_available_bytes, + get_integrated_gpu_free_memory, get_mla_context_workspace_kv_len_cap, ) from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window @@ -1966,3 +1969,217 @@ def test_create_kv_cache_manager_layer_specs_preserve_window_guards(guard: str) assert manager_config is kv_cache_config assert manager_config.max_attention_window == ([1024] if guard == "explicit_window" else None) + + +# --------------------------------------------------------------------------- +# Integrated-GPU (unified memory) free-memory accounting for the V2 manager +# --------------------------------------------------------------------------- + +_UTIL = "tensorrt_llm._torch.pyexecutor._util" +_GiB = 1 << 30 + + +def _write_meminfo(tmp_path, mem_available_kb): + path = tmp_path / "meminfo" + path.write_text( + "MemTotal: 125000000 kB\n" + "MemFree: 4500000 kB\n" + f"MemAvailable: {mem_available_kb} kB\n" + "Cached: 60000000 kB\n" + ) + return str(path) + + +def _write_cgroup(tmp_path, limit, current, file_bytes): + cgroup = tmp_path / "cgroup" + cgroup.mkdir() + (cgroup / "memory.max").write_text(f"{limit}\n") + (cgroup / "memory.current").write_text(f"{current}\n") + (cgroup / "memory.stat").write_text(f"anon 123\nfile {file_bytes}\nkernel 0\n") + return str(cgroup) + + +def test_read_meminfo_available_bytes(tmp_path): + assert _read_meminfo_available_bytes(_write_meminfo(tmp_path, 1000)) == 1000 * 1024 + + +def test_read_meminfo_available_bytes_missing_or_malformed(tmp_path): + assert _read_meminfo_available_bytes(str(tmp_path / "absent")) is None + bad = tmp_path / "bad" + bad.write_text("MemTotal: 1 kB\nMemAvailable: not-a-number kB\n") + assert _read_meminfo_available_bytes(str(bad)) is None + no_field = tmp_path / "no_field" + no_field.write_text("MemTotal: 1 kB\nMemFree: 1 kB\n") + assert _read_meminfo_available_bytes(str(no_field)) is None + + +def test_read_cgroup_headroom_counts_page_cache_as_reclaimable(tmp_path): + cgroup = _write_cgroup(tmp_path, limit=100, current=90, file_bytes=60) + # 100 - (90 - 60): the 60 bytes of page cache are reclaimable in the cgroup. + assert _read_cgroup_memory_headroom_bytes(cgroup) == 70 + + +def test_read_cgroup_headroom_unlimited_or_absent(tmp_path): + cgroup = tmp_path / "cgroup" + cgroup.mkdir() + (cgroup / "memory.max").write_text("max\n") + assert _read_cgroup_memory_headroom_bytes(str(cgroup)) is None + assert _read_cgroup_memory_headroom_bytes(str(tmp_path / "absent")) is None + + +@pytest.mark.parametrize( + "cuda_free,cuda_total,mem_available,cgroup,expected", + [ + # Cold image pull on DGX Spark: CUDA reports MemFree, page cache fills the rest. + (4 * _GiB, 120 * _GiB, 60 * _GiB, None, 60 * _GiB), + # Never report more than the device total. + (4 * _GiB, 120 * _GiB, 200 * _GiB, None, 120 * _GiB), + # A tighter cgroup limit wins. + (4 * _GiB, 120 * _GiB, 60 * _GiB, 20 * _GiB, 20 * _GiB), + # Never report less than CUDA does. + (8 * _GiB, 120 * _GiB, 6 * _GiB, None, 8 * _GiB), + (8 * _GiB, 120 * _GiB, 60 * _GiB, 2 * _GiB, 8 * _GiB), + # No estimate available: unchanged. + (4 * _GiB, 120 * _GiB, None, None, 4 * _GiB), + ], +) +def test_integrated_gpu_free_memory_bounds(cuda_free, cuda_total, mem_available, cgroup, expected): + with ( + patch(f"{_UTIL}._read_meminfo_available_bytes", return_value=mem_available), + patch(f"{_UTIL}._read_cgroup_memory_headroom_bytes", return_value=cgroup), + ): + assert get_integrated_gpu_free_memory(cuda_free, cuda_total) == expected + + +@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, integrated, expected_free): + 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"{_UTIL}._read_meminfo_available_bytes", return_value=60 * _GiB), + patch(f"{_UTIL}._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=None): + kv_cache_config = KvCacheConfig( + free_gpu_memory_fraction=0.9, + max_gpu_total_bytes=max_gpu_total_bytes, + enable_block_reuse=False, + ) + 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, user_cap, expected +): + 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"{_UTIL}._read_meminfo_available_bytes", return_value=60 * _GiB), + patch(f"{_UTIL}._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): + """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, total_gpu_memory, fraction, allocated_bytes): + 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"{_UTIL}._read_meminfo_available_bytes", return_value=60 * _GiB), + patch(f"{_UTIL}._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