Skip to content

GitLab

"...git@developer.sourcefind.cn:2222/OpenDAS/vllm_cscc.git" did not exist on "53fa45739176e81b9078f7e3516a8b11fda161de"
Commit d9ef7ce7 authored by zhuwenwen's avatar zhuwenwen
Browse files

update rocm.py

parent 9e053941
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 3 deletions
vllm/platforms/rocm.py
View file @ d9ef7ce7
...@@ -225,8 +225,9 @@ class RocmPlatform(Platform): ...@@ -225,8 +225,9 @@ class RocmPlatform(Platform):
device: Optional[torch.types.Device] = None device: Optional[torch.types.Device] = 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( # return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(
device)[0] # device)[0]
return torch.cuda.max_memory_allocated(device)
@classmethod @classmethod
def get_device_communicator_cls(cls) -> str: def get_device_communicator_cls(cls) -> str:
...@@ -247,4 +248,4 @@ class RocmPlatform(Platform): ...@@ -247,4 +248,4 @@ class RocmPlatform(Platform):
if cls.is_fp8_fnuz(): if cls.is_fp8_fnuz():
return torch.float8_e4m3fnuz return torch.float8_e4m3fnuz
else: else:
return torch.float8_e4m3fn return torch.float8_e4m3fn
\ No newline at end of file
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