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 10, 16:30 UTC
[Log] add log about gpu worker init snapshot and requested memory (#29493)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Ning Xie committed
Jan 7, 2026 at 01:32 UTC
6f5e65338346a0be4466bfb31423c2968b7363bd
3 files changed
+59
-42
vllm/utils/mem_utils.py
+20
-4
@@ -14,6 +14,10 @@ import torch.types
14
from .mem_constants import GiB_bytes
15
16
17
+def format_gib(b: int) -> float:
18
+ return round(b / GiB_bytes, 2)
19
+
20
+
21
@cache
22
def get_max_shared_memory_bytes(gpu: int = 0) -> int:
23
"""Returns the maximum shared memory per thread block in bytes."""
@@ -146,6 +150,18 @@ class MemorySnapshot:
150
auto_measure=False,
151
)
152
153
+ def __repr__(self) -> str:
154
+ return (
155
+ f"torch_peak={format_gib(self.torch_peak)}GiB, "
156
+ f"free_memory={format_gib(self.free_memory)}GiB, "
157
+ f"total_memory={format_gib(self.total_memory)}GiB, "
158
+ f"cuda_memory={format_gib(self.cuda_memory)}GiB, "
159
+ f"torch_memory={format_gib(self.torch_memory)}GiB, "
160
+ f"non_torch_memory={format_gib(self.non_torch_memory)}GiB, "
161
+ f"timestamp={self.timestamp}, "
162
+ f"auto_measure={self.auto_measure}"
163
+ )
164
+
165
166
@dataclass
167
class MemoryProfilingResult:
@@ -168,12 +184,12 @@ class MemoryProfilingResult:
184
return (
185
f"Memory profiling takes {self.profile_time:.2f} seconds. "
186
f"Total non KV cache memory: "
171
- f"{(self.non_kv_cache_memory / GiB_bytes):.2f}GiB; "
187
+ f"{format_gib(self.non_kv_cache_memory)}GiB; "
188
f"torch peak memory increase: "
173
- f"{(self.torch_peak_increase / GiB_bytes):.2f}GiB; "
189
+ f"{format_gib(self.torch_peak_increase)}GiB; "
190
f"non-torch forward increase memory: "
175
- f"{(self.non_torch_increase / GiB_bytes):.2f}GiB; "
176
- f"weights memory: {(self.weights_memory / GiB_bytes):.2f}GiB."
191
+ f"{format_gib(self.non_torch_increase)}GiB; "
192
+ f"weights memory: {format_gib(self.weights_memory)}GiB."
193
)
194
195
vllm/v1/worker/gpu_worker.py
+30
-30
@@ -40,8 +40,7 @@ from vllm.platforms import current_platform
40
from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper
41
from vllm.sequence import IntermediateTensors
42
from vllm.tasks import SupportedTask
43
-from vllm.utils.mem_constants import GiB_bytes
44
-from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
43
+from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
44
from vllm.utils.torch_utils import set_random_seed
45
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
46
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
@@ -132,9 +131,9 @@ class Worker(WorkerBase):
131
used_bytes = total - free_bytes_after_sleep
132
assert freed_bytes >= 0, "Memory usage increased after sleeping."
133
logger.info(
135
- "Sleep mode freed %.2f GiB memory, %.2f GiB memory is still in use.",
136
- freed_bytes / GiB_bytes,
137
- used_bytes / GiB_bytes,
134
+ "Sleep mode freed %f GiB memory, %f GiB memory is still in use.",
135
+ format_gib(freed_bytes),
136
+ format_gib(used_bytes),
137
)
138
139
def wake_up(self, tags: list[str] | None = None) -> None:
@@ -239,6 +238,10 @@ class Worker(WorkerBase):
238
# take current memory snapshot
239
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
240
self.requested_memory = request_memory(init_snapshot, self.cache_config)
241
+ logger.debug("worker init memory snapshot: %r", self.init_snapshot)
242
+ logger.debug(
243
+ "worker requested memory: %sGiB", format_gib(self.requested_memory)
244
+ )
245
else:
246
raise RuntimeError(f"Not support device type: {self.device_config.device}")
247
@@ -293,15 +296,14 @@ class Worker(WorkerBase):
296
You may limit the usage of GPU memory
297
by adjusting the `gpu_memory_utilization` parameter.
298
"""
296
- GiB = lambda b: b / GiB_bytes
299
if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:
300
# still need a profile run which compiles the model for
301
# max_num_batched_tokens
302
self.model_runner.profile_run()
303
304
msg = (
303
- f"Initial free memory {GiB(self.init_snapshot.free_memory):.2f} "
304
- f"GiB, reserved {GiB(kv_cache_memory_bytes):.2f} GiB memory for "
305
+ f"Initial free memory {format_gib(self.init_snapshot.free_memory)} "
306
+ f"GiB, reserved {format_gib(kv_cache_memory_bytes)} GiB memory for "
307
"KV Cache as specified by kv_cache_memory_bytes config and "
308
"skipped memory profiling. This does not respect the "
309
"gpu_memory_utilization config. Only use kv_cache_memory_bytes "
@@ -333,8 +335,8 @@ class Worker(WorkerBase):
335
# GPU did not change their memory usage during the profiling.
336
assert self.init_snapshot.free_memory > free_gpu_memory, (
337
"Error in memory profiling. "
336
- f"Initial free memory {GiB(self.init_snapshot.free_memory)} GiB, "
337
- f"current free memory {GiB(free_gpu_memory)} GiB. "
338
+ f"Initial free memory {format_gib(self.init_snapshot.free_memory)} GiB, "
339
+ f"current free memory {format_gib(free_gpu_memory)} GiB. "
340
"This happens when other processes sharing the same container "
341
"release GPU memory while vLLM is profiling during initialization. "
342
"To fix this, ensure consistent GPU memory allocation or "
@@ -346,21 +348,20 @@ class Worker(WorkerBase):
348
349
unrequested_memory = self.init_snapshot.free_memory - self.requested_memory
350
logger.debug(
349
- "Initial free memory: %.2f GiB; Requested memory: %.2f (util), %.2f GiB",
350
- GiB(self.init_snapshot.free_memory),
351
+ "Initial free memory: %f GiB; Requested memory: %f (util), %f GiB",
352
+ format_gib(self.init_snapshot.free_memory),
353
self.cache_config.gpu_memory_utilization,
352
- GiB(self.requested_memory),
354
+ format_gib(self.requested_memory),
355
)
356
logger.debug(
355
- "Free memory after profiling: %.2f GiB (total), "
356
- "%.2f GiB (within requested)",
357
- GiB(free_gpu_memory),
358
- GiB(free_gpu_memory - unrequested_memory),
357
+ "Free memory after profiling: %f GiB (total), %f GiB (within requested)",
358
+ format_gib(free_gpu_memory),
359
+ format_gib(free_gpu_memory - unrequested_memory),
360
)
361
logger.debug(profile_result)
362
logger.info_once(
362
- "Available KV cache memory: %.2f GiB",
363
- GiB(self.available_kv_cache_memory_bytes),
363
+ "Available KV cache memory: %f GiB",
364
+ format_gib(self.available_kv_cache_memory_bytes),
365
scope="local",
366
)
367
gc.collect()
@@ -467,7 +468,6 @@ class Worker(WorkerBase):
468
# CUDAGraph memory size and may not utilize all gpu memory.
469
# Users may want fine-grained control to specify kv cache
470
# memory size.
470
- GiB = lambda b: round(b / GiB_bytes, 2)
471
472
# empirically observed that the memory profiling may
473
# slightly underestimate the memory consumption.
@@ -492,24 +492,24 @@ class Worker(WorkerBase):
492
493
msg = (
494
f"Free memory on device "
495
- f"({GiB(self.init_snapshot.free_memory)}/"
496
- f"{GiB(self.init_snapshot.total_memory)} GiB) on startup. "
495
+ f"({format_gib(self.init_snapshot.free_memory)}/"
496
+ f"{format_gib(self.init_snapshot.total_memory)} GiB) on startup. "
497
f"Desired GPU memory utilization is "
498
f"({self.cache_config.gpu_memory_utilization}, "
499
- f"{GiB(self.requested_memory)} GiB). "
500
- f"Actual usage is {GiB(self.model_runner.model_memory_usage)} "
501
- f"GiB for weight, {GiB(self.peak_activation_memory)} GiB "
502
- f"for peak activation, {GiB(self.non_torch_memory)} GiB "
503
- f"for non-torch memory, and {GiB(cuda_graph_memory_bytes)} "
499
+ f"{format_gib(self.requested_memory)} GiB). "
500
+ f"Actual usage is {format_gib(self.model_runner.model_memory_usage)} "
501
+ f"GiB for weight, {format_gib(self.peak_activation_memory)} GiB "
502
+ f"for peak activation, {format_gib(self.non_torch_memory)} GiB "
503
+ f"for non-torch memory, and {format_gib(cuda_graph_memory_bytes)} "
504
f"GiB for CUDAGraph memory. Replace gpu_memory_utilization "
505
f"config with `--kv-cache-memory="
506
f"{kv_cache_memory_bytes_to_requested_limit}` "
507
- f"({GiB(kv_cache_memory_bytes_to_requested_limit)} GiB) to fit "
507
+ f"({format_gib(kv_cache_memory_bytes_to_requested_limit)} GiB) to fit "
508
f"into requested memory, or `--kv-cache-memory="
509
f"{kv_cache_memory_bytes_to_gpu_limit}` "
510
- f"({GiB(kv_cache_memory_bytes_to_gpu_limit)} GiB) to fully "
510
+ f"({format_gib(kv_cache_memory_bytes_to_gpu_limit)} GiB) to fully "
511
f"utilize gpu memory. Current kv cache memory in use is "
512
- f"{GiB(self.available_kv_cache_memory_bytes)} GiB."
512
+ f"{format_gib(self.available_kv_cache_memory_bytes)} GiB."
513
)
514
515
logger.debug(msg)
vllm/v1/worker/utils.py
+9
-8
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
from collections import defaultdict
5
from dataclasses import dataclass, field
6
@@ -15,8 +16,7 @@ from vllm.model_executor.models.utils import extract_layer_index
16
from vllm.multimodal.cache import processor_only_cache_from_config
17
from vllm.multimodal.registry import MultiModalRegistry
18
from vllm.platforms import current_platform
18
-from vllm.utils.mem_constants import GiB_bytes
19
-from vllm.utils.mem_utils import MemorySnapshot
19
+from vllm.utils.mem_utils import MemorySnapshot, format_gib
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
@@ -250,22 +250,23 @@ def gather_mm_placeholders(
250
return placeholders[is_embed]
251
252
253
-def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) -> float:
253
+def request_memory(init_snapshot: MemorySnapshot, cache_config: CacheConfig) -> int:
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
258
+ requested_memory = math.ceil(
259
+ init_snapshot.total_memory * cache_config.gpu_memory_utilization
260
+ )
261
262
if init_snapshot.free_memory < requested_memory:
261
- GiB = lambda b: round(b / GiB_bytes, 2)
263
raise ValueError(
264
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 "
265
+ f"({format_gib(init_snapshot.free_memory)}/"
266
+ f"{format_gib(init_snapshot.total_memory)} GiB) on startup "
267
f"is less than desired GPU memory utilization "
268
f"({cache_config.gpu_memory_utilization}, "
268
- f"{GiB(requested_memory)} GiB). Decrease GPU memory "
269
+ f"{format_gib(requested_memory)} GiB). Decrease GPU memory "
270
f"utilization or reduce GPU memory used by other processes."
271
)
272