anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-11 09:43:11 +08:00
1 parent c27623c29c
commit 5308222abb
10 files changed
+1032
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/kernels/wan_vae_data_movement.py
+405
View File
@@ -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