From 87f8a2782dcf7d2351007b4c2e5b5dd940592b66 Mon Sep 17 00:00:00 2001 From: William Lin Date: Fri, 26 Dec 2025 15:21:49 -0600 Subject: [PATCH] [docs] refactor attention docs (#964) --- docs/assets/custom.css | 8 +- docs/attention/developer/index.md | 179 ++++++++++++++++++++ docs/attention/index.md | 53 ++++++ docs/attention/sta/index.md | 36 ++++ docs/attention/vsa/index.md | 36 ++++ docs/sliding_tile_attention/demo.md | 16 -- docs/sliding_tile_attention/installation.md | 65 ------- docs/video_sparse_attention/installation.md | 65 ------- mkdocs.yml | 17 +- 9 files changed, 317 insertions(+), 158 deletions(-) create mode 100644 docs/attention/developer/index.md create mode 100644 docs/attention/index.md create mode 100644 docs/attention/sta/index.md create mode 100644 docs/attention/vsa/index.md delete mode 100644 docs/sliding_tile_attention/demo.md delete mode 100644 docs/sliding_tile_attention/installation.md delete mode 100644 docs/video_sparse_attention/installation.md diff --git a/docs/assets/custom.css b/docs/assets/custom.css index 7bf389a10..bb037f8c8 100644 --- a/docs/assets/custom.css +++ b/docs/assets/custom.css @@ -2,9 +2,11 @@ writing-mode: sideways-lr; white-space: nowrap; max-width: 0; - p { - margin: 0; - } +} + +/* Keep header cell paragraph content tight (avoid CSS nesting for compatibility) */ +.vertical-table-header th.head:not(.stub) p { + margin: 0; } /* Image sizing classes */ diff --git a/docs/attention/developer/index.md b/docs/attention/developer/index.md new file mode 100644 index 000000000..90f578c57 --- /dev/null +++ b/docs/attention/developer/index.md @@ -0,0 +1,179 @@ +# Adding a New Attention Backend + +FastVideo allows integrating new attention mechanisms easily. This guide walks you through adding a new backend (e.g., `MyNewAttn`). + +## 1. Implement the Backend (Python) + +Create a new file in `fastvideo/attention/backends/` (e.g., `mynew_attn.py`). + +Your implementation should inherit from `AttentionBackend` defined in `abstract.py`. + +```python +# fastvideo/attention/backends/mynew_attn.py +import torch +from .abstract import AttentionBackend +# Import the context manager to access metadata (optional) +from fastvideo.forward_context import get_forward_context + +# Import compiled kernel if applicable (see Section 2) +try: + # Import from the top-level package + from fastvideo_kernel import my_compiled_attn_func +except ImportError: + my_compiled_attn_func = None + +class MyNewAttnBackend(AttentionBackend): + def process_inputs(self, q, k, v, **kwargs): + # Pre-process inputs if necessary + return q, k, v + + def forward(self, q, k, v, **kwargs): + # Optional: Access extra metadata passed via ForwardContext + # Only needed if your backend requires global state (e.g. window_size) + try: + context = get_forward_context() + metadata = context.attn_metadata + # Example: window_size = metadata.window_size + except (AssertionError, AttributeError): + # Handle case where context is not set (e.g. standard inference) + pass + + if my_compiled_attn_func is not None: + return my_compiled_attn_func(q, k, v) + else: + # Fallback implementation (e.g., Triton or pure PyTorch) + return self.fallback_impl(q, k, v) +``` + +## 2. Passing Extra Information via ForwardContext (Optional) + +FastVideo uses a `ForwardContext` to pass global metadata (like current timestep, batch info, or custom attention configurations) to attention backends without changing the `forward` signature of every layer. **This is optional and only required if your backend needs dynamic per-step information.** + +To use this: +1. **Set Context**: In your pipeline or generation loop, use the `set_forward_context` context manager. +2. **Access Context**: Inside your attention backend, use `get_forward_context()`. + +See `docs/attention/sta/index.md` (Sliding Tile Attention) for an example of how complex configuration (window sizes) is passed this way. + +## 3. Adding Compiled Kernels (C++/CUDA) + +If your backend requires custom CUDA kernels, you need to add them to the `fastvideo-kernel` package. + +### A. Add Source Files +Place your kernel implementation files in `fastvideo-kernel/csrc/attention/`. +* `mynew_attn.cu` (CUDA implementation) +* `mynew_attn.h` (Optional headers) + +### B. Register in Extension +Update `fastvideo-kernel/csrc/common_extension.cpp` to expose your function to Python. + +```cpp +// 1. Declare external function +#ifdef COMPILE_MYNEW_ATTN +extern torch::Tensor mynew_attn_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v); +#endif + +// 2. Register in module +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + // ... other kernels ... + +#ifdef COMPILE_MYNEW_ATTN + m.def("mynew_attn_fwd", torch::wrap_pybind_function(mynew_attn_forward), "My New Attention Forward"); +#endif +} +``` + +### C. Update CMakeLists.txt +Update `fastvideo-kernel/CMakeLists.txt` to compile your new files. + +**Case 1: General CUDA Kernel (Runs on all GPUs)** +Add your source file directly to `EXTENSION_SOURCES` and define the compilation flag. + +```cmake +# Add to EXTENSION_SOURCES +list(APPEND EXTENSION_SOURCES csrc/attention/mynew_attn.cu) + +# Add compilation definition for common_extension.cpp +list(APPEND COMPILE_DEFS COMPILE_MYNEW_ATTN) +``` + +**Case 2: ThunderKittens Kernel (Hopper H100 Only)** +If your kernel uses ThunderKittens (TK), it requires specific architecture flags (`sm_90a`). Add it inside the `ENABLE_TK_KERNELS` block. + +```cmake +if(ENABLE_TK_KERNELS) + # Add source only if TK is enabled + list(APPEND EXTENSION_SOURCES csrc/attention/mynew_attn_tk.cu) + + # Add definition to guard registration + list(APPEND COMPILE_DEFS TK_COMPILE_MYNEW_ATTN) +endif() +``` + +### D. Expose in Python Ops +Update `fastvideo-kernel/python/fastvideo_kernel/ops.py` to make the function importable and handle fallbacks gracefully. + +```python +# fastvideo-kernel/python/fastvideo_kernel/ops.py + +# Try to load C++ extension symbols +try: + from fastvideo_kernel._C import fastvideo_kernel_ops + mynew_attn_fwd = getattr(fastvideo_kernel_ops, "mynew_attn_fwd", None) +except ImportError: + mynew_attn_fwd = None + +def my_compiled_attn_func(q, k, v): + # Runtime check: use C++ kernel if available, else fallback + if mynew_attn_fwd is not None: + return mynew_attn_fwd(q, k, v) + else: + # Call Triton/Python fallback + return mynew_attn_triton(q, k, v) +``` + +### E. Expose in Package Init +Update `fastvideo-kernel/python/fastvideo_kernel/__init__.py` to export the function. + +```python +from fastvideo_kernel.ops import ( + my_compiled_attn_func, + # ... +) + +__all__ = [ + "my_compiled_attn_func", + # ... +] +``` + +## 4. Register the Backend + +Update `fastvideo/attention/backends/__init__.py` to export your new class. + +```python +from .mynew_attn import MyNewAttnBackend +``` + +## 5. Platform Integration + +If your backend requires specific platform checks (e.g., checking for H100 support), handle that in `fastvideo/platforms/cuda.py` or within your backend's `__init__`. + +## 6. Add Documentation + +Create a new documentation page for your backend to explain its usage, installation (if custom kernels are needed), and features. + +1. **Create Directory**: `docs/attention/mynew_attn/` +2. **Create Index**: `docs/attention/mynew_attn/index.md` +3. **Update Navigation**: Add an entry to `mkdocs.yml` under the "Attention" tab. + +## Checklist + +* [ ] Created `fastvideo/attention/backends/mynew_attn.py`. +* [ ] (Optional) Added CUDA kernels in `fastvideo-kernel/csrc/attention/`. +* [ ] (Optional) Updated `common_extension.cpp` and `CMakeLists.txt`. +* [ ] (Optional) Exposed kernel in `fastvideo-kernel/python/fastvideo_kernel/ops.py`. +* [ ] (Optional) Exported kernel in `fastvideo-kernel/python/fastvideo_kernel/__init__.py`. +* [ ] Implemented `forward` method respecting the standard signature. +* [ ] Added unit tests in `tests/`. +* [ ] Added documentation in `docs/attention/` and updated `mkdocs.yml`. diff --git a/docs/attention/index.md b/docs/attention/index.md new file mode 100644 index 000000000..f77c7b742 --- /dev/null +++ b/docs/attention/index.md @@ -0,0 +1,53 @@ +# FastVideo Attention Kernels + +FastVideo provides highly optimized custom attention kernels to accelerate video generation. + +## Supported Kernels + +* **[Video Sparse Attention (VSA)](vsa/index.md)**: Sparse attention mechanism selecting top-k blocks. +* **[Sliding Tile Attention (STA)](sta/index.md)**: Optimized attention for window-based video generation. + +## General Build Instructions + +These instructions apply to building the `fastvideo-kernel` package from source, which includes both STA and VSA kernels. + +### Prerequisites + +* **PyTorch**: 2.5.0+ +* **CUDA**: 12.4+ (12.8 recommended for best performance) +* **C++ Compiler**: GCC 11+ (C++20 support required for ThunderKittens) + +Install system dependencies: + +```bash +sudo apt update +sudo apt install -y gcc-11 g++-11 clang-11 ninja-build + +# Set gcc-11 as default +sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 +``` + +Set up your CUDA environment variables (adjust version as needed): + +```bash +export CUDA_HOME=/usr/local/cuda-12.8 +export PATH=${CUDA_HOME}/bin:${PATH} +export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH +``` + +### Compile and Install + +Clone the repository and build the kernel: + +```bash +# Clone recursively to get ThunderKittens submodule +git clone --recursive https://github.com/hao-ai-lab/FastVideo.git +cd FastVideo/fastvideo-kernel + +# Build and install +./build.sh +``` + +The build script automatically detects your GPU architecture: +* **H100 (sm_90a)**: Compiles optimized C++ ThunderKittens kernels. +* **Other (A100, etc.)**: Skips C++ compilation; installs Python package with Triton kernels. diff --git a/docs/attention/sta/index.md b/docs/attention/sta/index.md new file mode 100644 index 000000000..a23460bbe --- /dev/null +++ b/docs/attention/sta/index.md @@ -0,0 +1,36 @@ +# Sliding Tile Attention (STA) + +Optimized attention for window-based video generation (e.g., HunyuanVideo). + +## Installation + +STA is included in the `fastvideo-kernel` package. See the [main Attention page](../index.md) for build instructions. + +## Usage + +```python +from fastvideo_kernel import sliding_tile_attention + +# q, k, v: [batch_size, num_heads, seq_length, head_dim] +# window_size: List of (t, h, w) tiles. Tile size is (6, 8, 8). +# text_length: Number of text tokens (0-256) + +out = sliding_tile_attention( + q, k, v, + window_size=[(3, 3, 3)], # Example window + text_length=256 +) +``` + +## Citation + +If you use Sliding Tile Attention in your research, please cite: + +```bibtex +@article{zhang2025fast, + title={Fast video generation with sliding tile attention}, + author={Zhang, Peiyuan and Chen, Yongqi and Su, Runlong and Ding, Hangliang and Stoica, Ion and Liu, Zhengzhong and Zhang, Hao}, + journal={arXiv preprint arXiv:2502.04507}, + year={2025} +} +``` diff --git a/docs/attention/vsa/index.md b/docs/attention/vsa/index.md new file mode 100644 index 000000000..6cc8b5e6f --- /dev/null +++ b/docs/attention/vsa/index.md @@ -0,0 +1,36 @@ +# Video Sparse Attention (VSA) + +Sparse attention mechanism selecting top-k blocks. + +## Installation + +VSA is included in the `fastvideo-kernel` package. See the [main Attention page](../index.md) for build instructions. + +## Usage + +```python +from fastvideo_kernel import video_sparse_attn + +# q, k, v: [batch_size, num_heads, seq_len, head_dim] +# variable_block_sizes: Number of valid tokens per block +# topk: Number of blocks to attend + +output = video_sparse_attn( + q, k, v, + variable_block_sizes=block_sizes, + topk=32 +) +``` + +## Citation + +If you use Video Sparse Attention in your research, please cite: + +```bibtex +@article{zhang2025vsa, + title={Vsa: Faster video diffusion with trainable sparse attention}, + author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao}, + journal={arXiv preprint arXiv:2505.13389}, + year={2025} +} +``` diff --git a/docs/sliding_tile_attention/demo.md b/docs/sliding_tile_attention/demo.md deleted file mode 100644 index 080b8e24c..000000000 --- a/docs/sliding_tile_attention/demo.md +++ /dev/null @@ -1,16 +0,0 @@ - -# 🔍 Demo -This is is a demo for 2D STA with window size (6,6) operating on a (10, 10) image. - -
- -
- -You can run STA using the following command: - -```bash -bash scripts/inference/v1_inference_wan_STA.sh -``` diff --git a/docs/sliding_tile_attention/installation.md b/docs/sliding_tile_attention/installation.md deleted file mode 100644 index 7da08aced..000000000 --- a/docs/sliding_tile_attention/installation.md +++ /dev/null @@ -1,65 +0,0 @@ - -# 🔧 Installation -You can install the Sliding Tile Attention package using - -``` -pip install st_attn -``` - -# Building from Source -We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently, we only have an implementation for H100s. -First, install C++20 for ThunderKittens: - -```bash -sudo apt update -sudo apt install gcc-11 g++-11 - -sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 - -sudo apt update -sudo apt install clang-11 -``` - -Set up CUDA environment (if using CUDA 12.4): - -```bash -export CUDA_HOME=/usr/local/cuda-12.4 -export PATH=${CUDA_HOME}/bin:${PATH} -export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH -``` - -Install FastVideo Kernel: - -```bash -cd fastvideo-kernel/ -git submodule update --init --recursive -./build.sh -``` - -# 🧪 Test - -```bash -pytest fastvideo-kernel/tests/test_correctness.py -``` - -# 📋 Usage - -```python -from fastvideo_kernel import sliding_tile_attention -# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. -# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256 -# a tile is a cube of size (6, 8, 8) -# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video. -# text_length: int ranging from 0 to 256 -# If your attention contains text token (Hunyuan) -out = sliding_tile_attention(q, k, v, window_size, text_length) -# If your attention does not contain text token (StepVideo) -out = sliding_tile_attention(q, k, v, window_size, 0, False) - -``` - -# 🚀Inference - -```bash -bash scripts/inference/v1_inference_wan_STA.sh -``` diff --git a/docs/video_sparse_attention/installation.md b/docs/video_sparse_attention/installation.md deleted file mode 100644 index 74b5ca7e6..000000000 --- a/docs/video_sparse_attention/installation.md +++ /dev/null @@ -1,65 +0,0 @@ - -# 🔧 Installation -You can install the Video Sparse Attention package using - -```bash -pip install vsa -``` - -# Building from Source -We support H100s (via ThunderKittens) and any other GPU (via Triton) for VSA. - -First, install C++20 for ThunderKittens (if using an H100): - -```bash -sudo apt update -sudo apt install gcc-11 g++-11 - -sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 - -sudo apt update -sudo apt install clang-11 -``` - -Set up CUDA environment (if using CUDA 12.8): - -```bash -export CUDA_HOME=/usr/local/cuda-12.8 -export PATH=${CUDA_HOME}/bin:${PATH} -export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH -``` - -Install FastVideo Kernel: - -```bash -cd fastvideo-kernel/ -git submodule update --init --recursive -./build.sh -``` - -# 🧪 Test - -```bash -pytest fastvideo-kernel/tests/test_correctness.py -``` - -# 📋 Usage - -```python -from fastvideo_kernel import video_sparse_attn - -# q, k, v: [batch_size, num_heads, seq_len, head_dim] -# variable_block_sizes: [num_blocks] - number of valid tokens in each block -# topk: int - number of top-k blocks to attend to -# block_size: int or tuple of 3 ints - size of each block (default: 64 tokens) -# compress_attn_weight: optional weight for compressed attention branch - -output = video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight) - -``` - -# 🚀Inference - -```bash -bash scripts/inference/v1_inference_wan_VSA.sh -``` diff --git a/mkdocs.yml b/mkdocs.yml index 409e2c347..09b1d8d64 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -89,9 +89,10 @@ markdown_extensions: - admonition - pymdownx.highlight: anchor_linenums: true - line_spans: __span + # line_spans: __span pygments_lang_class: true - pymdownx.inlinehilite + - pymdownx.snippets - pymdownx.superfences: custom_fences: - name: mermaid @@ -99,13 +100,10 @@ markdown_extensions: - pymdownx.tabbed: alternate_style: true - pymdownx.emoji - - pymdownx.superfences - pymdownx.arithmatex: generic: true - pymdownx.critic - pymdownx.details - - pymdownx.highlight - - pymdownx.inlinehilite - pymdownx.keys - pymdownx.mark - pymdownx.smartsymbols @@ -149,11 +147,11 @@ nav: - Distillation: - Data Preprocessing: distillation/data_preprocess.md - DMD: distillation/dmd.md - - Sliding Tile Attention: - - Installation: sliding_tile_attention/installation.md - - Demo: sliding_tile_attention/demo.md - - Video Sparse Attention: - - Installation: video_sparse_attention/installation.md + - Attention: + - Overview: attention/index.md + - Video Sparse Attention: attention/vsa/index.md + - Sliding Tile Attention: attention/sta/index.md + - Adding a New Attention Backend: attention/developer/index.md - Design: - Overview: design/overview.md - Developer Guide: @@ -164,6 +162,7 @@ nav: - RunPod: contributing/developer_env/runpod.md - Testing: contributing/testing.md - Profiling: contributing/profiling.md + - Adding a New Attention Backend: attention/developer/index.md - API Reference: - FastVideo: api/fastvideo.md