Skip to content

Commit 105bf03

Browse files
gshtrasshreyankg
authored andcommitted
[ROCm] Using a more precise memory profiling (vllm-project#12624)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
1 parent 1f618c3 commit 105bf03

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

vllm/platforms/rocm.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -169,4 +169,5 @@ def get_current_memory_usage(cls,
169169
device: Optional[torch.types.Device] = None
170170
) -> float:
171171
torch.cuda.reset_peak_memory_stats(device)
172-
return torch.cuda.max_memory_allocated(device)
172+
return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(
173+
device)[0]

0 commit comments

Comments
 (0)