hao-ai-lab-FastVideo/fastvideo-kernel/README.md

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 (VSA-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-256 fastpath)

The VSA-256 fastpath (tile volume 256, on NVIDIA Blackwell / sm_100) routes 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, flash_attn.cute.interface._flash_attn_fwd) are provided upstream by Dao-AILab/flash-attention. Pin to commit 940cd9680f3315f2f06b43ab5bea2c2cf2d96806, the revision FastVideo pins as the flash-attn-4 source in the repo-root pyproject.toml; other revisions may have an incompatible _flash_attn_fwd signature.

pip install "nvidia-cutlass-dsl>=4.5.0" torchvision
pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute"

The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against tests/test_vsa256_forward*.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

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

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}
}