diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 9de4fe587..88e140437 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -118,7 +118,7 @@ steps: - path: - "fastvideo/**" - "csrc/attn/video_sparse_attn/**" - - "csrc/attn/tk/**" + - "csrc/attn/video_sparse_attn/tk/**" - "csrc/attn/video_sparse_attn/setup.py" - "csrc/attn/video_sparse_attn/config_vsa.py" - "csrc/attn/video_sparse_attn/vsa.cpp" @@ -162,7 +162,7 @@ steps: queue: "default" - path: - "csrc/attn/video_sparse_attn/**" - - "csrc/attn/tk/**" + - "csrc/attn/video_sparse_attn/tk/**" - "csrc/attn/tests/test_vsa.py" - "csrc/attn/video_sparse_attn/setup.py" - "csrc/attn/video_sparse_attn/config_vsa.py" diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 3d8e8f667..428cb8a6e 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -105,12 +105,13 @@ jobs: - 'docker/Dockerfile.python3.12' sta-kernel-paths: &sta-kernel-paths - 'csrc/attn/sliding_tile_attn/**' + - 'csrc/attn/sliding_tile_attn/tk/**' - 'csrc/attn/sliding_tile_attn/setup.py' - 'csrc/attn/sliding_tile_attn/config_sta.py' - 'csrc/attn/sliding_tile_attn/st_attn.cpp' vsa-kernel-paths: &vsa-kernel-paths - 'csrc/attn/video_sparse_attn/**' - - 'csrc/attn/tk/**' + - 'csrc/attn/video_sparse_attn/tk/**' - 'csrc/attn/video_sparse_attn/setup.py' - 'csrc/attn/video_sparse_attn/config_vsa.py' - 'csrc/attn/video_sparse_attn/vsa.cpp' diff --git a/.gitmodules b/.gitmodules index e8b448235..bdb04686d 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,7 @@ -[submodule "csrc/attn/tk"] - path = csrc/attn/tk +[submodule "csrc/attn/video_sparse_attn/tk"] + path = csrc/attn/video_sparse_attn/tk + url = https://github.com/HazyResearch/ThunderKittens.git + +[submodule "csrc/attn/sliding_tile_attn/tk"] + path = csrc/attn/sliding_tile_attn/tk url = https://github.com/HazyResearch/ThunderKittens.git diff --git a/csrc/attn/MANIFEST.in b/csrc/attn/sliding_tile_attn/MANIFEST.in similarity index 51% rename from csrc/attn/MANIFEST.in rename to csrc/attn/sliding_tile_attn/MANIFEST.in index 386033830..cc38220c0 100644 --- a/csrc/attn/MANIFEST.in +++ b/csrc/attn/sliding_tile_attn/MANIFEST.in @@ -1,2 +1,2 @@ recursive-include tk * -include config.py +include config_sta.py diff --git a/csrc/attn/sliding_tile_attn/README.md b/csrc/attn/sliding_tile_attn/README.md new file mode 100644 index 000000000..3375c2068 --- /dev/null +++ b/csrc/attn/sliding_tile_attn/README.md @@ -0,0 +1,87 @@ + +# Attention Kernel Used in FastVideo + +## Sliding Tile Attention (STA) +We only support H100 for STA. + +### Installation +```bash +pip install st_attn +``` + +Install from source: + +```bash +git submodule update --init --recursive +python setup.py install +``` + +If you encounter error during installation, try below: +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 +``` +(If you use CUDA12.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 +``` + +### Usage +End-2-end inference with FastVideo: +```bash +bash scripts/inference/v1_inference_wan_STA.sh +``` + +If you want to use sliding tile attention in your custom model: +```python +from st_attn 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) +``` + + +### Test +```bash +python ../tests/test_sta.py # test STA +python ../tests/test_vsa.py # test VSA +``` +### Benchmark +```bash +python ../benchmarks/bench_sta.py +``` + + +### How Does STA Work? +We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. + + +https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747 + +## Why is STA Fast? +2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. + +STA removes mixed blocks. + + +
+ +
+ +## Acknowledgement + +We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens. diff --git a/csrc/attn/sliding_tile_attn/setup.py b/csrc/attn/sliding_tile_attn/setup.py index 09bf0432c..c36e71878 100644 --- a/csrc/attn/sliding_tile_attn/setup.py +++ b/csrc/attn/sliding_tile_attn/setup.py @@ -9,7 +9,7 @@ target = target.lower() # Package metadata PACKAGE_NAME = "st_attn" -VERSION = "0.0.5" +VERSION = "0.0.6" AUTHOR = "Hao AI Lab" DESCRIPTION = "Sliding Tile Atteniton Kernel Used in FastVideo" URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/sliding_tile_attention" diff --git a/csrc/attn/tk b/csrc/attn/sliding_tile_attn/tk similarity index 100% rename from csrc/attn/tk rename to csrc/attn/sliding_tile_attn/tk diff --git a/csrc/attn/video_sparse_attn/MANIFEST.in b/csrc/attn/video_sparse_attn/MANIFEST.in new file mode 100644 index 000000000..de701fcb9 --- /dev/null +++ b/csrc/attn/video_sparse_attn/MANIFEST.in @@ -0,0 +1,2 @@ +recursive-include tk * +include config_vsa.py diff --git a/csrc/attn/video_sparse_attn/README.md b/csrc/attn/video_sparse_attn/README.md new file mode 100644 index 000000000..c4a027592 --- /dev/null +++ b/csrc/attn/video_sparse_attn/README.md @@ -0,0 +1,61 @@ + + +# Attention Kernel Used in FastVideo + +## Video Sparse Attention (VSA) + +### Installation +We support H100 (via TK) and any other GPU (via triton) for VSA. + +```bash +pip install vsa +``` + +Install from source: + +```bash +git submodule update --init --recursive +python setup.py install +``` + + +If you encounter error during installation, try below: +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 +``` +(If you use CUDA12.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 +``` + +### Verify if you have successfully installed + +```bash +# test numerical +python ../tests/test_vsa.py +# (For H100) test speed +python ../benchmarks/bench_vsa_hopper.py +``` + +bench_vsa_hopper.py should print something like this: + +```bash +Using topk=76 kv blocks per q block (out of 768 total kv blocks) + +=== BLOCK SPARSE ATTENTION BENCHMARK === +Block Sparse Forward - TFLOPS: 5622.26 +Block Sparse Backward - TFLOPS: 3865.68 +``` + +## Acknowledgement + +We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens. diff --git a/csrc/attn/video_sparse_attn/setup.py b/csrc/attn/video_sparse_attn/setup.py index 909174033..02f683fc4 100644 --- a/csrc/attn/video_sparse_attn/setup.py +++ b/csrc/attn/video_sparse_attn/setup.py @@ -9,7 +9,7 @@ target = target.lower() # Package metadata PACKAGE_NAME = "vsa" -VERSION = "0.0.2" +VERSION = "0.0.3" AUTHOR = "Hao AI Lab" DESCRIPTION = "Video Sparse Attention Kernel Used in FastVideo" URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn/video_sparse_attn" diff --git a/csrc/attn/video_sparse_attn/tk b/csrc/attn/video_sparse_attn/tk new file mode 160000 index 000000000..6c27e28c8 --- /dev/null +++ b/csrc/attn/video_sparse_attn/tk @@ -0,0 +1 @@ +Subproject commit 6c27e28c8115d1839d9eeeb530913c184a75fc87