[Core] Size the KV cache from this process's own device memory, not the device-wide delta

tests/utils_/test_mem_utils.py

@@ -1,10 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import multiprocessing as mp
 from unittest.mock import MagicMock, patch
 
+import pytest
 import torch
 from vllm_test_utils.monitor import monitor
 
+from vllm.platforms import current_platform
 from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
 
 from ..utils import create_new_process_for_each_test
@@ -148,3 +151,121 @@ 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 _mock_measurements(mock_platform, mock_accelerator, free_values, process_values):
+    """Feed one (free memory, process memory) pair per MemorySnapshot.measure()."""
+    total = 80 * 1024**3
+    mock_accelerator.get_memory_info.side_effect = [(f, total) for f in free_values]
+    mock_accelerator.memory_stats.return_value = {
+        "allocated_bytes.all.peak": 0,
+        "allocated_bytes.all.current": 0,
+    }
+    mock_accelerator.memory_reserved.return_value = 0
+    mock_accelerator.current_device = lambda: "cuda:0"
+    mock_platform.is_integrated_gpu.return_value = False
+    mock_platform.get_process_memory_usage.side_effect = list(process_values)
+
+
+def test_memory_profiling_uses_process_scoped_consumption():
+    """Another process allocating during load/profile must not be charged
+    to this instance when per-process usage is available."""
+    gib = 1024**3
+    with (
+        patch("vllm.utils.mem_utils.current_platform") as mock_platform,
+        patch("torch.accelerator") as mock_accelerator,
+    ):
+        # baseline -> before_profile -> after_profile:
+        # this process grows 1 GiB then 1 GiB more; another process takes
+        # 2 GiB in the same window, so the device-wide delta is 4 GiB.
+        _mock_measurements(
+            mock_platform,
+            mock_accelerator,
+            free_values=[70 * gib, 69 * gib, 66 * gib],
+            process_values=[1 * gib, 2 * gib, 3 * gib],
+        )
+        baseline = MemorySnapshot(device="cuda:0")
+        with memory_profiling(baseline_snapshot=baseline) as result:
+            pass
+
+    assert baseline.process_memory == 1 * gib
+    assert result.process_scoped
+    assert result.total_consumed == 2 * gib
+    assert result.non_kv_cache_memory == 2 * gib
+
+
+def test_memory_profiling_falls_back_to_device_delta_without_process_usage():
+    gib = 1024**3
+    with (
+        patch("vllm.utils.mem_utils.current_platform") as mock_platform,
+        patch("torch.accelerator") as mock_accelerator,
+    ):
+        _mock_measurements(
+            mock_platform,
+            mock_accelerator,
+            free_values=[70 * gib, 69 * gib, 66 * gib],
+            process_values=[None, None, None],
+        )
+        baseline = MemorySnapshot(device="cuda:0")
+        with memory_profiling(baseline_snapshot=baseline) as result:
+            pass
+
+    assert baseline.process_memory is None
+    assert not result.process_scoped
+    assert result.total_consumed == 4 * gib
+
+
+def test_memory_snapshot_subtraction_keeps_process_memory():
+    a = MemorySnapshot(device="cuda:0", auto_measure=False)
+    b = MemorySnapshot(device="cuda:0", auto_measure=False)
+    a.process_memory, b.process_memory = 5, 2
+    assert (a - b).process_memory == 3
+    b.process_memory = None
+    assert (a - b).process_memory is None
+
+
+def _hold_device_memory(num_bytes: int, ready, release):
+    import torch
+
+    buf = torch.empty(num_bytes, dtype=torch.uint8, device="cuda")
+    torch.cuda.synchronize()
+    ready.set()
+    release.wait()
+    del buf
+
+
+@create_new_process_for_each_test()
+def test_memory_profiling_ignores_other_process_allocations():
+    """End-to-end on a real device: a second process allocates 512 MiB while
+    this one is profiling; total_consumed must only reflect our own 256 MiB."""
+    _warmup = torch.zeros(1, device="cuda")
+    del _warmup
+    torch.accelerator.empty_cache()
+    if current_platform.get_process_memory_usage(torch.cuda.current_device()) is None:
+        pytest.skip("platform cannot report per-process device memory usage")
+
+    baseline_snapshot = MemorySnapshot()
+    weights = torch.randn(64, 1024, 1024, device="cuda", dtype=torch.float32)
+    weights_memory = 64 * 1024 * 1024 * 4  # 256 MiB
+
+    ctx = mp.get_context("spawn")
+    ready, release = ctx.Event(), ctx.Event()
+    other = ctx.Process(
+        target=_hold_device_memory, args=(512 * 1024 * 1024, ready, release)
+    )
+    other.start()
+    try:
+        with memory_profiling(
+            baseline_snapshot=baseline_snapshot, weights_memory=weights_memory
+        ) as result:
+            assert ready.wait(timeout=120), "helper process did not allocate"
+    finally:
+        release.set()
+        other.join(timeout=60)
+
+    assert result.process_scoped
+    ratio = result.total_consumed / weights_memory
+    assert abs(ratio - 1) <= 0.05, (
+        f"total_consumed={result.total_consumed}, expected~{weights_memory}"
+    )
+    del weights

vllm/platforms/cuda.py

@@ -816,6 +816,34 @@ def get_device_total_memory(cls, device_id: int = 0) -> int:
         handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
         return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
 
+    @classmethod
+    @with_nvml_context
+    def get_process_memory_usage(cls, device_id: int = 0) -> int | None:
+        """Device memory used by this process on the visible device
+        ``device_id`` as reported by NVML.
+
+        Returns ``None`` whenever NVML cannot attribute memory to this
+        process (no entry for our PID, e.g. a container whose PID namespace
+        NVML does not see, WDDM, MIG, or an NVML error) so that callers fall
+        back to device-level accounting.
+
+        Keep this on the per-process query: ``nvmlDeviceGetMemoryInfo`` is
+        ``NVMLError_NotSupported`` on integrated parts such as GB10 (DGX Spark),
+        where this per-process path is exactly what still works.
+        """
+        try:
+            physical_device_id = cls.visible_device_id_to_physical_device_id(device_id)
+            handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
+            processes = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
+        except (pynvml.NVMLError, IndexError, ValueError):
+            return None
+        pid = os.getpid()
+        for proc in processes:
+            if proc.pid == pid:
+                used = proc.usedGpuMemory
+                return int(used) if isinstance(used, int) else None
+        return None
+
     @classmethod
     @with_nvml_context
     def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:

vllm/platforms/interface.py

@@ -1036,6 +1036,15 @@ def get_current_memory_usage(
         """
         raise NotImplementedError
 
+    @classmethod
+    def get_process_memory_usage(cls, device_id: int = 0) -> int | None:
+        """
+        Return the device memory in bytes used by the current process on the
+        visible device ``device_id``, or ``None`` when the platform cannot
+        attribute device memory to processes.
+        """
+        return None
+
     @classmethod
     def get_punica_wrapper(cls) -> str:
         """

vllm/utils/mem_utils.py

@@ -18,6 +18,10 @@
 
 logger = init_logger(__name__)
 
+# Device-wide vs. per-process consumption differences below this are not
+# reported (CUDA runtime bookkeeping, other processes' minor fluctuations).
+_OTHER_PROCESS_DELTA_WARN_BYTES = 128 * MiB_bytes
+
 
 def format_kib(b: int) -> str:
     return f"{round(b / KiB_bytes, 2)}"
@@ -116,6 +120,9 @@ class MemorySnapshot:
     cuda_memory: int = 0
     torch_memory: int = 0
     non_torch_memory: int = 0
+    # Memory used by this process alone (per-process accounting, e.g. NVML);
+    # None when the platform cannot attribute device memory to processes.
+    process_memory: int | None = None
     timestamp: float = 0.0
 
     device: torch.types.Device = None
@@ -162,6 +169,10 @@ def measure(self) -> None:
         self.torch_memory = torch.accelerator.memory_reserved(device)
 
         self.non_torch_memory = self.cuda_memory - self.torch_memory
+        process_memory = current_platform.get_process_memory_usage(device.index or 0)
+        self.process_memory = (
+            process_memory if isinstance(process_memory, int) else None
+        )
         self.timestamp = time.time()
 
     def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot":
@@ -179,11 +190,21 @@ def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot":
             cuda_memory=self.cuda_memory - other.cuda_memory,
             torch_memory=self.torch_memory - other.torch_memory,
             non_torch_memory=self.non_torch_memory - other.non_torch_memory,
+            process_memory=(
+                self.process_memory - other.process_memory
+                if self.process_memory is not None and other.process_memory is not None
+                else None
+            ),
             timestamp=self.timestamp - other.timestamp,
             device=self.device_,
             auto_measure=False,
         )
 
+    def _format_process_memory(self) -> str:
+        if self.process_memory is None:
+            return "n/a"
+        return f"{format_gib(self.process_memory)}GiB"
+
     def __repr__(self) -> str:
         return (
             f"torch_peak={format_gib(self.torch_peak)}GiB, "
@@ -193,6 +214,7 @@ def __repr__(self) -> str:
             f"{current_platform.device_name}_memory={format_gib(self.cuda_memory)}GiB, "
             f"torch_memory={format_gib(self.torch_memory)}GiB, "
             f"non_torch_memory={format_gib(self.non_torch_memory)}GiB, "
+            f"process_memory={self._format_process_memory()}, "
             f"timestamp={self.timestamp}, "
             f"auto_measure={self.auto_measure}"
         )
@@ -206,6 +228,9 @@ class MemoryProfilingResult:
     torch_peak_increase: int = 0
     non_torch_increase: int = 0
     total_consumed: int = 0
+    # True when total_consumed comes from this process's own device memory
+    # usage rather than from the device-wide free-memory delta.
+    process_scoped: bool = False
     transient_peak_headroom: int = 0
     weights_memory: int = 0
     before_create: MemorySnapshot = field(default_factory=MemorySnapshot)
@@ -314,9 +339,33 @@ def memory_profiling(
     # Measure total consumption via mem_get_info() instead of
     # memory_reserved(), which goes negative when pluggable allocators
     # (e.g. cumem) bypass PyTorch's tracking.
-    result.total_consumed = (
+    device_consumed = (
         result.before_create.free_memory - result.after_profile.free_memory
     )
+    before_process = result.before_create.process_memory
+    after_process = result.after_profile.process_memory
+    if before_process is not None and after_process is not None:
+        # Per-process accounting: memory that other processes on the same
+        # device allocate or release while this instance loads and profiles
+        # must not be charged to (or credited against) this instance.
+        result.total_consumed = after_process - before_process
+        result.process_scoped = True
+        other_processes_delta = device_consumed - result.total_consumed
+        if abs(other_processes_delta) >= _OTHER_PROCESS_DELTA_WARN_BYTES:
+            logger.warning(
+                "Other processes on %s changed their device memory usage by "
+                "%s GiB while this instance was loading and profiling. The "
+                "KV cache budget is based on this process's own usage "
+                "(%s GiB) rather than the device-wide change (%s GiB); make "
+                "sure the instances sharing this device do not request more "
+                "memory than it has in total.",
+                result.before_create.device_,
+                format_gib(other_processes_delta),
+                format_gib(result.total_consumed),
+                format_gib(device_consumed),
+            )
+    else:
+        result.total_consumed = device_consumed
 
     # total_consumed already covers persistent torch allocations; add only the
     # transient peak headroom to avoid double-counting.

vllm/v1/worker/gpu_worker.py

@@ -591,7 +591,17 @@ def determine_available_memory(self) -> int:
         init_free_memory = self.init_snapshot.free_memory
         free_gpu_memory = profile_result.after_profile.free_memory
         rocm_fallback = maybe_rocm_profiling_fallback(profile_result)
-        if rocm_fallback is None:
+        if rocm_fallback is None and profile_result.process_scoped:
+            # The budget is based on this process's own memory usage, so
+            # other processes releasing memory during profiling is harmless.
+            if init_free_memory < free_gpu_memory:
+                logger.warning(
+                    "Other processes released %s GiB on the device while this "
+                    "instance was profiling; ignored because the KV cache "
+                    "budget uses this process's own memory usage.",
+                    format_gib(free_gpu_memory - init_free_memory),
+                )
+        elif rocm_fallback is None:
             # NOTE(woosuk): Here we assume that the other processes using the same
             # GPU did not change their memory usage during the profiling.
             assert init_free_memory >= free_gpu_memory, (