nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[perf] PAN2 sampling: Wan2.2 VAE decode kernel, split decode, fused DiT kernels

Browse Source 
- 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.
This commit is contained in:
William Lin
2026-11-02 21:42:30 +00:00
parent 3fb291b496
commit 552d7dfadc
46 changed files with 6834 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/contributing/env_vars.md
+2
View File
@@ -262,6 +262,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
| `FASTVIDEO_H3_VSA_FP4_K_SMOOTH` | bool | `0` | attention | Subtract the per-head key mean before FP4 key quantization, matching the emulator and QAD training numerics. Exact for the softmax; it only changes quantization error. |
| `FASTVIDEO_H3_GATE_CHUNK_ROWS` | int | `0` | performance | Rows per chunk when the VSA FP4 gate branch expands the pooled compression output back to token rows. 0 keeps the single full-size index_select; a positive value caps the transient at chunk_rows\*heads\*head_dim and is bit-identical. |
| `FASTVIDEO_H3_VSA_FINE_ATTENTION` | one of , emulator-mix8, emulator-mixmx, emulator-qk, emulator-bf16 | `""` | attention | Replace tile-128 fine attention, preserving VSA selection and compression, with the Triton emulator of the sparse FP4 kernel: emulator-mix8 = NVFP4 Q/K video, per-token FP8 Q/K prefix; emulator-mixmx = MXFP8 prefix; emulator-qk = NVFP4 Q/K everywhere; emulator-bf16 = unquantized BF16. emulator-bf16 runs on any CUDA GPU. Nonempty overrides take priority over FASTVIDEO_H3_VSA_FP4. Empty keeps the native route. |
| `FASTVIDEO_WAN22_VAE_DECODE_DIR` | path | `""` | performance | Directory of the Wan2.2 VAE decoder kernels (fastvideo-kernel/wan22_vae_decode) that AutoencoderKLWan loads, and builds on first use, when its config sets decode_kernel. Empty takes the copy in this checkout. |
| `FASTVIDEO_H3_VSA_HEADS_FIRST_TILE` | bool | `0` | performance | VSA-H3 64/128-token tiles: scatter rows straight into the heads-first layout the sparse kernels read, removing their per-call transpose copy of query, key and value. Bit-identical; no-grad, uncompiled CUDA only. |
| `FASTVIDEO_H3_VSA_TILE_FIRST` | bool | `0` | attention | Single-rank MiniMax-H3 VSA with one tile gather of the block input instead of separate Q/K/V/gate scatters. |
| `FASTVIDEO_H3_VSA_SM89_KERNEL` | one of original, bf16, int8 | `original` | attention | Fine-attention kernel for MiniMax-H3 VSA on sm_89: original, bf16, or int8 (INT8 QK, BF16 PV). |
@@ -335,6 +336,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
| `FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER` | int | `0` | test | Transformer layer that the MiniMax-H3 golden-gate test checks. Deprecated names: `MINIMAX_H3_GATE_LAYER`. |
| `FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT` | str | unset | test | Local MiniMax-H3 checkpoint for the golden-gate test. Deprecated names: `MINIMAX_H3_MODEL_ROOT`. |
| `FASTVIDEO_TEST_SD35_MODEL_DIR` | str | `stabilityai/stable-diffusion-3.5-medium` | test | Model for the Stable Diffusion 3.5 SSIM test. Deprecated names: `SD35_MODEL_DIR`. |
| `FASTVIDEO_TEST_PAN2_MODEL_DIR` | path | unset | test | PAN2 checkpoint root (with vae/) for the Wan2.2 VAE decode-kernel GPU tests; unset skips them. |
| `FASTVIDEO_TEST_TAEH3_REFERENCE_DIR` | str | unset | test | Upstream taehv checkout for the MLX TAEH3 parity test. Deprecated names: `TAEH3_REFERENCE_DIR`. |
| `FASTVIDEO_TEST_WAN_ANIMATE_MODEL_DIR` | str | unset | test | Local Wan2.2-Animate-14B checkpoint for the Wan-Animate weight tests. Deprecated names: `WAN_ANIMATE_MODEL_PATH`. |
docs/inference/pan2.md
+15 -1
View File
@@ -19,7 +19,7 @@ from the `pan2_t2v` and `pan2_i2v` presets:
| Resolution | 704x1248 (height and width must be multiples of 32) |
| Frames | 313, 13 s at 24 fps (`num_frames - 1` must be a multiple of 4) |
| Steps | 50, uniform sigmas from 1 to 0 with the checkpoint's flow shift (9 for `IFM/PAN2-TEST`, the reference sampler's value) |
| Guidance | 3.0, classifier-free guidance. Image-to-video keeps the first latent frame out of guidance (`i2v_skip_first_latent_cfg`, on by default, as in the reference sampler) |
| Guidance | 2.5, classifier-free guidance, the reference launcher's default (3 before 2026-10-23). Image-to-video keeps the first latent frame out of guidance (`i2v_skip_first_latent_cfg`, on by default, as in the reference sampler) |
| Negative prompt | None unless the request carries one. The reference sampler always sends the release's structured negative caption (`negative_prompt.json` of the PAN2 repository) with its three `time_range_sec` values set to the clip length, `0 to 20` for 481 frames at 24 fps |
PAN 2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw
@@ -67,6 +67,20 @@ The trimmed sequence can use either FlashAttention or Torch SDPA.
PAN2_NUM_GPUS=4 python examples/inference/basic/basic_pan2.py
```
## Performance switches
PAN2 sampling turns these on by default (`PAN2PipelineConfig`); pass any of them as `False` to
`VideoGenerator.from_pretrained` to turn it off:
| Switch | What it does |
|---|---|
| `vae_decode_kernel` | Decodes with fastvideo-kernel's Wan2.2 decoder on B200 / GB200 / B300 (704x1248, batch 1): bf16 with fp32 accumulation, 7x faster than the fp32 decode on a B300 at PSNR 61-63 dB and SSIM 0.998 against it, at a peak of about 20 GiB for a 5 s video. It builds on first use (minutes) into `WAN22_VAE_DECODE_BUILD_DIR`, against fastvideo-kernel's CUTLASS submodule (`git submodule update --init fastvideo-kernel/include/cutlass`, or set `CUTLASS_DIR`). Other GPUs and sizes keep the fp32 decode. |
| `vae_decode_kernel_parallel` | With 2, 4 or 8 GPUs, splits that decode over them instead of decoding the whole video on each: a 20 s video takes 3.0 s on one B300, 1.55 s on two, 0.79 s on four and 0.47 s on eight. |
| `dit_fused_kernels` | Runs the DiT's fused Triton kernels for Q/K RMSNorm + RoPE and LayerNorm + modulation. Their rounding differs slightly from the unfused path. |
On GPUs with room for every component (B200, B300), keep the text encoder resident as well:
`text_encoder_cpu_offload=False`.
## Memory and offloading
The example keeps the DiT on the GPUs and offloads the text encoder between requests. Sequence parallelism shards