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, 04:00 UTC
[Misc] Make mem utils can be reused by other platforms (#32322)
Signed-off-by: shen-shanshan <467638484@qq.com>
Shanshan Shen committed
Jan 14, 2026 at 19:46 UTC
ce0946249d28f263930f2789186e49db242d1834
3 files changed
+21
-25
vllm/platforms/interface.py
+12
-8
@@ -589,14 +589,18 @@ class Platform:
589
def __getattr__(self, key: str):
590
device = getattr(torch, self.device_type, None)
591
if device is not None and hasattr(device, key):
592
- return getattr(device, key)
593
- else:
594
- logger.warning(
595
- "Current platform %s does not have '%s' attribute.",
596
- self.device_type,
597
- key,
598
- )
599
- return None
592
+ attr = getattr(device, key)
593
+ # NOTE: `hasattr(device, key)=True` can only avoid AttributeError,
594
+ # but the value of this attr could be `None`.
595
+ if attr is not None:
596
+ return attr
597
+
598
+ logger.warning(
599
+ "Current platform %s does not have '%s' attribute.",
600
+ self.device_type,
601
+ key,
602
+ )
603
+ return None
604
605
def get_global_graph_pool(self) -> Any:
606
"""
vllm/utils/mem_utils.py
+9
-13
@@ -11,6 +11,8 @@ import psutil
11
import torch
12
import torch.types
13
14
+from vllm.platforms import current_platform
15
+
16
from .mem_constants import GiB_bytes, MiB_bytes
17
18
@@ -45,8 +47,6 @@ class DeviceMemoryProfiler:
47
48
def current_memory_usage(self) -> float:
49
# Return the memory usage in bytes.
48
- from vllm.platforms import current_platform
49
-
50
gc.collect()
51
return current_platform.get_current_memory_usage(self.device)
52
@@ -80,8 +80,6 @@ class MemorySnapshot:
80
81
def __post_init__(self) -> None:
82
if self.device is None:
83
- from vllm.platforms import current_platform
84
-
83
device_fn = current_platform.current_device
84
assert device_fn is not None
85
self.device_ = torch.device(device_fn())
@@ -92,8 +90,6 @@ class MemorySnapshot:
90
self.measure()
91
92
def measure(self) -> None:
95
- from vllm.platforms import current_platform
96
-
93
device = self.device_
94
95
# we measure the torch peak memory usage via allocated_bytes,
@@ -101,11 +97,11 @@ class MemorySnapshot:
97
# After `torch.cuda.reset_peak_memory_stats()`,
98
# `torch.cuda.memory_reserved()` will keep growing, and only shrink
99
# when we call `torch.cuda.empty_cache()` or OOM happens.
104
- self.torch_peak = torch.cuda.memory_stats(device).get(
100
+ self.torch_peak = current_platform.memory_stats(device).get(
101
"allocated_bytes.all.peak", 0
102
)
103
108
- self.free_memory, self.total_memory = torch.cuda.mem_get_info(device)
104
+ self.free_memory, self.total_memory = current_platform.mem_get_info(device)
105
shared_sysmem_device_mem_sms = ((8, 7), (11, 0), (12, 1)) # Orin, Thor, Spark
106
if (
107
current_platform.is_cuda()
@@ -130,7 +126,7 @@ class MemorySnapshot:
126
# torch.cuda.memory_reserved() is how many bytes
127
# PyTorch gets from cuda (by calling cudaMalloc, etc.)
128
# this is used to measure the non-torch memory usage
133
- self.torch_memory = torch.cuda.memory_reserved(device)
129
+ self.torch_memory = current_platform.memory_reserved(device)
130
131
self.non_torch_memory = self.cuda_memory - self.torch_memory
132
self.timestamp = time.time()
@@ -159,7 +155,7 @@ class MemorySnapshot:
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, "
162
- f"cuda_memory={format_gib(self.cuda_memory)}GiB, "
158
+ f"{current_platform.device_name}_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}, "
@@ -254,8 +250,8 @@ def memory_profiling(
250
until after profiling to get (c.).
251
"""
252
gc.collect()
257
- torch.cuda.empty_cache()
258
- torch.cuda.reset_peak_memory_stats(baseline_snapshot.device_)
253
+ current_platform.empty_cache()
254
+ current_platform.reset_peak_memory_stats(baseline_snapshot.device_)
255
256
result = MemoryProfilingResult(
257
before_create=baseline_snapshot,
@@ -268,7 +264,7 @@ def memory_profiling(
264
yield result
265
266
gc.collect()
271
- torch.cuda.empty_cache()
267
+ current_platform.empty_cache()
268
269
result.after_profile.measure()
270
vllm/v1/worker/gpu_worker.py
-4
@@ -312,9 +312,6 @@ class Worker(WorkerBase):
312
logger.info(msg)
313
return kv_cache_memory_bytes
314
315
- torch.cuda.empty_cache()
316
- torch.cuda.reset_peak_memory_stats()
317
-
315
# Execute a forward pass with dummy inputs to profile the memory usage
316
# of the model.
317
with memory_profiling(
@@ -360,7 +357,6 @@ class Worker(WorkerBase):
357
format_gib(self.available_kv_cache_memory_bytes),
358
scope="local",
359
)
363
- gc.collect()
360
361
return int(self.available_kv_cache_memory_bytes)
362