[platform] Refactor current_memory_usage() function in DeviceMemoryProfiler into Platform

vllm/platforms/cuda.py

@@ -140,6 +140,11 @@
         cache_config = vllm_config.cache_config
         if cache_config and cache_config.block_size is None:
             cache_config.block_size = 16
+
+    @classmethod
+    def get_current_memory_usage(cls, device: Optional[torch.types.Device] = None) -> float:
+        torch.cuda.reset_peak_memory_stats(device)
+        return torch.cuda.max_memory_allocated(device)
 
 
 # NVML utils

vllm/platforms/interface.py

@@ -237,6 +237,13 @@
                            "This may slow down the performance.")
             return False
         return True
+
+    @classmethod
+    def get_current_memory_usage(cls, device: Optional[torch.types.Device] = None) -> float:
+        """
+        Return the memory usage in bytes.
+        """
+        raise NotImplementedError
 
 
 class UnspecifiedPlatform(Platform):

vllm/platforms/xpu.py

@@ -87,3 +87,8 @@
     def is_pin_memory_available(cls):
         logger.warning("Pin memory is not supported on XPU.")
         return False
+
+    @classmethod
+    def get_current_memory_usage(cls, device: Optional[torch.types.Device] = None) -> float:
+        torch.xpu.reset_peak_memory_stats(device)
+        return torch.xpu.max_memory_allocated(device)

vllm/utils.py

@@ -680,23 +680,13 @@
     def __init__(self, device: Optional[torch.types.Device] = None):
         self.device = device
 
-    def current_memory_usage(self) -> float:
-        # Return the memory usage in bytes.
-        if current_platform.is_cuda_alike():
-            torch.cuda.reset_peak_memory_stats(self.device)
-            mem = torch.cuda.max_memory_allocated(self.device)
-        elif current_platform.is_xpu():
-            torch.xpu.reset_peak_memory_stats(self.device)  # type: ignore
-            mem = torch.xpu.max_memory_allocated(self.device)  # type: ignore
-        return mem
-
     def __enter__(self):
-        self.initial_memory = self.current_memory_usage()
+        self.initial_memory = current_platform.get_current_memory_usage(self.device)
         # This allows us to call methods of the context manager if needed
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.final_memory = self.current_memory_usage()
+        self.final_memory = current_platform.get_current_memory_usage(self.device)
         self.consumed_memory = self.final_memory - self.initial_memory
 
         # Force garbage collection