[ROCm] Using a more precise memory profiling (#12624)

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
2025-02-11 08:47:10 -05:00 · 2025-02-11 08:47:10 -05:00 · 7539bbc6a6
commit 7539bbc6a6
parent 9cf4759493
1 changed files with 2 additions and 1 deletions
--- 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]