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/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -5,7 +5,7 @@ import pytest
 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
 
@@ -28,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
     current_omni_platform.empty_cache()
-    device_index = torch.cuda.current_device()
-    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+    device_index = current_omni_platform.current_device()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
 
     m = Omni(
@@ -44,7 +42,7 @@ def run_inference(
         flow_shift=5.0,
     )
 
-    torch.cuda.reset_peak_memory_stats(device=device_index)
+    current_omni_platform.reset_peak_memory_stats()
 
     # Refer to tests/e2e/offline_inference/test_t2v_model.py
     # Use minimal settings for testing
@@ -57,7 +55,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,