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