From 7539bbc6a6715dc8e5e71730e2377219b0e69e21 Mon Sep 17 00:00:00 2001 From: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com> Date: Tue, 11 Feb 2025 08:47:10 -0500 Subject: [PATCH] [ROCm] Using a more precise memory profiling (#12624) Signed-off-by: Gregory Shtrasberg --- vllm/platforms/rocm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 1f690b7111ee2..13aebc605af74 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -169,4 +169,5 @@ class RocmPlatform(Platform): device: Optional[torch.types.Device] = None ) -> float: torch.cuda.reset_peak_memory_stats(device) - return torch.cuda.max_memory_allocated(device) + return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info( + device)[0]