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 11, 07:30 UTC
[Chore] Factor out logic for requesting initial memory (#30868)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Dec 17, 2025 at 23:32 UTC
2497228ad4427310bc55427f6db404a00de4fd78
3 files changed
+56
-21
vllm/utils/mem_utils.py
+27
-4
@@ -66,27 +66,43 @@ class MemorySnapshot:
66
torch_memory: int = 0
67
non_torch_memory: int = 0
68
timestamp: float = 0.0
69
+
70
+ device: torch.types.Device = None
71
auto_measure: bool = True
72
73
def __post_init__(self) -> None:
74
+ if self.device is None:
75
+ from vllm.platforms import current_platform
76
+
77
+ device_fn = current_platform.current_device
78
+ assert device_fn is not None
79
+ self.device_ = torch.device(device_fn())
80
+ else:
81
+ self.device_ = torch.device(self.device)
82
+
83
if self.auto_measure:
84
self.measure()
85
86
def measure(self) -> None:
87
from vllm.platforms import current_platform
88
89
+ device = self.device_
90
+
91
# we measure the torch peak memory usage via allocated_bytes,
92
# rather than `torch.cuda.memory_reserved()` .
93
# After `torch.cuda.reset_peak_memory_stats()`,
94
# `torch.cuda.memory_reserved()` will keep growing, and only shrink
95
# when we call `torch.cuda.empty_cache()` or OOM happens.
83
- self.torch_peak = torch.cuda.memory_stats().get("allocated_bytes.all.peak", 0)
96
+ self.torch_peak = torch.cuda.memory_stats(device).get(
97
+ "allocated_bytes.all.peak", 0
98
+ )
99
85
- self.free_memory, self.total_memory = torch.cuda.mem_get_info()
100
+ self.free_memory, self.total_memory = torch.cuda.mem_get_info(device)
101
shared_sysmem_device_mem_sms = ((8, 7), (11, 0), (12, 1)) # Orin, Thor, Spark
102
if (
103
current_platform.is_cuda()
89
- and current_platform.get_device_capability() in shared_sysmem_device_mem_sms
104
+ and current_platform.get_device_capability(device.index)
105
+ in shared_sysmem_device_mem_sms
106
):
107
# On UMA (Orin, Thor and Spark) platform,
108
# where both CPU and GPU rely on system memory,
@@ -106,12 +122,18 @@ class MemorySnapshot:
122
# torch.cuda.memory_reserved() is how many bytes
123
# PyTorch gets from cuda (by calling cudaMalloc, etc.)
124
# this is used to measure the non-torch memory usage
109
- self.torch_memory = torch.cuda.memory_reserved()
125
+ self.torch_memory = torch.cuda.memory_reserved(device)
126
127
self.non_torch_memory = self.cuda_memory - self.torch_memory
128
self.timestamp = time.time()
129
130
def __sub__(self, other: "MemorySnapshot") -> "MemorySnapshot":
131
+ if self.device_ != other.device_:
132
+ raise ValueError(
133
+ "The two snapshots should be from the same device! "
134
+ f"Found: {self.device_} vs. {other.device_}"
135
+ )
136
+
137
return MemorySnapshot(
138
torch_peak=self.torch_peak - other.torch_peak,
139
free_memory=self.free_memory - other.free_memory,
@@ -120,6 +142,7 @@ class MemorySnapshot:
142
torch_memory=self.torch_memory - other.torch_memory,
143
non_torch_memory=self.non_torch_memory - other.non_torch_memory,
144
timestamp=self.timestamp - other.timestamp,
145
+ device=self.device_,
146
auto_measure=False,
147
)
148
vllm/v1/worker/gpu_worker.py
+4
-16
@@ -56,6 +56,8 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp
56
from vllm.v1.worker.worker_base import WorkerBase
57
from vllm.v1.worker.workspace import init_workspace_manager
58
59
+from .utils import request_memory
60
+
61
logger = init_logger(__name__)
62
63
if TYPE_CHECKING:
@@ -237,22 +239,8 @@ class Worker(WorkerBase):
239
torch.cuda.empty_cache()
240
241
# take current memory snapshot
240
- self.init_snapshot = MemorySnapshot()
241
- self.requested_memory = (
242
- self.init_snapshot.total_memory
243
- * self.cache_config.gpu_memory_utilization
244
- )
245
- if self.init_snapshot.free_memory < self.requested_memory:
246
- GiB = lambda b: round(b / GiB_bytes, 2)
247
- raise ValueError(
248
- f"Free memory on device "
249
- f"({GiB(self.init_snapshot.free_memory)}/"
250
- f"{GiB(self.init_snapshot.total_memory)} GiB) on startup "
251
- f"is less than desired GPU memory utilization "
252
- f"({self.cache_config.gpu_memory_utilization}, "
253
- f"{GiB(self.requested_memory)} GiB). Decrease GPU memory "
254
- f"utilization or reduce GPU memory used by other processes."
255
- )
242
+ self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
243
+ self.requested_memory = request_memory(init_snapshot, self.cache_config)
244
else:
245
raise RuntimeError(f"Not support device type: {self.device_config.device}")
246
vllm/v1/worker/utils.py
+25
-1
@@ -8,13 +8,15 @@ from typing_extensions import deprecated
8
9
from vllm.attention.backends.abstract import AttentionBackend
10
from vllm.attention.layer import Attention
11
-from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
11
+from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
12
from vllm.logger import init_logger
13
from vllm.model_executor.models.interfaces import MultiModalEmbeddings
14
from vllm.model_executor.models.utils import extract_layer_index
15
from vllm.multimodal.cache import processor_only_cache_from_config
16
from vllm.multimodal.registry import MultiModalRegistry
17
from vllm.platforms import current_platform
18
+from vllm.utils.mem_constants import GiB_bytes
19
+from vllm.utils.mem_utils import MemorySnapshot
20
from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
21
from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
22
from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec
@@ -248,6 +250,28 @@ def gather_mm_placeholders(
250
return placeholders[is_embed]
251
252
253
+def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) -> float:
254
+ """
255
+ Calculate the amount of memory required by vLLM, then validate
256
+ that the current amount of free memory is sufficient for that.
257
+ """
258
+ requested_memory = init_snapshot.total_memory * cache_config.gpu_memory_utilization
259
+
260
+ if init_snapshot.free_memory < requested_memory:
261
+ GiB = lambda b: round(b / GiB_bytes, 2)
262
+ raise ValueError(
263
+ f"Free memory on device {init_snapshot.device_} "
264
+ f"({GiB(init_snapshot.free_memory)}/"
265
+ f"{GiB(init_snapshot.total_memory)} GiB) on startup "
266
+ f"is less than desired GPU memory utilization "
267
+ f"({cache_config.gpu_memory_utilization}, "
268
+ f"{GiB(requested_memory)} GiB). Decrease GPU memory "
269
+ f"utilization or reduce GPU memory used by other processes."
270
+ )
271
+
272
+ return requested_memory
273
+
274
+
275
def add_kv_sharing_layers_to_kv_cache_groups(
276
shared_kv_cache_layers: dict[str, str],
277
kv_cache_groups: list[KVCacheGroupSpec],