diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index b0bec0852..38f438062 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -29,6 +29,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.offload import apply_offload_hooks from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -112,14 +113,20 @@ class DiffusionModelRunner: # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: - try: - self.pipeline.transformer = regionally_compile( - self.pipeline.transformer, - dynamic=True, + if current_omni_platform.supports_torch_inductor(): + try: + self.pipeline.transformer = regionally_compile( + self.pipeline.transformer, + dynamic=True, + ) + logger.info("Model runner: Model compiled with torch.compile.") + except Exception as e: + logger.warning(f"Model runner: torch.compile failed with error: {e}. Using eager mode.") + else: + logger.warning( + "Model runner: Platform %s does not support torch inductor, skipping torch.compile.", + current_omni_platform.get_torch_device(), ) - logger.info("Model runner: Model compiled with torch.compile.") - except Exception as e: - logger.warning(f"Model runner: torch.compile failed with error: {e}. Using eager mode.") # Setup cache backend self.cache_backend = get_cache_backend(self.od_config.cache_backend, self.od_config.cache_config) diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 89e403f99..6bf740a01 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -80,6 +80,10 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): logger.info("Defaulting to diffusion attention backend SDPA") return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + @classmethod + def supports_torch_inductor(cls) -> bool: + return True + @classmethod def get_torch_device(cls, local_rank: int | None = None) -> torch.device: if local_rank is None: diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index db541ebce..a04224063 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -73,6 +73,11 @@ class OmniPlatform(Platform): """ raise NotImplementedError + @classmethod + def supports_torch_inductor(cls) -> bool: + """Check if the platform supports torch.compile with inductor backend.""" + raise NotImplementedError + @classmethod def get_torch_device(cls, local_rank: int | None = None) -> torch.device: raise NotImplementedError diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 429297bbd..3c2495c3d 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -55,6 +55,10 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): logger.info("Falling back to diffusion attention backend SDPA") return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + @classmethod + def supports_torch_inductor(cls) -> bool: + return False + @classmethod def get_torch_device(cls, local_rank: int | None = None) -> torch.device: if local_rank is None: diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index f2c2d3a85..14534f3a1 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -66,6 +66,10 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): logger.info("Defaulting to diffusion attention backend SDPA") return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + @classmethod + def supports_torch_inductor(cls) -> bool: + return True + @classmethod def get_default_stage_config_path(cls) -> str: return "vllm_omni/platforms/rocm/stage_configs" diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 5c126c6d9..43d635702 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -43,6 +43,10 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): logger.info("Defaulting to diffusion attention backend SDPA") return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + @classmethod + def supports_torch_inductor(cls) -> bool: + return True + @classmethod def get_default_stage_config_path(cls) -> str: return "vllm_omni/platforms/xpu/stage_configs"