enable DeviceMemoryMonitor for all platforms

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

Author: xuechendi

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -5,7 +5,7 @@
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.utils import GPUMemoryMonitor, hardware_test
+from tests.utils import DeviceMemoryMonitor, hardware_test
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
@@ -21,11 +21,11 @@
 
 def inference(model_name: str, offload: bool = True):
     current_omni_platform.empty_cache()
-    device_index = torch.cuda.current_device()
-    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+    device_index = torch.accelerator.current_device_index()
+    monitor = DeviceMemoryMonitor.instantiate(device_index=device_index, interval=0.02)
     monitor.start()
     m = Omni(model=model_name, enable_cpu_offload=offload)
-    torch.cuda.reset_peak_memory_stats(device=device_index)
+    torch.accelerator.reset_peak_memory_stats()
     height = 256
     width = 256
 
@@ -36,7 +36,7 @@ def inference(model_name: str, offload: bool = True):
             width=width,
             num_inference_steps=9,
             guidance_scale=0.0,
-            generator=torch.Generator("cuda").manual_seed(42),
+            generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
         ),
     )
     peak = monitor.peak_used_mb

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -5,8 +5,9 @@
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.utils import GPUMemoryMonitor
+from tests.utils import DeviceMemoryMonitor
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -27,11 +28,9 @@ def run_inference(
     layerwise_offload: bool = False,
     num_inference_steps: int = 3,
 ) -> float:
-    # For now, only support on GPU, so apply torch.cuda operations here
-    # NPU / ROCm platforms are expected to be detected and skipped this test function
-    torch.cuda.empty_cache()
-    device_index = torch.cuda.current_device()
-    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+    current_omni_platform.empty_cache()
+    device_index = torch.accelerator.current_device_index()
+    monitor = DeviceMemoryMonitor.instantiate(device_index=device_index, interval=0.02)
     monitor.start()
 
     m = Omni(
@@ -41,7 +40,7 @@ def run_inference(
         flow_shift=5.0,
     )
 
-    torch.cuda.reset_peak_memory_stats(device=device_index)
+    torch.accelerator.reset_peak_memory_stats()
 
     # Refer to tests/e2e/offline_inference/test_t2v_model.py
     # Use minimal settings for testing
@@ -54,7 +53,7 @@ def run_inference(
         OmniDiffusionSamplingParams(
             height=height,
             width=width,
-            generator=torch.Generator("cuda").manual_seed(42),
+            generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
             guidance_scale=1.0,
             num_inference_steps=num_inference_steps,
             num_frames=num_frames,

tests/e2e/offline_inference/test_zimage_parallelism.py

@@ -22,7 +22,7 @@
 from PIL import Image
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.utils import GPUMemoryMonitor, hardware_test
+from tests.utils import DeviceMemoryMonitor, hardware_test
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -95,7 +95,7 @@ def _run_zimage_generate(
 
     torch.cuda.empty_cache()
     device_index = torch.cuda.current_device()
-    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor = DeviceMemoryMonitor.instantiate(device_index=device_index, interval=0.02)
     monitor.start()
     m = Omni(
         model=_get_zimage_model(),

tests/utils.py

@@ -20,6 +20,8 @@
 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,17 @@ def wrapper(f: Callable[_P, None]) -> Callable[_P, None]:
     return wrapper
 
 
-class GPUMemoryMonitor:
-    """Poll global device memory usage via CUDA APIs."""
+class DeviceMemoryMonitor:
+    """Poll global device memory usage."""
+
+    @classmethod
+    def instantiate(cls, **kwargs: Any) -> "DeviceMemoryMonitor":
+        if current_omni_platform.is_npu():
+            return NPUMemoryMonitor(**kwargs)
+        elif current_omni_platform.is_xpu():
+            return XPUMemoryMonitor(**kwargs)
+        else:
+            return cls(**kwargs)
 
     def __init__(self, device_index: int, interval: float = 0.05):
         self.device_index = device_index
@@ -543,3 +554,49 @@ def peak_used_mb(self) -> float:
 
     def __del__(self):
         self.stop()
+
+
+class NPUMemoryMonitor(DeviceMemoryMonitor):
+    def start(self) -> None:
+        def monitor_loop() -> None:
+            while not self._stop_event.is_set():
+                try:
+                    with torch.npu.device(self.device_index):
+                        free_bytes, total_bytes = torch.npu.mem_get_info()
+                    used_mb = (total_bytes - free_bytes) / (1024**2)
+                    self._peak_used_mb = max(self._peak_used_mb, used_mb)
+                except Exception:
+                    pass
+                time.sleep(self.interval)
+
+        self._thread = threading.Thread(target=monitor_loop, daemon=False)
+        self._thread.start()
+
+    @property
+    def peak_used_mb(self) -> float:
+        fallback_alloc = torch.npu.max_memory_allocated(device=self.device_index) / (1024**2)
+        fallback_reserved = torch.npu.max_memory_reserved(device=self.device_index) / (1024**2)
+        return max(self._peak_used_mb, fallback_alloc, fallback_reserved)
+
+
+class XPUMemoryMonitor(DeviceMemoryMonitor):
+    def start(self) -> None:
+        def monitor_loop() -> None:
+            while not self._stop_event.is_set():
+                try:
+                    with torch.xpu.device(self.device_index):
+                        free_bytes, total_bytes = torch.xpu.mem_get_info()
+                    used_mb = (total_bytes - free_bytes) / (1024**2)
+                    self._peak_used_mb = max(self._peak_used_mb, used_mb)
+                except Exception:
+                    pass
+                time.sleep(self.interval)
+
+        self._thread = threading.Thread(target=monitor_loop, daemon=False)
+        self._thread.start()
+
+    @property
+    def peak_used_mb(self) -> float:
+        fallback_alloc = torch.xpu.max_memory_allocated(device=self.device_index) / (1024**2)
+        fallback_reserved = torch.xpu.max_memory_reserved(device=self.device_index) / (1024**2)
+        return max(self._peak_used_mb, fallback_alloc, fallback_reserved)