From 6424fec32633c6fdfc32adc994aeec77682f70e7 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sun, 9 Aug 2026 09:50:56 +0800 Subject: [PATCH] [diffusion] Bit-exact data-movement elimination for the Wan causal VAE decoder (H200 LongLive2 704x1280x61f: decode 2.80->2.32 s lossless / 2.12->1.67 s quality=high, e2e -10.7%) (#34125) Co-authored-by: Claude Fable 5 --- .../ops/diffusion/triton/wan_causal_cache.py | 383 ++++++++++++++++++ .../runtime/models/vaes/wanvae.py | 230 ++++++----- .../ops/diffusion/test_wan_causal_cache.py | 164 ++++++++ 3 files changed, 666 insertions(+), 111 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/triton/wan_causal_cache.py create mode 100644 test/registered/kernels/ops/diffusion/test_wan_causal_cache.py diff --git a/python/sglang/kernels/ops/diffusion/triton/wan_causal_cache.py b/python/sglang/kernels/ops/diffusion/triton/wan_causal_cache.py new file mode 100644 index 000000000..0e00c7faa --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/triton/wan_causal_cache.py @@ -0,0 +1,383 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bit-exact data-movement kernels for the Wan causal VAE. + +Both kernels only move values (plus zero fill / one same-order addition), so +their outputs are bitwise identical to the aten op chains they replace: + +- :func:`cat_pad_channels_last_3d` builds a causal Conv3d input directly in + ``channels_last_3d`` layout from a strided hidden state and an optional + temporal feature cache, replacing ``cat + F.pad + contiguous`` (three full + tensor passes plus the cache ``clone``/``cat`` bookkeeping) with one pass. +- :func:`dup_up3d_add` evaluates ``main + DupUp3D(src)`` in one pass, + replacing ``repeat_interleave + permute().contiguous() + add`` (each a full + tensor pass over the upsampled tensor). +""" + +from __future__ import annotations + +import torch +import triton # type: ignore +import triton.language as tl # type: ignore + +_MAX_INT32 = 2**31 - 1 + + +@triton.jit +def _cat_pad_cl3d_kernel( + x_ptr, + cache_ptr, + out_ptr, + keep_ptr, + total, + C, + T, + H, + W, + cache_t, + out_t, + out_h, + out_w, + pad_t_zero, + pad_h, + pad_w, + sxb, + sxc, + sxt, + sxh, + sxw, + scb, + scc, + sct, + sch, + scw, + HAS_CACHE: tl.constexpr, + KEEP_T: tl.constexpr, + IDX64: tl.constexpr, + BLOCK: tl.constexpr, +): + if IDX64: + offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK).to(tl.int64) + else: + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < total + + # Output is channels_last_3d contiguous: linear index = (((b*T+t)*H+h)*W+w)*C+c + oc = offs % C + rest = offs // C + ow = rest % out_w + rest = rest // out_w + oh = rest % out_h + rest = rest // out_h + o_t = rest % out_t + ob = rest // out_t + + iw = ow - pad_w + ih = oh - pad_h + it = o_t - pad_t_zero + + spatial_ok = (iw >= 0) & (iw < W) & (ih >= 0) & (ih < H) + from_cache = spatial_ok & (it >= 0) & (it < cache_t) + from_x = spatial_ok & (it >= cache_t) & (it < cache_t + T) + + xt = it - cache_t + x_off = ob * sxb + oc * sxc + xt * sxt + ih * sxh + iw * sxw + vals = tl.load(x_ptr + x_off, mask=mask & from_x, other=0.0) + if HAS_CACHE: + c_off = ob * scb + oc * scc + it * sct + ih * sch + iw * scw + c_vals = tl.load(cache_ptr + c_off, mask=mask & from_cache, other=0.0) + vals = tl.where(from_cache, c_vals, vals) + tl.store(out_ptr + offs, vals, mask=mask) + if KEEP_T > 0: + # Second output: the compact next-chunk feature cache = unpadded + # interior of the last KEEP_T frames, written in the same pass + # (channels_last_3d contiguous, laid out (B, C, KEEP_T, H, W)). + ct = o_t - (out_t - KEEP_T) + keep = mask & spatial_ok & (ct >= 0) + k_off = (((ob * KEEP_T + ct) * H + ih) * W + iw) * C + oc + tl.store(keep_ptr + k_off, vals, mask=keep) + + +def cat_pad_channels_last_3d( + x: torch.Tensor, + cache_x: torch.Tensor | None, + padding: list[int] | tuple[int, ...], + keep_cache_t: int = 0, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: + """``contiguous_cl3d(F.pad(cat([cache_x, x], dim=2), padding))`` in one pass. + + ``padding`` follows the ``WanCausalConv3d._padding`` convention + ``(w_left, w_right, h_top, h_bottom, t_front, t_back)``; the temporal + front padding is consumed by ``cache_x`` frames first and any remainder is + zero filled (identical to the aten fallback). With ``keep_cache_t > 0`` + the same pass also emits the compact next-chunk feature cache (the + unpadded interior of the last ``keep_cache_t`` frames) and returns the + ``(conv_input, cache)`` pair. Returns ``None`` when the request is + unsupported so callers can fall back. + """ + pw_l, pw_r, ph_t, ph_b, pt_front, pt_back = padding + if pw_l != pw_r or ph_t != ph_b or pt_back != 0: + return None + if x.dim() != 5 or not x.is_cuda: + return None + cache_t = 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:] + ): + return None + cache_t = cache_x.shape[2] + pad_t_zero = pt_front - cache_t + if pad_t_zero < 0: + return None + + B, C, T, H, W = x.shape + out_t = pt_front + T + out_h = H + 2 * ph_t + out_w = W + 2 * pw_l + keep_t = min(keep_cache_t, out_t) + out = torch.empty( + (B, C, out_t, out_h, out_w), + device=x.device, + dtype=x.dtype, + memory_format=torch.channels_last_3d, + ) + total = out.numel() + if total == 0 or total > _MAX_INT32 * 4: + return None + if keep_t > 0: + keep_arg = torch.empty( + (B, C, keep_t, H, W), + device=x.device, + dtype=x.dtype, + memory_format=torch.channels_last_3d, + ) + else: + keep_arg = out # unused dummy pointer + + if cache_x is None: + cache_arg = x # unused dummy pointer + scb = scc = sct = sch = scw = 0 + else: + cache_arg = cache_x + scb, scc, sct, sch, scw = cache_x.stride() + sxb, sxc, sxt, sxh, sxw = x.stride() + + BLOCK = 512 + grid = (triton.cdiv(total, BLOCK),) + with torch.get_device_module().device(x.device): + _cat_pad_cl3d_kernel[grid]( + x, + cache_arg, + out, + keep_arg, + total, + C, + T, + H, + W, + cache_t, + out_t, + out_h, + out_w, + pad_t_zero, + ph_t, + pw_l, + sxb, + sxc, + sxt, + sxh, + sxw, + scb, + scc, + sct, + sch, + scw, + HAS_CACHE=cache_x is not None, + KEEP_T=keep_t, + IDX64=total >= _MAX_INT32, + BLOCK=BLOCK, + ) + if keep_cache_t > 0: + return out, keep_arg + return out + + +@triton.jit +def _dup_up3d_add_kernel( + main_ptr, + src_ptr, + out_ptr, + total, + C_out, + out_t, + out_h, + out_w, + t_offset, + smb, + smc, + smt, + smh, + smw, + ssb, + ssc, + sst, + ssh, + ssw, + sob, + soc, + sot, + soh, + sow, + FT: tl.constexpr, + FS: tl.constexpr, + REPEATS: tl.constexpr, + CHANNELS_INNER: tl.constexpr, + IDX64: tl.constexpr, + BLOCK: tl.constexpr, +): + if IDX64: + offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK).to(tl.int64) + else: + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < total + + # Logical (B, C_out, out_t, out_h, out_w) index; the output tensor keeps + # ``main``'s stride order (``empty_like`` preserve), matching what the + # aten add would produce, so downstream layout-sensitive reductions see + # the exact same memory format. FT/FS/REPEATS are constexpr powers of two, + # so the pixel-shuffle divisions compile to shifts. The traversal order + # follows the output's memory order (channels innermost for NHWC-style + # ``main``) so stores and ``main`` loads stay coalesced. + if CHANNELS_INNER: + oc = offs % C_out + rest = offs // C_out + ow = rest % out_w + rest = rest // out_w + oh = rest % out_h + rest = rest // out_h + o_t = rest % out_t + ob = rest // out_t + else: + ow = offs % out_w + rest = offs // out_w + oh = rest % out_h + rest = rest // out_h + o_t = rest % out_t + rest = rest // out_t + oc = rest % C_out + ob = rest // C_out + + # Undo the DupUp3D pixel-shuffle mapping (t_offset restores frames that + # were sliced away for the first chunk). + t2 = o_t + t_offset + ti = t2 // FT + rt = t2 % FT + hi = oh // FS + rh = oh % FS + wi = ow // FS + rw = ow % FS + ch_rep = ((oc * FT + rt) * FS + rh) * FS + rw + ci = ch_rep // REPEATS + + m_off = ob * smb + oc * smc + o_t * smt + oh * smh + ow * smw + s_off = ob * ssb + ci * ssc + ti * sst + hi * ssh + wi * ssw + o_off = ob * sob + oc * soc + o_t * sot + oh * soh + ow * sow + m = tl.load(main_ptr + m_off, mask=mask, other=0.0) + s = tl.load(src_ptr + s_off, mask=mask, other=0.0) + # Accumulate in fp32 and round once on store, matching aten's opmath + # behaviour for half-precision adds. + vals = m.to(tl.float32) + s.to(tl.float32) + tl.store(out_ptr + o_off, vals, mask=mask) + + +def dup_up3d_add( + main: torch.Tensor, + src: torch.Tensor, + factor_t: int, + factor_s: int, + repeats: int, + drop_first_frames: bool, +) -> torch.Tensor | None: + """``main + DupUp3D(src)`` in one pass (output layout follows ``main``). + + ``src`` is the DupUp3D input ``(B, C_in, T, H, W)``; ``main`` must match + the DupUp3D output shape. ``drop_first_frames`` mirrors the + ``first_chunk`` slicing (``x[:, :, factor_t - 1 :]``). Returns ``None`` + when unsupported so callers can fall back. + """ + if main.dim() != 5 or src.dim() != 5: + return None + # Power-of-two factors keep the constexpr pixel-shuffle math on the + # shift/mask path (all Wan-family VAEs use ft in {1, 2}, fs = 2). + if factor_t & (factor_t - 1) or factor_s & (factor_s - 1): + return None + if repeats <= 0 or repeats & (repeats - 1): + return None + if not main.is_cuda or not src.is_cuda: + return None + if main.dtype != src.dtype or main.device != src.device: + return None + B, C_in, T, H, W = src.shape + t_offset = factor_t - 1 if drop_first_frames else 0 + exp_shape = ( + B, + C_in * repeats // (factor_t * factor_s * factor_s), + T * factor_t - t_offset, + H * factor_s, + W * factor_s, + ) + if tuple(main.shape) != exp_shape: + return None + + # ``empty_like`` preserves the stride order of the dense ``main`` view — + # the same layout the aten ``main + dup`` would produce — so downstream + # layout-sensitive reductions see the exact same memory format. + out = torch.empty_like(main) + total = out.numel() + if total == 0 or total > _MAX_INT32 * 4: + return None + + smb, smc, smt, smh, smw = main.stride() + ssb, ssc, sst, ssh, ssw = src.stride() + sob, soc, sot, soh, sow = out.stride() + BLOCK = 512 + grid = (triton.cdiv(total, BLOCK),) + with torch.get_device_module().device(main.device): + _dup_up3d_add_kernel[grid]( + main, + src, + out, + total, + exp_shape[1], + exp_shape[2], + exp_shape[3], + exp_shape[4], + t_offset, + smb, + smc, + smt, + smh, + smw, + ssb, + ssc, + sst, + ssh, + ssw, + sob, + soc, + sot, + soh, + sow, + FT=factor_t, + FS=factor_s, + REPEATS=repeats, + CHANNELS_INNER=out.stride(1) == 1 and exp_shape[1] > 1, + IDX64=total >= _MAX_INT32, + BLOCK=BLOCK, + ) + return out diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py index 1bbad1241..cd2119b83 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py @@ -54,6 +54,19 @@ from sglang.multimodal_gen.runtime.models.vaes.common import ( ) from sglang.multimodal_gen.runtime.platforms import current_platform +if current_platform.is_cuda(): + try: + from sglang.kernels.ops.diffusion.triton.wan_causal_cache import ( + cat_pad_channels_last_3d, + dup_up3d_add, + ) + except ImportError: # pragma: no cover + cat_pad_channels_last_3d = None + dup_up3d_add = None +else: + cat_pad_channels_last_3d = None + dup_up3d_add = None + CACHE_T = 2 is_first_frame = contextvars.ContextVar("is_first_frame", default=False) @@ -82,6 +95,72 @@ def match_conv3d_input_format(x: torch.Tensor, weight: torch.Tensor) -> torch.Te return x +def _cache_payload(cache) -> torch.Tensor | None: + """Tensor payload of a feature-cache entry (``None`` for empty slots and + for the ``"Rep"`` marker).""" + return cache if isinstance(cache, torch.Tensor) else None + + +def _fused_conv_cache_supported(conv: nn.Module, x: torch.Tensor) -> bool: + return ( + cat_pad_channels_last_3d is not None + and type(conv) is WanCausalConv3d + and x.dim() == 5 + and x.is_cuda + and current_platform.is_amp_supported() + and _conv3d_weight_is_channels_last_3d(conv.weight) + and not torch.compiler.is_compiling() + ) + + +def _run_cached_causal_conv( + conv: nn.Module, + x: torch.Tensor, + cache_list: list, + idx: int, +) -> torch.Tensor: + """Run one causal conv, consuming and refreshing its feature-cache slot. + + Fast path (bit-exact with the aten chain, pure data movement plus zero + fill): build the conv input (cache frames + hidden state + padding) + directly in channels_last_3d with one kernel, and take the next cache + entry as one compact copy of that input's unpadded tail instead of the + per-chunk clone/cat bookkeeping (the compact copy holds exactly the + reference cache values, so fused and fallback chunks can interleave). + Falls back to the original op chain whenever the fused kernel does not + support the request. + """ + cache = cache_list[idx] + is_rep = isinstance(cache, str) # "Rep" marker from WanResample + payload = None if is_rep else _cache_payload(cache) + if _fused_conv_cache_supported(conv, x) and ( + payload is None or (payload.device == x.device and payload.dtype == x.dtype) + ): + # The same kernel pass emits the conv input and the compact + # next-chunk cache (so the conv-input buffer is freed after the conv + # instead of being pinned until the next chunk). + pair = cat_pad_channels_last_3d(x, payload, conv._padding, keep_cache_t=CACHE_T) + if pair is not None: + inp, cache_list[idx] = pair + return nn.Conv3d.forward(conv, inp) + # Original aten path (bit-identical bookkeeping). + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and payload is not None: + # cache last frame of last two chunk + cache_x = torch.cat( + [payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + elif cache_x.shape[2] < 2 and is_rep: + cache_x = torch.cat( + [torch.zeros_like(cache_x).to(cache_x.device), cache_x], + dim=2, + ) + out = conv(x) if payload is None else conv(x, payload) + cache_list[idx] = cache_x + return out + + class AvgDown3D(nn.Module): def __init__( self, @@ -218,6 +297,17 @@ class WanCausalConv3d(nn.Conv3d): def forward(self, x, cache_x=None): padding = list(self._padding) + if ( + any(padding) + and _fused_conv_cache_supported(self, x) + and ( + cache_x is None + or (cache_x.device == x.device and cache_x.dtype == x.dtype) + ) + ): + inp = cat_pad_channels_last_3d(x, cache_x, padding) + if inp is not None: + return super().forward(inp) x = causal_conv3d_cat_pad(x, cache_x, padding) x = ( x if current_platform.is_amp_supported() else x.to(self.weight.dtype) @@ -281,36 +371,7 @@ def resample_forward(self, x): _feat_cache[idx] = "Rep" _feat_idx += 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 last frame of last two chunk - 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 + x = _run_cached_causal_conv(self.time_conv, x, _feat_cache, idx) _feat_idx += 1 x = x.reshape(b, 2, c, t, h, w) @@ -360,18 +421,7 @@ def residual_block_forward(self, x): _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: - cache_x = torch.cat( - [ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - - x = self.conv1(x, _feat_cache[idx]) - _feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv1, x, _feat_cache, idx) _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) @@ -389,18 +439,7 @@ def residual_block_forward(self, x): _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: - cache_x = torch.cat( - [ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - - x = self.conv2(x, _feat_cache[idx]) - _feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv2, x, _feat_cache, idx) _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) @@ -478,7 +517,28 @@ def residual_up_block_forward(self, x): x = self.upsampler(x) if self.avg_shortcut is not None: - x = x + self.avg_shortcut(x_copy) + shortcut = self.avg_shortcut + if ( + dup_up3d_add is not None + and type(shortcut) is DupUp3D + and x.is_cuda + and x_copy.is_cuda + and x.dtype == x_copy.dtype + and not torch.compiler.is_compiling() + ): + # Bit-exact single-pass ``main + DupUp3D(src)`` (data movement + # plus one same-order fp32-accumulated add). + fused = dup_up3d_add( + x, + x_copy, + shortcut.factor_t, + shortcut.factor_s, + shortcut.repeats, + bool(first_chunk.get()), + ) + if fused is not None: + return fused + x = x + shortcut(x_copy) return x @@ -958,20 +1018,7 @@ class WanEncoder3d(nn.Module): _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat( - [ - _feat_cache[idx][:, :, -1, :, :] - .unsqueeze(2) - .to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv_in(x, _feat_cache[idx]) - _feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv_in, x, _feat_cache, idx) _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) @@ -995,20 +1042,7 @@ class WanEncoder3d(nn.Module): _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat( - [ - _feat_cache[idx][:, :, -1, :, :] - .unsqueeze(2) - .to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv_out(x, _feat_cache[idx]) - _feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv_out, x, _feat_cache, idx) _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) @@ -1316,20 +1350,7 @@ class WanDecoder3d(nn.Module): _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat( - [ - _feat_cache[idx][:, :, -1, :, :] - .unsqueeze(2) - .to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv_in(x, _feat_cache[idx]) - _feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv_in, x, _feat_cache, idx) _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) @@ -1350,20 +1371,7 @@ class WanDecoder3d(nn.Module): _feat_idx = feat_idx.get() if _feat_cache is not None: idx = _feat_idx - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and _feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat( - [ - _feat_cache[idx][:, :, -1, :, :] - .unsqueeze(2) - .to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv_out(x, _feat_cache[idx]) - _feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv_out, x, _feat_cache, idx) _feat_idx += 1 feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) diff --git a/test/registered/kernels/ops/diffusion/test_wan_causal_cache.py b/test/registered/kernels/ops/diffusion/test_wan_causal_cache.py new file mode 100644 index 000000000..530c295de --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_wan_causal_cache.py @@ -0,0 +1,164 @@ +"""Wan causal VAE data-movement kernels: the fused conv-input builder and the +fused DupUp3D shortcut add must be bitwise identical to the aten op chains +they replace (they are pure data movement plus zero fill / one fp32 add).""" + +import sys + +import pytest +import torch +import torch.nn.functional as F + +from sglang.kernels.ops.diffusion.triton.wan_causal_cache import ( + cat_pad_channels_last_3d, + dup_up3d_add, +) +from sglang.multimodal_gen.runtime.models.vaes import wanvae +from sglang.multimodal_gen.runtime.models.vaes.wanvae import ( + CACHE_T, + WanCausalConv3d, + _cache_payload, + _run_cached_causal_conv, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _cl3d(shape, dtype): + return torch.randn(shape, device="cuda", dtype=dtype).contiguous( + memory_format=torch.channels_last_3d + ) + + +def _ref_cat_pad(x, cache, padding): + p = list(padding) + if cache is not None: + x = torch.cat([cache, x], dim=2) + p[4] -= cache.shape[2] + if any(p): + x = F.pad(x, p) + return x.contiguous(memory_format=torch.channels_last_3d) + + +@torch.no_grad() +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "c,t,h,w,cache_t,pads", + [ + (96, 1, 10, 14, 0, (1, 1, 1, 1, 2, 0)), # first chunk, zero-fill front + (96, 1, 10, 14, 1, (1, 1, 1, 1, 2, 0)), # legacy 1-frame cache + (96, 1, 10, 14, 2, (1, 1, 1, 1, 2, 0)), # steady state k3 conv + (64, 1, 10, 14, 2, (0, 0, 0, 0, 2, 0)), # time_conv (temporal only) + (48, 4, 10, 14, 2, (1, 1, 1, 1, 2, 0)), # encoder-style T=4 chunk + ], +) +def test_cat_pad_bitwise(dtype, c, t, h, w, cache_t, pads) -> None: + torch.cuda.manual_seed(0) + x = _cl3d((1, c, t, h, w), dtype) + cache = None + if cache_t: + # Strided interior view: caches may arrive as non-contiguous slices. + ph, pw = pads[2], pads[0] + buf = _cl3d((1, c, cache_t, h + 2 * ph, w + 2 * pw), dtype) + cache = buf[:, :, :, ph : ph + h, pw : pw + w] + out = cat_pad_channels_last_3d(x, cache, pads) + ref = _ref_cat_pad(x, cache, pads) + assert out is not None and out.shape == ref.shape + assert out.is_contiguous(memory_format=torch.channels_last_3d) + assert torch.equal(out, ref) + + # Dual-output mode: the same pass also emits the compact feature cache + # (unpadded interior of the last frames), bitwise equal to the slice. + pair = cat_pad_channels_last_3d(x, cache, pads, keep_cache_t=2) + assert pair is not None + out2, keep = pair + assert torch.equal(out2, ref) + ph, pw = pads[2], pads[0] + keep_t = min(2, ref.shape[2]) + want = ref[:, :, ref.shape[2] - keep_t :, ph : ph + h, pw : pw + w] + assert keep.shape == want.shape + assert keep.is_contiguous(memory_format=torch.channels_last_3d) + assert torch.equal(keep, want) + + +@torch.no_grad() +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "c_in,c_out,t,h,w,ft,fs,drop", + [ + (128, 64, 1, 10, 14, 2, 2, False), + (128, 64, 1, 10, 14, 2, 2, True), # first_chunk slicing + (64, 32, 2, 10, 14, 1, 2, False), + ], +) +def test_dup_up3d_add_bitwise(dtype, c_in, c_out, t, h, w, ft, fs, drop) -> None: + torch.cuda.manual_seed(0) + repeats = c_out * ft * fs * fs // c_in + src = _cl3d((1, c_in, t, h, w), dtype) + t_out = t * ft - (ft - 1 if drop else 0) + # Main arm as a permuted view, like the WanResample 2D output. + main = torch.randn( + (1, t_out, c_out, h * fs, w * fs), device="cuda", dtype=dtype + ).permute(0, 2, 1, 3, 4) + + dup = src.repeat_interleave(repeats, dim=1) + dup = dup.view(1, c_out, ft, fs, fs, t, h, w) + dup = dup.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + dup = dup.view(1, c_out, t * ft, h * fs, w * fs) + if drop: + dup = dup[:, :, ft - 1 :, :, :] + ref = main + dup + + out = dup_up3d_add(main, src, ft, fs, repeats, drop) + assert out is not None and out.shape == ref.shape + # Layout must match the aten add output exactly (downstream reductions + # are layout-sensitive), and every value must be bitwise identical. + assert out.stride() == ref.stride() + assert torch.equal(out, ref) + + +@torch.no_grad() +@pytest.mark.parametrize("pads_temporal_only", [False, True]) +def test_cached_conv_chunk_loop_bitwise(pads_temporal_only) -> None: + """The fused conv-input/compact-cache scheme must reproduce the original + clone/cat bookkeeping bitwise across a chunked decode, including the + first-chunk zero fill and the "Rep" marker start used by WanResample.""" + torch.cuda.manual_seed(0) + c = 64 + if pads_temporal_only: + conv = WanCausalConv3d(c, 2 * c, (3, 1, 1), padding=(1, 0, 0)) + else: + conv = WanCausalConv3d(c, c, 3, padding=1) + conv = conv.to(device="cuda", dtype=torch.float32) + conv.weight.data = conv.weight.data.contiguous(memory_format=torch.channels_last_3d) + chunks = [_cl3d((1, c, 1, 10, 14), torch.float32) for _ in range(4)] + + def run(force_fallback, start): + cache = [start] + outs = [] + if force_fallback: + orig = wanvae.cat_pad_channels_last_3d + wanvae.cat_pad_channels_last_3d = None + try: + for x in chunks: + outs.append(_run_cached_causal_conv(conv, x, cache, 0)) + finally: + if force_fallback: + wanvae.cat_pad_channels_last_3d = orig + return outs, cache[0] + + for start in (None, "Rep"): + fused_outs, fused_cache = run(False, start) + ref_outs, ref_cache = run(True, start) + for got, want in zip(fused_outs, ref_outs, strict=True): + assert torch.equal(got, want) + got_payload = _cache_payload(fused_cache) + assert got_payload is not None and got_payload.shape[2] == CACHE_T + # Reference cache holds the last CACHE_T unpadded frames. + assert torch.equal(got_payload, ref_cache[:, :, -CACHE_T:]) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))