[Perf][Diffusion] Fuse lossless Wan VAE data movement

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

docs/design/feature/vae_parallel.md

@@ -449,21 +449,20 @@ Complete examples in the codebase:
 
 ---
 
-## Spatially-Sharded Decode (Wan and QwenImage)
+## Spatially-Sharded Decode (Wan)
 
-The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE and distributed QwenImage backend additionally support a **spatially-sharded decode** backend. `DiffusionParallelConfig.vae_parallel_mode="auto"` is the default; explicit `"tile"`, `"spatial_shard_height"`, and `"spatial_shard_width"` modes remain available.
+The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`).
 
 ### How it differs from tile parallel
 
-| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
-|--------|--------------------------|----------------------|-----------------------------------------------|
-| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W |
-| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
-| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0 and trim padding; QwenImage then broadcasts |
-| Activation memory | Bounded by tile size | Depends on the selected path | Grows with full output area divided across ranks |
-| Scope | Decode + encode | Wan/QwenImage decode selection; encode remains tiled | Decode only |
+| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
+|--------|--------------------------|-----------------------------------------------|
+| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
+| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
+| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
+| Scope | Decode + encode | Decode only |
 
-Spatial-shard decode swaps the decoder's spatial convolutions/padding for context-aware halo-exchanging variants so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Wan and QwenImage share the padding, 2D convolution, attention, halo, and trim primitives in `wan_spatial_shard.py`; `qwenimage_spatial_shard.py` supplies a causal-convolution subclass that remains discoverable by Diffusers' QwenImage cache accounting.
+Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`.
 
 ### Wiring
 
@@ -474,17 +473,15 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
   -> DiffusionParallelConfig.vae_parallel_mode
   -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...)
   -> DistributedVaeExecutor.parallel_mode
-  -> capability-based pipeline preparation
-  -> DistributedAutoencoderKLWan/DistributedAutoencoderKLQwenImage.tiled_decode dispatch
+  -> DistributedAutoencoderKLWan.tiled_decode dispatch
 ```
 
-The shared spatial policy gates the path: it requires distributed decode to be enabled, a 5D latent, a supported VAE capability, and `vae_patch_parallel_size == DiT group size`. The method is reached only after Diffusers has selected tiled decode for the request. In `auto` mode the longer axis that can cover every rank is selected, with width winning ties. This policy is cached only through normal kernel shape caches; it is reevaluated for every request.
+`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode.
 
 ### Notes
 
-- The decoder is patched **in place** the first time spatial-shard decode runs. Its wrappers consult a context-local split dimension and retain an exact direct fallback, allowing later requests to select tile, height, or width safely.
+- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
 - Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
-- QwenImage additionally covers cache discovery, compile compatibility, mixed tile/spatial requests, both axes, and all-rank result broadcast in `tests/diffusion/distributed/test_qwenimage_spatial_shard.py`.
 
 ---
 

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/QwenImage spatial-shard decode, which shards decoder feature maps along height or width.
+VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width.
 
 This is particularly useful for:
 - **High-resolution image generation** where VAE decode can become a memory bottleneck
@@ -113,7 +113,7 @@ In `DiffusionParallelConfig`:
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
-| `vae_parallel_mode` | str | `"auto"` | VAE parallel decode strategy: `"auto"` (shape/model/topology-aware selection), `"tile"`, `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding supports Wan and the distributed QwenImage backend. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan-and-qwenimage). |
+| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
 
 Additional requirements:
 
@@ -126,9 +126,9 @@ Additional requirements:
 
 ---
 
-## Spatially-Sharded Decode (Wan and QwenImage)
+## Spatially-Sharded Decode (Wan)
 
-The default `vae_parallel_mode="auto"` selects spatial-shard decode only after a Wan or distributed QwenImage request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. It splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests, partial groups, and other VAE families keep their existing path. Set `vae_parallel_mode="tile"` to force the previous behavior, or `"spatial_shard_height"` / `"spatial_shard_width"` to force an axis.
+The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`.
 
 Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance.
 
@@ -141,7 +141,7 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(
         tensor_parallel_size=2,
         vae_patch_parallel_size=2,               # must match the DiT group size
-        vae_parallel_mode="auto",                # default
+        vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height"
     ),
 )
 ```
@@ -152,18 +152,14 @@ Or from the CLI / serving entrypoint:
 vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \
     --tensor-parallel-size 2 \
     --vae-patch-parallel-size 2 \
-    --vae-parallel-mode auto
+    --vae-parallel-mode spatial_shard_width
 ```
 
 **Constraints and behavior:**
 
-- Spatial-shard decode is **decode-only** and implemented for the **Wan** VAE and the distributed **QwenImage** backend used by Qwen-Image and Krea2. Pipelines using a different VAE backend keep their existing behavior.
-- It requires `vae_patch_parallel_size` to **match the DiT process group size**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning.
-- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.
-- Wan preserves its rank-0-only decode result, while QwenImage preserves its existing all-rank result by broadcasting after rank-0 assembly.
-- Spatial sharding trades activation memory for latency as resolution grows. Use `vae_parallel_mode="tile"` when bounded VAE memory is more important than decode speed.
-
-The automatic policy is adapted from [SGLang #28071](https://github.com/sgl-project/sglang/pull/28071) and narrowed to vLLM-Omni's parity-tested Wan and QwenImage backends.
+- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes.
+- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime.
+- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension).
 
 For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:
 

tests/config/test_omni_config.py

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

tests/diffusion/distributed/test_autoencoder_kl_wan.py

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

tests/diffusion/distributed/test_qwenimage_spatial_shard.py

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

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -1,5 +1,4 @@
 import os
-from contextlib import contextmanager
 from types import SimpleNamespace
 
 import pytest
@@ -11,8 +10,6 @@ from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
     DistributedAutoencoderKLWan,
-    prepare_pipeline_wan_spatial_shard_decode,
-    select_auto_spatial_shard_split_dim,
 )
 from vllm_omni.platforms import current_omni_platform
 
@@ -20,23 +17,6 @@ from vllm_omni.platforms import current_omni_platform
 # the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel.
 
 
-@contextmanager
-def _active_spatial_shard(split_dim: str):
-    token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
-        wan_spatial_shard.SpatialShardContext(
-            input_extent=2,
-            local_input_extent=2,
-            split_dim=split_dim,
-            rank=0,
-            world_size=3,
-        )
-    )
-    try:
-        yield
-    finally:
-        wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_split_for_parallel_decode_pads_uneven_height():
@@ -275,14 +255,12 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
     mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
-    with _active_spatial_shard("height"):
-        mid = mid_rank_pad(x)
+    mid = mid_rank_pad(x)
     assert mid.shape == (1, 1, 2, 3)
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
     last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
-    with _active_spatial_shard("height"):
-        last = last_rank_pad(x)
+    last = last_rank_pad(x)
     assert last.shape == (1, 1, 3, 3)
 
 
@@ -292,15 +270,13 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
     x = torch.ones((1, 1, 2, 2))
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
-    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
-    with _active_spatial_shard("width"):
-        mid = mid_rank_pad(x)
+    mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    mid = mid_rank_pad(x)
     assert mid.shape == (1, 1, 2, 2)
 
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
-    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
-    with _active_spatial_shard("width"):
-        last = last_rank_pad(x)
+    last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
+    last = last_rank_pad(x)
     assert last.shape == (1, 1, 2, 3)
 
 
@@ -323,113 +299,6 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt
     assert vae._spatial_shard_decode_enabled(z) is False
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-@pytest.mark.parametrize(
-    "z_shape,world_size,expected",
-    [
-        ((1, 16, 1, 60, 104), 4, "width"),
-        ((1, 16, 5, 60, 104), 4, "width"),
-        ((1, 16, 5, 104, 60), 4, "height"),
-        ((1, 16, 5, 2, 2), 4, None),
-        ((1, 16, 5, 64, 64), 4, "width"),
-    ],
-)
-def test_auto_spatial_shard_selector_is_shape_aware(z_shape, world_size, expected):
-    assert select_auto_spatial_shard_split_dim(z_shape, world_size) == expected
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyPatch):
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
-        lambda group=None: 4,
-    )
-
-    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
-    vae.use_tiling = True
-    vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=4, parallel_mode="auto")
-    vae.is_distributed_enabled = lambda: True
-    large = torch.zeros((1, 16, 5, 60, 104))
-    unshardable = torch.zeros((1, 16, 1, 2, 2))
-
-    assert vae._spatial_shard_decode_split_dim(large, 4) == "width"
-    assert vae._spatial_shard_decode_enabled(large) is True
-    assert vae._spatial_shard_decode_enabled(unshardable) is False
-
-    vae.distributed_executor.parallel_size = 2
-    assert vae._spatial_shard_decode_enabled(large) is False
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch):
-    vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
-    object.__setattr__(
-        vae,
-        "distributed_executor",
-        SimpleNamespace(group=object(), parallel_size=2, parallel_mode="auto"),
-    )
-    installed = []
-
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
-        lambda pipeline: SimpleNamespace(vaes=[vae]),
-    )
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.is_initialized",
-        lambda: True,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
-        lambda group=None: 4,
-    )
-    monkeypatch.setattr(
-        wan_spatial_shard,
-        "install_wan_spatial_shard_decode",
-        lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)),
-    )
-
-    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
-    assert installed == []
-
-    vae.distributed_executor.parallel_size = 4
-    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
-    assert installed == [(vae, vae.distributed_executor.group, "height")]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-@torch.no_grad()
-def test_installed_spatial_decoder_keeps_direct_fallback_and_can_change_axis():
-    vae = DistributedAutoencoderKLWan(
-        base_dim=8,
-        decoder_base_dim=8,
-        z_dim=4,
-        dim_mult=[1, 2],
-        num_res_blocks=1,
-        attn_scales=[],
-        temperal_downsample=[False],
-        latents_mean=[0.0] * 4,
-        latents_std=[1.0] * 4,
-        scale_factor_temporal=1,
-        scale_factor_spatial=2,
-    ).eval()
-    latents = torch.randn(1, 4, 1, 4, 5)
-    parameter_names = {name for name, _ in vae.named_parameters()}
-    state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
-    reference = vae.decode(latents, return_dict=False)[0]
-
-    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="height")
-    assert any(getattr(module, "_vllm_omni_dynamic_spatial_shard_conv", False) for module in vae.decoder.modules())
-    wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="width")
-
-    assert {name for name, _ in vae.named_parameters()} == parameter_names
-    vae.load_state_dict(state_dict, strict=True)
-    assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
@@ -575,99 +444,3 @@ def test_spatial_shard_decode_matches_reference(split_dim: str):
     assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
         f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
     )
-
-
-def _auto_spatial_shard_decode_worker(rank: int, return_dict, master_port: str) -> None:
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        destroy_model_parallel,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = master_port
-    device = current_omni_platform.get_torch_device(rank)
-    current_omni_platform.set_device(device)
-    backend = current_omni_platform.dist_backend
-    init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
-    initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
-
-    try:
-        vae = DistributedAutoencoderKLWan.from_pretrained(
-            _SPATIAL_SHARD_MODEL,
-            subfolder=_SPATIAL_SHARD_SUBFOLDER,
-            torch_dtype=torch.float32,
-        ).to(device=device, dtype=torch.float32)
-        vae.eval()
-        generator = torch.Generator(device=device).manual_seed(1)
-        small = torch.randn((1, vae.config.z_dim, 1, 60, 104), generator=generator, device=device)
-        landscape = torch.randn((1, vae.config.z_dim, 2, 60, 104), generator=generator, device=device)
-        portrait = torch.randn((1, vae.config.z_dim, 2, 104, 60), generator=generator, device=device)
-
-        with torch.inference_mode():
-            vae.use_tiling = False
-            vae.set_parallel_size(1, mode="tile")
-            small_reference = vae.decode(small, return_dict=False)[0]
-
-            vae.use_tiling = True
-            vae.set_parallel_size(2, mode="tile")
-            tile_before_install = vae.decode(small, return_dict=False)[0]
-
-            vae.set_parallel_size(2, mode="auto")
-            assert vae._spatial_shard_decode_split_dim(small, 2) == "width"
-            small_sharded = vae.decode(small, return_dict=False)[0]
-            assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
-            vae.decode(landscape, return_dict=False)
-
-            vae.set_parallel_size(2, mode="tile")
-            tile_after_install = vae.decode(small, return_dict=False)[0]
-
-            vae.use_tiling = False
-            vae.set_parallel_size(1, mode="tile")
-            portrait_reference = vae.decode(portrait, return_dict=False)[0]
-
-            vae.use_tiling = True
-            vae.set_parallel_size(2, mode="auto")
-            assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
-            portrait_sharded = vae.decode(portrait, return_dict=False)[0]
-
-        assert all(
-            buffer is None
-            for module in vae.decoder.modules()
-            for name, buffer in module._buffers.items()
-            if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
-        )
-
-        if rank == 0:
-            small_diff = (small_sharded.float() - small_reference.float()).abs()
-            portrait_diff = (portrait_sharded.float() - portrait_reference.float()).abs()
-            return_dict["tile_fallback_exact"] = torch.equal(tile_after_install, tile_before_install)
-            return_dict["small_max_abs_diff"] = small_diff.max().item()
-            return_dict["small_mean_abs_diff"] = small_diff.mean().item()
-            return_dict["portrait_max_abs_diff"] = portrait_diff.max().item()
-            return_dict["portrait_mean_abs_diff"] = portrait_diff.mean().item()
-    finally:
-        destroy_model_parallel()
-        if dist.is_initialized():
-            dist.destroy_process_group()
-
-
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_auto_spatial_shard_switches_shape_paths_safely():
-    manager = mp.get_context("spawn").Manager()
-    return_dict = manager.dict()
-    mp.spawn(
-        _auto_spatial_shard_decode_worker,
-        args=(return_dict, "29502"),
-        nprocs=2,
-        join=True,
-    )
-
-    assert return_dict.get("tile_fallback_exact") is True
-    assert return_dict["small_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
-    assert return_dict["small_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
-    assert return_dict["portrait_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
-    assert return_dict["portrait_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE

tests/diffusion/kernels/test_wan_vae_data_movement.py

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

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -120,34 +120,3 @@ def test_tp_accepts_checkpoint_supported_sizes():
     arch = MiniMaxH3DiTArchConfig()
     for tp_size in (1, 2, 4, 7):
         model._validate_tp_config(arch=arch, tp_size=tp_size)
-
-
-@pytest.mark.parametrize(("parallel_size", "enabled"), [(1, False), (2, True)])
-def test_video_vae_auto_mode_preserves_native_tile_parallelism(monkeypatch, parallel_size, enabled):
-    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
-
-    parallel_state = {}
-    video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE)
-    nn.Module.__init__(video_vae)
-    video_vae.remote = nn.Identity()
-    video_vae.model = type("FakeModel", (), {"parallel_tiling": False})()
-
-    monkeypatch.setattr(vae_module, "get_dit_group", lambda: object())
-    monkeypatch.setattr(vae_module.dist, "get_world_size", lambda group: 2)
-    monkeypatch.setattr(vae_module.dist, "get_rank", lambda group: 0)
-    monkeypatch.setattr(
-        vae_module.importlib,
-        "import_module",
-        lambda package: type(
-            "FakeParallelModule",
-            (),
-            {"get_parallel_state": staticmethod(lambda: parallel_state)},
-        )(),
-    )
-
-    video_vae.set_parallel_size(parallel_size, mode="auto")
-
-    assert video_vae.parallel_size == parallel_size
-    assert video_vae.model.parallel_tiling is enabled
-    assert parallel_state["group_size"] == parallel_size
-    assert parallel_state["sp_enabled"] is enabled

tests/diffusion/test_diffusion_model_runner.py

@@ -728,18 +728,7 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
-    memory_pool_active = False
-    preparation_pool_states = []
-
-    @contextmanager
-    def _memory_pool_context(tag):
-        nonlocal memory_pool_active
-        assert tag == "weights"
-        memory_pool_active = True
-        try:
-            yield
-        finally:
-            memory_pool_active = False
+    prepared_pipelines = []
 
     class _DummyLoader:
         def __init__(self, load_config, od_config=None):
@@ -790,20 +779,20 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
     monkeypatch.setattr(
         model_runner_module,
-        "prepare_pipeline_spatial_shard_decode",
-        lambda pipeline: preparation_pool_states.append(memory_pool_active),
+        "prepare_pipeline_wan_vae_data_movement",
+        lambda pipeline: prepared_pipelines.append(pipeline),
     )
     monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
     monkeypatch.setattr(
         model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
     )
 
-    DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context)
+    DiffusionModelRunner.load_model(runner)
 
     assert runner.cache_backend is None
     assert runner.od_config.cache_backend is None
     assert dummy_cache_backend.enabled is False
-    assert preparation_pool_states == [True]
+    assert prepared_pipelines == [runner.pipeline]
 
 
 @pytest.mark.core_model

tests/entrypoints/test_serve.py

@@ -76,18 +76,6 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     assert args.cfg_parallel_size == 4
 
 
-def test_serve_parser_defaults_vae_parallel_mode_to_auto_and_accepts_tile() -> None:
-    parser = TrackingArgumentParser()
-    subparsers = parser.add_subparsers(dest="subcommand")
-    OmniServeCommand().subparser_init(subparsers)
-
-    default_args = parser.parse_args(["serve", "fake-model", "--omni"])
-    tile_args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-parallel-mode", "tile"])
-
-    assert default_args.vae_parallel_mode == "auto"
-    assert tile_args.vae_parallel_mode == "tile"
-
-
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",

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 = "auto"
+    vae_parallel_mode: str = "tile"
     use_hsdp: bool = False
     mask_sp_padding: bool = False
     hsdp_shard_size: int = -1
@@ -466,10 +466,9 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
             raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1")
         if self.ulysses_mode not in {"strict", "advanced_uaa"}:
             raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'")
-        if self.vae_parallel_mode not in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}:
             raise ValueError(
-                "vae_parallel_mode must be one of "
-                "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+                "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
                 f"but got {self.vae_parallel_mode!r}."
             )
 

vllm_omni/diffusion/data.py

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

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py

@@ -4,23 +4,16 @@
 from typing import Any
 
 import torch
-import torch.distributed as dist
 from diffusers.models.autoencoders import AutoencoderKLQwenImage
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.distributed.autoencoders import qwenimage_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
     DistributedOperator,
     DistributedVaeMixin,
     GridSpec,
     TileTask,
 )
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
-    prepare_spatial_shard_decode,
-    resolve_spatial_shard_split_dim,
-    spatial_shard_decode_enabled,
-)
 
 logger = init_logger(__name__)
 
@@ -111,39 +104,10 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
         ]
         return dec
 
-    def _spatial_shard_decode_split_dim(
-        self,
-        z: torch.Tensor | None = None,
-        world_size: int | None = None,
-    ) -> str | None:
-        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
-
-    def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        return spatial_shard_decode_enabled(self, z, model_name="QwenImage")
-
-    def _prepare_spatial_shard_decode(self) -> None:
-        prepare_spatial_shard_decode(
-            self,
-            install=qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode,
-        )
-
     def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
         if not self.is_distributed_enabled():
             return super().tiled_decode(z, return_dict=return_dict)
 
-        if self._spatial_shard_decode_enabled(z):
-            world_size = dist.get_world_size(group=self.distributed_executor.group)
-            split_dim = self._spatial_shard_decode_split_dim(z, world_size)
-            assert split_dim is not None
-            logger.debug("Decode running with QwenImage VAE spatial_shard_%s mode", split_dim)
-            return qwenimage_spatial_shard.spatial_shard_decode(
-                self,
-                z,
-                group=self.distributed_executor.group,
-                return_dict=return_dict,
-                split_dim=split_dim,
-            )
-
         logger.debug("Decode running with distributed executor")
         result = self.distributed_executor.execute(
             z,

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -18,27 +18,26 @@ 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.diffusion.distributed.autoencoders.wan_vae_data_movement import (
+    install_wan_vae_data_movement,
 )
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-def prepare_pipeline_wan_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
-    """Backward-compatible alias for the model-agnostic pipeline dispatcher."""
-    prepare_pipeline_spatial_shard_decode(pipeline)
+def prepare_pipeline_wan_vae_data_movement(pipeline: torch.nn.Module) -> None:
+    """Install Wan data-movement methods before offload or compilation setup."""
+    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+    for vae in ModuleDiscovery.discover(pipeline).vaes:
+        if isinstance(vae, OmniAutoencoderKLWan):
+            install_wan_vae_data_movement(vae)
 
 
 class OmniAutoencoderKLWan(AutoencoderKLWan):
     def _execution_context(self):
+        install_wan_vae_data_movement(self)
         try:
             first_param = next(self.parameters())
         except StopIteration:
@@ -284,30 +283,45 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         dec = torch.clamp(dec, min=-1.0, max=1.0)
         return dec
 
-    def _spatial_shard_decode_split_dim(
-        self,
-        z: torch.Tensor | None = None,
-        world_size: int | None = None,
-    ) -> str | None:
-        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
+    def _spatial_shard_decode_split_dim(self) -> str | None:
+        mode = self.distributed_executor.parallel_mode
+        if mode == "spatial_shard_width":
+            return "width"
+        if mode == "spatial_shard_height":
+            return "height"
+        return None
 
     def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        return spatial_shard_decode_enabled(self, z, model_name="Wan")
+        split_dim = self._spatial_shard_decode_split_dim()
+        if split_dim is None:
+            return False
+        if z.ndim != 5:
+            logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.")
+            return False
+        if not self.is_distributed_enabled():
+            return False
 
-    def _prepare_spatial_shard_decode(self) -> None:
-        prepare_spatial_shard_decode(
-            self,
-            install=wan_spatial_shard.install_wan_spatial_shard_decode,
-        )
+        group = self.distributed_executor.group
+        world_size = dist.get_world_size(group=group)
+        requested_size = int(self.distributed_executor.parallel_size)
+        if requested_size != world_size:
+            logger.warning(
+                "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size "
+                "to match the DIT group size; falling back to tiled decode. "
+                "requested=%s dit_group=%s split_dim=%s",
+                requested_size,
+                world_size,
+                split_dim,
+            )
+            return False
+        return True
 
     def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
         if not self.is_distributed_enabled():
             return super().tiled_decode(z, return_dict=return_dict)
 
         if self._spatial_shard_decode_enabled(z):
-            world_size = dist.get_world_size(group=self.distributed_executor.group)
-            split_dim = self._spatial_shard_decode_split_dim(z, world_size)
-            assert split_dim is not None
+            split_dim = self._spatial_shard_decode_split_dim()
             logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim)
             return wan_spatial_shard.spatial_shard_decode(
                 self,

vllm_omni/diffusion/distributed/autoencoders/qwenimage_spatial_shard.py

@@ -1,86 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""QwenImage adapter for the shared spatial-shard VAE decoder kernels."""
-
-from __future__ import annotations
-
-import torch
-import torch.distributed as dist
-import torch.nn as nn
-from diffusers.models.autoencoders.autoencoder_kl_qwenimage import QwenImageCausalConv3d
-from diffusers.models.autoencoders.vae import DecoderOutput
-
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE
-from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import (
-    SpatialShardCausalConv3dMixin,
-    install_spatial_shard_decode,
-    spatial_shard_decode_impl,
-)
-
-
-class QwenImageDistCausalConv3d(SpatialShardCausalConv3dMixin, QwenImageCausalConv3d):
-    """Context-aware QwenImage causal conv that remains cache-discoverable."""
-
-    def __init__(self, source: nn.Conv3d, group: dist.ProcessGroup) -> None:
-        # Initialize Conv3d directly so the replacement preserves source device,
-        # dtype, groups, and bias while remaining an isinstance() match for
-        # QwenImageCausalConv3d. Diffusers clear_cache() relies on that match.
-        nn.Conv3d.__init__(
-            self,
-            source.in_channels,
-            source.out_channels,
-            source.kernel_size,
-            stride=source.stride,
-            padding=0,
-            dilation=source.dilation,
-            groups=source.groups,
-            bias=source.bias is not None,
-            padding_mode=source.padding_mode,
-            device=source.weight.device,
-            dtype=source.weight.dtype,
-        )
-        self._init_spatial_shard_causal_conv(source, group)
-
-    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        return self._spatial_shard_causal_conv_forward(x, cache_x)
-
-
-def install_qwenimage_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    group: dist.ProcessGroup,
-    split_dim: str = "height",
-) -> None:
-    install_spatial_shard_decode(
-        vae,
-        group,
-        split_dim,
-        causal_conv_class_name="QwenImageCausalConv3d",
-        causal_conv_factory=QwenImageDistCausalConv3d,
-        attention_block_class_name="QwenImageAttentionBlock",
-        installed_attr="_vllm_omni_qwenimage_spatial_shard_installed",
-        model_name="QwenImage",
-    )
-
-
-def spatial_shard_decode(
-    vae: SpatialShardVAE,
-    z: torch.Tensor,
-    *,
-    group: dist.ProcessGroup,
-    return_dict: bool = True,
-    split_dim: str = "height",
-) -> DecoderOutput | tuple[torch.Tensor]:
-    """Decode QwenImage latents and preserve its all-rank output contract."""
-    return spatial_shard_decode_impl(
-        vae,
-        z,
-        group=group,
-        install=install_qwenimage_spatial_shard_decode,
-        model_name="QwenImage",
-        pass_first_chunk=False,
-        broadcast_result=True,
-        unpatchify_patch_size=None,
-        return_dict=return_dict,
-        split_dim=split_dim,
-    )

vllm_omni/diffusion/distributed/autoencoders/spatial_shard.py

@@ -1,142 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from typing import Protocol
-
-import torch
-import torch.distributed as dist
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class SpatialShardExecutor(Protocol):
-    group: dist.ProcessGroup
-    parallel_mode: str
-    parallel_size: int
-
-    def _sync_final_result(
-        self,
-        rank0_result: torch.Tensor,
-        output_ndim: int,
-        output_device: torch.device,
-        output_dtype: torch.dtype,
-    ) -> torch.Tensor: ...
-
-
-class SpatialShardVAE(Protocol):
-    distributed_executor: SpatialShardExecutor
-    decoder: torch.nn.Module
-    post_quant_conv: torch.nn.Module
-    dtype: torch.dtype
-    _conv_idx: list[int]
-    _feat_map: list[object]
-
-    def is_distributed_enabled(self) -> bool: ...
-
-    def clear_cache(self) -> None: ...
-
-
-def select_auto_spatial_shard_split_dim(
-    z_shape: tuple[int, ...],
-    world_size: int,
-) -> str | None:
-    """Select the longer spatial axis that can be split across every rank."""
-    if len(z_shape) != 5 or world_size <= 1:
-        return None
-
-    _, _, _, height, width = z_shape
-    candidates = [
-        (extent, split_dim) for extent, split_dim in ((height, "height"), (width, "width")) if extent >= world_size
-    ]
-    if not candidates:
-        return None
-    # Prefer width for square latents. It was modestly faster in the original
-    # four-GPU Wan benchmark and keeps the choice deterministic across models.
-    return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1]
-
-
-def resolve_spatial_shard_split_dim(
-    mode: str,
-    z: torch.Tensor | None = None,
-    world_size: int | None = None,
-) -> str | None:
-    if mode == "spatial_shard_width":
-        return "width"
-    if mode == "spatial_shard_height":
-        return "height"
-    if mode == "auto" and z is not None and world_size is not None:
-        return select_auto_spatial_shard_split_dim(tuple(z.shape), world_size)
-    return None
-
-
-def spatial_shard_decode_enabled(vae: SpatialShardVAE, z: torch.Tensor, *, model_name: str) -> bool:
-    """Apply the shared request- and topology-level spatial decode policy."""
-    executor = vae.distributed_executor
-    if executor.parallel_mode == "tile":
-        return False
-    if z.ndim != 5:
-        logger.warning(
-            "%s VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.", model_name
-        )
-        return False
-    if not vae.is_distributed_enabled():
-        return False
-
-    group = executor.group
-    world_size = dist.get_world_size(group=group)
-    requested_size = int(executor.parallel_size)
-    requested_split_dim = resolve_spatial_shard_split_dim(executor.parallel_mode)
-    if requested_size != world_size:
-        if executor.parallel_mode == "auto":
-            logger.debug(
-                "%s VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
-                model_name,
-                requested_size,
-                world_size,
-            )
-        elif requested_split_dim is not None:
-            logger.warning(
-                "%s VAE spatial sharded decode currently requires vae_patch_parallel_size "
-                "to match the DIT group size; falling back to tiled decode. "
-                "requested=%s dit_group=%s split_dim=%s",
-                model_name,
-                requested_size,
-                world_size,
-                requested_split_dim,
-            )
-        return False
-    return resolve_spatial_shard_split_dim(executor.parallel_mode, z, world_size) is not None
-
-
-def prepare_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    *,
-    install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
-) -> None:
-    """Install one VAE's dynamic wrappers while allocations are still tagged."""
-    executor = vae.distributed_executor
-    mode = executor.parallel_mode
-    if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}:
-        return
-    requested_size = int(executor.parallel_size)
-    world_size = dist.get_world_size(group=executor.group)
-    if requested_size <= 1 or requested_size != world_size:
-        return
-    split_dim = "width" if mode == "spatial_shard_width" else "height"
-    install(vae, executor.group, split_dim)
-
-
-def prepare_pipeline_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
-    """Prepare every VAE that advertises spatial-shard decode capability."""
-    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
-
-    if not dist.is_initialized():
-        return
-    for vae in ModuleDiscovery.discover(pipeline).vaes:
-        prepare = getattr(vae, "_prepare_spatial_shard_decode", None)
-        if callable(prepare):
-            prepare()

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -7,21 +7,19 @@
 #   (python/sglang/multimodal_gen/runtime/layers/parallel_conv.py)
 # which is in turn adapted from FastVideo (https://github.com/hao-ai-lab/FastVideo).
 # This version generalizes the height-only sharding to shard along height or
-# width and supports the causal-conv ``feat_cache`` handling used by Wan and
-# QwenImage.
-"""Shared spatial-shard kernels plus the Wan VAE adapter.
+# width and adds the Wan causal-conv ``feat_cache`` handling.
+"""Spatially-sharded Wan VAE decode.
 
-The halo, padding, stride-trim, attention-gather, and decode-lifecycle helpers
-are also reused by QwenImage's thin adapter. The backend shards decoder feature
-maps along height or width and exchanges boundary rows/columns before spatial
-convolutions. It is decode-only and keeps checkpoint loading unchanged by
-patching the already-loaded decoder.
+The existing distributed Wan VAE path shards *tiles*.  This module adds an
+opt-in decode backend that shards decoder feature maps along height or width and
+exchanges boundary rows/columns before spatial convolutions.  It is
+intentionally decode-only and keeps checkpoint loading unchanged by patching the
+already-loaded decoder.
 """
 
 from __future__ import annotations
 
 import math
-from collections.abc import Callable
 from contextlib import nullcontext
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -36,8 +34,6 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE
-
 logger = init_logger(__name__)
 
 
@@ -56,17 +52,12 @@ _SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar(
 )
 
 
-def _active_spatial_shard_split_dim() -> str | None:
-    context = _SPATIAL_SHARD_CONTEXT.get()
-    return context.split_dim if context is not None else None
-
-
 def _spatial_dim(split_dim: str) -> int:
     if split_dim == "height":
         return -2
     if split_dim == "width":
         return -1
-    raise ValueError(f"Unsupported VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
+    raise ValueError(f"Unsupported Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
 
 
 def _narrow_along_dim(x: torch.Tensor, dim: int, start: int, length: int) -> torch.Tensor:
@@ -164,9 +155,9 @@ def split_for_parallel_decode(
     rank = 0 if rank is None else int(rank)
     world_size = 1 if world_size is None else int(world_size)
     if world_size < 1:
-        raise ValueError(f"VAE world_size must be >= 1, got {world_size}.")
+        raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.")
     if not 0 <= rank < world_size:
-        raise ValueError(f"VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
+        raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
 
     dim = _spatial_dim(split_dim)
     expected_extent = x.shape[dim] * (2**upsample_count)
@@ -332,27 +323,21 @@ class WanDistZeroPad2d(nn.Module):
         padding: tuple[int, int, int, int],
         group: dist.ProcessGroup,
         *,
-        defer_split_after_padding: bool = False,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
     ) -> None:
         super().__init__()
         self.padding = padding
-        self.defer_split_after_padding = defer_split_after_padding
+        self.split_dim = split_dim
+        default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1])
+        self.split_padding = split_padding or default_split_padding
         self.group = group
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        if split_dim is None:
-            return F.pad(x, self.padding)
-
         rank, world_size = _rank_world(self.group)
         left, right, top, bottom = self.padding
-        if self.defer_split_after_padding:
-            if split_dim == "height":
-                bottom = 0
-            else:
-                right = 0
         if world_size > 1:
-            if split_dim == "height":
+            if self.split_dim == "height":
                 top = top if rank == 0 else 0
                 bottom = bottom if rank == world_size - 1 else 0
             else:
@@ -366,7 +351,8 @@ class WanDistConv2d(nn.Conv2d):
         self,
         source: nn.Conv2d,
         group: dist.ProcessGroup,
-        deferred_padding: tuple[int, int, int, int] | None = None,
+        split_dim: str = "height",
+        split_padding: tuple[int, int] | None = None,
     ):
         super().__init__(
             source.in_channels,
@@ -383,253 +369,176 @@ class WanDistConv2d(nn.Conv2d):
         )
         self.load_state_dict(source.state_dict())
         self.group = group
-        self._direct_padding = source.padding
-        self._direct_padding_mode = source.padding_mode
-        self._direct_reversed_padding = source._reversed_padding_repeated_twice
-        self._deferred_padding = deferred_padding
-        source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding)
-        if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding):
-            raise ValueError(f"Spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
-        self._spatial_pad_h = int(source_padding[-2])
-        self._spatial_pad_w = int(source_padding[-1])
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
+        self.split_dim = split_dim
+        self.split_tensor_dim = _spatial_dim(split_dim)
+        kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1]
+        self.halo_size = (kernel_extent - 1) // 2
+        pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding
+        pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding
+        if split_dim == "height":
+            if split_padding is None:
+                split_padding = (pad_h, pad_h)
+            self._non_split_padding = (pad_w, pad_w, 0, 0)
+            self.kernel_extent = self.kernel_size[-2]
+            self.stride_extent = self.stride[-2]
+        else:
+            if split_padding is None:
+                split_padding = (pad_w, pad_w)
+            self._non_split_padding = (0, 0, pad_h, pad_h)
+            self.kernel_extent = self.kernel_size[-1]
+            self.stride_extent = self.stride[-1]
+        self.split_pad_left, self.split_pad_right = split_padding
+        self._halo_recv_top_buf: torch.Tensor | None = None
+        self._halo_recv_bottom_buf: torch.Tensor | None = None
+        self._trim_cache: dict[int, tuple[int, int, int]] = {}
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        if split_dim is None:
-            if self._direct_padding_mode != "zeros":
-                x = F.pad(x, self._direct_reversed_padding, mode=self._direct_padding_mode)
-                padding: str | tuple[int, int] = (0, 0)
-            else:
-                padding = self._direct_padding
-            return F.conv2d(x, self.weight, self.bias, self.stride, padding, self.dilation, self.groups)
-
-        split_tensor_dim = _spatial_dim(split_dim)
-        if split_dim == "height":
-            kernel_extent = self.kernel_size[-2]
-            stride_extent = self.stride[-2]
-            non_split_padding = (self._spatial_pad_w, self._spatial_pad_w, 0, 0)
-            default_split_padding = (self._spatial_pad_h, self._spatial_pad_h)
-            deferred_split_padding = (
-                (self._deferred_padding[2], self._deferred_padding[3]) if self._deferred_padding is not None else None
-            )
-        else:
-            kernel_extent = self.kernel_size[-1]
-            stride_extent = self.stride[-1]
-            non_split_padding = (0, 0, self._spatial_pad_h, self._spatial_pad_h)
-            default_split_padding = (self._spatial_pad_w, self._spatial_pad_w)
-            deferred_split_padding = (
-                (self._deferred_padding[0], self._deferred_padding[1]) if self._deferred_padding is not None else None
-            )
-        split_pad_left, split_pad_right = deferred_split_padding or default_split_padding
-        halo_size = (kernel_extent - 1) // 2
-
-        x = F.pad(x, non_split_padding)
+        x = F.pad(x, self._non_split_padding)
         x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
             x,
             group=self.group,
-            halo_size=halo_size,
-            split_dim=split_dim,
+            halo_size=self.halo_size,
+            split_dim=self.split_dim,
             recv_top_buf=self._halo_recv_top_buf,
             recv_bottom_buf=self._halo_recv_bottom_buf,
         )
-        shift, start, upper_bound = self._get_trim_params(
-            x.shape[split_tensor_dim],
-            split_dim=split_dim,
-            halo_size=halo_size,
-            split_pad_left=split_pad_left,
-            split_pad_right=split_pad_right,
-            kernel_extent=kernel_extent,
-            stride_extent=stride_extent,
-        )
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                split_tensor_dim,
+                self.split_tensor_dim,
                 shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
             )
         out = super().forward(x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)
 
-    def _get_trim_params(
-        self,
-        local_extent: int,
-        *,
-        split_dim: str,
-        halo_size: int,
-        split_pad_left: int,
-        split_pad_right: int,
-        kernel_extent: int,
-        stride_extent: int,
-    ) -> tuple[int, int, int]:
-        cache_key = (split_dim, local_extent)
-        trim_params = self._trim_cache.get(cache_key)
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
         if trim_params is None:
             rank, world_size = _rank_world(self.group)
             trim_params = _compute_conv_trim_params(
                 local_extent=local_extent,
                 rank=rank,
                 world_size=world_size,
-                halo_size=halo_size,
-                pad_before=split_pad_left,
-                pad_after=split_pad_right,
-                kernel_extent=kernel_extent,
-                stride_extent=stride_extent,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
             )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = trim_params
         return trim_params
 
 
-# Model-neutral aliases let other causal VAEs reuse the implementation while
-# preserving the original Wan class identities for stacked PRs and downstreams.
-SpatialShardZeroPad2d = WanDistZeroPad2d
-SpatialShardConv2d = WanDistConv2d
-
-
-class SpatialShardCausalConv3dMixin:
-    """Shared causal-convolution behavior for supported Diffusers VAE classes."""
-
-    def _init_spatial_shard_causal_conv(
+class WanDistCausalConv3d(nn.Conv3d):
+    def __init__(
         self,
         source: nn.Conv3d,
         group: dist.ProcessGroup,
-    ) -> None:
+        split_dim: str = "height",
+    ):
+        super().__init__(
+            source.in_channels,
+            source.out_channels,
+            source.kernel_size,
+            stride=source.stride,
+            padding=0,
+            dilation=source.dilation,
+            groups=source.groups,
+            bias=source.bias is not None,
+            padding_mode=source.padding_mode,
+            device=source.weight.device,
+            dtype=source.weight.dtype,
+        )
         self.load_state_dict(source.state_dict())
         self.group = group
+        self.split_dim = split_dim
+        self.split_tensor_dim = _spatial_dim(split_dim)
         source_padding = getattr(source, "_padding", None)
         if source_padding is None:
             p_t, p_h, p_w = source.padding
             source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0)
         self._source_padding = tuple(source_padding)
-        self._padding = tuple(source_padding)
-        self.register_buffer("_halo_recv_top_buf", None, persistent=False)
-        self.register_buffer("_halo_recv_bottom_buf", None, persistent=False)
-        self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {}
+        if split_dim == "height":
+            self.split_pad_left = int(self._source_padding[2])
+            self.split_pad_right = int(self._source_padding[3])
+            self.kernel_extent = self.kernel_size[-2]
+            self.stride_extent = self.stride[-2]
+            padding = (
+                self._source_padding[0],
+                self._source_padding[1],
+                0,
+                0,
+                self._source_padding[4],
+                self._source_padding[5],
+            )
+        else:
+            self.split_pad_left = int(self._source_padding[0])
+            self.split_pad_right = int(self._source_padding[1])
+            self.kernel_extent = self.kernel_size[-1]
+            self.stride_extent = self.stride[-1]
+            padding = (
+                0,
+                0,
+                self._source_padding[2],
+                self._source_padding[3],
+                self._source_padding[4],
+                self._source_padding[5],
+            )
+        self.halo_size = (self.kernel_extent - 1) // 2
+        self._padding = padding if self.halo_size > 0 else self._source_padding
+        self._halo_recv_top_buf: torch.Tensor | None = None
+        self._halo_recv_bottom_buf: torch.Tensor | None = None
+        self._trim_cache: dict[int, tuple[int, int, int]] = {}
 
-    def _spatial_shard_causal_conv_forward(
-        self,
-        x: torch.Tensor,
-        cache_x: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        padding = list(self._source_padding)
+    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
+        padding = list(self._padding)
         if cache_x is not None and padding[4] > 0:
             cache_x = cache_x.to(x.device)
             x = torch.cat([cache_x, x], dim=2)
             padding[4] -= cache_x.shape[2]
 
-        if split_dim is None:
-            return nn.Conv3d.forward(self, F.pad(x, padding))
-
-        split_tensor_dim = _spatial_dim(split_dim)
-        if split_dim == "height":
-            split_pad_left = int(padding[2])
-            split_pad_right = int(padding[3])
-            kernel_extent = self.kernel_size[-2]
-            stride_extent = self.stride[-2]
-            padding[2] = 0
-            padding[3] = 0
-        else:
-            split_pad_left = int(padding[0])
-            split_pad_right = int(padding[1])
-            kernel_extent = self.kernel_size[-1]
-            stride_extent = self.stride[-1]
-            padding[0] = 0
-            padding[1] = 0
-        halo_size = (kernel_extent - 1) // 2
-
         x = F.pad(x, padding)
         x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
             x,
             group=self.group,
-            halo_size=halo_size,
-            split_dim=split_dim,
+            halo_size=self.halo_size,
+            split_dim=self.split_dim,
             recv_top_buf=self._halo_recv_top_buf,
             recv_bottom_buf=self._halo_recv_bottom_buf,
         )
-        shift, start, upper_bound = self._get_trim_params(
-            x.shape[split_tensor_dim],
-            split_dim=split_dim,
-            halo_size=halo_size,
-            split_pad_left=split_pad_left,
-            split_pad_right=split_pad_right,
-            kernel_extent=kernel_extent,
-            stride_extent=stride_extent,
-        )
+        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
-                split_tensor_dim,
+                self.split_tensor_dim,
                 shift,
-                x_padded.shape[split_tensor_dim] - shift,
+                x_padded.shape[self.split_tensor_dim] - shift,
             )
-        out = nn.Conv3d.forward(self, x_padded)
-        out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim)
-        return _zero_invalid_extent(out, split_dim=split_dim)
+        out = super().forward(x_padded)
+        out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
+        return _zero_invalid_extent(out, split_dim=self.split_dim)
 
-    def _get_trim_params(
-        self,
-        local_extent: int,
-        *,
-        split_dim: str,
-        halo_size: int,
-        split_pad_left: int,
-        split_pad_right: int,
-        kernel_extent: int,
-        stride_extent: int,
-    ) -> tuple[int, int, int]:
-        cache_key = (split_dim, local_extent)
-        trim_params = self._trim_cache.get(cache_key)
+    def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
+        trim_params = self._trim_cache.get(local_extent)
         if trim_params is None:
             rank, world_size = _rank_world(self.group)
             trim_params = _compute_conv_trim_params(
                 local_extent=local_extent,
                 rank=rank,
                 world_size=world_size,
-                halo_size=halo_size,
-                pad_before=split_pad_left,
-                pad_after=split_pad_right,
-                kernel_extent=kernel_extent,
-                stride_extent=stride_extent,
+                halo_size=self.halo_size,
+                pad_before=self.split_pad_left,
+                pad_after=self.split_pad_right,
+                kernel_extent=self.kernel_extent,
+                stride_extent=self.stride_extent,
             )
-            self._trim_cache[cache_key] = trim_params
+            self._trim_cache[local_extent] = trim_params
         return trim_params
 
 
-class WanDistCausalConv3d(SpatialShardCausalConv3dMixin, nn.Conv3d):
-    # Cross-PR compatibility contract for the lossless Wan data-movement
-    # installer: it may fuse the no-context direct path, but must leave this
-    # wrapper in control whenever a spatial-shard context is active.
-    _vllm_omni_dynamic_spatial_shard_conv = True
-
-    def __init__(
-        self,
-        source: nn.Conv3d,
-        group: dist.ProcessGroup,
-    ):
-        super().__init__(
-            source.in_channels,
-            source.out_channels,
-            source.kernel_size,
-            stride=source.stride,
-            padding=0,
-            dilation=source.dilation,
-            groups=source.groups,
-            bias=source.bias is not None,
-            padding_mode=source.padding_mode,
-            device=source.weight.device,
-            dtype=source.weight.dtype,
-        )
-        self._init_spatial_shard_causal_conv(source, group)
-
-    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        return self._spatial_shard_causal_conv_forward(x, cache_x)
-
-
 def _compute_conv_trim_params(
     *,
     local_extent: int,
@@ -673,15 +582,12 @@ def _trim_local_conv_output(
     return out
 
 
-def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> None:
+def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None:
     if getattr(module, "_vllm_omni_spatial_shard_attention", False):
         return
     orig_forward = module.forward
 
     def _forward(self: nn.Module, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
-        split_dim = _active_spatial_shard_split_dim()
-        if split_dim is None:
-            return orig_forward(x, *args, **kwargs)
         _, world_size = _rank_world(group)
         if world_size <= 1:
             return orig_forward(x, *args, **kwargs)
@@ -703,43 +609,60 @@ def _replace_child(
     name: str,
     child: nn.Module,
     group: dist.ProcessGroup,
-    *,
-    causal_conv_class_name: str,
-    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
+    split_dim: str,
 ) -> None:
-    if child.__class__.__name__ == causal_conv_class_name:
-        setattr(parent, name, causal_conv_factory(child, group))
-        return
-    if isinstance(child, nn.ZeroPad2d):
-        padding = tuple(int(p) for p in child.padding)
+    if child.__class__.__name__ == "WanCausalConv3d":
         setattr(
             parent,
             name,
-            SpatialShardZeroPad2d(
-                padding,
+            WanDistCausalConv3d(
+                child,
                 group,
-                defer_split_after_padding=parent.__class__.__name__ == "Sequential",
+                split_dim=split_dim,
+            ),
+        )
+        return
+    if isinstance(child, nn.ZeroPad2d):
+        padding = tuple(int(p) for p in child.padding)
+        module_padding = padding
+        if parent.__class__.__name__ == "Sequential":
+            # Let the following WanDistConv2d account for global after-edge
+            # padding; this module handles the non-split dimension and before edge.
+            if split_dim == "height":
+                module_padding = (padding[0], padding[1], padding[2], 0)
+            else:
+                module_padding = (padding[0], 0, padding[2], padding[3])
+        setattr(
+            parent,
+            name,
+            WanDistZeroPad2d(
+                module_padding,
+                group,
+                split_dim=split_dim,
+                split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]),
             ),
         )
         return
     if isinstance(child, nn.Conv2d):
-        deferred_padding = None
+        split_padding = None
         if name == "1" and parent.__class__.__name__ == "Sequential":
-            # Wan/QwenImage downsampling uses ZeroPad2d((0, 1, 0, 1)) before a
+            # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a
             # stride-2 conv with padding=0.  Only the last rank should see the
             # bottom/right global padding, which is approximated by split_padding.
             prev = getattr(parent, "0", None)
-            if isinstance(prev, SpatialShardZeroPad2d):
-                deferred_padding = prev.padding
+            if isinstance(prev, WanDistZeroPad2d):
+                split_padding = prev.split_padding
             elif isinstance(prev, nn.ZeroPad2d):
-                deferred_padding = tuple(int(value) for value in prev.padding)
+                pad = prev.padding
+                split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1]))
         setattr(
             parent,
             name,
-            SpatialShardConv2d(
+            WanDistConv2d(
                 child,
                 group,
-                deferred_padding=deferred_padding,
+                split_dim=split_dim,
+                split_padding=split_padding,
             ),
         )
 
@@ -747,45 +670,21 @@ def _replace_child(
 def _patch_decoder_modules(
     module: nn.Module,
     group: dist.ProcessGroup,
-    *,
-    causal_conv_class_name: str,
-    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
-    attention_block_class_name: str,
+    split_dim: str,
     inside_attention: bool = False,
 ) -> None:
-    if module.__class__.__name__ == attention_block_class_name:
-        _patch_attention_block(module, group)
+    if module.__class__.__name__ == "WanAttentionBlock":
+        _patch_attention_block(module, group, split_dim)
         inside_attention = True
 
     for name, child in list(module.named_children()):
-        if child.__class__.__name__ == causal_conv_class_name:
-            _replace_child(
-                module,
-                name,
-                child,
-                group,
-                causal_conv_class_name=causal_conv_class_name,
-                causal_conv_factory=causal_conv_factory,
-            )
+        if child.__class__.__name__ == "WanCausalConv3d":
+            _replace_child(module, name, child, group, split_dim)
             continue
         if isinstance(child, nn.Conv2d) and not inside_attention:
-            _replace_child(
-                module,
-                name,
-                child,
-                group,
-                causal_conv_class_name=causal_conv_class_name,
-                causal_conv_factory=causal_conv_factory,
-            )
+            _replace_child(module, name, child, group, split_dim)
             continue
-        _patch_decoder_modules(
-            child,
-            group,
-            causal_conv_class_name=causal_conv_class_name,
-            causal_conv_factory=causal_conv_factory,
-            attention_block_class_name=attention_block_class_name,
-            inside_attention=inside_attention,
-        )
+        _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention)
 
 
 def _decoder_upsample_count(decoder: nn.Module) -> int:
@@ -796,58 +695,47 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
     return count
 
 
-def clear_spatial_shard_runtime_buffers(vae: SpatialShardVAE) -> None:
-    decoder = getattr(vae, "decoder", None)
-    if decoder is None:
-        return
-    for module in decoder.modules():
-        if "_halo_recv_top_buf" in module._buffers:
-            module._halo_recv_top_buf = None
-        if "_halo_recv_bottom_buf" in module._buffers:
-            module._halo_recv_bottom_buf = None
-
-
-def install_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    group: dist.ProcessGroup,
-    split_dim: str,
-    *,
-    causal_conv_class_name: str,
-    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
-    attention_block_class_name: str,
-    installed_attr: str,
-    model_name: str,
-) -> None:
+def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None:
     """Patch ``vae.decoder`` once for spatially-sharded decode.
 
     This mutates the already-loaded decoder in place by swapping its spatial
-    convolutions/padding for context-aware variants and attaching a dedicated
-    spatial forward. Outside that forward the replacements execute their exact
-    local PyTorch behavior, so a later auto-selected tile request remains safe.
-    The same installed decoder can alternate height and width sharding.
+    convolutions/padding for halo-exchanging variants and wrapping
+    ``decoder.forward``. The patch is permanent for the lifetime of the VAE
+    instance and is applied only once (subsequent calls are no-ops). A given
+    instance is bound to a single ``split_dim``; switching between
+    ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here
+    otherwise.
 
-    Group-relative rank 0 assembles the final decoded frame. The model adapter
-    then preserves its prior contract: Wan keeps the result on rank 0, while
-    QwenImage broadcasts it to every rank.
+    Only group-relative rank 0 assembles the final decoded frame, mirroring the
+    distributed tiled-decode ``broadcast_result=False`` contract; the other ranks
+    take part in the collectives but return an empty placeholder.
     """
     _spatial_dim(split_dim)
-    if getattr(vae, installed_attr, False):
+    if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False):
+        installed_split_dim = getattr(vae, "_vllm_omni_wan_spatial_shard_split_dim", "height")
+        if installed_split_dim != split_dim:
+            raise ValueError(
+                "Wan spatial-shard VAE decoder was already patched for "
+                f"{installed_split_dim!r} split; create a fresh VAE instance to use {split_dim!r} split."
+            )
         return
     decoder = getattr(vae, "decoder", None)
     if decoder is None:
-        raise ValueError(f"{model_name} spatial-shard VAE decode requires a decoder module.")
+        raise ValueError("Wan spatial-shard VAE decode requires a decoder module.")
 
-    _patch_decoder_modules(
-        decoder,
-        group,
-        causal_conv_class_name=causal_conv_class_name,
-        causal_conv_factory=causal_conv_factory,
-        attention_block_class_name=attention_block_class_name,
-    )
+    _patch_decoder_modules(decoder, group, split_dim)
     upsample_count = _decoder_upsample_count(decoder)
     orig_forward = decoder.forward
 
-    def _forward(self: nn.Module, x: torch.Tensor, *args: object, split_dim: str, **kwargs: object) -> torch.Tensor:
+    def _forward(
+        self: nn.Module,
+        x: torch.Tensor,
+        feat_cache: list[torch.Tensor] | None = None,
+        feat_idx: list[int] | None = None,
+        first_chunk: bool = False,
+    ) -> torch.Tensor:
+        if feat_idx is None:
+            feat_idx = [0]
         tensor_dim = _spatial_dim(split_dim)
         input_extent = x.shape[tensor_dim]
         x, expected_extent = split_for_parallel_decode(
@@ -867,50 +755,29 @@ def install_spatial_shard_decode(
             )
         )
         try:
-            out = orig_forward(x, *args, **kwargs)
+            out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
         finally:
             _SPATIAL_SHARD_CONTEXT.reset(token)
         return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0)
 
-    decoder._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder)
-    setattr(vae, installed_attr, True)
-    logger.info("Installed %s VAE dynamic-axis spatial-shard decode.", model_name)
+    decoder.forward = MethodType(_forward, decoder)
+    vae._vllm_omni_wan_spatial_shard_installed = True
+    vae._vllm_omni_wan_spatial_shard_split_dim = split_dim
+    logger.info("Installed Wan VAE %s-sharded decode.", split_dim)
 
 
-def install_wan_spatial_shard_decode(
-    vae: SpatialShardVAE,
-    group: dist.ProcessGroup,
-    split_dim: str = "height",
-) -> None:
-    install_spatial_shard_decode(
-        vae,
-        group,
-        split_dim,
-        causal_conv_class_name="WanCausalConv3d",
-        causal_conv_factory=WanDistCausalConv3d,
-        attention_block_class_name="WanAttentionBlock",
-        installed_attr="_vllm_omni_wan_spatial_shard_installed",
-        model_name="Wan",
-    )
-
-
-def spatial_shard_decode_impl(
-    vae: SpatialShardVAE,
+def spatial_shard_decode(
+    vae: Any,
     z: torch.Tensor,
     *,
     group: dist.ProcessGroup,
-    install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
-    model_name: str,
-    pass_first_chunk: bool,
-    broadcast_result: bool,
-    unpatchify_patch_size: int | None,
     return_dict: bool = True,
     split_dim: str = "height",
 ) -> DecoderOutput | tuple[torch.Tensor]:
-    install(vae, group, split_dim)
+    install_wan_spatial_shard_decode(vae, group, split_dim=split_dim)
 
     if z.shape[2] == 0:
-        raise ValueError(f"{model_name} spatial-shard VAE decode expects at least one latent frame.")
+        raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.")
 
     # Non-rank-0 ranks must still run the decoder every chunk to stay in lockstep with
     # the halo/all-gather collectives; they just skip keeping/assembling the output.
@@ -919,72 +786,31 @@ def spatial_shard_decode_impl(
 
     vae.clear_cache()
     try:
-        context_factory = getattr(vae, "_execution_context", None)
-        context = context_factory() if callable(context_factory) else nullcontext()
+        context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext()
         with context:
             x = vae.post_quant_conv(z)
             decoded_chunks = []
-            spatial_forward = getattr(vae.decoder, "_vllm_omni_spatial_shard_forward")
             for i in range(z.shape[2]):
                 vae._conv_idx = [0]
-                if pass_first_chunk:
-                    chunk = spatial_forward(
-                        x[:, :, i : i + 1, :, :],
-                        feat_cache=vae._feat_map,
-                        feat_idx=vae._conv_idx,
-                        first_chunk=i == 0,
-                        split_dim=split_dim,
-                    )
-                else:
-                    chunk = spatial_forward(
-                        x[:, :, i : i + 1, :, :],
-                        feat_cache=vae._feat_map,
-                        feat_idx=vae._conv_idx,
-                        split_dim=split_dim,
-                    )
+                chunk = vae.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=vae._feat_map,
+                    feat_idx=vae._conv_idx,
+                    first_chunk=(i == 0),
+                )
                 if produce_output:
                     decoded_chunks.append(chunk)
 
             if produce_output:
                 out = torch.cat(decoded_chunks, dim=2)
-                if unpatchify_patch_size is not None:
-                    out = unpatchify(out, patch_size=unpatchify_patch_size)
+                if vae.config.patch_size is not None:
+                    out = unpatchify(out, patch_size=vae.config.patch_size)
                 out = torch.clamp(out, min=-1.0, max=1.0)
             else:
                 out = z.new_zeros(0)
     finally:
         vae.clear_cache()
-        # Halo buffers are request-shape scratch space allocated outside the
-        # tagged weight pool. Drop live references so sleep/offload never has
-        # to preserve or discard them as model state.
-        clear_spatial_shard_runtime_buffers(vae)
-
-    if broadcast_result:
-        out = vae.distributed_executor._sync_final_result(out, z.ndim, z.device, vae.dtype)
 
     if not return_dict:
         return (out,)
     return DecoderOutput(sample=out)
-
-
-def spatial_shard_decode(
-    vae: SpatialShardVAE,
-    z: torch.Tensor,
-    *,
-    group: dist.ProcessGroup,
-    return_dict: bool = True,
-    split_dim: str = "height",
-) -> DecoderOutput | tuple[torch.Tensor]:
-    """Decode a Wan latent with spatial activation sharding."""
-    return spatial_shard_decode_impl(
-        vae,
-        z,
-        group=group,
-        install=install_wan_spatial_shard_decode,
-        model_name="Wan",
-        pass_first_chunk=True,
-        broadcast_result=False,
-        unpatchify_patch_size=vae.config.patch_size,
-        return_dict=return_dict,
-        split_dim=split_dim,
-    )

vllm_omni/diffusion/distributed/autoencoders/wan_vae_data_movement.py

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

vllm_omni/diffusion/kernels/__init__.py

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

vllm_omni/diffusion/kernels/wan_vae_data_movement.py

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

vllm_omni/diffusion/models/helios/pipeline_helios.py

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

vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py

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

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -155,11 +155,6 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         parallel_size: int,
         mode: str = "tile",
     ) -> None:
-        # ``auto`` is the cross-model default. MiniMax H3 has no spatial-shard
-        # capability, so preserve its native tile path for that default while
-        # continuing to reject explicitly requested spatial modes.
-        if mode == "auto":
-            mode = "tile"
         if mode != "tile":
             raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}")
         group = get_dit_group()

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -30,8 +30,8 @@ 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.distributed.autoencoders.autoencoder_kl_wan import (
+    prepare_pipeline_wan_vae_data_movement,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -246,7 +246,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     custom_pipeline_name=custom_pipeline_name,
                     device=self.device,
                 )
-                prepare_pipeline_spatial_shard_decode(self.pipeline)
+                prepare_pipeline_wan_vae_data_movement(self.pipeline)
         time_after_load = time.perf_counter()
 
         logger.info(

vllm_omni/engine/arg_utils.py

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

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 "auto"
+            vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile"
             text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1
             enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
             use_hsdp = normalized_kwargs.get("use_hsdp", False)

vllm_omni/entrypoints/cli/serve.py

@@ -750,12 +750,10 @@ class OmniServeCommand(CLISubcommand):
         omni_config_group.add_argument(
             "--vae-parallel-mode",
             type=str,
-            default="auto",
-            choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"],
+            default="tile",
+            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
             help="VAE parallel decode strategy for diffusion models. "
-            "'auto' (default) selects the longer spatial-shard axis for validated Wan and distributed QwenImage "
-            "tiled decodes "
-            "on a full DiT group, then falls back to patch/tile decode; 'tile' always uses patch/tile decode; "
+            "'tile' (default) uses patch/tile parallel decode; "
             "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits "
             "decoder feature maps along height/width and exchanges halo regions. The "
             "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. "