From ba3d48e1c783083634b3f70acbb8f2512658e65f Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Thu, 24 Sep 2026 18:56:44 +0800 Subject: [PATCH] [LingBot World] Speed up streaming VAE decode and RGB frame delivery (#7749) Signed-off-by: tzhouam Co-authored-by: Claude Fable 5.1 --- .../distributed/test_wan_decoder_fast_path.py | 182 ++++++++++++ .../autoencoders/wan_decoder_fast_path.py | 192 ++++++++++++ .../autoencoders/wan_decoder_kernels.py | 278 ++++++++++++++++++ .../autoencoders/wan_decoder_utils.py | 114 +++++++ .../autoencoders/wan_spatial_shard.py | 128 ++++++-- .../models/lingbot_world/pipeline.py | 25 +- .../diffusion/models/lingbot_world/utils.py | 26 ++ 7 files changed, 915 insertions(+), 30 deletions(-) create mode 100644 tests/diffusion/distributed/test_wan_decoder_fast_path.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_decoder_fast_path.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_decoder_kernels.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_decoder_utils.py create mode 100644 vllm_omni/diffusion/models/lingbot_world/utils.py diff --git a/tests/diffusion/distributed/test_wan_decoder_fast_path.py b/tests/diffusion/distributed/test_wan_decoder_fast_path.py new file mode 100644 index 000000000..f9fb24976 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_decoder_fast_path.py b/vllm_omni/diffusion/distributed/autoencoders/wan_decoder_fast_path.py new file mode 100644 index 000000000..5e4a3c583 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_decoder_kernels.py b/vllm_omni/diffusion/distributed/autoencoders/wan_decoder_kernels.py new file mode 100644 index 000000000..33dccfd45 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_decoder_utils.py b/vllm_omni/diffusion/distributed/autoencoders/wan_decoder_utils.py new file mode 100644 index 000000000..cd857a8b3 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py index e3f0ef5c3..817aee6c5 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py +++ b/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) diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index fec175e69..b236e44b5 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/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", diff --git a/vllm_omni/diffusion/models/lingbot_world/utils.py b/vllm_omni/diffusion/models/lingbot_world/utils.py new file mode 100644 index 000000000..94156045a --- /dev/null +++ b/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}")