[Bugfix] Size the extensible KV cache from UMA-corrected free memory

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()

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] = []

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
 

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()
     )