iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Platform] Replace torch.cuda.mem_get_info with torch.accelerator.get_memory_info (#44825)

Browse Source 
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
This commit is contained in:
Kunshang Ji
2026-06-30 14:39:52 +08:00
committed by GitHub
parent 0feca7ffa8
commit fb42e5219e
17 changed files with 56 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/pre_commit/check_torch_cuda.py
+2 -1
View File
@@ -8,11 +8,12 @@ import regex as re
# Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
# --------------------------------------------------------------------------- #
_TORCH_CUDA_PATTERNS = [
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",
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|mem_get_info|set_device|device\()\b",
r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b",
r"\bwith\storch\.cuda\.device\b",
# Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally
r"\bcuda_device_count_stateless\(\)\b",
r"\bcurrent_platform\.mem_get_info\(\)\b",
]
ALLOWED_FILES = {