fix(diffusion): address Wan auto decode review feedback

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

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -11,7 +11,6 @@ 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
@@ -352,7 +351,7 @@ def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.Mon
         lambda x, *, group, dim, dst=None: torch.cat([x, x], dim=dim),
     )
     module = ShrinkingWanAttentionBlock()
-    wan_spatial_shard._patch_attention_block(module, group=object(), split_dim="width")
+    wan_spatial_shard._patch_attention_block(module, group=object())
     x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32)
     token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
         wan_spatial_shard.SpatialShardContext(
@@ -490,23 +489,18 @@ def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyP
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch):
+def test_set_parallel_size_prepares_spatial_wrappers_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"),
-    )
+    executor = SimpleNamespace(group=object(), parallel_size=1, parallel_mode="tile")
+
+    def set_parallel_size(parallel_size: int, mode: str = "tile") -> None:
+        executor.parallel_size = parallel_size
+        executor.parallel_mode = mode
+
+    executor.set_parallel_size = set_parallel_size
+    object.__setattr__(vae, "distributed_executor", executor)
     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,
@@ -517,11 +511,10 @@ def test_spatial_wrappers_are_prepared_only_for_a_full_group(monkeypatch: pytest
         lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)),
     )
 
-    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
+    vae.set_parallel_size(2, mode="auto")
     assert installed == []
 
-    vae.distributed_executor.parallel_size = 4
-    prepare_pipeline_wan_spatial_shard_decode(torch.nn.Identity())
+    vae.set_parallel_size(4, mode="auto")
     assert installed == [(vae, vae.distributed_executor.group, "height")]
 
 

tests/diffusion/test_diffusion_model_runner.py

@@ -732,19 +732,6 @@ 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
@@ -792,22 +779,16 @@ 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, memory_pool_context_fn=_memory_pool_context)
+    DiffusionModelRunner.load_model(runner)
 
     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

vllm_omni/config/omni_config.py

@@ -471,10 +471,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 {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}:
+        if self.vae_parallel_mode not in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}:
             raise ValueError(
                 "vae_parallel_mode must be one of "
-                "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+                "{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
                 f"but got {self.vae_parallel_mode!r}."
             )
 

vllm_omni/diffusion/data.py

@@ -210,11 +210,11 @@ class DiffusionParallelConfig:
     vae_parallel_mode: str = "tile"
     """VAE parallel decode strategy.
 
+    - "tile": Patch/tile parallel decode (default). Each rank decodes a subset
+      of spatial tiles and the results are stitched on rank 0.
     - "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 (default). 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
       convolutions.
@@ -254,9 +254,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 {"auto", "tile", "spatial_shard_height", "spatial_shard_width"}, (
+        assert self.vae_parallel_mode in {"tile", "auto", "spatial_shard_height", "spatial_shard_width"}, (
             "vae_parallel_mode must be one of "
-            "{'auto', 'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
+            "{'tile', 'auto', 'spatial_shard_height', 'spatial_shard_width'}, "
             f"but got {self.vae_parallel_mode!r}."
         )
         if self.allgather_degree > 1:

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -47,27 +47,6 @@ def select_auto_spatial_shard_split_dim(
     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:
@@ -101,6 +80,19 @@ class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
         model.init_distributed()
         return model
 
+    def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
+        super().set_parallel_size(parallel_size, mode=mode)
+        if mode not in {"auto", "spatial_shard_height", "spatial_shard_width"}:
+            return
+
+        executor = self.distributed_executor
+        world_size = dist.get_world_size(group=executor.group)
+        if parallel_size <= 1 or parallel_size != world_size:
+            return
+
+        split_dim = "width" if mode == "spatial_shard_width" else "height"
+        wan_spatial_shard.install_wan_spatial_shard_decode(self, executor.group, split_dim=split_dim)
+
     def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
         _, _, num_frames, height, width = z.shape
         sample_height = height * self.spatial_compression_ratio

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -32,9 +32,6 @@ from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
-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 (
@@ -269,7 +266,6 @@ 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(

vllm_omni/entrypoints/cli/serve.py

@@ -784,7 +784,7 @@ class OmniServeCommand(CLISubcommand):
             "--vae-parallel-mode",
             type=str,
             default="tile",
-            choices=["auto", "tile", "spatial_shard_height", "spatial_shard_width"],
+            choices=["tile", "auto", "spatial_shard_height", "spatial_shard_width"],
             help="VAE parallel decode strategy for diffusion models. "
             "'tile' (default) uses patch/tile decode; 'auto' selects the longer spatial-shard axis for "
             "validated Wan tiled decodes on a full DiT group, then falls back to patch/tile decode; "