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 6, 17:00 UTC
[Hardware] Replace memory related torch.cuda APIs (#37031)
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
Kunshang Ji committed
Mar 16, 2026 at 18:24 UTC
747b0681364aa53235b71a30488f450652cc316a
8 files changed
+20
-17
benchmarks/attention_benchmarks/runner.py
+2
-2
@@ -418,8 +418,8 @@ def _run_single_benchmark(
418
mem_stats = {}
419
if config.profile_memory:
420
mem_stats = {
421
- "allocated_mb": torch.cuda.memory_allocated(device) / 1024**2,
422
- "reserved_mb": torch.cuda.memory_reserved(device) / 1024**2,
421
+ "allocated_mb": torch.accelerator.memory_allocated(device) / 1024**2,
422
+ "reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
423
}
424
425
return times, mem_stats
benchmarks/benchmark_topk_topp.py
+5
-2
@@ -95,13 +95,16 @@ def create_logits(
95
def measure_memory() -> tuple[int, int]:
96
"""Return (allocated, reserved) memory in bytes."""
97
torch.accelerator.synchronize()
98
- return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
98
+ return (
99
+ torch.accelerator.memory_allocated(),
100
+ torch.accelerator.max_memory_allocated(),
101
+ )
102
103
104
def reset_memory_stats():
105
"""Reset peak memory statistics."""
106
reset_buffer_cache()
104
- torch.cuda.reset_peak_memory_stats()
107
+ torch.accelerator.reset_peak_memory_stats()
108
torch.accelerator.empty_cache()
109
gc.collect()
110
tests/test_regression.py
+1
-1
@@ -55,7 +55,7 @@ def test_gc():
55
# The memory allocated for model and KV cache should be released.
56
# The memory allocated for PyTorch and others should be less than 50MB.
57
# Usually, it's around 10MB.
58
- allocated = torch.cuda.memory_allocated()
58
+ allocated = torch.accelerator.memory_allocated()
59
assert allocated < 50 * 1024 * 1024
60
61
tests/utils_/test_mem_utils.py
+1
-1
@@ -29,7 +29,7 @@ def test_memory_profiling():
29
def measure_current_non_torch():
30
free, total = torch.cuda.mem_get_info()
31
current_used = total - free
32
- current_torch = torch.cuda.memory_reserved()
32
+ current_torch = torch.accelerator.memory_reserved()
33
current_non_torch = current_used - current_torch
34
return current_non_torch
35
tools/pre_commit/check_torch_cuda.py
+1
-1
@@ -8,7 +8,7 @@ import regex as re
8
# Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
9
# --------------------------------------------------------------------------- #
10
_TORCH_CUDA_PATTERNS = [
11
- r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|set_device|device\()\b",
11
+ r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|set_device|device\()\b",
12
r"\bwith\storch\.cuda\.device\b",
13
]
14
vllm/model_executor/model_loader/base_loader.py
+1
-1
@@ -64,7 +64,7 @@ class BaseModelLoader(ABC):
64
# Log peak GPU memory after loading weights. This is needed
65
# to have test coverage on peak memory for online quantization.
66
if current_platform.is_cuda():
67
- peak_memory = torch.cuda.max_memory_allocated()
67
+ peak_memory = torch.accelerator.max_memory_allocated()
68
logger.debug_once(
69
"Peak GPU memory after loading weights: %s GiB",
70
format_gib(peak_memory),
vllm/utils/mem_utils.py
+8
-8
@@ -93,11 +93,11 @@ class MemorySnapshot:
93
device = self.device_
94
95
# we measure the torch peak memory usage via allocated_bytes,
96
- # rather than `torch.cuda.memory_reserved()` .
97
- # After `torch.cuda.reset_peak_memory_stats()`,
98
- # `torch.cuda.memory_reserved()` will keep growing, and only shrink
96
+ # rather than `torch.accelerator.memory_reserved()` .
97
+ # After `torch.accelerator.reset_peak_memory_stats()`,
98
+ # `torch.accelerator.memory_reserved()` will keep growing, and only shrink
99
# when we call `torch.accelerator.empty_cache()` or OOM happens.
100
- self.torch_peak = current_platform.memory_stats(device).get(
100
+ self.torch_peak = torch.accelerator.memory_stats(device).get(
101
"allocated_bytes.all.peak", 0
102
)
103
@@ -123,10 +123,10 @@ class MemorySnapshot:
123
124
self.cuda_memory = self.total_memory - self.free_memory
125
126
- # torch.cuda.memory_reserved() is how many bytes
126
+ # torch.accelerator.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
129
- self.torch_memory = current_platform.memory_reserved(device)
129
+ self.torch_memory = torch.accelerator.memory_reserved(device)
130
131
self.non_torch_memory = self.cuda_memory - self.torch_memory
132
self.timestamp = time.time()
@@ -243,7 +243,7 @@ def memory_profiling(
243
The memory used for loading weights (a.) is directly given from the
244
argument `weights_memory`.
245
246
- The increase of `torch.cuda.memory_stats()["allocated_bytes.all.peak"]`
246
+ The increase of `torch.accelerator.memory_stats()["allocated_bytes.all.peak"]`
247
during profiling gives (b.).
248
249
The increase of `non_torch_memory` from creating the current vLLM instance
@@ -251,7 +251,7 @@ def memory_profiling(
251
"""
252
gc.collect()
253
torch.accelerator.empty_cache()
254
- current_platform.reset_peak_memory_stats(baseline_snapshot.device_)
254
+ torch.accelerator.reset_peak_memory_stats(baseline_snapshot.device_)
255
256
result = MemoryProfilingResult(
257
before_create=baseline_snapshot,
vllm/v1/worker/gpu_worker.py
+1
-1
@@ -387,7 +387,7 @@ class Worker(WorkerBase):
387
) as profile_result:
388
self.model_runner.profile_run()
389
390
- profile_torch_peak = current_platform.memory_stats(self.device).get(
390
+ profile_torch_peak = torch.accelerator.memory_stats(self.device).get(
391
"allocated_bytes.all.peak", 0
392
)
393