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
[Chore] Cleanup `mem_utils.py` (#31793)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Jan 6, 2026 at 19:55 UTC
14df02b4e15891c00a24b7fcf7d62076079c38d5
1 file changed
+27
-17
vllm/utils/mem_utils.py
+27
-17
@@ -22,7 +22,7 @@ def get_max_shared_memory_bytes(gpu: int = 0) -> int:
22
max_shared_mem = ops.get_max_shared_memory_per_block_device_attribute(gpu)
23
# value 0 will cause MAX_SEQ_LEN become negative and test_attention.py
24
# will fail
25
- assert max_shared_mem > 0, "max_shared_mem can not be zero"
25
+ assert max_shared_mem > 0, "max_shared_mem cannot be zero"
26
return int(max_shared_mem)
27
28
@@ -154,12 +154,16 @@ class MemoryProfilingResult:
154
non_kv_cache_memory: int = 0
155
torch_peak_increase: int = 0
156
non_torch_increase: int = 0
157
- weights_memory: float = 0
157
+ weights_memory: int = 0
158
before_create: MemorySnapshot = field(default_factory=MemorySnapshot)
159
- before_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
160
- after_profile: MemorySnapshot = field(default_factory=MemorySnapshot)
159
profile_time: float = 0.0
160
161
+ def __post_init__(self) -> None:
162
+ device = self.before_create.device_
163
+
164
+ self.before_profile = MemorySnapshot(device=device, auto_measure=False)
165
+ self.after_profile = MemorySnapshot(device=device, auto_measure=False)
166
+
167
def __repr__(self) -> str:
168
return (
169
f"Memory profiling takes {self.profile_time:.2f} seconds. "
@@ -175,9 +179,12 @@ class MemoryProfilingResult:
179
180
@contextlib.contextmanager
181
def memory_profiling(
178
- baseline_snapshot: MemorySnapshot, weights_memory: int
182
+ baseline_snapshot: MemorySnapshot,
183
+ weights_memory: int = 0,
184
) -> Generator[MemoryProfilingResult, None, None]:
180
- """Memory profiling context manager.
185
+ """
186
+ Memory profiling context manager.
187
+
188
baseline_snapshot: the memory snapshot before the current vLLM instance.
189
weights_memory: memory used by PyTorch when loading the model weights.
190
Note that, before loading the model weights, we also initialize the device
@@ -217,21 +224,24 @@ def memory_profiling(
224
b. 2 GiB reserved for the peak activation tensors (category 2)
225
c. 1 GiB used by non-torch components (category 3)
226
220
- The memory used for loading weights (a.) is directly given from the argument `weights_memory`.
227
+ The memory used for loading weights (a.) is directly given from the
228
+ argument `weights_memory`.
229
222
- The increase of `torch.cuda.memory_stats()["allocated_bytes.all.peak"]` during profiling gives (b.).
230
+ The increase of `torch.cuda.memory_stats()["allocated_bytes.all.peak"]`
231
+ during profiling gives (b.).
232
224
- The increase of `non_torch_memory` from creating the current vLLM instance until after profiling to get (c.).
225
- """ # noqa
233
+ The increase of `non_torch_memory` from creating the current vLLM instance
234
+ until after profiling to get (c.).
235
+ """
236
gc.collect()
237
torch.cuda.empty_cache()
228
- torch.cuda.reset_peak_memory_stats()
229
-
230
- result = MemoryProfilingResult()
238
+ torch.cuda.reset_peak_memory_stats(baseline_snapshot.device_)
239
232
- result.before_create = baseline_snapshot
233
- # the part of memory used for holding the model weights
234
- result.weights_memory = weights_memory
240
+ result = MemoryProfilingResult(
241
+ before_create=baseline_snapshot,
242
+ # the part of memory used for holding the model weights
243
+ weights_memory=weights_memory,
244
+ )
245
246
result.before_profile.measure()
247
@@ -252,4 +262,4 @@ def memory_profiling(
262
peak_activation_memory = result.torch_peak_increase
263
result.non_kv_cache_memory = (
264
non_torch_memory + peak_activation_memory + result.weights_memory
255
- ) # noqa
265
+ )