William Lin
86ef6bfadf
- 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.
135 lines
5.8 KiB
Python
135 lines
5.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""``AutoencoderKLWan.decode`` through fastvideo-kernel's Wan2.2 decoder (``WanVAEConfig.decode_kernel``).
|
|
|
|
CPU tests: the switch is off unless asked for (PAN2 sampling asks), and a decode it cannot take runs the module's own
|
|
decode without loading the kernel. GPU tests (sm_100 / sm_103, the PAN2 VAE's weights in
|
|
``FASTVIDEO_TEST_PAN2_MODEL_DIR``): the kernel's decode against the module's fp32 decode, within three times the error
|
|
of a bf16 decode (the kernel's own rule); another geometry decoding exactly as before; two ranks decoding on their own
|
|
GPUs with no collective.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import fastvideo.envs as envs
|
|
from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig
|
|
from fastvideo.models.wan import vae as wan_vae
|
|
from fastvideo.models.wan.vae_config import WanVAEArchConfig, WanVAEConfig
|
|
|
|
|
|
def _tiny_vae(decode_kernel: bool) -> wan_vae.AutoencoderKLWan:
|
|
torch.manual_seed(0)
|
|
arch = WanVAEArchConfig(base_dim=8, z_dim=4, dim_mult=(1, 1), num_res_blocks=1, temperal_downsample=(False, ),
|
|
latents_mean=(0.0, ) * 4, latents_std=(1.0, ) * 4)
|
|
config = WanVAEConfig(arch_config=arch)
|
|
config.load_encoder, config.load_decoder, config.decode_kernel = False, True, decode_kernel
|
|
return wan_vae.AutoencoderKLWan(config).eval()
|
|
|
|
|
|
def test_the_kernel_is_opt_in_and_pan2_sampling_opts_in() -> None:
|
|
assert WanVAEConfig().decode_kernel is False and WanVAEConfig().decode_kernel_parallel is False
|
|
pan2 = PAN2PipelineConfig()
|
|
assert pan2.vae_decode_kernel is True and pan2.vae_decode_kernel_parallel is True and pan2.dit_fused_kernels
|
|
|
|
|
|
def test_a_decode_the_kernel_cannot_take_runs_the_module(monkeypatch) -> None:
|
|
def no_kernel():
|
|
raise AssertionError("a CPU decode must not load the kernel")
|
|
|
|
monkeypatch.setattr(wan_vae, "_wan22_vae_decode", no_kernel)
|
|
z = torch.randn(1, 4, 2, 4, 4)
|
|
with torch.no_grad():
|
|
plain = _tiny_vae(False).decode(z)
|
|
routed = _tiny_vae(True).decode(z)
|
|
assert torch.equal(plain, routed)
|
|
|
|
|
|
_ON_KERNEL_GPU = pytest.mark.skipif(
|
|
not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)),
|
|
reason="needs an sm_100 / sm_103 GPU")
|
|
|
|
|
|
@pytest.fixture
|
|
def pan2_vae_dir() -> str:
|
|
root = envs.FASTVIDEO_TEST_PAN2_MODEL_DIR.get()
|
|
if not root:
|
|
pytest.skip("set FASTVIDEO_TEST_PAN2_MODEL_DIR to a PAN2 root with vae/")
|
|
return os.path.join(root, "vae")
|
|
|
|
|
|
def _pan2_vae(vae_dir: str) -> wan_vae.AutoencoderKLWan:
|
|
from fastvideo.models.loader.component_loader import VAELoader
|
|
|
|
config = WanVAEConfig()
|
|
config.load_encoder, config.load_decoder = False, True
|
|
loader_args = SimpleNamespace(model_paths={}, vae_cpu_offload=False,
|
|
pipeline_config=SimpleNamespace(vae_config=config, vae_precision="fp32"))
|
|
return VAELoader().load(vae_dir, loader_args).to("cuda").eval()
|
|
|
|
|
|
@_ON_KERNEL_GPU
|
|
def test_the_kernel_decodes_pan2_latents_within_the_bf16_error(pan2_vae_dir) -> None:
|
|
vae = _pan2_vae(pan2_vae_dir)
|
|
z = torch.randn((1, 48, 3, 44, 78), generator=torch.Generator().manual_seed(3)).cuda()
|
|
with torch.no_grad():
|
|
vae.config.decode_kernel = False
|
|
reference = vae.decode(z)
|
|
with torch.autocast("cuda", dtype=torch.bfloat16):
|
|
bf16 = vae.decode(z)
|
|
vae.config.decode_kernel = True
|
|
kernel = vae.decode(z)
|
|
assert list(vae._kernel_decoders) == [(z.device, None)], "the kernel path was not taken"
|
|
assert kernel.shape == reference.shape == (1, 3, 9, 704, 1248) and kernel.dtype == torch.float32
|
|
assert float(kernel.abs().max()) <= 1.0
|
|
d, d_bf16 = kernel - reference, bf16.float() - reference
|
|
assert float(d.abs().max()) <= 3 * float(d_bf16.abs().max())
|
|
assert float(d.pow(2).mean().sqrt()) <= 3 * float(d_bf16.pow(2).mean().sqrt())
|
|
|
|
|
|
@_ON_KERNEL_GPU
|
|
def test_another_geometry_decodes_as_before_on_the_gpu(pan2_vae_dir) -> None:
|
|
vae = _pan2_vae(pan2_vae_dir)
|
|
z = torch.randn((1, 48, 2, 32, 32), generator=torch.Generator().manual_seed(4)).cuda()
|
|
with torch.no_grad():
|
|
vae.config.decode_kernel = False
|
|
plain = vae.decode(z)
|
|
vae.config.decode_kernel = True
|
|
routed = vae.decode(z)
|
|
assert torch.equal(plain, routed) and not vae._kernel_decoders
|
|
|
|
|
|
def _decode_alone(rank: int, store_path: str, vae_dir: str, z: torch.Tensor, out_path: str) -> None:
|
|
import torch.distributed as dist
|
|
from safetensors.torch import load_file
|
|
|
|
torch.cuda.set_device(rank)
|
|
dist.init_process_group("nccl", init_method=f"file://{store_path}", rank=rank, world_size=2)
|
|
try:
|
|
weights = load_file(os.path.join(vae_dir, "diffusion_pytorch_model.safetensors"))
|
|
decoder = wan_vae._wan22_vae_decode().WanVAEDecoder(weights, device=f"cuda:{rank}", local=True)
|
|
video = decoder.decode(z)
|
|
torch.save(video.cpu(), f"{out_path}.{rank}")
|
|
dist.barrier(device_ids=[rank])
|
|
finally:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
@_ON_KERNEL_GPU
|
|
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two GPUs")
|
|
def test_two_ranks_decode_on_their_own_gpus(tmp_path, pan2_vae_dir) -> None:
|
|
import torch.multiprocessing as mp
|
|
|
|
# Build the extension here first, so the two ranks load it rather than both compiling into one directory.
|
|
wan_vae._wan22_vae_decode()._module()
|
|
z = torch.randn((1, 48, 2, 44, 78), generator=torch.Generator().manual_seed(5))
|
|
out = str(tmp_path / "video")
|
|
mp.start_processes(_decode_alone, args=(str(tmp_path / "store"), pan2_vae_dir, z, out), nprocs=2, join=True,
|
|
start_method="spawn")
|
|
first, second = torch.load(f"{out}.0"), torch.load(f"{out}.1")
|
|
assert first.shape == (1, 3, 5, 704, 1248) and torch.equal(first, second)
|