feat(diffusion): add QwenImage spatial-shard decode

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

docs/design/feature/vae_parallel.md

@@ -449,20 +449,21 @@ Complete examples in the codebase:
 
 ---
 
-## Spatially-Sharded Decode (Wan)
+## Spatially-Sharded Decode (Wan and QwenImage)
 
-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"`).
+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.
 
 ### How it differs from tile parallel
 
-| 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 |
+| 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 |
 
-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`.
+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.
 
 ### Wiring
 
@@ -473,15 +474,17 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
   -> DiffusionParallelConfig.vae_parallel_mode
   -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...)
   -> DistributedVaeExecutor.parallel_mode
-  -> DistributedAutoencoderKLWan.tiled_decode dispatch
+  -> capability-based pipeline preparation
+  -> DistributedAutoencoderKLWan/DistributedAutoencoderKLQwenImage.tiled_decode dispatch
 ```
 
-`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.
+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.
 
 ### Notes
 
-- 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.
+- 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.
 - 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`.
 
 ---
 

docs/user_guide/diffusion/parallelism/vae_parallelism.md

@@ -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 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/QwenImage 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 | `"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). |
+| `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). |
 
 Additional requirements:
 
@@ -126,9 +126,9 @@ Additional requirements:
 
 ---
 
-## Spatially-Sharded Decode (Wan)
+## Spatially-Sharded Decode (Wan and QwenImage)
 
-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"`.
+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.
 
 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="spatial_shard_width", # or "spatial_shard_height"
+        vae_parallel_mode="auto",                # default
     ),
 )
 ```
@@ -152,14 +152,18 @@ 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 spatial_shard_width
+    --vae-parallel-mode auto
 ```
 
 **Constraints and behavior:**
 
-- 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).
+- 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.
 
 For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:
 

tests/config/test_omni_config.py

@@ -579,6 +579,7 @@ 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
 
 

tests/diffusion/distributed/test_qwenimage_spatial_shard.py

@@ -0,0 +1,233 @@
+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

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -1,4 +1,5 @@
 import os
+from contextlib import contextmanager
 from types import SimpleNamespace
 
 import pytest
@@ -10,6 +11,8 @@ 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
 
@@ -17,6 +20,23 @@ 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():
@@ -255,12 +275,14 @@ 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())
-    mid = mid_rank_pad(x)
+    with _active_spatial_shard("height"):
+        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())
-    last = last_rank_pad(x)
+    with _active_spatial_shard("height"):
+        last = last_rank_pad(x)
     assert last.shape == (1, 1, 3, 3)
 
 
@@ -270,13 +292,15 @@ 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(), split_dim="width")
-    mid = mid_rank_pad(x)
+    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
+    with _active_spatial_shard("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(), split_dim="width")
-    last = last_rank_pad(x)
+    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
+    with _active_spatial_shard("width"):
+        last = last_rank_pad(x)
     assert last.shape == (1, 1, 2, 3)
 
 
@@ -299,6 +323,113 @@ 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):
@@ -444,3 +575,99 @@ 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

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -120,3 +120,34 @@ 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

tests/diffusion/test_diffusion_config_propagation.py

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

tests/diffusion/test_diffusion_model_runner.py

@@ -512,35 +512,6 @@ 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()."""
 
@@ -608,18 +579,6 @@ def test_execute_model_batch_rejects_output_count_mismatch(monkeypatch):
         DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_execute_model_batch_rejects_mixed_vae_quality() -> None:
-    runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")]))
-    sched = _make_scheduler_output(num_reqs=2)
-    sched.scheduled_new_reqs[0].req.sampling_params.quality = "lossless"
-    sched.scheduled_new_reqs[1].req.sampling_params.quality = "high"
-
-    with pytest.raises(RuntimeError, match="cannot mix request quality levels"):
-        DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config)
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_model_batch_routes_one_output_per_request(monkeypatch):
@@ -770,7 +729,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
-    optimization_pool_states = []
+    preparation_pool_states = []
 
     @contextmanager
     def _memory_pool_context(tag):
@@ -831,8 +790,8 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
     monkeypatch.setattr(
         model_runner_module,
-        "optimize_pipeline_vaes",
-        lambda pipeline, od_config: optimization_pool_states.append(memory_pool_active),
+        "prepare_pipeline_spatial_shard_decode",
+        lambda pipeline: preparation_pool_states.append(memory_pool_active),
     )
     monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
     monkeypatch.setattr(
@@ -844,7 +803,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 optimization_pool_states == [True]
+    assert preparation_pool_states == [True]
 
 
 @pytest.mark.core_model

tests/diffusion/test_diffusion_worker.py

@@ -102,7 +102,6 @@ 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]
@@ -129,7 +128,6 @@ class TestDiffusionWorkerSleep:
 
         # Verify sleep was called with correct tags
         mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
-        clear_vae_caches.assert_called_once_with(mock_gpu_worker.model_runner.pipeline)
         assert bool(result) is True
         # Verify buffers were NOT saved (level 1 doesn't save buffers)
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0

tests/diffusion/vae_optimizations/test_vae_fastpaths.py

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

tests/entrypoints/test_serve.py

@@ -76,15 +76,16 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     assert args.cfg_parallel_size == 4
 
 
-def test_serve_parser_accepts_wan_decode_precision() -> None:
+def test_serve_parser_defaults_vae_parallel_mode_to_auto_and_accepts_tile() -> None:
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="subcommand")
     OmniServeCommand().subparser_init(subparsers)
 
-    args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-decode-precision", "bf16"])
+    default_args = parser.parse_args(["serve", "fake-model", "--omni"])
+    tile_args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-parallel-mode", "tile"])
 
-    assert args.vae_decode_precision == "bf16"
-    assert args.get_explicit_kwargs_dict()["vae_decode_precision"] == "bf16"
+    assert default_args.vae_parallel_mode == "auto"
+    assert tile_args.vae_parallel_mode == "tile"
 
 
 def _make_headless_args(**kwargs) -> TrackingNamespace:

vllm_omni/config/omni_config.py

@@ -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 = "tile"
+    vae_parallel_mode: str = "auto"
     use_hsdp: bool = False
     mask_sp_padding: bool = False
     hsdp_shard_size: int = -1
@@ -466,9 +466,10 @@ 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 {"tile", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}:
             raise ValueError(
-                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+                "vae_parallel_mode must be one of "
+                "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
                 f"but got {self.vae_parallel_mode!r}."
             )
 
@@ -573,7 +574,6 @@ 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

vllm_omni/config/stage_config.py

@@ -401,7 +401,6 @@ class StageDeployConfig:
     step_execution: bool | None = None
     vae_use_slicing: bool | None = None
     vae_use_tiling: bool | None = None
-    vae_decode_precision: str | None = None
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/diffusion/data.py

@@ -206,10 +206,13 @@ class DiffusionParallelConfig:
     process group.
     """
 
-    vae_parallel_mode: str = "tile"
+    vae_parallel_mode: str = "auto"
     """VAE parallel decode strategy.
 
-    - "tile": Patch/tile parallel decode (default). Each rank decodes a subset
+    - "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
       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
@@ -250,8 +253,9 @@ 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 {"tile", "spatial_shard_height", "spatial_shard_width"}, (
-            "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+        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'}, "
             f"but got {self.vae_parallel_mode!r}."
         )
         if self.allgather_degree > 1:
@@ -711,9 +715,6 @@ 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
@@ -926,11 +927,6 @@ 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', "

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py

@@ -4,16 +4,23 @@
 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__)
 
@@ -104,10 +111,39 @@ 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,

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -18,11 +18,25 @@ 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:
@@ -106,7 +120,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=getattr(self, "_vllm_decode_dtype", self.dtype),
+            output_dtype=self.dtype,
         )
         return tiletask_list, grid_spec
 
@@ -270,45 +284,30 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         dec = torch.clamp(dec, min=-1.0, max=1.0)
         return dec
 
-    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_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:
-        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
+        return spatial_shard_decode_enabled(self, z, model_name="Wan")
 
-        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 _prepare_spatial_shard_decode(self) -> None:
+        prepare_spatial_shard_decode(
+            self,
+            install=wan_spatial_shard.install_wan_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):
-            split_dim = self._spatial_shard_decode_split_dim()
+            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 Wan VAE spatial_shard_%s mode", split_dim)
             return wan_spatial_shard.spatial_shard_decode(
                 self,

vllm_omni/diffusion/distributed/autoencoders/qwenimage_spatial_shard.py

@@ -0,0 +1,86 @@
+# 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,
+    )

vllm_omni/diffusion/distributed/autoencoders/spatial_shard.py

@@ -0,0 +1,142 @@
+# 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()

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -7,19 +7,21 @@
 #   (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 adds the Wan causal-conv ``feat_cache`` handling.
-"""Spatially-sharded Wan VAE decode.
+# width and supports the causal-conv ``feat_cache`` handling used by Wan and
+# QwenImage.
+"""Shared spatial-shard kernels plus the Wan VAE adapter.
 
-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.
+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.
 """
 
 from __future__ import annotations
 
 import math
+from collections.abc import Callable
 from contextlib import nullcontext
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -34,6 +36,8 @@ 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__)
 
 
@@ -52,12 +56,17 @@ _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 Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
+    raise ValueError(f"Unsupported 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:
@@ -155,9 +164,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"Wan VAE world_size must be >= 1, got {world_size}.")
+        raise ValueError(f"VAE world_size must be >= 1, got {world_size}.")
     if not 0 <= rank < world_size:
-        raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
+        raise ValueError(f"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)
@@ -323,21 +332,27 @@ class WanDistZeroPad2d(nn.Module):
         padding: tuple[int, int, int, int],
         group: dist.ProcessGroup,
         *,
-        split_dim: str = "height",
-        split_padding: tuple[int, int] | None = None,
+        defer_split_after_padding: bool = False,
     ) -> None:
         super().__init__()
         self.padding = 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.defer_split_after_padding = defer_split_after_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 self.split_dim == "height":
+            if split_dim == "height":
                 top = top if rank == 0 else 0
                 bottom = bottom if rank == world_size - 1 else 0
             else:
@@ -351,8 +366,7 @@ class WanDistConv2d(nn.Conv2d):
         self,
         source: nn.Conv2d,
         group: dist.ProcessGroup,
-        split_dim: str = "height",
-        split_padding: tuple[int, int] | None = None,
+        deferred_padding: tuple[int, int, int, int] | None = None,
     ):
         super().__init__(
             source.in_channels,
@@ -369,176 +383,253 @@ class WanDistConv2d(nn.Conv2d):
         )
         self.load_state_dict(source.state_dict())
         self.group = group
-        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]] = {}
+        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]] = {}
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.pad(x, self._non_split_padding)
+        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_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
             x,
             group=self.group,
-            halo_size=self.halo_size,
-            split_dim=self.split_dim,
+            halo_size=halo_size,
+            split_dim=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[self.split_tensor_dim])
+        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,
+        )
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                self.split_tensor_dim,
+                split_tensor_dim,
                 shift,
-                x_padded.shape[self.split_tensor_dim] - shift,
+                x_padded.shape[split_tensor_dim] - shift,
             )
         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)
+        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
+        return _zero_invalid_extent(out, split_dim=split_dim)
 
-    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
-        trim_params = self._trim_cache.get(local_extent)
+    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)
         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=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,
+                halo_size=halo_size,
+                pad_before=split_pad_left,
+                pad_after=split_pad_right,
+                kernel_extent=kernel_extent,
+                stride_extent=stride_extent,
             )
-            self._trim_cache[local_extent] = trim_params
+            self._trim_cache[cache_key] = trim_params
         return trim_params
 
 
-class WanDistCausalConv3d(nn.Conv3d):
-    def __init__(
+# 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(
         self,
         source: nn.Conv3d,
         group: dist.ProcessGroup,
-        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,
-        )
+    ) -> None:
         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)
-        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]] = {}
+        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]] = {}
 
-    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        padding = list(self._padding)
+    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)
         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=self.halo_size,
-            split_dim=self.split_dim,
+            halo_size=halo_size,
+            split_dim=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[self.split_tensor_dim])
+        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,
+        )
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                self.split_tensor_dim,
+                split_tensor_dim,
                 shift,
-                x_padded.shape[self.split_tensor_dim] - shift,
+                x_padded.shape[split_tensor_dim] - shift,
             )
-        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)
+        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)
 
-    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
-        trim_params = self._trim_cache.get(local_extent)
+    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)
         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=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,
+                halo_size=halo_size,
+                pad_before=split_pad_left,
+                pad_after=split_pad_right,
+                kernel_extent=kernel_extent,
+                stride_extent=stride_extent,
             )
-            self._trim_cache[local_extent] = trim_params
+            self._trim_cache[cache_key] = 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,
@@ -582,12 +673,15 @@ def _trim_local_conv_output(
     return out
 
 
-def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None:
+def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> 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)
@@ -609,60 +703,43 @@ def _replace_child(
     name: str,
     child: nn.Module,
     group: dist.ProcessGroup,
-    split_dim: str,
+    *,
+    causal_conv_class_name: str,
+    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
 ) -> None:
-    if child.__class__.__name__ == "WanCausalConv3d":
-        setattr(
-            parent,
-            name,
-            WanDistCausalConv3d(
-                child,
-                group,
-                split_dim=split_dim,
-            ),
-        )
+    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)
-        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,
+            SpatialShardZeroPad2d(
+                padding,
                 group,
-                split_dim=split_dim,
-                split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
+                defer_split_after_padding=parent.__class__.__name__ == "Sequential",
             ),
         )
         return
     if isinstance(child, nn.Conv2d):
-        split_padding = None
+        deferred_padding = None
         if name == "1" and parent.__class__.__name__ == "Sequential":
-            # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a
+            # Wan/QwenImage downsampling 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, WanDistZeroPad2d):
-                split_padding = prev.split_padding
+            if isinstance(prev, SpatialShardZeroPad2d):
+                deferred_padding = prev.padding
             elif isinstance(prev, nn.ZeroPad2d):
-                pad = prev.padding
-                split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1]))
+                deferred_padding = tuple(int(value) for value in prev.padding)
         setattr(
             parent,
             name,
-            WanDistConv2d(
+            SpatialShardConv2d(
                 child,
                 group,
-                split_dim=split_dim,
-                split_padding=split_padding,
+                deferred_padding=deferred_padding,
             ),
         )
 
@@ -670,21 +747,45 @@ def _replace_child(
 def _patch_decoder_modules(
     module: nn.Module,
     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,
     inside_attention: bool = False,
 ) -> None:
-    if module.__class__.__name__ == "WanAttentionBlock":
-        _patch_attention_block(module, group, split_dim)
+    if module.__class__.__name__ == attention_block_class_name:
+        _patch_attention_block(module, group)
         inside_attention = True
 
     for name, child in list(module.named_children()):
-        if child.__class__.__name__ == "WanCausalConv3d":
-            _replace_child(module, name, child, group, split_dim)
+        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,
+            )
             continue
         if isinstance(child, nn.Conv2d) and not inside_attention:
-            _replace_child(module, name, child, group, split_dim)
+            _replace_child(
+                module,
+                name,
+                child,
+                group,
+                causal_conv_class_name=causal_conv_class_name,
+                causal_conv_factory=causal_conv_factory,
+            )
             continue
-        _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)
+        _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,
+        )
 
 
 def _decoder_upsample_count(decoder: nn.Module) -> int:
@@ -695,47 +796,58 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
     return count
 
 
-def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None:
+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:
     """Patch ``vae.decoder`` once for spatially-sharded decode.
 
     This mutates the already-loaded decoder in place by swapping its spatial
-    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.
+    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.
 
-    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.
+    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.
     """
     _spatial_dim(split_dim)
-    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."
-            )
+    if getattr(vae, installed_attr, False):
         return
     decoder = getattr(vae, "decoder", None)
     if decoder is None:
-        raise ValueError("Wan spatial-shard VAE decode requires a decoder module.")
+        raise ValueError(f"{model_name} spatial-shard VAE decode requires a decoder module.")
 
-    _patch_decoder_modules(decoder, group, split_dim)
+    _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,
+    )
     upsample_count = _decoder_upsample_count(decoder)
     orig_forward = decoder.forward
 
-    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]
+    def _forward(self: nn.Module, x: torch.Tensor, *args: object, split_dim: str, **kwargs: object) -> torch.Tensor:
         tensor_dim = _spatial_dim(split_dim)
         input_extent = x.shape[tensor_dim]
         x, expected_extent = split_for_parallel_decode(
@@ -755,29 +867,50 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
             )
         )
         try:
-            out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
+            out = orig_forward(x, *args, **kwargs)
         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.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)
+    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)
 
 
-def spatial_shard_decode(
-    vae: Any,
+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,
     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_wan_spatial_shard_decode(vae, group, split_dim=split_dim)
+    install(vae, group, split_dim)
 
     if z.shape[2] == 0:
-        raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.")
+        raise ValueError(f"{model_name} 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.
@@ -786,31 +919,72 @@ def spatial_shard_decode(
 
     vae.clear_cache()
     try:
-        context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext()
+        context_factory = getattr(vae, "_execution_context", None)
+        context = context_factory() if callable(context_factory) 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]
-                chunk = vae.decoder(
-                    x[:, :, i : i + 1, :, :],
-                    feat_cache=vae._feat_map,
-                    feat_idx=vae._conv_idx,
-                    first_chunk=(i == 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,
+                    )
                 if produce_output:
                     decoded_chunks.append(chunk)
 
             if produce_output:
                 out = torch.cat(decoded_chunks, dim=2)
-                if vae.config.patch_size is not None:
-                    out = unpatchify(out, patch_size=vae.config.patch_size)
+                if unpatchify_patch_size is not None:
+                    out = unpatchify(out, patch_size=unpatchify_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,
+    )

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -155,6 +155,11 @@ 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()

vllm_omni/diffusion/vae_optimizations/__init__.py

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

vllm_omni/diffusion/vae_optimizations/gate.py

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

vllm_omni/diffusion/vae_optimizations/image.py

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

vllm_omni/diffusion/vae_optimizations/optimize.py

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

vllm_omni/diffusion/vae_optimizations/precision.py

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

vllm_omni/diffusion/vae_optimizations/triton_group_norm_silu.py

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

vllm_omni/diffusion/vae_optimizations/triton_wan_rmsnorm_silu.py

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

vllm_omni/diffusion/vae_optimizations/wan.py

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

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -30,6 +30,9 @@ 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 (
@@ -43,10 +46,6 @@ 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 (
@@ -247,9 +246,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     custom_pipeline_name=custom_pipeline_name,
                     device=self.device,
                 )
-                # Keep precision rematerialization in the tagged weights pool
-                # so sleep level 1 can offload every resulting parameter.
-                optimize_pipeline_vaes(self.pipeline, self.od_config)
+                prepare_pipeline_spatial_shard_decode(self.pipeline)
         time_after_load = time.perf_counter()
 
         logger.info(
@@ -506,11 +503,6 @@ 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.
@@ -535,8 +527,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
             with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config):
                 with record_function(record_name):
-                    with use_pipeline_vae_fast_path(self.pipeline, use_fast_vae):
-                        raw_outputs = self.pipeline.forward(batch)
+                    raw_outputs = self.pipeline.forward(batch)
                     outputs = _normalize_pipeline_outputs(
                         raw_outputs,
                         expected_count=len(reqs),
@@ -787,11 +778,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
                             if should_decode:
                                 clear_pipeline_stage_durations(self.pipeline)
-                                with use_pipeline_vae_fast_path(
-                                    self.pipeline,
-                                    getattr(req.sampling, "quality", None) == "high",
-                                ):
-                                    result = self.pipeline.post_decode(req)
+                                result = self.pipeline.post_decode(req)
                                 if result is not None:
                                     self._attach_stepwise_metrics(
                                         req,

vllm_omni/diffusion/worker/diffusion_worker.py

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

vllm_omni/engine/arg_utils.py

@@ -491,7 +491,6 @@ 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
@@ -506,7 +505,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 = "tile"
+    vae_parallel_mode: str = "auto"
     text_encoder_tp_size: int = 1
     default_sampling_params: str | None = None
     max_generated_image_size: int | None = None

vllm_omni/engine/async_omni_engine.py

@@ -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 "tile"
+            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "auto"
             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,7 +1026,6 @@ 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),

vllm_omni/entrypoints/cli/serve.py

@@ -629,15 +629,6 @@ 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(
@@ -759,10 +750,12 @@ class OmniServeCommand(CLISubcommand):
         omni_config_group.add_argument(
             "--vae-parallel-mode",
             type=str,
-            default="tile",
-            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
+            default="auto",
+            choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"],
             help="VAE parallel decode strategy for diffusion models. "
-            "'tile' (default) uses patch/tile parallel decode; "
+            "'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; "
             "'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. "