[LingBot World] Speed up streaming VAE decode and RGB frame delivery (#7749)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

tests/diffusion/distributed/test_wan_decoder_fast_path.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The exact Wan decoder fast path must be bit-identical to the plain streaming decode.
+
+Two decoders are driven frame by frame through ``WanStreamingDecoder`` over two chunks of one session and
+a fresh session (the persistent buffers' causal front frames go zeros -> cache -> zeros), on CPU: plain
+diffusers modules in one process, and the spatially sharded wrappers across two gloo ranks with real halo
+exchange, each in fp32 and under bf16 autocast (where the parameter cast applies). Every output must
+``torch.equal`` the untouched decoder's on every rank.
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d, WanDecoder3d
+
+from vllm_omni.diffusion.distributed.autoencoders.wan_decoder_fast_path import install_wan_decoder_fast_path
+from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import (
+    WanDistCausalConv3d,
+    WanDistConv2d,
+    install_wan_spatial_shard_decode,
+)
+from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
+
+
+def _make_vae(seed: int = 0) -> SimpleNamespace:
+    torch.manual_seed(seed)
+    # Three levels: the default temporal-upsample list gives one upsample2d and one upsample3d resample.
+    decoder = WanDecoder3d(dim=8, z_dim=4, dim_mult=[1, 2, 4], num_res_blocks=1, attn_scales=[])
+    post_quant_conv = WanCausalConv3d(4, 4, 1)
+    for p in list(decoder.parameters()) + list(post_quant_conv.parameters()):
+        p.data.uniform_(-0.5, 0.5)
+    decoder.eval()
+    count = sum(1 for m in decoder.modules() if m.__class__.__name__ == "WanCausalConv3d")
+    return SimpleNamespace(decoder=decoder, post_quant_conv=post_quant_conv, _cached_conv_counts={"decoder": count})
+
+
+def _run(vae, latents: list[torch.Tensor], autocast: bool) -> list[torch.Tensor]:
+    decoder = WanStreamingDecoder(vae)
+    outs = []
+    ctx = torch.autocast("cpu", dtype=torch.bfloat16) if autocast else torch.no_grad()
+    with torch.no_grad(), ctx:
+        state = decoder.new_decode_state("a")
+        outs.append(decoder.decode_chunk(latents[0], state))
+        outs.append(decoder.decode_chunk(latents[1], state))
+        fresh = decoder.new_decode_state("b")
+        outs.append(decoder.decode_chunk(latents[2], fresh))
+    return [o.float() for o in outs]
+
+
+def _latents() -> list[torch.Tensor]:
+    torch.manual_seed(1)
+    return [torch.randn(1, 4, 2, 6, 10), torch.randn(1, 4, 1, 6, 10), torch.randn(1, 4, 2, 6, 10)]
+
+
+def _check_pair(reference, candidate, autocast: bool, sharded: bool, level: str = "exact") -> None:
+    dtype = torch.bfloat16 if autocast else None
+    counts = install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level=level)
+    assert counts["level"] == level
+    assert counts["upsamples"] == 2
+    convs = sum(1 for m in candidate.decoder.modules() if isinstance(m, torch.nn.Conv3d | torch.nn.Conv2d)) + 1
+    assert counts["conv_params_cast"] == (convs if autocast else 0)
+    dist_convs = sum(1 for m in candidate.decoder.modules() if isinstance(m, WanDistCausalConv3d | WanDistConv2d))
+    assert counts["persistent_input_buffers"] == dist_convs
+    assert (dist_convs > 0) == sharded
+    if level == "fused":
+        assert (
+            counts["fused_norm_silu"]
+            == 2 * sum(1 for m in candidate.decoder.modules() if m.__class__.__name__ == "WanResidualBlock") + 1
+        )
+        assert counts["channels_last_convs"] == convs
+        assert counts["resample_forwards"] == 2
+        with pytest.raises(ValueError, match="already installed"):
+            install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level="exact")
+    # Idempotent: a second install is a no-op that reports the same counts.
+    assert install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level=level) is counts
+    latents = _latents()
+    expected = _run(reference, latents, autocast)
+    actual = _run(candidate, latents, autocast)
+    assert len(expected) == len(actual) == 3
+    for e, a in zip(expected, actual):
+        assert e.shape == a.shape
+        if level == "exact":
+            assert torch.equal(e, a)
+        else:
+            # The fused level changes layouts only on CPU (the kernels need CUDA); the eager norm's reduction
+            # order may differ with the layout, so the outputs are close, not identical.
+            assert torch.allclose(e, a, atol=1e-2 if autocast else 1e-5, rtol=1e-2 if autocast else 1e-4), (
+                (e - a).abs().max()
+            )
+    if sharded:
+        bufs = [m._input_buf for m in candidate.decoder.modules() if isinstance(m, WanDistCausalConv3d | WanDistConv2d)]
+        assert all(b is not None for b in bufs)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("autocast", [False, True])
+def test_fast_path_is_bit_identical(autocast: bool) -> None:
+    reference = _make_vae()
+    candidate = copy.deepcopy(reference)
+    _check_pair(reference, candidate, autocast, sharded=False)
+
+
+def _swap_norms_to_rmsnorm_vae(vae) -> int:
+    """Ensure the decoder uses RMSNormVAE and return the number of matching norms.
+
+    Collecting pipeline tests can already apply patch_wan_rms_norm process-wide before this test runs.
+    """
+    from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
+
+    count = 0
+    for module in list(vae.decoder.modules()):
+        for name, child in list(module.named_children()):
+            if isinstance(child, RMSNormVAE):
+                count += 1
+                continue
+            if child.__class__.__name__ == "WanRMS_norm":
+                images = child.gamma.dim() == 3  # (dim, 1, 1) for the attention block's 4D norm
+                norm = RMSNormVAE(child.gamma.shape[0], channel_first=child.channel_first, images=images, bias=False)
+                norm.gamma.data.copy_(child.gamma.data.reshape(norm.gamma.shape))
+                setattr(module, name, norm)
+                count += 1
+    return count
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("norm", ["default", "RMSNormVAE"])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_fused_level_plumbing_on_cpu(autocast: bool, norm: str) -> None:
+    reference = _make_vae()
+    if norm == "RMSNormVAE":
+        assert _swap_norms_to_rmsnorm_vae(reference) > 0
+    candidate = copy.deepcopy(reference)
+    _check_pair(reference, candidate, autocast, sharded=False, level="fused")
+
+
+def _sharded_worker(rank: int, world_size: int, port: str, autocast: bool, return_dict, level: str = "exact") -> None:
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = port
+    dist.init_process_group("gloo", rank=rank, world_size=world_size)
+    try:
+        torch.manual_seed(0)
+        reference = _make_vae()
+        candidate = copy.deepcopy(reference)
+        for vae in (reference, candidate):
+            install_wan_spatial_shard_decode(vae, dist.group.WORLD, split_dim="width", dst=None)
+        _check_pair(reference, candidate, autocast, sharded=True, level=level)
+        return_dict[rank] = "ok"
+    except BaseException as exc:  # report, the parent asserts
+        return_dict[rank] = f"{type(exc).__name__}: {exc}"
+        raise
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("level", ["exact", "fused"])
+@pytest.mark.parametrize("autocast", [False, True])
+def test_sharded_fast_path_on_every_rank(autocast: bool, level: str) -> None:
+    manager = mp.get_context("spawn").Manager()
+    return_dict = manager.dict()
+    port = str(29610 + int(autocast) + 2 * (level == "fused"))
+    mp.spawn(_sharded_worker, args=(2, port, autocast, return_dict, level), nprocs=2, join=True)
+    for rank in range(2):
+        assert return_dict.get(rank) == "ok", f"rank {rank}: {return_dict.get(rank)}"
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_conv_dtype_requires_half() -> None:
+    with pytest.raises(ValueError, match="conv_dtype"):
+        install_wan_decoder_fast_path(_make_vae(), conv_dtype=torch.float32)

vllm_omni/diffusion/distributed/autoencoders/wan_decoder_fast_path.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exact fast path for the Wan VAE decoder on a streaming (frame-at-a-time) decode.
+
+A kernel ledger of the LingBot-World served path (4xH200, 480x832, width-sharded bf16 decode) put the
+decoder at 91 ms of GPU time per chunk per rank, half of it in ~2,300 launches of glue around the
+convolutions: autocast re-casting every conv weight and bias on every call (the autocast weight cache only
+serves leaves that require grad, so under ``torch.no_grad`` it never hits), the nearest upsample's fp32
+round trip, and a freshly allocated, zero-filled, concatenated input tensor per conv call. None of that
+touches a value that reaches the convolution, so it can go without changing a single output bit:
+
+- ``conv_dtype``: the decoder's convolution parameters are cast once to the dtype the decode already runs
+  under (``decode_autocast_dtype``), the same cast autocast applied per call. Norm parameters stay in their
+  own dtype, so the normalisation arithmetic and its fp32 promotion are untouched.
+- ``WanUpsample`` runs its nearest-exact gather on the activation dtype directly instead of via
+  ``x.float()`` and ``type_as``: a gather moves values, so the result is identical.
+- The spatially sharded conv wrappers keep one input buffer per conv across calls
+  (``reuse_input_buffer``), writing only the activation interior and the halo slots into it.
+
+Numerics: bit-identical to the plain path (``tests/diffusion/distributed/test_wan_decoder_fast_path.py``).
+Memory: one persistent input buffer per sharded conv (about the size of that conv's input).
+
+Level ``"fused"`` adds the rest of what the ledger showed, at the cost of exactness (it is gated with the bf16
+decode option, whose quality gate it shares): the decoder runs channels_last_3d end to end so cuDNN uses its
+NHWC kernels without the NCHW<->NHWC transposes it otherwise inserts around every convolution, every
+``WanRMS_norm -> SiLU`` chain becomes one Triton kernel that reads and writes the activation dtype (the eager
+chain promotes to fp32 at ``* gamma`` and leaves the cast to the next convolution), the nearest upsample is a
+channels_last gather, and the ``upsample3d`` time-pair interleave writes channels_last_3d directly. Values
+differ from eager only through the norm's fp32 reduction order and the single rounding at the SiLU output.
+"""
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import Any
+
+import torch
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.wan_decoder_kernels import (
+    can_use_wan_rmsnorm_silu,
+    wan_rmsnorm_silu,
+    wan_rmsnorm_silu_reference,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_decoder_utils import (
+    _dense_channels_last_3d,
+    _is_nearest_upsample,
+    _resample_forward_channels_last,
+    _upsample_forward,
+    _upsample_forward_channels_last,
+)
+from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import WanDistCausalConv3d, WanDistConv2d
+
+logger = init_logger(__name__)
+
+_INSTALLED_ATTR = "_vllm_omni_wan_decoder_fast_path"
+LEVELS = ("exact", "fused")
+# diffusers' WanRMS_norm, and vLLM-Omni's RMSNormVAE that patch_wan_rms_norm swaps in for every Wan-family VAE.
+_NORM_CLASSES = ("WanRMS_norm", "RMSNormVAE")
+
+
+class FusedWanRMSNormSiLU(nn.Module):
+    """``WanRMS_norm`` followed by SiLU as one kernel on channels_last_3d CUDA tensors.
+
+    Keeps the norm's parameters registered under their own names (``...norm1.gamma``), so state dicts and
+    weight loading are unchanged. Off the kernel's domain (CPU, other layouts, grad) it runs the eager op
+    chain, and a channels-first CUDA input is converted to channels_last_3d once (the attention block's
+    gathered frame is the only such producer).
+    """
+
+    def __init__(self, norm: nn.Module) -> None:
+        super().__init__()
+        self.gamma = norm.gamma
+        bias = norm.bias
+        if isinstance(bias, torch.Tensor):
+            self.bias = bias
+            self._bias_value = 0.0
+        else:
+            self.bias = None
+            self._bias_value = float(bias) if bias is not None else 0.0
+        self.scale = float(norm.scale)
+        # diffusers' WanRMS_norm normalises x.float() with F.normalize's default eps; vLLM-Omni's RMSNormVAE
+        # (patched in for every Wan-family VAE it loads) normalises x itself with its own epsilon.
+        self.upcast = norm.__class__.__name__ == "WanRMS_norm"
+        self.eps = 1e-12 if self.upcast else float(getattr(norm, "epsilon", 1e-6))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() == 5 and x.is_cuda and not (torch.is_grad_enabled() and x.requires_grad):
+            if not _dense_channels_last_3d(x):
+                x = x.contiguous(memory_format=torch.channels_last_3d)
+            if can_use_wan_rmsnorm_silu(x, self.gamma, self.bias):
+                return wan_rmsnorm_silu(x, self.gamma, self.bias, self.scale, self.eps, self.upcast)
+        bias = self.bias if self.bias is not None else self._bias_value
+        return wan_rmsnorm_silu_reference(x, self.gamma, bias, self.scale, self.eps, self.upcast)
+
+
+def _install_fused(decoder: nn.Module, post_quant_conv: nn.Module | None, counts: dict[str, int]) -> None:
+    """Channels_last weights and buffers, fused norm+SiLU, layout-preserving resample. Fails closed."""
+    blocks = [m for m in decoder.modules() if m.__class__.__name__ == "WanResidualBlock"]
+    for m in blocks:
+        for name in ("norm1", "norm2"):
+            norm = getattr(m, name, None)
+            if norm.__class__.__name__ not in _NORM_CLASSES or not getattr(norm, "channel_first", False):
+                raise ValueError(
+                    f"fused fast path needs a channel-first Wan RMSNorm at {name}; got {type(norm).__name__}"
+                )
+        if not isinstance(m.nonlinearity, nn.SiLU) or m.nonlinearity.inplace:
+            raise ValueError("fused fast path needs a plain SiLU residual block nonlinearity")
+    head_norm = getattr(decoder, "norm_out", None)
+    if head_norm.__class__.__name__ not in _NORM_CLASSES or not isinstance(
+        getattr(decoder, "nonlinearity", None), nn.SiLU
+    ):
+        raise ValueError("fused fast path needs the standard WanRMS_norm + SiLU decoder head")
+    for m in blocks:
+        m.norm1 = FusedWanRMSNormSiLU(m.norm1)
+        m.norm2 = FusedWanRMSNormSiLU(m.norm2)
+        m.nonlinearity = nn.Identity()
+        counts["fused_norm_silu"] += 2
+    decoder.norm_out = FusedWanRMSNormSiLU(decoder.norm_out)
+    decoder.nonlinearity = nn.Identity()
+    counts["fused_norm_silu"] += 1
+    roots = [decoder] + ([post_quant_conv] if post_quant_conv is not None else [])
+    for root in roots:
+        for module in root.modules():
+            if isinstance(module, nn.Conv3d):
+                module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d)
+                counts["channels_last_convs"] += 1
+            elif isinstance(module, nn.Conv2d):
+                module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last)
+                counts["channels_last_convs"] += 1
+            if isinstance(module, WanDistCausalConv3d):
+                module.input_memory_format = torch.channels_last_3d
+            elif isinstance(module, WanDistConv2d):
+                module.input_memory_format = torch.channels_last
+            if _is_nearest_upsample(module):
+                module.forward = MethodType(_upsample_forward_channels_last, module)
+            if module.__class__.__name__ == "WanResample" and str(getattr(module, "mode", "")).startswith("upsample"):
+                module.forward = MethodType(_resample_forward_channels_last, module)
+                counts["resample_forwards"] += 1
+
+
+def install_wan_decoder_fast_path(vae: Any, *, conv_dtype: torch.dtype | None, level: str = "exact") -> dict[str, int]:
+    """Install the fast path on ``vae``'s decoder (and ``post_quant_conv``); idempotent.
+
+    ``conv_dtype`` must be the dtype the decode runs under (``decode_autocast_dtype``): with ``None`` the
+    parameters are left alone, because casting them without autocast would change the convolution's
+    arithmetic rather than remove a cast. ``level`` is ``"exact"`` or ``"fused"`` (see the module docstring).
+    """
+    if level not in LEVELS:
+        raise ValueError(f"level must be one of {LEVELS}; got {level!r}")
+    installed = getattr(vae, _INSTALLED_ATTR, None)
+    if installed is not None:
+        if installed.get("level") != level:
+            raise ValueError(f"Wan decoder fast path already installed at level {installed.get('level')!r}")
+        return installed
+    decoder = getattr(vae, "decoder", None)
+    if decoder is None:
+        raise ValueError("Wan decoder fast path requires a decoder module.")
+    if conv_dtype is not None and conv_dtype not in (torch.float16, torch.bfloat16):
+        raise ValueError(f"conv_dtype must be float16, bfloat16 or None; got {conv_dtype!r}")
+    roots = [decoder]
+    post_quant_conv = getattr(vae, "post_quant_conv", None)
+    if isinstance(post_quant_conv, nn.Module):
+        roots.append(post_quant_conv)
+    counts: dict[str, Any] = {
+        "level": level,
+        "conv_params_cast": 0,
+        "upsamples": 0,
+        "persistent_input_buffers": 0,
+        "fused_norm_silu": 0,
+        "channels_last_convs": 0,
+        "resample_forwards": 0,
+    }
+    for root in roots:
+        for module in root.modules():
+            if conv_dtype is not None and isinstance(module, (nn.Conv2d, nn.Conv3d)):
+                if any(p.dtype != conv_dtype for p in module.parameters(recurse=False)):
+                    module.to(dtype=conv_dtype)
+                    counts["conv_params_cast"] += 1
+            if _is_nearest_upsample(module) and not getattr(module, "_vllm_omni_fast_path_upsample", False):
+                module.forward = MethodType(_upsample_forward, module)
+                module._vllm_omni_fast_path_upsample = True
+                counts["upsamples"] += 1
+            if isinstance(module, (WanDistCausalConv3d, WanDistConv2d)):
+                module.reuse_input_buffer = True
+                counts["persistent_input_buffers"] += 1
+    if level == "fused":
+        _install_fused(decoder, post_quant_conv if isinstance(post_quant_conv, nn.Module) else None, counts)
+    setattr(vae, _INSTALLED_ATTR, counts)
+    logger.info("Installed the %s Wan decoder fast path (conv dtype %s): %s", level, conv_dtype, counts)
+    return counts

vllm_omni/diffusion/distributed/autoencoders/wan_decoder_kernels.py

@@ -0,0 +1,278 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused kernels for the Wan VAE decoder on channels_last_3d activations.
+
+Adapted from SGLang's ``sglang.kernels.ops.diffusion`` (``norm/wan_rmsnorm_silu_triton.py`` and
+``layout/nearest_upsample_nhwc_triton.py``, Apache-2.0), reduced to what the LingBot-World streaming decode
+uses. Each kernel comes with a ``can_use_*`` predicate and a torch reference of the eager op chain it
+replaces; the predicate decides, the kernel raises on an input it does not support.
+
+- :func:`wan_rmsnorm_silu`: ``SiLU(F.normalize(x, dim=1) * scale * gamma + bias)`` on a dense channels_last_3d
+  ``[B, C, T, H, W]`` tensor, one program per pixel. fp32 channel statistics; the intermediate dtype
+  boundaries of the eager chain are kept (normalise and ``* scale`` in ``x.dtype``, ``* gamma`` in the
+  promoted dtype, SiLU in fp32) and the result is stored in ``x.dtype``, which is the dtype the following
+  convolution reads under autocast anyway. Not bit-identical to aten (the channel reduction order differs),
+  so it is gated with the bf16 decode option rather than mounted unconditionally.
+- :func:`nearest_upsample_nhwc`: integer-factor nearest upsample of a dense channels_last ``[N, C, H, W]``
+  tensor as a gather, bit-exact vs ``nn.Upsample`` in ``nearest`` and ``nearest-exact`` modes; aten's own NHWC
+  nearest kernel is several times slower than its NCHW one, which is what a channels_last decoder would hit.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+try:  # Triton is optional: CPU tests and non-CUDA platforms use the torch references.
+    import triton
+    import triton.language as tl
+
+    _HAS_TRITON = True
+except ImportError:  # pragma: no cover
+    _HAS_TRITON = False
+
+_MAX_CHANNELS = 1024
+_MAX_INT32 = 2**31 - 1
+_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
+
+
+# --------------------------------------------------------------------------- RMSNorm + SiLU
+def wan_rmsnorm_silu_reference(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | float,
+    scale: float,
+    eps: float = 1e-12,
+    upcast: bool = True,
+) -> torch.Tensor:
+    """The eager norm chain followed by SiLU, cast back to ``x.dtype``.
+
+    ``upcast=True`` is diffusers' ``WanRMS_norm`` (normalise in fp32, round to ``x.dtype``); ``upcast=False``
+    is vLLM-Omni's ``RMSNormVAE`` (``F.normalize`` on ``x`` itself, so every step rounds to ``x.dtype``).
+    """
+    if upcast:
+        normalized = F.normalize(x.float(), dim=1, eps=eps).to(x.dtype)
+    else:
+        normalized = F.normalize(x, dim=1, eps=eps)
+    y = normalized * scale * gamma + bias
+    return F.silu(y).to(x.dtype)
+
+
+def _dense_channels_last_3d(x: torch.Tensor) -> bool:
+    if x.dim() != 5 or x.shape[1] <= 1:
+        return False
+    b, c, t, h, w = x.shape
+    return x.stride() == (t * h * w * c, 1, h * w * c, w * c, c)
+
+
+def can_use_wan_rmsnorm_silu(x: torch.Tensor, gamma: torch.Tensor, bias: torch.Tensor | float | None) -> bool:
+    return (
+        _HAS_TRITON
+        and x.is_cuda
+        and not (torch.is_grad_enabled() and x.requires_grad)
+        and x.dtype in _SUPPORTED_DTYPES
+        and x.numel() > 0
+        and 0 < x.shape[1] <= _MAX_CHANNELS
+        and _dense_channels_last_3d(x)
+        and gamma.is_cuda
+        and gamma.device == x.device
+        and gamma.numel() == x.shape[1]
+        and (gamma.dtype == x.dtype or gamma.dtype == torch.float32)
+        and (
+            not isinstance(bias, torch.Tensor)
+            or (bias.device == x.device and bias.numel() == x.shape[1] and bias.dtype in (x.dtype, torch.float32))
+        )
+    )
+
+
+if _HAS_TRITON:
+
+    @triton.jit
+    def _wan_rmsnorm_silu_kernel(
+        x_ptr,
+        gamma_ptr,
+        bias_ptr,
+        out_ptr,
+        channels: tl.constexpr,
+        rms_scale,
+        eps,
+        has_bias: tl.constexpr,
+        round_affine: tl.constexpr,
+        upcast: tl.constexpr,
+        block_c: tl.constexpr,
+    ):
+        row = tl.program_id(0).to(tl.int64)
+        offsets = tl.arange(0, block_c)
+        mask = offsets < channels
+        # Dense channels_last_3d stores each pixel as one contiguous channel row.
+        row_offsets = row * channels + offsets
+        x = tl.load(x_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32)
+        norm = tl.sqrt(tl.sum(x * x, axis=0))
+        if upcast:
+            # WanRMS_norm: fp32 normalise, one rounding to x.dtype.
+            y = (x / tl.maximum(norm, eps)).to(x_ptr.dtype.element_ty)
+        else:
+            # RMSNormVAE: F.normalize on x itself, so the norm, the clamp and the quotient each round to x.dtype.
+            denom = tl.maximum(norm.to(x_ptr.dtype.element_ty).to(tl.float32), eps)
+            denom = denom.to(x_ptr.dtype.element_ty).to(tl.float32)
+            y = (x / denom).to(x_ptr.dtype.element_ty)
+        # Then the eager boundaries: * scale in x.dtype, * gamma / + bias in the promoted dtype, SiLU in fp32,
+        # stored in x.dtype.
+        y = (y * rms_scale).to(x_ptr.dtype.element_ty)
+        gamma = tl.load(gamma_ptr + offsets, mask=mask, other=1.0).to(tl.float32)
+        y = y.to(tl.float32) * gamma
+        if round_affine:
+            # gamma (and bias) share x's dtype: the eager product and sum are rounded to it.
+            y = y.to(x_ptr.dtype.element_ty).to(tl.float32)
+        if has_bias:
+            bias = tl.load(bias_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+            y = y + bias
+            if round_affine:
+                y = y.to(x_ptr.dtype.element_ty).to(tl.float32)
+        y = y * tl.sigmoid(y)
+        tl.store(out_ptr + row_offsets, y, mask=mask)
+
+
+def wan_rmsnorm_silu(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | float | None,
+    scale: float,
+    eps: float = 1e-12,
+    upcast: bool = True,
+) -> torch.Tensor:
+    """Fused ``SiLU(normalize(x) * scale * gamma + bias)`` on a dense channels_last_3d tensor; raises otherwise."""
+    if not can_use_wan_rmsnorm_silu(x, gamma, bias):
+        raise ValueError(
+            "wan_rmsnorm_silu needs a CUDA, dense channels_last_3d, bf16/fp16/fp32 [B, C, T, H, W] tensor with "
+            f"C <= {_MAX_CHANNELS}; got {x.device.type} {tuple(x.shape)} {x.dtype} strides {tuple(x.stride())}"
+        )
+    b, c, t, h, w = x.shape
+    out = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=x.dtype)
+    has_bias = isinstance(bias, torch.Tensor)
+    bias_arg = bias.reshape(-1) if isinstance(bias, torch.Tensor) else gamma
+    # With affine parameters in x's dtype the eager chain rounds after ``* gamma`` and ``+ bias``; fp32
+    # parameters promote the chain to fp32 and the only rounding left is the store.
+    round_affine = gamma.dtype == x.dtype and x.dtype != torch.float32
+    block_c = triton.next_power_of_2(c)
+    num_warps = 1 if block_c <= 64 else 4 if block_c <= 512 else 8
+    with torch.get_device_module(x.device).device(x.device):
+        _wan_rmsnorm_silu_kernel[(b * t * h * w,)](
+            x,
+            gamma.reshape(-1),
+            bias_arg,
+            out,
+            c,
+            float(scale),
+            float(eps),
+            has_bias,
+            round_affine,
+            upcast,
+            block_c,
+            num_warps=num_warps,
+        )
+    return out
+
+
+# --------------------------------------------------------------------------- nearest upsample, NHWC
+def _integer_scale(scale) -> tuple[int, int] | None:
+    if isinstance(scale, bool):
+        return None
+    if isinstance(scale, (int, float)):
+        scale = (scale, scale)
+    if not isinstance(scale, (tuple, list)) or len(scale) != 2:
+        return None
+    out = []
+    for s in scale:
+        if isinstance(s, bool) or not isinstance(s, (int, float)):
+            return None
+        f = float(s)
+        if not math.isfinite(f) or f < 1.0 or f != int(f):
+            return None
+        out.append(int(f))
+    return out[0], out[1]
+
+
+def canonical_nhwc(x: torch.Tensor) -> bool:
+    """Dense channels_last with ``C > 1`` and canonical strides on every dim, size-1 dims included."""
+    if x.dim() != 4:
+        return False
+    _, c, h, w = x.shape
+    return c > 1 and x.stride() == (h * w * c, 1, w * c, c)
+
+
+def can_use_nearest_upsample_nhwc(x: torch.Tensor, scale_factor, mode: str) -> bool:
+    return (
+        _HAS_TRITON
+        and isinstance(x, torch.Tensor)
+        and x.is_cuda
+        and not (torch.is_grad_enabled() and x.requires_grad)
+        and mode in ("nearest", "nearest-exact")
+        and x.dim() == 4
+        and x.numel() > 0
+        and x.dtype in _SUPPORTED_DTYPES
+        and canonical_nhwc(x)
+        and _integer_scale(scale_factor) is not None
+    )
+
+
+if _HAS_TRITON:
+
+    @triton.jit
+    def _nearest_upsample_nhwc_kernel(
+        x_ptr, out_ptr, total, channels, out_h, out_w, fh, fw, sxn, sxh, sxw, idx64: tl.constexpr, block: tl.constexpr
+    ):
+        pid = tl.program_id(0)
+        if idx64:
+            offs = pid.to(tl.int64) * block + tl.arange(0, block).to(tl.int64)
+        else:
+            offs = pid * block + tl.arange(0, block)
+        mask = offs < total
+        # Output is dense NHWC: offs = ((n * out_h + h) * out_w + w) * channels + c.
+        c = offs % channels
+        t = offs // channels
+        w = t % out_w
+        t = t // out_w
+        h = t % out_h
+        n = t // out_h
+        src = n * sxn + (h // fh) * sxh + (w // fw) * sxw + c
+        vals = tl.load(x_ptr + src, mask=mask)
+        tl.store(out_ptr + offs, vals, mask=mask)
+
+
+def nearest_upsample_nhwc(x: torch.Tensor, scale_factor) -> torch.Tensor:
+    """Integer-factor nearest upsample of a canonical channels_last ``[N, C, H, W]`` tensor, dense channels_last out."""
+    factors = _integer_scale(scale_factor)
+    if factors is None or not can_use_nearest_upsample_nhwc(x, scale_factor, "nearest"):
+        raise ValueError(
+            "nearest_upsample_nhwc needs a CUDA, canonical channels_last 4D bf16/fp16/fp32 tensor and an integer "
+            f"scale; got {x.device.type} {tuple(x.shape)} strides {tuple(x.stride())} scale {scale_factor!r}"
+        )
+    fh, fw = factors
+    n, c, h, w = x.shape
+    out_h, out_w = h * fh, w * fw
+    out = torch.empty((n, c, out_h, out_w), device=x.device, dtype=x.dtype, memory_format=torch.channels_last)
+    total = out.numel()
+    if total == 0:
+        return out
+    sxn, _, sxh, sxw = x.stride()
+    block = 1024
+    with torch.get_device_module(x.device).device(x.device):
+        _nearest_upsample_nhwc_kernel[(triton.cdiv(total, block),)](
+            x,
+            out,
+            total,
+            c,
+            out_h,
+            out_w,
+            fh,
+            fw,
+            sxn,
+            sxh,
+            sxw,
+            idx64=total >= _MAX_INT32 or x.numel() >= _MAX_INT32,
+            block=block,
+        )
+    return out

vllm_omni/diffusion/distributed/autoencoders/wan_decoder_utils.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Utility functions for the streaming Wan decoder fast path."""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.distributed.autoencoders.wan_decoder_kernels import (
+    can_use_nearest_upsample_nhwc,
+    canonical_nhwc,
+    nearest_upsample_nhwc,
+)
+
+_CACHE_T = 2  # diffusers.models.autoencoders.autoencoder_kl_wan.CACHE_T
+
+
+def _upsample_forward(self: nn.Upsample, x: torch.Tensor) -> torch.Tensor:
+    # ``WanUpsample.forward`` is ``super().forward(x.float()).type_as(x)``; a nearest gather never changes a
+    # value, so running it on ``x`` itself is the same tensor without the two casts.
+    return nn.Upsample.forward(self, x)
+
+
+def _is_nearest_upsample(module: nn.Module) -> bool:
+    return (
+        module.__class__.__name__ == "WanUpsample"
+        and isinstance(module, nn.Upsample)
+        and str(getattr(module, "mode", "")) in ("nearest", "nearest-exact")
+    )
+
+
+def _upsample_forward_channels_last(self: nn.Upsample, x: torch.Tensor) -> torch.Tensor:
+    if x.dim() == 4 and self.size is None and can_use_nearest_upsample_nhwc(x, self.scale_factor, self.mode):
+        return nearest_upsample_nhwc(x, self.scale_factor)
+    return nn.Upsample.forward(self, x)
+
+
+def _dense_channels_last_3d(x: torch.Tensor) -> bool:
+    if x.dim() != 5 or x.shape[1] <= 1:
+        return False
+    b, c, t, h, w = x.shape
+    return x.stride() == (t * h * w * c, 1, h * w * c, w * c, c)
+
+
+def _interleave_time_pairs(x: torch.Tensor, b: int, c: int, t: int, h: int, w: int) -> torch.Tensor:
+    """``[B, 2C, T, H, W] -> [B, C, 2T, H, W]``: the time-conv's two channel halves interleaved along time.
+
+    Same values as the eager ``reshape / stack / reshape``; on a dense channels_last_3d input the result is
+    written channels_last_3d with one copy instead of materialising it channels-first.
+    """
+    if _dense_channels_last_3d(x):
+        out = torch.empty((b, c, t * 2, h, w), device=x.device, dtype=x.dtype, memory_format=torch.channels_last_3d)
+        out.view(b, c, t, 2, h, w).copy_(x.view(b, 2, c, t, h, w).permute(0, 2, 3, 1, 4, 5))
+        return out
+    x = x.reshape(b, 2, c, t, h, w)
+    x = torch.stack((x[:, 0], x[:, 1]), 3)
+    return x.reshape(b, c, t * 2, h, w)
+
+
+def _frames_nhwc(x: torch.Tensor) -> torch.Tensor:
+    """``[B, C, T, H, W] -> [B*T, C, H, W]`` for the 2D resample, with canonical NHWC strides when the data is
+    laid out that way (a size-1 batch leaves a degenerate stride that aten's layout test rejects)."""
+    b, c, t, h, w = x.shape
+    frames = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
+    if frames.is_contiguous(memory_format=torch.channels_last) and not canonical_nhwc(frames) and frames.stride(1) == 1:
+        frames = frames.as_strided((b * t, c, h, w), (h * w * c, 1, w * c, c))
+    return frames
+
+
+def _resample_forward_channels_last(self: nn.Module, x: torch.Tensor, feat_cache=None, feat_idx=[0]) -> torch.Tensor:  # noqa: B006
+    """``WanResample.forward`` for the upsample modes with the layout-preserving interleave and frame view."""
+    b, c, t, h, w = x.size()
+    if self.mode == "upsample3d" and feat_cache is not None:
+        idx = feat_idx[0]
+        if feat_cache[idx] is None:
+            feat_cache[idx] = "Rep"
+            feat_idx[0] += 1
+        else:
+            cache_x = x[:, :, -_CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
+                cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+            if feat_cache[idx] == "Rep":
+                x = self.time_conv(x)
+            else:
+                x = self.time_conv(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+            x = _interleave_time_pairs(x, b, c, t, h, w)
+    t = x.shape[2]
+    frames = self.resample(_frames_nhwc(x))
+    return frames.view(b, t, frames.size(1), frames.size(2), frames.size(3)).permute(0, 2, 1, 3, 4)
+
+
+def _persistent_input_buffer(module: nn.Module, shape: tuple[int, ...], reference: torch.Tensor) -> torch.Tensor:
+    """One zero-initialised input buffer per conv, reallocated only when the shape, dtype or device changes.
+
+    The callers write only the activation interior and the halo slots, so the padding rows and columns stay
+    zero for the buffer's lifetime and are never filled again.
+    """
+    memory_format = getattr(module, "input_memory_format", torch.contiguous_format)
+    buf = module._input_buf
+    if (
+        buf is None
+        or buf.shape != shape
+        or buf.dtype != reference.dtype
+        or buf.device != reference.device
+        or not buf.is_contiguous(memory_format=memory_format)
+    ):
+        buf = torch.empty(shape, dtype=reference.dtype, device=reference.device, memory_format=memory_format).zero_()
+        module._input_buf = buf
+    return buf

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -34,6 +34,7 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
 from diffusers.models.autoencoders.vae import DecoderOutput
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.distributed.autoencoders.wan_decoder_utils import _persistent_input_buffer
 from vllm_omni.diffusion.models.interface import DecodedChunkConsumer
 
 logger = init_logger(__name__)
@@ -241,7 +242,10 @@ def gather_and_trim_extent(
 
 
 def _ensure_recv_buf(recv_buf: torch.Tensor | None, reference: torch.Tensor) -> torch.Tensor:
-    memory_format = _halo_memory_format(reference)
+    # Send and receive buffers are dense in the default order whatever the activation's layout: every
+    # backend accepts that (gloo rejects a channels_last buffer), the halo is a single row or column, and the
+    # copy into the assembled input handles the layout.
+    memory_format = torch.contiguous_format
     if (
         recv_buf is None
         or recv_buf.shape != reference.shape
@@ -271,7 +275,7 @@ def _halo_exchange_p2p(
     p2p_ops = []
     if rank > 0:
         prev_rank = _global_rank(group, rank - 1)
-        top_row = top_row_ref.contiguous(memory_format=_halo_memory_format(top_row_ref))
+        top_row = top_row_ref.contiguous()
         p2p_ops.append(dist.P2POp(dist.irecv, recv_top_buf, prev_rank, group))
         p2p_ops.append(dist.P2POp(dist.isend, top_row, prev_rank, group))
     else:
@@ -279,7 +283,7 @@ def _halo_exchange_p2p(
 
     if rank < world_size - 1:
         next_rank = _global_rank(group, rank + 1)
-        bottom_row = bottom_row_ref.contiguous(memory_format=_halo_memory_format(bottom_row_ref))
+        bottom_row = bottom_row_ref.contiguous()
         p2p_ops.append(dist.P2POp(dist.isend, bottom_row, next_rank, group))
         p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group))
     else:
@@ -402,18 +406,27 @@ class WanDistConv2d(nn.Conv2d):
         self._halo_recv_top_buf: torch.Tensor | None = None
         self._halo_recv_bottom_buf: torch.Tensor | None = None
         self._trim_cache: dict[int, tuple[int, int, int]] = {}
+        # Opt-in (install_wan_decoder_fast_path): assemble the conv input in one buffer kept across calls
+        # instead of allocating, zero-filling and concatenating a fresh one every call. Same bytes.
+        self.reuse_input_buffer = False
+        self.input_memory_format = torch.contiguous_format
+        self._input_buf: torch.Tensor | None = None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.pad(x, self._non_split_padding)
-        x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
-            x,
-            group=self.group,
-            halo_size=self.halo_size,
-            split_dim=self.split_dim,
-            recv_top_buf=self._halo_recv_top_buf,
-            recv_bottom_buf=self._halo_recv_bottom_buf,
-        )
-        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
+        if self.reuse_input_buffer:
+            x_padded, local_extent = self._assemble_input(x)
+        else:
+            x = F.pad(x, self._non_split_padding)
+            x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
+                x,
+                group=self.group,
+                halo_size=self.halo_size,
+                split_dim=self.split_dim,
+                recv_top_buf=self._halo_recv_top_buf,
+                recv_bottom_buf=self._halo_recv_bottom_buf,
+            )
+            local_extent = x.shape[self.split_tensor_dim]
+        shift, start, upper_bound = self._get_trim_params(local_extent)
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
@@ -425,6 +438,47 @@ class WanDistConv2d(nn.Conv2d):
         out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
         return _zero_invalid_extent(out, split_dim=self.split_dim)
 
+    def _assemble_input(self, x: torch.Tensor) -> tuple[torch.Tensor, int]:
+        """The padded, halo-exchanged conv input built in one buffer that persists across calls.
+
+        Byte for byte what ``F.pad`` followed by ``halo_exchange`` produced -- the non-split zero padding
+        around the activation, the neighbours' halo rows or columns outside it -- without the two
+        allocations and the three copies those took per call. Returns the assembled tensor and the extent
+        along the split dimension before halos (what the trim parameters are keyed by).
+        """
+        left, right, top, bottom = self._non_split_padding
+        rank, world_size = _rank_world(self.group)
+        halo = self.halo_size if world_size > 1 else 0
+        dim = self.split_tensor_dim
+        halo_h = halo if self.split_dim == "height" else 0
+        halo_w = halo if self.split_dim == "width" else 0
+        batch, channels, height, width = x.shape
+        total_height = halo_h + top + height + bottom + halo_h
+        total_width = halo_w + left + width + right + halo_w
+        buf = _persistent_input_buffer(self, (batch, channels, total_height, total_width), x)
+        buf[:, :, halo_h + top : halo_h + top + height, halo_w + left : halo_w + left + width].copy_(x)
+        local_extent = buf.shape[dim] - 2 * halo
+        if halo:
+            # What halo_exchange sends: the padded activation's first and last ``halo`` rows along the
+            # split dimension, non-split padding included; the neighbours' rows land in the outer slots.
+            padded = _narrow_along_dim(buf, dim, halo, local_extent)
+            top_row_ref = _narrow_along_dim(padded, dim, 0, halo)
+            bottom_row_ref = _narrow_along_dim(padded, dim, local_extent - halo, halo)
+            self._halo_recv_top_buf = _ensure_recv_buf(self._halo_recv_top_buf, top_row_ref)
+            self._halo_recv_bottom_buf = _ensure_recv_buf(self._halo_recv_bottom_buf, bottom_row_ref)
+            _halo_exchange_p2p(
+                rank=rank,
+                world_size=world_size,
+                group=self.group,
+                top_row_ref=top_row_ref,
+                bottom_row_ref=bottom_row_ref,
+                recv_top_buf=self._halo_recv_top_buf,
+                recv_bottom_buf=self._halo_recv_bottom_buf,
+            )
+            _narrow_along_dim(buf, dim, 0, halo).copy_(self._halo_recv_top_buf)
+            _narrow_along_dim(buf, dim, buf.shape[dim] - halo, halo).copy_(self._halo_recv_bottom_buf)
+        return buf, local_extent
+
     def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
         trim_params = self._trim_cache.get(local_extent)
         if trim_params is None:
@@ -503,6 +557,11 @@ class WanDistCausalConv3d(nn.Conv3d):
         self._halo_recv_top_buf: torch.Tensor | None = None
         self._halo_recv_bottom_buf: torch.Tensor | None = None
         self._trim_cache: dict[int, tuple[int, int, int]] = {}
+        # Opt-in (install_wan_decoder_fast_path): assemble the conv input in one buffer kept across calls
+        # instead of allocating, zero-filling and concatenating a fresh one every call. Same bytes.
+        self.reuse_input_buffer = False
+        self.input_memory_format = torch.contiguous_format
+        self._input_buf: torch.Tensor | None = None
 
     def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
         x_padded, local_extent = self._assemble_input(x, cache_x)
@@ -550,22 +609,33 @@ class WanDistCausalConv3d(nn.Conv3d):
         total_frames = pad_t0 + cache_frames + frames + pad_t1
         total_height = top + height + pad_h1 + halo_h
         total_width = left + width + pad_w1 + halo_w
-        buf = torch.empty((batch, channels, total_frames, total_height, total_width), dtype=x.dtype, device=x.device)
-        # Zero only what the conv must read as zeros: the causal frames in
-        # front (and any trailing ones) and the spatial padding rows/columns.
-        # The halo columns are filled below.
-        if pad_t0:
-            buf[:, :, :pad_t0].zero_()
-        if pad_t1:
-            buf[:, :, total_frames - pad_t1 :].zero_()
-        if pad_h0:
-            buf[:, :, :, halo_h:top].zero_()
-        if pad_h1:
-            buf[:, :, :, total_height - halo_h - pad_h1 : total_height - halo_h].zero_()
-        if pad_w0:
-            buf[:, :, :, :, halo_w:left].zero_()
-        if pad_w1:
-            buf[:, :, :, :, total_width - halo_w - pad_w1 : total_width - halo_w].zero_()
+        shape = (batch, channels, total_frames, total_height, total_width)
+        if self.reuse_input_buffer:
+            # The copies below touch only the activation interior and the halo slots, so the spatial padding
+            # stays zero for the buffer's lifetime. The causal front frames alternate between zeros (a
+            # session's first frame) and cached frames, so they are re-zeroed whenever this call needs zeros.
+            buf = _persistent_input_buffer(self, shape, x)
+            if pad_t0:
+                buf[:, :, :pad_t0].zero_()
+            if pad_t1:
+                buf[:, :, total_frames - pad_t1 :].zero_()
+        else:
+            buf = torch.empty(shape, dtype=x.dtype, device=x.device)
+            # Zero only what the conv must read as zeros: the causal frames in
+            # front (and any trailing ones) and the spatial padding rows/columns.
+            # The halo columns are filled below.
+            if pad_t0:
+                buf[:, :, :pad_t0].zero_()
+            if pad_t1:
+                buf[:, :, total_frames - pad_t1 :].zero_()
+            if pad_h0:
+                buf[:, :, :, halo_h:top].zero_()
+            if pad_h1:
+                buf[:, :, :, total_height - halo_h - pad_h1 : total_height - halo_h].zero_()
+            if pad_w0:
+                buf[:, :, :, :, halo_w:left].zero_()
+            if pad_w1:
+                buf[:, :, :, :, total_width - halo_w - pad_w1 : total_width - halo_w].zero_()
         interior = buf[:, :, pad_t0 : pad_t0 + cache_frames + frames, top : top + height, left : left + width]
         if cache_frames:
             interior[:, :, :cache_frames].copy_(cache_x)

vllm_omni/diffusion/models/lingbot_world/pipeline.py

@@ -26,6 +26,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
+from vllm_omni.diffusion.distributed.autoencoders.wan_decoder_fast_path import install_wan_decoder_fast_path
 from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import install_wan_spatial_shard_decode
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.interaction.mixin import InteractionMixin
@@ -61,6 +62,7 @@ from vllm_omni.diffusion.models.lingbot_world.transformer import (
     LingBotAttentionCache,
     LingBotTransformerCache,
 )
+from vllm_omni.diffusion.models.lingbot_world.utils import _vae_decode_fast_path
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.utils import _load_json
@@ -338,7 +340,22 @@ def _uint8_frames(video: torch.Tensor) -> np.ndarray:
     # the decoder dtype, then widen, scale and round.
     frames = (video[0] / 2 + 0.5).clamp(0, 1).permute(1, 2, 3, 0)
     frames = frames.float().mul_(255).round_().to(torch.uint8)
-    return np.ascontiguousarray(frames.cpu().numpy())
+    pixels = frames.cpu().numpy()
+    frame_count, height, width, channels = pixels.shape
+    # A planar GPU result stays planar after the host copy. NumPy's generic
+    # contiguous conversion is slow for this three-channel interleave; copying
+    # each plane into its output channel preserves every byte and vectorizes it.
+    if channels == 3 and pixels.strides == (
+        height * width,
+        width,
+        1,
+        frame_count * height * width,
+    ):
+        contiguous = np.empty(pixels.shape, dtype=np.uint8)
+        for channel in range(3):
+            contiguous[..., channel] = pixels[..., channel]
+        return contiguous
+    return np.ascontiguousarray(pixels)
 
 
 @functools.lru_cache(maxsize=1)
@@ -695,6 +712,12 @@ class LingBotWorldCausalDMDPipeline(
         sequence_parallel_size = int(getattr(parallel_config, "sequence_parallel_size", 1) or 1)
         if sequence_parallel_size > 1 and vae_sharding:
             self._install_sharded_vae_decode(sequence_parallel_size)
+        fast_path_level = _vae_decode_fast_path(model_config)
+        if fast_path_level is not None:
+            # Install after sharding so its conv wrappers get persistent buffers. The VAE parameters
+            # already use the stage dtype, so no autocast parameter cast is needed.
+            counts = install_wan_decoder_fast_path(self.vae, conv_dtype=None, level=fast_path_level)
+            logger.info("LingBot World VAE decode fast path installed: %s", counts)
         self.setup_diffusion_pipeline_profiler(
             profiler_targets=[
                 "vae.encode",

vllm_omni/diffusion/models/lingbot_world/utils.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Configuration helpers for the LingBot World pipeline."""
+
+from __future__ import annotations
+
+
+def _vae_decode_fast_path(model_config: dict) -> str | None:
+    """Resolve ``lingbot_vae_decode_fast_path`` from the model config: ``None``, ``"exact"`` or ``"fused"``.
+
+    Off by default. ``true``/``exact`` installs the decoder's exact fast path after the decode dtype and the
+    spatial shard are configured (see ``wan_decoder_fast_path``): no output bit changes, only launches
+    disappear. ``fused`` adds the channels_last decoder with the fused RMSNorm+SiLU kernel, which is not
+    bit-exact and requires subjective video quality review before enabling.
+    """
+    value = model_config.get("lingbot_vae_decode_fast_path", False)
+    if isinstance(value, bool) or value is None:
+        return "exact" if value else None
+    name = str(value).strip().lower()
+    if name in ("", "0", "false", "no", "off", "none"):
+        return None
+    if name in ("1", "true", "yes", "on", "exact"):
+        return "exact"
+    if name == "fused":
+        return "fused"
+    raise ValueError(f"lingbot_vae_decode_fast_path must be off, exact or fused; got {value!r}")