[Hardware][Intel] Wan VAE: run wan_vae_fastpath on XPU

benchmarks/diffusion/bench_wan_vae_decode.py

@@ -81,6 +81,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--fast-path", default=",".join(VAE_FAST_PATH_LEVELS), help="Comma-separated levels")
     parser.add_argument("--warmup", type=int, default=1)
     parser.add_argument("--iters", type=int, default=2)
+    parser.add_argument("--tiling", action="store_true", help="Decode with the VAE's spatial tiling (single GPU)")
     parser.add_argument("--seed", type=int, default=0)
     parser.add_argument("--profile", action="store_true", help="Print a torch.profiler CUDA kernel table")
     parser.add_argument("--profile-rows", type=int, default=25)
@@ -132,6 +133,8 @@ def load_vae(
             vae.init_distributed()
     else:
         vae = cls.from_pretrained(args.model, subfolder=args.subfolder, torch_dtype=dtype).to(device).eval()
+    if args.tiling:
+        vae.use_tiling = True
     if parallel:
         # Same setup as the registry: parallel decode needs tiling on.
         vae.use_tiling = True
@@ -248,7 +251,8 @@ def run_level(
         "reason": report.reason,
         "fused_silu": report.fused_silu_dtypes,
         "patched": dict(report.patched),
-        "time_s": min(timings),
+        "time_s": statistics.median(timings),
+        "times_s": timings,
         "peak_gib": peak.item(),
     }
     if args.profile:
@@ -354,16 +358,19 @@ def main() -> None:
     args = parse_args()
     if args.latents and not args.first_frame_ablation:
         raise SystemExit("--latents requires --first-frame-ablation")
-    if not torch.cuda.is_available():
-        raise SystemExit("CUDA device required")
+    accelerator = torch.accelerator.current_accelerator()
+    if accelerator is None or accelerator.type not in ("cuda", "xpu"):
+        raise SystemExit("CUDA or XPU device required")
+    if args.profile and accelerator.type != "cuda":
+        raise SystemExit("--profile requires CUDA")
     rank, world_size = init_distributed()
     parallel = args.vae_patch_parallel_size > 1
     if parallel and world_size < args.vae_patch_parallel_size:
         raise SystemExit(
             f"--vae-patch-parallel-size {args.vae_patch_parallel_size} needs torchrun with at least that many "
             f"processes (WORLD_SIZE={world_size})"
         )
-    device = torch.device("cuda", torch.accelerator.current_device_index()) if world_size > 1 else torch.device("cuda")
+    device = torch.device(accelerator.type, torch.accelerator.current_device_index())
     dtype = DTYPES[args.dtype]
     if args.first_frame_ablation:
         if world_size != 1 or parallel:
@@ -401,8 +408,11 @@ def main() -> None:
                 del output
                 continue
             speedup = baseline_time / stats["time_s"] if baseline_time else float("nan")
+            times = stats["times_s"]
             line = (
-                f"[{level:13s}] decode {stats['time_s'] * 1e3:9.1f} ms  speedup {speedup:5.2f}x  "
+                f"[{level:13s}] decode median {stats['time_s'] * 1e3:9.1f} ms  "
+                f"min {min(times) * 1e3:9.1f}  max {max(times) * 1e3:9.1f}  "
+                f"stdev {statistics.pstdev(times) * 1e3:7.1f} (n={len(times)})  speedup {speedup:5.2f}x  "
                 f"peak {stats['peak_gib']:5.2f} GiB  installed={stats['installed']}"
             )
             if level != "off":
@@ -417,6 +427,7 @@ def main() -> None:
                 else:
                     line += f"  reason={stats['reason']}"
             print(line)
+            print(f"[{level:13s}] raw_ms " + json.dumps([round(t * 1e3, 1) for t in times]))
             del output
             torch.accelerator.empty_cache()
     finally:

docs/user_guide/diffusion/parallelism/vae_parallelism.md

@@ -121,7 +121,7 @@ Additional requirements:
 | `vae_use_tiling` | bool   | False     | Must be set to `True` when using VAE patch parallelism. |
 
 !!! note "Automatic VAE Tiling"
-    When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), tile and spatial-shard modes automatically enable `vae_use_tiling`. Batch mode preserves the requested tiling setting.
+    When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), tile mode automatically enables `vae_use_tiling`. Spatial-shard and batch modes preserve the requested tiling setting: spatial-shard decode is full-frame and does not need tiling. In spatial-shard modes the **encode** is distributed only through the tiled path, so set `vae_use_tiling=True` (`--vae-use-tiling`) explicitly if you also want a multi-rank encode; the decode stays spatially sharded either way.
 
 ---
 
@@ -218,7 +218,7 @@ 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.
-- 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.
+- 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 for latents larger than one tile, whether or not `vae_use_tiling` is set. Chunked (streaming) decode has no tile-parallel path: there it decodes locally on every rank, or raises if `vae_use_tiling` is set.
 - `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).
 
 For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:

docs/user_guide/diffusion/vae_fast_path.md

@@ -42,7 +42,7 @@ stages:
 
 ## Behavior and limitations
 
-- Only CUDA is supported. On other platforms the reference decoder runs.
+- Only CUDA and XPU are supported. On other platforms the reference decoder runs.
 - The fast path is installed once per VAE instance when the pipeline is
   initialized. It rebinds the forwards of the loaded decoder modules; parameter
   names, `state_dict` keys and weight loading are unchanged.
@@ -229,8 +229,11 @@ activations and the cuDNN workspace, not by the output assembly.
 
 `benchmarks/diffusion/bench_wan_vae_decode.py` decodes seeded latents with the
 real VAE at each level and reports decode time, peak memory, bitwise equality
-and PSNR against the `off` level. `--profile` prints a per-kernel table with the
-convolution share and any layout-transpose kernels:
+and PSNR against the `off` level. After `--warmup` decodes it times `--iters`
+decodes and prints their median, min, max, standard deviation and every raw
+sample; the speedup compares medians. `--tiling` decodes with the VAE's spatial
+tiling on one GPU. `--profile` prints a per-kernel table with the convolution
+share and any layout-transpose kernels:
 
 ```bash
 python benchmarks/diffusion/bench_wan_vae_decode.py --model nvidia/Cosmos3-Nano \

tests/diffusion/distributed/test_wan_chunked_decode.py

@@ -148,27 +148,68 @@ def test_distributed_small_latent_callback_ownership(vae, monkeypatch, rank):
     assert callback.call_count == (2 if rank == 0 else 0)
 
 
-@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width", "tile"])
-def test_distributed_chunked_routing(monkeypatch, mode):
+def _bare_distributed_vae(monkeypatch, *, mode, use_tiling, world_size=2):
     vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
     torch.nn.Module.__init__(vae)
-    vae.use_tiling = True
+    vae.use_tiling = use_tiling
     vae.use_slicing = False
     vae.tile_sample_min_height = vae.tile_sample_min_width = 2
     vae.spatial_compression_ratio = 2
-    vae.distributed_executor = SimpleNamespace(group=None, parallel_mode=mode, parallel_size=2)
+    vae.distributed_executor = SimpleNamespace(group=None, parallel_mode=mode, parallel_size=2, execute=Mock())
     monkeypatch.setattr(vae, "is_distributed_enabled", lambda: True)
-    monkeypatch.setattr(torch.distributed, "get_world_size", lambda **kwargs: 2)
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda **kwargs: world_size)
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda **kwargs: 0)
+    return vae
+
+
+@pytest.mark.parametrize("use_tiling", [True, False])
+@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width", "tile"])
+def test_distributed_chunked_routing(monkeypatch, mode, use_tiling):
+    # use_tiling=False is a fresh serving instance: the registry no longer forces tiling for spatial_shard_*.
+    vae = _bare_distributed_vae(monkeypatch, mode=mode, use_tiling=use_tiling)
     decode = Mock()
     monkeypatch.setattr(wan_spatial_shard, "spatial_shard_decode", decode)
+    temporal = Mock()
+    monkeypatch.setattr(vae, "_decode_temporal", temporal)
     callback = Mock()
     z = torch.zeros(1, 2, 2, 2, 2)
-    if mode == "tile":
+    if mode == "tile" and use_tiling:
         with pytest.raises(ValueError, match="require spatial-shard"):
             vae.decode_with_chunks(z, on_chunk=callback)
         decode.assert_not_called()
+    elif mode == "tile":
+        assert vae.decode_with_chunks(z, on_chunk=callback) is None
+        decode.assert_not_called()
+        temporal.assert_called_once_with(z, on_chunk=callback)
     else:
         assert vae.decode_with_chunks(z, on_chunk=callback) is None
         decode.assert_called_once_with(
             vae, z, group=None, return_dict=True, split_dim=mode.removeprefix("spatial_shard_"), on_chunk=callback
         )
+        temporal.assert_not_called()
+
+
+@pytest.mark.parametrize("large", [True, False])
+@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"])
+def test_partial_vae_group_decode_falls_back_to_tile_parallel(monkeypatch, mode, large):
+    # vae_patch_parallel_size=2 inside a 4-rank DiT group cannot shard; without native tiling a
+    # large latent must still take the tile-parallel executor rather than a full decode per rank.
+    vae = _bare_distributed_vae(monkeypatch, mode=mode, use_tiling=False, world_size=4)
+    decode = Mock()
+    monkeypatch.setattr(wan_spatial_shard, "spatial_shard_decode", decode)
+    temporal = Mock()
+    monkeypatch.setattr(vae, "_decode_temporal", temporal)
+    merged = torch.ones(1, 3, 1, 4, 4)
+    vae.distributed_executor.execute.return_value = merged
+    z = torch.zeros(1, 2, 1, 2, 2) if large else torch.zeros(1, 2, 1, 1, 1)
+
+    out = vae._decode(z, return_dict=False)
+
+    decode.assert_not_called()
+    if large:
+        vae.distributed_executor.execute.assert_called_once()
+        assert out[0] is merged
+        temporal.assert_not_called()
+    else:
+        vae.distributed_executor.execute.assert_not_called()
+        temporal.assert_called_once_with(z, return_dict=False)

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -834,3 +834,40 @@ 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}"
     )
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("mode", "auto_tiling"),
+    [("tile", True), ("spatial_shard_height", False), ("spatial_shard_width", False)],
+)
+def test_registry_auto_enables_vae_tiling_only_for_tile_mode(monkeypatch, mocker, mode, auto_tiling):
+    """Spatial-shard decode is full-frame, so loading must not force VAE tiling on it."""
+    from torch import nn
+
+    from vllm_omni.diffusion import registry
+    from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+    from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeExecutor
+
+    vae = DistributedAutoencoderKLWan(base_dim=8, z_dim=4, dim_mult=[1, 2], num_res_blocks=1, attn_scales=[]).eval()
+    vae.distributed_executor = mocker.Mock(spec=DistributedVaeExecutor)
+
+    class Pipeline(nn.Module):
+        def __init__(self, od_config):
+            super().__init__()
+            self.vae = vae
+
+    monkeypatch.setattr(registry.DiffusionModelRegistry, "_try_load_model_cls", lambda _: Pipeline)
+    config = OmniDiffusionConfig(
+        model_class_name="test_wan_vae_tiling",
+        parallel_config=DiffusionParallelConfig(
+            tensor_parallel_size=2, vae_patch_parallel_size=2, vae_parallel_mode=mode
+        ),
+        vae_use_tiling=False,
+        vae_fast_path="off",
+    )
+    model = registry.initialize_model(config)
+    assert config.vae_use_tiling is auto_tiling
+    assert model.vae.use_tiling is auto_tiling
+    vae.distributed_executor.set_parallel_size.assert_called_once_with(2, mode=mode)

tests/diffusion/distributed/test_wan_vae_encoder_fastpath.py

@@ -107,7 +107,7 @@ def test_encode_frames_assembles_1x1_latents_like_upstream(frames):
     """
     ref, fast = pair()
     assert install_wan_vae_encoder_fastpath(fast).installed
-    seen = {}
+    seen: dict[str, tuple[int, ...]] = {}
     for name, vae in (("expected", ref), ("actual", fast)):
         vae.quant_conv.register_forward_pre_hook(lambda _, args, name=name: seen.__setitem__(name, args[0].stride()))
     x = torch.randn(1, 3, frames, 16, 16)
@@ -398,7 +398,7 @@ def fail(modules, *, channels_last):
         original_setattr = type(vae.encoder).__setattr__
 
         def fail_binding(module, name, value):
-            original_setattr(module, name, value)
+            original_setattr(module, name, value)  # type: ignore[call-arg]
             if module is vae.encoder and name == fp.CFG_ATTR:
                 raise RuntimeError("injected failure")
 
@@ -613,7 +613,7 @@ def test_norm_cache_fusion_falls_back_in_order(monkeypatch, reason):
     elif reason == "dropout_hook":
         dropout.register_forward_hook(lambda *_: events.append("hook"))
     elif reason == "dropout_custom":
-        dropout.forward = lambda x: events.append("custom") or x
+        dropout.forward = lambda x: events.append("custom") or x  # type: ignore[func-returns-value]
     x = torch.randn(1, 8, 1, 3, 5)
     cache = [torch.randn(1, 8, 2, 3, 5)]
     previous = cache[0]
@@ -668,7 +668,7 @@ def test_global_hooks_disable_only_inlined_module_calls(register):
         assert fp._has_forward_hooks(dropout) and fp._has_forward_hooks(norm)
         assert not fp._can_bypass_dropout(dropout)
         # A module's own optimized forward still runs: its hooks see the real arguments and output.
-        assert fp._kernels_allowed(SimpleNamespace(is_cuda=True))
+        assert fp._kernels_allowed(SimpleNamespace(device=torch.device("cuda")))
         assert fp.rms_norm_fastpath(norm, x) is not None
     finally:
         handle.remove()