A high-throughput and memory-efficient inference and serving engine for LLMs
[Core] Release cached device memory under pressure on UMA GPUs during weight loading (#45179)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Michael Goin committed
Jun 11, 2026 at 15:49 UTC
5edf7ff489e83616c00c64d3ac6562f81dbe5638
2 files changed
+42
vllm/model_executor/model_loader/utils.py
+4
@@ -30,6 +30,7 @@ from vllm.model_executor.model_loader.reload import (
30
)
31
from vllm.model_executor.models.interfaces import SupportsQuant
32
from vllm.tracing import instrument
33
+from vllm.utils.mem_utils import release_device_memory_under_pressure
34
from vllm.utils.platform_utils import is_pin_memory_available
35
from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
36
@@ -109,6 +110,9 @@ def process_weights_after_loading(
110
# parameters onto device for processing and back off after.
111
with device_loading_context(module, target_device):
112
quant_method.process_weights_after_loading(module)
113
+ # Repacking transients above can leave large amounts of memory in
114
+ # the caching allocator, which starves the OS on UMA devices.
115
+ release_device_memory_under_pressure(target_device)
116
117
# Initialize post-load attention weights for Attention, MLA, and MM encoder.
118
# NOTE: Happens after other modules so we can easily decompress weights.
vllm/utils/mem_utils.py
+38
@@ -11,10 +11,13 @@ import psutil
11
import torch
12
import torch.types
13
14
+from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
17
from .mem_constants import GiB_bytes, KiB_bytes, MiB_bytes
18
19
+logger = init_logger(__name__)
20
+
21
22
def format_kib(b: int) -> str:
23
return f"{round(b / KiB_bytes, 2)}"
@@ -45,6 +48,41 @@ def get_cpu_memory() -> int:
48
return psutil.virtual_memory().total
49
50
51
+_UMA_PRESSURE_THRESHOLD = 0.8
52
+_UMA_MIN_RELEASE_BYTES = 512 * MiB_bytes
53
+
54
+
55
+def release_device_memory_under_pressure(device: torch.device) -> bool:
56
+ """On integrated (UMA) GPUs, release caching-allocator memory back to the
57
+ OS when system memory pressure is high. The OS may start thrashing before
58
+ an allocation failure would trigger PyTorch's own cache release.
59
+
60
+ Returns:
61
+ True if memory was released.
62
+ """
63
+ if device.type != "cuda" or not current_platform.is_integrated_gpu(device.index):
64
+ return False
65
+
66
+ releasable = torch.accelerator.memory_reserved(
67
+ device
68
+ ) - torch.accelerator.memory_allocated(device)
69
+ if releasable < _UMA_MIN_RELEASE_BYTES:
70
+ return False
71
+
72
+ # cudaMemGetInfo underreports free memory on UMA, see MemorySnapshot.measure
73
+ mem = psutil.virtual_memory()
74
+ if mem.available > (1 - _UMA_PRESSURE_THRESHOLD) * mem.total:
75
+ return False
76
+
77
+ torch.accelerator.synchronize(device)
78
+ torch.accelerator.empty_cache()
79
+ logger.debug(
80
+ "Released %sGiB of cached device memory under memory pressure",
81
+ format_gib(releasable),
82
+ )
83
+ return True
84
+
85
+
86
class DeviceMemoryProfiler:
87
def __init__(self, device: torch.types.Device | None = None):
88
self.device = device