[Perf][Diffusion] Fuse lossless Wan VAE data movement

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

tests/diffusion/distributed/test_autoencoder_kl_wan.py

@@ -1,3 +1,5 @@
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -41,3 +43,21 @@ def test_wan_vae_execution_context_uses_platform_autocast(mocker):
         dtype=torch.bfloat16,
         enabled=True,
     )
+
+
+def test_prepare_pipeline_installs_data_movement_before_compile(monkeypatch):
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
+    installed = []
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
+        lambda pipeline: SimpleNamespace(vaes=[model]),
+    )
+    monkeypatch.setattr(
+        wan_vae_module,
+        "install_wan_vae_data_movement",
+        lambda vae: installed.append(vae),
+    )
+
+    wan_vae_module.prepare_pipeline_wan_vae_data_movement(torch.nn.Identity())
+
+    assert installed == [model]

tests/diffusion/kernels/test_wan_vae_data_movement.py

@@ -0,0 +1,274 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn.functional as F
+from diffusers.models.autoencoders import AutoencoderKLWan
+
+from vllm_omni.diffusion.distributed.autoencoders import wan_vae_data_movement
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import (
+    _cache_payload,
+    _causal_padding,
+    _is_dynamic_spatial_conv,
+    _run_cached_causal_conv,
+    install_wan_vae_data_movement,
+)
+from vllm_omni.diffusion.kernels.wan_vae_data_movement import cat_pad_5d, dup_up3d_add
+
+pytestmark = pytest.mark.core_model
+cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _dense_5d(shape: tuple[int, ...], dtype: torch.dtype, channels_last: bool) -> torch.Tensor:
+    tensor = torch.randn(shape, device="cuda", dtype=dtype)
+    if channels_last:
+        return tensor.contiguous(memory_format=torch.channels_last_3d)
+    return tensor.contiguous()
+
+
+def _reference_cat_pad(
+    x: torch.Tensor,
+    cache: torch.Tensor | None,
+    padding: tuple[int, ...],
+) -> torch.Tensor:
+    reference_padding = list(padding)
+    if cache is not None:
+        x = torch.cat([cache, x], dim=2)
+        reference_padding[4] -= cache.shape[2]
+    return F.pad(x, reference_padding)
+
+
+@pytest.mark.cuda
+@cuda_only
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("channels_last", [False, True])
+@pytest.mark.parametrize(
+    "channels,frames,height,width,cache_frames,padding",
+    [
+        (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)),
+        (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+        (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)),
+        (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)),
+    ],
+)
+def test_cat_pad_5d_is_bitwise_and_layout_exact(
+    dtype: torch.dtype,
+    channels_last: bool,
+    channels: int,
+    frames: int,
+    height: int,
+    width: int,
+    cache_frames: int,
+    padding: tuple[int, ...],
+) -> None:
+    torch.cuda.manual_seed(0)
+    x = _dense_5d((1, channels, frames, height, width), dtype, channels_last)
+    cache = None
+    if cache_frames:
+        pad_height, pad_width = padding[2], padding[0]
+        buffer = _dense_5d(
+            (1, channels, cache_frames, height + 2 * pad_height, width + 2 * pad_width),
+            dtype,
+            channels_last,
+        )
+        cache = buffer[:, :, :, pad_height : pad_height + height, pad_width : pad_width + width]
+
+    reference = _reference_cat_pad(x, cache, padding)
+    output = cat_pad_5d(x, cache, padding)
+    assert output is not None
+    assert output.shape == reference.shape
+    assert output.stride() == reference.stride()
+    assert torch.equal(output, reference)
+
+    pair = cat_pad_5d(x, cache, padding, keep_cache_frames=2)
+    assert pair is not None
+    output_with_cache, next_cache = pair
+    assert torch.equal(output_with_cache, reference)
+    pad_height, pad_width = padding[2], padding[0]
+    keep_frames = min(2, reference.shape[2])
+    expected_cache = reference[
+        :,
+        :,
+        -keep_frames:,
+        pad_height : pad_height + height,
+        pad_width : pad_width + width,
+    ]
+    assert next_cache.shape == expected_cache.shape
+    assert torch.equal(next_cache, expected_cache)
+    assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last
+
+
+@pytest.mark.cuda
+@cuda_only
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("main_layout", ["contiguous", "channels_last", "permuted"])
+@pytest.mark.parametrize(
+    "in_channels,out_channels,frames,height,width,factor_temporal,factor_spatial,drop_first",
+    [
+        (128, 64, 1, 10, 14, 2, 2, False),
+        (128, 64, 1, 10, 14, 2, 2, True),
+        (64, 32, 2, 10, 14, 1, 2, False),
+    ],
+)
+def test_dup_up3d_add_is_bitwise_and_layout_exact(
+    dtype: torch.dtype,
+    main_layout: str,
+    in_channels: int,
+    out_channels: int,
+    frames: int,
+    height: int,
+    width: int,
+    factor_temporal: int,
+    factor_spatial: int,
+    drop_first: bool,
+) -> None:
+    torch.cuda.manual_seed(0)
+    repeats = out_channels * factor_temporal * factor_spatial * factor_spatial // in_channels
+    source = _dense_5d((1, in_channels, frames, height, width), dtype, main_layout == "channels_last")
+    out_frames = frames * factor_temporal - (factor_temporal - 1 if drop_first else 0)
+    shape = (1, out_channels, out_frames, height * factor_spatial, width * factor_spatial)
+    if main_layout == "channels_last":
+        main = _dense_5d(shape, dtype, True)
+    elif main_layout == "permuted":
+        main = torch.randn(
+            (shape[0], shape[2], shape[1], shape[3], shape[4]),
+            device="cuda",
+            dtype=dtype,
+        ).permute(0, 2, 1, 3, 4)
+    else:
+        main = _dense_5d(shape, dtype, False)
+
+    shortcut = source.repeat_interleave(repeats, dim=1)
+    shortcut = shortcut.view(
+        1,
+        out_channels,
+        factor_temporal,
+        factor_spatial,
+        factor_spatial,
+        frames,
+        height,
+        width,
+    )
+    shortcut = shortcut.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+    shortcut = shortcut.view(shape[0], shape[1], frames * factor_temporal, shape[3], shape[4])
+    if drop_first:
+        shortcut = shortcut[:, :, factor_temporal - 1 :]
+    reference = main + shortcut
+
+    output = dup_up3d_add(
+        main,
+        source,
+        factor_temporal,
+        factor_spatial,
+        repeats,
+        drop_first,
+    )
+    assert output is not None
+    assert output.shape == reference.shape
+    assert output.stride() == reference.stride()
+    assert torch.equal(output, reference)
+
+
+@pytest.mark.cuda
+@cuda_only
+@torch.no_grad()
+@pytest.mark.parametrize("channels_last", [False, True])
+@pytest.mark.parametrize("temporal_only", [False, True])
+def test_cached_conv_chunk_loop_matches_reference(monkeypatch, channels_last: bool, temporal_only: bool) -> None:
+    from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T, WanCausalConv3d
+
+    torch.cuda.manual_seed(0)
+    channels = 64
+    if temporal_only:
+        conv = WanCausalConv3d(channels, 2 * channels, (3, 1, 1), padding=(1, 0, 0))
+    else:
+        conv = WanCausalConv3d(channels, channels, 3, padding=1)
+    conv = conv.to(device="cuda", dtype=torch.float32)
+    chunks = [_dense_5d((1, channels, 1, 10, 14), torch.float32, channels_last) for _ in range(4)]
+
+    def run(force_reference: bool, start):
+        cache = [start]
+        outputs = []
+        original = wan_vae_data_movement.cat_pad_5d
+        if force_reference:
+            monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", None)
+        try:
+            for chunk in chunks:
+                outputs.append(_run_cached_causal_conv(conv, chunk, cache, 0))
+        finally:
+            monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", original)
+        return outputs, cache[0]
+
+    for start in (None, "Rep"):
+        fused_outputs, fused_cache = run(False, start)
+        reference_outputs, reference_cache = run(True, start)
+        for output, reference in zip(fused_outputs, reference_outputs, strict=True):
+            assert torch.equal(output, reference)
+        payload = _cache_payload(fused_cache)
+        assert payload is not None and payload.shape[2] == CACHE_T
+        assert torch.equal(payload, reference_cache[:, :, -CACHE_T:])
+
+
+@pytest.mark.cpu
+def test_installer_is_idempotent_and_preserves_state_dict_keys() -> None:
+    vae = OmniAutoencoderKLWan(
+        base_dim=8,
+        decoder_base_dim=8,
+        z_dim=4,
+        dim_mult=[1, 1],
+        num_res_blocks=1,
+        temperal_downsample=[False, False],
+        is_residual=True,
+    )
+    keys_before = set(vae.state_dict())
+    assert install_wan_vae_data_movement(vae)
+    assert install_wan_vae_data_movement(vae)
+    assert set(vae.state_dict()) == keys_before
+
+
+@pytest.mark.cpu
+def test_dynamic_spatial_conv_exposes_direct_path_compatibility_contract() -> None:
+    conv = torch.nn.Conv3d(2, 2, 3)
+    conv._vllm_omni_dynamic_spatial_shard_conv = True
+    conv._source_padding = (1, 1, 1, 1, 2, 0)
+
+    assert _is_dynamic_spatial_conv(conv)
+    assert _causal_padding(conv) == conv._source_padding
+
+
+@pytest.mark.cuda
+@cuda_only
+@torch.no_grad()
+@pytest.mark.parametrize("channels_last", [False, True])
+def test_tiny_wan_decoder_is_bitwise_exact(channels_last: bool) -> None:
+    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,
+    )
+    torch.manual_seed(0)
+    reference = AutoencoderKLWan(**config).eval().to("cuda")
+    candidate = OmniAutoencoderKLWan(**config).eval().to("cuda")
+    candidate.load_state_dict(reference.state_dict())
+    if channels_last:
+        for model in (reference, candidate):
+            for module in model.decoder.modules():
+                if isinstance(module, torch.nn.Conv3d):
+                    module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d)
+
+    latents = torch.randn(1, 4, 3, 4, 4, device="cuda")
+    if channels_last:
+        latents = latents.contiguous(memory_format=torch.channels_last_3d)
+    expected = reference.decode(latents, return_dict=False)[0]
+    output = candidate.decode(latents, return_dict=False)[0]
+    assert output.stride() == expected.stride()
+    assert torch.equal(output, expected)

tests/diffusion/test_diffusion_model_runner.py

@@ -728,6 +728,8 @@ 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):
+    prepared_pipelines = []
+
     class _DummyLoader:
         def __init__(self, load_config, od_config=None):
             del load_config, od_config
@@ -775,6 +777,11 @@ 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_vae_data_movement",
+        lambda pipeline: prepared_pipelines.append(pipeline),
+    )
     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
@@ -785,6 +792,7 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     assert runner.cache_backend is None
     assert runner.od_config.cache_backend is None
     assert dummy_cache_backend.enabled is False
+    assert prepared_pipelines == [runner.pipeline]
 
 
 @pytest.mark.core_model

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -18,13 +18,26 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     GridSpec,
     TileTask,
 )
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import (
+    install_wan_vae_data_movement,
+)
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
+def prepare_pipeline_wan_vae_data_movement(pipeline: torch.nn.Module) -> None:
+    """Install Wan data-movement methods before offload or compilation setup."""
+    from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+    for vae in ModuleDiscovery.discover(pipeline).vaes:
+        if isinstance(vae, OmniAutoencoderKLWan):
+            install_wan_vae_data_movement(vae)
+
+
 class OmniAutoencoderKLWan(AutoencoderKLWan):
     def _execution_context(self):
+        install_wan_vae_data_movement(self)
         try:
             first_param = next(self.parameters())
         except StopIteration:

vllm_omni/diffusion/distributed/autoencoders/wan_vae_data_movement.py

@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Guarded Wan VAE decoder data-movement fast paths.
+
+This module adapts the lossless Wan decoder optimizations from SGLang while
+keeping Diffusers' parameter names and module topology unchanged. Every
+optimized method falls back to the Diffusers 0.38.0 operation order when the
+kernel, device, layout, dtype, autograd state, or module type is unsupported.
+"""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.autoencoders.autoencoder_kl_wan import (
+    CACHE_T,
+    DupUp3D,
+    WanCausalConv3d,
+    WanDecoder3d,
+    WanResample,
+    WanResidualBlock,
+    WanResidualUpBlock,
+)
+
+from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
+
+try:
+    from vllm_omni.diffusion.kernels.wan_vae_data_movement import (
+        cat_pad_5d,
+        dup_up3d_add,
+    )
+except ImportError:  # pragma: no cover - Triton is optional on non-GPU installs.
+    cat_pad_5d = None
+    dup_up3d_add = None
+
+_INSTALL_ATTR = "_vllm_omni_wan_data_movement_installed"
+_DYNAMIC_SPATIAL_CONV_ATTR = "_vllm_omni_dynamic_spatial_shard_conv"
+
+
+def _cache_payload(cache: Any) -> torch.Tensor | None:
+    return cache if isinstance(cache, torch.Tensor) else None
+
+
+def _is_dynamic_spatial_conv(conv: nn.Module) -> bool:
+    return bool(getattr(conv, _DYNAMIC_SPATIAL_CONV_ATTR, False))
+
+
+def _spatial_shard_is_active() -> bool:
+    return wan_spatial_shard._SPATIAL_SHARD_CONTEXT.get() is not None
+
+
+def _causal_padding(conv: nn.Module) -> tuple[int, ...]:
+    if _is_dynamic_spatial_conv(conv):
+        return tuple(conv._source_padding)
+    return tuple(conv._padding)
+
+
+def _can_use_fused_cache_path(conv: nn.Module, x: torch.Tensor) -> bool:
+    return bool(
+        cat_pad_5d is not None
+        and (type(conv) is WanCausalConv3d or _is_dynamic_spatial_conv(conv))
+        and x.dim() == 5
+        and x.is_cuda
+        and not torch.is_grad_enabled()
+        and not torch.compiler.is_compiling()
+        and not _spatial_shard_is_active()
+    )
+
+
+def _run_cached_causal_conv(
+    conv: nn.Module,
+    x: torch.Tensor,
+    cache_list: list[Any],
+    index: int,
+) -> torch.Tensor:
+    """Run a causal conv and refresh its feature-cache slot."""
+    cache = cache_list[index]
+    is_repeat_marker = isinstance(cache, str)
+    payload = None if is_repeat_marker else _cache_payload(cache)
+    if _can_use_fused_cache_path(conv, x) and (
+        payload is None or (payload.device == x.device and payload.dtype == x.dtype)
+    ):
+        pair = cat_pad_5d(x, payload, _causal_padding(conv), keep_cache_frames=CACHE_T)
+        if pair is not None:
+            conv_input, cache_list[index] = pair
+            return nn.Conv3d.forward(conv, conv_input)
+
+    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+    if cache_x.shape[2] < CACHE_T and payload is not None:
+        cache_x = torch.cat(
+            [payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+            dim=2,
+        )
+    elif cache_x.shape[2] < CACHE_T and is_repeat_marker:
+        cache_x = torch.cat([torch.zeros_like(cache_x), cache_x], dim=2)
+
+    out = conv(x) if payload is None else conv(x, payload)
+    cache_list[index] = cache_x
+    return out
+
+
+def _causal_conv_forward(self: WanCausalConv3d, x: torch.Tensor, cache_x: torch.Tensor | None = None):
+    padding = list(self._padding)
+    if (
+        any(padding)
+        and _can_use_fused_cache_path(self, x)
+        and (cache_x is None or (cache_x.device == x.device and cache_x.dtype == x.dtype))
+    ):
+        conv_input = cat_pad_5d(x, cache_x, padding)
+        if conv_input is not None:
+            return nn.Conv3d.forward(self, conv_input)
+
+    if cache_x is not None and padding[4] > 0:
+        cache_x = cache_x.to(x.device)
+        x = torch.cat([cache_x, x], dim=2)
+        padding[4] -= cache_x.shape[2]
+    x = F.pad(x, padding)
+    return nn.Conv3d.forward(self, x)
+
+
+def _resample_forward(
+    self: WanResample,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+):
+    if feat_idx is None:
+        feat_idx = [0]
+    batch, channels, frames, height, width = x.size()
+    if self.mode == "upsample3d" and feat_cache is not None:
+        index = feat_idx[0]
+        if feat_cache[index] is None:
+            feat_cache[index] = "Rep"
+            feat_idx[0] += 1
+        else:
+            x = _run_cached_causal_conv(self.time_conv, x, feat_cache, index)
+            feat_idx[0] += 1
+            x = x.reshape(batch, 2, channels, frames, height, width)
+            x = torch.stack((x[:, 0], x[:, 1]), 3)
+            x = x.reshape(batch, channels, frames * 2, height, width)
+
+    frames = x.shape[2]
+    x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
+    x = self.resample(x)
+    x = x.view(batch, frames, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
+
+    if self.mode == "downsample3d" and feat_cache is not None:
+        index = feat_idx[0]
+        if feat_cache[index] is None:
+            feat_cache[index] = x.clone()
+            feat_idx[0] += 1
+        else:
+            cache_x = x[:, :, -1:, :, :].clone()
+            x = self.time_conv(torch.cat([feat_cache[index][:, :, -1:, :, :], x], dim=2))
+            feat_cache[index] = cache_x
+            feat_idx[0] += 1
+    return x
+
+
+def _residual_block_forward(
+    self: WanResidualBlock,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+):
+    if feat_idx is None:
+        feat_idx = [0]
+    residual = self.conv_shortcut(x)
+    x = self.nonlinearity(self.norm1(x))
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x = _run_cached_causal_conv(self.conv1, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv1(x)
+
+    x = self.dropout(self.nonlinearity(self.norm2(x)))
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x = _run_cached_causal_conv(self.conv2, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv2(x)
+    return x + residual
+
+
+def _residual_up_block_forward(
+    self: WanResidualUpBlock,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+    first_chunk: bool = False,
+):
+    if feat_idx is None:
+        feat_idx = [0]
+    shortcut_source = x.clone()
+    for resnet in self.resnets:
+        if feat_cache is None:
+            x = resnet(x)
+        else:
+            x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx)
+
+    if self.upsampler is not None:
+        if feat_cache is None:
+            x = self.upsampler(x)
+        else:
+            x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx)
+
+    shortcut = self.avg_shortcut
+    if shortcut is None:
+        return x
+    if (
+        dup_up3d_add is not None
+        and type(shortcut) is DupUp3D
+        and x.is_cuda
+        and shortcut_source.is_cuda
+        and x.dtype == shortcut_source.dtype
+        and not torch.is_grad_enabled()
+        and not torch.compiler.is_compiling()
+        and not _spatial_shard_is_active()
+    ):
+        fused = dup_up3d_add(
+            x,
+            shortcut_source,
+            shortcut.factor_t,
+            shortcut.factor_s,
+            shortcut.repeats,
+            first_chunk,
+        )
+        if fused is not None:
+            return fused
+    return x + shortcut(shortcut_source, first_chunk=first_chunk)
+
+
+def _decoder_forward(
+    self: WanDecoder3d,
+    x: torch.Tensor,
+    feat_cache: list[Any] | None = None,
+    feat_idx: list[int] | None = None,
+    first_chunk: bool = False,
+):
+    if feat_idx is None:
+        feat_idx = [0]
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x = _run_cached_causal_conv(self.conv_in, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv_in(x)
+
+    x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx)
+    for up_block in self.up_blocks:
+        x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)
+
+    x = self.nonlinearity(self.norm_out(x))
+    if feat_cache is not None:
+        index = feat_idx[0]
+        x = _run_cached_causal_conv(self.conv_out, x, feat_cache, index)
+        feat_idx[0] += 1
+    else:
+        x = self.conv_out(x)
+    return x
+
+
+def install_wan_vae_data_movement(vae: nn.Module) -> bool:
+    """Install guarded fast paths on a Diffusers Wan decoder once."""
+    if getattr(vae, _INSTALL_ATTR, False):
+        return True
+    if cat_pad_5d is None or dup_up3d_add is None:
+        return False
+
+    decoder = getattr(vae, "decoder", None)
+    if type(decoder) is not WanDecoder3d:
+        return False
+
+    decoder.forward = MethodType(_decoder_forward, decoder)
+    patched = 0
+    for module in decoder.modules():
+        module_type = type(module)
+        if module_type is WanCausalConv3d:
+            module.forward = MethodType(_causal_conv_forward, module)
+            patched += 1
+        elif _is_dynamic_spatial_conv(module):
+            # The dynamic wrapper owns halo exchange while a spatial request is
+            # active. Its direct path can still use the lossless fused cache
+            # movement through ``_run_cached_causal_conv``.
+            patched += 1
+        elif module_type is WanResidualBlock:
+            module.forward = MethodType(_residual_block_forward, module)
+        elif module_type is WanResidualUpBlock:
+            module.forward = MethodType(_residual_up_block_forward, module)
+        elif module_type is WanResample:
+            module.forward = MethodType(_resample_forward, module)
+
+    if patched == 0:
+        return False
+    setattr(vae, _INSTALL_ATTR, True)
+    return True

vllm_omni/diffusion/kernels/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/kernels/wan_vae_data_movement.py

@@ -0,0 +1,405 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bit-exact data-movement kernels for the Wan causal VAE decoder.
+
+Adapted from SGLang PR #34125:
+https://github.com/sgl-project/sglang/pull/34125
+
+The kernels in this module only move values, fill zeros, or perform the same
+two-operand addition as the PyTorch reference path. They preserve the input
+layout instead of selecting a new Conv3d layout, so enabling them does not
+change the decoder's cuDNN accumulation order.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_MAX_INT32 = 2**31 - 1
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+
+
+@triton.jit
+def _cat_pad_5d_kernel(
+    x_ptr,
+    cache_ptr,
+    out_ptr,
+    keep_ptr,
+    total,
+    channels,
+    frames,
+    height,
+    width,
+    cache_frames,
+    out_frames,
+    out_height,
+    out_width,
+    zero_front_frames,
+    pad_height,
+    pad_width,
+    stride_x_batch,
+    stride_x_channel,
+    stride_x_frame,
+    stride_x_height,
+    stride_x_width,
+    stride_cache_batch,
+    stride_cache_channel,
+    stride_cache_frame,
+    stride_cache_height,
+    stride_cache_width,
+    has_cache: tl.constexpr,
+    keep_frames_const: tl.constexpr,
+    channels_inner: tl.constexpr,
+    index_64bit: tl.constexpr,
+    block_size_const: tl.constexpr,
+):
+    if index_64bit:
+        offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
+    else:
+        offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
+    mask = offsets < total
+
+    if channels_inner:
+        out_channel = offsets % channels
+        remaining = offsets // channels
+        out_width_idx = remaining % out_width
+        remaining = remaining // out_width
+        out_height_idx = remaining % out_height
+        remaining = remaining // out_height
+        out_frame = remaining % out_frames
+        out_batch = remaining // out_frames
+    else:
+        out_width_idx = offsets % out_width
+        remaining = offsets // out_width
+        out_height_idx = remaining % out_height
+        remaining = remaining // out_height
+        out_frame = remaining % out_frames
+        remaining = remaining // out_frames
+        out_channel = remaining % channels
+        out_batch = remaining // channels
+
+    input_width = out_width_idx - pad_width
+    input_height = out_height_idx - pad_height
+    input_frame = out_frame - zero_front_frames
+
+    spatial_valid = (input_width >= 0) & (input_width < width) & (input_height >= 0) & (input_height < height)
+    from_cache = spatial_valid & (input_frame >= 0) & (input_frame < cache_frames)
+    from_x = spatial_valid & (input_frame >= cache_frames) & (input_frame < cache_frames + frames)
+
+    x_frame = input_frame - cache_frames
+    x_offset = (
+        out_batch * stride_x_batch
+        + out_channel * stride_x_channel
+        + x_frame * stride_x_frame
+        + input_height * stride_x_height
+        + input_width * stride_x_width
+    )
+    values = tl.load(x_ptr + x_offset, mask=mask & from_x, other=0.0)
+    if has_cache:
+        cache_offset = (
+            out_batch * stride_cache_batch
+            + out_channel * stride_cache_channel
+            + input_frame * stride_cache_frame
+            + input_height * stride_cache_height
+            + input_width * stride_cache_width
+        )
+        cache_values = tl.load(cache_ptr + cache_offset, mask=mask & from_cache, other=0.0)
+        values = tl.where(from_cache, cache_values, values)
+    tl.store(out_ptr + offsets, values, mask=mask)
+
+    if keep_frames_const > 0:
+        keep_frame = out_frame - (out_frames - keep_frames_const)
+        keep_mask = mask & spatial_valid & (keep_frame >= 0)
+        if channels_inner:
+            keep_offset = (
+                ((out_batch * keep_frames_const + keep_frame) * height + input_height) * width + input_width
+            ) * channels + out_channel
+        else:
+            keep_offset = (
+                ((out_batch * channels + out_channel) * keep_frames_const + keep_frame) * height + input_height
+            ) * width + input_width
+        tl.store(keep_ptr + keep_offset, values, mask=keep_mask)
+
+
+def _uses_channels_last_3d(tensor: torch.Tensor) -> bool:
+    return tensor.shape[1] > 1 and tensor.stride(1) == 1
+
+
+def cat_pad_5d(
+    x: torch.Tensor,
+    cache_x: torch.Tensor | None,
+    padding: list[int] | tuple[int, ...],
+    keep_cache_frames: int = 0,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None:
+    """Fuse temporal concat, constant padding, and optional cache refresh.
+
+    ``padding`` follows the Wan causal Conv3d convention
+    ``(w_left, w_right, h_top, h_bottom, t_front, t_back)``. The output uses
+    the same dense layout that the reference ``cat``/``pad`` path selects.
+    Unsupported inputs return ``None`` so callers can execute the reference
+    implementation.
+    """
+    if len(padding) != 6:
+        return None
+    pad_width_left, pad_width_right, pad_height_top, pad_height_bottom, pad_front, pad_back = padding
+    if any(value < 0 for value in padding):
+        return None
+    if pad_width_left != pad_width_right or pad_height_top != pad_height_bottom or pad_back != 0:
+        return None
+    if x.dim() != 5 or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES:
+        return None
+    if keep_cache_frames < 0:
+        return None
+
+    channels_inner = _uses_channels_last_3d(x)
+    cache_frames = 0
+    if cache_x is not None:
+        if (
+            cache_x.dim() != 5
+            or cache_x.dtype != x.dtype
+            or cache_x.device != x.device
+            or cache_x.shape[0] != x.shape[0]
+            or cache_x.shape[1] != x.shape[1]
+            or cache_x.shape[3:] != x.shape[3:]
+            or _uses_channels_last_3d(cache_x) != channels_inner
+        ):
+            return None
+        cache_frames = cache_x.shape[2]
+
+    zero_front_frames = pad_front - cache_frames
+    if zero_front_frames < 0:
+        return None
+
+    batch, channels, frames, height, width = x.shape
+    out_frames = pad_front + frames
+    out_height = height + 2 * pad_height_top
+    out_width = width + 2 * pad_width_left
+    keep_frames = min(keep_cache_frames, out_frames)
+    memory_format = torch.channels_last_3d if channels_inner else torch.contiguous_format
+    out = torch.empty(
+        (batch, channels, out_frames, out_height, out_width),
+        device=x.device,
+        dtype=x.dtype,
+        memory_format=memory_format,
+    )
+    total = out.numel()
+    if total == 0 or total > _MAX_INT32 * 4:
+        return None
+
+    if keep_frames:
+        keep_arg = torch.empty(
+            (batch, channels, keep_frames, height, width),
+            device=x.device,
+            dtype=x.dtype,
+            memory_format=memory_format,
+        )
+    else:
+        keep_arg = out
+
+    if cache_x is None:
+        cache_arg = x
+        cache_strides = (0, 0, 0, 0, 0)
+    else:
+        cache_arg = cache_x
+        cache_strides = cache_x.stride()
+
+    block_size = 512
+    grid = (triton.cdiv(total, block_size),)
+    with torch.get_device_module().device(x.device):
+        _cat_pad_5d_kernel[grid](
+            x,
+            cache_arg,
+            out,
+            keep_arg,
+            total,
+            channels,
+            frames,
+            height,
+            width,
+            cache_frames,
+            out_frames,
+            out_height,
+            out_width,
+            zero_front_frames,
+            pad_height_top,
+            pad_width_left,
+            *x.stride(),
+            *cache_strides,
+            has_cache=cache_x is not None,
+            keep_frames_const=keep_frames,
+            channels_inner=channels_inner,
+            index_64bit=total >= _MAX_INT32,
+            block_size_const=block_size,
+        )
+
+    if keep_cache_frames:
+        return out, keep_arg
+    return out
+
+
+@triton.jit
+def _dup_up3d_add_kernel(
+    main_ptr,
+    source_ptr,
+    out_ptr,
+    total,
+    out_channels,
+    out_frames,
+    out_height,
+    out_width,
+    frame_offset,
+    stride_main_batch,
+    stride_main_channel,
+    stride_main_frame,
+    stride_main_height,
+    stride_main_width,
+    stride_source_batch,
+    stride_source_channel,
+    stride_source_frame,
+    stride_source_height,
+    stride_source_width,
+    stride_out_batch,
+    stride_out_channel,
+    stride_out_frame,
+    stride_out_height,
+    stride_out_width,
+    factor_temporal_const: tl.constexpr,
+    factor_spatial_const: tl.constexpr,
+    repeats_const: tl.constexpr,
+    channels_inner: tl.constexpr,
+    index_64bit: tl.constexpr,
+    block_size_const: tl.constexpr,
+):
+    if index_64bit:
+        offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
+    else:
+        offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
+    mask = offsets < total
+
+    if channels_inner:
+        out_channel = offsets % out_channels
+        remaining = offsets // out_channels
+        out_width_idx = remaining % out_width
+        remaining = remaining // out_width
+        out_height_idx = remaining % out_height
+        remaining = remaining // out_height
+        out_frame = remaining % out_frames
+        out_batch = remaining // out_frames
+    else:
+        out_width_idx = offsets % out_width
+        remaining = offsets // out_width
+        out_height_idx = remaining % out_height
+        remaining = remaining // out_height
+        out_frame = remaining % out_frames
+        remaining = remaining // out_frames
+        out_channel = remaining % out_channels
+        out_batch = remaining // out_channels
+
+    unsliced_frame = out_frame + frame_offset
+    source_frame = unsliced_frame // factor_temporal_const
+    temporal_remainder = unsliced_frame % factor_temporal_const
+    source_height = out_height_idx // factor_spatial_const
+    height_remainder = out_height_idx % factor_spatial_const
+    source_width = out_width_idx // factor_spatial_const
+    width_remainder = out_width_idx % factor_spatial_const
+    repeated_channel = (
+        (out_channel * factor_temporal_const + temporal_remainder) * factor_spatial_const + height_remainder
+    ) * factor_spatial_const + width_remainder
+    source_channel = repeated_channel // repeats_const
+
+    main_offset = (
+        out_batch * stride_main_batch
+        + out_channel * stride_main_channel
+        + out_frame * stride_main_frame
+        + out_height_idx * stride_main_height
+        + out_width_idx * stride_main_width
+    )
+    source_offset = (
+        out_batch * stride_source_batch
+        + source_channel * stride_source_channel
+        + source_frame * stride_source_frame
+        + source_height * stride_source_height
+        + source_width * stride_source_width
+    )
+    out_offset = (
+        out_batch * stride_out_batch
+        + out_channel * stride_out_channel
+        + out_frame * stride_out_frame
+        + out_height_idx * stride_out_height
+        + out_width_idx * stride_out_width
+    )
+    main_value = tl.load(main_ptr + main_offset, mask=mask, other=0.0)
+    source_value = tl.load(source_ptr + source_offset, mask=mask, other=0.0)
+    value = main_value.to(tl.float32) + source_value.to(tl.float32)
+    tl.store(out_ptr + out_offset, value, mask=mask)
+
+
+def dup_up3d_add(
+    main: torch.Tensor,
+    source: torch.Tensor,
+    factor_temporal: int,
+    factor_spatial: int,
+    repeats: int,
+    drop_first_frames: bool,
+) -> torch.Tensor | None:
+    """Evaluate ``main + DupUp3D(source)`` without materializing the shortcut."""
+    if main.dim() != 5 or source.dim() != 5:
+        return None
+    if factor_temporal <= 0 or factor_spatial <= 0 or repeats <= 0:
+        return None
+    if factor_temporal & (factor_temporal - 1) or factor_spatial & (factor_spatial - 1):
+        return None
+    if repeats & (repeats - 1):
+        return None
+    if not main.is_cuda or not source.is_cuda:
+        return None
+    if main.dtype != source.dtype or main.dtype not in _SUPPORTED_DTYPES or main.device != source.device:
+        return None
+
+    batch, source_channels, frames, height, width = source.shape
+    frame_offset = factor_temporal - 1 if drop_first_frames else 0
+    numerator = source_channels * repeats
+    denominator = factor_temporal * factor_spatial * factor_spatial
+    if numerator % denominator:
+        return None
+    expected_shape = (
+        batch,
+        numerator // denominator,
+        frames * factor_temporal - frame_offset,
+        height * factor_spatial,
+        width * factor_spatial,
+    )
+    if tuple(main.shape) != expected_shape:
+        return None
+
+    out = torch.empty_like(main)
+    total = out.numel()
+    if total == 0 or total > _MAX_INT32 * 4:
+        return None
+
+    block_size = 512
+    grid = (triton.cdiv(total, block_size),)
+    with torch.get_device_module().device(main.device):
+        _dup_up3d_add_kernel[grid](
+            main,
+            source,
+            out,
+            total,
+            expected_shape[1],
+            expected_shape[2],
+            expected_shape[3],
+            expected_shape[4],
+            frame_offset,
+            *main.stride(),
+            *source.stride(),
+            *out.stride(),
+            factor_temporal_const=factor_temporal,
+            factor_spatial_const=factor_spatial,
+            repeats_const=repeats,
+            channels_inner=out.stride(1) == 1 and expected_shape[1] > 1,
+            index_64bit=total >= _MAX_INT32,
+            block_size_const=block_size,
+        )
+    return out

vllm_omni/diffusion/models/helios/pipeline_helios.py

@@ -21,6 +21,7 @@ from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -215,6 +216,7 @@ class HeliosPipeline(
         self.vae = AutoencoderKLWan.from_pretrained(
             model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
         ).to(self.device)
+        install_wan_vae_data_movement(self.vae)
 
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
         self.transformer = create_transformer_from_config(

vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py

@@ -19,6 +19,7 @@ from torch import nn
 from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.interface import (
     SupportImageInput,
@@ -374,6 +375,7 @@ class LingBotVideoPipeline(
             torch_dtype=vae_dtype,
             local_files_only=local_files_only,
         ).to(self.device)
+        install_wan_vae_data_movement(self.vae)
         self.scheduler = FlowUniPCMultistepScheduler.from_pretrained(
             model,
             subfolder=scheduler_subfolder,

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -30,6 +30,9 @@ 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_vae_data_movement,
+)
 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 (
@@ -243,6 +246,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     custom_pipeline_name=custom_pipeline_name,
                     device=self.device,
                 )
+                prepare_pipeline_wan_vae_data_movement(self.pipeline)
         time_after_load = time.perf_counter()
 
         logger.info(