From 5308222abbf4c1a61a11e8a846407ade8d8a2b50 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Tue, 11 Aug 2026 06:47:24 +0800 Subject: [PATCH] [Perf][Diffusion] Fuse lossless Wan VAE data movement Signed-off-by: Hongsheng Liu --- .../distributed/test_autoencoder_kl_wan.py | 20 + .../kernels/test_wan_vae_data_movement.py | 274 ++++++++++++ .../diffusion/test_diffusion_model_runner.py | 8 + .../autoencoders/autoencoder_kl_wan.py | 13 + .../autoencoders/wan_vae_data_movement.py | 302 +++++++++++++ vllm_omni/diffusion/kernels/__init__.py | 2 + .../kernels/wan_vae_data_movement.py | 405 ++++++++++++++++++ .../models/helios/pipeline_helios.py | 2 + .../lingbot_video/pipeline_lingbot_video.py | 2 + .../worker/diffusion_model_runner.py | 4 + 10 files changed, 1032 insertions(+) create mode 100644 tests/diffusion/kernels/test_wan_vae_data_movement.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/wan_vae_data_movement.py create mode 100644 vllm_omni/diffusion/kernels/__init__.py create mode 100644 vllm_omni/diffusion/kernels/wan_vae_data_movement.py diff --git a/tests/diffusion/distributed/test_autoencoder_kl_wan.py b/tests/diffusion/distributed/test_autoencoder_kl_wan.py index 2ea1c1214..209a72e66 100644 --- a/tests/diffusion/distributed/test_autoencoder_kl_wan.py +++ b/tests/diffusion/distributed/test_autoencoder_kl_wan.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import pytest import torch @@ -41,3 +43,21 @@ def test_wan_vae_execution_context_uses_platform_autocast(mocker): dtype=torch.bfloat16, enabled=True, ) + + +def test_prepare_pipeline_installs_data_movement_before_compile(monkeypatch): + model = _DummyOmniAutoencoderKLWan(dtype=torch.float32) + installed = [] + monkeypatch.setattr( + "vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover", + lambda pipeline: SimpleNamespace(vaes=[model]), + ) + monkeypatch.setattr( + wan_vae_module, + "install_wan_vae_data_movement", + lambda vae: installed.append(vae), + ) + + wan_vae_module.prepare_pipeline_wan_vae_data_movement(torch.nn.Identity()) + + assert installed == [model] diff --git a/tests/diffusion/kernels/test_wan_vae_data_movement.py b/tests/diffusion/kernels/test_wan_vae_data_movement.py new file mode 100644 index 000000000..8305b31c1 --- /dev/null +++ b/tests/diffusion/kernels/test_wan_vae_data_movement.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F +from diffusers.models.autoencoders import AutoencoderKLWan + +from vllm_omni.diffusion.distributed.autoencoders import wan_vae_data_movement +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan +from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import ( + _cache_payload, + _causal_padding, + _is_dynamic_spatial_conv, + _run_cached_causal_conv, + install_wan_vae_data_movement, +) +from vllm_omni.diffusion.kernels.wan_vae_data_movement import cat_pad_5d, dup_up3d_add + +pytestmark = pytest.mark.core_model +cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _dense_5d(shape: tuple[int, ...], dtype: torch.dtype, channels_last: bool) -> torch.Tensor: + tensor = torch.randn(shape, device="cuda", dtype=dtype) + if channels_last: + return tensor.contiguous(memory_format=torch.channels_last_3d) + return tensor.contiguous() + + +def _reference_cat_pad( + x: torch.Tensor, + cache: torch.Tensor | None, + padding: tuple[int, ...], +) -> torch.Tensor: + reference_padding = list(padding) + if cache is not None: + x = torch.cat([cache, x], dim=2) + reference_padding[4] -= cache.shape[2] + return F.pad(x, reference_padding) + + +@pytest.mark.cuda +@cuda_only +@torch.no_grad() +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("channels_last", [False, True]) +@pytest.mark.parametrize( + "channels,frames,height,width,cache_frames,padding", + [ + (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)), + (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)), + (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)), + (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)), + (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)), + ], +) +def test_cat_pad_5d_is_bitwise_and_layout_exact( + dtype: torch.dtype, + channels_last: bool, + channels: int, + frames: int, + height: int, + width: int, + cache_frames: int, + padding: tuple[int, ...], +) -> None: + torch.cuda.manual_seed(0) + x = _dense_5d((1, channels, frames, height, width), dtype, channels_last) + cache = None + if cache_frames: + pad_height, pad_width = padding[2], padding[0] + buffer = _dense_5d( + (1, channels, cache_frames, height + 2 * pad_height, width + 2 * pad_width), + dtype, + channels_last, + ) + cache = buffer[:, :, :, pad_height : pad_height + height, pad_width : pad_width + width] + + reference = _reference_cat_pad(x, cache, padding) + output = cat_pad_5d(x, cache, padding) + assert output is not None + assert output.shape == reference.shape + assert output.stride() == reference.stride() + assert torch.equal(output, reference) + + pair = cat_pad_5d(x, cache, padding, keep_cache_frames=2) + assert pair is not None + output_with_cache, next_cache = pair + assert torch.equal(output_with_cache, reference) + pad_height, pad_width = padding[2], padding[0] + keep_frames = min(2, reference.shape[2]) + expected_cache = reference[ + :, + :, + -keep_frames:, + pad_height : pad_height + height, + pad_width : pad_width + width, + ] + assert next_cache.shape == expected_cache.shape + assert torch.equal(next_cache, expected_cache) + assert next_cache.is_contiguous(memory_format=torch.channels_last_3d) is channels_last + + +@pytest.mark.cuda +@cuda_only +@torch.no_grad() +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("main_layout", ["contiguous", "channels_last", "permuted"]) +@pytest.mark.parametrize( + "in_channels,out_channels,frames,height,width,factor_temporal,factor_spatial,drop_first", + [ + (128, 64, 1, 10, 14, 2, 2, False), + (128, 64, 1, 10, 14, 2, 2, True), + (64, 32, 2, 10, 14, 1, 2, False), + ], +) +def test_dup_up3d_add_is_bitwise_and_layout_exact( + dtype: torch.dtype, + main_layout: str, + in_channels: int, + out_channels: int, + frames: int, + height: int, + width: int, + factor_temporal: int, + factor_spatial: int, + drop_first: bool, +) -> None: + torch.cuda.manual_seed(0) + repeats = out_channels * factor_temporal * factor_spatial * factor_spatial // in_channels + source = _dense_5d((1, in_channels, frames, height, width), dtype, main_layout == "channels_last") + out_frames = frames * factor_temporal - (factor_temporal - 1 if drop_first else 0) + shape = (1, out_channels, out_frames, height * factor_spatial, width * factor_spatial) + if main_layout == "channels_last": + main = _dense_5d(shape, dtype, True) + elif main_layout == "permuted": + main = torch.randn( + (shape[0], shape[2], shape[1], shape[3], shape[4]), + device="cuda", + dtype=dtype, + ).permute(0, 2, 1, 3, 4) + else: + main = _dense_5d(shape, dtype, False) + + shortcut = source.repeat_interleave(repeats, dim=1) + shortcut = shortcut.view( + 1, + out_channels, + factor_temporal, + factor_spatial, + factor_spatial, + frames, + height, + width, + ) + shortcut = shortcut.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + shortcut = shortcut.view(shape[0], shape[1], frames * factor_temporal, shape[3], shape[4]) + if drop_first: + shortcut = shortcut[:, :, factor_temporal - 1 :] + reference = main + shortcut + + output = dup_up3d_add( + main, + source, + factor_temporal, + factor_spatial, + repeats, + drop_first, + ) + assert output is not None + assert output.shape == reference.shape + assert output.stride() == reference.stride() + assert torch.equal(output, reference) + + +@pytest.mark.cuda +@cuda_only +@torch.no_grad() +@pytest.mark.parametrize("channels_last", [False, True]) +@pytest.mark.parametrize("temporal_only", [False, True]) +def test_cached_conv_chunk_loop_matches_reference(monkeypatch, channels_last: bool, temporal_only: bool) -> None: + from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T, WanCausalConv3d + + torch.cuda.manual_seed(0) + channels = 64 + if temporal_only: + conv = WanCausalConv3d(channels, 2 * channels, (3, 1, 1), padding=(1, 0, 0)) + else: + conv = WanCausalConv3d(channels, channels, 3, padding=1) + conv = conv.to(device="cuda", dtype=torch.float32) + chunks = [_dense_5d((1, channels, 1, 10, 14), torch.float32, channels_last) for _ in range(4)] + + def run(force_reference: bool, start): + cache = [start] + outputs = [] + original = wan_vae_data_movement.cat_pad_5d + if force_reference: + monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", None) + try: + for chunk in chunks: + outputs.append(_run_cached_causal_conv(conv, chunk, cache, 0)) + finally: + monkeypatch.setattr(wan_vae_data_movement, "cat_pad_5d", original) + return outputs, cache[0] + + for start in (None, "Rep"): + fused_outputs, fused_cache = run(False, start) + reference_outputs, reference_cache = run(True, start) + for output, reference in zip(fused_outputs, reference_outputs, strict=True): + assert torch.equal(output, reference) + payload = _cache_payload(fused_cache) + assert payload is not None and payload.shape[2] == CACHE_T + assert torch.equal(payload, reference_cache[:, :, -CACHE_T:]) + + +@pytest.mark.cpu +def test_installer_is_idempotent_and_preserves_state_dict_keys() -> None: + vae = OmniAutoencoderKLWan( + base_dim=8, + decoder_base_dim=8, + z_dim=4, + dim_mult=[1, 1], + num_res_blocks=1, + temperal_downsample=[False, False], + is_residual=True, + ) + keys_before = set(vae.state_dict()) + assert install_wan_vae_data_movement(vae) + assert install_wan_vae_data_movement(vae) + assert set(vae.state_dict()) == keys_before + + +@pytest.mark.cpu +def test_dynamic_spatial_conv_exposes_direct_path_compatibility_contract() -> None: + conv = torch.nn.Conv3d(2, 2, 3) + conv._vllm_omni_dynamic_spatial_shard_conv = True + conv._source_padding = (1, 1, 1, 1, 2, 0) + + assert _is_dynamic_spatial_conv(conv) + assert _causal_padding(conv) == conv._source_padding + + +@pytest.mark.cuda +@cuda_only +@torch.no_grad() +@pytest.mark.parametrize("channels_last", [False, True]) +def test_tiny_wan_decoder_is_bitwise_exact(channels_last: bool) -> None: + config = dict( + base_dim=8, + decoder_base_dim=8, + z_dim=4, + dim_mult=[1, 1], + num_res_blocks=1, + temperal_downsample=[False, True], + is_residual=True, + ) + torch.manual_seed(0) + reference = AutoencoderKLWan(**config).eval().to("cuda") + candidate = OmniAutoencoderKLWan(**config).eval().to("cuda") + candidate.load_state_dict(reference.state_dict()) + if channels_last: + for model in (reference, candidate): + for module in model.decoder.modules(): + if isinstance(module, torch.nn.Conv3d): + module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d) + + latents = torch.randn(1, 4, 3, 4, 4, device="cuda") + if channels_last: + latents = latents.contiguous(memory_format=torch.channels_last_3d) + expected = reference.decode(latents, return_dict=False)[0] + output = candidate.decode(latents, return_dict=False)[0] + assert output.stride() == expected.stride() + assert torch.equal(output, expected) diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index a1f8440e3..e8ef49ffd 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -728,6 +728,8 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch): @pytest.mark.core_model @pytest.mark.cpu def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): + prepared_pipelines = [] + class _DummyLoader: def __init__(self, load_config, od_config=None): del load_config, od_config @@ -775,6 +777,11 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object()) monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler) + monkeypatch.setattr( + model_runner_module, + "prepare_pipeline_wan_vae_data_movement", + lambda pipeline: prepared_pipelines.append(pipeline), + ) monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None) monkeypatch.setattr( model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend @@ -785,6 +792,7 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): assert runner.cache_backend is None assert runner.od_config.cache_backend is None assert dummy_cache_backend.enabled is False + assert prepared_pipelines == [runner.pipeline] @pytest.mark.core_model diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 44c6efe66..37a4069b3 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -18,13 +18,26 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor GridSpec, TileTask, ) +from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import ( + install_wan_vae_data_movement, +) from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) +def prepare_pipeline_wan_vae_data_movement(pipeline: torch.nn.Module) -> None: + """Install Wan data-movement methods before offload or compilation setup.""" + from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + + for vae in ModuleDiscovery.discover(pipeline).vaes: + if isinstance(vae, OmniAutoencoderKLWan): + install_wan_vae_data_movement(vae) + + class OmniAutoencoderKLWan(AutoencoderKLWan): def _execution_context(self): + install_wan_vae_data_movement(self) try: first_param = next(self.parameters()) except StopIteration: diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_vae_data_movement.py b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_data_movement.py new file mode 100644 index 000000000..3770fa991 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_vae_data_movement.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Guarded Wan VAE decoder data-movement fast paths. + +This module adapts the lossless Wan decoder optimizations from SGLang while +keeping Diffusers' parameter names and module topology unchanged. Every +optimized method falls back to the Diffusers 0.38.0 operation order when the +kernel, device, layout, dtype, autograd state, or module type is unsupported. +""" + +from __future__ import annotations + +from types import MethodType +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.autoencoders.autoencoder_kl_wan import ( + CACHE_T, + DupUp3D, + WanCausalConv3d, + WanDecoder3d, + WanResample, + WanResidualBlock, + WanResidualUpBlock, +) + +from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard + +try: + from vllm_omni.diffusion.kernels.wan_vae_data_movement import ( + cat_pad_5d, + dup_up3d_add, + ) +except ImportError: # pragma: no cover - Triton is optional on non-GPU installs. + cat_pad_5d = None + dup_up3d_add = None + +_INSTALL_ATTR = "_vllm_omni_wan_data_movement_installed" +_DYNAMIC_SPATIAL_CONV_ATTR = "_vllm_omni_dynamic_spatial_shard_conv" + + +def _cache_payload(cache: Any) -> torch.Tensor | None: + return cache if isinstance(cache, torch.Tensor) else None + + +def _is_dynamic_spatial_conv(conv: nn.Module) -> bool: + return bool(getattr(conv, _DYNAMIC_SPATIAL_CONV_ATTR, False)) + + +def _spatial_shard_is_active() -> bool: + return wan_spatial_shard._SPATIAL_SHARD_CONTEXT.get() is not None + + +def _causal_padding(conv: nn.Module) -> tuple[int, ...]: + if _is_dynamic_spatial_conv(conv): + return tuple(conv._source_padding) + return tuple(conv._padding) + + +def _can_use_fused_cache_path(conv: nn.Module, x: torch.Tensor) -> bool: + return bool( + cat_pad_5d is not None + and (type(conv) is WanCausalConv3d or _is_dynamic_spatial_conv(conv)) + and x.dim() == 5 + and x.is_cuda + and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() + and not _spatial_shard_is_active() + ) + + +def _run_cached_causal_conv( + conv: nn.Module, + x: torch.Tensor, + cache_list: list[Any], + index: int, +) -> torch.Tensor: + """Run a causal conv and refresh its feature-cache slot.""" + cache = cache_list[index] + is_repeat_marker = isinstance(cache, str) + payload = None if is_repeat_marker else _cache_payload(cache) + if _can_use_fused_cache_path(conv, x) and ( + payload is None or (payload.device == x.device and payload.dtype == x.dtype) + ): + pair = cat_pad_5d(x, payload, _causal_padding(conv), keep_cache_frames=CACHE_T) + if pair is not None: + conv_input, cache_list[index] = pair + return nn.Conv3d.forward(conv, conv_input) + + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < CACHE_T and payload is not None: + cache_x = torch.cat( + [payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + elif cache_x.shape[2] < CACHE_T and is_repeat_marker: + cache_x = torch.cat([torch.zeros_like(cache_x), cache_x], dim=2) + + out = conv(x) if payload is None else conv(x, payload) + cache_list[index] = cache_x + return out + + +def _causal_conv_forward(self: WanCausalConv3d, x: torch.Tensor, cache_x: torch.Tensor | None = None): + padding = list(self._padding) + if ( + any(padding) + and _can_use_fused_cache_path(self, x) + and (cache_x is None or (cache_x.device == x.device and cache_x.dtype == x.dtype)) + ): + conv_input = cat_pad_5d(x, cache_x, padding) + if conv_input is not None: + return nn.Conv3d.forward(self, conv_input) + + if cache_x is not None and padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return nn.Conv3d.forward(self, x) + + +def _resample_forward( + self: WanResample, + x: torch.Tensor, + feat_cache: list[Any] | None = None, + feat_idx: list[int] | None = None, +): + if feat_idx is None: + feat_idx = [0] + batch, channels, frames, height, width = x.size() + if self.mode == "upsample3d" and feat_cache is not None: + index = feat_idx[0] + if feat_cache[index] is None: + feat_cache[index] = "Rep" + feat_idx[0] += 1 + else: + x = _run_cached_causal_conv(self.time_conv, x, feat_cache, index) + feat_idx[0] += 1 + x = x.reshape(batch, 2, channels, frames, height, width) + x = torch.stack((x[:, 0], x[:, 1]), 3) + x = x.reshape(batch, channels, frames * 2, height, width) + + frames = x.shape[2] + x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width) + x = self.resample(x) + x = x.view(batch, frames, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + if self.mode == "downsample3d" and feat_cache is not None: + index = feat_idx[0] + if feat_cache[index] is None: + feat_cache[index] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[index][:, :, -1:, :, :], x], dim=2)) + feat_cache[index] = cache_x + feat_idx[0] += 1 + return x + + +def _residual_block_forward( + self: WanResidualBlock, + x: torch.Tensor, + feat_cache: list[Any] | None = None, + feat_idx: list[int] | None = None, +): + if feat_idx is None: + feat_idx = [0] + residual = self.conv_shortcut(x) + x = self.nonlinearity(self.norm1(x)) + if feat_cache is not None: + index = feat_idx[0] + x = _run_cached_causal_conv(self.conv1, x, feat_cache, index) + feat_idx[0] += 1 + else: + x = self.conv1(x) + + x = self.dropout(self.nonlinearity(self.norm2(x))) + if feat_cache is not None: + index = feat_idx[0] + x = _run_cached_causal_conv(self.conv2, x, feat_cache, index) + feat_idx[0] += 1 + else: + x = self.conv2(x) + return x + residual + + +def _residual_up_block_forward( + self: WanResidualUpBlock, + x: torch.Tensor, + feat_cache: list[Any] | None = None, + feat_idx: list[int] | None = None, + first_chunk: bool = False, +): + if feat_idx is None: + feat_idx = [0] + shortcut_source = x.clone() + for resnet in self.resnets: + if feat_cache is None: + x = resnet(x) + else: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + + if self.upsampler is not None: + if feat_cache is None: + x = self.upsampler(x) + else: + x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx) + + shortcut = self.avg_shortcut + if shortcut is None: + return x + if ( + dup_up3d_add is not None + and type(shortcut) is DupUp3D + and x.is_cuda + and shortcut_source.is_cuda + and x.dtype == shortcut_source.dtype + and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() + and not _spatial_shard_is_active() + ): + fused = dup_up3d_add( + x, + shortcut_source, + shortcut.factor_t, + shortcut.factor_s, + shortcut.repeats, + first_chunk, + ) + if fused is not None: + return fused + return x + shortcut(shortcut_source, first_chunk=first_chunk) + + +def _decoder_forward( + self: WanDecoder3d, + x: torch.Tensor, + feat_cache: list[Any] | None = None, + feat_idx: list[int] | None = None, + first_chunk: bool = False, +): + if feat_idx is None: + feat_idx = [0] + if feat_cache is not None: + index = feat_idx[0] + x = _run_cached_causal_conv(self.conv_in, x, feat_cache, index) + feat_idx[0] += 1 + else: + x = self.conv_in(x) + + x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + for up_block in self.up_blocks: + x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) + + x = self.nonlinearity(self.norm_out(x)) + if feat_cache is not None: + index = feat_idx[0] + x = _run_cached_causal_conv(self.conv_out, x, feat_cache, index) + feat_idx[0] += 1 + else: + x = self.conv_out(x) + return x + + +def install_wan_vae_data_movement(vae: nn.Module) -> bool: + """Install guarded fast paths on a Diffusers Wan decoder once.""" + if getattr(vae, _INSTALL_ATTR, False): + return True + if cat_pad_5d is None or dup_up3d_add is None: + return False + + decoder = getattr(vae, "decoder", None) + if type(decoder) is not WanDecoder3d: + return False + + decoder.forward = MethodType(_decoder_forward, decoder) + patched = 0 + for module in decoder.modules(): + module_type = type(module) + if module_type is WanCausalConv3d: + module.forward = MethodType(_causal_conv_forward, module) + patched += 1 + elif _is_dynamic_spatial_conv(module): + # The dynamic wrapper owns halo exchange while a spatial request is + # active. Its direct path can still use the lossless fused cache + # movement through ``_run_cached_causal_conv``. + patched += 1 + elif module_type is WanResidualBlock: + module.forward = MethodType(_residual_block_forward, module) + elif module_type is WanResidualUpBlock: + module.forward = MethodType(_residual_up_block_forward, module) + elif module_type is WanResample: + module.forward = MethodType(_resample_forward, module) + + if patched == 0: + return False + setattr(vae, _INSTALL_ATTR, True) + return True diff --git a/vllm_omni/diffusion/kernels/__init__.py b/vllm_omni/diffusion/kernels/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/diffusion/kernels/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/diffusion/kernels/wan_vae_data_movement.py b/vllm_omni/diffusion/kernels/wan_vae_data_movement.py new file mode 100644 index 000000000..8d80147bb --- /dev/null +++ b/vllm_omni/diffusion/kernels/wan_vae_data_movement.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bit-exact data-movement kernels for the Wan causal VAE decoder. + +Adapted from SGLang PR #34125: +https://github.com/sgl-project/sglang/pull/34125 + +The kernels in this module only move values, fill zeros, or perform the same +two-operand addition as the PyTorch reference path. They preserve the input +layout instead of selecting a new Conv3d layout, so enabling them does not +change the decoder's cuDNN accumulation order. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +_MAX_INT32 = 2**31 - 1 +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +@triton.jit +def _cat_pad_5d_kernel( + x_ptr, + cache_ptr, + out_ptr, + keep_ptr, + total, + channels, + frames, + height, + width, + cache_frames, + out_frames, + out_height, + out_width, + zero_front_frames, + pad_height, + pad_width, + stride_x_batch, + stride_x_channel, + stride_x_frame, + stride_x_height, + stride_x_width, + stride_cache_batch, + stride_cache_channel, + stride_cache_frame, + stride_cache_height, + stride_cache_width, + has_cache: tl.constexpr, + keep_frames_const: tl.constexpr, + channels_inner: tl.constexpr, + index_64bit: tl.constexpr, + block_size_const: tl.constexpr, +): + if index_64bit: + offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64) + else: + offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const) + mask = offsets < total + + if channels_inner: + out_channel = offsets % channels + remaining = offsets // channels + out_width_idx = remaining % out_width + remaining = remaining // out_width + out_height_idx = remaining % out_height + remaining = remaining // out_height + out_frame = remaining % out_frames + out_batch = remaining // out_frames + else: + out_width_idx = offsets % out_width + remaining = offsets // out_width + out_height_idx = remaining % out_height + remaining = remaining // out_height + out_frame = remaining % out_frames + remaining = remaining // out_frames + out_channel = remaining % channels + out_batch = remaining // channels + + input_width = out_width_idx - pad_width + input_height = out_height_idx - pad_height + input_frame = out_frame - zero_front_frames + + spatial_valid = (input_width >= 0) & (input_width < width) & (input_height >= 0) & (input_height < height) + from_cache = spatial_valid & (input_frame >= 0) & (input_frame < cache_frames) + from_x = spatial_valid & (input_frame >= cache_frames) & (input_frame < cache_frames + frames) + + x_frame = input_frame - cache_frames + x_offset = ( + out_batch * stride_x_batch + + out_channel * stride_x_channel + + x_frame * stride_x_frame + + input_height * stride_x_height + + input_width * stride_x_width + ) + values = tl.load(x_ptr + x_offset, mask=mask & from_x, other=0.0) + if has_cache: + cache_offset = ( + out_batch * stride_cache_batch + + out_channel * stride_cache_channel + + input_frame * stride_cache_frame + + input_height * stride_cache_height + + input_width * stride_cache_width + ) + cache_values = tl.load(cache_ptr + cache_offset, mask=mask & from_cache, other=0.0) + values = tl.where(from_cache, cache_values, values) + tl.store(out_ptr + offsets, values, mask=mask) + + if keep_frames_const > 0: + keep_frame = out_frame - (out_frames - keep_frames_const) + keep_mask = mask & spatial_valid & (keep_frame >= 0) + if channels_inner: + keep_offset = ( + ((out_batch * keep_frames_const + keep_frame) * height + input_height) * width + input_width + ) * channels + out_channel + else: + keep_offset = ( + ((out_batch * channels + out_channel) * keep_frames_const + keep_frame) * height + input_height + ) * width + input_width + tl.store(keep_ptr + keep_offset, values, mask=keep_mask) + + +def _uses_channels_last_3d(tensor: torch.Tensor) -> bool: + return tensor.shape[1] > 1 and tensor.stride(1) == 1 + + +def cat_pad_5d( + x: torch.Tensor, + cache_x: torch.Tensor | None, + padding: list[int] | tuple[int, ...], + keep_cache_frames: int = 0, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: + """Fuse temporal concat, constant padding, and optional cache refresh. + + ``padding`` follows the Wan causal Conv3d convention + ``(w_left, w_right, h_top, h_bottom, t_front, t_back)``. The output uses + the same dense layout that the reference ``cat``/``pad`` path selects. + Unsupported inputs return ``None`` so callers can execute the reference + implementation. + """ + if len(padding) != 6: + return None + pad_width_left, pad_width_right, pad_height_top, pad_height_bottom, pad_front, pad_back = padding + if any(value < 0 for value in padding): + return None + if pad_width_left != pad_width_right or pad_height_top != pad_height_bottom or pad_back != 0: + return None + if x.dim() != 5 or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES: + return None + if keep_cache_frames < 0: + return None + + channels_inner = _uses_channels_last_3d(x) + cache_frames = 0 + if cache_x is not None: + if ( + cache_x.dim() != 5 + or cache_x.dtype != x.dtype + or cache_x.device != x.device + or cache_x.shape[0] != x.shape[0] + or cache_x.shape[1] != x.shape[1] + or cache_x.shape[3:] != x.shape[3:] + or _uses_channels_last_3d(cache_x) != channels_inner + ): + return None + cache_frames = cache_x.shape[2] + + zero_front_frames = pad_front - cache_frames + if zero_front_frames < 0: + return None + + batch, channels, frames, height, width = x.shape + out_frames = pad_front + frames + out_height = height + 2 * pad_height_top + out_width = width + 2 * pad_width_left + keep_frames = min(keep_cache_frames, out_frames) + memory_format = torch.channels_last_3d if channels_inner else torch.contiguous_format + out = torch.empty( + (batch, channels, out_frames, out_height, out_width), + device=x.device, + dtype=x.dtype, + memory_format=memory_format, + ) + total = out.numel() + if total == 0 or total > _MAX_INT32 * 4: + return None + + if keep_frames: + keep_arg = torch.empty( + (batch, channels, keep_frames, height, width), + device=x.device, + dtype=x.dtype, + memory_format=memory_format, + ) + else: + keep_arg = out + + if cache_x is None: + cache_arg = x + cache_strides = (0, 0, 0, 0, 0) + else: + cache_arg = cache_x + cache_strides = cache_x.stride() + + block_size = 512 + grid = (triton.cdiv(total, block_size),) + with torch.get_device_module().device(x.device): + _cat_pad_5d_kernel[grid]( + x, + cache_arg, + out, + keep_arg, + total, + channels, + frames, + height, + width, + cache_frames, + out_frames, + out_height, + out_width, + zero_front_frames, + pad_height_top, + pad_width_left, + *x.stride(), + *cache_strides, + has_cache=cache_x is not None, + keep_frames_const=keep_frames, + channels_inner=channels_inner, + index_64bit=total >= _MAX_INT32, + block_size_const=block_size, + ) + + if keep_cache_frames: + return out, keep_arg + return out + + +@triton.jit +def _dup_up3d_add_kernel( + main_ptr, + source_ptr, + out_ptr, + total, + out_channels, + out_frames, + out_height, + out_width, + frame_offset, + stride_main_batch, + stride_main_channel, + stride_main_frame, + stride_main_height, + stride_main_width, + stride_source_batch, + stride_source_channel, + stride_source_frame, + stride_source_height, + stride_source_width, + stride_out_batch, + stride_out_channel, + stride_out_frame, + stride_out_height, + stride_out_width, + factor_temporal_const: tl.constexpr, + factor_spatial_const: tl.constexpr, + repeats_const: tl.constexpr, + channels_inner: tl.constexpr, + index_64bit: tl.constexpr, + block_size_const: tl.constexpr, +): + if index_64bit: + offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64) + else: + offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const) + mask = offsets < total + + if channels_inner: + out_channel = offsets % out_channels + remaining = offsets // out_channels + out_width_idx = remaining % out_width + remaining = remaining // out_width + out_height_idx = remaining % out_height + remaining = remaining // out_height + out_frame = remaining % out_frames + out_batch = remaining // out_frames + else: + out_width_idx = offsets % out_width + remaining = offsets // out_width + out_height_idx = remaining % out_height + remaining = remaining // out_height + out_frame = remaining % out_frames + remaining = remaining // out_frames + out_channel = remaining % out_channels + out_batch = remaining // out_channels + + unsliced_frame = out_frame + frame_offset + source_frame = unsliced_frame // factor_temporal_const + temporal_remainder = unsliced_frame % factor_temporal_const + source_height = out_height_idx // factor_spatial_const + height_remainder = out_height_idx % factor_spatial_const + source_width = out_width_idx // factor_spatial_const + width_remainder = out_width_idx % factor_spatial_const + repeated_channel = ( + (out_channel * factor_temporal_const + temporal_remainder) * factor_spatial_const + height_remainder + ) * factor_spatial_const + width_remainder + source_channel = repeated_channel // repeats_const + + main_offset = ( + out_batch * stride_main_batch + + out_channel * stride_main_channel + + out_frame * stride_main_frame + + out_height_idx * stride_main_height + + out_width_idx * stride_main_width + ) + source_offset = ( + out_batch * stride_source_batch + + source_channel * stride_source_channel + + source_frame * stride_source_frame + + source_height * stride_source_height + + source_width * stride_source_width + ) + out_offset = ( + out_batch * stride_out_batch + + out_channel * stride_out_channel + + out_frame * stride_out_frame + + out_height_idx * stride_out_height + + out_width_idx * stride_out_width + ) + main_value = tl.load(main_ptr + main_offset, mask=mask, other=0.0) + source_value = tl.load(source_ptr + source_offset, mask=mask, other=0.0) + value = main_value.to(tl.float32) + source_value.to(tl.float32) + tl.store(out_ptr + out_offset, value, mask=mask) + + +def dup_up3d_add( + main: torch.Tensor, + source: torch.Tensor, + factor_temporal: int, + factor_spatial: int, + repeats: int, + drop_first_frames: bool, +) -> torch.Tensor | None: + """Evaluate ``main + DupUp3D(source)`` without materializing the shortcut.""" + if main.dim() != 5 or source.dim() != 5: + return None + if factor_temporal <= 0 or factor_spatial <= 0 or repeats <= 0: + return None + if factor_temporal & (factor_temporal - 1) or factor_spatial & (factor_spatial - 1): + return None + if repeats & (repeats - 1): + return None + if not main.is_cuda or not source.is_cuda: + return None + if main.dtype != source.dtype or main.dtype not in _SUPPORTED_DTYPES or main.device != source.device: + return None + + batch, source_channels, frames, height, width = source.shape + frame_offset = factor_temporal - 1 if drop_first_frames else 0 + numerator = source_channels * repeats + denominator = factor_temporal * factor_spatial * factor_spatial + if numerator % denominator: + return None + expected_shape = ( + batch, + numerator // denominator, + frames * factor_temporal - frame_offset, + height * factor_spatial, + width * factor_spatial, + ) + if tuple(main.shape) != expected_shape: + return None + + out = torch.empty_like(main) + total = out.numel() + if total == 0 or total > _MAX_INT32 * 4: + return None + + block_size = 512 + grid = (triton.cdiv(total, block_size),) + with torch.get_device_module().device(main.device): + _dup_up3d_add_kernel[grid]( + main, + source, + out, + total, + expected_shape[1], + expected_shape[2], + expected_shape[3], + expected_shape[4], + frame_offset, + *main.stride(), + *source.stride(), + *out.stride(), + factor_temporal_const=factor_temporal, + factor_spatial_const=factor_spatial, + repeats_const=repeats, + channels_inner=out.stride(1) == 1 and expected_shape[1] > 1, + index_64bit=total >= _MAX_INT32, + block_size_const=block_size, + ) + return out diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index cb11eb847..ff49bb6e5 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -21,6 +21,7 @@ from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -215,6 +216,7 @@ class HeliosPipeline( self.vae = AutoencoderKLWan.from_pretrained( model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only ).to(self.device) + install_wan_vae_data_movement(self.vae) transformer_config = load_transformer_config(model, "transformer", local_files_only) self.transformer = create_transformer_from_config( diff --git a/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py index 1c94fecc6..9115b36e1 100644 --- a/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py +++ b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py @@ -19,6 +19,7 @@ from torch import nn from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.wan_vae_data_movement import install_wan_vae_data_movement from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.interface import ( SupportImageInput, @@ -374,6 +375,7 @@ class LingBotVideoPipeline( torch_dtype=vae_dtype, local_files_only=local_files_only, ).to(self.device) + install_wan_vae_data_movement(self.vae) self.scheduler = FlowUniPCMultistepScheduler.from_pretrained( model, subfolder=scheduler_subfolder, diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 52c2c646d..d4a593738 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -30,6 +30,9 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import ( from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + prepare_pipeline_wan_vae_data_movement, +) from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import ( @@ -243,6 +246,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): custom_pipeline_name=custom_pipeline_name, device=self.device, ) + prepare_pipeline_wan_vae_data_movement(self.pipeline) time_after_load = time.perf_counter() logger.info(