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.
79 lines
4.5 KiB
Plaintext
79 lines
4.5 KiB
Plaintext
// CUTLASS 4.2 Blackwell implicit convolution: 2-CTA 256x256x64 tiles,
|
|
// TMA im2col loads, tcgen05 MMA, fp32 accumulation, bf16 epilogue.
|
|
// Every included CUTLASS/CuTe dependency is vendored beside this source.
|
|
#include <ATen/ATen.h>
|
|
#include <ATen/cuda/CUDAContext.h>
|
|
#include "cutlass/cutlass.h"
|
|
#include "cute/tensor.hpp"
|
|
#include "cutlass/kernel_hardware_info.hpp"
|
|
#include "cutlass/conv/convolution.h"
|
|
#include "cutlass/conv/convnd_problem_shape.hpp"
|
|
#include "cutlass/tensor_ref.h"
|
|
#include "cutlass/conv/dispatch_policy.hpp"
|
|
#include "cutlass/conv/collective/collective_builder.hpp"
|
|
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
|
#include "cutlass/conv/device/conv_universal_adapter.hpp"
|
|
#include "cutlass/conv/kernel/conv_universal.hpp"
|
|
using namespace cute;
|
|
using ElementBF16=cutlass::bfloat16_t;
|
|
using LayoutNDHWC=cutlass::layout::TensorNDHWC;
|
|
using ArchTag=cutlass::arch::Sm100;
|
|
using OperatorClass=cutlass::arch::OpClassTensorOp;
|
|
constexpr auto Fprop=cutlass::conv::Operator::kFprop;
|
|
|
|
template<int M,int N,int CL>
|
|
at::Tensor launch(const at::Tensor& x,const at::Tensor& w,int ph,int pw,int phu=-1,int pwu=-1) {
|
|
if(phu<0)phu=ph;if(pwu<0)pwu=pw;
|
|
using Tile=Shape<Int<M>,Int<N>,Shape<_64>>;
|
|
using Cluster=Shape<Int<CL>,_1,_1>;
|
|
using Epilogue=typename cutlass::epilogue::collective::CollectiveBuilder<
|
|
ArchTag,OperatorClass,Tile,Cluster,cutlass::epilogue::collective::EpilogueTileAuto,
|
|
float,float,ElementBF16,LayoutNDHWC,8,ElementBF16,LayoutNDHWC,8,cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
|
|
using Mainloop=typename cutlass::conv::collective::CollectiveBuilder<
|
|
ArchTag,OperatorClass,Fprop,ElementBF16,LayoutNDHWC,8,ElementBF16,LayoutNDHWC,8,float,Tile,Cluster,
|
|
cutlass::conv::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename Epilogue::SharedStorage))>,
|
|
cutlass::conv::collective::KernelScheduleAuto>::CollectiveOp;
|
|
using Problem=cutlass::conv::ConvProblemShape<Fprop,3>;
|
|
using Kernel=cutlass::conv::kernel::ConvUniversal<Problem,Mainloop,Epilogue>;
|
|
using Conv=cutlass::conv::device::ConvUniversalAdapter<Kernel>;
|
|
int C=x.size(1),D=x.size(2),H=x.size(3),W=x.size(4);
|
|
int K=w.size(0),KT=w.size(2),KH=w.size(3),KW=w.size(4);
|
|
auto y=at::empty({1,K,D-KT+1,H+ph+phu-KH+1,W+pw+pwu-KW+1},x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
|
|
Problem problem(cutlass::conv::Mode::kCrossCorrelation,
|
|
{1,D,H,W,C},{x.stride(0),x.stride(2),x.stride(3),x.stride(4),x.stride(1)},
|
|
{K,KT,KH,KW,C},{w.stride(0),w.stride(2),w.stride(3),w.stride(4),w.stride(1)},
|
|
{0,ph,pw},{0,phu,pwu},{1,1,1},{1,1,1},1);
|
|
typename Kernel::StrideC sc{};
|
|
typename Kernel::StrideD sd{};
|
|
cute::for_each(cute::make_seq<cute::rank<0>(sc)>{},[&](auto i){cute::get<0,i>(sc)=problem.stride_C[Problem::RankT-2-i];});
|
|
cute::for_each(cute::make_seq<cute::rank<0>(sd)>{},[&](auto i){cute::get<0,i>(sd)=problem.stride_C[Problem::RankT-2-i];});
|
|
typename Conv::Arguments args{problem,{(const ElementBF16*)x.data_ptr(),(const ElementBF16*)w.data_ptr()},
|
|
{{1.f,0.f},(const ElementBF16*)y.data_ptr(),sc,(ElementBF16*)y.data_ptr(),sd}};
|
|
Conv conv;
|
|
TORCH_CHECK(conv.can_implement(args)==cutlass::Status::kSuccess,"CUTLASS cannot implement shape");
|
|
size_t size=Conv::get_workspace_size(args);
|
|
auto workspace=at::empty({(int64_t)size},x.options().dtype(at::kByte));
|
|
auto stream=at::cuda::getCurrentCUDAStream();
|
|
auto status=conv.initialize(args,workspace.data_ptr(),stream);
|
|
TORCH_CHECK(status==cutlass::Status::kSuccess,"CUTLASS initialization failed ",int(status));
|
|
status=conv.run(stream);
|
|
TORCH_CHECK(status==cutlass::Status::kSuccess,"CUTLASS convolution failed ",int(status));
|
|
return y;
|
|
}
|
|
at::Tensor custom_conv(const at::Tensor& x,const at::Tensor& w,int64_t ph,int64_t pw,int64_t variant) {
|
|
return launch<256,256,2>(x,w,ph,pw);
|
|
}
|
|
|
|
// The output head's 3x1x1 convolution: N=128 because the head has 112 output channels, and a
|
|
// 256-wide tile would spend 144 of its 256 columns on nothing. Measured at one clock (1951-2022
|
|
// MHz, both arms in J1044 and J1343): N=128 is 0.83 PFLOP/s = 34% of peak, N=256 is 0.69 = 28%.
|
|
// The time is what decides it and it is not close: 6.30 ms against 15.2 ms.
|
|
at::Tensor conv_n128(const at::Tensor& x,const at::Tensor& w) {
|
|
return launch<256,128,2>(x,w,0,0);
|
|
}
|
|
|
|
at::Tensor phase_conv(const at::Tensor& x,const at::Tensor& w,int a,int b,bool halo) {
|
|
at::Tensor v=halo?x.narrow(3,a,x.size(3)-1):x;
|
|
return launch<256,256,2>(v,w,halo?0:1-a,1-b,halo?0:a,b);
|
|
}
|