From f5a3fa8c3acae2fc18925785feacc8e01a87ad07 Mon Sep 17 00:00:00 2001 From: Alberto Massidda Date: Tue, 3 Nov 2026 11:59:57 +0100 Subject: [PATCH 1/2] [Bugfix] Size the extensible KV cache from UMA-corrected free memory On integrated GPUs (DGX Spark, GH200, Jetson) cudaMemGetInfo does not count reclaimable OS memory (page cache, buffers) as free, although a device allocation reclaims it. MemorySnapshot already corrects for this with psutil, but the extensible KV cache read torch.accelerator.get_memory_info directly in committable_blocks, measure_kv_cache_blocks and the weight-update release check, so after the weights had been read through the page cache (e.g. safetensors_load_strategy=prefetch) the measured headroom was a fraction of the real one. On a 119.7 GiB DGX Spark with a 27B NVFP4 model this sized the cache to 216 blocks instead of the 71 GiB the classic path finds, and to zero blocks in disaggregated serving, where the larger transient peak exhausted the headroom term. Commit via cuMemCreate/cuMemMap on the same part reclaimed 106 GiB with 59 GiB reported free, so the psutil figure is the right one there too. Move the correction into a shared device_memory_info() helper in mem_utils and use it from MemorySnapshot.measure and the three extensible KV cache call sites. Fixes #65378 Signed-off-by: Alberto Massidda Co-Authored-By: Claude Fable 5.1 --- tests/utils_/test_mem_utils.py | 43 +++++++++++++- tests/v1/worker/test_extensible_kv_cache.py | 65 +++++++++++++++++++++ vllm/utils/mem_utils.py | 34 ++++++++--- vllm/v1/worker/extensible_kv_cache.py | 8 +-- 4 files changed, 136 insertions(+), 14 deletions(-) diff --git a/tests/utils_/test_mem_utils.py b/tests/utils_/test_mem_utils.py index 19690d1f34f8..3cf7858bc260 100644 --- a/tests/utils_/test_mem_utils.py +++ b/tests/utils_/test_mem_utils.py @@ -5,7 +5,11 @@ import torch from vllm_test_utils.monitor import monitor -from vllm.utils.mem_utils import MemorySnapshot, memory_profiling +from vllm.utils.mem_utils import ( + MemorySnapshot, + device_memory_info, + memory_profiling, +) from ..utils import create_new_process_for_each_test @@ -148,3 +152,40 @@ def test_memory_snapshot_uses_cuda_on_discrete_gpu(): assert snapshot.free_memory == mock_cuda_free assert snapshot.total_memory == mock_cuda_total mock_psutil.virtual_memory.assert_not_called() + + +def test_device_memory_info_uses_psutil_on_integrated_gpu(): + """The shared helper applies the UMA correction, so every KV cache sizing + path (profiling and the extensible cache) sees the same free memory.""" + mock_cuda_free = 40 * 1024**3 + mock_cuda_total = 120 * 1024**3 + mock_psutil_available = 100 * 1024**3 + + with ( + patch("vllm.utils.mem_utils.current_platform") as mock_platform, + patch("vllm.utils.mem_utils.psutil") as mock_psutil, + patch("torch.accelerator") as mock_accelerator, + ): + mock_accelerator.get_memory_info.return_value = ( + mock_cuda_free, + mock_cuda_total, + ) + mock_platform.is_integrated_gpu.return_value = True + mock_vmem = MagicMock() + mock_vmem.available = mock_psutil_available + mock_psutil.virtual_memory.return_value = mock_vmem + + assert device_memory_info("cuda:0") == (mock_psutil_available, mock_cuda_total) + mock_platform.is_integrated_gpu.assert_called_once_with(0) + + mock_platform.is_integrated_gpu.return_value = False + assert device_memory_info("cuda:0") == (mock_cuda_free, mock_cuda_total) + + # A CPU device is never an integrated GPU. + mock_platform.is_integrated_gpu.return_value = True + mock_platform.is_integrated_gpu.reset_mock() + assert device_memory_info(torch.device("cpu")) == ( + mock_cuda_free, + mock_cuda_total, + ) + mock_platform.is_integrated_gpu.assert_not_called() diff --git a/tests/v1/worker/test_extensible_kv_cache.py b/tests/v1/worker/test_extensible_kv_cache.py index 35004b5bd4f8..b7892bed78d3 100644 --- a/tests/v1/worker/test_extensible_kv_cache.py +++ b/tests/v1/worker/test_extensible_kv_cache.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy +import functools from types import SimpleNamespace import pytest @@ -226,6 +227,70 @@ def test_measured_sizing_caps_torch_memory(monkeypatch): assert limits == [2 * gib + left_free] +def test_measured_sizing_counts_reclaimable_memory_on_integrated_gpu(monkeypatch): + """On UMA parts the page cache (e.g. prefetched weights) is reclaimable + but cudaMemGetInfo reports it as used; sizing from the raw figure left + a 120 GiB DGX Spark with zero KV cache blocks (vllm-project/vllm#65378).""" + gib = 1 << 30 + kv_cache = SimpleNamespace( + reserved_headroom_bytes=7 * gib, + physical_bytes=gib, + bytes_per_block=gib, + blocks_within=lambda num_bytes: num_bytes // gib, + physical_bytes_for=lambda num_blocks: num_blocks * gib, + set_torch_memory_limit=lambda limit: None, + ) + accelerator = extensible_kv_cache.torch.accelerator + monkeypatch.setattr(accelerator, "synchronize", lambda *args: None) + monkeypatch.setattr(accelerator, "empty_cache", lambda: None) + monkeypatch.setattr(accelerator, "memory_reserved", lambda *args: 20 * gib) + # cudaMemGetInfo: 10 GiB free of 120; the OS has 100 GiB available. + monkeypatch.setattr( + accelerator, "get_memory_info", lambda *args: (10 * gib, 120 * gib) + ) + monkeypatch.setattr(accelerator, "current_device_index", lambda: 0) + from vllm.utils import mem_utils + + monkeypatch.setattr( + mem_utils.current_platform, "is_integrated_gpu", lambda *args: True + ) + monkeypatch.setattr( + mem_utils.psutil, + "virtual_memory", + lambda: SimpleNamespace(available=100 * gib), + ) + monkeypatch.setattr( + extensible_kv_cache, + "reserve_mm_ipc_gpu_memory", + lambda available, *args: available, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(multimodal_config=None), + parallel_config=SimpleNamespace(), + kv_transfer_config=None, + ) + measure = functools.partial( + extensible_kv_cache.measure_kv_cache_blocks, + kv_cache, # type: ignore[arg-type] + vllm_config, # type: ignore[arg-type] + torch.device("cuda", 0), + init_free_memory=110 * gib, + requested_memory=96 * gib, + ) + + num_blocks = measure() + + # headroom 101 GiB, margin 0.75 * 7, non-KV use 110 - 100 - 1 = 9 GiB: + # min(96 - 9, 101 - 5.25) - 7 = 80 GiB. + assert num_blocks == 80 + + # The raw figure would have sized it to nothing. + monkeypatch.setattr( + mem_utils.current_platform, "is_integrated_gpu", lambda *args: False + ) + assert measure() == 0 + + def test_freeing_the_cache_lifts_the_torch_memory_cap(monkeypatch): """The process may start another engine, which must not inherit the cap.""" fractions: list[float] = [] diff --git a/vllm/utils/mem_utils.py b/vllm/utils/mem_utils.py index 1b449058b27d..f5233620736d 100644 --- a/vllm/utils/mem_utils.py +++ b/vllm/utils/mem_utils.py @@ -101,6 +101,30 @@ def __exit__(self, exc_type, exc_val, exc_tb): gc.collect() +def device_memory_info(device: torch.types.Device = None) -> tuple[int, int]: + """Free and total device memory in bytes, like + `torch.accelerator.get_memory_info`, with the free memory corrected on + integrated GPUs. + + On UMA (Unified Memory Architecture) platforms where CPU and GPU share + physical memory (e.g. GH200, DGX Spark, Jetson Orin), cudaMemGetInfo + underreports free memory because it does not account for reclaimable OS + memory (page cache, buffers), which a device allocation does reclaim. + psutil reports the true available memory there. + https://docs.nvidia.com/cuda/cuda-for-tegra-appnote/#estimating-total-allocatable-device-memory-on-an-integrated-gpu-device + """ + free_memory, total_memory = torch.accelerator.get_memory_info(device) + torch_device = torch.device(device) if device is not None else None + if torch_device is not None and torch_device.type == "cpu": + return free_memory, total_memory + index = torch_device.index if torch_device is not None else None + if index is None: + index = torch.accelerator.current_device_index() + if current_platform.is_integrated_gpu(index): + free_memory = psutil.virtual_memory().available + return free_memory, total_memory + + @dataclass class MemorySnapshot: """Memory snapshot.""" @@ -140,15 +164,7 @@ def measure(self) -> None: self.torch_peak = stats.get("allocated_bytes.all.peak", 0) self.torch_allocated = stats.get("allocated_bytes.all.current", 0) - self.free_memory, self.total_memory = torch.accelerator.get_memory_info(device) - if current_platform.is_integrated_gpu(device.index): - # On UMA (Unified Memory Architecture) platforms where CPU and - # GPU share physical memory (e.g. GH200, DGX Spark, Jetson Orin), - # cudaMemGetInfo underreports free memory because it does not - # account for reclaimable OS memory (page cache, buffers). - # Use psutil to get the true available memory. - # https://docs.nvidia.com/cuda/cuda-for-tegra-appnote/#estimating-total-allocatable-device-memory-on-an-integrated-gpu-device - self.free_memory = psutil.virtual_memory().available + self.free_memory, self.total_memory = device_memory_info(device) self.cuda_memory = self.total_memory - self.free_memory diff --git a/vllm/v1/worker/extensible_kv_cache.py b/vllm/v1/worker/extensible_kv_cache.py index 6b2cd75bc7de..30f240a4d0eb 100644 --- a/vllm/v1/worker/extensible_kv_cache.py +++ b/vllm/v1/worker/extensible_kv_cache.py @@ -24,7 +24,7 @@ from vllm.platforms import current_platform from vllm.utils.extensible_tensor import ExtensibleTensor from vllm.utils.math_utils import round_up -from vllm.utils.mem_utils import format_gib +from vllm.utils.mem_utils import device_memory_info, format_gib from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheTensor, @@ -246,7 +246,7 @@ def committable_blocks(self) -> int: a share of the headroom, never below what is already committed nor above the capacity. """ - free_memory, _ = torch.accelerator.get_memory_info(self.buffer.device) + free_memory, _ = device_memory_info(self.buffer.device) headroom = free_memory + self.physical_bytes reserve = max( self.reserved_headroom_bytes, @@ -548,7 +548,7 @@ def measure_kv_cache_blocks( torch.accelerator.synchronize() gc.collect() torch.accelerator.empty_cache() - free_memory, _ = torch.accelerator.get_memory_info(device) + free_memory, _ = device_memory_info(device) transient_peak = kv_cache.reserved_headroom_bytes available = measure_kv_cache_bytes( init_free_memory=init_free_memory, @@ -649,7 +649,7 @@ def _release_for_weight_update( layer_bytes = max_layer_reload_bytes(model) needed = WEIGHT_UPDATE_PEAK_LAYER_FACTOR * layer_bytes free = ( - torch.accelerator.get_memory_info()[0] + device_memory_info()[0] + torch.accelerator.memory_reserved() - torch.accelerator.memory_allocated() ) From 5cc15d8740d9abfa0d5f0e56f60d5d7610d6b561 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 3 Nov 2026 09:08:42 -0800 Subject: [PATCH 2/2] minor simplification Signed-off-by: Nick Hill --- tests/utils_/test_mem_utils.py | 9 --------- tests/v1/worker/test_extensible_kv_cache.py | 1 - vllm/platforms/cuda.py | 2 +- vllm/platforms/interface.py | 4 ++-- vllm/utils/mem_utils.py | 7 +------ 5 files changed, 4 insertions(+), 19 deletions(-) diff --git a/tests/utils_/test_mem_utils.py b/tests/utils_/test_mem_utils.py index 3cf7858bc260..78ccba8ad895 100644 --- a/tests/utils_/test_mem_utils.py +++ b/tests/utils_/test_mem_utils.py @@ -180,12 +180,3 @@ def test_device_memory_info_uses_psutil_on_integrated_gpu(): mock_platform.is_integrated_gpu.return_value = False assert device_memory_info("cuda:0") == (mock_cuda_free, mock_cuda_total) - - # A CPU device is never an integrated GPU. - mock_platform.is_integrated_gpu.return_value = True - mock_platform.is_integrated_gpu.reset_mock() - assert device_memory_info(torch.device("cpu")) == ( - mock_cuda_free, - mock_cuda_total, - ) - mock_platform.is_integrated_gpu.assert_not_called() diff --git a/tests/v1/worker/test_extensible_kv_cache.py b/tests/v1/worker/test_extensible_kv_cache.py index b7892bed78d3..8f8e2bf25211 100644 --- a/tests/v1/worker/test_extensible_kv_cache.py +++ b/tests/v1/worker/test_extensible_kv_cache.py @@ -248,7 +248,6 @@ def test_measured_sizing_counts_reclaimable_memory_on_integrated_gpu(monkeypatch monkeypatch.setattr( accelerator, "get_memory_info", lambda *args: (10 * gib, 120 * gib) ) - monkeypatch.setattr(accelerator, "current_device_index", lambda: 0) from vllm.utils import mem_utils monkeypatch.setattr( diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 87330e6d5b11..4c73e3283d87 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -701,7 +701,7 @@ def support_deep_gemm(cls) -> bool: ) @classmethod - def is_integrated_gpu(cls, device_id: int = 0) -> bool: + def is_integrated_gpu(cls, device_id: int | None = 0) -> bool: return bool(torch.cuda.get_device_properties(device_id).is_integrated) @classmethod diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 533f23d13eb0..61fc98c7a993 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -1229,9 +1229,9 @@ def get_device_communicator_cls(cls) -> str: return "vllm.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa @classmethod - def is_integrated_gpu(cls, device_id: int = 0) -> bool: + def is_integrated_gpu(cls, device_id: int | None = 0) -> bool: """Returns whether the GPU is an integrated (UMA) device that shares - system memory with the CPU. + system memory with the CPU. device_id None means the current device. On UMA systems (e.g. NVIDIA GH200, DGX Spark, Jetson Orin), cudaMemGetInfo may underreport free memory because it does not diff --git a/vllm/utils/mem_utils.py b/vllm/utils/mem_utils.py index f5233620736d..042f9f390292 100644 --- a/vllm/utils/mem_utils.py +++ b/vllm/utils/mem_utils.py @@ -114,12 +114,7 @@ def device_memory_info(device: torch.types.Device = None) -> tuple[int, int]: https://docs.nvidia.com/cuda/cuda-for-tegra-appnote/#estimating-total-allocatable-device-memory-on-an-integrated-gpu-device """ free_memory, total_memory = torch.accelerator.get_memory_info(device) - torch_device = torch.device(device) if device is not None else None - if torch_device is not None and torch_device.type == "cpu": - return free_memory, total_memory - index = torch_device.index if torch_device is not None else None - if index is None: - index = torch.accelerator.current_device_index() + index = torch.device(device).index if device is not None else None if current_platform.is_integrated_gpu(index): free_memory = psutil.virtual_memory().available return free_memory, total_memory