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.
52 lines
3.7 KiB
C++
52 lines
3.7 KiB
C++
#pragma once
|
|
// Resampler stage of the Wan2.2 VAE decoder: the four-phase spatial upsamplers of up_blocks 0,1,2.
|
|
//
|
|
// A resampler is "nearest x2 in space, then a 3x3x3 convolution", which equals four 2x2 low-resolution
|
|
// convolutions, one per output parity (their repeated taps already summed into per-phase kernels by
|
|
// prepare()). Those convolutions stay direct (4 taps, no transform: F(2x2,2x2) would cut the MACs 1.78x
|
|
// but materialise 4 x 9/4 = 9 x the low-resolution tensor of transformed input and 2.25x the output of
|
|
// products, which prices out as a wash - measured and posted).
|
|
//
|
|
// What is fixed here is the pass after them: the seed's interleave kernel walked the four phase tensors
|
|
// and the DupUp3D source with four 64-bit division/modulo chains per 16 bytes and ran at 3.9 TB/s, half
|
|
// of HBM. This one walks tiles and channel groups with the grid carrying the coordinates, so the pass
|
|
// is bandwidth bound: it reads the four phase outputs and the DupUp3D source once and writes the
|
|
// interleaved high-resolution tensor once, with bias and the shortcut added in bf16 at the same
|
|
// rounding points the seed used.
|
|
#include <ATen/ATen.h>
|
|
#include <vector>
|
|
using at::Tensor;
|
|
|
|
// One resampler over one latent chunk. x is the low-resolution input, already padded by one row top
|
|
// and bottom ([1,C,D,H+2,W] channels-last) with the halo installed at world > 1 (`halo`), and by the two
|
|
// causal history frames in time; pw are the four per-phase 2x2 kernels ([C_out,C_in,1,2,2], phase =
|
|
// 2*row_parity + col_parity), bias [C_out], and skip the DupUp3D source ([1,TC,Ds,HL,WL]). `temporal`
|
|
// is true for up_blocks 0 and 1 (the shortcut and the phase kernels are nearest-x2 in time as well),
|
|
// false for up block 2, whose shortcut is the channel-pair-to-row rearrangement. `first` marks the
|
|
// chunk that holds latent frame 0.
|
|
// Returns the interleaved high-resolution tensor [1,C_out,D,2*HL,2*WL] channels-last.
|
|
//
|
|
// The pass moves 100.5 GB of logical traffic for f121x1 (the four phase tensors read 42.5, the DupUp3D
|
|
// source 15.5, the output written 42.5) in 14.8 ms = 6.8 TB/s.
|
|
//
|
|
// THAT 6.8 IS NOT A BOUND AND MUST NOT BE TRACED TO ONE. It is solved from 100.5 / 14.8, so quoting it
|
|
// against a ceiling is an identity - w06 M18769/M18811, w22 M18463. The bounds on this board that could
|
|
// apply are: the MIX ceiling for a read+write stage, 5.07-5.59 TB/s at this stage's own f_r (w08's
|
|
// f_r-form, J2654/J2944); w14's per-thread store-loop write arm, 4.86-5.15 (J3082); and w12's 4-byte-
|
|
// per-lane burst arm, 7.35-7.50 (J3029), which is BURST-FLATTERED at 2 KB against the tree's 128 B and
|
|
// which w12 withdrew. 6.9 TB/s was a deflation of that withdrawn arm, so it inherits the arm's defect:
|
|
// it has no f_r, it belongs to a limb with no read, and this is a read+write stage.
|
|
//
|
|
// So the honest statement is that the stage moves 100.5 GB at 6.8 TB/s and that NO CEILING FOR ITS OWN
|
|
// STREAM HAS BEEN MEASURED. The counters on this pass would supply one; that run has not been made.
|
|
// Do not write "at its byte bound", and do not write "the model is conservative" (106.7 GB / 14.8 ms is
|
|
// 7.21 TB/s, which is ABOVE the mixed ceiling - the same fact as "the model is 6 % low", stated with the
|
|
// sign that does not flatter it).
|
|
//
|
|
// J2574 also counts 1.37x the sectors of its logical read. That is sector waste on the access pattern
|
|
// and not work the algorithm needs: this kernel reads each phase tensor and the shortcut at the
|
|
// position's own address, with no neighbour re-read and no stencil anywhere in it.
|
|
|
|
Tensor resample_upsample(const Tensor& x, const std::vector<Tensor>& pw, const Tensor& bias,
|
|
const Tensor& skip, bool halo, bool temporal, bool first);
|