diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index e46d7ec96..bf36b78ff 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/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: diff --git a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py index 8791bb02d..998e6232e 100644 --- a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py +++ b/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() diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index f4fd78c7d..7987da887 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/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 diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 30edc8256..945cfc413 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/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 diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index a04224063..f90c69e3a 100644 --- a/vllm_omni/platforms/interface.py +++ b/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 diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 63efcdd18..a77aa83a4 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/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