[Hardware][Ascend] Optimize Wan VAE spatial padding

.buildkite/npu/test-npu-ready.yml

@@ -33,6 +33,7 @@ steps:
             pytest -s -v
             tests/diffusion/layers/test_activation.py
             tests/diffusion/layers/test_fused_qk_norm_rope_npu.py
+            tests/diffusion/distributed/test_wan_vae_fastpath_npu.py
             tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py
             tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py
             -m 'core_model and npu and A3'

benchmarks/diffusion/bench_wan_vae_decode.py

@@ -172,6 +172,20 @@ def sync() -> None:
         dist.barrier()
 
 
+def reset_peak_memory_stats(device: torch.device) -> None:
+    """Reset peak allocation through the active device module.
+
+    PyTorch's generic accelerator memory API does not yet dispatch to all
+    out-of-tree accelerators, including torch-npu.
+    """
+    torch.get_device_module(device).reset_peak_memory_stats(device)
+
+
+def max_memory_allocated(device: torch.device) -> int:
+    """Return peak allocation through the active device module."""
+    return int(torch.get_device_module(device).max_memory_allocated(device))
+
+
 def psnr(a: torch.Tensor, b: torch.Tensor) -> float:
     mse = torch.mean((a.float() - b.float()) ** 2).item()
     return math.inf if mse == 0 else 10 * math.log10(4.0 / mse)  # outputs live in [-1, 1]
@@ -234,7 +248,7 @@ def run_level(
     for _ in range(args.warmup):
         decode(vae, latents)
     sync()
-    torch.accelerator.reset_peak_memory_stats()
+    reset_peak_memory_stats(device)
     timings = []
     output = None
     for _ in range(args.iters):
@@ -243,7 +257,7 @@ def run_level(
         output = decode(vae, latents)
         sync()
         timings.append(time.perf_counter() - start)
-    peak = torch.tensor(torch.accelerator.max_memory_allocated() / 2**30, device=device)
+    peak = torch.tensor(max_memory_allocated(device) / 2**30, device=device)
     if dist.is_initialized():
         dist.all_reduce(peak, op=dist.ReduceOp.MAX)  # the largest rank decides whether the config fits
     stats = {
@@ -310,12 +324,12 @@ def counted(conv, x, bias):
         for arm in arms:
             with patch.object(forwards, "_first_frame_conv2d", helpers[arm]):
                 sync()
-                torch.accelerator.reset_peak_memory_stats()
+                reset_peak_memory_stats(device)
                 start = time.perf_counter()
                 output = decode(vae, latents)
                 sync()
                 elapsed = (time.perf_counter() - start) * 1000
-                peak = torch.accelerator.max_memory_allocated()
+                peak = max_memory_allocated(device)
             measurements.append({"arm": arm, "wall_ms": elapsed, "peak_allocated_bytes": peak})
             if len(outputs[arm]) < 2:
                 outputs[arm].append(output.cpu())
@@ -359,8 +373,8 @@ def main() -> None:
     if args.latents and not args.first_frame_ablation:
         raise SystemExit("--latents requires --first-frame-ablation")
     accelerator = torch.accelerator.current_accelerator()
-    if accelerator is None or accelerator.type not in ("cuda", "xpu"):
-        raise SystemExit("CUDA or XPU device required")
+    if accelerator is None or accelerator.type not in ("cuda", "xpu", "npu"):
+        raise SystemExit("CUDA, XPU, or NPU device required")
     if args.profile and accelerator.type != "cuda":
         raise SystemExit("--profile requires CUDA")
     rank, world_size = init_distributed()
@@ -370,9 +384,11 @@ def main() -> None:
             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(accelerator.type, torch.accelerator.current_device_index())
+    device = torch.device(accelerator.type, torch.accelerator.current_device_index()) if world_size > 1 else accelerator
     dtype = DTYPES[args.dtype]
     if args.first_frame_ablation:
+        if accelerator.type != "cuda":
+            raise SystemExit("--first-frame-ablation requires CUDA")
         if world_size != 1 or parallel:
             raise SystemExit("--first-frame-ablation requires one GPU")
         with torch.backends.cudnn.flags(benchmark=False, allow_tf32=False):
@@ -384,7 +400,6 @@ def main() -> None:
             raise SystemExit(f"unknown fast-path level {level!r}; choose from {VAE_FAST_PATH_LEVELS}")
     if "off" not in levels:
         levels = ["off", *levels]
-
     if rank == 0:
         parallel_desc = f" vae_parallel={args.vae_parallel_mode}x{args.vae_patch_parallel_size}" if parallel else ""
         print(

docs/user_guide/diffusion/vae_fast_path.md

@@ -6,8 +6,10 @@ end-to-end time decoding latents to pixels. The decoder is a chunked 3D
 convolutional network whose reference implementation in diffusers moves every
 activation through memory several times per layer (normalization, activation,
 causal padding, feature-cache bookkeeping, shortcut upsampling). vLLM-Omni
-installs a fast path on every loaded Wan VAE that fuses this data movement into
-a handful of Triton kernels. The opt-in `channels_last` level also specializes
+installs a fast path on every supported Wan VAE. CUDA and XPU fuse this data
+movement into a handful of Triton kernels. Ascend delegates spatial padding to
+Conv3D after checking bitwise equivalence, avoiding the materialized padded
+tensor. The opt-in `channels_last` level on CUDA and XPU also specializes
 first-frame convolutions that have no history.
 
 ## Decoder levels
@@ -17,8 +19,8 @@ The fast path is controlled by `--vae-fast-path` (engine argument
 
 | Level | Default | Output vs. diffusers | What it does |
 | --- | --- | --- | --- |
-| `lossless` | yes | bit-identical | Fused RMSNorm epilogue, fused causal-conv input assembly and cache refresh, fused shortcut upsampling and residual adds (with the neighbouring convolution biases folded in), single-pass nearest 2x upsampling, preallocated output assembly. |
-| `channels_last` | no | within tolerance (PSNR typically > 60 dB) | Everything in `lossless`, plus channels-last convolution weights, a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the preceding `conv1` bias, and Conv2d for supported single-frame causal convolutions without history. |
+| `lossless` | yes | bit-identical | On CUDA and XPU: fused RMSNorm epilogue, causal-conv input assembly and cache refresh, shortcut upsampling and residual adds, nearest 2x upsampling, and preallocated output assembly. On Ascend: internal causal-Conv3D spatial padding. |
+| `channels_last` | no | within tolerance (PSNR typically > 60 dB) | CUDA/XPU only: everything in `lossless`, plus channels-last convolution weights, a single-pass channels-last RMSNorm+SiLU kernel that also absorbs the preceding `conv1` bias, and Conv2d for supported single-frame causal convolutions without history. |
 | `off` | no | bit-identical | Reference diffusers decoder. |
 
 ```bash
@@ -42,7 +44,10 @@ stages:
 
 ## Behavior and limitations
 
-- Only CUDA and XPU are supported. On other platforms the reference decoder runs.
+- CUDA and XPU support the complete fast path. On Ascend, `lossless` avoids
+  materializing spatial padding for Wan causal Conv3D after a one-time bitwise
+  check for each module and input shape; the remaining decoder operations use
+  their reference implementations. Other platforms use the reference decoder.
 - 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.
@@ -258,6 +263,15 @@ The default input is a seeded latent tensor. `--latents /path/to/latents.pt`
 can instead load a saved VAE input after latent mean/std scaling; the benchmark
 selects the temporal prefix requested by `--frames` and checks its shape.
 
+On Ascend, compare the reference decoder with the lossless internal-padding
+path using the same command and an installed Wan-based pipeline:
+
+```bash
+python benchmarks/diffusion/bench_wan_vae_decode.py \
+    --model /path/to/Helios-Distilled --size 640x384 --frames 33 \
+    --dtype fp32 --fast-path off,lossless --warmup 1 --iters 5
+```
+
 The same script benchmarks multi-GPU VAE decode when launched with `torchrun`;
 the decode is timed across all ranks and rank 0 reports:
 

tests/diffusion/distributed/test_wan_vae_fastpath_install.py

@@ -727,6 +727,7 @@ def test_registry_hook_installs_on_supported_platform(mocker, device_type: str,
     platform = mocker.Mock()
     platform.is_cuda.return_value = device_type == "cuda"
     platform.is_xpu.return_value = device_type == "xpu"
+    platform.is_npu.return_value = False
     mocker.patch.object(registry_module, "current_omni_platform", platform)
 
     registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(vae), SimpleNamespace(vae_fast_path=level))
@@ -735,13 +736,26 @@ def test_registry_hook_installs_on_supported_platform(mocker, device_type: str,
         assert getattr(vae, REPORT_ATTR).level == level
 
 
-def test_registry_hook_skips_unsupported_platform_and_non_wan_vaes(mocker) -> None:
+def test_registry_hook_supports_npu_and_skips_other_platforms_and_non_wan_vaes(mocker) -> None:
     _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
     platform = mocker.Mock()
     platform.is_cuda.return_value = False
     platform.is_xpu.return_value = False
+    platform.is_npu.return_value = True
     mocker.patch.object(registry_module, "current_omni_platform", platform)
     registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(vae), SimpleNamespace(vae_fast_path="lossless"))
+    assert is_installed(vae)
+    assert set(getattr(vae, REPORT_ATTR).patched) == {"WanCausalConv3d"}
+    assert not is_encoder_installed(vae)
+
+    uninstall_wan_vae_fastpath(vae)
+    registry_module._apply_wan_vae_fastpath_if_enabled(
+        _StubPipeline(vae), SimpleNamespace(vae_fast_path="channels_last")
+    )
+    assert not is_installed(vae)
+
+    platform.is_npu.return_value = False
+    registry_module._apply_wan_vae_fastpath_if_enabled(_StubPipeline(vae), SimpleNamespace(vae_fast_path="lossless"))
     assert not is_installed(vae)
 
     platform.is_cuda.return_value = True
@@ -758,6 +772,8 @@ def test_registry_encoder_decoder_settings_are_independent(mocker, decode_level,
     vae = AutoencoderKLWan(**CONFIG).eval()
     platform = mocker.Mock()
     platform.is_cuda.return_value = True
+    platform.is_xpu.return_value = False
+    platform.is_npu.return_value = False
     mocker.patch.object(registry_module, "current_omni_platform", platform)
     registry_module._apply_wan_vae_fastpath_if_enabled(
         _StubPipeline(vae), SimpleNamespace(vae_fast_path=decode_level, vae_encode_fast_path=encode_level)
@@ -774,6 +790,8 @@ def test_registry_component_failure_does_not_disable_other_component(mocker, fai
     vae = AutoencoderKLWan(**CONFIG).eval()
     platform = mocker.Mock()
     platform.is_cuda.return_value = True
+    platform.is_xpu.return_value = False
+    platform.is_npu.return_value = False
     mocker.patch.object(registry_module, "current_omni_platform", platform)
     install_name = "install_wan_vae_encoder_fastpath" if failed_component == "encoder" else "install_wan_vae_fastpath"
     mocker.patch.object(wan_vae_fastpath, install_name, side_effect=RuntimeError("injected failure"))

tests/diffusion/distributed/test_wan_vae_fastpath_npu.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Ascend exactness tests for Wan VAE causal-Conv3D spatial padding."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from diffusers.models.autoencoders import AutoencoderKLWan
+from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fp
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import install_wan_vae_fastpath
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+TINY_VAE_CONFIG = dict(
+    base_dim=8,
+    decoder_base_dim=8,
+    z_dim=4,
+    dim_mult=[1, 1],
+    num_res_blocks=1,
+    temperal_downsample=[False, True],
+    is_residual=True,
+)
+
+
+@hardware_test(res={"npu": "A3"}, num_cards=1)
+def test_installer_limits_npu_fast_path_to_causal_conv() -> None:
+    vae = AutoencoderKLWan(**TINY_VAE_CONFIG).eval().to(device="npu", dtype=torch.float32)
+    report = install_wan_vae_fastpath(vae, level="lossless")
+
+    assert report.installed, report.reason
+    assert set(report.patched) == {"WanCausalConv3d"}
+    assert report.patched["WanCausalConv3d"] > 1
+
+
+@hardware_test(res={"npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize(
+    "in_channels,out_channels,frames,height,width,cache_frames",
+    [
+        (3, 5, 4, 7, 9, 0),
+        (3, 5, 4, 7, 9, 2),
+        (16, 32, 1, 48, 80, 0),
+        (16, 32, 1, 48, 80, 2),
+        (128, 128, 1, 48, 80, 0),
+        (128, 128, 1, 48, 80, 2),
+    ],
+)
+@torch.no_grad()
+def test_causal_conv_internal_spatial_padding_is_bitwise_exact(
+    dtype: torch.dtype,
+    in_channels: int,
+    out_channels: int,
+    frames: int,
+    height: int,
+    width: int,
+    cache_frames: int,
+) -> None:
+    torch.manual_seed(1101)
+    conv = WanCausalConv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1).to(device="npu", dtype=dtype)
+    x = torch.randn(1, in_channels, frames, height, width, device="npu", dtype=dtype)
+    cache = (
+        None
+        if cache_frames == 0
+        else torch.randn(1, in_channels, cache_frames, height, width, device="npu", dtype=dtype)
+    )
+
+    expected = WanCausalConv3d.forward(conv, x, cache)
+    actual = fp.causal_conv_forward(conv, x, cache)
+
+    assert torch.equal(actual, expected)
+    verdicts = fp._SPATIAL_PAD_VERDICTS[conv]
+    assert len(verdicts) == 1
+    assert next(iter(verdicts.values()))
+    assert torch.equal(fp.causal_conv_forward(conv, x, cache), expected)

vllm_omni/diffusion/distributed/autoencoders/wan_vae_fastpath/forwards.py

@@ -212,6 +212,20 @@ def causal_conv_forward(self: WanCausalConv3d, x: torch.Tensor, cache_x: torch.T
         assembled = dm.cat_pad_5d(x, history, padding)
         if assembled is not None:
             return nn.Conv3d.forward(self, assembled)
+    if x.device.type == "npu" and not torch.is_grad_enabled() and not torch.compiler.is_compiling():
+        history = cache_x if padding[4] > 0 else None
+        cache_frames = 0 if history is None else history.shape[2]
+        key = _conv_verdict_key(self, x, cache_frames)
+        verdicts = _SPATIAL_PAD_VERDICTS.setdefault(self, {})
+        if verdicts.get(key, True):
+            temporal_padding = list(padding)
+            if history is not None:
+                x = torch.cat([history.to(x.device), x], dim=2)
+                temporal_padding[4] -= cache_frames
+            temporal_left, temporal_right = temporal_padding[4], temporal_padding[5]
+            if temporal_left != 0 or temporal_right != 0:
+                x = F.pad(x, (0, 0, 0, 0, temporal_left, temporal_right))
+            return _conv_with_spatial_padding(self, x, self.bias, verdicts, key)
     return WanCausalConv3d.forward(self, x, cache_x)
 
 
@@ -248,7 +262,7 @@ def _deferred_conv_bias(conv: nn.Module, output: torch.Tensor) -> torch.Tensor |
 
 # Each cache retains at most this many verdicts per live convolution module.
 _MAX_VERDICTS_PER_CONV = 128
-# Per conv module: {(x.shape, cache_frames, dtype, layout): cuDNN padding is bitwise
+# Per conv module: {(shape, cache frames, dtype, layout, device/backend settings): internal padding is bitwise
 # identical to a pre-padded input}. ``False`` routes that shape to the padded path.
 _SPATIAL_PAD_VERDICTS: WeakKeyDictionary[nn.Module, dict[tuple, bool]] = WeakKeyDictionary()
 # Per conv module: {(x.shape, cache_frames, dtype): the convolution of the channels-last
@@ -289,7 +303,7 @@ def _equal_on_random_operands(
     """Compare two convolution formulations on seeded random stand-ins for ``x`` and ``weight``.
 
     Agreement on the real operands proves little: all-zero activations or
-    weights give the same bytes under any kernel. Which kernel cuDNN runs
+    weights give the same bytes under any kernel. Which backend kernel runs
     depends only on the verdict key's metadata, never on values, so a match on
     random data is what makes a verdict reusable for every later input and
     weight update; it is also the only comparison needed. A private generator
@@ -317,16 +331,16 @@ def _conv_with_spatial_padding(
     verdicts: dict[tuple, bool],
     key: tuple,
 ) -> torch.Tensor:
-    """Convolve the time-assembled input, letting cuDNN apply the spatial zero padding.
+    """Convolve the time-assembled input with operator-provided spatial padding.
 
     Upstream pads the input tensor explicitly and convolves with ``padding=0``.
-    Handing the spatial padding to cuDNN avoids materializing the padded copy,
-    but is only bit-identical if cuDNN selects the same kernel for both
-    formulations. Exact paths compare both on seeded random operands for each
-    (conv, shape) before running the selected formulation; later calls reuse
-    the verdict until it is evicted (see :func:`_equal_on_random_operands`).
-    Internal channels-last decoder convolutions skip this bitwise check under
-    their tolerance-based contract.
+    Handing the spatial padding to the device convolution avoids materializing
+    the padded copy, but is only bit-identical if the backend selects the same
+    kernel for both formulations. Exact paths compare both on seeded random
+    operands for each (conv, shape) before running the selected formulation;
+    later calls reuse the verdict until it is evicted (see
+    :func:`_equal_on_random_operands`). Internal channels-last decoder
+    convolutions skip this bitwise check under their tolerance-based contract.
     """
     pad_height, pad_width = conv._padding[2], conv._padding[0]
     pads = (pad_width, pad_width, pad_height, pad_height, 0, 0)
@@ -353,7 +367,7 @@ def padded(x_padded: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
         _record_verdict(verdicts, key, verdict)
         if not verdict:
             logger.info(
-                "cuDNN spatial padding is not bitwise identical to pre-padded input for %s %s; "
+                "Internal spatial padding is not bitwise identical to pre-padded input for %s %s; "
                 "using the pre-padded path for this shape",
                 type(conv).__name__,
                 tuple(assembled.shape),