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.
20 lines
1.3 KiB
C
20 lines
1.3 KiB
C
#pragma once
|
|
#include <ATen/ATen.h>
|
|
|
|
// The output head, epilogue half: conv_out's 3x3 spatial stencil, its bias and the clamp to
|
|
// [-1,1], from the 112-channel plane the 3x1x1 convolution writes (conv.cu's `conv_n128`) to the
|
|
// 12-channel plane write_video consumes.
|
|
// g: [1,112,D,rows+2*g_off,W] channels-last-3d bf16, n = (i*3+j)*12 + c for the tap (i,j);
|
|
// channels 108..111 are the zero-weight lanes and are never read.
|
|
// out: [1,12,D,rows,W] channels-last-3d bf16, written whole.
|
|
// H_global is the height of the plane this stencil filters (352 = 44 latent rows x 8; the pixel
|
|
// plane is 704 tall and run()'s video row is twice this one). lo is the first row of this rank's
|
|
// band IN THAT PLANE, 8*(44*rank/world), and it must stay that: run() passes lo*16 to write_video,
|
|
// the same row on the 704-tall video grid, and the two derivations have to agree.
|
|
void head_stencil(const at::Tensor& g,const at::Tensor& bias,const at::Tensor& out,
|
|
int64_t H_global,int64_t g_off,int64_t lo);
|
|
|
|
// conv.cu: the head's 3x1x1 convolution, at the tile its 112 output channels want. Declared here
|
|
// so the two-sided declaration of a CUTLASS tile cannot drift from its definition in conv.cu.
|
|
at::Tensor conv_n128(const at::Tensor& x,const at::Tensor& w);
|