227 lines
11 KiB
Markdown
227 lines
11 KiB
Markdown
# 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:
|
|
|
|
```bash
|
|
export CUDA_HOME=/usr/local/cuda
|
|
export CUDACXX=$CUDA_HOME/bin/nvcc
|
|
```
|
|
|
|
```bash
|
|
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.
|
|
|
|
```bash
|
|
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:
|
|
|
|
```bash
|
|
./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](https://github.com/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.
|
|
|
|
```bash
|
|
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](../docs/attention/index.md).
|
|
|
|
```python
|
|
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`):
|
|
|
|
```bash
|
|
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:
|
|
|
|
```bash
|
|
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](https://github.com/thu-ml/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](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) from the SGLang project.
|
|
|
|
The implementation of `turbodiffusion` kernels is adapted from [TurboDiffusion](https://github.com/thu-ml/TurboDiffusion). If you use these kernels, please cite:
|
|
|
|
```bibtex
|
|
@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}
|
|
}
|
|
```
|