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.
7.8 KiB
7.8 KiB
How the decoder got here
The solution is the main of a multi-agent optimization run (bifrost, 2026-10-25/26): 32 DeepSeek V4.1 Flash peers
and two GPT-6 Astra architects on GB200, starting from a first solution written before the run. Every merge passed the
run's CI (correct on every workload and short clip at worlds 1, 2 and 4, and at least 1 % faster than main by the
geomean in the same job, unless waived) and an independent review. The comments in solution/ cite the run's pull
requests (PRn), issues (In) and jobs (Jn).
The design
The decoder runs layer by layer, as kernels of its own and CUTLASS kernels instantiated in solution/ (the
headers are fastvideo-kernel's CUTLASS submodule). There is no megakernel: see "What did not pay" below.
- Frames in passes. A pass carries several latent frames per GPU instead of the reference's one (4 by default,
frames_per_pass; the run measured its speed at 16, see "Memory"), so its GEMMs are large, and the causal history (the last two frames of each 3 x 3 x 3 convolution's input) is carried from pass to pass as an offset, never by concatenation. - The 18 convolutions of up blocks 1-3 (most of the FLOPs) are F(4x4, 3x3) Winograd in space, keeping the causal
3-tap convolution in time, in three steps:
wino_in.cu: the input transform, two blocks per tile;- 36 component products, each a CUTLASS SM100 implicit-GEMM 3 x 1 x 1 convolution (
conv.cu: 2-CTA 256 x 256 x 64 tiles, TMA im2col,tcgen05MMA, fp32 accumulation); wino_epilogue.cu: the inverse transform with the bias and the residual add, and the next convolution's RMS norm, SiLU and causal padding, so that convolution reads its input ready-made.
- The other convolutions at 44 x 78 are direct CUTLASS implicit GEMMs; a few small ones still go through cuDNN and cuBLAS, under 1 % of the GPU time.
- The upsamplers never materialize the upsampled tensor: nearest x2 followed by a 3 x 3 convolution is four 2 x 2
phase convolutions at the low resolution, and
resample.cuinterleaves them with the bias and the DupUp3D shortcut in one pass. The temporal upsampler (timeconv.cu) writes its two frames straight into place. - Pointwise and the rest:
norm.cu(RMS norm + SiLU + causal pad in one pass, where the epilogue does not already apply it),attn.cu(the mid block's single-head attention over each frame's 3,432 positions), andhead.cu(the output head:conv_outas a 3 x 1 x 1 convolution to 112 channels, then a 9-tap stencil with the bias, the clamp and the unpatchify write). - Over GPUs (
comm.cu,prefix.cu): the decoder's prefix (post_quant_conv, conv_in, the mid block's first residual block and its attention) is split by time, each rank computing its frames with exact causal context from the full latents, then transposed to bands of rows (prefix.cu). From there each rank owns 44 / world latent rows; every 3 x 3 convolution exchanges one halo row with the neighbours only (NCCL send/recv in one group), and rank 0 receives the other ranks' pixels per chunk.
The merges
| Merge | Change | CI speedup | Geomean after |
|---|---|---|---|
| kernel#0 | the first solution: Winograd up blocks, phase upsamplers, fused pointwise, row bands | 1043 ms | |
| PR3 | 1 x 1 channel shortcuts as own kernels, their bias folded into the residual | x1.017 | 1025 ms |
| PR24 | mid-block attention as own kernels instead of ATen SDPA (8.5 -> 2.6 ms of f121x1) | x1.013 | 1012 ms |
| PR15 | RMS norm + SiLU + pad in one pass | x1.110 | 913 ms |
| PR204 | halo exchange with the neighbours only, one coalesced group | x1.014 | 901 ms |
| PR392 | the fused Winograd input transform and the fused inverse epilogue (with the next norm) | x1.189 | 757 ms |
| PR90 | output head as a temporal convolution and a 9-tap stencil (waived at parity) | x1.003 | 755 ms |
| PR438 | the prefix split by time before the row bands | x1.019 | 740 ms |
| PR452 | world 8 uses all-gather, not the ring (worlds 1/2/4 unchanged; waived) | x1.001 | 740 ms |
| PR488 | the upsampler's phases interleaved in one pass, and the temporal upsampler in place | x1.017 | 728 ms |
| PR513 | chunks of min(16 x world, 61) latent frames instead of 8 x world | x1.023 | 711 ms |
| PR490 | a check that no fused input leaks across chunks (waived) | x0.999 | 712 ms |
| PR592 | no dead store of y at each residual block's first convolution | x1.010 | 704 ms |
| PR517 | the input transform at two blocks per tile (waived at x1.007) | x1.007 | 699 ms |
Memory
The run optimized time alone, and its main decoded 16 latent frames per GPU per pass at a peak of 100-102 GiB on one GPU, against 13-16 GiB for the release's decode, which carries one latent frame. After the run:
- Tensors are freed at their last use. At the peak, up block 3's first convolution (512 -> 256 channels at 352 x 624) held its padded input after the Winograd transform had read it, the transform (2.25x the input) after every product was enqueued, and the block's input after its shortcut and first norm had read it. Releasing each of them then takes the peak from 101 to 60 GiB at 16 frames per pass, at the same speed, with a bit-identical video.
- The pass is 4 latent frames per GPU by default (
frames_per_pass), since memory grows with it at 3.5-4.2 GiB per latent frame at 704 x 1248: 21-22 GiB on one GPU and 18-20 GiB per GPU at 2 and 4 GPUs, for 8-9 % more time than 16 frames per pass. The time a smaller pass costs is in the low resolutions (44 x 78 and 88 x 156), whose GEMMs do not fill the GPU with few frames, while the memory is in the two high ones (176 x 312 and 352 x 624). Passes of 16 frames through up block 1 and of 2-4 through up blocks 2 and 3 should keep most of the speed at this memory; it is not built.
What did not pay
- Megakernels. Before the run, a persistent kernel per residual block measured 4.1-7.5x slower than two fused per-convolution kernels. During the run, a cooperative kernel for one whole Winograd convolution (forward transform, products and inverse with the next norm, in one launch of 152 CTAs in 2-CTA clusters) was correct but took 6.2-7.2 ms at T 8, 352 x 624, 256 channels, against 3.3 ms for the same work in three launches (0.46-0.53x). The gaps between the kernels of a decode are about 5 us each, so fewer launches buy little; what a single kernel loses is the occupancy and the tiling each stage picks for itself.
- FP8. e4m3 operands in the Winograd component products cut the products from 262 to 162 ms at f121x1 (10-11 % of the decode), but fail the error budget on all three clips (RMS at 2.3-3.3x the allowance, even with a power-of-two scale per component): the transformed operands need about one more mantissa bit. Before the run, FP8 in the direct convolutions passed in up blocks 1-3 but bought only 1 % of the decode, since quantization was a pass of its own.
- Splitting by time throughout. A chunk needs 38 latent frames of causal context, so recomputing it is a poor fit; only the prefix is split by time (PR438).
Where to go next
- The products are about 60 % of the GPU time (59.6 % at f121x1: 1,383 CUTLASS launches, 384 ms of 644), at 65-68 % of the tensor cores' peak, bound by operand fetch. The transforms and the epilogues are another 30 %. A grouped products kernel (all 36 components of a convolution, persistent, in 2-CTA clusters) passed correctness during the run; the step after it is to compute the input transform inside the products' operand producer, which removes a write and a read of the transformed tensor (2.25x the activation) per convolution.
- Eight GPUs across two trays: the code accepts up to 8 ranks (PR452 fixed world 8's all-gather), but world 8 was never in CI. A frame pipeline that passes each layer's two-frame context beat row bands at eight GPUs before the run (1.40-1.55 s against 2.12 s on f241, with library convolutions).