diff --git a/docs/design/feature/vae_parallel.md b/docs/design/feature/vae_parallel.md index c465449f8..3fe165328 100644 --- a/docs/design/feature/vae_parallel.md +++ b/docs/design/feature/vae_parallel.md @@ -451,16 +451,16 @@ Complete examples in the codebase: ## Spatially-Sharded Decode (Wan) -The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`). +The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE 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. ### How it differs from tile parallel -| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) | -|--------|--------------------------|-----------------------------------------------| -| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W | -| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) | -| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) | -| Scope | Decode + encode | Decode only | +| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) | +|--------|--------------------------|----------------------|-----------------------------------------------| +| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W | +| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) | +| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) | +| Scope | Decode + encode | Wan 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`. @@ -476,11 +476,11 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs -> DistributedAutoencoderKLWan.tiled_decode dispatch ``` -`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode. +`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. ### Notes -- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance. +- The decoder is patched **in place** the first time spatial-shard decode runs. Its wrappers consult a context-local split dimension and retain an exact direct fallback, allowing later requests to select tile, height, or width safely. - Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`). --- diff --git a/docs/user_guide/diffusion/parallelism/vae_parallelism.md b/docs/user_guide/diffusion/parallelism/vae_parallelism.md index ca0c428e1..a45d69ec2 100644 --- a/docs/user_guide/diffusion/parallelism/vae_parallelism.md +++ b/docs/user_guide/diffusion/parallelism/vae_parallelism.md @@ -113,7 +113,7 @@ In `DiffusionParallelConfig`: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. | -| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). | +| `vae_parallel_mode` | str | `"auto"` | VAE parallel decode strategy: `"auto"` (shape/model/topology-aware selection), `"tile"`, `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding is currently Wan-only. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). | Additional requirements: @@ -128,7 +128,7 @@ Additional requirements: ## Spatially-Sharded Decode (Wan) -The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`. +The default `vae_parallel_mode="auto"` selects spatial-shard decode only after a Wan request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. It splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests, partial groups, and other VAE families keep their existing path. Set `vae_parallel_mode="tile"` to force the previous behavior, or `"spatial_shard_height"` / `"spatial_shard_width"` to force an axis. Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance. @@ -141,7 +141,7 @@ omni = Omni( parallel_config=DiffusionParallelConfig( tensor_parallel_size=2, vae_patch_parallel_size=2, # must match the DiT group size - vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height" + vae_parallel_mode="auto", # default ), ) ``` @@ -152,14 +152,16 @@ Or from the CLI / serving entrypoint: vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \ --tensor-parallel-size 2 \ --vae-patch-parallel-size 2 \ - --vae-parallel-mode spatial_shard_width + --vae-parallel-mode auto ``` **Constraints and behavior:** - Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes. -- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime. -- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension). +- 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. + +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. For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark: diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 3d2caeabf..e89e75889 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -579,6 +579,7 @@ def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface(): assert cfg.data_parallel_size == 3 assert cfg.cfg_parallel_size == 3 assert cfg.mask_sp_padding is True + assert cfg.vae_parallel_mode == "auto" assert cfg.world_size == 72 diff --git a/tests/diffusion/distributed/test_wan_spatial_shard.py b/tests/diffusion/distributed/test_wan_spatial_shard.py index 72a2a45f5..a7983a734 100644 --- a/tests/diffusion/distributed/test_wan_spatial_shard.py +++ b/tests/diffusion/distributed/test_wan_spatial_shard.py @@ -1,4 +1,5 @@ import os +from contextlib import contextmanager from types import SimpleNamespace import pytest @@ -10,6 +11,8 @@ from tests.helpers.mark import hardware_test from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( DistributedAutoencoderKLWan, + prepare_pipeline_wan_spatial_shard_decode, + select_auto_spatial_shard_split_dim, ) from vllm_omni.platforms import current_omni_platform @@ -17,6 +20,23 @@ from vllm_omni.platforms import current_omni_platform # the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel. +@contextmanager +def _active_spatial_shard(split_dim: str): + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=2, + local_input_extent=2, + split_dim=split_dim, + rank=0, + world_size=3, + ) + ) + try: + yield + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + @pytest.mark.core_model @pytest.mark.cpu def test_split_for_parallel_decode_pads_uneven_height(): @@ -255,12 +275,14 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3)) mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object()) - mid = mid_rank_pad(x) + with _active_spatial_shard("height"): + mid = mid_rank_pad(x) assert mid.shape == (1, 1, 2, 3) monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object()) - last = last_rank_pad(x) + with _active_spatial_shard("height"): + last = last_rank_pad(x) assert last.shape == (1, 1, 3, 3) @@ -270,13 +292,15 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke x = torch.ones((1, 1, 2, 2)) monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3)) - mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width") - mid = mid_rank_pad(x) + mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object()) + with _active_spatial_shard("width"): + mid = mid_rank_pad(x) assert mid.shape == (1, 1, 2, 2) monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) - last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width") - last = last_rank_pad(x) + last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object()) + with _active_spatial_shard("width"): + last = last_rank_pad(x) assert last.shape == (1, 1, 2, 3) @@ -299,6 +323,113 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt assert vae._spatial_shard_decode_enabled(z) is False +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + "z_shape,world_size,expected", + [ + ((1, 16, 1, 60, 104), 4, "width"), + ((1, 16, 5, 60, 104), 4, "width"), + ((1, 16, 5, 104, 60), 4, "height"), + ((1, 16, 5, 2, 2), 4, None), + ((1, 16, 5, 64, 64), 4, "width"), + ], +) +def test_auto_spatial_shard_selector_is_shape_aware(z_shape, world_size, expected): + assert select_auto_spatial_shard_split_dim(z_shape, world_size) == expected + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size", + lambda group=None: 4, + ) + + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + vae.use_tiling = True + vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=4, parallel_mode="auto") + vae.is_distributed_enabled = lambda: True + large = torch.zeros((1, 16, 5, 60, 104)) + unshardable = torch.zeros((1, 16, 1, 2, 2)) + + assert vae._spatial_shard_decode_split_dim(large, 4) == "width" + assert vae._spatial_shard_decode_enabled(large) is True + assert vae._spatial_shard_decode_enabled(unshardable) is False + + vae.distributed_executor.parallel_size = 2 + assert vae._spatial_shard_decode_enabled(large) is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch): + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + object.__setattr__( + vae, + "distributed_executor", + SimpleNamespace(group=object(), parallel_size=2, parallel_mode="auto"), + ) + installed = [] + + monkeypatch.setattr( + "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover", + lambda pipeline: SimpleNamespace(vaes=[vae]), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.is_initialized", + lambda: True, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size", + lambda group=None: 4, + ) + monkeypatch.setattr( + wan_spatial_shard, + "install_wan_spatial_shard_decode", + lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)), + ) + + prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity()) + assert installed == [] + + vae.distributed_executor.parallel_size = 4 + prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity()) + assert installed == [(vae, vae.distributed_executor.group, "height")] + + +@pytest.mark.core_model +@pytest.mark.cpu +@torch.no_grad() +def test_installed_spatial_decoder_keeps_direct_fallback_and_can_change_axis(): + vae = DistributedAutoencoderKLWan( + base_dim=8, + decoder_base_dim=8, + z_dim=4, + dim_mult=[1, 2], + num_res_blocks=1, + attn_scales=[], + temperal_downsample=[False], + latents_mean=[0.0] * 4, + latents_std=[1.0] * 4, + scale_factor_temporal=1, + scale_factor_spatial=2, + ).eval() + latents = torch.randn(1, 4, 1, 4, 5) + parameter_names = {name for name, _ in vae.named_parameters()} + state_dict = {name: value.clone() for name, value in vae.state_dict().items()} + reference = vae.decode(latents, return_dict=False)[0] + + wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="height") + assert any(getattr(module, "_vllm_omni_dynamic_spatial_shard_conv", False) for module in vae.decoder.modules()) + wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="width") + + assert {name for name, _ in vae.named_parameters()} == parameter_names + vae.load_state_dict(state_dict, strict=True) + assert torch.equal(vae.decode(latents, return_dict=False)[0], reference) + + @pytest.mark.core_model @pytest.mark.cpu def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch): @@ -444,3 +575,99 @@ def test_spatial_shard_decode_matches_reference(split_dim: str): assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, ( f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}" ) + + +def _auto_spatial_shard_decode_worker(rank: int, return_dict, master_port: str) -> None: + from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, + ) + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = master_port + device = current_omni_platform.get_torch_device(rank) + current_omni_platform.set_device(device) + backend = current_omni_platform.dist_backend + init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend) + initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend) + + try: + vae = DistributedAutoencoderKLWan.from_pretrained( + _SPATIAL_SHARD_MODEL, + subfolder=_SPATIAL_SHARD_SUBFOLDER, + torch_dtype=torch.float32, + ).to(device=device, dtype=torch.float32) + vae.eval() + generator = torch.Generator(device=device).manual_seed(1) + small = torch.randn((1, vae.config.z_dim, 1, 60, 104), generator=generator, device=device) + landscape = torch.randn((1, vae.config.z_dim, 2, 60, 104), generator=generator, device=device) + portrait = torch.randn((1, vae.config.z_dim, 2, 104, 60), generator=generator, device=device) + + with torch.inference_mode(): + vae.use_tiling = False + vae.set_parallel_size(1, mode="tile") + small_reference = vae.decode(small, return_dict=False)[0] + + vae.use_tiling = True + vae.set_parallel_size(2, mode="tile") + tile_before_install = vae.decode(small, return_dict=False)[0] + + vae.set_parallel_size(2, mode="auto") + assert vae._spatial_shard_decode_split_dim(small, 2) == "width" + small_sharded = vae.decode(small, return_dict=False)[0] + assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width" + vae.decode(landscape, return_dict=False) + + vae.set_parallel_size(2, mode="tile") + tile_after_install = vae.decode(small, return_dict=False)[0] + + vae.use_tiling = False + vae.set_parallel_size(1, mode="tile") + portrait_reference = vae.decode(portrait, return_dict=False)[0] + + vae.use_tiling = True + vae.set_parallel_size(2, mode="auto") + assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height" + portrait_sharded = vae.decode(portrait, return_dict=False)[0] + + assert all( + buffer is None + for module in vae.decoder.modules() + for name, buffer in module._buffers.items() + if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"} + ) + + if rank == 0: + small_diff = (small_sharded.float() - small_reference.float()).abs() + portrait_diff = (portrait_sharded.float() - portrait_reference.float()).abs() + return_dict["tile_fallback_exact"] = torch.equal(tile_after_install, tile_before_install) + return_dict["small_max_abs_diff"] = small_diff.max().item() + return_dict["small_mean_abs_diff"] = small_diff.mean().item() + return_dict["portrait_max_abs_diff"] = portrait_diff.max().item() + return_dict["portrait_mean_abs_diff"] = portrait_diff.mean().item() + finally: + destroy_model_parallel() + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "H100"}, num_cards=2) +def test_auto_spatial_shard_switches_shape_paths_safely(): + manager = mp.get_context("spawn").Manager() + return_dict = manager.dict() + mp.spawn( + _auto_spatial_shard_decode_worker, + args=(return_dict, "29502"), + nprocs=2, + join=True, + ) + + assert return_dict.get("tile_fallback_exact") is True + assert return_dict["small_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE + assert return_dict["small_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE + assert return_dict["portrait_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE + assert return_dict["portrait_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py index 2b168499e..7c406c1fb 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py @@ -120,3 +120,34 @@ def test_tp_accepts_checkpoint_supported_sizes(): arch = MiniMaxH3DiTArchConfig() for tp_size in (1, 2, 4, 7): model._validate_tp_config(arch=arch, tp_size=tp_size) + + +@pytest.mark.parametrize(("parallel_size", "enabled"), [(1, False), (2, True)]) +def test_video_vae_auto_mode_preserves_native_tile_parallelism(monkeypatch, parallel_size, enabled): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + parallel_state = {} + video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE) + nn.Module.__init__(video_vae) + video_vae.remote = nn.Identity() + video_vae.model = type("FakeModel", (), {"parallel_tiling": False})() + + monkeypatch.setattr(vae_module, "get_dit_group", lambda: object()) + monkeypatch.setattr(vae_module.dist, "get_world_size", lambda group: 2) + monkeypatch.setattr(vae_module.dist, "get_rank", lambda group: 0) + monkeypatch.setattr( + vae_module.importlib, + "import_module", + lambda package: type( + "FakeParallelModule", + (), + {"get_parallel_state": staticmethod(lambda: parallel_state)}, + )(), + ) + + video_vae.set_parallel_size(parallel_size, mode="auto") + + assert video_vae.parallel_size == parallel_size + assert video_vae.model.parallel_tiling is enabled + assert parallel_state["group_size"] == parallel_size + assert parallel_state["sp_enabled"] is enabled diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index a1f8440e3..be8a81331 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -728,6 +728,19 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch): @pytest.mark.core_model @pytest.mark.cpu def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): + memory_pool_active = False + preparation_pool_states = [] + + @contextmanager + def _memory_pool_context(tag): + nonlocal memory_pool_active + assert tag == "weights" + memory_pool_active = True + try: + yield + finally: + memory_pool_active = False + class _DummyLoader: def __init__(self, load_config, od_config=None): del load_config, od_config @@ -775,16 +788,22 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object()) monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler) + monkeypatch.setattr( + model_runner_module, + "prepare_pipeline_wan_spatial_shard_decode", + lambda pipeline: preparation_pool_states.append(memory_pool_active), + ) monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None) monkeypatch.setattr( model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend ) - DiffusionModelRunner.load_model(runner) + DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context) assert runner.cache_backend is None assert runner.od_config.cache_backend is None assert dummy_cache_backend.enabled is False + assert preparation_pool_states == [True] @pytest.mark.core_model diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 04a29edd1..35c155155 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -76,6 +76,18 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None: assert args.cfg_parallel_size == 4 +def test_serve_parser_defaults_vae_parallel_mode_to_auto_and_accepts_tile() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + default_args = parser.parse_args(["serve", "fake-model", "--omni"]) + tile_args = parser.parse_args(["serve", "fake-model", "--omni", "--vae-parallel-mode", "tile"]) + + assert default_args.vae_parallel_mode == "auto" + assert tile_args.vae_parallel_mode == "tile" + + def _make_headless_args(**kwargs) -> TrackingNamespace: defaults = { "model": "fake-model", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 8997b12e1..01115c4b1 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -452,7 +452,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): cfg_parallel_size: int = Field(default=1, ge=1) vae_patch_parallel_size: int = Field(default=1, ge=1) text_encoder_tp_size: int = Field(default=1, ge=1) - vae_parallel_mode: str = "tile" + vae_parallel_mode: str = "auto" use_hsdp: bool = False mask_sp_padding: bool = False hsdp_shard_size: int = -1 @@ -466,9 +466,10 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1") if self.ulysses_mode not in {"strict", "advanced_uaa"}: raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'") - if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}: + if self.vae_parallel_mode not in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}: raise ValueError( - "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, " + "vae_parallel_mode must be one of " + "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, " f"but got {self.vae_parallel_mode!r}." ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 73051692d..561c0e87f 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -206,10 +206,13 @@ class DiffusionParallelConfig: process group. """ - vae_parallel_mode: str = "tile" + vae_parallel_mode: str = "auto" """VAE parallel decode strategy. - - "tile": Patch/tile parallel decode (default). Each rank decodes a subset + - "auto": Select the longer spatial-shard axis for validated Wan tiled + decodes on a full DiT-group topology; otherwise use patch/tile parallel + decode. + - "tile": Patch/tile parallel decode. Each rank decodes a subset of spatial tiles and the results are stitched on rank 0. - "spatial_shard_height": Spatially-sharded decode that splits decoder feature maps along height and exchanges halo rows around spatial @@ -250,8 +253,9 @@ class DiffusionParallelConfig: assert self.allgather_degree > 0, "AllGather degree must be > 0" assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0" assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0" - assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, ( - "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, " + assert self.vae_parallel_mode in {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}, ( + "vae_parallel_mode must be one of " + "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, " f"but got {self.vae_parallel_mode!r}." ) if self.allgather_degree > 1: diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 44c6efe66..4b8f0b965 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -23,6 +23,51 @@ 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) + + class OmniAutoencoderKLWan(AutoencoderKLWan): def _execution_context(self): try: @@ -270,20 +315,26 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): dec = torch.clamp(dec, min=-1.0, max=1.0) return dec - def _spatial_shard_decode_split_dim(self) -> str | None: + def _spatial_shard_decode_split_dim( + self, + 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 def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool: - split_dim = self._spatial_shard_decode_split_dim() - if split_dim is None: + if self.distributed_executor.parallel_mode == "tile": return False if z.ndim != 5: - logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.") + 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 @@ -292,23 +343,32 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): world_size = dist.get_world_size(group=group) requested_size = int(self.distributed_executor.parallel_size) if requested_size != world_size: - logger.warning( - "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size " - "to match the DIT group size; falling back to tiled decode. " - "requested=%s dit_group=%s split_dim=%s", - requested_size, - world_size, - split_dim, - ) + 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 True + return self._spatial_shard_decode_split_dim(z, world_size) is not None def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): if not self.is_distributed_enabled(): return super().tiled_decode(z, return_dict=return_dict) if self._spatial_shard_decode_enabled(z): - split_dim = self._spatial_shard_decode_split_dim() + world_size = dist.get_world_size(group=self.distributed_executor.group) + split_dim = self._spatial_shard_decode_split_dim(z, world_size) + assert split_dim is not None logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim) return wan_spatial_shard.spatial_shard_decode( self, diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py index bda728540..984aff203 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py @@ -52,6 +52,11 @@ _SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar( ) +def _active_spatial_shard_split_dim() -> str | None: + context = _SPATIAL_SHARD_CONTEXT.get() + return context.split_dim if context is not None else None + + def _spatial_dim(split_dim: str) -> int: if split_dim == "height": return -2 @@ -323,21 +328,27 @@ class WanDistZeroPad2d(nn.Module): padding: tuple[int, int, int, int], group: dist.ProcessGroup, *, - split_dim: str = "height", - split_padding: tuple[int, int] | None = None, + defer_split_after_padding: bool = False, ) -> None: super().__init__() self.padding = padding - self.split_dim = split_dim - default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]) - self.split_padding = split_padding or default_split_padding + self.defer_split_after_padding = defer_split_after_padding self.group = group def forward(self, x: torch.Tensor) -> torch.Tensor: + split_dim = _active_spatial_shard_split_dim() + if split_dim is None: + return F.pad(x, self.padding) + rank, world_size = _rank_world(self.group) left, right, top, bottom = self.padding + if self.defer_split_after_padding: + if split_dim == "height": + bottom = 0 + else: + right = 0 if world_size > 1: - if self.split_dim == "height": + if split_dim == "height": top = top if rank == 0 else 0 bottom = bottom if rank == world_size - 1 else 0 else: @@ -351,8 +362,7 @@ class WanDistConv2d(nn.Conv2d): self, source: nn.Conv2d, group: dist.ProcessGroup, - split_dim: str = "height", - split_padding: tuple[int, int] | None = None, + deferred_padding: tuple[int, int, int, int] | None = None, ): super().__init__( source.in_channels, @@ -369,75 +379,117 @@ class WanDistConv2d(nn.Conv2d): ) self.load_state_dict(source.state_dict()) self.group = group - self.split_dim = split_dim - self.split_tensor_dim = _spatial_dim(split_dim) - kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1] - self.halo_size = (kernel_extent - 1) // 2 - pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding - pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding - if split_dim == "height": - if split_padding is None: - split_padding = (pad_h, pad_h) - self._non_split_padding = (pad_w, pad_w, 0, 0) - self.kernel_extent = self.kernel_size[-2] - self.stride_extent = self.stride[-2] - else: - if split_padding is None: - split_padding = (pad_w, pad_w) - self._non_split_padding = (0, 0, pad_h, pad_h) - self.kernel_extent = self.kernel_size[-1] - self.stride_extent = self.stride[-1] - self.split_pad_left, self.split_pad_right = split_padding - self._halo_recv_top_buf: torch.Tensor | None = None - self._halo_recv_bottom_buf: torch.Tensor | None = None - self._trim_cache: dict[int, tuple[int, int, int]] = {} + self._direct_padding = source.padding + self._direct_padding_mode = source.padding_mode + self._direct_reversed_padding = source._reversed_padding_repeated_twice + self._deferred_padding = deferred_padding + source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding) + if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding): + raise ValueError(f"Wan spatial-shard Conv2d requires integer padding, got {source.padding!r}.") + self._spatial_pad_h = int(source_padding[-2]) + self._spatial_pad_w = int(source_padding[-1]) + self.register_buffer("_halo_recv_top_buf", None, persistent=False) + self.register_buffer("_halo_recv_bottom_buf", None, persistent=False) + self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {} def forward(self, x: torch.Tensor) -> torch.Tensor: - x = F.pad(x, self._non_split_padding) + split_dim = _active_spatial_shard_split_dim() + if split_dim is None: + if self._direct_padding_mode != "zeros": + x = F.pad(x, self._direct_reversed_padding, mode=self._direct_padding_mode) + padding: str | tuple[int, int] = (0, 0) + else: + padding = self._direct_padding + return F.conv2d(x, self.weight, self.bias, self.stride, padding, self.dilation, self.groups) + + split_tensor_dim = _spatial_dim(split_dim) + if split_dim == "height": + kernel_extent = self.kernel_size[-2] + stride_extent = self.stride[-2] + non_split_padding = (self._spatial_pad_w, self._spatial_pad_w, 0, 0) + default_split_padding = (self._spatial_pad_h, self._spatial_pad_h) + deferred_split_padding = ( + (self._deferred_padding[2], self._deferred_padding[3]) if self._deferred_padding is not None else None + ) + else: + kernel_extent = self.kernel_size[-1] + stride_extent = self.stride[-1] + non_split_padding = (0, 0, self._spatial_pad_h, self._spatial_pad_h) + default_split_padding = (self._spatial_pad_w, self._spatial_pad_w) + deferred_split_padding = ( + (self._deferred_padding[0], self._deferred_padding[1]) if self._deferred_padding is not None else None + ) + split_pad_left, split_pad_right = deferred_split_padding or default_split_padding + halo_size = (kernel_extent - 1) // 2 + + x = F.pad(x, non_split_padding) x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange( x, group=self.group, - halo_size=self.halo_size, - split_dim=self.split_dim, + halo_size=halo_size, + split_dim=split_dim, recv_top_buf=self._halo_recv_top_buf, recv_bottom_buf=self._halo_recv_bottom_buf, ) - shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim]) + shift, start, upper_bound = self._get_trim_params( + x.shape[split_tensor_dim], + split_dim=split_dim, + halo_size=halo_size, + split_pad_left=split_pad_left, + split_pad_right=split_pad_right, + kernel_extent=kernel_extent, + stride_extent=stride_extent, + ) if shift: x_padded = _narrow_along_dim( x_padded, - self.split_tensor_dim, + split_tensor_dim, shift, - x_padded.shape[self.split_tensor_dim] - shift, + x_padded.shape[split_tensor_dim] - shift, ) out = super().forward(x_padded) - out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim) - return _zero_invalid_extent(out, split_dim=self.split_dim) + out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim) + return _zero_invalid_extent(out, split_dim=split_dim) - def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]: - trim_params = self._trim_cache.get(local_extent) + def _get_trim_params( + self, + local_extent: int, + *, + split_dim: str, + halo_size: int, + split_pad_left: int, + split_pad_right: int, + kernel_extent: int, + stride_extent: int, + ) -> tuple[int, int, int]: + cache_key = (split_dim, local_extent) + trim_params = self._trim_cache.get(cache_key) if trim_params is None: rank, world_size = _rank_world(self.group) trim_params = _compute_conv_trim_params( local_extent=local_extent, rank=rank, world_size=world_size, - halo_size=self.halo_size, - pad_before=self.split_pad_left, - pad_after=self.split_pad_right, - kernel_extent=self.kernel_extent, - stride_extent=self.stride_extent, + halo_size=halo_size, + pad_before=split_pad_left, + pad_after=split_pad_right, + kernel_extent=kernel_extent, + stride_extent=stride_extent, ) - self._trim_cache[local_extent] = trim_params + self._trim_cache[cache_key] = trim_params return trim_params class WanDistCausalConv3d(nn.Conv3d): + # 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, - split_dim: str = "height", ): super().__init__( source.in_channels, @@ -454,88 +506,98 @@ class WanDistCausalConv3d(nn.Conv3d): ) self.load_state_dict(source.state_dict()) self.group = group - self.split_dim = split_dim - self.split_tensor_dim = _spatial_dim(split_dim) source_padding = getattr(source, "_padding", None) if source_padding is None: p_t, p_h, p_w = source.padding source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0) self._source_padding = tuple(source_padding) - if split_dim == "height": - self.split_pad_left = int(self._source_padding[2]) - self.split_pad_right = int(self._source_padding[3]) - self.kernel_extent = self.kernel_size[-2] - self.stride_extent = self.stride[-2] - padding = ( - self._source_padding[0], - self._source_padding[1], - 0, - 0, - self._source_padding[4], - self._source_padding[5], - ) - else: - self.split_pad_left = int(self._source_padding[0]) - self.split_pad_right = int(self._source_padding[1]) - self.kernel_extent = self.kernel_size[-1] - self.stride_extent = self.stride[-1] - padding = ( - 0, - 0, - self._source_padding[2], - self._source_padding[3], - self._source_padding[4], - self._source_padding[5], - ) - self.halo_size = (self.kernel_extent - 1) // 2 - self._padding = padding if self.halo_size > 0 else self._source_padding - self._halo_recv_top_buf: torch.Tensor | None = None - self._halo_recv_bottom_buf: torch.Tensor | None = None - self._trim_cache: dict[int, tuple[int, int, int]] = {} + self.register_buffer("_halo_recv_top_buf", None, persistent=False) + self.register_buffer("_halo_recv_bottom_buf", None, persistent=False) + self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {} def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: - padding = list(self._padding) + split_dim = _active_spatial_shard_split_dim() + padding = list(self._source_padding) if cache_x is not None and padding[4] > 0: cache_x = cache_x.to(x.device) x = torch.cat([cache_x, x], dim=2) padding[4] -= cache_x.shape[2] + if split_dim is None: + return super().forward(F.pad(x, padding)) + + split_tensor_dim = _spatial_dim(split_dim) + if split_dim == "height": + split_pad_left = int(padding[2]) + split_pad_right = int(padding[3]) + kernel_extent = self.kernel_size[-2] + stride_extent = self.stride[-2] + padding[2] = 0 + padding[3] = 0 + else: + split_pad_left = int(padding[0]) + split_pad_right = int(padding[1]) + kernel_extent = self.kernel_size[-1] + stride_extent = self.stride[-1] + padding[0] = 0 + padding[1] = 0 + halo_size = (kernel_extent - 1) // 2 + x = F.pad(x, padding) x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange( x, group=self.group, - halo_size=self.halo_size, - split_dim=self.split_dim, + halo_size=halo_size, + split_dim=split_dim, recv_top_buf=self._halo_recv_top_buf, recv_bottom_buf=self._halo_recv_bottom_buf, ) - shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim]) + shift, start, upper_bound = self._get_trim_params( + x.shape[split_tensor_dim], + split_dim=split_dim, + halo_size=halo_size, + split_pad_left=split_pad_left, + split_pad_right=split_pad_right, + kernel_extent=kernel_extent, + stride_extent=stride_extent, + ) if shift: x_padded = _narrow_along_dim( x_padded, - self.split_tensor_dim, + split_tensor_dim, shift, - x_padded.shape[self.split_tensor_dim] - shift, + x_padded.shape[split_tensor_dim] - shift, ) out = super().forward(x_padded) - out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim) - return _zero_invalid_extent(out, split_dim=self.split_dim) + out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim) + return _zero_invalid_extent(out, split_dim=split_dim) - def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]: - trim_params = self._trim_cache.get(local_extent) + def _get_trim_params( + self, + local_extent: int, + *, + split_dim: str, + halo_size: int, + split_pad_left: int, + split_pad_right: int, + kernel_extent: int, + stride_extent: int, + ) -> tuple[int, int, int]: + cache_key = (split_dim, local_extent) + trim_params = self._trim_cache.get(cache_key) if trim_params is None: rank, world_size = _rank_world(self.group) trim_params = _compute_conv_trim_params( local_extent=local_extent, rank=rank, world_size=world_size, - halo_size=self.halo_size, - pad_before=self.split_pad_left, - pad_after=self.split_pad_right, - kernel_extent=self.kernel_extent, - stride_extent=self.stride_extent, + halo_size=halo_size, + pad_before=split_pad_left, + pad_after=split_pad_right, + kernel_extent=kernel_extent, + stride_extent=stride_extent, ) - self._trim_cache[local_extent] = trim_params + self._trim_cache[cache_key] = trim_params return trim_params @@ -582,12 +644,15 @@ def _trim_local_conv_output( return out -def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None: +def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup) -> None: if getattr(module, "_vllm_omni_spatial_shard_attention", False): return orig_forward = module.forward def _forward(self: nn.Module, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + split_dim = _active_spatial_shard_split_dim() + if split_dim is None: + return orig_forward(x, *args, **kwargs) _, world_size = _rank_world(group) if world_size <= 1: return orig_forward(x, *args, **kwargs) @@ -609,7 +674,6 @@ def _replace_child( name: str, child: nn.Module, group: dist.ProcessGroup, - split_dim: str, ) -> None: if child.__class__.__name__ == "WanCausalConv3d": setattr( @@ -618,51 +682,39 @@ def _replace_child( WanDistCausalConv3d( child, group, - split_dim=split_dim, ), ) return if isinstance(child, nn.ZeroPad2d): padding = tuple(int(p) for p in child.padding) - module_padding = padding - if parent.__class__.__name__ == "Sequential": - # Let the following WanDistConv2d account for global after-edge - # padding; this module handles the non-split dimension and before edge. - if split_dim == "height": - module_padding = (padding[0], padding[1], padding[2], 0) - else: - module_padding = (padding[0], 0, padding[2], padding[3]) setattr( parent, name, WanDistZeroPad2d( - module_padding, + padding, group, - split_dim=split_dim, - split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]), + defer_split_after_padding=parent.__class__.__name__ == "Sequential", ), ) return if isinstance(child, nn.Conv2d): - split_padding = None + deferred_padding = None if name == "1" and parent.__class__.__name__ == "Sequential": # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a # stride-2 conv with padding=0. Only the last rank should see the # bottom/right global padding, which is approximated by split_padding. prev = getattr(parent, "0", None) if isinstance(prev, WanDistZeroPad2d): - split_padding = prev.split_padding + deferred_padding = prev.padding elif isinstance(prev, nn.ZeroPad2d): - pad = prev.padding - split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1])) + deferred_padding = tuple(int(value) for value in prev.padding) setattr( parent, name, WanDistConv2d( child, group, - split_dim=split_dim, - split_padding=split_padding, + deferred_padding=deferred_padding, ), ) @@ -670,21 +722,20 @@ def _replace_child( def _patch_decoder_modules( module: nn.Module, group: dist.ProcessGroup, - split_dim: str, inside_attention: bool = False, ) -> None: if module.__class__.__name__ == "WanAttentionBlock": - _patch_attention_block(module, group, split_dim) + _patch_attention_block(module, group) inside_attention = True for name, child in list(module.named_children()): if child.__class__.__name__ == "WanCausalConv3d": - _replace_child(module, name, child, group, split_dim) + _replace_child(module, name, child, group) continue if isinstance(child, nn.Conv2d) and not inside_attention: - _replace_child(module, name, child, group, split_dim) + _replace_child(module, name, child, group) continue - _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention) + _patch_decoder_modules(child, group, inside_attention=inside_attention) def _decoder_upsample_count(decoder: nn.Module) -> int: @@ -695,16 +746,25 @@ def _decoder_upsample_count(decoder: nn.Module) -> int: return count +def _clear_wan_spatial_shard_runtime_buffers(vae: Any) -> None: + decoder = getattr(vae, "decoder", None) + if decoder is None: + return + for module in decoder.modules(): + if "_halo_recv_top_buf" in module._buffers: + module._halo_recv_top_buf = None + if "_halo_recv_bottom_buf" in module._buffers: + module._halo_recv_bottom_buf = None + + def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None: """Patch ``vae.decoder`` once for spatially-sharded decode. This mutates the already-loaded decoder in place by swapping its spatial - convolutions/padding for halo-exchanging variants and wrapping - ``decoder.forward``. The patch is permanent for the lifetime of the VAE - instance and is applied only once (subsequent calls are no-ops). A given - instance is bound to a single ``split_dim``; switching between - ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here - otherwise. + convolutions/padding for context-aware variants and attaching a dedicated + spatial forward. Outside that forward the replacements execute their exact + local PyTorch behavior, so a later auto-selected tile request remains safe. + The same installed decoder can alternate height and width sharding. Only group-relative rank 0 assembles the final decoded frame, mirroring the distributed tiled-decode ``broadcast_result=False`` contract; the other ranks @@ -712,18 +772,12 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d """ _spatial_dim(split_dim) if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False): - installed_split_dim = getattr(vae, "_vllm_omni_wan_spatial_shard_split_dim", "height") - if installed_split_dim != split_dim: - raise ValueError( - "Wan spatial-shard VAE decoder was already patched for " - f"{installed_split_dim!r} split; create a fresh VAE instance to use {split_dim!r} split." - ) return decoder = getattr(vae, "decoder", None) if decoder is None: raise ValueError("Wan spatial-shard VAE decode requires a decoder module.") - _patch_decoder_modules(decoder, group, split_dim) + _patch_decoder_modules(decoder, group) upsample_count = _decoder_upsample_count(decoder) orig_forward = decoder.forward @@ -733,6 +787,8 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d 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] @@ -760,10 +816,9 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d _SPATIAL_SHARD_CONTEXT.reset(token) return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0) - decoder.forward = MethodType(_forward, decoder) + decoder._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder) vae._vllm_omni_wan_spatial_shard_installed = True - vae._vllm_omni_wan_spatial_shard_split_dim = split_dim - logger.info("Installed Wan VAE %s-sharded decode.", split_dim) + logger.info("Installed Wan VAE dynamic-axis spatial-shard decode.") def spatial_shard_decode( @@ -792,11 +847,12 @@ def spatial_shard_decode( decoded_chunks = [] for i in range(z.shape[2]): vae._conv_idx = [0] - chunk = vae.decoder( + 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 produce_output: decoded_chunks.append(chunk) @@ -810,6 +866,10 @@ def spatial_shard_decode( out = z.new_zeros(0) finally: vae.clear_cache() + # Halo buffers are request-shape scratch space allocated outside the + # tagged weight pool. Drop live references so sleep/offload never has + # to preserve or discard them as model state. + _clear_wan_spatial_shard_runtime_buffers(vae) if not return_dict: return (out,) diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index e28450c90..44b1064ff 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -155,6 +155,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): parallel_size: int, mode: str = "tile", ) -> None: + # ``auto`` is the cross-model default. 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. + if mode == "auto": + mode = "tile" if mode != "tile": raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}") group = get_dit_group() diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 52c2c646d..e905debbf 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -30,6 +30,9 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import ( from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + prepare_pipeline_wan_spatial_shard_decode, +) from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import ( @@ -243,6 +246,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): custom_pipeline_name=custom_pipeline_name, device=self.device, ) + prepare_pipeline_wan_spatial_shard_decode(self.pipeline) time_after_load = time.perf_counter() logger.info( diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 02ec3587b..43ce46694 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -505,7 +505,7 @@ class OrchestratorArgs: diffusion_kv_cache_skip_layers: str | None = None cfg_parallel_size: int = 1 vae_patch_parallel_size: int = 1 - vae_parallel_mode: str = "tile" + vae_parallel_mode: str = "auto" text_encoder_tp_size: int = 1 default_sampling_params: str | None = None max_generated_image_size: int | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index fc80cfa93..ad6b437e5 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -972,7 +972,7 @@ class AsyncOmniEngine: cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 - vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile" + vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "auto" text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1 enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False use_hsdp = normalized_kwargs.get("use_hsdp", False) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 7fb9cb8cd..bb3bc4456 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -750,10 +750,11 @@ class OmniServeCommand(CLISubcommand): omni_config_group.add_argument( "--vae-parallel-mode", type=str, - default="tile", - choices=["tile", "spatial_shard_height", "spatial_shard_width"], + default="auto", + choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"], help="VAE parallel decode strategy for diffusion models. " - "'tile' (default) uses patch/tile parallel decode; " + "'auto' (default) selects the longer spatial-shard axis for validated Wan tiled decodes " + "on a full DiT group, then falls back to patch/tile decode; 'tile' always uses patch/tile decode; " "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits " "decoder feature maps along height/width and exchanges halo regions. The " "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. "