Skip to content

README.md

Latest commit

 

History

History
204 lines (152 loc) · 9.39 KB

README.md

File metadata and controls

204 lines (152 loc) · 9.39 KB
Raw

FastVideo Kernel

CUDA kernels for FastVideo video generation.

Kernel inventory

Compiled CUDA extensions (CMake, see the build summary printed at the end of every configure):

Extension Kernels Sources GPU arch Build gate
fastvideo_kernel._C.fastvideo_kernel_ops TurboDiffusion INT8 GEMM, quant, RMSNorm, LayerNorm csrc/turbodiffusion/ every arch in TORCH_CUDA_ARCH_LIST always built
same extension, optional part ThunderKittens sliding-tile attention (sta_fwd) and VSA block-sparse (block_sparse_fwd/bwd) csrc/attention/*_h100.cu Hopper sm_90a only FASTVIDEO_KERNEL_BUILD_TK (AUTO = ON iff 9.0a is in the arch list; always OFF on aarch64 hosts — TK headers don't compile there)
fp4attn_cuda, fp4quant_cuda FP4 attention + quantization ("attn_qat_infer", modified SageAttention3) attn_qat_infer/ consumer Blackwell sm_120a only, CUDA ≥ 12.8 FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER (AUTO = ON iff 12.0a is in the arch list)

Runtime-JIT kernels (no build step, ship in every wheel/image):

Kernels Where Used when
Triton: STA, VSA block-sparse, SLA, fused compress+topk, FP4 QAT training, quant/norm utils python/fastvideo_kernel/triton_kernels/ automatic fallback when the matching C++ op is absent (ops.py, turbodiffusion_ops.py)
FA4 CuTe-DSL block-sparse forward/backward (VSA-128/256 fastpath on sm_100) block_sparse_attn_cute_fwd.py optional flash_attn.cute dependency, see below
VMoBA moba_attn_varlen vmoba.py wraps flash-attn varlen

What gets built where, and when

Surface Trigger Leg TORCH_CUDA_ARCH_LIST TK FP4
PyPI wheels (.github/workflows/publish-kernel.yml) version bump in fastvideo-kernel/pyproject.toml on main, or manual dispatch x86_64 cu126 9.0a ON — (CUDA < 12.8)
x86_64 cu130 9.0a;12.0a ON ON
aarch64 cu130 10.0a;12.0a — ON
Docker images ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev (.github/workflows/infra-build-image.yml) docker/Dockerfile changes on main, or manual dispatch amd64 cuda12.6.3 + cuda13.0.0 9.0a ON —
arm64 cuda12.6.3 (GH200) 9.0a — (aarch64) —
arm64 cuda13.0.0 (GB10 / DGX Spark) 12.1 — —
Local ./build.sh manual probes the visible GPU via torch detected ON iff sm_90 (non-aarch64 host) ON iff sm_120

Notes:

  • No Docker image ships the FP4 kernels; only the x86_64/aarch64 cu130 wheels do.
  • On arm64 images (GH200 included) STA/VSA run on the Triton fallbacks, since TK never builds on aarch64.
  • Kernel tests run on Buildkite GPU CI for PRs touching fastvideo-kernel/** (see .buildkite/pipeline.yml).

Installation

Standard Installation (Local Development)

This will automatically detect your GPU architecture. If an NVIDIA Hopper (H100/sm_90a) GPU is detected, ThunderKittens kernels will be enabled. Otherwise, they will be skipped, and the package will use Triton fallbacks at runtime.

Before installation, set CUDA toolchain paths:

export CUDA_HOME=/usr/local/cuda
export CUDACXX=$CUDA_HOME/bin/nvcc
git submodule update --init --recursive
cd fastvideo-kernel
./build.sh

Rocm Build

If you are in a rocm environment without the compilation toolchaine of CUDA.

cd fastvideo-kernel
./build.sh --rocm

Optional: FA4 CuTe block-sparse backend (VSA-128/256 fastpath)

The VSA-128/256 fastpaths (tile volume 128 or 256, on NVIDIA Blackwell / sm_100) route to the FlashAttention-4 CuTe-DSL block-sparse kernel exposed as flash_attn.cute. This is an optional dependency: it is imported lazily, and video_sparse_attn transparently falls back to the Triton backend when it is absent (so the package is fully usable without it).

The symbols the fastpath needs (flash_attn.cute.block_sparsity.BlockSparseTensorsTorch and the public/private forward-backward bridges in flash_attn.cute.interface) are provided upstream by Dao-AILab/flash-attention. Pin to commit 14c377950125c70b7a9dabf9c561fca53715ac7d, the revision FastVideo pins as the flash-attn-4 source in the repo-root pyproject.toml; other revisions may have incompatible block-sparse forward/backward interfaces.

Install it under its distribution name so its own runtime stack resolves with it. Do not pre-install nvidia-cutlass-dsl by hand: this revision pins nvidia-cutlass-dsl==4.6.0.dev0 exactly, and a hand-installed 4.5.x floor either gets silently upgraded or, if something else holds it back, leaves the CuTe kernels broken.

pip install torchvision
pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@14c377950125c70b7a9dabf9c561fca53715ac7d#subdirectory=flash_attn/cute"

That resolves nvidia-cutlass-dsl to 4.6.0.dev0 and quack-kernels to 0.5.3, a combination this revision works with. A mismatched CuTe DSL only surfaces when the kernel JIT-compiles, so the error points at CuTe internals rather than at the install:

Error on first VSA-128/256 CuTe call Cause
TypeError: fmax() missing 1 required positional argument: 'b' nvidia-cutlass-dsl 4.5.x
AttributeError: module 'cutlass.cute.core' has no attribute 'ThrMma' quack-kernels older than 0.5.1
ImportError: cannot import name 'alloc_reserved_mbarrier' quack-kernels 0.6.2 or newer

An environment whose flash_attn.cute came from a prebuilt flash-attn wheel rather than from this pin hits the first row; that is what the overlay step in docker/Dockerfile works around.

The CuTe kernels JIT-compile on first use. Forward and backward are verified on Blackwell (sm_100) against tests/test_vsa128_*.py and tests/test_vsa256_*.py.

Usage

Sliding Tile Attention (STA) & Video Sparse Attention (VSA)

For detailed usage, please check the Attention Documentation.

from fastvideo_kernel import sliding_tile_attention, video_sparse_attn, moba_attn_varlen

# Example: Sliding Tile Attention
out = sliding_tile_attention(q, k, v, window_sizes, text_len)

# Example: Video Sparse Attention (with Triton fallback)
out = video_sparse_attn(q, k, v, block_sizes, block_sizes, topk=5)

# Example: VMoBA
out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)

Benchmark

Attn-QAT training

The default shape matches one sequence-parallel rank of the 4-GPU Wan2.1-T2V-1.3B MixKit recipe (B=1, H=3, L=31200, D=128):

cd fastvideo-kernel
python benchmarks/benchmark_attn_qat_train.py

The benchmark reports both conventional attention FLOPs and the extra matrix multiplications executed by the QAT straight-through path. Override --peak-tflops when running on a GPU other than RTX 5090.

The QAT kernel is entirely Triton and routes by architecture at runtime. SM100 uses a large-tile forward and split 64x64 backward for the production non-causal, head-dimension-128 configuration with a 16-aligned KV length. SM120 (including RTX 5090) keeps the previous tiling but joins the quantized and STE P@V operations and uses a shallower backward pipeline for long sequences. Set FASTVIDEO_ATTN_QAT_SM120_JOIN_QAT_PV=0 to compare against the split P@V path. Unsupported configurations retain the previous implementation. Set FASTVIDEO_ATTN_QAT_SM100_OPTIMIZED=0 to benchmark that previous path on SM100. Forward tuning is available through FASTVIDEO_ATTN_QAT_FWD_MODE=fast|balanced|reference; exact reference-order softmax statistics are controlled by FASTVIDEO_ATTN_QAT_FWD_EXACT_M and are disabled by default for maximum throughput. Set it to 1 for reference-order statistics and bitwise-compatible dV.

VSA (block-sparse) TFLOPs

After building/installing fastvideo-kernel, run:

cd fastvideo-kernel
python benchmarks/bench_vsa.py --batch_size 1 --num_heads 16 --head_dim 128 --q_seq_lens 49152 --topk 64

# VSA-256 FA4 CuTe forward/backward on Blackwell
python benchmarks/bench_vsa.py --block_size 256 --use_cute \
  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 20

# VSA-128 FA4 CuTe forward/backward on Blackwell
python benchmarks/bench_vsa.py --block_size 128 --use_cute \
  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 40

TurboDiffusion Kernels

This package also includes kernels from TurboDiffusion, including INT8 GEMM, Quantization, RMSNorm and LayerNorm.

Requirements

  • Runtime:
    • NVIDIA H100 (sm_90a) for C++ optimized kernels.
    • Any CUDA GPU for Triton-based fallbacks.
  • Build:
    • CUDA Toolkit 12.3+
    • CUDA_HOME must be set (for example, /usr/local/cuda)
    • CUDACXX must be set (for example, $CUDA_HOME/bin/nvcc)
    • C++20 compatible compiler (GCC 10+, Clang 11+)

Acknowledgement

This package structure and build system are based on sgl-kernel from the SGLang project.

The implementation of turbodiffusion kernels is adapted from TurboDiffusion. If you use these kernels, please cite:

@article{zhang2025turbodiffusion,
  title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times},
  author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun},
  journal={arXiv preprint arXiv:2512.16093},
  year={2025}
}