[docs] refactor attention docs (#964)

2025-12-26 15:21:49 -06:00
parent cbbb09d7b8
commit 87f8a2782d
9 changed files with 317 additions and 158 deletions
@@ -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 */
@@ -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`.
@@ -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.
@@ -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}
+}
+```
@@ -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}
+}
+```
@@ -1,16 +0,0 @@
-
-# 🔍 Demo
-This is is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
-
-<div style="text-align: center;">
-  <video controls width="800">
-    <source src="https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747" type="video/mp4">
-    Your browser does not support the video tag.
-  </video>
-</div>
-
-You can run STA using the following command:
-
-```bash
-bash scripts/inference/v1_inference_wan_STA.sh
-```
@@ -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
-```
@@ -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
-```
@@ -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