| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | import gc |
| 5 | import time |
| 6 | from collections.abc import Generator |
| 7 | from dataclasses import dataclass, field |
| 8 | from functools import cache |
| 9 | |
| 10 | 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)}" |
| 24 | |
| 25 | |
| 26 | def format_mib(b: int) -> str: |
| 27 | return f"{round(b / MiB_bytes, 2)}" |
| 28 | |
| 29 | |
| 30 | def format_gib(b: int) -> str: |
| 31 | return f"{round(b / GiB_bytes, 2)}" |
| 32 | |
| 33 | |
| 34 | @cache |
| 35 | def get_max_shared_memory_bytes(gpu: int = 0) -> int: |
| 36 | """Returns the maximum shared memory per thread block in bytes.""" |
| 37 | from vllm import _custom_ops as ops |
| 38 | |
| 39 | max_shared_mem = ops.get_max_shared_memory_per_block_device_attribute(gpu) |
| 40 | # value 0 will cause MAX_SEQ_LEN become negative and test_attention.py |
| 41 | # will fail |
| 42 | assert max_shared_mem > 0, "max_shared_mem cannot be zero" |
| 43 | return int(max_shared_mem) |
| 44 | |
| 45 | |
| 46 | _UMA_PRESSURE_THRESHOLD = 0.8 |
| 47 | _UMA_MIN_RELEASE_BYTES = 512 * MiB_bytes |
| 48 | |
| 49 | |
| 50 | def release_device_memory_under_pressure(device: torch.device) -> bool: |
| 51 | """On integrated (UMA) GPUs, release caching-allocator memory back to the |
| 52 | OS when system memory pressure is high. The OS may start thrashing before |
| 53 | an allocation failure would trigger PyTorch's own cache release. |
| 54 | |
| 55 | Returns: |
| 56 | True if memory was released. |
| 57 | |
| 58 | """ |
| 59 | if device.type != "cuda" or not current_platform.is_integrated_gpu(device.index): |
| 60 | return False |
| 61 | |
| 62 | releasable = torch.accelerator.memory_reserved( |
| 63 | device |
| 64 | ) - torch.accelerator.memory_allocated(device) |
| 65 | if releasable < _UMA_MIN_RELEASE_BYTES: |
| 66 | return False |
| 67 | |
| 68 | # cudaMemGetInfo underreports free memory on UMA, see MemorySnapshot.measure |
| 69 | mem = psutil.virtual_memory() |
| 70 | if mem.available > (1 - _UMA_PRESSURE_THRESHOLD) * mem.total: |
| 71 | return False |
| 72 | |
| 73 | torch.accelerator.synchronize(device) |
| 74 | torch.accelerator.empty_cache() |
| 75 | logger.debug( |
| 76 | "Released %sGiB of cached device memory under memory pressure", |
| 77 | format_gib(releasable), |
| 78 | ) |
| 79 | return True |
| 80 | |
| 81 | |
| 82 | class DeviceMemoryProfiler: |
| 83 | def __init__(self, device: torch.types.Device | None = None): |
| 84 | self.device = device |
| 85 | |
| 86 | def current_memory_usage(self) -> float: |
| 87 | # Return the memory usage in bytes. |
| 88 | gc.collect() |
| 89 | return current_platform.get_current_memory_usage(self.device) |
| 90 | |
| 91 | def __enter__(self): |
| 92 | self.initial_memory = self.current_memory_usage() |
| 93 | # This allows us to call methods of the context manager if needed |
| 94 | return self |
| 95 | |
| 96 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 97 | self.final_memory = self.current_memory_usage() |
| 98 | self.consumed_memory = self.final_memory - self.initial_memory |
| 99 | |
| 100 | # Force garbage collection |
| 101 | gc.collect() |
| 102 | |
| 103 | |
| 104 | @dataclass |
| 105 | class MemorySnapshot: |
| 106 | """Memory snapshot.""" |
| 107 | |
| 108 | torch_peak: int = 0 |
| 109 | torch_allocated: int = 0 |
| 110 | free_memory: int = 0 |
| 111 | total_memory: int = 0 |
| 112 | cuda_memory: int = 0 |
| 113 | torch_memory: int = 0 |
| 114 | non_torch_memory: int = 0 |
| 115 | timestamp: float = 0.0 |
| 116 | |
| 117 | device: torch.types.Device = None |
| 118 | auto_measure: bool = True |
| 119 | |
| 120 | def __post_init__(self) -> None: |
| 121 | if self.device is None: |
| 122 | device_fn = current_platform.current_device |
| 123 | assert device_fn is not None |
| 124 | self.device_ = torch.device(device_fn()) |
| 125 | else: |
| 126 | self.device_ = torch.device(self.device) |
| 127 | |
| 128 | if self.auto_measure: |
| 129 | self.measure() |
| 130 | |
| 131 | def measure(self) -> None: |
| 132 | device = self.device_ |
| 133 | |
| 134 | # we measure the torch peak memory usage via allocated_bytes, |
| 135 | # rather than `torch.accelerator.memory_reserved()` . |
| 136 | # After `torch.accelerator.reset_peak_memory_stats()`, |
| 137 | # `torch.accelerator.memory_reserved()` will keep growing, and only shrink |
| 138 | # when we call `torch.accelerator.empty_cache()` or OOM happens. |
| 139 | stats = torch.accelerator.memory_stats(device) |
| 140 | self.torch_peak = stats.get("allocated_bytes.all.peak", 0) |
| 141 | self.torch_allocated = stats.get("allocated_bytes.all.current", 0) |
| 142 | |
| 143 | self.free_memory, self.total_memory = torch.accelerator.get_memory_info(device) |
| 144 | if current_platform.is_integrated_gpu(device.index): |
| 145 | # On UMA (Unified Memory Architecture) platforms where CPU and |
| 146 | # GPU share physical memory (e.g. GH200, DGX Spark, Jetson Orin), |
| 147 | # cudaMemGetInfo underreports free memory because it does not |
| 148 | # account for reclaimable OS memory (page cache, buffers). |
| 149 | # Use psutil to get the true available memory. |
| 150 | # https://docs.nvidia.com/cuda/cuda-for-tegra-appnote/#estimating-total-allocatable-device-memory-on-an-integrated-gpu-device |
| 151 | self.free_memory = psutil.virtual_memory().available |
| 152 | |
| 153 | self.cuda_memory = self.total_memory - self.free_memory |
| 154 | |
| 155 | # torch.accelerator.memory_reserved() is how many bytes |
| 156 | # PyTorch gets from cuda (by calling cudaMalloc, etc.) |
| 157 | # this is used to measure the non-torch memory usage |
| 158 | self.torch_memory = torch.accelerator.memory_reserved(device) |
| 159 | |
| 160 | self.non_torch_memory = self.cuda_memory - self.torch_memory |
| 161 | self.timestamp = time.time() |
| 162 | |
| 163 | def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot": |
| 164 | if self.device_ != other.device_: |
| 165 | raise ValueError( |
| 166 | "The two snapshots should be from the same device! " |
| 167 | f"Found: {self.device_} vs. {other.device_}" |
| 168 | ) |
| 169 | |
| 170 | return MemorySnapshot( |
| 171 | torch_peak=self.torch_peak - other.torch_peak, |
| 172 | torch_allocated=self.torch_allocated - other.torch_allocated, |
| 173 | free_memory=self.free_memory - other.free_memory, |
| 174 | total_memory=self.total_memory - other.total_memory, |
| 175 | cuda_memory=self.cuda_memory - other.cuda_memory, |
| 176 | torch_memory=self.torch_memory - other.torch_memory, |
| 177 | non_torch_memory=self.non_torch_memory - other.non_torch_memory, |
| 178 | timestamp=self.timestamp - other.timestamp, |
| 179 | device=self.device_, |
| 180 | auto_measure=False, |
| 181 | ) |
| 182 | |
| 183 | def __repr__(self) -> str: |
| 184 | return ( |
| 185 | f"torch_peak={format_gib(self.torch_peak)}GiB, " |
| 186 | f"torch_allocated={format_gib(self.torch_allocated)}GiB, " |
| 187 | f"free_memory={format_gib(self.free_memory)}GiB, " |
| 188 | f"total_memory={format_gib(self.total_memory)}GiB, " |
| 189 | f"{current_platform.device_name}_memory={format_gib(self.cuda_memory)}GiB, " |
| 190 | f"torch_memory={format_gib(self.torch_memory)}GiB, " |
| 191 | f"non_torch_memory={format_gib(self.non_torch_memory)}GiB, " |
| 192 | f"timestamp={self.timestamp}, " |
| 193 | f"auto_measure={self.auto_measure}" |
| 194 | ) |
| 195 | |
| 196 | |
| 197 | @dataclass |
| 198 | class MemoryProfilingResult: |
| 199 | """Memory profiling result. All numbers are in bytes.""" |
| 200 | |
| 201 | non_kv_cache_memory: int = 0 |
| 202 | torch_peak_increase: int = 0 |
| 203 | non_torch_increase: int = 0 |
| 204 | total_consumed: int = 0 |
| 205 | transient_peak_headroom: int = 0 |
| 206 | weights_memory: int = 0 |
| 207 | before_create: MemorySnapshot = field(default_factory=MemorySnapshot) |
| 208 | profile_time: float = 0.0 |
| 209 | |
| 210 | def __post_init__(self) -> None: |
| 211 | device = self.before_create.device_ |
| 212 | |
| 213 | self.before_profile = MemorySnapshot(device=device, auto_measure=False) |
| 214 | self.after_profile = MemorySnapshot(device=device, auto_measure=False) |
| 215 | |
| 216 | def __repr__(self) -> str: |
| 217 | return ( |
| 218 | f"Memory profiling takes {self.profile_time:.2f} seconds. " |
| 219 | f"Total non KV cache memory: " |
| 220 | f"{format_gib(self.non_kv_cache_memory)}GiB; " |
| 221 | f"torch peak memory increase: " |
| 222 | f"{format_gib(self.torch_peak_increase)}GiB; " |
| 223 | f"total consumed (from mem_get_info): " |
| 224 | f"{format_gib(self.total_consumed)}GiB; " |
| 225 | f"weights memory: {format_gib(self.weights_memory)}GiB." |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | @contextlib.contextmanager |
| 230 | def memory_profiling( |
| 231 | baseline_snapshot: MemorySnapshot, |
| 232 | weights_memory: int = 0, |
| 233 | ) -> Generator[MemoryProfilingResult, None, None]: |
| 234 | """Memory profiling context manager. |
| 235 | |
| 236 | baseline_snapshot: the memory snapshot before the current vLLM instance. |
| 237 | weights_memory: memory used by PyTorch when loading the model weights. |
| 238 | Note that, before loading the model weights, we also initialize the device |
| 239 | and distributed environment, which may consume some memory. This part is not |
| 240 | included in the weights_memory because PyTorch does not control it. |
| 241 | |
| 242 | The memory in one GPU can be classified into 3 categories: |
| 243 | 1. memory used by anything other than the current vLLM instance. |
| 244 | 2. memory used by torch in the current vLLM instance. |
| 245 | 3. memory used in the current vLLM instance, but not by torch. |
| 246 | |
| 247 | A quantitive example: |
| 248 | |
| 249 | Before creating the current vLLM instance: |
| 250 | category 1: 1 GiB |
| 251 | category 2: 0 GiB |
| 252 | category 3: 0 GiB |
| 253 | |
| 254 | After creating the current vLLM instance and loading the model, |
| 255 | (i.e. before profiling): |
| 256 | category 1: 1 GiB |
| 257 | category 2: 2 GiB (model weights take 2 GiB) |
| 258 | category 3: 0.5 GiB (memory used by NCCL) |
| 259 | |
| 260 | During profiling (peak): |
| 261 | category 1: 1 GiB |
| 262 | category 2: 4 GiB (peak activation tensors take 2 GiB) |
| 263 | category 3: 1 GiB (memory used by NCCL + buffers for some attention backends) |
| 264 | |
| 265 | After profiling: |
| 266 | category 1: 1 GiB |
| 267 | category 2: 3 GiB (after garbage-collecting activation tensors) |
| 268 | category 3: 1 GiB (memory used by NCCL + buffers for some attention backends) |
| 269 | |
| 270 | In this case, non-kv cache takes 5 GiB in total, including: |
| 271 | a. 2 GiB used by the model weights (category 2) |
| 272 | b. 2 GiB reserved for the peak activation tensors (category 2) |
| 273 | c. 1 GiB used by non-torch components (category 3) |
| 274 | |
| 275 | The memory used for loading weights (a.) is directly given from the |
| 276 | argument `weights_memory`. |
| 277 | |
| 278 | The increase of `torch.accelerator.memory_stats()["allocated_bytes.all.peak"]` |
| 279 | during profiling gives (b.). |
| 280 | |
| 281 | The increase of `non_torch_memory` from creating the current vLLM instance |
| 282 | until after profiling to get (c.). |
| 283 | """ |
| 284 | gc.collect() |
| 285 | torch.accelerator.empty_cache() |
| 286 | torch.accelerator.reset_peak_memory_stats(baseline_snapshot.device_) |
| 287 | |
| 288 | result = MemoryProfilingResult( |
| 289 | before_create=baseline_snapshot, |
| 290 | # the part of memory used for holding the model weights |
| 291 | weights_memory=weights_memory, |
| 292 | ) |
| 293 | |
| 294 | result.before_profile.measure() |
| 295 | |
| 296 | yield result |
| 297 | |
| 298 | gc.collect() |
| 299 | torch.accelerator.empty_cache() |
| 300 | |
| 301 | result.after_profile.measure() |
| 302 | |
| 303 | diff_profile = result.after_profile - result.before_profile |
| 304 | diff_from_create = result.after_profile - result.before_create |
| 305 | result.torch_peak_increase = diff_profile.torch_peak |
| 306 | result.non_torch_increase = diff_from_create.non_torch_memory |
| 307 | result.profile_time = diff_profile.timestamp |
| 308 | |
| 309 | # Measure total consumption via mem_get_info() instead of |
| 310 | # memory_reserved(), which goes negative when pluggable allocators |
| 311 | # (e.g. cumem) bypass PyTorch's tracking. |
| 312 | result.total_consumed = ( |
| 313 | result.before_create.free_memory - result.after_profile.free_memory |
| 314 | ) |
| 315 | |
| 316 | # total_consumed already covers persistent torch allocations; add only the |
| 317 | # transient peak headroom to avoid double-counting. |
| 318 | result.transient_peak_headroom = ( |
| 319 | result.after_profile.torch_peak - result.after_profile.torch_allocated |
| 320 | ) |
| 321 | result.non_kv_cache_memory = result.total_consumed + result.transient_peak_headroom |