[Perf][Diffusion] Compile and prewarm VAE decode

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/compile.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
 from typing import Any
 
 import torch
@@ -9,6 +10,59 @@ from vllm.logger import init_logger
 logger = init_logger(__name__)
 
 
+class _FallbackCompiledCallable:
+    """Run a compiled callable until it fails, then permanently use eager."""
+
+    def __init__(
+        self,
+        eager: Callable[..., Any],
+        compiled: Callable[..., Any],
+        *,
+        label: str,
+    ) -> None:
+        self._eager = eager
+        self._compiled: Callable[..., Any] | None = compiled
+        self._label = label
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        compiled = self._compiled
+        if compiled is None:
+            return self._eager(*args, **kwargs)
+        try:
+            return compiled(*args, **kwargs)
+        except Exception as exc:
+            self._compiled = None
+            logger.warning(
+                "%s failed during lazy torch.compile execution (%s); disabling compilation and retrying eagerly.",
+                self._label,
+                exc,
+            )
+            return self._eager(*args, **kwargs)
+
+
+def compile_callable_with_fallback(
+    fn: Callable[..., Any],
+    *,
+    label: str,
+    **compile_kwargs: Any,
+) -> tuple[Callable[..., Any], bool]:
+    """Compile ``fn`` with a fail-closed eager fallback.
+
+    The boolean return value reports whether lazy compilation was activated.
+    Synchronous setup failures return the original callable unchanged.
+    """
+    try:
+        compiled = torch.compile(fn, **compile_kwargs)
+    except Exception as exc:
+        logger.warning(
+            "%s torch.compile setup failed (%s); continuing eagerly.",
+            label,
+            exc,
+        )
+        return fn, False
+    return _FallbackCompiledCallable(fn, compiled, label=label), True
+
+
 def _matches_repeated_block(
     name: str,
     module: nn.Module,

vllm_omni/diffusion/data.py

@@ -726,6 +726,9 @@ class OmniDiffusionConfig:
     # provide its own setup_compile() implementation.
     diffusion_compile_granularity: str = "regional"
     diffusion_compile_dynamic: bool = True
+    # Compile discovered VAE ``decode`` callables whenever generic diffusion
+    # torch.compile is enabled. Runtime failures permanently fall back to eager.
+    diffusion_compile_vae: bool = True
 
     # Parallel weight loading (for faster diffusion model startup)
     enable_multithread_weight_load: bool = True
@@ -930,6 +933,8 @@ class OmniDiffusionConfig:
             )
         if not isinstance(self.diffusion_compile_dynamic, bool):
             raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
+        if not isinstance(self.diffusion_compile_vae, bool):
+            raise TypeError(f"diffusion_compile_vae must be a bool, got {type(self.diffusion_compile_vae)!r}")
         self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
 
         if self.omni_kv_config is None:

vllm_omni/diffusion/diffusion_engine.py

@@ -56,6 +56,7 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 _ASYNC_OUTPUT_TIMEOUT = 30.0  # seconds
+_VAE_COMPILE_WARMUP_REQUEST_ID = "vae_compile_warmup_req_id"
 
 __all__ = [
     "DiffusionEngine",
@@ -851,28 +852,26 @@ class DiffusionEngine:
         num_inference_steps = 1
         height = 512
         width = 512
-        prompt: OmniTextPrompt = {"prompt": "dummy run"}
 
         supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
-        if supports_image_input:
-            # Provide a dummy image input if the model supports it
-            color_format = image_color_format(self.od_config.model_class_name)
-            dummy_image = PIL.Image.new(color_format, (width, height))
-            prompt.setdefault("multi_modal_data", {})["image"] = dummy_image
-
-        if supports_audio_input:
-            audio_sr = 16000
-            dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
-            prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
-
         num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
         if num_frames <= 0:
             logger.info("Skipping dummy warmup run (num_frames=0)")
             return
-        req = OmniDiffusionRequest(
-            prompt=prompt,
-            request_id=DUMMY_DIFFUSION_REQUEST_ID,
-            sampling_params=OmniDiffusionSamplingParams(
+
+        def build_request(request_id: str, prompt_text: str, *, real_path: bool) -> OmniDiffusionRequest:
+            prompt: OmniTextPrompt = {"prompt": prompt_text}
+            if supports_image_input:
+                color_format = image_color_format(self.od_config.model_class_name)
+                prompt.setdefault("multi_modal_data", {})["image"] = PIL.Image.new(
+                    color_format,
+                    (width, height),
+                )
+            if supports_audio_input:
+                audio_sr = 16000
+                prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32)
+
+            sampling_params = OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
                 num_inference_steps=num_inference_steps,
@@ -885,13 +884,47 @@ class DiffusionEngine:
                 # Disable CFG for warmup to avoid triggering CFG parallel
                 # validation when cfg_parallel_size > 1.
                 extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
-            ),
-        )
+            )
+            if real_path:
+                # Internal pretraffic must exercise ordinary pipeline branches
+                # without waiting for KV produced by an upstream stage.
+                sampling_params.need_kv_receive = False
+            return OmniDiffusionRequest(
+                prompt=prompt,
+                request_id=request_id,
+                sampling_params=sampling_params,
+            )
+
+        req = build_request(DUMMY_DIFFUSION_REQUEST_ID, "dummy run", real_path=False)
         logger.info("dummy run to warm up the model")
         output = self.add_req_and_wait_for_response(req)
         if output.error:
             raise RuntimeError(f"Dummy run failed: {output.error}")
 
+        compile_real_path = bool(
+            getattr(self.od_config, "diffusion_compile_vae", False)
+            and not getattr(self.od_config, "enforce_eager", True)
+        )
+        if not compile_real_path:
+            return
+
+        logger.info("real request-path warmup for compiled VAE decode")
+        real_path_req = build_request(
+            _VAE_COMPILE_WARMUP_REQUEST_ID,
+            "A detailed image.",
+            real_path=True,
+        )
+        try:
+            real_path_output = self.add_req_and_wait_for_response(real_path_req)
+        except Exception as exc:
+            logger.warning("Real request-path compile warmup failed (%s); serving will continue.", exc)
+            return
+        if real_path_output.error:
+            logger.warning(
+                "Real request-path compile warmup returned an error (%s); serving will continue.",
+                real_path_output.error,
+            )
+
     def _submit_rpc(
         self,
         method: str,

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -28,7 +28,7 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
     resolve_prompt_embed_cache_config,
 )
 from vllm_omni.diffusion.cache.selector import get_cache_backend
-from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -40,6 +40,7 @@ from vllm_omni.diffusion.models.interface import (
     supports_step_execution,
 )
 from vllm_omni.diffusion.offloader import get_offload_backend
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
@@ -188,6 +189,39 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             compile_dynamic,
         )
 
+    def _compile_vae_decodes(self) -> None:
+        """Compile declared VAE decode callables with eager fallback."""
+        assert self.pipeline is not None
+        modules = ModuleDiscovery.discover(self.pipeline)
+        compiled_count = 0
+        for name, vae in zip(modules.vae_names, modules.vaes, strict=True):
+            decode = getattr(vae, "decode", None)
+            if not callable(decode):
+                logger.debug("Model runner: VAE component %s has no decode callable; skipping compile.", name)
+                continue
+
+            label = f"VAE component {name}.decode"
+            compiled_decode, activated = compile_callable_with_fallback(
+                decode,
+                label=label,
+                mode="default",
+                fullgraph=False,
+                dynamic=None,
+            )
+            if not activated:
+                continue
+            setattr(vae, "decode", compiled_decode)
+            compiled_count += 1
+
+        if compiled_count:
+            logger.info(
+                "Model runner: configured lazy torch.compile for %d VAE decode callable(s); "
+                "startup warmup will compile the real request path.",
+                compiled_count,
+            )
+        else:
+            logger.info("Model runner: no compile-compatible VAE decode callable was discovered.")
+
     def load_model(
         self,
         memory_pool_context_fn: Callable[[str], AbstractContextManager[Any]] | None = None,
@@ -277,7 +311,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         # Apply torch.compile if not in eager mode
         if not self.od_config.enforce_eager:
             if current_omni_platform.supports_torch_inductor():
-                if hasattr(self.pipeline, "setup_compile"):
+                pipeline_owns_compile = hasattr(self.pipeline, "setup_compile")
+                if pipeline_owns_compile:
                     try:
                         self.pipeline.setup_compile()
                     except Exception as exc:
@@ -291,6 +326,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                         transformer_attrs = ("transformer", "transformer_2")
                     for attr_name in transformer_attrs:
                         self._compile_transformer(attr_name)
+
+                if self.od_config.diffusion_compile_vae and not pipeline_owns_compile:
+                    self._compile_vae_decodes()
             else:
                 logger.warning(
                     "Model runner: Platform %s does not support torch inductor, skipping torch.compile.",