[Platform] Add supports_torch_inductor interface (#1108)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

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)

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:

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

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:

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"

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"