Files
hao-ai-lab-FastVideo/fastvideo-kernel/wan22_vae_decode/solution/kernels.h
T
William Lin cf3ea82e74 [perf] PAN2 sampling: Wan2.2 VAE decode kernel, split decode, fused DiT kernels
- 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.
2026-10-28 19:43:04 +00:00

22 lines
1.2 KiB
C++

#pragma once
#include <ATen/ATen.h>
using at::Tensor;
Tensor nhwc_empty(int64_t c,int64_t t,int64_t h,int64_t w,const Tensor& like);
Tensor norm_pad(const Tensor& x, const Tensor& gamma, const Tensor& history, int pt, int ph, bool silu);
Tensor last_frames(const Tensor& padded, int n, int ph);
void finish(Tensor& x, const Tensor& bias, const Tensor& residual);
void write_video(const Tensor& x, Tensor& out, int frame_offset, int row_offset);
Tensor winograd_input(const Tensor& x,int ph);
Tensor winograd_input_fast(const Tensor& x,int ph);
// winograd_output / winograd_output_kernel are deleted: the fused inverse in wino_epilogue.cu
// replaced their only caller, and they had no live user left (PR172 round-1 review).
// Fused inverse transform + bias/residual + RMS norm/SiLU. y = A^T M A + bias + residual;
// n (when gamma is given) = silu(rmsnorm(y)*sqrt(C)*gain) in the (D+pt, H+2ph, W, C) padded
// layout the next convolution consumes, with its first pt frames taken from `hist`.
struct WinoOut { Tensor y; Tensor n; };
WinoOut wino_epilogue(const std::vector<Tensor>& m,int H,int W,const Tensor& bias,
const Tensor& residual,const Tensor& gamma,const Tensor& hist,int pt,int ph,
bool write_y=true);