hao-ai-lab-FastVideo/fastvideo-kernel

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)
same extension, optional part	MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks)	csrc/attention/block_sparse*_sm100a.cu	Data-center Blackwell sm_100a/sm_103a	ON iff 10.0a or 10.3a is in TORCH_CUDA_ARCH_LIST
same extension, optional part	fused NVLink Ulysses all-to-all	csrc/comm/ulysses_all_to_all.cu	CUDA	FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A (AUTO = ON with NCCL 2.29+ device headers and library; always OFF on ROCm)
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	Ulysses	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	AUTO	— (CUDA < 12.8)
		x86_64 cu130	9.0a;10.0a;12.0a	ON	AUTO	ON
		aarch64 cu130	10.0a;12.0a	—	AUTO	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	AUTO	—
		arm64 cuda12.6.3 (GH200)	9.0a	— (aarch64)	AUTO	—
		arm64 cuda13.0.0 (GB10 / DGX Spark)	12.1	—	AUTO	—
Local ./build.sh	manual	probes the visible GPU via torch	detected	ON iff sm_90 (non-aarch64 host)	AUTO	ON iff sm_120

Notes:

• No Docker image ships the FP4 kernels; only the x86_64/aarch64 cu130 wheels do.
• Both cu130 PyPI wheels ship native sm_100a and sm_103a images for the MiniMax-H3 VSA forward.
• On arm64 images (GH200 included) STA/VSA run on the Triton fallbacks, since TK never builds on aarch64.
• Ulysses AUTO builds only when CMake finds a NCCL library and device-API headers with the 2.29 initializers. Use -DFASTVIDEO_KERNEL_BUILD_ULYSSES_A2A=ON to require it or OFF to test the portable build.
• 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

The compiled extension needs the HIP CMake toolchain (hip-lang). Images that install ROCm as a pip SDK, such as the rocm/pytorch 7.14 images, ship hipcc but not that package, so build.sh --rocm falls back to the Python + Triton package there and says so in the configure summary. That package is all the video sparse attention path needs on ROCm. To ask for it explicitly, or to fail instead of falling back:

./build.sh --rocm --python-only                                    # Python + Triton only
CMAKE_ARGS="-DFASTVIDEO_KERNEL_BUILD_EXTENSION=ON" ./build.sh --rocm  # error if HIP cannot be configured

FASTVIDEO_KERNEL_BUILD_EXTENSION (AUTO/ON/OFF) is also read from the environment by a plain pip install --no-build-isolation ..

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. It resolves dense BF16 peak throughput for the full-GPU variants listed in benchmarks/device_specs.py. Unknown or partitioned devices still report timing and achieved TFLOPS, with MFU shown as N/A; pass a positive, finite --peak-tflops value to report MFU for those devices.

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