A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 06:00 UTC
[Bugfix] Use is_integrated to detect UMA GPUs for memory reporting (#35356)
Signed-off-by: haosdent <haosdent@gmail.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>
haosdent committed
Apr 14, 2026 at 02:07 UTC
995e9a209e68a95ffa03c73f3401472837a4072b
4 files changed
+83
-16
tests/utils_/test_mem_utils.py
+61
@@ -1,5 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from unittest.mock import MagicMock, patch
4
+
5
import torch
6
from vllm_test_utils.monitor import monitor
7
@@ -61,3 +63,62 @@ def test_memory_profiling():
63
del weights
64
lib.cudaFree(handle1)
65
lib.cudaFree(handle2)
66
+
67
+
68
+def test_memory_snapshot_uses_psutil_on_integrated_gpu():
69
+ """On integrated (UMA) GPUs, free_memory should come from psutil."""
70
+ mock_cuda_free = 40 * 1024**3
71
+ mock_cuda_total = 120 * 1024**3
72
+ mock_psutil_available = 100 * 1024**3
73
+
74
+ with (
75
+ patch("vllm.utils.mem_utils.current_platform") as mock_platform,
76
+ patch("vllm.utils.mem_utils.psutil") as mock_psutil,
77
+ ):
78
+ mock_platform.mem_get_info.return_value = (
79
+ mock_cuda_free,
80
+ mock_cuda_total,
81
+ )
82
+ mock_platform.is_integrated_gpu.return_value = True
83
+ mock_platform.memory_stats.return_value = {
84
+ "allocated_bytes.all.peak": 0,
85
+ }
86
+ mock_platform.memory_reserved.return_value = 0
87
+ mock_platform.current_device = lambda: "cuda:0"
88
+
89
+ mock_vmem = MagicMock()
90
+ mock_vmem.available = mock_psutil_available
91
+ mock_psutil.virtual_memory.return_value = mock_vmem
92
+
93
+ snapshot = MemorySnapshot(device="cuda:0")
94
+
95
+ assert snapshot.free_memory == mock_psutil_available
96
+ assert snapshot.total_memory == mock_cuda_total
97
+ mock_psutil.virtual_memory.assert_called_once()
98
+
99
+
100
+def test_memory_snapshot_uses_cuda_on_discrete_gpu():
101
+ """On discrete GPUs, free_memory should come from CUDA mem_get_info."""
102
+ mock_cuda_free = 70 * 1024**3
103
+ mock_cuda_total = 80 * 1024**3
104
+
105
+ with (
106
+ patch("vllm.utils.mem_utils.current_platform") as mock_platform,
107
+ patch("vllm.utils.mem_utils.psutil") as mock_psutil,
108
+ ):
109
+ mock_platform.mem_get_info.return_value = (
110
+ mock_cuda_free,
111
+ mock_cuda_total,
112
+ )
113
+ mock_platform.is_integrated_gpu.return_value = False
114
+ mock_platform.memory_stats.return_value = {
115
+ "allocated_bytes.all.peak": 0,
116
+ }
117
+ mock_platform.memory_reserved.return_value = 0
118
+ mock_platform.current_device = lambda: "cuda:0"
119
+
120
+ snapshot = MemorySnapshot(device="cuda:0")
121
+
122
+ assert snapshot.free_memory == mock_cuda_free
123
+ assert snapshot.total_memory == mock_cuda_total
124
+ mock_psutil.virtual_memory.assert_not_called()
vllm/platforms/cuda.py
+4
@@ -548,6 +548,10 @@ class CudaPlatformBase(Platform):
548
"""Currently, only Hopper and Blackwell GPUs are supported."""
549
return cls.is_device_capability(90) or cls.is_device_capability_family(100)
550
551
+ @classmethod
552
+ def is_integrated_gpu(cls, device_id: int = 0) -> bool:
553
+ return bool(torch.cuda.get_device_properties(device_id).is_integrated)
554
+
555
@classmethod
556
def num_compute_units(cls, device_id: int = 0) -> int:
557
return torch.cuda.get_device_properties(device_id).multi_processor_count
vllm/platforms/interface.py
+12
@@ -734,6 +734,18 @@ class Platform:
734
"""
735
return "vllm.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa
736
737
+ @classmethod
738
+ def is_integrated_gpu(cls, device_id: int = 0) -> bool:
739
+ """
740
+ Returns whether the GPU is an integrated (UMA) device that shares
741
+ system memory with the CPU.
742
+
743
+ On UMA systems (e.g. NVIDIA GH200, DGX Spark, Jetson Orin),
744
+ cudaMemGetInfo may underreport free memory because it does not
745
+ account for reclaimable OS memory (page cache, buffers).
746
+ """
747
+ return False
748
+
749
@classmethod
750
def supports_mx(cls) -> bool:
751
"""
vllm/utils/mem_utils.py
+6
-16
@@ -106,22 +106,12 @@ class MemorySnapshot:
106
)
107
108
self.free_memory, self.total_memory = current_platform.mem_get_info(device)
109
- shared_sysmem_device_mem_sms = ((8, 7), (11, 0), (12, 1)) # Orin, Thor, Spark
110
- if (
111
- current_platform.is_cuda()
112
- and current_platform.get_device_capability(device.index)
113
- in shared_sysmem_device_mem_sms
114
- ):
115
- # On UMA (Orin, Thor and Spark) platform,
116
- # where both CPU and GPU rely on system memory,
117
- # the cudaMemGetInfo function shows the amount of free system memory
118
- # rather than what’s actually available.
119
- # In the case,
120
- # torch.cuda.mem_get_info() only reports "free" memory,
121
- # which can be lower than what is actually
122
- # available due to not including cache memory.
123
- # There’s also a comprehensive reference page
124
- # that explains how you can compute the proper value yourself.
109
+ if current_platform.is_integrated_gpu(device.index):
110
+ # On UMA (Unified Memory Architecture) platforms where CPU and
111
+ # GPU share physical memory (e.g. GH200, DGX Spark, Jetson Orin),
112
+ # cudaMemGetInfo underreports free memory because it does not
113
+ # account for reclaimable OS memory (page cache, buffers).
114
+ # Use psutil to get the true available memory.
115
# https://docs.nvidia.com/cuda/cuda-for-tegra-appnote/#estimating-total-allocatable-device-memory-on-an-integrated-gpu-device
116
self.free_memory = psutil.virtual_memory().available
117