From fd82e41fff8490278a551b9e67e8ba51ea88763e Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Tue, 11 Aug 2026 10:49:10 +0800 Subject: [PATCH] feat(diffusion): add QwenImage spatial-shard decode Signed-off-by: Hongsheng Liu --- docs/design/feature/vae_parallel.md | 17 +- .../diffusion/parallelism/vae_parallelism.md | 14 +- .../test_qwenimage_spatial_shard.py | 233 ++++++++++++++ .../diffusion/test_diffusion_model_runner.py | 2 +- vllm_omni/diffusion/data.py | 6 +- .../autoencoders/autoencoder_kl_qwenimage.py | 36 +++ .../autoencoders/autoencoder_kl_wan.py | 97 ++---- .../autoencoders/qwenimage_spatial_shard.py | 86 ++++++ .../distributed/autoencoders/spatial_shard.py | 142 +++++++++ .../autoencoders/wan_spatial_shard.py | 292 ++++++++++++------ vllm_omni/diffusion/models/minimax_h3/vae.py | 6 +- .../worker/diffusion_model_runner.py | 6 +- vllm_omni/entrypoints/cli/serve.py | 3 +- 13 files changed, 748 insertions(+), 192 deletions(-) create mode 100644 tests/diffusion/distributed/test_qwenimage_spatial_shard.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/qwenimage_spatial_shard.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/spatial_shard.py diff --git a/docs/design/feature/vae_parallel.md b/docs/design/feature/vae_parallel.md index 3fe165328..6072cb4b9 100644 --- a/docs/design/feature/vae_parallel.md +++ b/docs/design/feature/vae_parallel.md @@ -449,9 +449,9 @@ Complete examples in the codebase: --- -## Spatially-Sharded Decode (Wan) +## Spatially-Sharded Decode (Wan and QwenImage) -The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend. `DiffusionParallelConfig.vae_parallel_mode="auto"` is the default; explicit `"tile"`, `"spatial_shard_height"`, and `"spatial_shard_width"` modes remain available. +The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE and distributed QwenImage backend additionally support a **spatially-sharded decode** backend. `DiffusionParallelConfig.vae_parallel_mode="auto"` is the default; explicit `"tile"`, `"spatial_shard_height"`, and `"spatial_shard_width"` modes remain available. ### How it differs from tile parallel @@ -459,10 +459,11 @@ The tile-parallel executor above assigns independent spatial **tiles** to ranks. |--------|--------------------------|----------------------|-----------------------------------------------| | Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W | | Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) | -| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) | -| Scope | Decode + encode | Wan decode selection; encode remains tiled | Decode only | +| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0 and trim padding; QwenImage then broadcasts | +| Activation memory | Bounded by tile size | Depends on the selected path | Grows with full output area divided across ranks | +| Scope | Decode + encode | Wan/QwenImage decode selection; encode remains tiled | Decode only | -Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`. +Spatial-shard decode swaps the decoder's spatial convolutions/padding for context-aware halo-exchanging variants so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Wan and QwenImage share the padding, 2D convolution, attention, halo, and trim primitives in `wan_spatial_shard.py`; `qwenimage_spatial_shard.py` supplies a causal-convolution subclass that remains discoverable by Diffusers' QwenImage cache accounting. ### Wiring @@ -473,15 +474,17 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs -> DiffusionParallelConfig.vae_parallel_mode -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...) -> DistributedVaeExecutor.parallel_mode - -> DistributedAutoencoderKLWan.tiled_decode dispatch + -> capability-based pipeline preparation + -> DistributedAutoencoderKLWan/DistributedAutoencoderKLQwenImage.tiled_decode dispatch ``` -`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. The method is reached only after Diffusers has selected tiled decode for the request. In `auto` mode the longer axis that can cover every rank is selected, with width winning ties. This policy is cached only through normal kernel shape caches; it is reevaluated for every request. +The shared spatial policy gates the path: it requires distributed decode to be enabled, a 5D latent, a supported VAE capability, and `vae_patch_parallel_size == DiT group size`. The method is reached only after Diffusers has selected tiled decode for the request. In `auto` mode the longer axis that can cover every rank is selected, with width winning ties. This policy is cached only through normal kernel shape caches; it is reevaluated for every request. ### Notes - The decoder is patched **in place** the first time spatial-shard decode runs. Its wrappers consult a context-local split dimension and retain an exact direct fallback, allowing later requests to select tile, height, or width safely. - Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`). +- QwenImage additionally covers cache discovery, compile compatibility, mixed tile/spatial requests, both axes, and all-rank result broadcast in `tests/diffusion/distributed/test_qwenimage_spatial_shard.py`. --- diff --git a/docs/user_guide/diffusion/parallelism/vae_parallelism.md b/docs/user_guide/diffusion/parallelism/vae_parallelism.md index a45d69ec2..d5886c235 100644 --- a/docs/user_guide/diffusion/parallelism/vae_parallelism.md +++ b/docs/user_guide/diffusion/parallelism/vae_parallelism.md @@ -15,7 +15,7 @@ ## Overview -VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width. +VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan/QwenImage spatial-shard decode, which shards decoder feature maps along height or width. This is particularly useful for: - **High-resolution image generation** where VAE decode can become a memory bottleneck @@ -113,7 +113,7 @@ In `DiffusionParallelConfig`: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. | -| `vae_parallel_mode` | str | `"auto"` | VAE parallel decode strategy: `"auto"` (shape/model/topology-aware selection), `"tile"`, `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding is currently Wan-only. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). | +| `vae_parallel_mode` | str | `"auto"` | VAE parallel decode strategy: `"auto"` (shape/model/topology-aware selection), `"tile"`, `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding supports Wan and the distributed QwenImage backend. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan-and-qwenimage). | Additional requirements: @@ -126,9 +126,9 @@ Additional requirements: --- -## Spatially-Sharded Decode (Wan) +## Spatially-Sharded Decode (Wan and QwenImage) -The default `vae_parallel_mode="auto"` selects spatial-shard decode only after a Wan request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. It splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests, partial groups, and other VAE families keep their existing path. Set `vae_parallel_mode="tile"` to force the previous behavior, or `"spatial_shard_height"` / `"spatial_shard_width"` to force an axis. +The default `vae_parallel_mode="auto"` selects spatial-shard decode only after a Wan or distributed QwenImage request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. It splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests, partial groups, and other VAE families keep their existing path. Set `vae_parallel_mode="tile"` to force the previous behavior, or `"spatial_shard_height"` / `"spatial_shard_width"` to force an axis. Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance. @@ -157,11 +157,13 @@ vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \ **Constraints and behavior:** -- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes. +- Spatial-shard decode is **decode-only** and implemented for the **Wan** VAE and the distributed **QwenImage** backend used by Qwen-Image and Krea2. Pipelines using a different VAE backend keep their existing behavior. - It requires `vae_patch_parallel_size` to **match the DiT process group size**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning. - The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely. +- Wan preserves its rank-0-only decode result, while QwenImage preserves its existing all-rank result by broadcasting after rank-0 assembly. +- Spatial sharding trades activation memory for latency as resolution grows. Use `vae_parallel_mode="tile"` when bounded VAE memory is more important than decode speed. -The automatic policy is adapted from [SGLang #28071](https://github.com/sgl-project/sglang/pull/28071) and narrowed to vLLM-Omni's parity-tested Wan backend. +The automatic policy is adapted from [SGLang #28071](https://github.com/sgl-project/sglang/pull/28071) and narrowed to vLLM-Omni's parity-tested Wan and QwenImage backends. For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark: diff --git a/tests/diffusion/distributed/test_qwenimage_spatial_shard.py b/tests/diffusion/distributed/test_qwenimage_spatial_shard.py new file mode 100644 index 000000000..a1236b9ad --- /dev/null +++ b/tests/diffusion/distributed/test_qwenimage_spatial_shard.py @@ -0,0 +1,233 @@ +import os +from dataclasses import dataclass + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from diffusers.models.autoencoders.autoencoder_kl_qwenimage import QwenImageCausalConv3d + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import get_open_port +from vllm_omni.diffusion.distributed.autoencoders import qwenimage_spatial_shard +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import ( + DistributedAutoencoderKLQwenImage, +) +from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import ( + prepare_pipeline_spatial_shard_decode, +) +from vllm_omni.diffusion.offloader.module_collector import PipelineModules +from vllm_omni.platforms import current_omni_platform + + +@dataclass +class _FakeExecutor: + group: object + parallel_size: int + parallel_mode: str + + +def _tiny_qwenimage_vae() -> DistributedAutoencoderKLQwenImage: + return DistributedAutoencoderKLQwenImage( + base_dim=8, + z_dim=4, + dim_mult=[1, 2, 2, 2], + num_res_blocks=1, + attn_scales=[], + temperal_downsample=[False, True, True], + input_channels=3, + latents_mean=[0.0] * 4, + latents_std=[1.0] * 4, + ) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_qwenimage_auto_spatial_shard_gate_is_shape_and_topology_aware(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.autoencoders.spatial_shard.dist.get_world_size", + lambda group=None: 4, + ) + vae = DistributedAutoencoderKLQwenImage.__new__(DistributedAutoencoderKLQwenImage) + vae.use_tiling = True + vae.distributed_executor = _FakeExecutor(group=object(), parallel_size=4, parallel_mode="auto") + vae.is_distributed_enabled = lambda: True + + landscape = torch.zeros((1, 16, 1, 60, 104)) + portrait = torch.zeros((1, 16, 1, 104, 60)) + + assert vae._spatial_shard_decode_split_dim(landscape, 4) == "width" + assert vae._spatial_shard_decode_split_dim(portrait, 4) == "height" + assert vae._spatial_shard_decode_enabled(landscape) is True + + vae.distributed_executor.parallel_size = 2 + assert vae._spatial_shard_decode_enabled(landscape) is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_pipeline_dispatches_qwenimage_spatial_preparation(monkeypatch: pytest.MonkeyPatch): + vae = DistributedAutoencoderKLQwenImage.__new__(DistributedAutoencoderKLQwenImage) + prepared = [] + vae._prepare_spatial_shard_decode = lambda: prepared.append(vae) + + monkeypatch.setattr( + "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover", + lambda pipeline: PipelineModules( + dits=[], + dit_names=[], + encoders=[], + encoder_names=[], + vaes=[torch.nn.Identity(), vae], + ), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.autoencoders.spatial_shard.dist.is_initialized", + lambda: True, + ) + + prepare_pipeline_spatial_shard_decode(torch.nn.Identity()) + + assert prepared == [vae] + + +@pytest.mark.core_model +@pytest.mark.cpu +@torch.no_grad() +def test_qwenimage_installed_decoder_preserves_cache_discovery_state_and_direct_path(): + torch.manual_seed(0) + vae = _tiny_qwenimage_vae().eval() + latents = torch.randn(1, 4, 2, 4, 5) + parameter_names = {name for name, _ in vae.named_parameters()} + state_dict = {name: value.clone() for name, value in vae.state_dict().items()} + reference = vae.decode(latents, return_dict=False)[0] + causal_conv_count = sum(isinstance(module, QwenImageCausalConv3d) for module in vae.decoder.modules()) + + qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object(), split_dim="height") + qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object(), split_dim="width") + vae.clear_cache() + + assert causal_conv_count > 0 + assert len(vae._feat_map) == causal_conv_count + assert sum(isinstance(module, QwenImageCausalConv3d) for module in vae.decoder.modules()) == causal_conv_count + assert any( + isinstance(module, qwenimage_spatial_shard.QwenImageDistCausalConv3d) for module in vae.decoder.modules() + ) + assert {name for name, _ in vae.named_parameters()} == parameter_names + vae.load_state_dict(state_dict, strict=True) + assert torch.equal(vae.decode(latents, return_dict=False)[0], reference) + + +@pytest.mark.core_model +@pytest.mark.cpu +@torch.no_grad() +def test_qwenimage_installed_decoder_direct_path_can_be_compiled(): + torch.manual_seed(1) + vae = _tiny_qwenimage_vae().eval() + latents = torch.randn(1, 4, 1, 3, 5) + qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode(vae, group=object()) + eager = vae.decode(latents, return_dict=False)[0] + + compiled_decode = torch.compile(vae.decode, backend="eager", fullgraph=False) + compiled = compiled_decode(latents, return_dict=False)[0] + + torch.testing.assert_close(compiled, eager, rtol=0, atol=0) + + +def _qwenimage_spatial_shard_worker(rank: int, master_port: str, results) -> None: + from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, + ) + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = master_port + device = current_omni_platform.get_torch_device(rank) + current_omni_platform.set_device(device) + backend = current_omni_platform.dist_backend + init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend) + initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend) + + try: + torch.manual_seed(2) + vae = _tiny_qwenimage_vae().to(device=device, dtype=torch.float32).eval() + vae.init_distributed() + generator = torch.Generator(device=device).manual_seed(3) + landscape = torch.randn((1, 4, 2, 5, 7), generator=generator, device=device) + portrait = torch.randn((1, 4, 2, 7, 5), generator=generator, device=device) + + with torch.inference_mode(): + vae.use_tiling = False + vae.set_parallel_size(1, mode="tile") + landscape_reference = vae.decode(landscape, return_dict=False)[0] + portrait_reference = vae.decode(portrait, return_dict=False)[0] + + vae.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_stride_height=24, + tile_sample_stride_width=24, + ) + vae.set_parallel_size(2, mode="tile") + tile_before_install = vae.decode(landscape, return_dict=False)[0] + + vae.set_parallel_size(2, mode="auto") + assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width" + landscape_auto = vae.decode(landscape, return_dict=False)[0] + assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height" + portrait_auto = vae.decode(portrait).sample + + vae.set_parallel_size(2, mode="spatial_shard_height") + landscape_height = vae.decode(landscape, return_dict=False)[0] + vae.set_parallel_size(2, mode="spatial_shard_width") + portrait_width = vae.decode(portrait, return_dict=False)[0] + + vae.set_parallel_size(2, mode="tile") + tile_after_install = vae.decode(landscape, return_dict=False)[0] + + assert tile_after_install.shape == tile_before_install.shape + assert torch.equal(tile_after_install, tile_before_install) + assert all( + buffer is None + for module in vae.decoder.modules() + for name, buffer in module._buffers.items() + if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"} + ) + + diffs = { + "landscape_auto": (landscape_auto - landscape_reference).abs().max().item(), + "portrait_auto": (portrait_auto - portrait_reference).abs().max().item(), + "landscape_height": (landscape_height - landscape_reference).abs().max().item(), + "portrait_width": (portrait_width - portrait_reference).abs().max().item(), + } + results[rank] = { + "diffs": diffs, + "landscape_shape": tuple(landscape_auto.shape), + "portrait_shape": tuple(portrait_auto.shape), + } + finally: + destroy_model_parallel() + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=2) +def test_qwenimage_spatial_shard_matches_reference_broadcasts_and_switches_paths(): + manager = mp.get_context("spawn").Manager() + results = manager.dict() + mp.spawn( + _qwenimage_spatial_shard_worker, + args=(str(get_open_port()), results), + nprocs=2, + join=True, + ) + + assert set(results) == {0, 1} + assert results[0]["landscape_shape"] == results[1]["landscape_shape"] == (1, 3, 5, 40, 56) + assert results[0]["portrait_shape"] == results[1]["portrait_shape"] == (1, 3, 5, 56, 40) + for rank_results in results.values(): + assert max(rank_results["diffs"].values()) <= 2e-5 diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index be8a81331..905e71dfe 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -790,7 +790,7 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler) monkeypatch.setattr( model_runner_module, - "prepare_pipeline_wan_spatial_shard_decode", + "prepare_pipeline_spatial_shard_decode", lambda pipeline: preparation_pool_states.append(memory_pool_active), ) monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 561c0e87f..7f7cbef11 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -209,9 +209,9 @@ class DiffusionParallelConfig: vae_parallel_mode: str = "auto" """VAE parallel decode strategy. - - "auto": Select the longer spatial-shard axis for validated Wan tiled - decodes on a full DiT-group topology; otherwise use patch/tile parallel - decode. + - "auto": Select the longer spatial-shard axis for validated Wan and + distributed QwenImage tiled decodes on a full DiT-group topology; + otherwise use patch/tile parallel decode. - "tile": Patch/tile parallel decode. Each rank decodes a subset of spatial tiles and the results are stitched on rank 0. - "spatial_shard_height": Spatially-sharded decode that splits decoder diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py index f9dea8a36..684ba4727 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py @@ -4,16 +4,23 @@ from typing import Any import torch +import torch.distributed as dist from diffusers.models.autoencoders import AutoencoderKLQwenImage from diffusers.models.autoencoders.vae import DecoderOutput from vllm.logger import init_logger +from vllm_omni.diffusion.distributed.autoencoders import qwenimage_spatial_shard from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( DistributedOperator, DistributedVaeMixin, GridSpec, TileTask, ) +from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import ( + prepare_spatial_shard_decode, + resolve_spatial_shard_split_dim, + spatial_shard_decode_enabled, +) logger = init_logger(__name__) @@ -104,10 +111,39 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi ] return dec + def _spatial_shard_decode_split_dim( + self, + z: torch.Tensor | None = None, + world_size: int | None = None, + ) -> str | None: + return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size) + + def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool: + return spatial_shard_decode_enabled(self, z, model_name="QwenImage") + + def _prepare_spatial_shard_decode(self) -> None: + prepare_spatial_shard_decode( + self, + install=qwenimage_spatial_shard.install_qwenimage_spatial_shard_decode, + ) + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): if not self.is_distributed_enabled(): return super().tiled_decode(z, return_dict=return_dict) + if self._spatial_shard_decode_enabled(z): + world_size = dist.get_world_size(group=self.distributed_executor.group) + split_dim = self._spatial_shard_decode_split_dim(z, world_size) + assert split_dim is not None + logger.debug("Decode running with QwenImage VAE spatial_shard_%s mode", split_dim) + return qwenimage_spatial_shard.spatial_shard_decode( + self, + z, + group=self.distributed_executor.group, + return_dict=return_dict, + split_dim=split_dim, + ) + logger.debug("Decode running with distributed executor") result = self.distributed_executor.execute( z, diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 4b8f0b965..6b8b9c5d1 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -18,54 +18,23 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor GridSpec, TileTask, ) +from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import ( + prepare_pipeline_spatial_shard_decode, + prepare_spatial_shard_decode, + resolve_spatial_shard_split_dim, + spatial_shard_decode_enabled, +) +from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import ( + select_auto_spatial_shard_split_dim as select_auto_spatial_shard_split_dim, +) from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) -def select_auto_spatial_shard_split_dim( - z_shape: tuple[int, ...], - world_size: int, -) -> str | None: - """Select a profitable Wan spatial-shard axis for one latent shape. - - This selector runs only after Diffusers has already chosen ``tiled_decode`` - for a spatially large request. Splitting the longer eligible axis reduces - halo traffic for rectangular requests and improves per-rank balance. - """ - if len(z_shape) != 5 or world_size <= 1: - return None - - _, _, _, height, width = z_shape - candidates = [ - (extent, split_dim) for extent, split_dim in ((height, "height"), (width, "width")) if extent >= world_size - ] - if not candidates: - return None - # Prefer width for square latents: it was modestly faster in the existing - # four-GPU Wan benchmark and keeps the choice deterministic. - return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1] - - def prepare_pipeline_wan_spatial_shard_decode(pipeline: torch.nn.Module) -> None: - """Install eligible Wan wrappers while weight allocations are still tagged.""" - from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery - - if not dist.is_initialized(): - return - for vae in ModuleDiscovery.discover(pipeline).vaes: - if not isinstance(vae, DistributedAutoencoderKLWan): - continue - executor = vae.distributed_executor - mode = executor.parallel_mode - if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}: - continue - requested_size = int(executor.parallel_size) - world_size = dist.get_world_size(group=executor.group) - if requested_size <= 1 or requested_size != world_size: - continue - split_dim = "width" if mode == "spatial_shard_width" else "height" - wan_spatial_shard.install_wan_spatial_shard_decode(vae, executor.group, split_dim=split_dim) + """Backward-compatible alias for the model-agnostic pipeline dispatcher.""" + prepare_pipeline_spatial_shard_decode(pipeline) class OmniAutoencoderKLWan(AutoencoderKLWan): @@ -320,46 +289,16 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin): z: torch.Tensor | None = None, world_size: int | None = None, ) -> str | None: - mode = self.distributed_executor.parallel_mode - if mode == "spatial_shard_width": - return "width" - if mode == "spatial_shard_height": - return "height" - if mode == "auto" and z is not None and world_size is not None: - return select_auto_spatial_shard_split_dim(tuple(z.shape), world_size) - return None + return resolve_spatial_shard_split_dim(self.distributed_executor.parallel_mode, z, world_size) def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool: - if self.distributed_executor.parallel_mode == "tile": - return False - if z.ndim != 5: - if self.distributed_executor.parallel_mode != "tile": - logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.") - return False - if not self.is_distributed_enabled(): - return False + return spatial_shard_decode_enabled(self, z, model_name="Wan") - group = self.distributed_executor.group - world_size = dist.get_world_size(group=group) - requested_size = int(self.distributed_executor.parallel_size) - if requested_size != world_size: - if self.distributed_executor.parallel_mode == "auto": - logger.debug( - "Wan VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)", - requested_size, - world_size, - ) - elif self._spatial_shard_decode_split_dim() is not None: - logger.warning( - "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size " - "to match the DIT group size; falling back to tiled decode. " - "requested=%s dit_group=%s split_dim=%s", - requested_size, - world_size, - self._spatial_shard_decode_split_dim(), - ) - return False - return self._spatial_shard_decode_split_dim(z, world_size) is not None + def _prepare_spatial_shard_decode(self) -> None: + prepare_spatial_shard_decode( + self, + install=wan_spatial_shard.install_wan_spatial_shard_decode, + ) def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): if not self.is_distributed_enabled(): diff --git a/vllm_omni/diffusion/distributed/autoencoders/qwenimage_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/qwenimage_spatial_shard.py new file mode 100644 index 000000000..0e5364369 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/qwenimage_spatial_shard.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""QwenImage adapter for the shared spatial-shard VAE decoder kernels.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +import torch.nn as nn +from diffusers.models.autoencoders.autoencoder_kl_qwenimage import QwenImageCausalConv3d +from diffusers.models.autoencoders.vae import DecoderOutput + +from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE +from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import ( + SpatialShardCausalConv3dMixin, + install_spatial_shard_decode, + spatial_shard_decode_impl, +) + + +class QwenImageDistCausalConv3d(SpatialShardCausalConv3dMixin, QwenImageCausalConv3d): + """Context-aware QwenImage causal conv that remains cache-discoverable.""" + + def __init__(self, source: nn.Conv3d, group: dist.ProcessGroup) -> None: + # Initialize Conv3d directly so the replacement preserves source device, + # dtype, groups, and bias while remaining an isinstance() match for + # QwenImageCausalConv3d. Diffusers clear_cache() relies on that match. + nn.Conv3d.__init__( + self, + source.in_channels, + source.out_channels, + source.kernel_size, + stride=source.stride, + padding=0, + dilation=source.dilation, + groups=source.groups, + bias=source.bias is not None, + padding_mode=source.padding_mode, + device=source.weight.device, + dtype=source.weight.dtype, + ) + self._init_spatial_shard_causal_conv(source, group) + + def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + return self._spatial_shard_causal_conv_forward(x, cache_x) + + +def install_qwenimage_spatial_shard_decode( + vae: SpatialShardVAE, + group: dist.ProcessGroup, + split_dim: str = "height", +) -> None: + install_spatial_shard_decode( + vae, + group, + split_dim, + causal_conv_class_name="QwenImageCausalConv3d", + causal_conv_factory=QwenImageDistCausalConv3d, + attention_block_class_name="QwenImageAttentionBlock", + installed_attr="_vllm_omni_qwenimage_spatial_shard_installed", + model_name="QwenImage", + ) + + +def spatial_shard_decode( + vae: SpatialShardVAE, + z: torch.Tensor, + *, + group: dist.ProcessGroup, + return_dict: bool = True, + split_dim: str = "height", +) -> DecoderOutput | tuple[torch.Tensor]: + """Decode QwenImage latents and preserve its all-rank output contract.""" + return spatial_shard_decode_impl( + vae, + z, + group=group, + install=install_qwenimage_spatial_shard_decode, + model_name="QwenImage", + pass_first_chunk=False, + broadcast_result=True, + unpatchify_patch_size=None, + return_dict=return_dict, + split_dim=split_dim, + ) diff --git a/vllm_omni/diffusion/distributed/autoencoders/spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/spatial_shard.py new file mode 100644 index 000000000..fbee19444 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/spatial_shard.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Callable +from typing import Protocol + +import torch +import torch.distributed as dist +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +class SpatialShardExecutor(Protocol): + group: dist.ProcessGroup + parallel_mode: str + parallel_size: int + + def _sync_final_result( + self, + rank0_result: torch.Tensor, + output_ndim: int, + output_device: torch.device, + output_dtype: torch.dtype, + ) -> torch.Tensor: ... + + +class SpatialShardVAE(Protocol): + distributed_executor: SpatialShardExecutor + decoder: torch.nn.Module + post_quant_conv: torch.nn.Module + dtype: torch.dtype + _conv_idx: list[int] + _feat_map: list[object] + + def is_distributed_enabled(self) -> bool: ... + + def clear_cache(self) -> None: ... + + +def select_auto_spatial_shard_split_dim( + z_shape: tuple[int, ...], + world_size: int, +) -> str | None: + """Select the longer spatial axis that can be split across every rank.""" + if len(z_shape) != 5 or world_size <= 1: + return None + + _, _, _, height, width = z_shape + candidates = [ + (extent, split_dim) for extent, split_dim in ((height, "height"), (width, "width")) if extent >= world_size + ] + if not candidates: + return None + # Prefer width for square latents. It was modestly faster in the original + # four-GPU Wan benchmark and keeps the choice deterministic across models. + return max(candidates, key=lambda candidate: (candidate[0], candidate[1] == "width"))[1] + + +def resolve_spatial_shard_split_dim( + mode: str, + z: torch.Tensor | None = None, + world_size: int | None = None, +) -> str | None: + if mode == "spatial_shard_width": + return "width" + if mode == "spatial_shard_height": + return "height" + if mode == "auto" and z is not None and world_size is not None: + return select_auto_spatial_shard_split_dim(tuple(z.shape), world_size) + return None + + +def spatial_shard_decode_enabled(vae: SpatialShardVAE, z: torch.Tensor, *, model_name: str) -> bool: + """Apply the shared request- and topology-level spatial decode policy.""" + executor = vae.distributed_executor + if executor.parallel_mode == "tile": + return False + if z.ndim != 5: + logger.warning( + "%s VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.", model_name + ) + return False + if not vae.is_distributed_enabled(): + return False + + group = executor.group + world_size = dist.get_world_size(group=group) + requested_size = int(executor.parallel_size) + requested_split_dim = resolve_spatial_shard_split_dim(executor.parallel_mode) + if requested_size != world_size: + if executor.parallel_mode == "auto": + logger.debug( + "%s VAE auto decode selected tile mode for a partial DiT group (requested=%s dit_group=%s)", + model_name, + requested_size, + world_size, + ) + elif requested_split_dim is not None: + logger.warning( + "%s VAE spatial sharded decode currently requires vae_patch_parallel_size " + "to match the DIT group size; falling back to tiled decode. " + "requested=%s dit_group=%s split_dim=%s", + model_name, + requested_size, + world_size, + requested_split_dim, + ) + return False + return resolve_spatial_shard_split_dim(executor.parallel_mode, z, world_size) is not None + + +def prepare_spatial_shard_decode( + vae: SpatialShardVAE, + *, + install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None], +) -> None: + """Install one VAE's dynamic wrappers while allocations are still tagged.""" + executor = vae.distributed_executor + mode = executor.parallel_mode + if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}: + return + requested_size = int(executor.parallel_size) + world_size = dist.get_world_size(group=executor.group) + if requested_size <= 1 or requested_size != world_size: + return + split_dim = "width" if mode == "spatial_shard_width" else "height" + install(vae, executor.group, split_dim) + + +def prepare_pipeline_spatial_shard_decode(pipeline: torch.nn.Module) -> None: + """Prepare every VAE that advertises spatial-shard decode capability.""" + from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + + if not dist.is_initialized(): + return + for vae in ModuleDiscovery.discover(pipeline).vaes: + prepare = getattr(vae, "_prepare_spatial_shard_decode", None) + if callable(prepare): + prepare() diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py index 984aff203..bb9562173 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py @@ -7,19 +7,21 @@ # (python/sglang/multimodal_gen/runtime/layers/parallel_conv.py) # which is in turn adapted from FastVideo (https://github.com/hao-ai-lab/FastVideo). # This version generalizes the height-only sharding to shard along height or -# width and adds the Wan causal-conv ``feat_cache`` handling. -"""Spatially-sharded Wan VAE decode. +# width and supports the causal-conv ``feat_cache`` handling used by Wan and +# QwenImage. +"""Shared spatial-shard kernels plus the Wan VAE adapter. -The existing distributed Wan VAE path shards *tiles*. This module adds an -opt-in decode backend that shards decoder feature maps along height or width and -exchanges boundary rows/columns before spatial convolutions. It is -intentionally decode-only and keeps checkpoint loading unchanged by patching the -already-loaded decoder. +The halo, padding, stride-trim, attention-gather, and decode-lifecycle helpers +are also reused by QwenImage's thin adapter. The backend shards decoder feature +maps along height or width and exchanges boundary rows/columns before spatial +convolutions. It is decode-only and keeps checkpoint loading unchanged by +patching the already-loaded decoder. """ from __future__ import annotations import math +from collections.abc import Callable from contextlib import nullcontext from contextvars import ContextVar from dataclasses import dataclass @@ -34,6 +36,8 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify from diffusers.models.autoencoders.vae import DecoderOutput from vllm.logger import init_logger +from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import SpatialShardVAE + logger = init_logger(__name__) @@ -62,7 +66,7 @@ def _spatial_dim(split_dim: str) -> int: return -2 if split_dim == "width": return -1 - raise ValueError(f"Unsupported Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.") + raise ValueError(f"Unsupported VAE split_dim={split_dim!r}; expected 'height' or 'width'.") def _narrow_along_dim(x: torch.Tensor, dim: int, start: int, length: int) -> torch.Tensor: @@ -160,9 +164,9 @@ def split_for_parallel_decode( rank = 0 if rank is None else int(rank) world_size = 1 if world_size is None else int(world_size) if world_size < 1: - raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.") + raise ValueError(f"VAE world_size must be >= 1, got {world_size}.") if not 0 <= rank < world_size: - raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.") + raise ValueError(f"VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.") dim = _spatial_dim(split_dim) expected_extent = x.shape[dim] * (2**upsample_count) @@ -385,7 +389,7 @@ class WanDistConv2d(nn.Conv2d): self._deferred_padding = deferred_padding source_padding = source.padding if isinstance(source.padding, tuple) else (source.padding, source.padding) if len(source_padding) != 2 or not all(isinstance(value, int) for value in source_padding): - raise ValueError(f"Wan spatial-shard Conv2d requires integer padding, got {source.padding!r}.") + raise ValueError(f"Spatial-shard Conv2d requires integer padding, got {source.padding!r}.") self._spatial_pad_h = int(source_padding[-2]) self._spatial_pad_w = int(source_padding[-1]) self.register_buffer("_halo_recv_top_buf", None, persistent=False) @@ -480,30 +484,20 @@ class WanDistConv2d(nn.Conv2d): return trim_params -class WanDistCausalConv3d(nn.Conv3d): - # Cross-PR compatibility contract for the lossless Wan data-movement - # installer: it may fuse the no-context direct path, but must leave this - # wrapper in control whenever a spatial-shard context is active. - _vllm_omni_dynamic_spatial_shard_conv = True +# Model-neutral aliases let other causal VAEs reuse the implementation while +# preserving the original Wan class identities for stacked PRs and downstreams. +SpatialShardZeroPad2d = WanDistZeroPad2d +SpatialShardConv2d = WanDistConv2d - def __init__( + +class SpatialShardCausalConv3dMixin: + """Shared causal-convolution behavior for supported Diffusers VAE classes.""" + + def _init_spatial_shard_causal_conv( self, source: nn.Conv3d, group: dist.ProcessGroup, - ): - super().__init__( - source.in_channels, - source.out_channels, - source.kernel_size, - stride=source.stride, - padding=0, - dilation=source.dilation, - groups=source.groups, - bias=source.bias is not None, - padding_mode=source.padding_mode, - device=source.weight.device, - dtype=source.weight.dtype, - ) + ) -> None: self.load_state_dict(source.state_dict()) self.group = group source_padding = getattr(source, "_padding", None) @@ -511,11 +505,16 @@ class WanDistCausalConv3d(nn.Conv3d): p_t, p_h, p_w = source.padding source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0) self._source_padding = tuple(source_padding) + self._padding = tuple(source_padding) self.register_buffer("_halo_recv_top_buf", None, persistent=False) self.register_buffer("_halo_recv_bottom_buf", None, persistent=False) self._trim_cache: dict[tuple[str, int], tuple[int, int, int]] = {} - def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + def _spatial_shard_causal_conv_forward( + self, + x: torch.Tensor, + cache_x: torch.Tensor | None = None, + ) -> torch.Tensor: split_dim = _active_spatial_shard_split_dim() padding = list(self._source_padding) if cache_x is not None and padding[4] > 0: @@ -524,7 +523,7 @@ class WanDistCausalConv3d(nn.Conv3d): padding[4] -= cache_x.shape[2] if split_dim is None: - return super().forward(F.pad(x, padding)) + return nn.Conv3d.forward(self, F.pad(x, padding)) split_tensor_dim = _spatial_dim(split_dim) if split_dim == "height": @@ -568,7 +567,7 @@ class WanDistCausalConv3d(nn.Conv3d): shift, x_padded.shape[split_tensor_dim] - shift, ) - out = super().forward(x_padded) + out = nn.Conv3d.forward(self, x_padded) out = _trim_local_conv_output(out, halo_size, start, upper_bound, split_dim=split_dim) return _zero_invalid_extent(out, split_dim=split_dim) @@ -601,6 +600,36 @@ class WanDistCausalConv3d(nn.Conv3d): return trim_params +class WanDistCausalConv3d(SpatialShardCausalConv3dMixin, nn.Conv3d): + # Cross-PR compatibility contract for the lossless Wan data-movement + # installer: it may fuse the no-context direct path, but must leave this + # wrapper in control whenever a spatial-shard context is active. + _vllm_omni_dynamic_spatial_shard_conv = True + + def __init__( + self, + source: nn.Conv3d, + group: dist.ProcessGroup, + ): + super().__init__( + source.in_channels, + source.out_channels, + source.kernel_size, + stride=source.stride, + padding=0, + dilation=source.dilation, + groups=source.groups, + bias=source.bias is not None, + padding_mode=source.padding_mode, + device=source.weight.device, + dtype=source.weight.dtype, + ) + self._init_spatial_shard_causal_conv(source, group) + + def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + return self._spatial_shard_causal_conv_forward(x, cache_x) + + def _compute_conv_trim_params( *, local_extent: int, @@ -674,23 +703,19 @@ def _replace_child( name: str, child: nn.Module, group: dist.ProcessGroup, + *, + causal_conv_class_name: str, + causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module], ) -> None: - if child.__class__.__name__ == "WanCausalConv3d": - setattr( - parent, - name, - WanDistCausalConv3d( - child, - group, - ), - ) + if child.__class__.__name__ == causal_conv_class_name: + setattr(parent, name, causal_conv_factory(child, group)) return if isinstance(child, nn.ZeroPad2d): padding = tuple(int(p) for p in child.padding) setattr( parent, name, - WanDistZeroPad2d( + SpatialShardZeroPad2d( padding, group, defer_split_after_padding=parent.__class__.__name__ == "Sequential", @@ -700,18 +725,18 @@ def _replace_child( if isinstance(child, nn.Conv2d): deferred_padding = None if name == "1" and parent.__class__.__name__ == "Sequential": - # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a + # Wan/QwenImage downsampling uses ZeroPad2d((0, 1, 0, 1)) before a # stride-2 conv with padding=0. Only the last rank should see the # bottom/right global padding, which is approximated by split_padding. prev = getattr(parent, "0", None) - if isinstance(prev, WanDistZeroPad2d): + if isinstance(prev, SpatialShardZeroPad2d): deferred_padding = prev.padding elif isinstance(prev, nn.ZeroPad2d): deferred_padding = tuple(int(value) for value in prev.padding) setattr( parent, name, - WanDistConv2d( + SpatialShardConv2d( child, group, deferred_padding=deferred_padding, @@ -722,20 +747,45 @@ def _replace_child( def _patch_decoder_modules( module: nn.Module, group: dist.ProcessGroup, + *, + causal_conv_class_name: str, + causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module], + attention_block_class_name: str, inside_attention: bool = False, ) -> None: - if module.__class__.__name__ == "WanAttentionBlock": + if module.__class__.__name__ == attention_block_class_name: _patch_attention_block(module, group) inside_attention = True for name, child in list(module.named_children()): - if child.__class__.__name__ == "WanCausalConv3d": - _replace_child(module, name, child, group) + if child.__class__.__name__ == causal_conv_class_name: + _replace_child( + module, + name, + child, + group, + causal_conv_class_name=causal_conv_class_name, + causal_conv_factory=causal_conv_factory, + ) continue if isinstance(child, nn.Conv2d) and not inside_attention: - _replace_child(module, name, child, group) + _replace_child( + module, + name, + child, + group, + causal_conv_class_name=causal_conv_class_name, + causal_conv_factory=causal_conv_factory, + ) continue - _patch_decoder_modules(child, group, inside_attention=inside_attention) + _patch_decoder_modules( + child, + group, + causal_conv_class_name=causal_conv_class_name, + causal_conv_factory=causal_conv_factory, + attention_block_class_name=attention_block_class_name, + inside_attention=inside_attention, + ) def _decoder_upsample_count(decoder: nn.Module) -> int: @@ -746,7 +796,7 @@ def _decoder_upsample_count(decoder: nn.Module) -> int: return count -def _clear_wan_spatial_shard_runtime_buffers(vae: Any) -> None: +def clear_spatial_shard_runtime_buffers(vae: SpatialShardVAE) -> None: decoder = getattr(vae, "decoder", None) if decoder is None: return @@ -757,7 +807,17 @@ def _clear_wan_spatial_shard_runtime_buffers(vae: Any) -> None: module._halo_recv_bottom_buf = None -def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None: +def install_spatial_shard_decode( + vae: SpatialShardVAE, + group: dist.ProcessGroup, + split_dim: str, + *, + causal_conv_class_name: str, + causal_conv_factory: Callable[[nn.Conv3d, dist.ProcessGroup], nn.Module], + attention_block_class_name: str, + installed_attr: str, + model_name: str, +) -> None: """Patch ``vae.decoder`` once for spatially-sharded decode. This mutates the already-loaded decoder in place by swapping its spatial @@ -766,32 +826,28 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d local PyTorch behavior, so a later auto-selected tile request remains safe. The same installed decoder can alternate height and width sharding. - Only group-relative rank 0 assembles the final decoded frame, mirroring the - distributed tiled-decode ``broadcast_result=False`` contract; the other ranks - take part in the collectives but return an empty placeholder. + Group-relative rank 0 assembles the final decoded frame. The model adapter + then preserves its prior contract: Wan keeps the result on rank 0, while + QwenImage broadcasts it to every rank. """ _spatial_dim(split_dim) - if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False): + if getattr(vae, installed_attr, False): return decoder = getattr(vae, "decoder", None) if decoder is None: - raise ValueError("Wan spatial-shard VAE decode requires a decoder module.") + raise ValueError(f"{model_name} spatial-shard VAE decode requires a decoder module.") - _patch_decoder_modules(decoder, group) + _patch_decoder_modules( + decoder, + group, + causal_conv_class_name=causal_conv_class_name, + causal_conv_factory=causal_conv_factory, + attention_block_class_name=attention_block_class_name, + ) upsample_count = _decoder_upsample_count(decoder) orig_forward = decoder.forward - def _forward( - self: nn.Module, - x: torch.Tensor, - feat_cache: list[torch.Tensor] | None = None, - feat_idx: list[int] | None = None, - first_chunk: bool = False, - *, - split_dim: str, - ) -> torch.Tensor: - if feat_idx is None: - feat_idx = [0] + def _forward(self: nn.Module, x: torch.Tensor, *args: object, split_dim: str, **kwargs: object) -> torch.Tensor: tensor_dim = _spatial_dim(split_dim) input_extent = x.shape[tensor_dim] x, expected_extent = split_for_parallel_decode( @@ -811,28 +867,50 @@ def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_d ) ) try: - out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) + out = orig_forward(x, *args, **kwargs) finally: _SPATIAL_SHARD_CONTEXT.reset(token) return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0) decoder._vllm_omni_spatial_shard_forward = MethodType(_forward, decoder) - vae._vllm_omni_wan_spatial_shard_installed = True - logger.info("Installed Wan VAE dynamic-axis spatial-shard decode.") + setattr(vae, installed_attr, True) + logger.info("Installed %s VAE dynamic-axis spatial-shard decode.", model_name) -def spatial_shard_decode( - vae: Any, +def install_wan_spatial_shard_decode( + vae: SpatialShardVAE, + group: dist.ProcessGroup, + split_dim: str = "height", +) -> None: + install_spatial_shard_decode( + vae, + group, + split_dim, + causal_conv_class_name="WanCausalConv3d", + causal_conv_factory=WanDistCausalConv3d, + attention_block_class_name="WanAttentionBlock", + installed_attr="_vllm_omni_wan_spatial_shard_installed", + model_name="Wan", + ) + + +def spatial_shard_decode_impl( + vae: SpatialShardVAE, z: torch.Tensor, *, group: dist.ProcessGroup, + install: Callable[[SpatialShardVAE, dist.ProcessGroup, str], None], + model_name: str, + pass_first_chunk: bool, + broadcast_result: bool, + unpatchify_patch_size: int | None, return_dict: bool = True, split_dim: str = "height", ) -> DecoderOutput | tuple[torch.Tensor]: - install_wan_spatial_shard_decode(vae, group, split_dim=split_dim) + install(vae, group, split_dim) if z.shape[2] == 0: - raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.") + raise ValueError(f"{model_name} spatial-shard VAE decode expects at least one latent frame.") # Non-rank-0 ranks must still run the decoder every chunk to stay in lockstep with # the halo/all-gather collectives; they just skip keeping/assembling the output. @@ -841,26 +919,36 @@ def spatial_shard_decode( vae.clear_cache() try: - context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext() + context_factory = getattr(vae, "_execution_context", None) + context = context_factory() if callable(context_factory) else nullcontext() with context: x = vae.post_quant_conv(z) decoded_chunks = [] + spatial_forward = getattr(vae.decoder, "_vllm_omni_spatial_shard_forward") for i in range(z.shape[2]): vae._conv_idx = [0] - chunk = vae.decoder._vllm_omni_spatial_shard_forward( - x[:, :, i : i + 1, :, :], - feat_cache=vae._feat_map, - feat_idx=vae._conv_idx, - first_chunk=(i == 0), - split_dim=split_dim, - ) + if pass_first_chunk: + chunk = spatial_forward( + x[:, :, i : i + 1, :, :], + feat_cache=vae._feat_map, + feat_idx=vae._conv_idx, + first_chunk=i == 0, + split_dim=split_dim, + ) + else: + chunk = spatial_forward( + x[:, :, i : i + 1, :, :], + feat_cache=vae._feat_map, + feat_idx=vae._conv_idx, + split_dim=split_dim, + ) if produce_output: decoded_chunks.append(chunk) if produce_output: out = torch.cat(decoded_chunks, dim=2) - if vae.config.patch_size is not None: - out = unpatchify(out, patch_size=vae.config.patch_size) + if unpatchify_patch_size is not None: + out = unpatchify(out, patch_size=unpatchify_patch_size) out = torch.clamp(out, min=-1.0, max=1.0) else: out = z.new_zeros(0) @@ -869,8 +957,34 @@ def spatial_shard_decode( # Halo buffers are request-shape scratch space allocated outside the # tagged weight pool. Drop live references so sleep/offload never has # to preserve or discard them as model state. - _clear_wan_spatial_shard_runtime_buffers(vae) + clear_spatial_shard_runtime_buffers(vae) + + if broadcast_result: + out = vae.distributed_executor._sync_final_result(out, z.ndim, z.device, vae.dtype) if not return_dict: return (out,) return DecoderOutput(sample=out) + + +def spatial_shard_decode( + vae: SpatialShardVAE, + z: torch.Tensor, + *, + group: dist.ProcessGroup, + return_dict: bool = True, + split_dim: str = "height", +) -> DecoderOutput | tuple[torch.Tensor]: + """Decode a Wan latent with spatial activation sharding.""" + return spatial_shard_decode_impl( + vae, + z, + group=group, + install=install_wan_spatial_shard_decode, + model_name="Wan", + pass_first_chunk=True, + broadcast_result=False, + unpatchify_patch_size=vae.config.patch_size, + return_dict=return_dict, + split_dim=split_dim, + ) diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index 44b1064ff..0cfec47b6 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -155,9 +155,9 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): parallel_size: int, mode: str = "tile", ) -> None: - # ``auto`` is the cross-model default. Spatial selection is Wan-only, - # so preserve MiniMax H3's native tile path for that default while - # continuing to reject explicitly requested Wan spatial modes. + # ``auto`` is the cross-model default. MiniMax H3 has no spatial-shard + # capability, so preserve its native tile path for that default while + # continuing to reject explicitly requested spatial modes. if mode == "auto": mode = "tile" if mode != "tile": diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index e905debbf..ec4368a52 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -30,8 +30,8 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import ( from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( - prepare_pipeline_wan_spatial_shard_decode, +from vllm_omni.diffusion.distributed.autoencoders.spatial_shard import ( + prepare_pipeline_spatial_shard_decode, ) from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -246,7 +246,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): custom_pipeline_name=custom_pipeline_name, device=self.device, ) - prepare_pipeline_wan_spatial_shard_decode(self.pipeline) + prepare_pipeline_spatial_shard_decode(self.pipeline) time_after_load = time.perf_counter() logger.info( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index bb3bc4456..ede3997aa 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -753,7 +753,8 @@ class OmniServeCommand(CLISubcommand): default="auto", choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"], help="VAE parallel decode strategy for diffusion models. " - "'auto' (default) selects the longer spatial-shard axis for validated Wan tiled decodes " + "'auto' (default) selects the longer spatial-shard axis for validated Wan and distributed QwenImage " + "tiled decodes " "on a full DiT group, then falls back to patch/tile decode; 'tile' always uses patch/tile decode; " "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits " "decoder feature maps along height/width and exchanges halo regions. The "