[Perf][Diffusion] Compile and prewarm VAE decode

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

tests/config/test_omni_config.py

@@ -751,6 +751,7 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
                 "  - stage_id: 0",
                 "    diffusion_compile_granularity: regional",
                 "    diffusion_compile_dynamic: false",
+                "    diffusion_compile_vae: false",
             ]
         )
     )
@@ -762,13 +763,16 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
         cli_overrides={
             "diffusion_compile_granularity": "full",
             "diffusion_compile_dynamic": True,
+            "diffusion_compile_vae": True,
         },
     ).stage_by_id(0)
 
     assert configured_stage.diffusion_config.diffusion_compile_granularity == "regional"
     assert configured_stage.diffusion_config.diffusion_compile_dynamic is False
+    assert configured_stage.diffusion_config.diffusion_compile_vae is False
     assert overridden_stage.diffusion_config.diffusion_compile_granularity == "full"
     assert overridden_stage.diffusion_config.diffusion_compile_dynamic is True
+    assert overridden_stage.diffusion_config.diffusion_compile_vae is True
 
 
 def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
@@ -776,6 +780,11 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
         omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
 
 
+def test_structured_diffusion_config_rejects_non_boolean_compile_vae():
+    with pytest.raises(ValidationError, match="diffusion_compile_vae"):
+        omni_config_module._DiffusionConfigProjection(diffusion_compile_vae="false")
+
+
 def test_structured_diffusion_config_rejects_invalid_compile_granularity():
     with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
         omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")

tests/diffusion/distributed/test_autoencoder_kl_wan.py

@@ -1,5 +1,3 @@
-from types import SimpleNamespace
-
 import pytest
 import torch
 
@@ -43,21 +41,3 @@ def test_wan_vae_execution_context_uses_platform_autocast(mocker):
         dtype=torch.bfloat16,
         enabled=True,
     )
-
-
-def test_prepare_pipeline_installs_data_movement_before_compile(monkeypatch):
-    model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
-    installed = []
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
-        lambda pipeline: SimpleNamespace(vaes=[model]),
-    )
-    monkeypatch.setattr(
-        wan_vae_module,
-        "install_wan_vae_data_movement",
-        lambda vae: installed.append(vae),
-    )
-
-    wan_vae_module.prepare_pipeline_wan_vae_data_movement(torch.nn.Identity())
-
-    assert installed == [model]

tests/diffusion/kernels/test_wan_vae_data_movement.py

@@ -1,274 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-import torch.nn.functional as F
-from diffusers.models.autoencoders import AutoencoderKLWan
-
-from vllm_omni.diffusion.distributed.autoencoders import wan_vae_data_movement
-from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
-from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import (
-    _cache_payload,
-    _causal_padding,
-    _is_dynamic_spatial_conv,
-    _run_cached_causal_conv,
-    install_wan_vae_data_movement,
-)
-from vllm_omni.diffusion.kernels.wan_vae_data_movement import cat_pad_5d, dup_up3d_add
-
-pytestmark = pytest.mark.core_model
-cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-def _dense_5d(shape: tuple[int, ...], dtype: torch.dtype, channels_last: bool) -> torch.Tensor:
-    tensor = torch.randn(shape, device="cuda", dtype=dtype)
-    if channels_last:
-        return tensor.contiguous(memory_format=torch.channels_last_3d)
-    return tensor.contiguous()
-
-
-def _reference_cat_pad(
-    x: torch.Tensor,
-    cache: torch.Tensor | None,
-    padding: tuple[int, ...],
-) -> torch.Tensor:
-    reference_padding = list(padding)
-    if cache is not None:
-        x = torch.cat([cache, x], dim=2)
-        reference_padding[4] -= cache.shape[2]
-    return F.pad(x, reference_padding)
-
-
-@pytest.mark.cuda
-@cuda_only
-@torch.no_grad()
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
-@pytest.mark.parametrize("channels_last", [False, True])
-@pytest.mark.parametrize(
-    "channels,frames,height,width,cache_frames,padding",
-    [
-        (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)),
-        (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)),
-        (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
-        (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)),
-        (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
-    ],
-)
-def test_cat_pad_5d_is_bitwise_and_layout_exact(
-    dtype: torch.dtype,
-    channels_last: bool,
-    channels: int,
-    frames: int,
-    height: int,
-    width: int,
-    cache_frames: int,
-    padding: tuple[int, ...],
-) -> None:
-    torch.cuda.manual_seed(0)
-    x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
-    cache = None
-    if cache_frames:
-        pad_height, pad_width = padding[2], padding[0]
-        buffer = _dense_5d(
-            (1, channels, cache_frames, height + 2 * pad_height, width + 2 * pad_width),
-            dtype,
-            channels_last,
-        )
-        cache = buffer[:, :, :, pad_height : pad_height + height, pad_width : pad_width + width]
-
-    reference = _reference_cat_pad(x, cache, padding)
-    output = cat_pad_5d(x, cache, padding)
-    assert output is not None
-    assert output.shape == reference.shape
-    assert output.stride() == reference.stride()
-    assert torch.equal(output, reference)
-
-    pair = cat_pad_5d(x, cache, padding, keep_cache_frames=2)
-    assert pair is not None
-    output_with_cache, next_cache = pair
-    assert torch.equal(output_with_cache, reference)
-    pad_height, pad_width = padding[2], padding[0]
-    keep_frames = min(2, reference.shape[2])
-    expected_cache = reference[
-        :,
-        :,
-        -keep_frames:,
-        pad_height : pad_height + height,
-        pad_width : pad_width + width,
-    ]
-    assert next_cache.shape == expected_cache.shape
-    assert torch.equal(next_cache, expected_cache)
-    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
-
-
-@pytest.mark.cuda
-@cuda_only
-@torch.no_grad()
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
-@pytest.mark.parametrize("main_layout", ["contiguous", "channels_last", "permuted"])
-@pytest.mark.parametrize(
-    "in_channels,out_channels,frames,height,width,factor_temporal,factor_spatial,drop_first",
-    [
-        (128, 64, 1, 10, 14, 2, 2, False),
-        (128, 64, 1, 10, 14, 2, 2, True),
-        (64, 32, 2, 10, 14, 1, 2, False),
-    ],
-)
-def test_dup_up3d_add_is_bitwise_and_layout_exact(
-    dtype: torch.dtype,
-    main_layout: str,
-    in_channels: int,
-    out_channels: int,
-    frames: int,
-    height: int,
-    width: int,
-    factor_temporal: int,
-    factor_spatial: int,
-    drop_first: bool,
-) -> None:
-    torch.cuda.manual_seed(0)
-    repeats = out_channels * factor_temporal * factor_spatial * factor_spatial // in_channels
-    source = _dense_5d((1, in_channels, frames, height, width), dtype, main_layout == "channels_last")
-    out_frames = frames * factor_temporal - (factor_temporal - 1 if drop_first else 0)
-    shape = (1, out_channels, out_frames, height * factor_spatial, width * factor_spatial)
-    if main_layout == "channels_last":
-        main = _dense_5d(shape, dtype, True)
-    elif main_layout == "permuted":
-        main = torch.randn(
-            (shape[0], shape[2], shape[1], shape[3], shape[4]),
-            device="cuda",
-            dtype=dtype,
-        ).permute(0, 2, 1, 3, 4)
-    else:
-        main = _dense_5d(shape, dtype, False)
-
-    shortcut = source.repeat_interleave(repeats, dim=1)
-    shortcut = shortcut.view(
-        1,
-        out_channels,
-        factor_temporal,
-        factor_spatial,
-        factor_spatial,
-        frames,
-        height,
-        width,
-    )
-    shortcut = shortcut.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
-    shortcut = shortcut.view(shape[0], shape[1], frames * factor_temporal, shape[3], shape[4])
-    if drop_first:
-        shortcut = shortcut[:, :, factor_temporal - 1 :]
-    reference = main + shortcut
-
-    output = dup_up3d_add(
-        main,
-        source,
-        factor_temporal,
-        factor_spatial,
-        repeats,
-        drop_first,
-    )
-    assert output is not None
-    assert output.shape == reference.shape
-    assert output.stride() == reference.stride()
-    assert torch.equal(output, reference)
-
-
-@pytest.mark.cuda
-@cuda_only
-@torch.no_grad()
-@pytest.mark.parametrize("channels_last", [False, True])
-@pytest.mark.parametrize("temporal_only", [False, True])
-def test_cached_conv_chunk_loop_matches_reference(monkeypatch, channels_last: bool, temporal_only: bool) -> None:
-    from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T, WanCausalConv3d
-
-    torch.cuda.manual_seed(0)
-    channels = 64
-    if temporal_only:
-        conv = WanCausalConv3d(channels, 2 * channels, (3, 1, 1), padding=(1, 0, 0))
-    else:
-        conv = WanCausalConv3d(channels, channels, 3, padding=1)
-    conv = conv.to(device="cuda", dtype=torch.float32)
-    chunks = [_dense_5d((1, channels, 1, 10, 14), torch.float32, channels_last) for _ in range(4)]
-
-    def run(force_reference: bool, start):
-        cache = [start]
-        outputs = []
-        original = wan_vae_data_movement.cat_pad_5d
-        if force_reference:
-            monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", None)
-        try:
-            for chunk in chunks:
-                outputs.append(_run_cached_causal_conv(conv, chunk, cache, 0))
-        finally:
-            monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", original)
-        return outputs, cache[0]
-
-    for start in (None, "Rep"):
-        fused_outputs, fused_cache = run(False, start)
-        reference_outputs, reference_cache = run(True, start)
-        for output, reference in zip(fused_outputs, reference_outputs, strict=True):
-            assert torch.equal(output, reference)
-        payload = _cache_payload(fused_cache)
-        assert payload is not None and payload.shape[2] == CACHE_T
-        assert torch.equal(payload, reference_cache[:, :, -CACHE_T:])
-
-
-@pytest.mark.cpu
-def test_installer_is_idempotent_and_preserves_state_dict_keys() -> None:
-    vae = OmniAutoencoderKLWan(
-        base_dim=8,
-        decoder_base_dim=8,
-        z_dim=4,
-        dim_mult=[1, 1],
-        num_res_blocks=1,
-        temperal_downsample=[False, False],
-        is_residual=True,
-    )
-    keys_before = set(vae.state_dict())
-    assert install_wan_vae_data_movement(vae)
-    assert install_wan_vae_data_movement(vae)
-    assert set(vae.state_dict()) == keys_before
-
-
-@pytest.mark.cpu
-def test_dynamic_spatial_conv_exposes_direct_path_compatibility_contract() -> None:
-    conv = torch.nn.Conv3d(2, 2, 3)
-    conv._vllm_omni_dynamic_spatial_shard_conv = True
-    conv._source_padding = (1, 1, 1, 1, 2, 0)
-
-    assert _is_dynamic_spatial_conv(conv)
-    assert _causal_padding(conv) == conv._source_padding
-
-
-@pytest.mark.cuda
-@cuda_only
-@torch.no_grad()
-@pytest.mark.parametrize("channels_last", [False, True])
-def test_tiny_wan_decoder_is_bitwise_exact(channels_last: bool) -> None:
-    config = dict(
-        base_dim=8,
-        decoder_base_dim=8,
-        z_dim=4,
-        dim_mult=[1, 1],
-        num_res_blocks=1,
-        temperal_downsample=[False, True],
-        is_residual=True,
-    )
-    torch.manual_seed(0)
-    reference = AutoencoderKLWan(**config).eval().to("cuda")
-    candidate = OmniAutoencoderKLWan(**config).eval().to("cuda")
-    candidate.load_state_dict(reference.state_dict())
-    if channels_last:
-        for model in (reference, candidate):
-            for module in model.decoder.modules():
-                if isinstance(module, torch.nn.Conv3d):
-                    module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d)
-
-    latents = torch.randn(1, 4, 3, 4, 4, device="cuda")
-    if channels_last:
-        latents = latents.contiguous(memory_format=torch.channels_last_3d)
-    expected = reference.decode(latents, return_dict=False)[0]
-    output = candidate.decode(latents, return_dict=False)[0]
-    assert output.stride() == expected.stride()
-    assert torch.equal(output, expected)

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):
@@ -728,8 +770,6 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
-    prepared_pipelines = []
-
     class _DummyLoader:
         def __init__(self, load_config, od_config=None):
             del load_config, od_config
@@ -777,11 +817,6 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
     monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
     monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
-    monkeypatch.setattr(
-        model_runner_module,
-        "prepare_pipeline_wan_vae_data_movement",
-        lambda pipeline: prepared_pipelines.append(pipeline),
-    )
     monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
     monkeypatch.setattr(
         model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
@@ -792,7 +827,6 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     assert runner.cache_backend is None
     assert runner.od_config.cache_backend is None
     assert dummy_cache_backend.enabled is False
-    assert prepared_pipelines == [runner.pipeline]
 
 
 @pytest.mark.core_model

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:

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -87,6 +87,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
             "enforce_eager": None,
             "diffusion_compile_granularity": None,
             "diffusion_compile_dynamic": None,
+            "diffusion_compile_vae": None,
         }
     )[0]
 
@@ -98,6 +99,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
     assert stage_cfg["engine_args"]["enforce_eager"] is False
     assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "regional"
     assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is True
+    assert stage_cfg["engine_args"]["diffusion_compile_vae"] is True
 
 
 def test_default_stage_config_propagates_ulysses_mode():
@@ -315,7 +317,7 @@ def test_serve_cli_forwards_distributed_offload_residency():
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():
-    """Ensure both compile controls reach the diffusion stage."""
+    """Ensure all compile controls reach the diffusion stage."""
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
@@ -328,6 +330,7 @@ def test_serve_cli_accepts_diffusion_compile_controls():
             "--diffusion-compile-granularity",
             "full",
             "--no-diffusion-compile-dynamic",
+            "--no-diffusion-compile-vae",
         ]
     )
 
@@ -336,8 +339,10 @@ def test_serve_cli_accepts_diffusion_compile_controls():
 
     assert args.diffusion_compile_granularity == "full"
     assert args.diffusion_compile_dynamic is False
+    assert args.diffusion_compile_vae is False
     assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "full"
     assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False
+    assert stage_cfg["engine_args"]["diffusion_compile_vae"] is False
 
 
 def test_serve_cli_accepts_diffusion_attention_backend():

vllm_omni/config/omni_config.py

@@ -570,6 +570,7 @@ class _DiffusionConfigProjection:
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
+    diffusion_compile_vae: bool = Field(default=True, strict=True)
     fa_deterministic: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False

vllm_omni/config/stage_config.py

@@ -394,6 +394,7 @@ class StageDeployConfig:
     # Diffusion execution, cache, and VAE behavior.
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
+    diffusion_compile_vae: bool | None = None
     fa_deterministic: bool | None = None
     cache_backend: str | None = None
     cache_config: dict[str, Any] | None = None

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/distributed/autoencoders/autoencoder_kl_wan.py

@@ -18,26 +18,13 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     GridSpec,
     TileTask,
 )
-from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import (
-    install_wan_vae_data_movement,
-)
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-def prepare_pipeline_wan_vae_data_movement(pipeline: torch.nn.Module) -> None:
-    """Install Wan data-movement methods before offload or compilation setup."""
-    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
-
-    for vae in ModuleDiscovery.discover(pipeline).vaes:
-        if isinstance(vae, OmniAutoencoderKLWan):
-            install_wan_vae_data_movement(vae)
-
-
 class OmniAutoencoderKLWan(AutoencoderKLWan):
     def _execution_context(self):
-        install_wan_vae_data_movement(self)
         try:
             first_param = next(self.parameters())
         except StopIteration:

vllm_omni/diffusion/distributed/autoencoders/wan_vae_data_movement.py

@@ -1,302 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Guarded Wan VAE decoder data-movement fast paths.
-
-This module adapts the lossless Wan decoder optimizations from SGLang while
-keeping Diffusers' parameter names and module topology unchanged. Every
-optimized method falls back to the Diffusers 0.38.0 operation order when the
-kernel, device, layout, dtype, autograd state, or module type is unsupported.
-"""
-
-from __future__ import annotations
-
-from types import MethodType
-from typing import Any
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.models.autoencoders.autoencoder_kl_wan import (
-    CACHE_T,
-    DupUp3D,
-    WanCausalConv3d,
-    WanDecoder3d,
-    WanResample,
-    WanResidualBlock,
-    WanResidualUpBlock,
-)
-
-from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
-
-try:
-    from vllm_omni.diffusion.kernels.wan_vae_data_movement import (
-        cat_pad_5d,
-        dup_up3d_add,
-    )
-except ImportError:  # pragma: no cover - Triton is optional on non-GPU installs.
-    cat_pad_5d = None
-    dup_up3d_add = None
-
-_INSTALL_ATTR = "_vllm_omni_wan_data_movement_installed"
-_DYNAMIC_SPATIAL_CONV_ATTR = "_vllm_omni_dynamic_spatial_shard_conv"
-
-
-def _cache_payload(cache: Any) -> torch.Tensor | None:
-    return cache if isinstance(cache, torch.Tensor) else None
-
-
-def _is_dynamic_spatial_conv(conv: nn.Module) -> bool:
-    return bool(getattr(conv, _DYNAMIC_SPATIAL_CONV_ATTR, False))
-
-
-def _spatial_shard_is_active() -> bool:
-    return wan_spatial_shard._SPATIAL_SHARD_CONTEXT.get() is not None
-
-
-def _causal_padding(conv: nn.Module) -> tuple[int, ...]:
-    if _is_dynamic_spatial_conv(conv):
-        return tuple(conv._source_padding)
-    return tuple(conv._padding)
-
-
-def _can_use_fused_cache_path(conv: nn.Module, x: torch.Tensor) -> bool:
-    return bool(
-        cat_pad_5d is not None
-        and (type(conv) is WanCausalConv3d or _is_dynamic_spatial_conv(conv))
-        and x.dim() == 5
-        and x.is_cuda
-        and not torch.is_grad_enabled()
-        and not torch.compiler.is_compiling()
-        and not _spatial_shard_is_active()
-    )
-
-
-def _run_cached_causal_conv(
-    conv: nn.Module,
-    x: torch.Tensor,
-    cache_list: list[Any],
-    index: int,
-) -> torch.Tensor:
-    """Run a causal conv and refresh its feature-cache slot."""
-    cache = cache_list[index]
-    is_repeat_marker = isinstance(cache, str)
-    payload = None if is_repeat_marker else _cache_payload(cache)
-    if _can_use_fused_cache_path(conv, x) and (
-        payload is None or (payload.device == x.device and payload.dtype == x.dtype)
-    ):
-        pair = cat_pad_5d(x, payload, _causal_padding(conv), keep_cache_frames=CACHE_T)
-        if pair is not None:
-            conv_input, cache_list[index] = pair
-            return nn.Conv3d.forward(conv, conv_input)
-
-    cache_x = x[:, :, -CACHE_T:, :, :].clone()
-    if cache_x.shape[2] < CACHE_T and payload is not None:
-        cache_x = torch.cat(
-            [payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
-            dim=2,
-        )
-    elif cache_x.shape[2] < CACHE_T and is_repeat_marker:
-        cache_x = torch.cat([torch.zeros_like(cache_x), cache_x], dim=2)
-
-    out = conv(x) if payload is None else conv(x, payload)
-    cache_list[index] = cache_x
-    return out
-
-
-def _causal_conv_forward(self: WanCausalConv3d, x: torch.Tensor, cache_x: torch.Tensor | None = None):
-    padding = list(self._padding)
-    if (
-        any(padding)
-        and _can_use_fused_cache_path(self, x)
-        and (cache_x is None or (cache_x.device == x.device and cache_x.dtype == x.dtype))
-    ):
-        conv_input = cat_pad_5d(x, cache_x, padding)
-        if conv_input is not None:
-            return nn.Conv3d.forward(self, conv_input)
-
-    if cache_x is not None and padding[4] > 0:
-        cache_x = cache_x.to(x.device)
-        x = torch.cat([cache_x, x], dim=2)
-        padding[4] -= cache_x.shape[2]
-    x = F.pad(x, padding)
-    return nn.Conv3d.forward(self, x)
-
-
-def _resample_forward(
-    self: WanResample,
-    x: torch.Tensor,
-    feat_cache: list[Any] | None = None,
-    feat_idx: list[int] | None = None,
-):
-    if feat_idx is None:
-        feat_idx = [0]
-    batch, channels, frames, height, width = x.size()
-    if self.mode == "upsample3d" and feat_cache is not None:
-        index = feat_idx[0]
-        if feat_cache[index] is None:
-            feat_cache[index] = "Rep"
-            feat_idx[0] += 1
-        else:
-            x = _run_cached_causal_conv(self.time_conv, x, feat_cache, index)
-            feat_idx[0] += 1
-            x = x.reshape(batch, 2, channels, frames, height, width)
-            x = torch.stack((x[:, 0], x[:, 1]), 3)
-            x = x.reshape(batch, channels, frames * 2, height, width)
-
-    frames = x.shape[2]
-    x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
-    x = self.resample(x)
-    x = x.view(batch, frames, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
-
-    if self.mode == "downsample3d" and feat_cache is not None:
-        index = feat_idx[0]
-        if feat_cache[index] is None:
-            feat_cache[index] = x.clone()
-            feat_idx[0] += 1
-        else:
-            cache_x = x[:, :, -1:, :, :].clone()
-            x = self.time_conv(torch.cat([feat_cache[index][:, :, -1:, :, :], x], dim=2))
-            feat_cache[index] = cache_x
-            feat_idx[0] += 1
-    return x
-
-
-def _residual_block_forward(
-    self: WanResidualBlock,
-    x: torch.Tensor,
-    feat_cache: list[Any] | None = None,
-    feat_idx: list[int] | None = None,
-):
-    if feat_idx is None:
-        feat_idx = [0]
-    residual = self.conv_shortcut(x)
-    x = self.nonlinearity(self.norm1(x))
-    if feat_cache is not None:
-        index = feat_idx[0]
-        x = _run_cached_causal_conv(self.conv1, x, feat_cache, index)
-        feat_idx[0] += 1
-    else:
-        x = self.conv1(x)
-
-    x = self.dropout(self.nonlinearity(self.norm2(x)))
-    if feat_cache is not None:
-        index = feat_idx[0]
-        x = _run_cached_causal_conv(self.conv2, x, feat_cache, index)
-        feat_idx[0] += 1
-    else:
-        x = self.conv2(x)
-    return x + residual
-
-
-def _residual_up_block_forward(
-    self: WanResidualUpBlock,
-    x: torch.Tensor,
-    feat_cache: list[Any] | None = None,
-    feat_idx: list[int] | None = None,
-    first_chunk: bool = False,
-):
-    if feat_idx is None:
-        feat_idx = [0]
-    shortcut_source = x.clone()
-    for resnet in self.resnets:
-        if feat_cache is None:
-            x = resnet(x)
-        else:
-            x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx)
-
-    if self.upsampler is not None:
-        if feat_cache is None:
-            x = self.upsampler(x)
-        else:
-            x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx)
-
-    shortcut = self.avg_shortcut
-    if shortcut is None:
-        return x
-    if (
-        dup_up3d_add is not None
-        and type(shortcut) is DupUp3D
-        and x.is_cuda
-        and shortcut_source.is_cuda
-        and x.dtype == shortcut_source.dtype
-        and not torch.is_grad_enabled()
-        and not torch.compiler.is_compiling()
-        and not _spatial_shard_is_active()
-    ):
-        fused = dup_up3d_add(
-            x,
-            shortcut_source,
-            shortcut.factor_t,
-            shortcut.factor_s,
-            shortcut.repeats,
-            first_chunk,
-        )
-        if fused is not None:
-            return fused
-    return x + shortcut(shortcut_source, first_chunk=first_chunk)
-
-
-def _decoder_forward(
-    self: WanDecoder3d,
-    x: torch.Tensor,
-    feat_cache: list[Any] | None = None,
-    feat_idx: list[int] | None = None,
-    first_chunk: bool = False,
-):
-    if feat_idx is None:
-        feat_idx = [0]
-    if feat_cache is not None:
-        index = feat_idx[0]
-        x = _run_cached_causal_conv(self.conv_in, x, feat_cache, index)
-        feat_idx[0] += 1
-    else:
-        x = self.conv_in(x)
-
-    x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx)
-    for up_block in self.up_blocks:
-        x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
-
-    x = self.nonlinearity(self.norm_out(x))
-    if feat_cache is not None:
-        index = feat_idx[0]
-        x = _run_cached_causal_conv(self.conv_out, x, feat_cache, index)
-        feat_idx[0] += 1
-    else:
-        x = self.conv_out(x)
-    return x
-
-
-def install_wan_vae_data_movement(vae: nn.Module) -> bool:
-    """Install guarded fast paths on a Diffusers Wan decoder once."""
-    if getattr(vae, _INSTALL_ATTR, False):
-        return True
-    if cat_pad_5d is None or dup_up3d_add is None:
-        return False
-
-    decoder = getattr(vae, "decoder", None)
-    if type(decoder) is not WanDecoder3d:
-        return False
-
-    decoder.forward = MethodType(_decoder_forward, decoder)
-    patched = 0
-    for module in decoder.modules():
-        module_type = type(module)
-        if module_type is WanCausalConv3d:
-            module.forward = MethodType(_causal_conv_forward, module)
-            patched += 1
-        elif _is_dynamic_spatial_conv(module):
-            # The dynamic wrapper owns halo exchange while a spatial request is
-            # active. Its direct path can still use the lossless fused cache
-            # movement through ``_run_cached_causal_conv``.
-            patched += 1
-        elif module_type is WanResidualBlock:
-            module.forward = MethodType(_residual_block_forward, module)
-        elif module_type is WanResidualUpBlock:
-            module.forward = MethodType(_residual_up_block_forward, module)
-        elif module_type is WanResample:
-            module.forward = MethodType(_resample_forward, module)
-
-    if patched == 0:
-        return False
-    setattr(vae, _INSTALL_ATTR, True)
-    return True

vllm_omni/diffusion/kernels/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/kernels/wan_vae_data_movement.py

@@ -1,405 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Bit-exact data-movement kernels for the Wan causal VAE decoder.
-
-Adapted from SGLang PR #34125:
-https://github.com/sgl-project/sglang/pull/34125
-
-The kernels in this module only move values, fill zeros, or perform the same
-two-operand addition as the PyTorch reference path. They preserve the input
-layout instead of selecting a new Conv3d layout, so enabling them does not
-change the decoder's cuDNN accumulation order.
-"""
-
-from __future__ import annotations
-
-import torch
-import triton
-import triton.language as tl
-
-_MAX_INT32 = 2**31 - 1
-_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
-
-
-@triton.jit
-def _cat_pad_5d_kernel(
-    x_ptr,
-    cache_ptr,
-    out_ptr,
-    keep_ptr,
-    total,
-    channels,
-    frames,
-    height,
-    width,
-    cache_frames,
-    out_frames,
-    out_height,
-    out_width,
-    zero_front_frames,
-    pad_height,
-    pad_width,
-    stride_x_batch,
-    stride_x_channel,
-    stride_x_frame,
-    stride_x_height,
-    stride_x_width,
-    stride_cache_batch,
-    stride_cache_channel,
-    stride_cache_frame,
-    stride_cache_height,
-    stride_cache_width,
-    has_cache: tl.constexpr,
-    keep_frames_const: tl.constexpr,
-    channels_inner: tl.constexpr,
-    index_64bit: tl.constexpr,
-    block_size_const: tl.constexpr,
-):
-    if index_64bit:
-        offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
-    else:
-        offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
-    mask = offsets < total
-
-    if channels_inner:
-        out_channel = offsets % channels
-        remaining = offsets // channels
-        out_width_idx = remaining % out_width
-        remaining = remaining // out_width
-        out_height_idx = remaining % out_height
-        remaining = remaining // out_height
-        out_frame = remaining % out_frames
-        out_batch = remaining // out_frames
-    else:
-        out_width_idx = offsets % out_width
-        remaining = offsets // out_width
-        out_height_idx = remaining % out_height
-        remaining = remaining // out_height
-        out_frame = remaining % out_frames
-        remaining = remaining // out_frames
-        out_channel = remaining % channels
-        out_batch = remaining // channels
-
-    input_width = out_width_idx - pad_width
-    input_height = out_height_idx - pad_height
-    input_frame = out_frame - zero_front_frames
-
-    spatial_valid = (input_width >= 0) & (input_width < width) & (input_height >= 0) & (input_height < height)
-    from_cache = spatial_valid & (input_frame >= 0) & (input_frame < cache_frames)
-    from_x = spatial_valid & (input_frame >= cache_frames) & (input_frame < cache_frames + frames)
-
-    x_frame = input_frame - cache_frames
-    x_offset = (
-        out_batch * stride_x_batch
-        + out_channel * stride_x_channel
-        + x_frame * stride_x_frame
-        + input_height * stride_x_height
-        + input_width * stride_x_width
-    )
-    values = tl.load(x_ptr + x_offset, mask=mask & from_x, other=0.0)
-    if has_cache:
-        cache_offset = (
-            out_batch * stride_cache_batch
-            + out_channel * stride_cache_channel
-            + input_frame * stride_cache_frame
-            + input_height * stride_cache_height
-            + input_width * stride_cache_width
-        )
-        cache_values = tl.load(cache_ptr + cache_offset, mask=mask & from_cache, other=0.0)
-        values = tl.where(from_cache, cache_values, values)
-    tl.store(out_ptr + offsets, values, mask=mask)
-
-    if keep_frames_const > 0:
-        keep_frame = out_frame - (out_frames - keep_frames_const)
-        keep_mask = mask & spatial_valid & (keep_frame >= 0)
-        if channels_inner:
-            keep_offset = (
-                ((out_batch * keep_frames_const + keep_frame) * height + input_height) * width + input_width
-            ) * channels + out_channel
-        else:
-            keep_offset = (
-                ((out_batch * channels + out_channel) * keep_frames_const + keep_frame) * height + input_height
-            ) * width + input_width
-        tl.store(keep_ptr + keep_offset, values, mask=keep_mask)
-
-
-def _uses_channels_last_3d(tensor: torch.Tensor) -> bool:
-    return tensor.shape[1] > 1 and tensor.stride(1) == 1
-
-
-def cat_pad_5d(
-    x: torch.Tensor,
-    cache_x: torch.Tensor | None,
-    padding: list[int] | tuple[int, ...],
-    keep_cache_frames: int = 0,
-) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None:
-    """Fuse temporal concat, constant padding, and optional cache refresh.
-
-    ``padding`` follows the Wan causal Conv3d convention
-    ``(w_left, w_right, h_top, h_bottom, t_front, t_back)``. The output uses
-    the same dense layout that the reference ``cat``/``pad`` path selects.
-    Unsupported inputs return ``None`` so callers can execute the reference
-    implementation.
-    """
-    if len(padding) != 6:
-        return None
-    pad_width_left, pad_width_right, pad_height_top, pad_height_bottom, pad_front, pad_back = padding
-    if any(value < 0 for value in padding):
-        return None
-    if pad_width_left != pad_width_right or pad_height_top != pad_height_bottom or pad_back != 0:
-        return None
-    if x.dim() != 5 or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES:
-        return None
-    if keep_cache_frames < 0:
-        return None
-
-    channels_inner = _uses_channels_last_3d(x)
-    cache_frames = 0
-    if cache_x is not None:
-        if (
-            cache_x.dim() != 5
-            or cache_x.dtype != x.dtype
-            or cache_x.device != x.device
-            or cache_x.shape[0] != x.shape[0]
-            or cache_x.shape[1] != x.shape[1]
-            or cache_x.shape[3:] != x.shape[3:]
-            or _uses_channels_last_3d(cache_x) != channels_inner
-        ):
-            return None
-        cache_frames = cache_x.shape[2]
-
-    zero_front_frames = pad_front - cache_frames
-    if zero_front_frames < 0:
-        return None
-
-    batch, channels, frames, height, width = x.shape
-    out_frames = pad_front + frames
-    out_height = height + 2 * pad_height_top
-    out_width = width + 2 * pad_width_left
-    keep_frames = min(keep_cache_frames, out_frames)
-    memory_format = torch.channels_last_3d if channels_inner else torch.contiguous_format
-    out = torch.empty(
-        (batch, channels, out_frames, out_height, out_width),
-        device=x.device,
-        dtype=x.dtype,
-        memory_format=memory_format,
-    )
-    total = out.numel()
-    if total == 0 or total > _MAX_INT32 * 4:
-        return None
-
-    if keep_frames:
-        keep_arg = torch.empty(
-            (batch, channels, keep_frames, height, width),
-            device=x.device,
-            dtype=x.dtype,
-            memory_format=memory_format,
-        )
-    else:
-        keep_arg = out
-
-    if cache_x is None:
-        cache_arg = x
-        cache_strides = (0, 0, 0, 0, 0)
-    else:
-        cache_arg = cache_x
-        cache_strides = cache_x.stride()
-
-    block_size = 512
-    grid = (triton.cdiv(total, block_size),)
-    with torch.get_device_module().device(x.device):
-        _cat_pad_5d_kernel[grid](
-            x,
-            cache_arg,
-            out,
-            keep_arg,
-            total,
-            channels,
-            frames,
-            height,
-            width,
-            cache_frames,
-            out_frames,
-            out_height,
-            out_width,
-            zero_front_frames,
-            pad_height_top,
-            pad_width_left,
-            *x.stride(),
-            *cache_strides,
-            has_cache=cache_x is not None,
-            keep_frames_const=keep_frames,
-            channels_inner=channels_inner,
-            index_64bit=total >= _MAX_INT32,
-            block_size_const=block_size,
-        )
-
-    if keep_cache_frames:
-        return out, keep_arg
-    return out
-
-
-@triton.jit
-def _dup_up3d_add_kernel(
-    main_ptr,
-    source_ptr,
-    out_ptr,
-    total,
-    out_channels,
-    out_frames,
-    out_height,
-    out_width,
-    frame_offset,
-    stride_main_batch,
-    stride_main_channel,
-    stride_main_frame,
-    stride_main_height,
-    stride_main_width,
-    stride_source_batch,
-    stride_source_channel,
-    stride_source_frame,
-    stride_source_height,
-    stride_source_width,
-    stride_out_batch,
-    stride_out_channel,
-    stride_out_frame,
-    stride_out_height,
-    stride_out_width,
-    factor_temporal_const: tl.constexpr,
-    factor_spatial_const: tl.constexpr,
-    repeats_const: tl.constexpr,
-    channels_inner: tl.constexpr,
-    index_64bit: tl.constexpr,
-    block_size_const: tl.constexpr,
-):
-    if index_64bit:
-        offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
-    else:
-        offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
-    mask = offsets < total
-
-    if channels_inner:
-        out_channel = offsets % out_channels
-        remaining = offsets // out_channels
-        out_width_idx = remaining % out_width
-        remaining = remaining // out_width
-        out_height_idx = remaining % out_height
-        remaining = remaining // out_height
-        out_frame = remaining % out_frames
-        out_batch = remaining // out_frames
-    else:
-        out_width_idx = offsets % out_width
-        remaining = offsets // out_width
-        out_height_idx = remaining % out_height
-        remaining = remaining // out_height
-        out_frame = remaining % out_frames
-        remaining = remaining // out_frames
-        out_channel = remaining % out_channels
-        out_batch = remaining // out_channels
-
-    unsliced_frame = out_frame + frame_offset
-    source_frame = unsliced_frame // factor_temporal_const
-    temporal_remainder = unsliced_frame % factor_temporal_const
-    source_height = out_height_idx // factor_spatial_const
-    height_remainder = out_height_idx % factor_spatial_const
-    source_width = out_width_idx // factor_spatial_const
-    width_remainder = out_width_idx % factor_spatial_const
-    repeated_channel = (
-        (out_channel * factor_temporal_const + temporal_remainder) * factor_spatial_const + height_remainder
-    ) * factor_spatial_const + width_remainder
-    source_channel = repeated_channel // repeats_const
-
-    main_offset = (
-        out_batch * stride_main_batch
-        + out_channel * stride_main_channel
-        + out_frame * stride_main_frame
-        + out_height_idx * stride_main_height
-        + out_width_idx * stride_main_width
-    )
-    source_offset = (
-        out_batch * stride_source_batch
-        + source_channel * stride_source_channel
-        + source_frame * stride_source_frame
-        + source_height * stride_source_height
-        + source_width * stride_source_width
-    )
-    out_offset = (
-        out_batch * stride_out_batch
-        + out_channel * stride_out_channel
-        + out_frame * stride_out_frame
-        + out_height_idx * stride_out_height
-        + out_width_idx * stride_out_width
-    )
-    main_value = tl.load(main_ptr + main_offset, mask=mask, other=0.0)
-    source_value = tl.load(source_ptr + source_offset, mask=mask, other=0.0)
-    value = main_value.to(tl.float32) + source_value.to(tl.float32)
-    tl.store(out_ptr + out_offset, value, mask=mask)
-
-
-def dup_up3d_add(
-    main: torch.Tensor,
-    source: torch.Tensor,
-    factor_temporal: int,
-    factor_spatial: int,
-    repeats: int,
-    drop_first_frames: bool,
-) -> torch.Tensor | None:
-    """Evaluate ``main + DupUp3D(source)`` without materializing the shortcut."""
-    if main.dim() != 5 or source.dim() != 5:
-        return None
-    if factor_temporal <= 0 or factor_spatial <= 0 or repeats <= 0:
-        return None
-    if factor_temporal & (factor_temporal - 1) or factor_spatial & (factor_spatial - 1):
-        return None
-    if repeats & (repeats - 1):
-        return None
-    if not main.is_cuda or not source.is_cuda:
-        return None
-    if main.dtype != source.dtype or main.dtype not in _SUPPORTED_DTYPES or main.device != source.device:
-        return None
-
-    batch, source_channels, frames, height, width = source.shape
-    frame_offset = factor_temporal - 1 if drop_first_frames else 0
-    numerator = source_channels * repeats
-    denominator = factor_temporal * factor_spatial * factor_spatial
-    if numerator % denominator:
-        return None
-    expected_shape = (
-        batch,
-        numerator // denominator,
-        frames * factor_temporal - frame_offset,
-        height * factor_spatial,
-        width * factor_spatial,
-    )
-    if tuple(main.shape) != expected_shape:
-        return None
-
-    out = torch.empty_like(main)
-    total = out.numel()
-    if total == 0 or total > _MAX_INT32 * 4:
-        return None
-
-    block_size = 512
-    grid = (triton.cdiv(total, block_size),)
-    with torch.get_device_module().device(main.device):
-        _dup_up3d_add_kernel[grid](
-            main,
-            source,
-            out,
-            total,
-            expected_shape[1],
-            expected_shape[2],
-            expected_shape[3],
-            expected_shape[4],
-            frame_offset,
-            *main.stride(),
-            *source.stride(),
-            *out.stride(),
-            factor_temporal_const=factor_temporal,
-            factor_spatial_const=factor_spatial,
-            repeats_const=repeats,
-            channels_inner=out.stride(1) == 1 and expected_shape[1] > 1,
-            index_64bit=total >= _MAX_INT32,
-            block_size_const=block_size,
-        )
-    return out

vllm_omni/diffusion/models/helios/pipeline_helios.py

@@ -21,7 +21,6 @@ from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -216,7 +215,6 @@ class HeliosPipeline(
         self.vae = AutoencoderKLWan.from_pretrained(
             model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
         ).to(self.device)
-        install_wan_vae_data_movement(self.vae)
 
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
         self.transformer = create_transformer_from_config(

vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py

@@ -19,7 +19,6 @@ from torch import nn
 from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.interface import (
     SupportImageInput,
@@ -375,7 +374,6 @@ class LingBotVideoPipeline(
             torch_dtype=vae_dtype,
             local_files_only=local_files_only,
         ).to(self.device)
-        install_wan_vae_data_movement(self.vae)
         self.scheduler = FlowUniPCMultistepScheduler.from_pretrained(
             model,
             subfolder=scheduler_subfolder,

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -28,11 +28,8 @@ 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.distributed.autoencoders.autoencoder_kl_wan import (
-    prepare_pipeline_wan_vae_data_movement,
-)
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import (
@@ -43,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
@@ -191,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,
@@ -246,7 +277,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     custom_pipeline_name=custom_pipeline_name,
                     device=self.device,
                 )
-                prepare_pipeline_wan_vae_data_movement(self.pipeline)
         time_after_load = time.perf_counter()
 
         logger.info(
@@ -281,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:
@@ -295,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.",

vllm_omni/engine/arg_utils.py

@@ -485,6 +485,7 @@ class OrchestratorArgs:
     diffusion_attention_config: str | None = None
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
+    diffusion_compile_vae: bool | None = None
     cache_backend: str = "none"
     cache_config: str | None = None
     enable_cache_dit_summary: bool = False

vllm_omni/engine/async_omni_engine.py

@@ -1043,6 +1043,9 @@ class AsyncOmniEngine:
             "diffusion_compile_dynamic": (
                 True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"]
             ),
+            "diffusion_compile_vae": (
+                True if kwargs.get("diffusion_compile_vae") is None else kwargs["diffusion_compile_vae"]
+            ),
             "fa_deterministic": bool(kwargs.get("fa_deterministic", False)),
             "boundary_ratio": kwargs.get("boundary_ratio", None),
             "flow_shift": kwargs.get("flow_shift", None),

vllm_omni/entrypoints/cli/serve.py

@@ -446,6 +446,15 @@ class OmniServeCommand(CLISubcommand):
                 "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
             ),
         )
+        omni_config_group.add_argument(
+            "--diffusion-compile-vae",
+            action=argparse.BooleanOptionalAction,
+            default=None,
+            help=(
+                "Compile discovered VAE decode callables when diffusion torch.compile is enabled. "
+                "Use --no-diffusion-compile-vae to keep VAE decode eager."
+            ),
+        )
         omni_config_group.add_argument(
             "--fa-deterministic",
             dest="fa_deterministic",