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.
97 lines
4.7 KiB
Plaintext
97 lines
4.7 KiB
Plaintext
// Resampler stage: the four-phase spatial upsamplers of up_blocks 0,1,2 of the Wan2.2 VAE decoder.
|
|
// See resample.h for why the phase convolutions stay direct and what this pass replaces.
|
|
#include "resample.h"
|
|
#include <ATen/cuda/CUDAContext.h>
|
|
#include <c10/cuda/CUDAException.h>
|
|
#include <cuda_bf16.h>
|
|
using B = __nv_bfloat16;
|
|
// conv.cu: a 2x2 low-resolution convolution of one output parity, padded per phase; `halo` is set when
|
|
// the input carries a one-row halo from the neighbour rank, so the input is narrowed instead of padded.
|
|
at::Tensor phase_conv(const at::Tensor&, const at::Tensor&, int, int, bool);
|
|
|
|
static cudaStream_t stream() { return at::cuda::getCurrentCUDAStream(); }
|
|
static const B* pc(const Tensor& t) { return (const B*)t.data_ptr<at::BFloat16>(); }
|
|
static B* pm(Tensor& t) { return (B*)t.data_ptr<at::BFloat16>(); }
|
|
|
|
union P8 { uint4 v; B b[8]; __nv_bfloat162 h[4]; };
|
|
|
|
// One thread: eight channels of one 2x2 low-resolution tile, all four phases. The grid carries the
|
|
// coordinates (tx, t, ty) so there is no integer division on the hot path, and every access is a
|
|
// 16-byte vector: four phase outputs and the DupUp3D source in, four interleaved positions out.
|
|
__global__ void upsampler_kernel(const B* __restrict__ p0, const B* __restrict__ p1,
|
|
const B* __restrict__ p2, const B* __restrict__ p3,
|
|
B* __restrict__ y, const B* __restrict__ bias,
|
|
const B* __restrict__ skip, int C, int HL, int WL, int TC,
|
|
bool temporal, bool first) {
|
|
const int cg = threadIdx.x; // 8-channel group
|
|
const int tx = blockIdx.x, t = blockIdx.y, ty = blockIdx.z;
|
|
const int CP = C / 8, SP = TC / 8;
|
|
const P8 b = {((const uint4*)bias)[cg]};
|
|
const int st = temporal ? (first ? (t + 1) / 2 : t / 2) : t;
|
|
const int64_t wC = 2 * (int64_t)WL * C; // one row of the interleaved output
|
|
const int64_t tC = 2 * (int64_t)HL * wC; // one frame of it
|
|
#pragma unroll
|
|
for (int i = 0; i < 2; ++i) {
|
|
const int yy = 2 * ty + i;
|
|
if (yy >= HL) continue;
|
|
#pragma unroll
|
|
for (int j = 0; j < 2; ++j) {
|
|
const int xx = 2 * tx + j;
|
|
if (xx >= WL) continue;
|
|
const int64_t q0 = ((int64_t)t * HL + yy) * WL + xx;
|
|
P8 ph[4];
|
|
ph[0].v = ((const uint4*)p0)[q0 * CP + cg];
|
|
ph[1].v = ((const uint4*)p1)[q0 * CP + cg];
|
|
ph[2].v = ((const uint4*)p2)[q0 * CP + cg];
|
|
ph[3].v = ((const uint4*)p3)[q0 * CP + cg];
|
|
const int64_t s = ((int64_t)st * HL + yy) * WL + xx;
|
|
P8 s0, s1;
|
|
if (temporal) {
|
|
s0.v = ((const uint4*)skip)[s * SP + cg];
|
|
s1.v = s0.v; // nearest x2 in time: both row parities alike
|
|
} else {
|
|
// up block 2: output channel J takes source channel 2*J + (row parity), so read both parities.
|
|
s0.v = ((const uint4*)skip)[s * SP + 2 * cg];
|
|
s1.v = ((const uint4*)skip)[s * SP + 2 * cg + 1];
|
|
}
|
|
B* out = y + (int64_t)t * tC + (int64_t)(2 * yy) * wC + (int64_t)(2 * xx) * C + 8 * cg;
|
|
if (!temporal) {
|
|
// channel J of the output takes source channel 2J + a: even channels from s0, odd from s1.
|
|
P8 e, o;
|
|
#pragma unroll
|
|
for (int k = 0; k < 4; ++k) {
|
|
e.b[k] = s0.b[2 * k];
|
|
e.b[k + 4] = s1.b[2 * k];
|
|
o.b[k] = s0.b[2 * k + 1];
|
|
o.b[k + 4] = s1.b[2 * k + 1];
|
|
}
|
|
s0.v = e.v;
|
|
s1.v = o.v;
|
|
}
|
|
P8 r;
|
|
#pragma unroll
|
|
for (int p = 0; p < 4; ++p) {
|
|
const P8& sv = (p & 2) ? s1 : s0; // row parity of this phase
|
|
#pragma unroll
|
|
for (int k = 0; k < 4; ++k) r.h[k] = __hadd2(__hadd2(ph[p].h[k], b.h[k]), sv.h[k]);
|
|
const int a = p >> 1, bcol = p & 1;
|
|
((uint4*)(out + (int64_t)a * wC + (int64_t)bcol * C))[0] = r.v;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Tensor resample_upsample(const Tensor& x, const std::vector<Tensor>& pw, const Tensor& bias,
|
|
const Tensor& skip, bool halo, bool temporal, bool first) {
|
|
TORCH_CHECK(pw.size() == 4, "resample_upsample needs the four phase kernels");
|
|
std::vector<Tensor> q(4);
|
|
for (int a = 0; a < 2; ++a) for (int b = 0; b < 2; ++b) q[2 * a + b] = phase_conv(x, pw[2 * a + b], a, b, halo);
|
|
const int C = q[0].size(1), D = q[0].size(2), HL = q[0].size(3), WL = q[0].size(4);
|
|
Tensor y = at::empty({1, C, D, 2 * HL, 2 * WL}, x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
|
|
dim3 block(C / 8, 1), grid((WL + 1) / 2, D, (HL + 1) / 2);
|
|
upsampler_kernel<<<grid, block, 0, stream()>>>(pc(q[0]), pc(q[1]), pc(q[2]), pc(q[3]), pm(y), pc(bias),
|
|
pc(skip), C, HL, WL, (int)skip.size(1), temporal, first);
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
return y;
|
|
}
|