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