update GpuMemoryMonitor to DeviceMemoryMonitor for all HW (#1526)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Signed-off-by: Chendi.Xue <chendi.xue@intel.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

tests/utils.py

@@ -20,6 +20,8 @@ from typing_extensions import ParamSpec
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import cuda_device_count_stateless
 
+from vllm_omni.platforms import current_omni_platform
+
 _P = ParamSpec("_P")
 
 if current_platform.is_rocm():
@@ -504,8 +506,8 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
     return wrapper
 
 
-class GPUMemoryMonitor:
-    """Poll global device memory usage via CUDA APIs."""
+class DeviceMemoryMonitor:
+    """Poll global device memory usage."""
 
     def __init__(self, device_index: int, interval: float = 0.05):
         self.device_index = device_index
@@ -518,8 +520,8 @@ class GPUMemoryMonitor:
         def monitor_loop() -> None:
             while not self._stop_event.is_set():
                 try:
-                    with torch.cuda.device(self.device_index):
-                        free_bytes, total_bytes = torch.cuda.mem_get_info()
+                    with current_omni_platform.device(self.device_index):
+                        free_bytes, total_bytes = current_omni_platform.mem_get_info()
                     used_mb = (total_bytes - free_bytes) / (1024**2)
                     self._peak_used_mb = max(self._peak_used_mb, used_mb)
                 except Exception:
@@ -537,8 +539,8 @@ class GPUMemoryMonitor:
 
     @property
     def peak_used_mb(self) -> float:
-        fallback_alloc = torch.cuda.max_memory_allocated(device=self.device_index) / (1024**2)
-        fallback_reserved = torch.cuda.max_memory_reserved(device=self.device_index) / (1024**2)
+        fallback_alloc = current_omni_platform.max_memory_allocated(device=self.device_index) / (1024**2)
+        fallback_reserved = current_omni_platform.max_memory_reserved(device=self.device_index) / (1024**2)
         return max(self._peak_used_mb, fallback_alloc, fallback_reserved)
 
     def __del__(self):