[LingBot World] Shard the streaming VAE decode across the Ulysses ranks (#7651)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -493,6 +493,84 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
     assert last.shape == (1, 1, 2, 3)
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+class _PlainDecoder(torch.nn.Module):
+    """A decoder with nothing to patch, so the install only wraps ``forward``."""
+
+    def __init__(self):
+        super().__init__()
+        self.calls: list[dict] = []
+
+    def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False):
+        self.calls.append({"feat_cache": feat_cache, "feat_idx": feat_idx, "first_chunk": first_chunk})
+        return x
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size, dst", [(4, -1), (4, 4), (4, 5), (1, 1), (4, 1.5)])
+def test_install_rejects_invalid_destination_before_mutating_decoder(monkeypatch, world_size, dst):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, world_size))
+    vae = SimpleNamespace(decoder=torch.nn.Identity())
+    original_forward = vae.decoder.forward
+
+    def unexpected_patch(*args, **kwargs):
+        pytest.fail("invalid destination must fail before patching modules")
+
+    monkeypatch.setattr(wan_spatial_shard, "_patch_decoder_modules", unexpected_patch)
+    with pytest.raises(ValueError, match="dst must be None or an integer"):
+        wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), dst=dst)
+    assert vae.decoder.forward == original_forward
+    assert not getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False)
+
+
+def _install_on_plain_decoder(monkeypatch: pytest.MonkeyPatch, *, dst):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 2))
+    monkeypatch.setattr(
+        wan_spatial_shard,
+        "split_for_parallel_decode",
+        lambda x, *, upsample_count, split_dim, group: (x, x.shape[wan_spatial_shard._spatial_dim(split_dim)]),
+    )
+    gathers: list[dict] = []
+
+    def gather(x, *, expected_extent, split_dim, group, dst=None):
+        gathers.append({"expected_extent": expected_extent, "split_dim": split_dim, "dst": dst})
+        return x
+
+    monkeypatch.setattr(wan_spatial_shard, "gather_and_trim_extent", gather)
+    vae = SimpleNamespace(decoder=_PlainDecoder())
+    wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), split_dim="width", dst=dst)
+    return vae, gathers
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("dst", [0, 1, None])
+def test_install_passes_the_assembling_rank_through_to_the_gather(monkeypatch: pytest.MonkeyPatch, dst):
+    vae, gathers = _install_on_plain_decoder(monkeypatch, dst=dst)
+    cache = [None]
+
+    out = vae.decoder(torch.zeros(1, 3, 1, 4, 6), feat_cache=cache, feat_idx=[0], first_chunk=True)
+
+    assert out.shape == (1, 3, 1, 4, 6)
+    assert gathers == [{"expected_extent": 6, "split_dim": "width", "dst": dst}]
+    # The streaming arguments reach the wrapped decoder untouched.
+    assert vae.decoder.calls == [{"feat_cache": cache, "feat_idx": [0], "first_chunk": True}]
+    assert vae._vllm_omni_wan_spatial_shard_dst == dst
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_install_refuses_to_change_the_assembling_rank(monkeypatch: pytest.MonkeyPatch):
+    vae, _ = _install_on_plain_decoder(monkeypatch, dst=None)
+
+    # Same settings again: a no-op, as before.
+    wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), split_dim="width", dst=None)
+    with pytest.raises(ValueError, match="already patched to assemble on None"):
+        wan_spatial_shard.install_wan_spatial_shard_decode(vae, object(), split_dim="width", dst=0)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pytest.MonkeyPatch):
@@ -626,6 +704,105 @@ def _spatial_shard_decode_worker(rank: int, split_dim: str, return_dict, master_
             dist.destroy_process_group()
 
 
+_STREAMING_SHARD_CHUNKS = 3
+
+
+def _streaming_shard_worker(rank: int, return_dict, master_port: str) -> None:
+    """Sharded streaming decode with every rank keeping the frame, against an unsharded streaming reference.
+
+    Two sessions are interleaved chunk by chunk, so the per-session temporal cache
+    and the halo exchange are exercised across successive calls, then one session
+    is reset and must reproduce its opening chunk. Every rank checks every chunk.
+    """
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        destroy_model_parallel,
+        get_sp_group,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
+
+    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)
+    dtype = torch.float32
+    backend = current_omni_platform.dist_backend
+    init_distributed_environment(world_size=_SPATIAL_SHARD_WORLD_SIZE, rank=rank, local_rank=rank, backend=backend)
+    initialize_model_parallel(
+        sequence_parallel_size=_SPATIAL_SHARD_WORLD_SIZE, ulysses_degree=_SPATIAL_SHARD_WORLD_SIZE, backend=backend
+    )
+    # VLLM_OMNI_TEST_WAN_VAE may point at a local diffusers checkpoint (with or
+    # without a ``vae`` subfolder) for machines without hub access.
+    model = os.environ.get("VLLM_OMNI_TEST_WAN_VAE", _SPATIAL_SHARD_MODEL)
+    local_vae_only = os.path.isdir(model) and not os.path.isdir(os.path.join(model, "vae"))
+    load = {"torch_dtype": dtype} if local_vae_only else {"subfolder": _SPATIAL_SHARD_SUBFOLDER, "torch_dtype": dtype}
+    try:
+        reference_vae = DistributedAutoencoderKLWan.from_pretrained(model, **load).to(device=device, dtype=dtype).eval()
+        sharded_vae = DistributedAutoencoderKLWan.from_pretrained(model, **load).to(device=device, dtype=dtype).eval()
+        wan_spatial_shard.install_wan_spatial_shard_decode(
+            sharded_vae, get_sp_group().device_group, split_dim="width", dst=None
+        )
+        reference = WanStreamingDecoder(reference_vae, bytes_per_pixel_fp32=1.0)
+        sharded = WanStreamingDecoder(sharded_vae, bytes_per_pixel_fp32=1.0)
+
+        generator = torch.Generator(device=device).manual_seed(0)
+        # Width 104 latents = 832 px: not divisible by two ranks' halo-padded shards without padding.
+        shape = (1, reference_vae.config.z_dim, 1, _SPATIAL_SHARD_LATENT_HEIGHT, _SPATIAL_SHARD_LATENT_WIDTH)
+        chunks = {
+            name: [
+                torch.randn(shape, generator=generator, device=device, dtype=dtype)
+                for _ in range(_STREAMING_SHARD_CHUNKS)
+            ]
+            for name in ("a", "b")
+        }
+        max_diff = 0.0
+        with torch.inference_mode():
+            ref_states = {name: reference.new_decode_state(name) for name in chunks}
+            shard_states = {name: sharded.new_decode_state(name) for name in chunks}
+            for index in range(_STREAMING_SHARD_CHUNKS):
+                for name in ("a", "b"):
+                    expected = reference.decode_chunk(chunks[name][index], ref_states[name]).float()
+                    got = sharded.decode_chunk(chunks[name][index], shard_states[name]).float()
+                    if got.shape != expected.shape:
+                        raise AssertionError(
+                            f"rank {rank} {name}[{index}]: {tuple(got.shape)} != {tuple(expected.shape)}"
+                        )
+                    max_diff = max(max_diff, (got - expected).abs().max().item())
+            # Reset session a: a fresh state must reproduce its opening chunk, on this rank too.
+            shard_states["a"].release()
+            ref_states["a"].release()
+            shard_states["a"] = sharded.new_decode_state("a")
+            ref_states["a"] = reference.new_decode_state("a")
+            expected = reference.decode_chunk(chunks["a"][0], ref_states["a"]).float()
+            got = sharded.decode_chunk(chunks["a"][0], shard_states["a"]).float()
+            max_diff = max(max_diff, (got - expected).abs().max().item())
+        return_dict[f"max_abs_diff_rank{rank}"] = max_diff
+        return_dict[f"shape_rank{rank}"] = tuple(got.shape)
+    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", "B200"]}, num_cards=_SPATIAL_SHARD_WORLD_SIZE)
+def test_streaming_shard_decode_matches_unsharded_streaming_on_every_rank():
+    manager = mp.get_context("spawn").Manager()
+    return_dict = manager.dict()
+    mp.spawn(_streaming_shard_worker, args=(return_dict, "29510"), nprocs=_SPATIAL_SHARD_WORLD_SIZE, join=True)
+
+    for rank in range(_SPATIAL_SHARD_WORLD_SIZE):
+        assert f"max_abs_diff_rank{rank}" in return_dict, f"rank {rank} did not report"
+        diff = return_dict[f"max_abs_diff_rank{rank}"]
+        shape = return_dict[f"shape_rank{rank}"]
+        print(f"streaming shard rank {rank}: max_abs_diff={diff:.6e} shape={shape}")
+        assert shape[-1] == _SPATIAL_SHARD_LATENT_WIDTH * 8 and shape[-2] == _SPATIAL_SHARD_LATENT_HEIGHT * 8
+        assert diff <= _SPATIAL_SHARD_TOLERANCE, f"rank {rank}: {diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
+
+
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py

@@ -737,6 +737,99 @@ def test_component_discovery_uses_official_checkpoint_contract() -> None:
     assert module._loader_state.prefetch_calls == [("checkpoint", ("tokenizer", "text_encoder", "vae"), False)]
 
 
+def _ulysses_config(size: int):
+    parallel_config = _od_config().parallel_config
+    parallel_config.sequence_parallel_size = size
+    parallel_config.ulysses_degree = size
+    return _od_config(parallel_config=parallel_config)
+
+
+def _record_shard_install(monkeypatch, module, *, world_size: int):
+    """Stand in for the Ulysses group and the decoder patch; return what the install was asked for."""
+    installs: list[dict] = []
+    group = object()
+
+    def install(vae, group_arg, split_dim, *, dst):
+        installs.append({"vae": vae, "group": group_arg, "split_dim": split_dim, "dst": dst})
+
+    monkeypatch.setattr(module, "install_wan_spatial_shard_decode", install)
+    monkeypatch.setattr(module.LingBotWorldCausalDMDPipeline, "_vae_shard_group", lambda self: (group, world_size))
+    return installs, group
+
+
+def test_a_multi_rank_deployment_shards_the_decoder_across_the_ulysses_ranks(monkeypatch) -> None:
+    """Multi-rank deployments shard decode by default."""
+    module = _load_pipeline_module()
+    installs, group = _record_shard_install(monkeypatch, module, world_size=2)
+
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(2))
+
+    # Along the width, and assembled on every rank: each rank's post_decode consumes the frame.
+    assert installs == [{"vae": pipeline.vae, "group": group, "split_dim": "width", "dst": None}]
+    assert pipeline._vae_shard_split_dim == "width"
+
+
+def test_a_single_rank_deployment_leaves_the_decoder_alone(monkeypatch) -> None:
+    module = _load_pipeline_module()
+    installs, _ = _record_shard_install(monkeypatch, module, world_size=1)
+
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(1))
+
+    assert installs == [] and pipeline._vae_shard_split_dim is None
+
+
+@pytest.mark.parametrize("enabled", [True, False])
+@pytest.mark.parametrize("width, local_width", [(832, 208), (840, 216)])
+def test_streaming_decode_reservation_uses_padded_shards(monkeypatch, enabled, width, local_width):
+    module = _load_pipeline_module()
+    installs, _ = _record_shard_install(monkeypatch, module, world_size=4)
+    config = _ulysses_config(4)
+    config.model_config.update(lingbot_vae_spatial_sharding=enabled, ar_diffusion_height=480, ar_diffusion_width=width)
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=config)
+    from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
+
+    # Exercise the real byte estimator without constructing a checkpoint VAE.
+    decoder = object.__new__(WanStreamingDecoder)
+    decoder._bytes_per_pixel_fp32 = 16.0
+    monkeypatch.setattr(pipeline, "_streaming_decoder", lambda: decoder)
+    expected_width = local_width if enabled else width
+    assert pipeline._streaming_decode_bytes_per_session() == 16 * 480 * expected_width
+    assert len(installs) == int(enabled)
+    assert pipeline._vae_shard_world_size == (4 if enabled else 1)
+    assert config.parallel_config.sequence_parallel_size == config.parallel_config.ulysses_degree == 4
+
+
+def test_disabling_vae_sharding_does_not_require_a_shard_group(monkeypatch):
+    module = _load_pipeline_module()
+    config = _ulysses_config(4)
+    config.model_config["lingbot_vae_spatial_sharding"] = False
+
+    def unexpected_group(self):
+        pytest.fail("disabled VAE sharding must not access its process group")
+
+    monkeypatch.setattr(module.LingBotWorldCausalDMDPipeline, "_vae_shard_group", unexpected_group)
+    pipeline = module.LingBotWorldCausalDMDPipeline(od_config=config)
+    assert pipeline._vae_shard_split_dim is None
+
+
+@pytest.mark.parametrize("value", ["false", 0, None])
+def test_vae_sharding_switch_requires_a_boolean(value):
+    module = _load_pipeline_module()
+    config = _ulysses_config(4)
+    config.model_config["lingbot_vae_spatial_sharding"] = value
+    with pytest.raises(ValueError, match="lingbot_vae_spatial_sharding must be a boolean"):
+        module.LingBotWorldCausalDMDPipeline(od_config=config)
+
+
+def test_vae_shard_refuses_a_group_of_the_wrong_size(monkeypatch) -> None:
+    module = _load_pipeline_module()
+    installs, _ = _record_shard_install(monkeypatch, module, world_size=4)
+
+    with pytest.raises(RuntimeError, match="sequence_parallel_size=2 but the Ulysses group has 4 ranks"):
+        module.LingBotWorldCausalDMDPipeline(od_config=_ulysses_config(2))
+    assert installs == []
+
+
 @pytest.mark.parametrize(
     ("field", "value", "feature"),
     [
@@ -758,8 +851,9 @@ def test_unsupported_parallel_modes_fail_before_component_loading(field: str, va
     assert module._loader_state.prefetch_calls == []
 
 
-def test_pure_ulysses_parallel_config_is_supported() -> None:
+def test_pure_ulysses_parallel_config_is_supported(monkeypatch) -> None:
     module = _load_pipeline_module()
+    _record_shard_install(monkeypatch, module, world_size=2)
     parallel_config = _od_config().parallel_config
     parallel_config.sequence_parallel_size = 2
     parallel_config.ulysses_degree = 2

vllm_omni/deploy/lingbot_world_v2_stepwise.yaml

@@ -34,6 +34,9 @@ stages:
       extra_args:
         flow_shift: 5.0
     model_config:
+      # Multi-rank width sharding changes floating-point numerics at shard borders.
+      # Set false for full-frame VAE decode while keeping DiT sequence parallelism.
+      lingbot_vae_spatial_sharding: true
       ar_diffusion_height: 480
       ar_diffusion_width: 832
       ar_diffusion_kv_config:

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -781,22 +781,34 @@ def _decoder_upsample_count(decoder: nn.Module) -> int:
     return count
 
 
-def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None:
+def install_wan_spatial_shard_decode(
+    vae: Any,
+    group: dist.ProcessGroup,
+    split_dim: str = "height",
+    *,
+    dst: int | None = 0,
+) -> None:
     """Patch ``vae.decoder`` once for spatially-sharded decode.
 
     This mutates the already-loaded decoder in place by swapping its spatial
     convolutions/padding for halo-exchanging variants and wrapping
     ``decoder.forward``. The patch is permanent for the lifetime of the VAE
     instance and is applied only once (subsequent calls are no-ops). A given
-    instance is bound to a single ``split_dim``; switching between
-    ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here
-    otherwise.
+    instance is bound to a single ``split_dim`` and ``dst``; switching either
+    requires a fresh VAE instance and raises here otherwise.
 
-    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.
+    With ``dst=0`` (the default) 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. ``dst=None`` lets every rank
+    keep the assembled frame, for a caller whose every rank goes on to
+    consume the decoded pixels; the gather is an all-gather either way, so
+    this costs no extra communication.
     """
     _spatial_dim(split_dim)
+    _, world_size = _rank_world(group)
+    if dst is not None and (not isinstance(dst, int) or not 0 <= dst < world_size):
+        raise ValueError(f"Wan spatial-shard dst must be None or an integer in [0, {world_size}), got {dst!r}.")
     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:
@@ -804,6 +816,12 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
                 "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."
             )
+        installed_dst = getattr(vae, "_vllm_omni_wan_spatial_shard_dst", 0)
+        if installed_dst != dst:
+            raise ValueError(
+                "Wan spatial-shard VAE decoder was already patched to assemble on "
+                f"{installed_dst!r}; create a fresh VAE instance to assemble on {dst!r}."
+            )
         return
     decoder = getattr(vae, "decoder", None)
     if decoder is None:
@@ -844,11 +862,12 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d
             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)
+        return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=dst)
 
     decoder.forward = MethodType(_forward, decoder)
     vae._vllm_omni_wan_spatial_shard_installed = True
     vae._vllm_omni_wan_spatial_shard_split_dim = split_dim
+    vae._vllm_omni_wan_spatial_shard_dst = dst
     logger.info("Installed Wan VAE %s-sharded decode.", split_dim)
 
 

vllm_omni/diffusion/models/lingbot_world/pipeline.py

@@ -26,6 +26,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
+from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import install_wan_spatial_shard_decode
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.interaction.mixin import InteractionMixin
 from vllm_omni.diffusion.interaction.modality_handlers.camera import CameraSession
@@ -92,6 +93,10 @@ logger = init_logger(__name__)
 # measured to hold.
 _STREAMING_DECODE_BYTES_PER_PIXEL_FP32 = 37832 * 1024 / (64 * 64)
 
+# The frame dimension the decoder is sharded along when the pipeline runs on
+# more than one Ulysses rank; see ``_install_sharded_vae_decode``.
+_VAE_SHARD_SPLIT_DIM = "width"
+
 if TYPE_CHECKING:
     from tqdm.std import tqdm as TqdmProgressBar
 
@@ -580,6 +585,7 @@ class LingBotWorldCausalDMDPipeline(
         super().__init__()
         del prefix
         _validate_parallel_config(od_config)
+        parallel_config = getattr(od_config, "parallel_config", None)
         self.od_config = od_config
         model_config = getattr(od_config, "model_config", None) or {}
         reuse_last_step_kv = model_config.get("lingbot_reuse_last_step_kv", False)
@@ -683,6 +689,14 @@ class LingBotWorldCausalDMDPipeline(
         self._streaming_decode_states: dict[str, StreamingDecodeState] = {}
         self._cached_streaming_decoder: WanStreamingDecoder | None = None
         self._streaming_decode_unsupported = False
+        self._vae_shard_split_dim: str | None = None
+        self._vae_shard_world_size = 1
+        vae_sharding = model_config.get("lingbot_vae_spatial_sharding", True)
+        if not isinstance(vae_sharding, bool):
+            raise ValueError("lingbot_vae_spatial_sharding must be a boolean.")
+        sequence_parallel_size = int(getattr(parallel_config, "sequence_parallel_size", 1) or 1)
+        if sequence_parallel_size > 1 and vae_sharding:
+            self._install_sharded_vae_decode(sequence_parallel_size)
         self.setup_diffusion_pipeline_profiler(
             profiler_targets=[
                 "vae.encode",
@@ -803,9 +817,24 @@ class LingBotWorldCausalDMDPipeline(
         decoder = self._streaming_decoder()
         if decoder is None:
             return 0
+        height, width = self._ar_height, self._ar_width
+        if self._vae_shard_split_dim is not None:
+            # The decoder pads the latent extent to a multiple of the group
+            # size before splitting. Account for that padding on every rank,
+            # including ranks whose final pixels are trimmed after gathering.
+            scale = self.vae_scale_factor_spatial
+            extent = width if self._vae_shard_split_dim == "width" else height
+            latent_extent = (extent + scale - 1) // scale
+            local_extent = (latent_extent + self._vae_shard_world_size - 1) // self._vae_shard_world_size
+            if self._vae_shard_split_dim == "width":
+                width = local_extent * scale
+            else:
+                height = local_extent * scale
+            # Temporal caches retain local inputs, before spatial halo exchange.
+            # Halo receive buffers belong to the decoder, not to each session.
         return decoder.declared_state_bytes(
-            height=self._ar_height,
-            width=self._ar_width,
+            height=height,
+            width=width,
             dtype=self.vae.dtype,
         )
 
@@ -1659,6 +1688,51 @@ class LingBotWorldCausalDMDPipeline(
         tile_latent_min_width = int(getattr(self.vae, "tile_sample_min_width", None) or 0) // compression
         return bool(latents.shape[-2] > tile_latent_min_height or latents.shape[-1] > tile_latent_min_width)
 
+    def _vae_shard_group(self) -> tuple[Any, int]:
+        """The Ulysses ranks the decoder is sharded across, and how many there are."""
+        from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
+
+        group = get_sp_group().device_group
+        return group, torch.distributed.get_world_size(group=group)
+
+    def _install_sharded_vae_decode(self, sequence_parallel_size: int) -> None:
+        """Shard the Wan decoder across the Ulysses ranks along the width.
+
+        The DiT runs on ``sequence_parallel_size`` ranks, and the decode of
+        one AR block used to run on every one of them in full: each rank
+        decoded the whole frame so that its own ``post_decode`` could carry
+        on. Splitting the frame across those same ranks with halo exchange at
+        the convolution borders makes each decode a quarter of the work at
+        four ranks; every rank still ends the block with the assembled frame
+        because the gather is an all-gather, so nothing downstream changes.
+        The streaming decoder is unaffected: ``feat_cache``, ``feat_idx`` and
+        ``first_chunk`` pass straight through to the patched decoder, and the
+        per-session temporal cache shrinks by the same factor.
+
+        This is the default for multi-rank deployments; setting
+        ``model_config.lingbot_vae_spatial_sharding=False`` keeps full-frame
+        decode on each rank without changing DiT parallelism.
+        ``vae_patch_parallel_size`` stays rejected because it would make the
+        registry set up the tiled distributed executor and force
+        tiling on a decoder that is already sharded. Width rather than height
+        because at 480x832 it is the longer side, so each rank keeps a
+        208-pixel shard against the same halo, and it is the split this was
+        measured on: 269.9 -> 98.7 ms per block on 4xH200. Numerically it is
+        not bit-exact against a single-rank decode: the halo-exchanged
+        convolutions accumulate in a different order at the shard borders
+        (max abs diff 0.0195 on [-1, 1] at 480x832).
+        """
+        group, world_size = self._vae_shard_group()
+        if world_size != sequence_parallel_size:
+            raise RuntimeError(
+                f"LingBot World sequence_parallel_size={sequence_parallel_size} but the Ulysses group has "
+                f"{world_size} ranks."
+            )
+        install_wan_spatial_shard_decode(self.vae, group, split_dim=_VAE_SHARD_SPLIT_DIM, dst=None)
+        self._vae_shard_split_dim = _VAE_SHARD_SPLIT_DIM
+        self._vae_shard_world_size = world_size
+        logger.info("LingBot World VAE decode sharded along %s across %d ranks.", _VAE_SHARD_SPLIT_DIM, world_size)
+
     def _streaming_decoder(self) -> WanStreamingDecoder | None:
         """The session-owned streaming decoder, or ``None`` on a VAE that cannot stream.