feat(diffusion): add QwenImage spatial-shard decode

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

docs/design/feature/vae_parallel.md

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

docs/user_guide/diffusion/parallelism/vae_parallelism.md

@@ -15,7 +15,7 @@
 
 ## Overview
 
-VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width.
+VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan/QwenImage spatial-shard decode, which shards decoder feature maps along height or width.
 
 This is particularly useful for:
 - **High-resolution image generation** where VAE decode can become a memory bottleneck
@@ -113,7 +113,7 @@ In `DiffusionParallelConfig`:
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
-| `vae_parallel_mode` | str | `"auto"` | VAE parallel decode strategy: `"auto"` (shape/model/topology-aware selection), `"tile"`, `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding is currently Wan-only. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
+| `vae_parallel_mode` | str | `"auto"` | VAE parallel decode strategy: `"auto"` (shape/model/topology-aware selection), `"tile"`, `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding supports Wan and the distributed QwenImage backend. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan-and-qwenimage). |
 
 Additional requirements:
 
@@ -126,9 +126,9 @@ Additional requirements:
 
 ---
 
-## Spatially-Sharded Decode (Wan)
+## Spatially-Sharded Decode (Wan and QwenImage)
 
-The default `vae_parallel_mode="auto"` selects spatial-shard decode only after a Wan 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="auto"` selects spatial-shard decode only after a Wan or distributed QwenImage request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. It splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests, partial groups, and other VAE families keep their existing path. Set `vae_parallel_mode="tile"` to force the previous behavior, or `"spatial_shard_height"` / `"spatial_shard_width"` to force an axis.
 
 Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance.
 
@@ -157,11 +157,13 @@ vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \
 
 **Constraints and behavior:**
 
-- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes.
+- 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 backend.
+The automatic policy is adapted from [SGLang #28071](https://github.com/sgl-project/sglang/pull/28071) and narrowed to vLLM-Omni's parity-tested Wan and QwenImage backends.
 
 For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:
 

tests/diffusion/distributed/test_qwenimage_spatial_shard.py

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

tests/diffusion/test_diffusion_model_runner.py

@@ -790,7 +790,7 @@ 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_wan_spatial_shard_decode",
+        "prepare_pipeline_spatial_shard_decode",
         lambda pipeline: preparation_pool_states.append(memory_pool_active),
     )
     monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)

vllm_omni/diffusion/data.py

@@ -209,9 +209,9 @@ class DiffusionParallelConfig:
     vae_parallel_mode: str = "auto"
     """VAE parallel decode strategy.
 
-    - "auto": Select the longer spatial-shard axis for validated Wan tiled
-      decodes on a full DiT-group topology; otherwise use patch/tile parallel
-      decode.
+    - "auto": Select the longer spatial-shard axis for validated Wan and
+      distributed QwenImage tiled decodes on a full DiT-group topology;
+      otherwise use patch/tile parallel decode.
     - "tile": Patch/tile parallel decode. Each rank decodes a subset
       of spatial tiles and the results are stitched on rank 0.
     - "spatial_shard_height": Spatially-sharded decode that splits decoder

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py

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

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -18,54 +18,23 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     GridSpec,
     TileTask,
 )
+from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
+    prepare_pipeline_spatial_shard_decode,
+    prepare_spatial_shard_decode,
+    resolve_spatial_shard_split_dim,
+    spatial_shard_decode_enabled,
+)
+from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import (
+    select_auto_spatial_shard_split_dim as select_auto_spatial_shard_split_dim,
+)
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-def select_auto_spatial_shard_split_dim(
-    z_shape: tuple[int, ...],
-    world_size: int,
-) -> str | None:
-    """Select a profitable Wan spatial-shard axis for one latent shape.
-
-    This selector runs only after Diffusers has already chosen ``tiled_decode``
-    for a spatially large request. Splitting the longer eligible axis reduces
-    halo traffic for rectangular requests and improves per-rank balance.
-    """
-    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 existing
-    # four-GPU Wan benchmark and keeps the choice deterministic.
-    return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1]
-
-
 def prepare_pipeline_wan_spatial_shard_decode(pipeline: torch.nn.Module) -> None:
-    """Install eligible Wan wrappers while weight allocations are still tagged."""
-    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
-
-    if not dist.is_initialized():
-        return
-    for vae in ModuleDiscovery.discover(pipeline).vaes:
-        if not isinstance(vae, DistributedAutoencoderKLWan):
-            continue
-        executor = vae.distributed_executor
-        mode = executor.parallel_mode
-        if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}:
-            continue
-        requested_size = int(executor.parallel_size)
-        world_size = dist.get_world_size(group=executor.group)
-        if requested_size <= 1 or requested_size != world_size:
-            continue
-        split_dim = "width" if mode == "spatial_shard_width" else "height"
-        wan_spatial_shard.install_wan_spatial_shard_decode(vae, executor.group, split_dim=split_dim)
+    """Backward-compatible alias for the model-agnostic pipeline dispatcher."""
+    prepare_pipeline_spatial_shard_decode(pipeline)
 
 
 class OmniAutoencoderKLWan(AutoencoderKLWan):
@@ -320,46 +289,16 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         z: torch.Tensor | None = None,
         world_size: int | None = None,
     ) -> str | None:
-        mode = self.distributed_executor.parallel_mode
-        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
+        return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size)
 
     def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool:
-        if self.distributed_executor.parallel_mode == "tile":
-            return False
-        if z.ndim != 5:
-            if self.distributed_executor.parallel_mode != "tile":
-                logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.")
-            return False
-        if not self.is_distributed_enabled():
-            return False
+        return spatial_shard_decode_enabled(self, z, model_name="Wan")
 
-        group = self.distributed_executor.group
-        world_size = dist.get_world_size(group=group)
-        requested_size = int(self.distributed_executor.parallel_size)
-        if requested_size != world_size:
-            if self.distributed_executor.parallel_mode == "auto":
-                logger.debug(
-                    "Wan VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)",
-                    requested_size,
-                    world_size,
-                )
-            elif self._spatial_shard_decode_split_dim() is not None:
-                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,
-                    self._spatial_shard_decode_split_dim(),
-                )
-            return False
-        return self._spatial_shard_decode_split_dim(z, world_size) is not None
+    def _prepare_spatial_shard_decode(self) -> None:
+        prepare_spatial_shard_decode(
+            self,
+            install=wan_spatial_shard.install_wan_spatial_shard_decode,
+        )
 
     def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
         if not self.is_distributed_enabled():

vllm_omni/diffusion/distributed/autoencoders/qwenimage_spatial_shard.py

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

vllm_omni/diffusion/distributed/autoencoders/spatial_shard.py

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

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -7,19 +7,21 @@
 #   (python/sglang/multimodal_gen/runtime/layers/parallel_conv.py)
 # which is in turn adapted from FastVideo (https://github.com/hao-ai-lab/FastVideo).
 # This version generalizes the height-only sharding to shard along height or
-# width and adds the Wan causal-conv ``feat_cache`` handling.
-"""Spatially-sharded Wan VAE decode.
+# width and supports the causal-conv ``feat_cache`` handling used by Wan and
+# QwenImage.
+"""Shared spatial-shard kernels plus the Wan VAE adapter.
 
-The existing distributed Wan VAE path shards *tiles*.  This module adds an
-opt-in decode backend that shards decoder feature maps along height or width and
-exchanges boundary rows/columns before spatial convolutions.  It is
-intentionally decode-only and keeps checkpoint loading unchanged by patching the
-already-loaded decoder.
+The halo, padding, stride-trim, attention-gather, and decode-lifecycle helpers
+are also reused by QwenImage's thin adapter. The backend shards decoder feature
+maps along height or width and exchanges boundary rows/columns before spatial
+convolutions. It is decode-only and keeps checkpoint loading unchanged by
+patching the already-loaded decoder.
 """
 
 from __future__ import annotations
 
 import math
+from collections.abc import Callable
 from contextlib import nullcontext
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -34,6 +36,8 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE
+
 logger = init_logger(__name__)
 
 
@@ -62,7 +66,7 @@ def _spatial_dim(split_dim: str) -> int:
         return -2
     if split_dim == "width":
         return -1
-    raise ValueError(f"Unsupported Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
+    raise ValueError(f"Unsupported VAE split_dim={split_dim!r}; expected 'height' or 'width'.")
 
 
 def _narrow_along_dim(x: torch.Tensor, dim: int, start: int, length: int) -> torch.Tensor:
@@ -160,9 +164,9 @@ def split_for_parallel_decode(
     rank = 0 if rank is None else int(rank)
     world_size = 1 if world_size is None else int(world_size)
     if world_size < 1:
-        raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.")
+        raise ValueError(f"VAE world_size must be >= 1, got {world_size}.")
     if not 0 <= rank < world_size:
-        raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
+        raise ValueError(f"VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.")
 
     dim = _spatial_dim(split_dim)
     expected_extent = x.shape[dim] * (2**upsample_count)
@@ -385,7 +389,7 @@ class WanDistConv2d(nn.Conv2d):
         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"Wan spatial-shard Conv2d requires integer padding, got {source.padding!r}.")
+            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)
@@ -480,30 +484,20 @@ class WanDistConv2d(nn.Conv2d):
         return trim_params
 
 
-class WanDistCausalConv3d(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
+# 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
 
-    def __init__(
+
+class SpatialShardCausalConv3dMixin:
+    """Shared causal-convolution behavior for supported Diffusers VAE classes."""
+
+    def _init_spatial_shard_causal_conv(
         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,
-        )
+    ) -> None:
         self.load_state_dict(source.state_dict())
         self.group = group
         source_padding = getattr(source, "_padding", None)
@@ -511,11 +505,16 @@ class WanDistCausalConv3d(nn.Conv3d):
             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]] = {}
 
-    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
+    def _spatial_shard_causal_conv_forward(
+        self,
+        x: torch.Tensor,
+        cache_x: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         split_dim = _active_spatial_shard_split_dim()
         padding = list(self._source_padding)
         if cache_x is not None and padding[4] > 0:
@@ -524,7 +523,7 @@ class WanDistCausalConv3d(nn.Conv3d):
             padding[4] -= cache_x.shape[2]
 
         if split_dim is None:
-            return super().forward(F.pad(x, padding))
+            return nn.Conv3d.forward(self, F.pad(x, padding))
 
         split_tensor_dim = _spatial_dim(split_dim)
         if split_dim == "height":
@@ -568,7 +567,7 @@ class WanDistCausalConv3d(nn.Conv3d):
                 shift,
                 x_padded.shape[split_tensor_dim] - shift,
             )
-        out = super().forward(x_padded)
+        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)
 
@@ -601,6 +600,36 @@ class WanDistCausalConv3d(nn.Conv3d):
         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,
@@ -674,23 +703,19 @@ 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],
 ) -> None:
-    if child.__class__.__name__ == "WanCausalConv3d":
-        setattr(
-            parent,
-            name,
-            WanDistCausalConv3d(
-                child,
-                group,
-            ),
-        )
+    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)
         setattr(
             parent,
             name,
-            WanDistZeroPad2d(
+            SpatialShardZeroPad2d(
                 padding,
                 group,
                 defer_split_after_padding=parent.__class__.__name__ == "Sequential",
@@ -700,18 +725,18 @@ def _replace_child(
     if isinstance(child, nn.Conv2d):
         deferred_padding = None
         if name == "1" and parent.__class__.__name__ == "Sequential":
-            # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a
+            # Wan/QwenImage downsampling uses ZeroPad2d((0, 1, 0, 1)) before a
             # stride-2 conv with padding=0.  Only the last rank should see the
             # bottom/right global padding, which is approximated by split_padding.
             prev = getattr(parent, "0", None)
-            if isinstance(prev, WanDistZeroPad2d):
+            if isinstance(prev, SpatialShardZeroPad2d):
                 deferred_padding = prev.padding
             elif isinstance(prev, nn.ZeroPad2d):
                 deferred_padding = tuple(int(value) for value in prev.padding)
         setattr(
             parent,
             name,
-            WanDistConv2d(
+            SpatialShardConv2d(
                 child,
                 group,
                 deferred_padding=deferred_padding,
@@ -722,20 +747,45 @@ 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,
     inside_attention: bool = False,
 ) -> None:
-    if module.__class__.__name__ == "WanAttentionBlock":
+    if module.__class__.__name__ == attention_block_class_name:
         _patch_attention_block(module, group)
         inside_attention = True
 
     for name, child in list(module.named_children()):
-        if child.__class__.__name__ == "WanCausalConv3d":
-            _replace_child(module, name, child, group)
+        if child.__class__.__name__ == causal_conv_class_name:
+            _replace_child(
+                module,
+                name,
+                child,
+                group,
+                causal_conv_class_name=causal_conv_class_name,
+                causal_conv_factory=causal_conv_factory,
+            )
             continue
         if isinstance(child, nn.Conv2d) and not inside_attention:
-            _replace_child(module, name, child, group)
+            _replace_child(
+                module,
+                name,
+                child,
+                group,
+                causal_conv_class_name=causal_conv_class_name,
+                causal_conv_factory=causal_conv_factory,
+            )
             continue
-        _patch_decoder_modules(child, group, inside_attention=inside_attention)
+        _patch_decoder_modules(
+            child,
+            group,
+            causal_conv_class_name=causal_conv_class_name,
+            causal_conv_factory=causal_conv_factory,
+            attention_block_class_name=attention_block_class_name,
+            inside_attention=inside_attention,
+        )
 
 
 def _decoder_upsample_count(decoder: nn.Module) -> int:
@@ -746,7 +796,7 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
     return count
 
 
-def _clear_wan_spatial_shard_runtime_buffers(vae: Any) -> None:
+def clear_spatial_shard_runtime_buffers(vae: SpatialShardVAE) -> None:
     decoder = getattr(vae, "decoder", None)
     if decoder is None:
         return
@@ -757,7 +807,17 @@ def _clear_wan_spatial_shard_runtime_buffers(vae: Any) -> None:
             module._halo_recv_bottom_buf = None
 
 
-def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None:
+def install_spatial_shard_decode(
+    vae: SpatialShardVAE,
+    group: dist.ProcessGroup,
+    split_dim: str,
+    *,
+    causal_conv_class_name: str,
+    causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module],
+    attention_block_class_name: str,
+    installed_attr: str,
+    model_name: str,
+) -> None:
     """Patch ``vae.decoder`` once for spatially-sharded decode.
 
     This mutates the already-loaded decoder in place by swapping its spatial
@@ -766,32 +826,28 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
     local PyTorch behavior, so a later auto-selected tile request remains safe.
     The same installed decoder can alternate height and width sharding.
 
-    Only group-relative rank 0 assembles the final decoded frame, mirroring the
-    distributed tiled-decode ``broadcast_result=False`` contract; the other ranks
-    take part in the collectives but return an empty placeholder.
+    Group-relative rank 0 assembles the final decoded frame. The model adapter
+    then preserves its prior contract: Wan keeps the result on rank 0, while
+    QwenImage broadcasts it to every rank.
     """
     _spatial_dim(split_dim)
-    if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False):
+    if getattr(vae, installed_attr, False):
         return
     decoder = getattr(vae, "decoder", None)
     if decoder is None:
-        raise ValueError("Wan spatial-shard VAE decode requires a decoder module.")
+        raise ValueError(f"{model_name} spatial-shard VAE decode requires a decoder module.")
 
-    _patch_decoder_modules(decoder, group)
+    _patch_decoder_modules(
+        decoder,
+        group,
+        causal_conv_class_name=causal_conv_class_name,
+        causal_conv_factory=causal_conv_factory,
+        attention_block_class_name=attention_block_class_name,
+    )
     upsample_count = _decoder_upsample_count(decoder)
     orig_forward = decoder.forward
 
-    def _forward(
-        self: nn.Module,
-        x: torch.Tensor,
-        feat_cache: list[torch.Tensor] | None = None,
-        feat_idx: list[int] | None = None,
-        first_chunk: bool = False,
-        *,
-        split_dim: str,
-    ) -> torch.Tensor:
-        if feat_idx is None:
-            feat_idx = [0]
+    def _forward(self: nn.Module, x: torch.Tensor, *args: object, split_dim: str, **kwargs: object) -> torch.Tensor:
         tensor_dim = _spatial_dim(split_dim)
         input_extent = x.shape[tensor_dim]
         x, expected_extent = split_for_parallel_decode(
@@ -811,28 +867,50 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
             )
         )
         try:
-            out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
+            out = orig_forward(x, *args, **kwargs)
         finally:
             _SPATIAL_SHARD_CONTEXT.reset(token)
         return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0)
 
     decoder._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder)
-    vae._vllm_omni_wan_spatial_shard_installed = True
-    logger.info("Installed Wan VAE dynamic-axis spatial-shard decode.")
+    setattr(vae, installed_attr, True)
+    logger.info("Installed %s VAE dynamic-axis spatial-shard decode.", model_name)
 
 
-def spatial_shard_decode(
-    vae: Any,
+def install_wan_spatial_shard_decode(
+    vae: SpatialShardVAE,
+    group: dist.ProcessGroup,
+    split_dim: str = "height",
+) -> None:
+    install_spatial_shard_decode(
+        vae,
+        group,
+        split_dim,
+        causal_conv_class_name="WanCausalConv3d",
+        causal_conv_factory=WanDistCausalConv3d,
+        attention_block_class_name="WanAttentionBlock",
+        installed_attr="_vllm_omni_wan_spatial_shard_installed",
+        model_name="Wan",
+    )
+
+
+def spatial_shard_decode_impl(
+    vae: SpatialShardVAE,
     z: torch.Tensor,
     *,
     group: dist.ProcessGroup,
+    install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None],
+    model_name: str,
+    pass_first_chunk: bool,
+    broadcast_result: bool,
+    unpatchify_patch_size: int | None,
     return_dict: bool = True,
     split_dim: str = "height",
 ) -> DecoderOutput | tuple[torch.Tensor]:
-    install_wan_spatial_shard_decode(vae, group, split_dim=split_dim)
+    install(vae, group, split_dim)
 
     if z.shape[2] == 0:
-        raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.")
+        raise ValueError(f"{model_name} spatial-shard VAE decode expects at least one latent frame.")
 
     # Non-rank-0 ranks must still run the decoder every chunk to stay in lockstep with
     # the halo/all-gather collectives; they just skip keeping/assembling the output.
@@ -841,26 +919,36 @@ def spatial_shard_decode(
 
     vae.clear_cache()
     try:
-        context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext()
+        context_factory = getattr(vae, "_execution_context", None)
+        context = context_factory() if callable(context_factory) else nullcontext()
         with context:
             x = vae.post_quant_conv(z)
             decoded_chunks = []
+            spatial_forward = getattr(vae.decoder, "_vllm_omni_spatial_shard_forward")
             for i in range(z.shape[2]):
                 vae._conv_idx = [0]
-                chunk = vae.decoder._vllm_omni_spatial_shard_forward(
-                    x[:, :, i : i + 1, :, :],
-                    feat_cache=vae._feat_map,
-                    feat_idx=vae._conv_idx,
-                    first_chunk=(i == 0),
-                    split_dim=split_dim,
-                )
+                if pass_first_chunk:
+                    chunk = spatial_forward(
+                        x[:, :, i : i + 1, :, :],
+                        feat_cache=vae._feat_map,
+                        feat_idx=vae._conv_idx,
+                        first_chunk=i == 0,
+                        split_dim=split_dim,
+                    )
+                else:
+                    chunk = spatial_forward(
+                        x[:, :, i : i + 1, :, :],
+                        feat_cache=vae._feat_map,
+                        feat_idx=vae._conv_idx,
+                        split_dim=split_dim,
+                    )
                 if produce_output:
                     decoded_chunks.append(chunk)
 
             if produce_output:
                 out = torch.cat(decoded_chunks, dim=2)
-                if vae.config.patch_size is not None:
-                    out = unpatchify(out, patch_size=vae.config.patch_size)
+                if unpatchify_patch_size is not None:
+                    out = unpatchify(out, patch_size=unpatchify_patch_size)
                 out = torch.clamp(out, min=-1.0, max=1.0)
             else:
                 out = z.new_zeros(0)
@@ -869,8 +957,34 @@ def spatial_shard_decode(
         # 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_wan_spatial_shard_runtime_buffers(vae)
+        clear_spatial_shard_runtime_buffers(vae)
+
+    if broadcast_result:
+        out = vae.distributed_executor._sync_final_result(out, z.ndim, z.device, vae.dtype)
 
     if not return_dict:
         return (out,)
     return DecoderOutput(sample=out)
+
+
+def spatial_shard_decode(
+    vae: SpatialShardVAE,
+    z: torch.Tensor,
+    *,
+    group: dist.ProcessGroup,
+    return_dict: bool = True,
+    split_dim: str = "height",
+) -> DecoderOutput | tuple[torch.Tensor]:
+    """Decode a Wan latent with spatial activation sharding."""
+    return spatial_shard_decode_impl(
+        vae,
+        z,
+        group=group,
+        install=install_wan_spatial_shard_decode,
+        model_name="Wan",
+        pass_first_chunk=True,
+        broadcast_result=False,
+        unpatchify_patch_size=vae.config.patch_size,
+        return_dict=return_dict,
+        split_dim=split_dim,
+    )

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -155,9 +155,9 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         parallel_size: int,
         mode: str = "tile",
     ) -> None:
-        # ``auto`` is the cross-model default. Spatial selection is Wan-only,
-        # so preserve MiniMax H3's native tile path for that default while
-        # continuing to reject explicitly requested Wan spatial modes.
+        # ``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":

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

vllm_omni/entrypoints/cli/serve.py

@@ -753,7 +753,8 @@ class OmniServeCommand(CLISubcommand):
             default="auto",
             choices=["auto", "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 tiled decodes "
+            "'auto' (default) selects the longer spatial-shard axis for validated Wan and distributed QwenImage "
+            "tiled decodes "
             "on a full DiT group, then falls back to patch/tile decode; 'tile' always uses patch/tile decode; "
             "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits "
             "decoder feature maps along height/width and exchanges halo regions. The "