[XPU][NPU][ROCM] enable cpu_offloading flag for non_cuda (#1488)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -48,7 +48,6 @@ def inference(model_name: str, offload: bool = True):
 @pytest.mark.core_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"})
-@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
 @pytest.mark.parametrize("model_name", models)
 def test_cpu_offload_diffusion_model(model_name: str):
     try:

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -7,6 +7,7 @@ from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
 from tests.utils import GPUMemoryMonitor
 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]
@@ -29,7 +30,7 @@ def run_inference(
 ) -> 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()
+    current_omni_platform.empty_cache()
     device_index = torch.cuda.current_device()
     monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()

vllm_omni/diffusion/offloader/__init__.py

@@ -55,8 +55,11 @@ def get_offload_backend(
         return None
 
     # Validate platform (CUDA required for now)
-    if not current_omni_platform.is_cuda() or current_omni_platform.get_device_count() < 1:
-        logger.warning("CPU offloading requires CUDA devices. Skipping offloading.")
+    if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1:
+        logger.warning(
+            "Current device: %s does not support CPU offloading. Skipping offloading.",
+            current_omni_platform.get_device_name(),
+        )
         return None
 
     # Detect device if not provided

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -76,9 +76,11 @@ class SequentialOffloadHook(ModelHook):
         current_omni_platform.synchronize()
 
         logger.debug(
-            "Swapped: %s -> CPU, %s -> GPU",
+            "Swapped: %s -> CPU, %s -> %s, free memory: %.4f GB",
             [t.__class__.__name__ for t in self.offload_targets],
             module.__class__.__name__,
+            f"{self.device.type}:{self.device.index}",
+            current_omni_platform.get_free_memory() / 1024 / 1024 / 1024,
         )
 
         return args, kwargs

vllm_omni/platforms/interface.py

@@ -98,6 +98,10 @@ class OmniPlatform(Platform):
     def get_free_memory(cls, device: torch.device | None = None) -> int:
         raise NotImplementedError
 
+    @classmethod
+    def supports_cpu_offload(cls) -> bool:
+        return True
+
 
 class UnspecifiedOmniPlatform(OmniPlatform):
     _omni_enum = OmniPlatformEnum.UNSPECIFIED

vllm_omni/platforms/xpu/platform.py

@@ -73,9 +73,5 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
 
     @classmethod
     def get_free_memory(cls, device: torch.device | None = None) -> int:
-        if device is None:
-            device_id = 0
-        else:
-            device_id = device.index if device.index is not None else 0
-        props = torch.xpu.get_device_properties(device_id)
-        return props.total_memory
+        free, _ = torch.xpu.mem_get_info(device)
+        return free