William Lin
552d7dfadc
- fastvideo-kernel/wan22_vae_decode: the Wan2.2 (TI2V-5B) VAE decoder on sm_100a / sm_103a at 704x1248, built on first use against fastvideo-kernel's CUTLASS submodule. It frees intermediates early and carries 4 latent frames per pass: a 5 s video peaks at 20.3 GiB, down from 94.6 GiB in the kernel's first version, with the same output. - WanVAEConfig.decode_kernel routes AutoencoderKLWan.decode through it where it applies (bf16 with fp32 accumulation, about 7x faster than the fp32 decode on a B300, PSNR 61-63 dB, SSIM 0.998 against it); decode_kernel_parallel splits the decode over the sequence-parallel group (2, 4 or 8 GPUs) and broadcasts the video, with the ranks agreeing on whether every one of them could build the kernel. - PAN2's DiT runs fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation (pan2_fusions, set_fused_kernels). - PAN2PipelineConfig turns all three on (vae_decode_kernel, vae_decode_kernel_parallel, dit_fused_kernels). - VideoGenerator copies a decoded RGB video's frames to the host once instead of once per frame. - The PAN2 presets sample at guidance 2.5, the reference launcher's default since 2026-10-23.
98 lines
4.6 KiB
Python
98 lines
4.6 KiB
Python
"""The Wan2.2 VAE decoder (fastvideo-kernel/wan22_vae_decode) through its wrapper, against the reference.
|
|
|
|
Needs an sm_100a or sm_103a GPU (B200 / GB200, B300) and the VAE's weights: WAN22_VAE_WEIGHTS names either a diffusers AutoencoderKLWan
|
|
directory (its diffusion_pytorch_model.safetensors) or a .pt file holding the 110 tensors in weight_names() order. The
|
|
first test builds the extension (about a minute). The kernel's own checks are wan22_vae_decode/tools/bench.py (nine
|
|
timed workloads and six short clips at 1, 2 and 4 GPUs, on real latents); these tests cover the wrapper on one GPU:
|
|
the weights it reads from a state dict, and decodes against the fp32 reference, each within three times what the bf16
|
|
reference errs by on the same latents (the bench's rule, calibrated per input instead of per clip).
|
|
"""
|
|
|
|
import functools
|
|
import os
|
|
import sys
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
PACKAGE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "wan22_vae_decode")
|
|
sys.path.insert(0, PACKAGE)
|
|
sys.path.insert(0, os.path.join(PACKAGE, "tools"))
|
|
|
|
WEIGHTS = os.environ.get("WAN22_VAE_WEIGHTS", "")
|
|
|
|
pytestmark = [
|
|
pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)),
|
|
reason="the Wan2.2 VAE decoder is built for sm_100a / sm_103a (B200 / GB200, B300)"),
|
|
pytest.mark.skipif(not WEIGHTS or not os.path.exists(WEIGHTS),
|
|
reason="set WAN22_VAE_WEIGHTS to the VAE's diffusers directory or a weights .pt file"),
|
|
]
|
|
|
|
|
|
@functools.lru_cache(maxsize=None)
|
|
def _state_dict():
|
|
"""name -> tensor of the VAE's parameters, under diffusers' names."""
|
|
from wan22_vae_decode import weight_names
|
|
|
|
if os.path.isdir(WEIGHTS):
|
|
from safetensors.torch import load_file
|
|
|
|
return load_file(os.path.join(WEIGHTS, "diffusion_pytorch_model.safetensors"))
|
|
tensors = torch.load(WEIGHTS, map_location="cpu", weights_only=True)
|
|
return tensors if isinstance(tensors, dict) else dict(zip(weight_names(), tensors))
|
|
|
|
|
|
@functools.lru_cache(maxsize=None)
|
|
def _decoder():
|
|
from wan22_vae_decode import WanVAEDecoder
|
|
|
|
return WanVAEDecoder(_state_dict(), device="cuda")
|
|
|
|
|
|
def test_weights_follow_diffusers_names():
|
|
from wan22_vae_decode import decoder_weights, weight_names
|
|
|
|
names = weight_names()
|
|
assert len(names) == 110 and names[0] == "post_quant_conv.weight" and names[-1] == "decoder.conv_out.bias"
|
|
weights = decoder_weights(_state_dict(), "cpu")
|
|
assert all(w.dtype == torch.bfloat16 and w.is_contiguous() for w in weights)
|
|
assert all(torch.equal(w, _state_dict()[n].to(torch.bfloat16)) for n, w in zip(names, weights))
|
|
|
|
|
|
@pytest.mark.parametrize("latent_frames", [1, 2, 5])
|
|
def test_decode_against_reference(latent_frames):
|
|
import reference
|
|
from wan22_vae_decode import weight_names
|
|
|
|
generator = torch.Generator().manual_seed(latent_frames)
|
|
z = torch.randn((1, 48, latent_frames, 44, 78), generator=generator).to("cuda", torch.bfloat16)
|
|
video = _decoder().decode(z)
|
|
assert video.shape == (1, 3, 4 * latent_frames - 3, 704, 1248) and video.dtype == torch.bfloat16
|
|
|
|
weights = {n: _state_dict()[n].to("cuda", torch.bfloat16).float() for n in weight_names()}
|
|
ref = reference.decode(z.float(), weights, torch.float32)
|
|
release = reference.decode(z, weights, torch.bfloat16).float()
|
|
d, d_release = video.float() - ref, release - ref
|
|
assert bool(torch.isfinite(video).all()) and float(video.abs().max()) <= 1.0
|
|
assert float(d.abs().max()) <= 3 * float(d_release.abs().max())
|
|
assert float(d.pow(2).mean().sqrt()) <= 3 * float(d_release.pow(2).mean().sqrt())
|
|
|
|
|
|
def test_decodes_are_independent():
|
|
"""No state carries from one call to the next: the same latents decode to the same video after another clip."""
|
|
z = torch.randn((1, 48, 2, 44, 78), generator=torch.Generator().manual_seed(7)).to("cuda", torch.bfloat16)
|
|
first = _decoder().decode(z)
|
|
_decoder().decode(torch.randn((1, 48, 3, 44, 78), generator=torch.Generator().manual_seed(8)).to("cuda",
|
|
torch.bfloat16))
|
|
assert torch.equal(_decoder().decode(z), first)
|
|
|
|
|
|
def test_pass_size_does_not_change_the_video():
|
|
"""The causal history crosses passes, so any number of latent frames per pass decodes the same video."""
|
|
from wan22_vae_decode import WanVAEDecoder
|
|
|
|
z = torch.randn((1, 48, 6, 44, 78), generator=torch.Generator().manual_seed(9)).to("cuda", torch.bfloat16)
|
|
reference_pass = _decoder().decode(z)
|
|
for frames in (1, 16):
|
|
assert torch.equal(WanVAEDecoder(_state_dict(), device="cuda", frames_per_pass=frames).decode(z), reference_pass)
|