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.
83 lines
4.4 KiB
Plaintext
83 lines
4.4 KiB
Plaintext
// The output head, epilogue half: conv_out's 3x3 spatial stencil, its bias and the clamp to
|
|
// [-1,1], from the 112-channel plane conv.cu's `conv_n128` (a 3x1x1 convolution) writes to the
|
|
// 12-channel plane write_video consumes.
|
|
//
|
|
// Why the convolution writes 108 channels and this pass sums 9 of them. conv_out is 256x27 MACs
|
|
// and its per-position INPUT WINDOW is 6912*2 B (the weight is 12x that and stays resident), i.e.
|
|
// 12 FLOP/byte against the machine's ~375, so any form that reads the 27 taps per position is
|
|
// operand-bound. Moving the nine spatial taps into the output channel dimension of a 3x1x1
|
|
// (temporal-only) convolution makes the operand 768*2 B for 108 outputs instead of 13.8 KB for
|
|
// 12, and the nine shifted copies are summed here.
|
|
//
|
|
// The value path, so that the rounding is on the record: the three temporal taps are accumulated
|
|
// in the convolution's fp32 accumulator and rounded to bf16 in g; the nine spatial taps are then
|
|
// summed here in fp32 from nine bf16 g values, the bias is added in fp32, and the result is
|
|
// rounded once to bf16 on the store. Main's path is also two roundings (cuDNN's bf16 conv_out,
|
|
// then finish_scalar's bf16 of x+bias); the count is the same and what differs is WHERE the first
|
|
// one sits - after three taps instead of after twenty-seven - which is why max|d| moves and rms
|
|
// does not. The reference accumulates all 27 taps in one fp32 accumulator.
|
|
#include <ATen/ATen.h>
|
|
#include <ATen/cuda/CUDAContext.h>
|
|
#include <c10/cuda/CUDAGuard.h>
|
|
#include <cuda_bf16.h>
|
|
#include "head.h"
|
|
|
|
// One thread per output position. Its 12 accumulators take the 12 channels of the same (i,j)
|
|
// tap from nine neighbouring positions; the nine shifted reads of each 216-byte channel vector
|
|
// are served from L1, so the pass moves each element of g about once.
|
|
static __global__ void head_stencil_kernel(const __nv_bfloat16* __restrict__ g,
|
|
const __nv_bfloat16* __restrict__ bias,
|
|
__nv_bfloat16* __restrict__ out,
|
|
int D,int rows,int W,int H,int g_off,int lo) {
|
|
int idx=blockIdx.x*blockDim.x+threadIdx.x;
|
|
if(idx>=D*rows*W)return;
|
|
int w=idx%W, t=idx/W, r=t%rows, d=t/rows;
|
|
int gh0=lo+r;
|
|
const int gH=rows+2*g_off;
|
|
float acc[12];
|
|
#pragma unroll
|
|
for(int c=0;c<12;++c)acc[c]=__bfloat162float(bias[c]);
|
|
for(int i=0;i<3;++i) {
|
|
int gh=gh0+i-1;
|
|
if(gh<0||gh>=H)continue; // conv3d's zero padding, in the plane
|
|
int gr=r+i-1+g_off;
|
|
if(gr<0||gr>=gH)continue;
|
|
for(int j=0;j<3;++j) {
|
|
int ww=w+j-1;
|
|
if(ww<0||ww>=W)continue; // conv3d's zero padding, across the row
|
|
const __nv_bfloat162* p=(const __nv_bfloat162*)(g+((((size_t)d*gH+gr)*W)+ww)*112+(i*3+j)*12);
|
|
#pragma unroll
|
|
for(int k=0;k<6;++k) {
|
|
float2 f=__bfloat1622float2(p[k]);
|
|
acc[2*k]+=f.x;
|
|
acc[2*k+1]+=f.y;
|
|
}
|
|
}
|
|
}
|
|
__nv_bfloat162* o=(__nv_bfloat162*)(out+((((size_t)d*rows+r)*W)+w)*12);
|
|
#pragma unroll
|
|
for(int k=0;k<6;++k)
|
|
o[k]=__floats2bfloat162_rn(fminf(1.f,fmaxf(-1.f,acc[2*k])),fminf(1.f,fmaxf(-1.f,acc[2*k+1])));
|
|
}
|
|
|
|
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) {
|
|
int D=out.size(2),rows=out.size(3),W=out.size(4);
|
|
TORCH_CHECK(g.is_contiguous(at::MemoryFormat::ChannelsLast3d)&&g.size(1)==112,
|
|
"head_stencil: g must be [1,112,D,rows+2*g_off,W] channels-last");
|
|
TORCH_CHECK(out.is_contiguous(at::MemoryFormat::ChannelsLast3d)&&out.size(1)==12,
|
|
"head_stencil: out must be [1,12,D,rows,W] channels-last");
|
|
// The kernel strides g by gH = rows+2*g_off taken from `out` and reads D*rows*W positions of it,
|
|
// so a g of the wrong depth or height reads out of bounds rather than failing here.
|
|
TORCH_CHECK(g.size(2)==out.size(2)&&g.size(3)==out.size(3)+2*g_off&&g.size(4)==out.size(4),
|
|
"head_stencil: g must be [1,112,D,rows+2*g_off,W] of out's [1,12,D,rows,W]");
|
|
c10::cuda::CUDAGuard guard(g.device());
|
|
int64_t total=(int64_t)D*rows*W;
|
|
int threads=256;
|
|
int blocks=(int)((total+threads-1)/threads);
|
|
head_stencil_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(
|
|
(const __nv_bfloat16*)g.data_ptr(),(const __nv_bfloat16*)bias.data_ptr(),
|
|
(__nv_bfloat16*)out.data_ptr(),D,rows,W,(int)H_global,(int)g_off,(int)lo);
|
|
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
|
}
|