diff --git a/tests/utils_/test_mem_utils.py b/tests/utils_/test_mem_utils.py index 19690d1f34f8..78ccba8ad895 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,31 @@ 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) diff --git a/tests/v1/worker/test_extensible_kv_cache.py b/tests/v1/worker/test_extensible_kv_cache.py index 35004b5bd4f8..8f8e2bf25211 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,69 @@ 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) + ) + 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/platforms/cuda.py b/vllm/platforms/cuda.py index 692e719564be..27fbef7dc165 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -696,7 +696,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 1b449058b27d..042f9f390292 100644 --- a/vllm/utils/mem_utils.py +++ b/vllm/utils/mem_utils.py @@ -101,6 +101,25 @@ 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) + 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 + + @dataclass class MemorySnapshot: """Memory snapshot.""" @@ -140,15 +159,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 fdeadee49565..e559c003d050 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.hisparse.binding import rebind_hisparse_device_caches from vllm.v1.kv_cache_interface import ( KVCacheConfig, @@ -255,7 +255,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, @@ -570,7 +570,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, @@ -671,7 +671,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() )