From 5356357cf0bf95a7818ba2452b2248b5a0e59af2 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Mon, 21 Sep 2026 19:00:46 +0800 Subject: [PATCH] [LingBot World] Shard the streaming VAE decode across the Ulysses ranks (#7651) Signed-off-by: tzhouam Co-authored-by: Claude Fable 5.1 --- .../distributed/test_wan_spatial_shard.py | 177 ++++++++++++++++++ .../test_pipeline_lingbot_world.py | 96 +++++++++- .../deploy/lingbot_world_v2_stepwise.yaml | 3 + .../autoencoders/wan_spatial_shard.py | 35 +++- .../models/lingbot_world/pipeline.py | 78 +++++++- 5 files changed, 378 insertions(+), 11 deletions(-) diff --git a/tests/diffusion/distributed/test_wan_spatial_shard.py b/tests/diffusion/distributed/test_wan_spatial_shard.py index 23c29c938..27177745e 100644 --- a/tests/diffusion/distributed/test_wan_spatial_shard.py +++ b/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 diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index 1ccc7aadc..b5a57ce75 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/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 diff --git a/vllm_omni/deploy/lingbot_world_v2_stepwise.yaml b/vllm_omni/deploy/lingbot_world_v2_stepwise.yaml index 3fc950e4c..9a00b4d6c 100644 --- a/vllm_omni/deploy/lingbot_world_v2_stepwise.yaml +++ b/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: diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py index 16fdd13a4..e3f0ef5c3 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py +++ b/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) diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index ae24c0138..8bb83ef80 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/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.