biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Hardware] Replace memory related torch.cuda APIs (#37031)

Browse Source 
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
This commit is contained in:
Kunshang Ji
2026-03-16 18:24:48 +08:00
committed by GitHub
parent 122f75d939
commit 747b068136
8 changed files with 20 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
benchmarks/attention_benchmarks/runner.py
View File

@@ -418,8 +418,8 @@ def _run_single_benchmark(
mem_stats = {}
if config.profile_memory:
mem_stats = {
"allocated_mb": torch.cuda.memory_allocated(device) / 1024**2,
"reserved_mb": torch.cuda.memory_reserved(device) / 1024**2,
"allocated_mb": torch.accelerator.memory_allocated(device) / 1024**2,
"reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
}
return times, mem_stats