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

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-11 09:31:27 +08:00
33 changed files with 1850 additions and 1246 deletions
@@ -449,21 +449,20 @@ Complete examples in the codebase:

 ---

-## Spatially-Sharded Decode (Wan and QwenImage)
+## Spatially-Sharded Decode (Wan)

-The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE and distributed QwenImage backend additionally support a **spatially-sharded decode** backend. `DiffusionParallelConfig.vae_parallel_mode="auto"` is the default; explicit `"tile"`, `"spatial_shard_height"`, and `"spatial_shard_width"` modes remain available.
+The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`).

 ### How it differs from tile parallel

-| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
-|--------|--------------------------|----------------------|-----------------------------------------------|
-| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W |
-| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
-| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0 and trim padding; QwenImage then broadcasts |
-| Activation memory | Bounded by tile size | Depends on the selected path | Grows with full output area divided across ranks |
-| Scope | Decode + encode | Wan/QwenImage decode selection; encode remains tiled | Decode only |
+| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
+|--------|--------------------------|-----------------------------------------------|
+| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
+| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
+| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
+| Scope | Decode + encode | Decode only |

-Spatial-shard decode swaps the decoder's spatial convolutions/padding for context-aware halo-exchanging variants so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Wan and QwenImage share the padding, 2D convolution, attention, halo, and trim primitives in `wan_spatial_shard.py`; `qwenimage_spatial_shard.py` supplies a causal-convolution subclass that remains discoverable by Diffusers' QwenImage cache accounting.
+Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`.

 ### Wiring

@@ -474,17 +473,15 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
  -> DiffusionParallelConfig.vae_parallel_mode
  -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...)
  -> DistributedVaeExecutor.parallel_mode
-  -> capability-based pipeline preparation
-  -> DistributedAutoencoderKLWan/DistributedAutoencoderKLQwenImage.tiled_decode dispatch
+  -> DistributedAutoencoderKLWan.tiled_decode dispatch
 ```

-The shared spatial policy gates the path: it requires distributed decode to be enabled, a 5D latent, a supported VAE capability, and `vae_patch_parallel_size == DiT group size`. The method is reached only after Diffusers has selected tiled decode for the request. In `auto` mode the longer axis that can cover every rank is selected, with width winning ties. This policy is cached only through normal kernel shape caches; it is reevaluated for every request.
+`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode.

 ### Notes

- The decoder is patched **in place** the first time spatial-shard decode runs. Its wrappers consult a context-local split dimension and retain an exact direct fallback, allowing later requests to select tile, height, or width safely.
+- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
 - Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
- QwenImage additionally covers cache discovery, compile compatibility, mixed tile/spatial requests, both axes, and all-rank result broadcast in `tests/diffusion/distributed/test_qwenimage_spatial_shard.py`.

 ---

@@ -15,7 +15,7 @@

 ## Overview

-VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan/QwenImage spatial-shard decode, which shards decoder feature maps along height or width.
+VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width.

 This is particularly useful for:
 - **High-resolution image generation** where VAE decode can become a memory bottleneck
@@ -113,7 +113,7 @@ In `DiffusionParallelConfig`:
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
-| `vae_parallel_mode` | str | `"auto"` | VAE parallel decode strategy: `"auto"` (shape/model/topology-aware selection), `"tile"`, `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding supports Wan and the distributed QwenImage backend. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan-and-qwenimage). |
+| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |

 Additional requirements:

@@ -126,9 +126,9 @@ Additional requirements:

 ---

-## Spatially-Sharded Decode (Wan and QwenImage)
+## Spatially-Sharded Decode (Wan)

-The default `vae_parallel_mode="auto"` selects spatial-shard decode only after a Wan or distributed QwenImage request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. It splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests, partial groups, and other VAE families keep their existing path. Set `vae_parallel_mode="tile"` to force the previous behavior, or `"spatial_shard_height"` / `"spatial_shard_width"` to force an axis.
+The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`.

 Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance.

@@ -141,7 +141,7 @@ omni = Omni(
    parallel_config=DiffusionParallelConfig(
        tensor_parallel_size=2,
        vae_patch_parallel_size=2,               # must match the DiT group size
-        vae_parallel_mode="auto",                # default
+        vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height"
    ),
 )
 ```
@@ -152,18 +152,14 @@ Or from the CLI / serving entrypoint:
 vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \
    --tensor-parallel-size 2 \
    --vae-patch-parallel-size 2 \
-    --vae-parallel-mode auto
+    --vae-parallel-mode spatial_shard_width
 ```

 **Constraints and behavior:**

- Spatial-shard decode is **decode-only** and implemented for the **Wan** VAE and the distributed **QwenImage** backend used by Qwen-Image and Krea2. Pipelines using a different VAE backend keep their existing behavior.
- It requires `vae_patch_parallel_size` to **match the DiT process group size**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning.
- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.
- Wan preserves its rank-0-only decode result, while QwenImage preserves its existing all-rank result by broadcasting after rank-0 assembly.
- Spatial sharding trades activation memory for latency as resolution grows. Use `vae_parallel_mode="tile"` when bounded VAE memory is more important than decode speed.
-
-The automatic policy is adapted from [SGLang #28071](https://github.com/sgl-project/sglang/pull/28071) and narrowed to vLLM-Omni's parity-tested Wan and QwenImage backends.
+- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes.
+- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime.
+- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension).

 For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:

@@ -579,7 +579,6 @@ def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
    assert cfg.data_parallel_size == 3
    assert cfg.cfg_parallel_size == 3
    assert cfg.mask_sp_padding is True
-    assert cfg.vae_parallel_mode == "auto"
    assert cfg.world_size == 72


@@ -1,233 +0,0 @@
-import os
-from dataclasses import dataclass
-
-import pytest
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-from diffusers.models.autoencoders.autoencoder_kl_qwenimage import QwenImageCausalConv3d
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import get_open_port
-from vllm_omni.diffusion.distributed.autoencoders import qwenimage_spatial_shard
-from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import (
-    DistributedAutoencoderKLQwenImage,
-)
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    prepare_pipeline_spatial_shard_decode,
-)
-from vllm_omni.diffusion.offloader.module_collector import PipelineModules
-from vllm_omni.platforms import current_omni_platform
-
-
-@dataclass
-class _FakeExecutor:
-    group: object
-    parallel_size: int
-    parallel_mode: str
-
-
-def _tiny_qwenimage_vae() -> DistributedAutoencoderKLQwenImage:
-    return DistributedAutoencoderKLQwenImage(
-        base_dim=8,
-        z_dim=4,
-        dim_mult=[1, 2, 2, 2],
-        num_res_blocks=1,
-        attn_scales=[],
-        temperal_downsample=[False, True, True],
-        input_channels=3,
-        latents_mean=[0.0] * 4,
-        latents_std=[1.0] * 4,
-    )
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_qwenimage_auto_spatial_shard_gate_is_shape_and_topology_aware(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.spatial_shard.dist.get_world_size",
-        lambda group=None: 4,
-    )
-    vae = DistributedAutoencoderKLQwenImage.__new__(DistributedAutoencoderKLQwenImage)
-    vae.use_tiling = True
-    vae.distributed_executor = _FakeExecutor(group=object(), parallel_size=4, parallel_mode="auto")
-    vae.is_distributed_enabled = lambda: True
-
-    landscape = torch.zeros((1, 16, 1, 60, 104))
-    portrait = torch.zeros((1, 16, 1, 104, 60))
-
-    assert vae._spatial_shard_decode_split_dim(landscape, 4) == "width"
-    assert vae._spatial_shard_decode_split_dim(portrait, 4) == "height"
-    assert vae._spatial_shard_decode_enabled(landscape) is True
-
-    vae.distributed_executor.parallel_size = 2
-    assert vae._spatial_shard_decode_enabled(landscape) is False
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_pipeline_dispatches_qwenimage_spatial_preparation(monkeypatch: pytest.MonkeyPatch):
-    vae = DistributedAutoencoderKLQwenImage.__new__(DistributedAutoencoderKLQwenImage)
-    prepared = []
-    vae._prepare_spatial_shard_decode = lambda: prepared.append(vae)
-
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
-        lambda pipeline: PipelineModules(
-            dits=[],
-            dit_names=[],
-            encoders=[],
-            encoder_names=[],
-            vaes=[torch.nn.Identity(), vae],
-        ),
-    )
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.spatial_shard.dist.is_initialized",
-        lambda: True,
-    )
-
-    prepare_pipeline_spatial_shard_decode(torch.nn.Identity())
-
-    assert prepared == [vae]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-@torch.no_grad()
-def test_qwenimage_installed_decoder_preserves_cache_discovery_state_and_direct_path():
-    torch.manual_seed(0)
-    vae = _tiny_qwenimage_vae().eval()
-    latents = torch.randn(1, 4, 2, 4, 5)
-    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, return_dict=False)[0]
-    causal_conv_count = sum(isinstance(module, QwenImageCausalConv3d) for module in vae.decoder.modules())
-
-    qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object(), split_dim="height")
-    qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object(), split_dim="width")
-    vae.clear_cache()
-
-    assert causal_conv_count > 0
-    assert len(vae._feat_map) == causal_conv_count
-    assert sum(isinstance(module, QwenImageCausalConv3d) for module in vae.decoder.modules()) == causal_conv_count
-    assert any(
-        isinstance(module, qwenimage_spatial_shard.QwenImageDistCausalConv3d) for module in vae.decoder.modules()
-    )
-    assert {name for name, _ in vae.named_parameters()} == parameter_names
-    vae.load_state_dict(state_dict, strict=True)
-    assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-@torch.no_grad()
-def test_qwenimage_installed_decoder_direct_path_can_be_compiled():
-    torch.manual_seed(1)
-    vae = _tiny_qwenimage_vae().eval()
-    latents = torch.randn(1, 4, 1, 3, 5)
-    qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object())
-    eager = vae.decode(latents, return_dict=False)[0]
-
-    compiled_decode = torch.compile(vae.decode, backend="eager", fullgraph=False)
-    compiled = compiled_decode(latents, return_dict=False)[0]
-
-    torch.testing.assert_close(compiled, eager, rtol=0, atol=0)
-
-
-def _qwenimage_spatial_shard_worker(rank: int, master_port: str, results) -> None:
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        destroy_model_parallel,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = master_port
-    device = current_omni_platform.get_torch_device(rank)
-    current_omni_platform.set_device(device)
-    backend = current_omni_platform.dist_backend
-    init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
-    initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
-
-    try:
-        torch.manual_seed(2)
-        vae = _tiny_qwenimage_vae().to(device=device, dtype=torch.float32).eval()
-        vae.init_distributed()
-        generator = torch.Generator(device=device).manual_seed(3)
-        landscape = torch.randn((1, 4, 2, 5, 7), generator=generator, device=device)
-        portrait = torch.randn((1, 4, 2, 7, 5), generator=generator, device=device)
-
-        with torch.inference_mode():
-            vae.use_tiling = False
-            vae.set_parallel_size(1, mode="tile")
-            landscape_reference = vae.decode(landscape, return_dict=False)[0]
-            portrait_reference = vae.decode(portrait, return_dict=False)[0]
-
-            vae.enable_tiling(
-                tile_sample_min_height=32,
-                tile_sample_min_width=32,
-                tile_sample_stride_height=24,
-                tile_sample_stride_width=24,
-            )
-            vae.set_parallel_size(2, mode="tile")
-            tile_before_install = vae.decode(landscape, return_dict=False)[0]
-
-            vae.set_parallel_size(2, mode="auto")
-            assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
-            landscape_auto = vae.decode(landscape, return_dict=False)[0]
-            assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
-            portrait_auto = vae.decode(portrait).sample
-
-            vae.set_parallel_size(2, mode="spatial_shard_height")
-            landscape_height = vae.decode(landscape, return_dict=False)[0]
-            vae.set_parallel_size(2, mode="spatial_shard_width")
-            portrait_width = vae.decode(portrait, return_dict=False)[0]
-
-            vae.set_parallel_size(2, mode="tile")
-            tile_after_install = vae.decode(landscape, return_dict=False)[0]
-
-        assert tile_after_install.shape == tile_before_install.shape
-        assert torch.equal(tile_after_install, tile_before_install)
-        assert all(
-            buffer is None
-            for module in vae.decoder.modules()
-            for name, buffer in module._buffers.items()
-            if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
-        )
-
-        diffs = {
-            "landscape_auto": (landscape_auto - landscape_reference).abs().max().item(),
-            "portrait_auto": (portrait_auto - portrait_reference).abs().max().item(),
-            "landscape_height": (landscape_height - landscape_reference).abs().max().item(),
-            "portrait_width": (portrait_width - portrait_reference).abs().max().item(),
-        }
-        results[rank] = {
-            "diffs": diffs,
-            "landscape_shape": tuple(landscape_auto.shape),
-            "portrait_shape": tuple(portrait_auto.shape),
-        }
-    finally:
-        destroy_model_parallel()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
-def test_qwenimage_spatial_shard_matches_reference_broadcasts_and_switches_paths():
-    manager = mp.get_context("spawn").Manager()
-    results = manager.dict()
-    mp.spawn(
-        _qwenimage_spatial_shard_worker,
-        args=(str(get_open_port()), results),
-        nprocs=2,
-        join=True,
-    )
-
-    assert set(results) == {0, 1}
-    assert results[0]["landscape_shape"] == results[1]["landscape_shape"] == (1, 3, 5, 40, 56)
-    assert results[0]["portrait_shape"] == results[1]["portrait_shape"] == (1, 3, 5, 56, 40)
-    for rank_results in results.values():
-        assert max(rank_results["diffs"].values()) <= 2e-5
@@ -1,5 +1,4 @@
 import os
-from contextlib import contextmanager
 from types import SimpleNamespace

 import pytest
@@ -11,8 +10,6 @@ from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
    DistributedAutoencoderKLWan,
-    prepare_pipeline_wan_spatial_shard_decode,
-    select_auto_spatial_shard_split_dim,
 )
 from vllm_omni.platforms import current_omni_platform

@@ -20,23 +17,6 @@ from vllm_omni.platforms import current_omni_platform
 # the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel.


-@contextmanager
-def _active_spatial_shard(split_dim: str):
-    token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
-        wan_spatial_shard.SpatialShardContext(
-            input_extent=2,
-            local_input_extent=2,
-            split_dim=split_dim,
-            rank=0,
-            world_size=3,
-        )
-    )
-    try:
-        yield
-    finally:
-        wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
-
-
@pytest.mark.core_model
@pytest.mark.cpu
 def test_split_for_parallel_decode_pads_uneven_height():
@@ -275,14 +255,12 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk

    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
-    with _active_spatial_shard("height"):
-        mid = mid_rank_pad(x)
+    mid = mid_rank_pad(x)
    assert mid.shape == (1, 1, 2, 3)

    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
-    with _active_spatial_shard("height"):
-        last = last_rank_pad(x)
+    last = last_rank_pad(x)
    assert last.shape == (1, 1, 3, 3)


@@ -292,15 +270,13 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
    x = torch.ones((1, 1, 2, 2))

    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
-    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
-    with _active_spatial_shard("width"):
-        mid = mid_rank_pad(x)
+    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    mid = mid_rank_pad(x)
    assert mid.shape == (1, 1, 2, 2)

    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
-    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
-    with _active_spatial_shard("width"):
-        last = last_rank_pad(x)
+    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    last = last_rank_pad(x)
    assert last.shape == (1, 1, 2, 3)


@@ -323,113 +299,6 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt
    assert vae._spatial_shard_decode_enabled(z) is False


-@pytest.mark.core_model
-@pytest.mark.cpu
-@pytest.mark.parametrize(
-    "z_shape,world_size,expected",
-    [
-        ((1, 16, 1, 60, 104), 4, "width"),
-        ((1, 16, 5, 60, 104), 4, "width"),
-        ((1, 16, 5, 104, 60), 4, "height"),
-        ((1, 16, 5, 2, 2), 4, None),
-        ((1, 16, 5, 64, 64), 4, "width"),
-    ],
-)
-def test_auto_spatial_shard_selector_is_shape_aware(z_shape, world_size, expected):
-    assert select_auto_spatial_shard_split_dim(z_shape, world_size) == expected
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
-        lambda group=None: 4,
-    )
-
-    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
-    vae.use_tiling = True
-    vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=4, parallel_mode="auto")
-    vae.is_distributed_enabled = lambda: True
-    large = torch.zeros((1, 16, 5, 60, 104))
-    unshardable = torch.zeros((1, 16, 1, 2, 2))
-
-    assert vae._spatial_shard_decode_split_dim(large, 4) == "width"
-    assert vae._spatial_shard_decode_enabled(large) is True
-    assert vae._spatial_shard_decode_enabled(unshardable) is False
-
-    vae.distributed_executor.parallel_size = 2
-    assert vae._spatial_shard_decode_enabled(large) is False
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch):
-    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
-    object.__setattr__(
-        vae,
-        "distributed_executor",
-        SimpleNamespace(group=object(), parallel_size=2, parallel_mode="auto"),
-    )
-    installed = []
-
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
-        lambda pipeline: SimpleNamespace(vaes=[vae]),
-    )
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.is_initialized",
-        lambda: True,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
-        lambda group=None: 4,
-    )
-    monkeypatch.setattr(
-        wan_spatial_shard,
-        "install_wan_spatial_shard_decode",
-        lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)),
-    )
-
-    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
-    assert installed == []
-
-    vae.distributed_executor.parallel_size = 4
-    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
-    assert installed == [(vae, vae.distributed_executor.group, "height")]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-@torch.no_grad()
-def test_installed_spatial_decoder_keeps_direct_fallback_and_can_change_axis():
-    vae = DistributedAutoencoderKLWan(
-        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,
-    ).eval()
-    latents = torch.randn(1, 4, 1, 4, 5)
-    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, return_dict=False)[0]
-
-    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="height")
-    assert any(getattr(module, "_vllm_omni_dynamic_spatial_shard_conv", False) for module in vae.decoder.modules())
-    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="width")
-
-    assert {name for name, _ in vae.named_parameters()} == parameter_names
-    vae.load_state_dict(state_dict, strict=True)
-    assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
-
-
@pytest.mark.core_model
@pytest.mark.cpu
 def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
@@ -575,99 +444,3 @@ def test_spatial_shard_decode_matches_reference(split_dim: str):
    assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
        f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
    )
-
-
-def _auto_spatial_shard_decode_worker(rank: int, return_dict, master_port: str) -> None:
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        destroy_model_parallel,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = master_port
-    device = current_omni_platform.get_torch_device(rank)
-    current_omni_platform.set_device(device)
-    backend = current_omni_platform.dist_backend
-    init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
-    initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
-
-    try:
-        vae = DistributedAutoencoderKLWan.from_pretrained(
-            _SPATIAL_SHARD_MODEL,
-            subfolder=_SPATIAL_SHARD_SUBFOLDER,
-            torch_dtype=torch.float32,
-        ).to(device=device, dtype=torch.float32)
-        vae.eval()
-        generator = torch.Generator(device=device).manual_seed(1)
-        small = torch.randn((1, vae.config.z_dim, 1, 60, 104), generator=generator, device=device)
-        landscape = torch.randn((1, vae.config.z_dim, 2, 60, 104), generator=generator, device=device)
-        portrait = torch.randn((1, vae.config.z_dim, 2, 104, 60), generator=generator, device=device)
-
-        with torch.inference_mode():
-            vae.use_tiling = False
-            vae.set_parallel_size(1, mode="tile")
-            small_reference = vae.decode(small, return_dict=False)[0]
-
-            vae.use_tiling = True
-            vae.set_parallel_size(2, mode="tile")
-            tile_before_install = vae.decode(small, return_dict=False)[0]
-
-            vae.set_parallel_size(2, mode="auto")
-            assert vae._spatial_shard_decode_split_dim(small, 2) == "width"
-            small_sharded = vae.decode(small, return_dict=False)[0]
-            assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
-            vae.decode(landscape, return_dict=False)
-
-            vae.set_parallel_size(2, mode="tile")
-            tile_after_install = vae.decode(small, return_dict=False)[0]
-
-            vae.use_tiling = False
-            vae.set_parallel_size(1, mode="tile")
-            portrait_reference = vae.decode(portrait, return_dict=False)[0]
-
-            vae.use_tiling = True
-            vae.set_parallel_size(2, mode="auto")
-            assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
-            portrait_sharded = vae.decode(portrait, return_dict=False)[0]
-
-        assert all(
-            buffer is None
-            for module in vae.decoder.modules()
-            for name, buffer in module._buffers.items()
-            if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
-        )
-
-        if rank == 0:
-            small_diff = (small_sharded.float() - small_reference.float()).abs()
-            portrait_diff = (portrait_sharded.float() - portrait_reference.float()).abs()
-            return_dict["tile_fallback_exact"] = torch.equal(tile_after_install, tile_before_install)
-            return_dict["small_max_abs_diff"] = small_diff.max().item()
-            return_dict["small_mean_abs_diff"] = small_diff.mean().item()
-            return_dict["portrait_max_abs_diff"] = portrait_diff.max().item()
-            return_dict["portrait_mean_abs_diff"] = portrait_diff.mean().item()
-    finally:
-        destroy_model_parallel()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-
-
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_auto_spatial_shard_switches_shape_paths_safely():
-    manager = mp.get_context("spawn").Manager()
-    return_dict = manager.dict()
-    mp.spawn(
-        _auto_spatial_shard_decode_worker,
-        args=(return_dict, "29502"),
-        nprocs=2,
-        join=True,
-    )
-
-    assert return_dict.get("tile_fallback_exact") is True
-    assert return_dict["small_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
-    assert return_dict["small_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
-    assert return_dict["portrait_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
-    assert return_dict["portrait_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
@@ -120,34 +120,3 @@ def test_tp_accepts_checkpoint_supported_sizes():
    arch = MiniMaxH3DiTArchConfig()
    for tp_size in (1, 2, 4, 7):
        model._validate_tp_config(arch=arch, tp_size=tp_size)
-
-
-@pytest.mark.parametrize(("parallel_size", "enabled"), [(1, False), (2, True)])
-def test_video_vae_auto_mode_preserves_native_tile_parallelism(monkeypatch, parallel_size, enabled):
-    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
-
-    parallel_state = {}
-    video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE)
-    nn.Module.__init__(video_vae)
-    video_vae.remote = nn.Identity()
-    video_vae.model = type("FakeModel", (), {"parallel_tiling": False})()
-
-    monkeypatch.setattr(vae_module, "get_dit_group", lambda: object())
-    monkeypatch.setattr(vae_module.dist, "get_world_size", lambda group: 2)
-    monkeypatch.setattr(vae_module.dist, "get_rank", lambda group: 0)
-    monkeypatch.setattr(
-        vae_module.importlib,
-        "import_module",
-        lambda package: type(
-            "FakeParallelModule",
-            (),
-            {"get_parallel_state": staticmethod(lambda: parallel_state)},
-        )(),
-    )
-
-    video_vae.set_parallel_size(parallel_size, mode="auto")
-
-    assert video_vae.parallel_size == parallel_size
-    assert video_vae.model.parallel_tiling is enabled
-    assert parallel_state["group_size"] == parallel_size
-    assert parallel_state["sp_enabled"] is enabled
@@ -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"
@@ -512,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()."""

@@ -579,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):
@@ -729,7 +770,7 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
@pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
    memory_pool_active = False
-    preparation_pool_states = []
+    optimization_pool_states = []

    @contextmanager
    def _memory_pool_context(tag):
@@ -790,8 +831,8 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
    monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
    monkeypatch.setattr(
        model_runner_module,
-        "prepare_pipeline_spatial_shard_decode",
-        lambda pipeline: preparation_pool_states.append(memory_pool_active),
+        "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(
@@ -803,7 +844,7 @@ 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 preparation_pool_states == [True]
+    assert optimization_pool_states == [True]


@pytest.mark.core_model
@@ -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
@@ -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
@@ -76,16 +76,15 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
    assert args.cfg_parallel_size == 4


-def test_serve_parser_defaults_vae_parallel_mode_to_auto_and_accepts_tile() -> None:
+def test_serve_parser_accepts_wan_decode_precision() -> None:
    parser = TrackingArgumentParser()
    subparsers = parser.add_subparsers(dest="subcommand")
    OmniServeCommand().subparser_init(subparsers)

-    default_args = parser.parse_args(["serve", "fake-model", "--omni"])
-    tile_args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-parallel-mode", "tile"])
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-decode-precision", "bf16"])

-    assert default_args.vae_parallel_mode == "auto"
-    assert tile_args.vae_parallel_mode == "tile"
+    assert args.vae_decode_precision == "bf16"
+    assert args.get_explicit_kwargs_dict()["vae_decode_precision"] == "bf16"


 def _make_headless_args(**kwargs) -> TrackingNamespace:
@@ -452,7 +452,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
    cfg_parallel_size: int = Field(default=1, ge=1)
    vae_patch_parallel_size: int = Field(default=1, ge=1)
    text_encoder_tp_size: int = Field(default=1, ge=1)
-    vae_parallel_mode: str = "auto"
+    vae_parallel_mode: str = "tile"
    use_hsdp: bool = False
    mask_sp_padding: bool = False
    hsdp_shard_size: int = -1
@@ -466,10 +466,9 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
            raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1")
        if self.ulysses_mode not in {"strict", "advanced_uaa"}:
            raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'")
-        if self.vae_parallel_mode not in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:
            raise ValueError(
-                "vae_parallel_mode must be one of "
-                "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
                f"but got {self.vae_parallel_mode!r}."
            )

@@ -574,6 +573,7 @@ class _DiffusionConfigProjection:
    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
@@ -401,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
@@ -206,13 +206,10 @@ class DiffusionParallelConfig:
    process group.
    """

-    vae_parallel_mode: str = "auto"
+    vae_parallel_mode: str = "tile"
    """VAE parallel decode strategy.

-    - "auto": Select the longer spatial-shard axis for validated Wan and
-      distributed QwenImage tiled decodes on a full DiT-group topology;
-      otherwise use patch/tile parallel decode.
-    - "tile": Patch/tile parallel decode. Each rank decodes a subset
+    - "tile": Patch/tile parallel decode (default). Each rank decodes a subset
      of spatial tiles and the results are stitched on rank 0.
    - "spatial_shard_height": Spatially-sharded decode that splits decoder
      feature maps along height and exchanges halo rows around spatial
@@ -253,9 +250,8 @@ class DiffusionParallelConfig:
        assert self.allgather_degree > 0, "AllGather degree must be > 0"
        assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
        assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0"
-        assert self.vae_parallel_mode in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}, (
-            "vae_parallel_mode must be one of "
-            "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+        assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, (
+            "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
            f"but got {self.vae_parallel_mode!r}."
        )
        if self.allgather_degree > 1:
@@ -715,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
@@ -927,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', "
@@ -4,23 +4,16 @@
 from typing import Any

 import torch
-import torch.distributed as dist
 from diffusers.models.autoencoders import AutoencoderKLQwenImage
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger

-from vllm_omni.diffusion.distributed.autoencoders import qwenimage_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
    DistributedOperator,
    DistributedVaeMixin,
    GridSpec,
    TileTask,
 )
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    prepare_spatial_shard_decode,
-    resolve_spatial_shard_split_dim,
-    spatial_shard_decode_enabled,
-)

 logger = init_logger(__name__)

@@ -111,39 +104,10 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
        ]
        return dec

-    def _spatial_shard_decode_split_dim(
-        self,
-        z: torch.Tensor | None = None,
-        world_size: int | None = None,
-    ) -> str | None:
-        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
-
-    def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        return spatial_shard_decode_enabled(self, z, model_name="QwenImage")
-
-    def _prepare_spatial_shard_decode(self) -> None:
-        prepare_spatial_shard_decode(
-            self,
-            install=qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode,
-        )
-
    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
        if not self.is_distributed_enabled():
            return super().tiled_decode(z, return_dict=return_dict)

-        if self._spatial_shard_decode_enabled(z):
-            world_size = dist.get_world_size(group=self.distributed_executor.group)
-            split_dim = self._spatial_shard_decode_split_dim(z, world_size)
-            assert split_dim is not None
-            logger.debug("Decode running with QwenImage VAE spatial_shard_%s mode", split_dim)
-            return qwenimage_spatial_shard.spatial_shard_decode(
-                self,
-                z,
-                group=self.distributed_executor.group,
-                return_dict=return_dict,
-                split_dim=split_dim,
-            )
-
        logger.debug("Decode running with distributed executor")
        result = self.distributed_executor.execute(
            z,
@@ -18,25 +18,11 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
    GridSpec,
    TileTask,
 )
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    prepare_pipeline_spatial_shard_decode,
-    prepare_spatial_shard_decode,
-    resolve_spatial_shard_split_dim,
-    spatial_shard_decode_enabled,
-)
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    select_auto_spatial_shard_split_dim as select_auto_spatial_shard_split_dim,
-)
 from vllm_omni.platforms import current_omni_platform

 logger = init_logger(__name__)


-def prepare_pipeline_wan_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
-    """Backward-compatible alias for the model-agnostic pipeline dispatcher."""
-    prepare_pipeline_spatial_shard_decode(pipeline)
-
-
 class OmniAutoencoderKLWan(AutoencoderKLWan):
    def _execution_context(self):
        try:
@@ -120,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

@@ -284,30 +270,45 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
        dec = torch.clamp(dec, min=-1.0, max=1.0)
        return dec

-    def _spatial_shard_decode_split_dim(
-        self,
-        z: torch.Tensor | None = None,
-        world_size: int | None = None,
-    ) -> str | None:
-        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
+    def _spatial_shard_decode_split_dim(self) -> str | None:
+        mode = self.distributed_executor.parallel_mode
+        if mode == "spatial_shard_width":
+            return "width"
+        if mode == "spatial_shard_height":
+            return "height"
+        return None

    def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        return spatial_shard_decode_enabled(self, z, model_name="Wan")
+        split_dim = self._spatial_shard_decode_split_dim()
+        if split_dim is None:
+            return False
+        if z.ndim != 5:
+            logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.")
+            return False
+        if not self.is_distributed_enabled():
+            return False

-    def _prepare_spatial_shard_decode(self) -> None:
-        prepare_spatial_shard_decode(
-            self,
-            install=wan_spatial_shard.install_wan_spatial_shard_decode,
-        )
+        group = self.distributed_executor.group
+        world_size = dist.get_world_size(group=group)
+        requested_size = int(self.distributed_executor.parallel_size)
+        if requested_size != world_size:
+            logger.warning(
+                "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size "
+                "to match the DIT group size; falling back to tiled decode. "
+                "requested=%s dit_group=%s split_dim=%s",
+                requested_size,
+                world_size,
+                split_dim,
+            )
+            return False
+        return True

    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
        if not self.is_distributed_enabled():
            return super().tiled_decode(z, return_dict=return_dict)

        if self._spatial_shard_decode_enabled(z):
-            world_size = dist.get_world_size(group=self.distributed_executor.group)
-            split_dim = self._spatial_shard_decode_split_dim(z, world_size)
-            assert split_dim is not None
+            split_dim = self._spatial_shard_decode_split_dim()
            logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim)
            return wan_spatial_shard.spatial_shard_decode(
                self,
@@ -1,86 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""QwenImage adapter for the shared spatial-shard VAE decoder kernels."""
-
-from __future__ import annotations
-
-import torch
-import torch.distributed as dist
-import torch.nn as nn
-from diffusers.models.autoencoders.autoencoder_kl_qwenimage import QwenImageCausalConv3d
-from diffusers.models.autoencoders.vae import DecoderOutput
-
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE
-from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import (
-    SpatialShardCausalConv3dMixin,
-    install_spatial_shard_decode,
-    spatial_shard_decode_impl,
-)
-
-
-class QwenImageDistCausalConv3d(SpatialShardCausalConv3dMixin, QwenImageCausalConv3d):
-    """Context-aware QwenImage causal conv that remains cache-discoverable."""
-
-    def __init__(self, source: nn.Conv3d, group: dist.ProcessGroup) -> None:
-        # Initialize Conv3d directly so the replacement preserves source device,
-        # dtype, groups, and bias while remaining an isinstance() match for
-        # QwenImageCausalConv3d. Diffusers clear_cache() relies on that match.
-        nn.Conv3d.__init__(
-            self,
-            source.in_channels,
-            source.out_channels,
-            source.kernel_size,
-            stride=source.stride,
-            padding=0,
-            dilation=source.dilation,
-            groups=source.groups,
-            bias=source.bias is not None,
-            padding_mode=source.padding_mode,
-            device=source.weight.device,
-            dtype=source.weight.dtype,
-        )
-        self._init_spatial_shard_causal_conv(source, group)
-
-    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        return self._spatial_shard_causal_conv_forward(x, cache_x)
-
-
-def install_qwenimage_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    group: dist.ProcessGroup,
-    split_dim: str = "height",
-) -> None:
-    install_spatial_shard_decode(
-        vae,
-        group,
-        split_dim,
-        causal_conv_class_name="QwenImageCausalConv3d",
-        causal_conv_factory=QwenImageDistCausalConv3d,
-        attention_block_class_name="QwenImageAttentionBlock",
-        installed_attr="_vllm_omni_qwenimage_spatial_shard_installed",
-        model_name="QwenImage",
-    )
-
-
-def spatial_shard_decode(
-    vae: SpatialShardVAE,
-    z: torch.Tensor,
-    *,
-    group: dist.ProcessGroup,
-    return_dict: bool = True,
-    split_dim: str = "height",
-) -> DecoderOutput | tuple[torch.Tensor]:
-    """Decode QwenImage latents and preserve its all-rank output contract."""
-    return spatial_shard_decode_impl(
-        vae,
-        z,
-        group=group,
-        install=install_qwenimage_spatial_shard_decode,
-        model_name="QwenImage",
-        pass_first_chunk=False,
-        broadcast_result=True,
-        unpatchify_patch_size=None,
-        return_dict=return_dict,
-        split_dim=split_dim,
-    )
@@ -1,142 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from typing import Protocol
-
-import torch
-import torch.distributed as dist
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class SpatialShardExecutor(Protocol):
-    group: dist.ProcessGroup
-    parallel_mode: str
-    parallel_size: int
-
-    def _sync_final_result(
-        self,
-        rank0_result: torch.Tensor,
-        output_ndim: int,
-        output_device: torch.device,
-        output_dtype: torch.dtype,
-    ) -> torch.Tensor: ...
-
-
-class SpatialShardVAE(Protocol):
-    distributed_executor: SpatialShardExecutor
-    decoder: torch.nn.Module
-    post_quant_conv: torch.nn.Module
-    dtype: torch.dtype
-    _conv_idx: list[int]
-    _feat_map: list[object]
-
-    def is_distributed_enabled(self) -> bool: ...
-
-    def clear_cache(self) -> None: ...
-
-
-def select_auto_spatial_shard_split_dim(
-    z_shape: tuple[int, ...],
-    world_size: int,
-) -> str | None:
-    """Select the longer spatial axis that can be split across every rank."""
-    if len(z_shape) != 5 or world_size <= 1:
-        return None
-
-    _, _, _, height, width = z_shape
-    candidates = [
-        (extent, split_dim) for extent, split_dim in ((height, "height"), (width, "width")) if extent >= world_size
-    ]
-    if not candidates:
-        return None
-    # Prefer width for square latents. It was modestly faster in the original
-    # four-GPU Wan benchmark and keeps the choice deterministic across models.
-    return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1]
-
-
-def resolve_spatial_shard_split_dim(
-    mode: str,
-    z: torch.Tensor | None = None,
-    world_size: int | None = None,
-) -> str | None:
-    if mode == "spatial_shard_width":
-        return "width"
-    if mode == "spatial_shard_height":
-        return "height"
-    if mode == "auto" and z is not None and world_size is not None:
-        return select_auto_spatial_shard_split_dim(tuple(z.shape), world_size)
-    return None
-
-
-def spatial_shard_decode_enabled(vae: SpatialShardVAE, z: torch.Tensor, *, model_name: str) -> bool:
-    """Apply the shared request- and topology-level spatial decode policy."""
-    executor = vae.distributed_executor
-    if executor.parallel_mode == "tile":
-        return False
-    if z.ndim != 5:
-        logger.warning(
-            "%s VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.", model_name
-        )
-        return False
-    if not vae.is_distributed_enabled():
-        return False
-
-    group = executor.group
-    world_size = dist.get_world_size(group=group)
-    requested_size = int(executor.parallel_size)
-    requested_split_dim = resolve_spatial_shard_split_dim(executor.parallel_mode)
-    if requested_size != world_size:
-        if executor.parallel_mode == "auto":
-            logger.debug(
-                "%s VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
-                model_name,
-                requested_size,
-                world_size,
-            )
-        elif requested_split_dim is not None:
-            logger.warning(
-                "%s VAE spatial sharded decode currently requires vae_patch_parallel_size "
-                "to match the DIT group size; falling back to tiled decode. "
-                "requested=%s dit_group=%s split_dim=%s",
-                model_name,
-                requested_size,
-                world_size,
-                requested_split_dim,
-            )
-        return False
-    return resolve_spatial_shard_split_dim(executor.parallel_mode, z, world_size) is not None
-
-
-def prepare_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    *,
-    install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
-) -> None:
-    """Install one VAE's dynamic wrappers while allocations are still tagged."""
-    executor = vae.distributed_executor
-    mode = executor.parallel_mode
-    if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}:
-        return
-    requested_size = int(executor.parallel_size)
-    world_size = dist.get_world_size(group=executor.group)
-    if requested_size <= 1 or requested_size != world_size:
-        return
-    split_dim = "width" if mode == "spatial_shard_width" else "height"
-    install(vae, executor.group, split_dim)
-
-
-def prepare_pipeline_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
-    """Prepare every VAE that advertises spatial-shard decode capability."""
-    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
-
-    if not dist.is_initialized():
-        return
-    for vae in ModuleDiscovery.discover(pipeline).vaes:
-        prepare = getattr(vae, "_prepare_spatial_shard_decode", None)
-        if callable(prepare):
-            prepare()
@@ -7,21 +7,19 @@
 #   (python/sglang/multimodal_gen/runtime/layers/parallel_conv.py)
 # which is in turn adapted from FastVideo (https://github.com/hao-ai-lab/FastVideo).
 # This version generalizes the height-only sharding to shard along height or
-# width and supports the causal-conv ``feat_cache`` handling used by Wan and
-# QwenImage.
-"""Shared spatial-shard kernels plus the Wan VAE adapter.
+# width and adds the Wan causal-conv ``feat_cache`` handling.
+"""Spatially-sharded Wan VAE decode.

-The halo, padding, stride-trim, attention-gather, and decode-lifecycle helpers
-are also reused by QwenImage's thin adapter. The backend shards decoder feature
-maps along height or width and exchanges boundary rows/columns before spatial
-convolutions. It is decode-only and keeps checkpoint loading unchanged by
-patching the already-loaded decoder.
+The existing distributed Wan VAE path shards *tiles*.  This module adds an
+opt-in decode backend that shards decoder feature maps along height or width and
+exchanges boundary rows/columns before spatial convolutions.  It is
+intentionally decode-only and keeps checkpoint loading unchanged by patching the
+already-loaded decoder.
 """

 from __future__ import annotations

 import math
-from collections.abc import Callable
 from contextlib import nullcontext
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -36,8 +34,6 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger

-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE
-
 logger = init_logger(__name__)


@@ -56,17 +52,12 @@ _SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar(
 )


-def _active_spatial_shard_split_dim() -> str | None:
-    context = _SPATIAL_SHARD_CONTEXT.get()
-    return context.split_dim if context is not None else None
-
-
 def _spatial_dim(split_dim: str) -> int:
    if split_dim == "height":
        return -2
    if split_dim == "width":
        return -1
-    raise ValueError(f"Unsupported VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
+    raise ValueError(f"Unsupported Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.")


 def _narrow_along_dim(x: torch.Tensor, dim: int, start: int, length: int) -> torch.Tensor:
@@ -164,9 +155,9 @@ def split_for_parallel_decode(
    rank = 0 if rank is None else int(rank)
    world_size = 1 if world_size is None else int(world_size)
    if world_size < 1:
-        raise ValueError(f"VAE world_size must be >= 1, got {world_size}.")
+        raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.")
    if not 0 <= rank < world_size:
-        raise ValueError(f"VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
+        raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")

    dim = _spatial_dim(split_dim)
    expected_extent = x.shape[dim] * (2**upsample_count)
@@ -332,27 +323,21 @@ class WanDistZeroPad2d(nn.Module):
        padding: tuple[int, int, int, int],
        group: dist.ProcessGroup,
        *,
-        defer_split_after_padding: bool = False,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
    ) -> None:
        super().__init__()
        self.padding = padding
-        self.defer_split_after_padding = defer_split_after_padding
+        self.split_dim = split_dim
+        default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1])
+        self.split_padding = split_padding or default_split_padding
        self.group = group

    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        if split_dim is None:
-            return F.pad(x, self.padding)
-
        rank, world_size = _rank_world(self.group)
        left, right, top, bottom = self.padding
-        if self.defer_split_after_padding:
-            if split_dim == "height":
-                bottom = 0
-            else:
-                right = 0
        if world_size > 1:
-            if split_dim == "height":
+            if self.split_dim == "height":
                top = top if rank == 0 else 0
                bottom = bottom if rank == world_size - 1 else 0
            else:
@@ -366,7 +351,8 @@ class WanDistConv2d(nn.Conv2d):
        self,
        source: nn.Conv2d,
        group: dist.ProcessGroup,
-        deferred_padding: tuple[int, int, int, int] | None = None,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
    ):
        super().__init__(
            source.in_channels,
@@ -383,253 +369,176 @@ class WanDistConv2d(nn.Conv2d):
        )
        self.load_state_dict(source.state_dict())
        self.group = group
-        self._direct_padding = source.padding
-        self._direct_padding_mode = source.padding_mode
-        self._direct_reversed_padding = source._reversed_padding_repeated_twice
-        self._deferred_padding = deferred_padding
-        source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding)
-        if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding):
-            raise ValueError(f"Spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
-        self._spatial_pad_h = int(source_padding[-2])
-        self._spatial_pad_w = int(source_padding[-1])
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
+        self.split_dim = split_dim
+        self.split_tensor_dim = _spatial_dim(split_dim)
+        kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1]
+        self.halo_size = (kernel_extent - 1) // 2
+        pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding
+        pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding
+        if split_dim == "height":
+            if split_padding is None:
+                split_padding = (pad_h, pad_h)
+            self._non_split_padding = (pad_w, pad_w, 0, 0)
+            self.kernel_extent = self.kernel_size[-2]
+            self.stride_extent = self.stride[-2]
+        else:
+            if split_padding is None:
+                split_padding = (pad_w, pad_w)
+            self._non_split_padding = (0, 0, pad_h, pad_h)
+            self.kernel_extent = self.kernel_size[-1]
+            self.stride_extent = self.stride[-1]
+        self.split_pad_left, self.split_pad_right = split_padding
+        self._halo_recv_top_buf: torch.Tensor | None = None
+        self._halo_recv_bottom_buf: torch.Tensor | None = None
+        self._trim_cache: dict[int, tuple[int, int, int]] = {}

    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        if split_dim is None:
-            if self._direct_padding_mode != "zeros":
-                x = F.pad(x, self._direct_reversed_padding, mode=self._direct_padding_mode)
-                padding: str | tuple[int, int] = (0, 0)
-            else:
-                padding = self._direct_padding
-            return F.conv2d(x, self.weight, self.bias, self.stride, padding, self.dilation, self.groups)
-
-        split_tensor_dim = _spatial_dim(split_dim)
-        if split_dim == "height":
-            kernel_extent = self.kernel_size[-2]
-            stride_extent = self.stride[-2]
-            non_split_padding = (self._spatial_pad_w, self._spatial_pad_w, 0, 0)
-            default_split_padding = (self._spatial_pad_h, self._spatial_pad_h)
-            deferred_split_padding = (
-                (self._deferred_padding[2], self._deferred_padding[3]) if self._deferred_padding is not None else None
-            )
-        else:
-            kernel_extent = self.kernel_size[-1]
-            stride_extent = self.stride[-1]
-            non_split_padding = (0, 0, self._spatial_pad_h, self._spatial_pad_h)
-            default_split_padding = (self._spatial_pad_w, self._spatial_pad_w)
-            deferred_split_padding = (
-                (self._deferred_padding[0], self._deferred_padding[1]) if self._deferred_padding is not None else None
-            )
-        split_pad_left, split_pad_right = deferred_split_padding or default_split_padding
-        halo_size = (kernel_extent - 1) // 2
-
-        x = F.pad(x, non_split_padding)
+        x = F.pad(x, self._non_split_padding)
        x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
            x,
            group=self.group,
-            halo_size=halo_size,
-            split_dim=split_dim,
+            halo_size=self.halo_size,
+            split_dim=self.split_dim,
            recv_top_buf=self._halo_recv_top_buf,
            recv_bottom_buf=self._halo_recv_bottom_buf,
        )
-        shift, start, upper_bound = self._get_trim_params(
-            x.shape[split_tensor_dim],
-            split_dim=split_dim,
-            halo_size=halo_size,
-            split_pad_left=split_pad_left,
-            split_pad_right=split_pad_right,
-            kernel_extent=kernel_extent,
-            stride_extent=stride_extent,
-        )
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
        if shift:
            x_padded = _narrow_along_dim(
                x_padded,
-                split_tensor_dim,
+                self.split_tensor_dim,
                shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
            )
        out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)

-    def _get_trim_params(
-        self,
-        local_extent: int,
-        *,
-        split_dim: str,
-        halo_size: int,
-        split_pad_left: int,
-        split_pad_right: int,
-        kernel_extent: int,
-        stride_extent: int,
-    ) -> tuple[int, int, int]:
-        cache_key = (split_dim, local_extent)
-        trim_params = self._trim_cache.get(cache_key)
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
        if trim_params is None:
            rank, world_size = _rank_world(self.group)
            trim_params = _compute_conv_trim_params(
                local_extent=local_extent,
                rank=rank,
                world_size=world_size,
-                halo_size=halo_size,
-                pad_before=split_pad_left,
-                pad_after=split_pad_right,
-                kernel_extent=kernel_extent,
-                stride_extent=stride_extent,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
            )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = trim_params
        return trim_params


-# Model-neutral aliases let other causal VAEs reuse the implementation while
-# preserving the original Wan class identities for stacked PRs and downstreams.
-SpatialShardZeroPad2d = WanDistZeroPad2d
-SpatialShardConv2d = WanDistConv2d
-
-
-class SpatialShardCausalConv3dMixin:
-    """Shared causal-convolution behavior for supported Diffusers VAE classes."""
-
-    def _init_spatial_shard_causal_conv(
+class WanDistCausalConv3d(nn.Conv3d):
+    def __init__(
        self,
        source: nn.Conv3d,
        group: dist.ProcessGroup,
-    ) -> None:
+        split_dim: str = "height",
+    ):
+        super().__init__(
+            source.in_channels,
+            source.out_channels,
+            source.kernel_size,
+            stride=source.stride,
+            padding=0,
+            dilation=source.dilation,
+            groups=source.groups,
+            bias=source.bias is not None,
+            padding_mode=source.padding_mode,
+            device=source.weight.device,
+            dtype=source.weight.dtype,
+        )
        self.load_state_dict(source.state_dict())
        self.group = group
+        self.split_dim = split_dim
+        self.split_tensor_dim = _spatial_dim(split_dim)
        source_padding = getattr(source, "_padding", None)
        if source_padding is None:
            p_t, p_h, p_w = source.padding
            source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0)
        self._source_padding = tuple(source_padding)
-        self._padding = tuple(source_padding)
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
+        if split_dim == "height":
+            self.split_pad_left = int(self._source_padding[2])
+            self.split_pad_right = int(self._source_padding[3])
+            self.kernel_extent = self.kernel_size[-2]
+            self.stride_extent = self.stride[-2]
+            padding = (
+                self._source_padding[0],
+                self._source_padding[1],
+                0,
+                0,
+                self._source_padding[4],
+                self._source_padding[5],
+            )
+        else:
+            self.split_pad_left = int(self._source_padding[0])
+            self.split_pad_right = int(self._source_padding[1])
+            self.kernel_extent = self.kernel_size[-1]
+            self.stride_extent = self.stride[-1]
+            padding = (
+                0,
+                0,
+                self._source_padding[2],
+                self._source_padding[3],
+                self._source_padding[4],
+                self._source_padding[5],
+            )
+        self.halo_size = (self.kernel_extent - 1) // 2
+        self._padding = padding if self.halo_size > 0 else self._source_padding
+        self._halo_recv_top_buf: torch.Tensor | None = None
+        self._halo_recv_bottom_buf: torch.Tensor | None = None
+        self._trim_cache: dict[int, tuple[int, int, int]] = {}

-    def _spatial_shard_causal_conv_forward(
-        self,
-        x: torch.Tensor,
-        cache_x: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        padding = list(self._source_padding)
+    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
+        padding = list(self._padding)
        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]

-        if split_dim is None:
-            return nn.Conv3d.forward(self, F.pad(x, padding))
-
-        split_tensor_dim = _spatial_dim(split_dim)
-        if split_dim == "height":
-            split_pad_left = int(padding[2])
-            split_pad_right = int(padding[3])
-            kernel_extent = self.kernel_size[-2]
-            stride_extent = self.stride[-2]
-            padding[2] = 0
-            padding[3] = 0
-        else:
-            split_pad_left = int(padding[0])
-            split_pad_right = int(padding[1])
-            kernel_extent = self.kernel_size[-1]
-            stride_extent = self.stride[-1]
-            padding[0] = 0
-            padding[1] = 0
-        halo_size = (kernel_extent - 1) // 2
-
        x = F.pad(x, padding)
        x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
            x,
            group=self.group,
-            halo_size=halo_size,
-            split_dim=split_dim,
+            halo_size=self.halo_size,
+            split_dim=self.split_dim,
            recv_top_buf=self._halo_recv_top_buf,
            recv_bottom_buf=self._halo_recv_bottom_buf,
        )
-        shift, start, upper_bound = self._get_trim_params(
-            x.shape[split_tensor_dim],
-            split_dim=split_dim,
-            halo_size=halo_size,
-            split_pad_left=split_pad_left,
-            split_pad_right=split_pad_right,
-            kernel_extent=kernel_extent,
-            stride_extent=stride_extent,
-        )
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
        if shift:
            x_padded = _narrow_along_dim(
                x_padded,
-                split_tensor_dim,
+                self.split_tensor_dim,
                shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
            )
-        out = nn.Conv3d.forward(self, x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        out = super().forward(x_padded)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)

-    def _get_trim_params(
-        self,
-        local_extent: int,
-        *,
-        split_dim: str,
-        halo_size: int,
-        split_pad_left: int,
-        split_pad_right: int,
-        kernel_extent: int,
-        stride_extent: int,
-    ) -> tuple[int, int, int]:
-        cache_key = (split_dim, local_extent)
-        trim_params = self._trim_cache.get(cache_key)
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
        if trim_params is None:
            rank, world_size = _rank_world(self.group)
            trim_params = _compute_conv_trim_params(
                local_extent=local_extent,
                rank=rank,
                world_size=world_size,
-                halo_size=halo_size,
-                pad_before=split_pad_left,
-                pad_after=split_pad_right,
-                kernel_extent=kernel_extent,
-                stride_extent=stride_extent,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
            )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = trim_params
        return trim_params


-class WanDistCausalConv3d(SpatialShardCausalConv3dMixin, nn.Conv3d):
-    # Cross-PR compatibility contract for the lossless Wan data-movement
-    # installer: it may fuse the no-context direct path, but must leave this
-    # wrapper in control whenever a spatial-shard context is active.
-    _vllm_omni_dynamic_spatial_shard_conv = True
-
-    def __init__(
-        self,
-        source: nn.Conv3d,
-        group: dist.ProcessGroup,
-    ):
-        super().__init__(
-            source.in_channels,
-            source.out_channels,
-            source.kernel_size,
-            stride=source.stride,
-            padding=0,
-            dilation=source.dilation,
-            groups=source.groups,
-            bias=source.bias is not None,
-            padding_mode=source.padding_mode,
-            device=source.weight.device,
-            dtype=source.weight.dtype,
-        )
-        self._init_spatial_shard_causal_conv(source, group)
-
-    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        return self._spatial_shard_causal_conv_forward(x, cache_x)
-
-
 def _compute_conv_trim_params(
    *,
    local_extent: int,
@@ -673,15 +582,12 @@ def _trim_local_conv_output(
    return out


-def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> None:
+def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None:
    if getattr(module, "_vllm_omni_spatial_shard_attention", False):
        return
    orig_forward = module.forward

    def _forward(self: nn.Module, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        if split_dim is None:
-            return orig_forward(x, *args, **kwargs)
        _, world_size = _rank_world(group)
        if world_size <= 1:
            return orig_forward(x, *args, **kwargs)
@@ -703,43 +609,60 @@ def _replace_child(
    name: str,
    child: nn.Module,
    group: dist.ProcessGroup,
-    *,
-    causal_conv_class_name: str,
-    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
+    split_dim: str,
 ) -> None:
-    if child.__class__.__name__ == causal_conv_class_name:
-        setattr(parent, name, causal_conv_factory(child, group))
-        return
-    if isinstance(child, nn.ZeroPad2d):
-        padding = tuple(int(p) for p in child.padding)
+    if child.__class__.__name__ == "WanCausalConv3d":
        setattr(
            parent,
            name,
-            SpatialShardZeroPad2d(
-                padding,
+            WanDistCausalConv3d(
+                child,
                group,
-                defer_split_after_padding=parent.__class__.__name__ == "Sequential",
+                split_dim=split_dim,
+            ),
+        )
+        return
+    if isinstance(child, nn.ZeroPad2d):
+        padding = tuple(int(p) for p in child.padding)
+        module_padding = padding
+        if parent.__class__.__name__ == "Sequential":
+            # Let the following WanDistConv2d account for global after-edge
+            # padding; this module handles the non-split dimension and before edge.
+            if split_dim == "height":
+                module_padding = (padding[0], padding[1], padding[2], 0)
+            else:
+                module_padding = (padding[0], 0, padding[2], padding[3])
+        setattr(
+            parent,
+            name,
+            WanDistZeroPad2d(
+                module_padding,
+                group,
+                split_dim=split_dim,
+                split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
            ),
        )
        return
    if isinstance(child, nn.Conv2d):
-        deferred_padding = None
+        split_padding = None
        if name == "1" and parent.__class__.__name__ == "Sequential":
-            # Wan/QwenImage downsampling uses ZeroPad2d((0, 1, 0, 1)) before a
+            # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a
            # stride-2 conv with padding=0.  Only the last rank should see the
            # bottom/right global padding, which is approximated by split_padding.
            prev = getattr(parent, "0", None)
-            if isinstance(prev, SpatialShardZeroPad2d):
-                deferred_padding = prev.padding
+            if isinstance(prev, WanDistZeroPad2d):
+                split_padding = prev.split_padding
            elif isinstance(prev, nn.ZeroPad2d):
-                deferred_padding = tuple(int(value) for value in prev.padding)
+                pad = prev.padding
+                split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1]))
        setattr(
            parent,
            name,
-            SpatialShardConv2d(
+            WanDistConv2d(
                child,
                group,
-                deferred_padding=deferred_padding,
+                split_dim=split_dim,
+                split_padding=split_padding,
            ),
        )

@@ -747,45 +670,21 @@ def _replace_child(
 def _patch_decoder_modules(
    module: nn.Module,
    group: dist.ProcessGroup,
-    *,
-    causal_conv_class_name: str,
-    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
-    attention_block_class_name: str,
+    split_dim: str,
    inside_attention: bool = False,
 ) -> None:
-    if module.__class__.__name__ == attention_block_class_name:
-        _patch_attention_block(module, group)
+    if module.__class__.__name__ == "WanAttentionBlock":
+        _patch_attention_block(module, group, split_dim)
        inside_attention = True

    for name, child in list(module.named_children()):
-        if child.__class__.__name__ == causal_conv_class_name:
-            _replace_child(
-                module,
-                name,
-                child,
-                group,
-                causal_conv_class_name=causal_conv_class_name,
-                causal_conv_factory=causal_conv_factory,
-            )
+        if child.__class__.__name__ == "WanCausalConv3d":
+            _replace_child(module, name, child, group, split_dim)
            continue
        if isinstance(child, nn.Conv2d) and not inside_attention:
-            _replace_child(
-                module,
-                name,
-                child,
-                group,
-                causal_conv_class_name=causal_conv_class_name,
-                causal_conv_factory=causal_conv_factory,
-            )
+            _replace_child(module, name, child, group, split_dim)
            continue
-        _patch_decoder_modules(
-            child,
-            group,
-            causal_conv_class_name=causal_conv_class_name,
-            causal_conv_factory=causal_conv_factory,
-            attention_block_class_name=attention_block_class_name,
-            inside_attention=inside_attention,
-        )
+        _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)


 def _decoder_upsample_count(decoder: nn.Module) -> int:
@@ -796,58 +695,47 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
    return count


-def clear_spatial_shard_runtime_buffers(vae: SpatialShardVAE) -> None:
-    decoder = getattr(vae, "decoder", None)
-    if decoder is None:
-        return
-    for module in decoder.modules():
-        if "_halo_recv_top_buf" in module._buffers:
-            module._halo_recv_top_buf = None
-        if "_halo_recv_bottom_buf" in module._buffers:
-            module._halo_recv_bottom_buf = None
-
-
-def install_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    group: dist.ProcessGroup,
-    split_dim: str,
-    *,
-    causal_conv_class_name: str,
-    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
-    attention_block_class_name: str,
-    installed_attr: str,
-    model_name: str,
-) -> None:
+def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None:
    """Patch ``vae.decoder`` once for spatially-sharded decode.

    This mutates the already-loaded decoder in place by swapping its spatial
-    convolutions/padding for context-aware variants and attaching a dedicated
-    spatial forward. Outside that forward the replacements execute their exact
-    local PyTorch behavior, so a later auto-selected tile request remains safe.
-    The same installed decoder can alternate height and width sharding.
+    convolutions/padding for halo-exchanging variants and wrapping
+    ``decoder.forward``. The patch is permanent for the lifetime of the VAE
+    instance and is applied only once (subsequent calls are no-ops). A given
+    instance is bound to a single ``split_dim``; switching between
+    ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here
+    otherwise.

-    Group-relative rank 0 assembles the final decoded frame. The model adapter
-    then preserves its prior contract: Wan keeps the result on rank 0, while
-    QwenImage broadcasts it to every rank.
+    Only group-relative rank 0 assembles the final decoded frame, mirroring the
+    distributed tiled-decode ``broadcast_result=False`` contract; the other ranks
+    take part in the collectives but return an empty placeholder.
    """
    _spatial_dim(split_dim)
-    if getattr(vae, installed_attr, False):
+    if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False):
+        installed_split_dim = getattr(vae, "_vllm_omni_wan_spatial_shard_split_dim", "height")
+        if installed_split_dim != split_dim:
+            raise ValueError(
+                "Wan spatial-shard VAE decoder was already patched for "
+                f"{installed_split_dim!r} split; create a fresh VAE instance to use {split_dim!r} split."
+            )
        return
    decoder = getattr(vae, "decoder", None)
    if decoder is None:
-        raise ValueError(f"{model_name} spatial-shard VAE decode requires a decoder module.")
+        raise ValueError("Wan spatial-shard VAE decode requires a decoder module.")

-    _patch_decoder_modules(
-        decoder,
-        group,
-        causal_conv_class_name=causal_conv_class_name,
-        causal_conv_factory=causal_conv_factory,
-        attention_block_class_name=attention_block_class_name,
-    )
+    _patch_decoder_modules(decoder, group, split_dim)
    upsample_count = _decoder_upsample_count(decoder)
    orig_forward = decoder.forward

-    def _forward(self: nn.Module, x: torch.Tensor, *args: object, split_dim: str, **kwargs: object) -> torch.Tensor:
+    def _forward(
+        self: nn.Module,
+        x: torch.Tensor,
+        feat_cache: list[torch.Tensor] | None = None,
+        feat_idx: list[int] | None = None,
+        first_chunk: bool = False,
+    ) -> torch.Tensor:
+        if feat_idx is None:
+            feat_idx = [0]
        tensor_dim = _spatial_dim(split_dim)
        input_extent = x.shape[tensor_dim]
        x, expected_extent = split_for_parallel_decode(
@@ -867,50 +755,29 @@ def install_spatial_shard_decode(
            )
        )
        try:
-            out = orig_forward(x, *args, **kwargs)
+            out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
        finally:
            _SPATIAL_SHARD_CONTEXT.reset(token)
        return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0)

-    decoder._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder)
-    setattr(vae, installed_attr, True)
-    logger.info("Installed %s VAE dynamic-axis spatial-shard decode.", model_name)
+    decoder.forward = MethodType(_forward, decoder)
+    vae._vllm_omni_wan_spatial_shard_installed = True
+    vae._vllm_omni_wan_spatial_shard_split_dim = split_dim
+    logger.info("Installed Wan VAE %s-sharded decode.", split_dim)


-def install_wan_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    group: dist.ProcessGroup,
-    split_dim: str = "height",
-) -> None:
-    install_spatial_shard_decode(
-        vae,
-        group,
-        split_dim,
-        causal_conv_class_name="WanCausalConv3d",
-        causal_conv_factory=WanDistCausalConv3d,
-        attention_block_class_name="WanAttentionBlock",
-        installed_attr="_vllm_omni_wan_spatial_shard_installed",
-        model_name="Wan",
-    )
-
-
-def spatial_shard_decode_impl(
-    vae: SpatialShardVAE,
+def spatial_shard_decode(
+    vae: Any,
    z: torch.Tensor,
    *,
    group: dist.ProcessGroup,
-    install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
-    model_name: str,
-    pass_first_chunk: bool,
-    broadcast_result: bool,
-    unpatchify_patch_size: int | None,
    return_dict: bool = True,
    split_dim: str = "height",
 ) -> DecoderOutput | tuple[torch.Tensor]:
-    install(vae, group, split_dim)
+    install_wan_spatial_shard_decode(vae, group, split_dim=split_dim)

    if z.shape[2] == 0:
-        raise ValueError(f"{model_name} spatial-shard VAE decode expects at least one latent frame.")
+        raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.")

    # Non-rank-0 ranks must still run the decoder every chunk to stay in lockstep with
    # the halo/all-gather collectives; they just skip keeping/assembling the output.
@@ -919,72 +786,31 @@ def spatial_shard_decode_impl(

    vae.clear_cache()
    try:
-        context_factory = getattr(vae, "_execution_context", None)
-        context = context_factory() if callable(context_factory) else nullcontext()
+        context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext()
        with context:
            x = vae.post_quant_conv(z)
            decoded_chunks = []
-            spatial_forward = getattr(vae.decoder, "_vllm_omni_spatial_shard_forward")
            for i in range(z.shape[2]):
                vae._conv_idx = [0]
-                if pass_first_chunk:
-                    chunk = spatial_forward(
-                        x[:, :, i : i + 1, :, :],
-                        feat_cache=vae._feat_map,
-                        feat_idx=vae._conv_idx,
-                        first_chunk=i == 0,
-                        split_dim=split_dim,
-                    )
-                else:
-                    chunk = spatial_forward(
-                        x[:, :, i : i + 1, :, :],
-                        feat_cache=vae._feat_map,
-                        feat_idx=vae._conv_idx,
-                        split_dim=split_dim,
-                    )
+                chunk = vae.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=vae._feat_map,
+                    feat_idx=vae._conv_idx,
+                    first_chunk=(i == 0),
+                )
                if produce_output:
                    decoded_chunks.append(chunk)

            if produce_output:
                out = torch.cat(decoded_chunks, dim=2)
-                if unpatchify_patch_size is not None:
-                    out = unpatchify(out, patch_size=unpatchify_patch_size)
+                if vae.config.patch_size is not None:
+                    out = unpatchify(out, patch_size=vae.config.patch_size)
                out = torch.clamp(out, min=-1.0, max=1.0)
            else:
                out = z.new_zeros(0)
    finally:
        vae.clear_cache()
-        # Halo buffers are request-shape scratch space allocated outside the
-        # tagged weight pool. Drop live references so sleep/offload never has
-        # to preserve or discard them as model state.
-        clear_spatial_shard_runtime_buffers(vae)
-
-    if broadcast_result:
-        out = vae.distributed_executor._sync_final_result(out, z.ndim, z.device, vae.dtype)

    if not return_dict:
        return (out,)
    return DecoderOutput(sample=out)
-
-
-def spatial_shard_decode(
-    vae: SpatialShardVAE,
-    z: torch.Tensor,
-    *,
-    group: dist.ProcessGroup,
-    return_dict: bool = True,
-    split_dim: str = "height",
-) -> DecoderOutput | tuple[torch.Tensor]:
-    """Decode a Wan latent with spatial activation sharding."""
-    return spatial_shard_decode_impl(
-        vae,
-        z,
-        group=group,
-        install=install_wan_spatial_shard_decode,
-        model_name="Wan",
-        pass_first_chunk=True,
-        broadcast_result=False,
-        unpatchify_patch_size=vae.config.patch_size,
-        return_dict=return_dict,
-        split_dim=split_dim,
-    )
@@ -155,11 +155,6 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
        parallel_size: int,
        mode: str = "tile",
    ) -> None:
-        # ``auto`` is the cross-model default. MiniMax H3 has no spatial-shard
-        # capability, so preserve its native tile path for that default while
-        # continuing to reject explicitly requested spatial modes.
-        if mode == "auto":
-            mode = "tile"
        if mode != "tile":
            raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}")
        group = get_dit_group()
@@ -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",
+]
@@ -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",
+]
@@ -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",
+]
@@ -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"]
@@ -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"]
@@ -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"]
@@ -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"]
@@ -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"]
@@ -30,9 +30,6 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    prepare_pipeline_spatial_shard_decode,
-)
 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 (
@@ -46,6 +43,10 @@ from vllm_omni.diffusion.offloader import get_offload_backend
 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 (
@@ -246,7 +247,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                    custom_pipeline_name=custom_pipeline_name,
                    device=self.device,
                )
-                prepare_pipeline_spatial_shard_decode(self.pipeline)
+                # 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(
@@ -503,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.
@@ -527,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),
@@ -778,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,
@@ -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()
@@ -491,6 +491,7 @@ class OrchestratorArgs:
    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
@@ -505,7 +506,7 @@ class OrchestratorArgs:
    diffusion_kv_cache_skip_layers: str | None = None
    cfg_parallel_size: int = 1
    vae_patch_parallel_size: int = 1
-    vae_parallel_mode: str = "auto"
+    vae_parallel_mode: str = "tile"
    text_encoder_tp_size: int = 1
    default_sampling_params: str | None = None
    max_generated_image_size: int | None = None
@@ -972,7 +972,7 @@ class AsyncOmniEngine:
            cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
            pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
            vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
-            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "auto"
+            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile"
            text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1
            enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
            use_hsdp = normalized_kwargs.get("use_hsdp", False)
@@ -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),
@@ -629,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(
@@ -750,12 +759,10 @@ class OmniServeCommand(CLISubcommand):
        omni_config_group.add_argument(
            "--vae-parallel-mode",
            type=str,
-            default="auto",
-            choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"],
+            default="tile",
+            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
            help="VAE parallel decode strategy for diffusion models. "
-            "'auto' (default) selects the longer spatial-shard axis for validated Wan and distributed QwenImage "
-            "tiled decodes "
-            "on a full DiT group, then falls back to patch/tile decode; 'tile' always uses patch/tile decode; "
+            "'tile' (default) uses patch/tile parallel decode; "
            "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits "
            "decoder feature maps along height/width and exchanges halo regions. The "
            "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. "