[Perf][Diffusion] Compile and prewarm VAE decode

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

tests/diffusion/test_compile.py

@@ -5,7 +5,7 @@ import pytest
 import torch.nn as nn
 
 import vllm_omni.diffusion.compile as compile_module
-from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -74,3 +74,50 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat
         regionally_compile(model, dynamic=True)
 
     assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
+
+
+def test_compile_callable_falls_back_once_after_lazy_failure(monkeypatch, caplog):
+    compiled_calls = 0
+    eager_calls = 0
+
+    def eager(value):
+        nonlocal eager_calls
+        eager_calls += 1
+        return value + 1
+
+    def compile_fn(fn, **kwargs):
+        assert fn is eager
+        assert kwargs == {"mode": "default"}
+
+        def compiled(value):
+            nonlocal compiled_calls
+            compiled_calls += 1
+            raise RuntimeError("lazy failure")
+
+        return compiled
+
+    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
+    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable", mode="default")
+
+    assert activated
+    assert callable_fn(1) == 2
+    assert callable_fn(2) == 3
+    assert compiled_calls == 1
+    assert eager_calls == 2
+    assert "disabling compilation and retrying eagerly" in caplog.text
+
+
+def test_compile_callable_keeps_eager_on_setup_failure(monkeypatch, caplog):
+    def eager(value):
+        return value + 1
+
+    def compile_fn(*args, **kwargs):
+        raise RuntimeError("setup failure")
+
+    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
+    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable")
+
+    assert not activated
+    assert callable_fn is eager
+    assert callable_fn(1) == 2
+    assert "continuing eagerly" in caplog.text

tests/diffusion/test_diffusion_engine.py

@@ -515,16 +515,19 @@ class TestDiffusionCompileConfig:
 
         assert config.diffusion_compile_granularity == "regional"
         assert config.diffusion_compile_dynamic is True
+        assert config.diffusion_compile_vae is True
 
     def test_from_kwargs_preserves_compile_controls(self) -> None:
         config = OmniDiffusionConfig.from_kwargs(
             model="test",
             diffusion_compile_granularity="full",
             diffusion_compile_dynamic=False,
+            diffusion_compile_vae=False,
         )
 
         assert config.diffusion_compile_granularity == "full"
         assert config.diffusion_compile_dynamic is False
+        assert config.diffusion_compile_vae is False
 
     def test_config_rejects_invalid_compile_granularity(self) -> None:
         with pytest.raises(ValueError, match="diffusion_compile_granularity"):
@@ -534,6 +537,10 @@ class TestDiffusionCompileConfig:
         with pytest.raises(TypeError, match="diffusion_compile_dynamic"):
             OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false")
 
+    def test_config_rejects_non_boolean_compile_vae(self) -> None:
+        with pytest.raises(TypeError, match="diffusion_compile_vae"):
+            OmniDiffusionConfig(model="test", diffusion_compile_vae="false")
+
     @pytest.mark.parametrize(
         "kwargs, feature",
         [

tests/diffusion/test_diffusion_model_runner.py

@@ -130,6 +130,23 @@ class _CompileTrackingModel:
         return self
 
 
+class _DecodeVAE(torch.nn.Module):
+    def decode(self, value):
+        return value + 1
+
+
+class _PipelineWithVAE(torch.nn.Module):
+    _dit_modules = []
+    _encoder_modules = []
+    _vae_modules = ["vae", "vae_alias"]
+    _resident_modules = []
+
+    def __init__(self):
+        super().__init__()
+        self.vae = _DecodeVAE()
+        self.vae_alias = self.vae
+
+
 def _make_request():
     sampling_params = SimpleNamespace(
         generator=None,
@@ -320,6 +337,31 @@ def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypa
     assert "lazy compilation errors" in caplog.text
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_compile_vae_decodes_uses_discovery_and_deduplicates_aliases(monkeypatch):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.pipeline = _PipelineWithVAE()
+    compile_calls = []
+
+    def compile_fn(fn, **kwargs):
+        compile_calls.append((fn, kwargs))
+        return lambda *args, **call_kwargs: fn(*args, **call_kwargs) + 10
+
+    monkeypatch.setattr(model_runner_module.torch, "compile", compile_fn)
+
+    DiffusionModelRunner._compile_vae_decodes(runner)
+
+    assert len(compile_calls) == 1
+    assert compile_calls[0][1] == {
+        "mode": "default",
+        "fullgraph": False,
+        "dynamic": None,
+    }
+    assert runner.pipeline.vae.decode(1) == 12
+    assert runner.pipeline.vae_alias.decode is runner.pipeline.vae.decode
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):

tests/diffusion/test_diffusion_scheduler.py

@@ -971,6 +971,44 @@ class TestDiffusionEngine:
         admitted_request = engine.add_req_and_wait_for_response.call_args.args[0]
         assert admitted_request.request_id == "dummy_req_id"
 
+    def test_compile_warmup_repeats_through_normal_request_path(self, mocker: MockerFixture) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.od_config = SimpleNamespace(
+            model_class_name="mock_model",
+            diffusion_load_format="default",
+            diffusion_compile_vae=True,
+            enforce_eager=False,
+        )
+        engine.add_req_and_wait_for_response = mocker.Mock(
+            side_effect=[DiffusionOutput(output=None), DiffusionOutput(output=None)]
+        )
+
+        engine._dummy_run()
+
+        requests = [call.args[0] for call in engine.add_req_and_wait_for_response.call_args_list]
+        assert len(requests) == 2
+        assert requests[0].is_dummy_run()
+        assert not requests[1].is_dummy_run()
+        assert requests[1].prompt == {"prompt": "A detailed image."}
+        assert requests[1].sampling_params.need_kv_receive is False
+        assert requests[1].sampling_params.num_inference_steps == 1
+
+    def test_compile_real_path_warmup_failure_is_non_fatal(self, mocker: MockerFixture, caplog) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.od_config = SimpleNamespace(
+            model_class_name="mock_model",
+            diffusion_load_format="default",
+            diffusion_compile_vae=True,
+            enforce_eager=False,
+        )
+        engine.add_req_and_wait_for_response = mocker.Mock(
+            side_effect=[DiffusionOutput(output=None), RuntimeError("prewarm failure")]
+        )
+
+        engine._dummy_run()
+
+        assert "serving will continue" in caplog.text
+
 
 class TestStepScheduler:
     def setup_method(self) -> None: