feat(diffusion): add request-gated VAE fast paths

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

tests/config/test_omni_config.py

@@ -751,7 +751,6 @@ 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",
             ]
         )
     )
@@ -763,16 +762,13 @@ 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():
@@ -780,11 +776,6 @@ 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/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 compile_callable_with_fallback, regionally_compile
+from vllm_omni.diffusion.compile import regionally_compile
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -74,50 +74,3 @@ 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_config_propagation.py

@@ -113,6 +113,14 @@ class TestCreateDefaultDiffusion:
         stages = StageConfigFactory.create_default_diffusion({"model": "x"})
         assert stages[0]["engine_args"]["cache_backend"] == "none"
 
+    def test_vae_decode_precision_roundtrip(self):
+        od = _roundtrip_diffusion_config(model="x", vae_decode_precision="bf16")
+        assert od.vae_decode_precision == "bf16"
+
+    def test_invalid_vae_decode_precision_is_rejected(self):
+        with pytest.raises(ValueError, match="vae_decode_precision"):
+            OmniDiffusionConfig(model="x", vae_decode_precision="int8")
+
     def test_single_gpu_default_devices(self):
         stages = StageConfigFactory.create_default_diffusion({"model": "x"})
         assert stages[0]["runtime"]["devices"] == "0"

tests/diffusion/test_diffusion_engine.py

@@ -515,19 +515,16 @@ 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"):
@@ -537,10 +534,6 @@ 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,23 +130,6 @@ 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,
@@ -337,31 +320,6 @@ 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):
@@ -554,6 +512,35 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
     assert runner.pipeline.last_req.num_reqs == 1
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("quality", "expected_enabled"),
+    [(None, False), ("lossless", False), ("high", True)],
+)
+def test_execute_model_scopes_vae_fast_path_to_request_quality(monkeypatch, quality, expected_enabled):
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    req = _make_request()
+    req.sampling_params.quality = quality
+    gate_calls = []
+
+    @contextmanager
+    def _record_gate(pipeline, enabled):
+        gate_calls.append((pipeline, enabled, "enter"))
+        yield
+        gate_calls.append((pipeline, enabled, "exit"))
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "use_pipeline_vae_fast_path", _record_gate)
+
+    DiffusionModelRunner.execute_model(runner, req)
+
+    assert gate_calls == [
+        (runner.pipeline, expected_enabled, "enter"),
+        (runner.pipeline, expected_enabled, "exit"),
+    ]
+
+
 class _BatchPipeline:
     """Pipeline returning a configurable list of outputs from forward()."""
 
@@ -621,6 +608,18 @@ def test_execute_model_batch_rejects_output_count_mismatch(monkeypatch):
         DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_batch_rejects_mixed_vae_quality() -> None:
+    runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")]))
+    sched = _make_scheduler_output(num_reqs=2)
+    sched.scheduled_new_reqs[0].req.sampling_params.quality = "lossless"
+    sched.scheduled_new_reqs[1].req.sampling_params.quality = "high"
+
+    with pytest.raises(RuntimeError, match="cannot mix request quality levels"):
+        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_model_batch_routes_one_output_per_request(monkeypatch):
@@ -770,6 +769,19 @@ 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):
+    memory_pool_active = False
+    optimization_pool_states = []
+
+    @contextmanager
+    def _memory_pool_context(tag):
+        nonlocal memory_pool_active
+        assert tag == "weights"
+        memory_pool_active = True
+        try:
+            yield
+        finally:
+            memory_pool_active = False
+
     class _DummyLoader:
         def __init__(self, load_config, od_config=None):
             del load_config, od_config
@@ -817,16 +829,22 @@ 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,
+        "optimize_pipeline_vaes",
+        lambda pipeline, od_config: optimization_pool_states.append(memory_pool_active),
+    )
     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
     )
 
-    DiffusionModelRunner.load_model(runner)
+    DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context)
 
     assert runner.cache_backend is None
     assert runner.od_config.cache_backend is None
     assert dummy_cache_backend.enabled is False
+    assert optimization_pool_states == [True]
 
 
 @pytest.mark.core_model

tests/diffusion/test_diffusion_scheduler.py

@@ -971,44 +971,6 @@ 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/diffusion/test_diffusion_worker.py

@@ -102,6 +102,7 @@ class TestDiffusionWorkerSleep:
 
     def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
+        clear_vae_caches = mocker.patch("vllm_omni.diffusion.vae_optimizations.clear_pipeline_vae_fast_path_caches")
         mock_allocator_class = patch_cumem_allocator(mocker)
         mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
         mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
@@ -128,6 +129,7 @@ class TestDiffusionWorkerSleep:
 
         # Verify sleep was called with correct tags
         mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
+        clear_vae_caches.assert_called_once_with(mock_gpu_worker.model_runner.pipeline)
         assert bool(result) is True
         # Verify buffers were NOT saved (level 1 doesn't save buffers)
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0

tests/diffusion/vae_optimizations/test_vae_fastpaths.py

@@ -0,0 +1,300 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness contracts for request-gated VAE decoder fast paths."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.vae_optimizations.gate import (
+    VaeFastPathGate,
+    register_vae_fast_path_gate,
+    use_pipeline_vae_fast_path,
+    use_vae_fast_path,
+)
+
+pytestmark = [pytest.mark.diffusion]
+
+
+class _PipelineWithVae(nn.Module):
+    def __init__(self, vae: nn.Module) -> None:
+        super().__init__()
+        self.vae = vae
+
+
+class _DistributedExecutorConfig:
+    def __init__(self, parallel_mode: str) -> None:
+        self.parallel_mode = parallel_mode
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_vae_fast_path_gate_restores_nested_state() -> None:
+    vae = nn.Identity()
+    gate = VaeFastPathGate()
+    register_vae_fast_path_gate(vae, gate)
+
+    with use_vae_fast_path(vae, True):
+        assert gate.enabled
+        with use_vae_fast_path(vae, False):
+            assert not gate.enabled
+        assert gate.enabled
+
+    assert not gate.enabled
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_pipeline_vae_fast_path_controls_discovered_vae() -> None:
+    vae = nn.Identity()
+    pipeline = _PipelineWithVae(vae)
+    gate = VaeFastPathGate()
+    register_vae_fast_path_gate(vae, gate)
+
+    with use_pipeline_vae_fast_path(pipeline, True):
+        assert gate.enabled
+
+    assert not gate.enabled
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_image_group_norm_fast_path_uses_compile_safe_fallback(monkeypatch) -> None:
+    from vllm_omni.diffusion.vae_optimizations import image
+
+    gate = VaeFastPathGate()
+    gate.enabled = True
+    norm = image.FusedGroupNormSiLU(nn.GroupNorm(2, 4), gate)
+    inputs = torch.randn(1, 4, 3, 3)
+
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+    monkeypatch.setattr(
+        image,
+        "group_norm_silu_4d",
+        lambda *args, **kwargs: pytest.fail("eager Triton kernel entered during compilation"),
+    )
+
+    actual = norm(inputs)
+    expected = F.silu(F.group_norm(inputs, 2, norm.weight, norm.bias, norm.eps))
+    torch.testing.assert_close(actual, expected)
+
+
+def _small_image_vae() -> nn.Module:
+    from diffusers.models.autoencoders import AutoencoderKL
+
+    return AutoencoderKL(
+        in_channels=3,
+        out_channels=3,
+        down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"),
+        up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"),
+        block_out_channels=(32, 32),
+        layers_per_block=1,
+        latent_channels=4,
+        norm_num_groups=8,
+        sample_size=16,
+    )
+
+
+def _small_wan_vae(vae_cls: type[nn.Module] | None = None) -> nn.Module:
+    if vae_cls is None:
+        from diffusers.models.autoencoders import AutoencoderKLWan
+
+        vae_cls = AutoencoderKLWan
+
+    return vae_cls(
+        base_dim=8,
+        decoder_base_dim=8,
+        z_dim=4,
+        dim_mult=[1, 2],
+        num_res_blocks=1,
+        attn_scales=[],
+        temperal_downsample=[False],
+        latents_mean=[0.0] * 4,
+        latents_std=[1.0] * 4,
+        scale_factor_temporal=1,
+        scale_factor_spatial=2,
+    )
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+def test_image_vae_install_preserves_lossless_path_and_state_dict() -> None:
+    from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches
+    from vllm_omni.diffusion.vae_optimizations.image import (
+        FusedGroupNormSiLU,
+        FusedUpsample2xConv2d,
+        maybe_optimize_image_vae,
+    )
+
+    torch.manual_seed(0)
+    vae = _small_image_vae().to(device="cuda", dtype=torch.bfloat16).eval()
+    latents = torch.randn(1, 4, 8, 8, device="cuda", dtype=torch.bfloat16)
+    parameter_names = {name for name, _ in vae.named_parameters()}
+    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
+    reference = vae.decode(latents).sample
+
+    maybe_optimize_image_vae(vae)
+    maybe_optimize_image_vae(vae)
+
+    assert {name for name, _ in vae.named_parameters()} == parameter_names
+    vae.load_state_dict(state_dict, strict=True)
+    assert any(isinstance(module, FusedGroupNormSiLU) for module in vae.modules())
+    assert any(isinstance(module, FusedUpsample2xConv2d) for module in vae.modules())
+    assert torch.equal(vae.decode(latents).sample, reference)
+
+    with use_vae_fast_path(vae, True):
+        fast = vae.decode(latents).sample
+
+    torch.testing.assert_close(fast.float(), reference.float(), atol=0.1, rtol=0)
+    assert any(
+        buffer is not None
+        for module in vae.modules()
+        for name, buffer in module._buffers.items()
+        if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
+    )
+    clear_pipeline_vae_fast_path_caches(_PipelineWithVae(vae))
+    assert all(
+        buffer is None
+        for module in vae.modules()
+        for name, buffer in module._buffers.items()
+        if name in {"_fused_weight", "_vllm_folded_value_weight", "_vllm_folded_value_bias"}
+    )
+    with use_vae_fast_path(vae, True):
+        torch.testing.assert_close(vae.decode(latents).sample.float(), reference.float(), atol=0.1, rtol=0)
+    assert torch.equal(vae.decode(latents).sample, reference)
+    assert vae.decoder.conv_in.weight.is_contiguous()
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_channels_last_group_norm_silu_matches_reference(dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import group_norm_silu_4d
+
+    torch.manual_seed(0)
+    x = torch.randn(1, 64, 16, 16, device="cuda", dtype=dtype).contiguous(memory_format=torch.channels_last)
+    weight = torch.randn(64, device="cuda", dtype=dtype)
+    bias = torch.randn(64, device="cuda", dtype=dtype)
+
+    actual = group_norm_silu_4d(x, weight, bias, num_groups=8, eps=1e-5)
+    expected = F.silu(F.group_norm(x, 8, weight, bias))
+
+    assert actual is not None
+    atol, rtol = (1e-5, 1e-5) if dtype == torch.float32 else (7e-2, 2e-2)
+    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+@pytest.mark.parametrize(
+    "activation_dtype,affine_dtype,atol,rtol",
+    [
+        (torch.float32, torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, torch.float32, 1.5e-1, 3e-2),
+    ],
+)
+def test_wan_rmsnorm_silu_matches_reference(
+    activation_dtype: torch.dtype,
+    affine_dtype: torch.dtype,
+    atol: float,
+    rtol: float,
+) -> None:
+    from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import wan_rmsnorm_silu
+
+    torch.manual_seed(0)
+    x = torch.randn(1, 96, 3, 10, 14, device="cuda", dtype=activation_dtype).contiguous(
+        memory_format=torch.channels_last_3d
+    )
+    gamma = torch.randn(96, 1, 1, 1, device="cuda", dtype=affine_dtype)
+    bias = torch.randn_like(gamma)
+    normalized = F.normalize(x.float() if activation_dtype != torch.float32 else x, dim=1).to(activation_dtype)
+    expected = F.silu(normalized * 96**0.5 * gamma + bias)
+
+    actual = wan_rmsnorm_silu(x, gamma, bias)
+
+    assert actual is not None
+    assert actual.stride() == x.stride()
+    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+def test_wan_vae_install_preserves_lossless_path_and_state_dict() -> None:
+    from vllm_omni.diffusion.vae_optimizations.wan import (
+        FusedWanRMSNormSiLU,
+        maybe_optimize_wan_vae,
+    )
+
+    torch.manual_seed(0)
+    vae = _small_wan_vae().cuda().eval()
+    latents = torch.randn(1, 4, 1, 8, 8, device="cuda")
+    parameter_names = {name for name, _ in vae.named_parameters()}
+    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
+    reference = vae.decode(latents).sample
+
+    maybe_optimize_wan_vae(vae)
+    maybe_optimize_wan_vae(vae)
+
+    assert {name for name, _ in vae.named_parameters()} == parameter_names
+    vae.load_state_dict(state_dict, strict=True)
+    assert any(isinstance(module, FusedWanRMSNormSiLU) for module in vae.modules())
+    assert torch.equal(vae.decode(latents).sample, reference)
+
+    with use_vae_fast_path(vae, True):
+        fast = vae.decode(latents).sample
+
+    torch.testing.assert_close(fast, reference, atol=5e-4, rtol=1e-4)
+    assert torch.equal(vae.decode(latents).sample, reference)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("parallel_mode", ["auto", "spatial_shard_height", "spatial_shard_width"])
+def test_wan_fast_path_skips_spatial_shard_mode(monkeypatch: pytest.MonkeyPatch, parallel_mode: str) -> None:
+    import vllm_omni.diffusion.vae_optimizations.wan as wan_optimizations
+
+    monkeypatch.setattr(wan_optimizations, "_HAS_TRITON", True)
+    vae = _small_wan_vae()
+    vae.distributed_executor = _DistributedExecutorConfig(parallel_mode)
+
+    wan_optimizations.maybe_optimize_wan_vae(vae)
+
+    assert not any(isinstance(module, wan_optimizations.FusedWanRMSNormSiLU) for module in vae.modules())
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.no_grad()
+@pytest.mark.parametrize(
+    ("precision", "decode_dtype"),
+    [("fp32", torch.float32), ("fp16", torch.float16), ("bf16", torch.bfloat16)],
+)
+def test_wan_decode_precision_keeps_encode_fp32(precision: str, decode_dtype: torch.dtype) -> None:
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
+    from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision
+
+    torch.manual_seed(0)
+    vae = _small_wan_vae(DistributedAutoencoderKLWan).to(device="cuda", dtype=torch.bfloat16).eval()
+    configure_wan_decode_precision(vae, precision)
+
+    assert next(vae.encoder.parameters()).dtype == torch.float32
+    assert next(vae.quant_conv.parameters()).dtype == torch.float32
+    assert next(vae.decoder.parameters()).dtype == decode_dtype
+    assert next(vae.post_quant_conv.parameters()).dtype == decode_dtype
+
+    _, grid_spec = vae.tile_split(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=decode_dtype))
+    assert grid_spec.output_dtype == decode_dtype
+
+    encoded = vae.encode(torch.randn(1, 3, 1, 16, 16, device="cuda", dtype=torch.bfloat16))
+    decoded = vae.decode(torch.randn(1, 4, 1, 8, 8, device="cuda", dtype=torch.float32))
+
+    assert encoded.latent_dist.parameters.dtype == torch.float32
+    assert decoded.sample.dtype == decode_dtype

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -87,7 +87,6 @@ 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]
 
@@ -99,7 +98,6 @@ 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():
@@ -317,7 +315,7 @@ def test_serve_cli_forwards_distributed_offload_residency():
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():
-    """Ensure all compile controls reach the diffusion stage."""
+    """Ensure both compile controls reach the diffusion stage."""
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
@@ -330,7 +328,6 @@ def test_serve_cli_accepts_diffusion_compile_controls():
             "--diffusion-compile-granularity",
             "full",
             "--no-diffusion-compile-dynamic",
-            "--no-diffusion-compile-vae",
         ]
     )
 
@@ -339,10 +336,8 @@ 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():

tests/entrypoints/test_serve.py

@@ -76,6 +76,17 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     assert args.cfg_parallel_size == 4
 
 
+def test_serve_parser_accepts_wan_decode_precision() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-decode-precision", "bf16"])
+
+    assert args.vae_decode_precision == "bf16"
+    assert args.get_explicit_kwargs_dict()["vae_decode_precision"] == "bf16"
+
+
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",

vllm_omni/config/omni_config.py

@@ -570,10 +570,10 @@ 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
+    vae_decode_precision: Literal["fp32", "fp16", "bf16"] | None = None
     mask_strategy_file_path: str | None = None
     skip_time_steps: int = 15
     VSA_sparsity: float = 0.0

vllm_omni/config/stage_config.py

@@ -394,7 +394,6 @@ 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
@@ -402,6 +401,7 @@ class StageDeployConfig:
     step_execution: bool | None = None
     vae_use_slicing: bool | None = None
     vae_use_tiling: bool | None = None
+    vae_decode_precision: str | None = None
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/diffusion/compile.py

@@ -1,6 +1,5 @@
 # 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
@@ -10,59 +9,6 @@ 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

@@ -711,6 +711,9 @@ class OmniDiffusionConfig:
     # VAE memory optimization parameters
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
+    # Optional Wan decode-only precision. When set, the VAE encoder remains
+    # FP32 while decoder-side modules use the selected precision.
+    vae_decode_precision: str | None = None
 
     # STA (Sliding Tile Attention) parameters
     mask_strategy_file_path: str | None = None
@@ -726,9 +729,6 @@ 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
@@ -926,6 +926,11 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        if self.vae_decode_precision not in (None, "fp32", "fp16", "bf16"):
+            raise ValueError(
+                "vae_decode_precision must be one of {'fp32', 'fp16', 'bf16'} or None, "
+                f"got {self.vae_decode_precision!r}"
+            )
         if self.diffusion_compile_granularity not in {"regional", "full"}:
             raise ValueError(
                 "diffusion_compile_granularity must be 'regional' or 'full', "
@@ -933,8 +938,6 @@ 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,7 +56,6 @@ 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",
@@ -852,26 +851,28 @@ 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(
-        def build_request(request_id: str, prompt_text: str, *, real_path: bool) -> OmniDiffusionRequest:
+            prompt=prompt,
-            prompt: OmniTextPrompt = {"prompt": prompt_text}
+            request_id=DUMMY_DIFFUSION_REQUEST_ID,
-            if supports_image_input:
+            sampling_params=OmniDiffusionSamplingParams(
-                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,
@@ -884,47 +885,13 @@ 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

@@ -106,7 +106,7 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
             split_dims=(3, 4),
             grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
             tile_spec=tile_spec,
-            output_dtype=self.dtype,
+            output_dtype=getattr(self, "_vllm_decode_dtype", self.dtype),
         )
         return tiletask_list, grid_spec
 

vllm_omni/diffusion/vae_optimizations/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Optional VAE decoder optimizations."""
+
+from vllm_omni.diffusion.vae_optimizations.gate import (
+    use_pipeline_vae_fast_path,
+    use_vae_fast_path,
+)
+from vllm_omni.diffusion.vae_optimizations.optimize import (
+    clear_pipeline_vae_fast_path_caches,
+    optimize_pipeline_vaes,
+)
+
+__all__ = [
+    "clear_pipeline_vae_fast_path_caches",
+    "optimize_pipeline_vaes",
+    "use_pipeline_vae_fast_path",
+    "use_vae_fast_path",
+]

vllm_omni/diffusion/vae_optimizations/gate.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Decode-scoped gates for optional, non-bit-exact VAE fast paths."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import ExitStack, contextmanager
+from weakref import WeakKeyDictionary
+
+from torch import nn
+
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+
+class VaeFastPathGate:
+    """Mutable flag shared by wrappers installed on one VAE instance."""
+
+    __slots__ = ("enabled",)
+
+    def __init__(self) -> None:
+        self.enabled = False
+
+
+_VAE_FAST_PATH_GATES: WeakKeyDictionary[nn.Module, VaeFastPathGate] = WeakKeyDictionary()
+
+
+def register_vae_fast_path_gate(vae: nn.Module, gate: VaeFastPathGate) -> None:
+    """Associate an installed fast path with its owning VAE."""
+
+    _VAE_FAST_PATH_GATES[vae] = gate
+
+
+@contextmanager
+def use_vae_fast_path(vae: nn.Module, enabled: bool) -> Iterator[None]:
+    """Enable an installed fast path for one decode and always restore it."""
+
+    gate = _VAE_FAST_PATH_GATES.get(vae)
+    if gate is None:
+        yield
+        return
+
+    previous_enabled = gate.enabled
+    gate.enabled = enabled
+    try:
+        yield
+    finally:
+        gate.enabled = previous_enabled
+
+
+@contextmanager
+def use_pipeline_vae_fast_path(pipeline: nn.Module, enabled: bool) -> Iterator[None]:
+    """Apply one request's quality choice to every declared VAE component."""
+
+    vaes = ModuleDiscovery.discover(pipeline).vaes
+    with ExitStack() as stack:
+        for vae in vaes:
+            stack.enter_context(use_vae_fast_path(vae, enabled))
+        yield
+
+
+__all__ = [
+    "VaeFastPathGate",
+    "register_vae_fast_path_gate",
+    "use_pipeline_vae_fast_path",
+    "use_vae_fast_path",
+]

vllm_omni/diffusion/vae_optimizations/image.py

@@ -0,0 +1,354 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quality-gated CUDA fast paths for Diffusers-style 2D KL VAE decoders."""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.vae_optimizations.gate import (
+    VaeFastPathGate,
+    register_vae_fast_path_gate,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from vllm_omni.diffusion.vae_optimizations.triton_group_norm_silu import (
+        group_norm_silu_4d,
+        group_norm_silu_rows,
+    )
+
+    _HAS_TRITON = True
+except ImportError:  # pragma: no cover - exercised on builds without Triton
+    _HAS_TRITON = False
+
+
+class FusedGroupNormSiLU(nn.Module):
+    """Quality-gated GroupNorm + SiLU with an exact eager fallback."""
+
+    def __init__(self, norm: nn.GroupNorm, gate: VaeFastPathGate) -> None:
+        super().__init__()
+        self.num_groups = norm.num_groups
+        self.num_channels = norm.num_channels
+        self.eps = norm.eps
+        self.affine = norm.affine
+        # Register parameters directly so state-dict FQNs do not change.
+        self.weight = norm.weight
+        self.bias = norm.bias
+        self._vllm_gate = gate
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._vllm_gate.enabled and x.dim() == 4 and not torch.compiler.is_compiling():
+            output = group_norm_silu_4d(
+                x,
+                self.weight,
+                self.bias,
+                self.num_groups,
+                self.eps,
+                apply_silu=True,
+            )
+            if output is not None:
+                return output
+        return F.silu(F.group_norm(x, self.num_groups, self.weight, self.bias, self.eps))
+
+
+def _install_norm_silu(decoder: nn.Module, resnet_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
+    def is_fusable(norm: object) -> bool:
+        return type(norm) is nn.GroupNorm and norm.affine and norm.weight is not None and norm.bias is not None
+
+    count = 0
+    for module in decoder.modules():
+        if (
+            type(module) is resnet_cls
+            and module.time_emb_proj is None
+            and module.time_embedding_norm in ("default", "group")
+            and module.upsample is None
+            and module.downsample is None
+            and isinstance(module.nonlinearity, nn.SiLU)
+            and is_fusable(module.norm1)
+            and is_fusable(module.norm2)
+        ):
+            module.norm1 = FusedGroupNormSiLU(module.norm1, gate)
+            module.norm2 = FusedGroupNormSiLU(module.norm2, gate)
+            module.nonlinearity = nn.Identity()
+            count += 2
+
+    if (
+        type(getattr(decoder, "conv_norm_out", None)) is nn.GroupNorm
+        and isinstance(getattr(decoder, "conv_act", None), nn.SiLU)
+        and is_fusable(decoder.conv_norm_out)
+    ):
+        decoder.conv_norm_out = FusedGroupNormSiLU(decoder.conv_norm_out, gate)
+        decoder.conv_act = nn.Identity()
+        count += 1
+    return count
+
+
+_UPSAMPLE_TAP_MAP = {0: (2,), 1: (1, 2), 2: (0, 1), 3: (0,)}
+
+
+def _fold_upsample2x_conv2d_weight(conv: nn.Conv2d) -> torch.Tensor:
+    weight = conv.weight.detach().float()
+    output_channels, input_channels = weight.shape[:2]
+    folded = weight.new_zeros(input_channels, output_channels, 4, 4)
+    for row in range(4):
+        for column in range(4):
+            value = weight.new_zeros(output_channels, input_channels)
+            for source_row in _UPSAMPLE_TAP_MAP[row]:
+                for source_column in _UPSAMPLE_TAP_MAP[column]:
+                    value += weight[:, :, source_row, source_column]
+            folded[:, :, row, column] = value.t()
+    folded = folded.to(conv.weight.dtype)
+    if conv.weight.is_contiguous(memory_format=torch.channels_last):
+        folded = folded.contiguous(memory_format=torch.channels_last)
+    return folded.to(conv.weight.device)
+
+
+class FusedUpsample2xConv2d(nn.Module):
+    """Quality-gated nearest-2x + Conv2d replacement using ConvTranspose2d."""
+
+    def __init__(self, upsample: nn.Module, gate: VaeFastPathGate) -> None:
+        super().__init__()
+        # Keep only ``conv`` registered so checkpoint parameter names stay
+        # ``...upsamplers.N.conv.*``. The original module is an eager fallback.
+        object.__setattr__(self, "_original", upsample)
+        self.conv = upsample.conv
+        self.channels = upsample.channels
+        self._vllm_gate = gate
+        self.register_buffer("_fused_weight", None, persistent=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        output_size: int | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        if not self._vllm_gate.enabled or output_size is not None or hidden_states.shape[1] != self.channels:
+            return self._original(hidden_states, output_size=output_size, *args, **kwargs)
+
+        weight = self._fused_weight
+        if weight is None or weight.device != self.conv.weight.device or weight.dtype != self.conv.weight.dtype:
+            weight = _fold_upsample2x_conv2d_weight(self.conv)
+            self._fused_weight = weight
+        return F.conv_transpose2d(hidden_states, weight, self.conv.bias, stride=2, padding=1)
+
+
+def _install_fused_upsample(decoder: nn.Module, upsample_cls: type[nn.Module], gate: VaeFastPathGate) -> int:
+    count = 0
+    for block in decoder.up_blocks:
+        upsamplers = getattr(block, "upsamplers", None)
+        if not upsamplers:
+            continue
+        for index, upsample in enumerate(upsamplers):
+            if type(upsample) is not upsample_cls:
+                continue
+            conv = getattr(upsample, "conv", None)
+            if (
+                upsample.use_conv
+                and not upsample.use_conv_transpose
+                and upsample.interpolate
+                and upsample.norm is None
+                and upsample.name == "conv"
+                and type(conv) is nn.Conv2d
+                and conv.kernel_size == (3, 3)
+                and conv.stride == (1, 1)
+                and conv.padding == (1, 1)
+                and conv.dilation == (1, 1)
+                and conv.groups == 1
+                and conv.padding_mode == "zeros"
+            ):
+                upsamplers[index] = FusedUpsample2xConv2d(upsample, gate)
+                count += 1
+    return count
+
+
+def _fold_attention_value_projection(module: nn.Module) -> tuple[torch.Tensor, torch.Tensor]:
+    value_weight = module.to_v.weight.detach().float()
+    value_bias = module.to_v.bias.detach().float()
+    output_weight = module.to_out[0].weight.detach().float()
+    output_bias = module.to_out[0].bias.detach().float()
+    dtype = module.to_v.weight.dtype
+    return (
+        (output_weight @ value_weight).to(dtype).contiguous(),
+        (output_weight @ value_bias + output_bias).to(dtype),
+    )
+
+
+def _attention_fast_forward(
+    self: nn.Module,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor | None = None,
+    attention_mask: torch.Tensor | None = None,
+    temb: torch.Tensor | None = None,
+    **cross_attention_kwargs: Any,
+) -> torch.Tensor:
+    if (
+        not self._vllm_gate.enabled
+        or encoder_hidden_states is not None
+        or attention_mask is not None
+        or temb is not None
+        or hidden_states.ndim != 4
+    ):
+        return type(self).forward(
+            self,
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=attention_mask,
+            temb=temb,
+            **cross_attention_kwargs,
+        )
+
+    residual = hidden_states
+    batch_size, channels, height, width = hidden_states.shape
+    hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
+
+    if self.group_norm is not None:
+        norm = self.group_norm
+        normalized = None
+        if not torch.compiler.is_compiling():
+            normalized = group_norm_silu_rows(
+                hidden_states,
+                norm.weight,
+                norm.bias,
+                norm.num_groups,
+                norm.eps,
+                apply_silu=False,
+            )
+        hidden_states = normalized if normalized is not None else norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+    query = self.to_q(hidden_states)
+    key = self.to_k(hidden_states)
+    folded_weight = self._vllm_folded_value_weight
+    folded_bias = self._vllm_folded_value_bias
+    if (
+        folded_weight is None
+        or folded_bias is None
+        or folded_weight.device != hidden_states.device
+        or folded_weight.dtype != self.to_v.weight.dtype
+    ):
+        folded_weight, folded_bias = _fold_attention_value_projection(self)
+        self._vllm_folded_value_weight = folded_weight
+        self._vllm_folded_value_bias = folded_bias
+    value = F.linear(hidden_states, folded_weight, folded_bias)
+
+    output = F.scaled_dot_product_attention(query.unsqueeze(1), key.unsqueeze(1), value.unsqueeze(1))
+    output = output.squeeze(1).to(query.dtype)
+    output = self.to_out[1](output)
+    output = output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+    if self.residual_connection:
+        output = output + residual
+    return output / self.rescale_output_factor
+
+
+def _attention_fast_compatible(module: nn.Module, attention_cls: type, processor_cls: type) -> bool:
+    return (
+        type(module) is attention_cls
+        and isinstance(module.processor, processor_cls)
+        and module.heads == 1
+        and module.scale_qk
+        and module.spatial_norm is None
+        and not module.norm_cross
+        and getattr(module, "norm_q", None) is None
+        and getattr(module, "norm_k", None) is None
+        and getattr(module, "add_k_proj", None) is None
+        and type(module.to_q) is nn.Linear
+        and type(module.to_k) is nn.Linear
+        and type(module.to_v) is nn.Linear
+        and type(module.to_out[0]) is nn.Linear
+        and isinstance(module.to_out[1], nn.Dropout)
+        and module.to_v.bias is not None
+        and module.to_out[0].bias is not None
+    )
+
+
+def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
+    use_channels_last = self._vllm_gate.enabled
+    if use_channels_last != self._vllm_channels_last:
+        memory_format = torch.channels_last if use_channels_last else torch.contiguous_format
+        self.to(memory_format=memory_format)
+        self._vllm_channels_last = use_channels_last
+        logger.debug(
+            "%s decoder switched to %s layout",
+            self._vllm_label,
+            "channels_last" if use_channels_last else "contiguous",
+        )
+    return type(self).forward(self, *args, **kwargs)
+
+
+def maybe_optimize_image_vae(vae: nn.Module) -> nn.Module:
+    """Install image VAE decoder fast paths when the structure is fully supported."""
+
+    if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False):
+        return vae
+
+    from diffusers.models.attention_processor import Attention, AttnProcessor2_0
+    from diffusers.models.autoencoders import AutoencoderKL
+    from diffusers.models.autoencoders.vae import Decoder
+    from diffusers.models.resnet import ResnetBlock2D
+    from diffusers.models.upsampling import Upsample2D
+
+    supported_types: tuple[type[nn.Module], ...] = (AutoencoderKL,)
+    try:
+        from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
+
+        supported_types += (AutoencoderKLFlux2,)
+    except ImportError:  # pragma: no cover - older supported Diffusers builds
+        pass
+
+    if not isinstance(vae, supported_types):
+        return vae
+    decoder = getattr(vae, "decoder", None)
+    if type(decoder) is not Decoder:
+        return vae
+
+    attention_modules = [
+        module for module in decoder.modules() if _attention_fast_compatible(module, Attention, AttnProcessor2_0)
+    ]
+    attention_count = sum(isinstance(module, Attention) for module in decoder.modules())
+    if len(attention_modules) != attention_count:
+        logger.info(
+            "%s has %d/%d attention blocks without a layout-safe rewrite; skipping VAE fast paths",
+            type(vae).__name__,
+            attention_count - len(attention_modules),
+            attention_count,
+        )
+        return vae
+
+    gate = VaeFastPathGate()
+    decoder._vllm_gate = gate
+    decoder._vllm_label = type(vae).__name__
+    decoder._vllm_channels_last = False
+    decoder.forward = MethodType(_decoder_layout_forward, decoder)
+    upsample_count = _install_fused_upsample(decoder, Upsample2D, gate)
+    for module in attention_modules:
+        module._vllm_gate = gate
+        module.register_buffer("_vllm_folded_value_weight", None, persistent=False)
+        module.register_buffer("_vllm_folded_value_bias", None, persistent=False)
+        module.forward = MethodType(_attention_fast_forward, module)
+    norm_count = _install_norm_silu(decoder, ResnetBlock2D, gate)
+    register_vae_fast_path_gate(vae, gate)
+    vae._vllm_vae_fast_path_installed = True
+    logger.info(
+        "%s installed quality-gated VAE decoder fast paths "
+        "(%d upsamplers, %d attention blocks, %d GroupNorm+SiLU fusions)",
+        type(vae).__name__,
+        upsample_count,
+        len(attention_modules),
+        norm_count,
+    )
+    return vae
+
+
+__all__ = [
+    "FusedGroupNormSiLU",
+    "FusedUpsample2xConv2d",
+    "maybe_optimize_image_vae",
+]

vllm_omni/diffusion/vae_optimizations/optimize.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pipeline-level installation for optional VAE optimizations."""
+
+from __future__ import annotations
+
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.vae_optimizations.image import maybe_optimize_image_vae
+from vllm_omni.diffusion.vae_optimizations.precision import configure_wan_decode_precision
+from vllm_omni.diffusion.vae_optimizations.wan import maybe_optimize_wan_vae
+from vllm_omni.platforms import current_omni_platform
+
+logger = init_logger(__name__)
+
+_FOLDED_CACHE_NAMES = (
+    "_fused_weight",
+    "_vllm_folded_value_weight",
+    "_vllm_folded_value_bias",
+)
+
+
+def optimize_pipeline_vaes(pipeline: nn.Module, od_config: OmniDiffusionConfig) -> None:
+    """Configure declared VAE components once after checkpoint loading."""
+
+    precision = getattr(od_config, "vae_decode_precision", None)
+    vaes = ModuleDiscovery.discover(pipeline).vaes
+    for vae in vaes:
+        # An explicit precision override is a serving contract, so failures on
+        # a recognized Wan VAE must abort startup instead of being hidden.
+        configure_wan_decode_precision(vae, precision)
+
+    if not current_omni_platform.is_cuda():
+        return
+
+    for vae in vaes:
+        for optimizer in (maybe_optimize_image_vae, maybe_optimize_wan_vae):
+            try:
+                optimizer(vae)
+            except Exception:
+                # A wrapper may already have installed its exact fallback
+                # before a later optional rewrite failed. With no registered
+                # enabled gate, those wrappers remain on the reference path.
+                logger.warning(
+                    "Failed to install optional VAE fast paths on %s; using the reference path",
+                    type(vae).__name__,
+                    exc_info=True,
+                )
+
+
+def clear_pipeline_vae_fast_path_caches(pipeline: nn.Module) -> None:
+    """Release lazily folded buffers before the worker enters sleep mode."""
+
+    cleared = 0
+    for vae in ModuleDiscovery.discover(pipeline).vaes:
+        if not getattr(vae, "_vllm_vae_fast_path_installed", False):
+            continue
+        for module in vae.modules():
+            for name in _FOLDED_CACHE_NAMES:
+                if name in module._buffers and module._buffers[name] is not None:
+                    setattr(module, name, None)
+                    cleared += 1
+    if cleared:
+        logger.debug("Released %d lazy VAE fast-path buffers before sleep", cleared)
+
+
+__all__ = ["clear_pipeline_vae_fast_path_caches", "optimize_pipeline_vaes"]

vllm_omni/diffusion/vae_optimizations/precision.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Decode-only precision configuration for structurally split Wan VAEs."""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+import torch
+from torch import nn
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+PRECISION_TO_DTYPE = {
+    "fp32": torch.float32,
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+}
+
+
+def _decode_with_precision(
+    self: nn.Module,
+    latents: torch.Tensor,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return type(self).decode(
+        self,
+        latents.to(dtype=self._vllm_decode_dtype),
+        *args,
+        **kwargs,
+    )
+
+
+def _encode_in_fp32(
+    self: nn.Module,
+    samples: torch.Tensor,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    return type(self).encode(
+        self,
+        samples.to(dtype=torch.float32),
+        *args,
+        **kwargs,
+    )
+
+
+def configure_wan_decode_precision(vae: nn.Module, precision: str | None) -> nn.Module:
+    """Keep Wan encode in FP32 while materializing only decode modules lower."""
+
+    if precision is None:
+        return vae
+
+    from diffusers.models.autoencoders import AutoencoderKLWan
+
+    if not isinstance(vae, AutoencoderKLWan):
+        return vae
+    try:
+        decode_dtype = PRECISION_TO_DTYPE[precision]
+    except KeyError as exc:
+        raise ValueError(
+            f"Unsupported vae_decode_precision={precision!r}; expected one of {sorted(PRECISION_TO_DTYPE)}"
+        ) from exc
+
+    encoder = getattr(vae, "encoder", None)
+    quant_conv = getattr(vae, "quant_conv", None)
+    decoder = getattr(vae, "decoder", None)
+    post_quant_conv = getattr(vae, "post_quant_conv", None)
+    if not all(isinstance(module, nn.Module) for module in (encoder, decoder, post_quant_conv)):
+        logger.info("Wan VAE lacks separate encoder/decoder modules; skipping decode-only precision")
+        return vae
+
+    encoder.to(dtype=torch.float32)
+    if isinstance(quant_conv, nn.Module):
+        quant_conv.to(dtype=torch.float32)
+    decoder.to(dtype=decode_dtype)
+    post_quant_conv.to(dtype=decode_dtype)
+    vae.encode = MethodType(_encode_in_fp32, vae)
+    vae._vllm_decode_dtype = decode_dtype
+    vae.decode = MethodType(_decode_with_precision, vae)
+    logger.info(
+        "Wan VAE configured with FP32 encode and %s decode",
+        precision,
+    )
+    return vae
+
+
+__all__ = ["PRECISION_TO_DTYPE", "configure_wan_decode_precision"]

vllm_omni/diffusion/vae_optimizations/triton_group_norm_silu.py

@@ -0,0 +1,237 @@
+# ruff: noqa: N803
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Channels-last two-pass GroupNorm(+SiLU) Triton kernels.
+
+The kernels keep the channel dimension innermost throughout a 2D VAE
+decoder. Statistics are accumulated in FP32, then a separate apply pass
+performs the affine transform and optional SiLU epilogue. Unsupported inputs
+return ``None`` so callers can fail closed to the reference implementation.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+_MAX_CHANNELS = 2048
+
+
+@triton.jit
+def _gn_partial_rows_kernel(
+    x_ptr,
+    psum_ptr,
+    psq_ptr,
+    rows,
+    rows_per_prog,
+    C: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+):
+    chunk = tl.program_id(0).to(tl.int64)
+    batch = tl.program_id(1).to(tl.int64)
+    num_chunks = tl.num_programs(0)
+    row0 = chunk * rows_per_prog
+    cols = tl.arange(0, C)
+    acc_sum = tl.zeros((C,), tl.float32)
+    acc_sq = tl.zeros((C,), tl.float32)
+    x_base = x_ptr + batch * rows * C
+    for row_offset in range(0, rows_per_prog, BLOCK_R):
+        row_indices = row0 + row_offset + tl.arange(0, BLOCK_R)
+        mask = row_indices < rows
+        x = tl.load(
+            x_base + row_indices[:, None] * C + cols[None, :],
+            mask=mask[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        acc_sum += tl.sum(x, 0)
+        acc_sq += tl.sum(x * x, 0)
+    output_offsets = (batch * num_chunks + chunk) * C + cols
+    tl.store(psum_ptr + output_offsets, acc_sum)
+    tl.store(psq_ptr + output_offsets, acc_sq)
+
+
+@triton.jit
+def _gn_finalize_kernel(
+    psum_ptr,
+    psq_ptr,
+    weight_ptr,
+    bias_ptr,
+    scale_shift_ptr,
+    num_chunks,
+    group_numel,
+    eps,
+    channels,
+    CPG: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    group = tl.program_id(0).to(tl.int64)
+    batch = tl.program_id(1).to(tl.int64)
+    cols = group * CPG + tl.arange(0, CPG)
+    acc_sum = tl.zeros((), tl.float32)
+    acc_sq = tl.zeros((), tl.float32)
+    for chunk0 in range(0, num_chunks, BLOCK_K):
+        chunks = chunk0 + tl.arange(0, BLOCK_K)
+        mask = chunks < num_chunks
+        offsets = (batch * num_chunks + chunks)[:, None] * channels + cols[None, :]
+        acc_sum += tl.sum(tl.load(psum_ptr + offsets, mask=mask[:, None], other=0.0))
+        acc_sq += tl.sum(tl.load(psq_ptr + offsets, mask=mask[:, None], other=0.0))
+    mean = acc_sum / group_numel
+    variance = acc_sq / group_numel - mean * mean
+    variance = tl.maximum(variance, 0.0)
+    reciprocal_std = tl.rsqrt(variance + eps)
+    weight = tl.load(weight_ptr + cols).to(tl.float32)
+    bias = tl.load(bias_ptr + cols).to(tl.float32)
+    scale = weight * reciprocal_std
+    shift = bias - mean * scale
+    tl.store(scale_shift_ptr + batch * 2 * channels + cols, scale)
+    tl.store(scale_shift_ptr + (batch * 2 + 1) * channels + cols, shift)
+
+
+@triton.jit
+def _gn_apply_rows_kernel(
+    x_ptr,
+    scale_shift_ptr,
+    y_ptr,
+    rows,
+    C: tl.constexpr,
+    BLOCK_R: tl.constexpr,
+    SILU: tl.constexpr,
+):
+    program = tl.program_id(0).to(tl.int64)
+    batch = tl.program_id(1).to(tl.int64)
+    cols = tl.arange(0, C)
+    scale = tl.load(scale_shift_ptr + batch * 2 * C + cols)
+    shift = tl.load(scale_shift_ptr + (batch * 2 + 1) * C + cols)
+    row_indices = program * BLOCK_R + tl.arange(0, BLOCK_R)
+    mask = row_indices < rows
+    offsets = batch * rows * C + row_indices[:, None] * C + cols[None, :]
+    x = tl.load(x_ptr + offsets, mask=mask[:, None], other=0.0).to(tl.float32)
+    y = x * scale[None, :] + shift[None, :]
+    if SILU:
+        y = y * tl.sigmoid(y)
+    tl.store(y_ptr + offsets, y.to(y_ptr.dtype.element_ty), mask=mask[:, None])
+
+
+def _group_norm_silu_rows(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+    apply_silu: bool,
+) -> torch.Tensor:
+    batch_size, rows, channels = x.shape
+    channels_per_group = channels // num_groups
+    block_rows = max(1, 8192 // channels)
+    rows_per_program = block_rows * 32
+    num_chunks = triton.cdiv(rows, rows_per_program)
+    partial_sum = torch.empty((batch_size, num_chunks, channels), device=x.device, dtype=torch.float32)
+    partial_sq = torch.empty_like(partial_sum)
+    with torch.cuda.device(x.device):
+        _gn_partial_rows_kernel[(num_chunks, batch_size)](
+            x,
+            partial_sum,
+            partial_sq,
+            rows,
+            rows_per_program,
+            C=channels,
+            BLOCK_R=block_rows,
+            num_warps=4,
+        )
+        scale_shift = torch.empty((batch_size, 2, channels), device=x.device, dtype=torch.float32)
+        block_chunks = max(1, min(4096 // channels_per_group, triton.next_power_of_2(num_chunks)))
+        _gn_finalize_kernel[(num_groups, batch_size)](
+            partial_sum,
+            partial_sq,
+            weight,
+            bias,
+            scale_shift,
+            num_chunks,
+            rows * channels_per_group,
+            eps,
+            channels,
+            CPG=channels_per_group,
+            BLOCK_K=block_chunks,
+            num_warps=4,
+        )
+        output = torch.empty_like(x)
+        _gn_apply_rows_kernel[(triton.cdiv(rows, block_rows), batch_size)](
+            x,
+            scale_shift,
+            output,
+            rows,
+            C=channels,
+            BLOCK_R=block_rows,
+            SILU=apply_silu,
+            num_warps=4,
+        )
+    return output
+
+
+def _is_supported(
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    num_groups: int,
+) -> bool:
+    if not (x.is_cuda and x.numel() > 0 and not torch.is_grad_enabled()):
+        return False
+    if x.requires_grad or x.dtype not in _SUPPORTED_DTYPES:
+        return False
+    if weight is None or bias is None:
+        return False
+    channels = x.shape[1] if x.dim() == 4 else x.shape[-1]
+    if weight.shape != (channels,) or bias.shape != (channels,):
+        return False
+    if not (
+        weight.device == bias.device == x.device
+        and weight.dtype == bias.dtype == x.dtype
+        and weight.is_contiguous()
+        and bias.is_contiguous()
+    ):
+        return False
+    if num_groups < 1 or channels % num_groups != 0:
+        return False
+    return triton.next_power_of_2(channels) == channels and channels <= _MAX_CHANNELS
+
+
+def group_norm_silu_4d(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+    apply_silu: bool = True,
+) -> torch.Tensor | None:
+    """Run GroupNorm(+SiLU) on a channels-last ``(N, C, H, W)`` tensor."""
+
+    if x.dim() != 4 or not _is_supported(x, weight, bias, num_groups):
+        return None
+    batch_size, channels, height, width = x.shape
+    if not (channels > 1 and (height > 1 or width > 1) and x.is_contiguous(memory_format=torch.channels_last)):
+        return None
+    rows = x.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
+    output = _group_norm_silu_rows(rows, weight, bias, num_groups, eps, apply_silu)
+    return output.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+
+
+def group_norm_silu_rows(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+    apply_silu: bool = True,
+) -> torch.Tensor | None:
+    """Run GroupNorm(+SiLU) over contiguous ``(N, rows, C)`` input."""
+
+    if x.dim() != 3 or not x.is_contiguous():
+        return None
+    if not _is_supported(x, weight, bias, num_groups):
+        return None
+    return _group_norm_silu_rows(x, weight, bias, num_groups, eps, apply_silu)
+
+
+__all__ = ["group_norm_silu_4d", "group_norm_silu_rows"]

vllm_omni/diffusion/vae_optimizations/triton_wan_rmsnorm_silu.py

@@ -0,0 +1,160 @@
+# ruff: noqa: N803
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Channels-last-3D Wan VAE RMSNorm(+SiLU) Triton kernel."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+_MAX_CHANNELS = 1024
+
+
+@triton.jit
+def _wan_rmsnorm_silu_kernel(
+    x_ptr,
+    gamma_ptr,
+    bias_ptr,
+    output_ptr,
+    channels: tl.constexpr,
+    time_size,
+    height_size,
+    width_size,
+    x_stride_batch,
+    x_stride_channel,
+    x_stride_time,
+    x_stride_height,
+    x_stride_width,
+    output_stride_batch,
+    output_stride_channel,
+    output_stride_time,
+    output_stride_height,
+    output_stride_width,
+    rms_scale,
+    eps,
+    has_bias: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    channel_offsets = tl.arange(0, BLOCK_C)
+    channel_mask = channel_offsets < channels
+
+    width = row % width_size
+    remainder = row // width_size
+    height = remainder % height_size
+    remainder = remainder // height_size
+    time = remainder % time_size
+    batch = remainder // time_size
+
+    x_base = batch * x_stride_batch + time * x_stride_time + height * x_stride_height + width * x_stride_width
+    output_base = (
+        batch * output_stride_batch
+        + time * output_stride_time
+        + height * output_stride_height
+        + width * output_stride_width
+    )
+
+    x = tl.load(
+        x_ptr + x_base + channel_offsets * x_stride_channel,
+        mask=channel_mask,
+        other=0.0,
+    ).to(tl.float32)
+    norm = tl.sqrt(tl.sum(x * x, axis=0))
+    inverse_norm = 1.0 / tl.maximum(norm, eps)
+
+    # Match the eager dtype boundaries: normalize and scale in activation
+    # dtype, affine in the promoted dtype, then SiLU in FP32.
+    y = (x * inverse_norm).to(x_ptr.dtype.element_ty)
+    gamma = tl.load(gamma_ptr + channel_offsets, mask=channel_mask, other=1.0)
+    y = (y * rms_scale).to(x_ptr.dtype.element_ty)
+    y = (y.to(tl.float32) * gamma.to(tl.float32)).to(output_ptr.dtype.element_ty)
+    if has_bias:
+        bias = tl.load(bias_ptr + channel_offsets, mask=channel_mask, other=0.0)
+        y = (y.to(tl.float32) + bias.to(tl.float32)).to(output_ptr.dtype.element_ty)
+    y = y.to(tl.float32)
+    y = y * tl.sigmoid(y)
+    tl.store(
+        output_ptr + output_base + channel_offsets * output_stride_channel,
+        y,
+        mask=channel_mask,
+    )
+
+
+def _affine_supported(x: torch.Tensor, affine: torch.Tensor) -> bool:
+    return (
+        affine.is_cuda
+        and affine.device == x.device
+        and (affine.dtype == x.dtype or affine.dtype == torch.float32)
+        and affine.numel() == x.shape[1]
+    )
+
+
+def can_use_wan_rmsnorm_silu(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None,
+) -> bool:
+    """Return whether ``wan_rmsnorm_silu`` supports this exact input."""
+
+    return (
+        x.is_cuda
+        and not torch.is_grad_enabled()
+        and not x.requires_grad
+        and x.dtype in _SUPPORTED_DTYPES
+        and x.ndim == 5
+        and x.numel() > 0
+        and 0 < x.shape[1] <= _MAX_CHANNELS
+        and x.is_contiguous(memory_format=torch.channels_last_3d)
+        and _affine_supported(x, gamma)
+        and (bias is None or _affine_supported(x, bias))
+    )
+
+
+def wan_rmsnorm_silu(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    rms_scale: float | None = None,
+    eps: float = 1e-12,
+) -> torch.Tensor | None:
+    """Fuse ``SiLU(F.normalize(x, dim=1) * scale * gamma + bias)``."""
+
+    if not can_use_wan_rmsnorm_silu(x, gamma, bias):
+        return None
+
+    batch_size, channels, time_size, height_size, width_size = x.shape
+    output_dtype = torch.promote_types(x.dtype, gamma.dtype)
+    output = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=output_dtype)
+    block_channels = triton.next_power_of_2(channels)
+    num_warps = 1 if block_channels <= 64 else 4 if block_channels <= 512 else 8
+    gamma = gamma.reshape(channels).contiguous()
+    has_bias = bias is not None
+    bias = gamma if bias is None else bias.reshape(channels).contiguous()
+    if rms_scale is None:
+        rms_scale = channels**0.5
+
+    with torch.cuda.device(x.device):
+        _wan_rmsnorm_silu_kernel[(batch_size * time_size * height_size * width_size,)](
+            x,
+            gamma,
+            bias,
+            output,
+            channels,
+            time_size,
+            height_size,
+            width_size,
+            *x.stride(),
+            *output.stride(),
+            float(rms_scale),
+            eps,
+            has_bias,
+            BLOCK_C=block_channels,
+            num_warps=num_warps,
+        )
+    return output
+
+
+__all__ = ["can_use_wan_rmsnorm_silu", "wan_rmsnorm_silu"]

vllm_omni/diffusion/vae_optimizations/wan.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quality-gated CUDA fast path for Diffusers Wan VAE decoders."""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.vae_optimizations.gate import (
+    VaeFastPathGate,
+    register_vae_fast_path_gate,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from vllm_omni.diffusion.vae_optimizations.triton_wan_rmsnorm_silu import (
+        wan_rmsnorm_silu,
+    )
+
+    _HAS_TRITON = True
+except ImportError:  # pragma: no cover - exercised on builds without Triton
+    _HAS_TRITON = False
+
+
+class FusedWanRMSNormSiLU(nn.Module):
+    """Quality-gated Wan RMSNorm + SiLU with an exact eager fallback."""
+
+    def __init__(self, norm: nn.Module, gate: VaeFastPathGate) -> None:
+        super().__init__()
+        self.gamma = norm.gamma
+        self.bias = norm.bias
+        self.scale = float(norm.scale)
+        self._vllm_gate = gate
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._vllm_gate.enabled and not torch.compiler.is_compiling():
+            bias = self.bias if isinstance(self.bias, torch.Tensor) else None
+            output = wan_rmsnorm_silu(x, self.gamma, bias, rms_scale=self.scale)
+            if output is not None:
+                return output
+
+        needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any(
+            token in str(x.dtype) for token in ("float4_", "float8_")
+        )
+        normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=1).to(x.dtype)
+        return F.silu(normalized * self.scale * self.gamma + self.bias)
+
+
+def _plain_silu(activation: object) -> bool:
+    return isinstance(activation, nn.SiLU) and not activation.inplace
+
+
+def _fusable_norm(norm: object, norm_cls: type) -> bool:
+    return (
+        type(norm) is norm_cls
+        and getattr(norm, "channel_first", False)
+        and isinstance(getattr(norm, "gamma", None), torch.Tensor)
+    )
+
+
+def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanResidualBlock, WanRMS_norm
+
+    residual_blocks = [module for module in decoder.modules() if isinstance(module, WanResidualBlock)]
+    eligible = [
+        module
+        for module in residual_blocks
+        if type(module) is WanResidualBlock
+        and _plain_silu(module.nonlinearity)
+        and _fusable_norm(module.norm1, WanRMS_norm)
+        and _fusable_norm(module.norm2, WanRMS_norm)
+    ]
+    if len(eligible) != len(residual_blocks):
+        logger.info(
+            "Wan VAE has %d/%d non-standard residual blocks; skipping fast path",
+            len(residual_blocks) - len(eligible),
+            len(residual_blocks),
+        )
+        return None
+    if not (
+        _fusable_norm(getattr(decoder, "norm_out", None), WanRMS_norm)
+        and _plain_silu(getattr(decoder, "nonlinearity", None))
+    ):
+        logger.info("Wan VAE has a non-standard output head; skipping fast path")
+        return None
+
+    count = 0
+    for module in eligible:
+        module.norm1 = FusedWanRMSNormSiLU(module.norm1, gate)
+        module.norm2 = FusedWanRMSNormSiLU(module.norm2, gate)
+        module.nonlinearity = nn.Identity()
+        count += 2
+    decoder.norm_out = FusedWanRMSNormSiLU(decoder.norm_out, gate)
+    decoder.nonlinearity = nn.Identity()
+    return count + 1
+
+
+def _set_wan_decoder_memory_format(decoder: nn.Module, *, channels_last: bool) -> None:
+    for module in decoder.modules():
+        if isinstance(module, nn.Conv3d):
+            module.to(memory_format=torch.channels_last_3d if channels_last else torch.contiguous_format)
+        elif isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)):
+            module.to(memory_format=torch.channels_last if channels_last else torch.contiguous_format)
+
+
+def _decoder_layout_forward(self: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
+    use_channels_last = self._vllm_gate.enabled
+    if use_channels_last != self._vllm_channels_last:
+        _set_wan_decoder_memory_format(self, channels_last=use_channels_last)
+        self._vllm_channels_last = use_channels_last
+        logger.debug(
+            "Wan VAE decoder switched to %s layout",
+            "channels_last_3d" if use_channels_last else "contiguous",
+        )
+    return type(self).forward(self, *args, **kwargs)
+
+
+def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module:
+    """Install the Wan decoder fast path when its full structure is supported."""
+
+    if not _HAS_TRITON or getattr(vae, "_vllm_vae_fast_path_installed", False):
+        return vae
+
+    from diffusers.models.autoencoders import AutoencoderKLWan
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanDecoder3d
+
+    if not isinstance(vae, AutoencoderKLWan):
+        return vae
+    decoder = getattr(vae, "decoder", None)
+    if type(decoder) is not WanDecoder3d:
+        return vae
+
+    distributed_executor = getattr(vae, "distributed_executor", None)
+    parallel_mode = getattr(distributed_executor, "parallel_mode", "tile")
+    if parallel_mode in ("auto", "spatial_shard_height", "spatial_shard_width"):
+        logger.info("Wan spatial-shard decode is configured; skipping quality-gated VAE fast path")
+        return vae
+
+    gate = VaeFastPathGate()
+    norm_count = _install_norm_silu(decoder, gate)
+    if norm_count is None:
+        return vae
+    decoder._vllm_gate = gate
+    decoder._vllm_channels_last = False
+    decoder.forward = MethodType(_decoder_layout_forward, decoder)
+    register_vae_fast_path_gate(vae, gate)
+    vae._vllm_vae_fast_path_installed = True
+    logger.info(
+        "Wan VAE installed quality-gated decoder fast path (%d RMSNorm+SiLU fusions)",
+        norm_count,
+    )
+    return vae
+
+
+__all__ = ["FusedWanRMSNormSiLU", "maybe_optimize_wan_vae"]

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 compile_callable_with_fallback, regionally_compile
+from vllm_omni.diffusion.compile import 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,10 +40,13 @@ 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
+from vllm_omni.diffusion.vae_optimizations import (
+    optimize_pipeline_vaes,
+    use_pipeline_vae_fast_path,
+)
 from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.diffusion.worker.utils import (
@@ -189,39 +192,6 @@ 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,6 +247,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     custom_pipeline_name=custom_pipeline_name,
                     device=self.device,
                 )
+                # Keep precision rematerialization in the tagged weights pool
+                # so sleep level 1 can offload every resulting parameter.
+                optimize_pipeline_vaes(self.pipeline, self.od_config)
         time_after_load = time.perf_counter()
 
         logger.info(
@@ -311,8 +284,7 @@ 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():
-                pipeline_owns_compile = hasattr(self.pipeline, "setup_compile")
+                if hasattr(self.pipeline, "setup_compile"):
-                if pipeline_owns_compile:
                     try:
                         self.pipeline.setup_compile()
                     except Exception as exc:
@@ -326,9 +298,6 @@ 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.",
@@ -537,6 +506,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if require_request_batch_support and not getattr(self.pipeline, "supports_request_batch", False):
             raise RuntimeError(f"{type(self.pipeline).__name__} does not support request-batch forward.")
 
+        quality = getattr(reqs[0].sampling_params, "quality", None)
+        if any(getattr(req.sampling_params, "quality", None) != quality for req in reqs[1:]):
+            raise RuntimeError("A diffusion worker batch cannot mix request quality levels.")
+        use_fast_vae = quality == "high"
+
         # Use no_grad() for HSDP compatibility, inference_mode() otherwise for
         # better perf. HSDP2's fully_shard pre-forward hooks need tensor version
         # counters, which inference tensors do not track.
@@ -561,7 +535,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
             with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config):
                 with record_function(record_name):
-                    raw_outputs = self.pipeline.forward(batch)
+                    with use_pipeline_vae_fast_path(self.pipeline, use_fast_vae):
+                        raw_outputs = self.pipeline.forward(batch)
                     outputs = _normalize_pipeline_outputs(
                         raw_outputs,
                         expected_count=len(reqs),
@@ -812,7 +787,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
                             if should_decode:
                                 clear_pipeline_stage_durations(self.pipeline)
-                                result = self.pipeline.post_decode(req)
+                                with use_pipeline_vae_fast_path(
+                                    self.pipeline,
+                                    getattr(req.sampling, "quality", None) == "high",
+                                ):
+                                    result = self.pipeline.post_decode(req)
                                 if result is not None:
                                     self._attach_stepwise_metrics(
                                         req,

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -576,6 +576,14 @@ class DiffusionWorker:
 
         usage_before = allocator.get_current_usage()
 
+        if self.model_runner is not None and self.model_runner.pipeline is not None:
+            from vllm_omni.diffusion.vae_optimizations import clear_pipeline_vae_fast_path_caches
+
+            # Folded weights are allocated lazily on the request path rather
+            # than in the tagged weight pool. Drop them before sleep so they
+            # neither remain resident nor need a separate CPU backup.
+            clear_pipeline_vae_fast_path_caches(self.model_runner.pipeline)
+
         if level == 2 and self.model_runner is not None:
             if hasattr(self.model_runner, "graph_runners"):
                 self.model_runner.graph_runners.clear()

vllm_omni/engine/arg_utils.py

@@ -485,13 +485,13 @@ 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
     step_execution: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
+    vae_decode_precision: str | None = None
     enable_multithread_weight_load: bool = True
     num_weight_load_threads: int = 4
     enable_cpu_offload: bool = False

vllm_omni/engine/async_omni_engine.py

@@ -1026,6 +1026,7 @@ class AsyncOmniEngine:
             "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0),
             "vae_use_slicing": kwargs.get("vae_use_slicing", False),
             "vae_use_tiling": kwargs.get("vae_use_tiling", False),
+            "vae_decode_precision": kwargs.get("vae_decode_precision", None),
             "cache_backend": cache_backend,
             "cache_config": cache_config,
             "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
@@ -1043,9 +1044,6 @@ 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,15 +446,6 @@ 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",
@@ -638,6 +629,15 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
         )
+        omni_config_group.add_argument(
+            "--vae-decode-precision",
+            choices=("fp32", "fp16", "bf16"),
+            default=None,
+            help=(
+                "Optional Wan decode-only VAE precision. The encoder remains FP32; "
+                "unset preserves the pipeline's existing precision behavior."
+            ),
+        )
 
         # Parallel weight loading (faster diffusion startup)
         omni_config_group.add_argument(