Skip to content

GitLab

Unverified Commit 157caf51 authored by Paco Xu's avatar Paco Xu Committed by GitHub
Browse files

[Perf] avoid duplicate mem_get_info() call in get_current_memory_usage (#33064) 


Signed-off-by: default avatarPaco Xu <paco.xu@daocloud.io>
parent 0b53bec6
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
vllm/platforms/rocm.py
View file @ 157caf51
...@@ -561,7 +561,8 @@ class RocmPlatform(Platform): ...@@ -561,7 +561,8 @@ class RocmPlatform(Platform):
cls, device: torch.types.Device | None = None cls, device: torch.types.Device | None = None
) -> float: ) -> float:
torch.cuda.reset_peak_memory_stats(device) torch.cuda.reset_peak_memory_stats(device)
return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(device)[0] free_mem, total_mem = torch.cuda.mem_get_info(device)
return total_mem - free_mem
@classmethod @classmethod
def get_device_communicator_cls(cls) -> str: def get_device_communicator_cls(cls) -> str:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment