From 09ab4526106d07899a468a0e4cee109308efb160 Mon Sep 17 00:00:00 2001 From: Zhang Peiyuan Date: Thu, 20 Feb 2025 22:26:26 -0800 Subject: [PATCH] Update STA README.md (#206) --- README.md | 8 +++---- csrc/sliding_tile_attention/README.md | 24 +++++++++++++++---- .../st_attn/__init__.py | 2 +- .../st_attn/st_attn_h100.cu | 24 ++++++++++++------- scripts/inference/inference_hunyuan_STA.sh | 2 +- 5 files changed, 41 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index 3c08ca387..ce35191ea 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ FastVideo is a lightweight framework for accelerating large video diffusion mode

- 🤗 FastHunyuan | 🤗 FastMochi | 🎮 Discord | 🕹ī¸ Replicate + 🤗 FastHunyuan | 🤗 FastMochi | đŸŸŖđŸ’Ŧ Slack

@@ -59,12 +59,12 @@ sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo ### Inference HunyuanVideo with Sliding Tile Attention First, download the model: -``` +```bash python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model ``` We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only. -``` -scripts/inference/inference_hunyuan_STA.sh +```bash +sh scripts/inference/inference_hunyuan_STA.sh ``` ### Video Demos using STA + Teacache Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100. diff --git a/csrc/sliding_tile_attention/README.md b/csrc/sliding_tile_attention/README.md index 44a161c43..0f32d9e49 100644 --- a/csrc/sliding_tile_attention/README.md +++ b/csrc/sliding_tile_attention/README.md @@ -2,9 +2,6 @@ # Sliding Tile Atteniton Kernel -
- -
## Installation We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100. @@ -37,7 +34,11 @@ from st_attn import sliding_tile_attention # 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) + ``` @@ -46,9 +47,22 @@ out = sliding_tile_attention(q, k, v, window_size, text_length) python test/test_sta.py ``` -## Use STA with FastVideo +## 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. \ No newline at end of file +We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens. diff --git a/csrc/sliding_tile_attention/st_attn/__init__.py b/csrc/sliding_tile_attention/st_attn/__init__.py index 195acc976..9d1e0a238 100644 --- a/csrc/sliding_tile_attention/st_attn/__init__.py +++ b/csrc/sliding_tile_attention/st_attn/__init__.py @@ -31,5 +31,5 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text) if has_text: - _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True) + _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True) return hidden_states[:, :, :seq_length] diff --git a/csrc/sliding_tile_attention/st_attn/st_attn_h100.cu b/csrc/sliding_tile_attention/st_attn/st_attn_h100.cu index e0e8c4091..4675b0748 100644 --- a/csrc/sliding_tile_attention/st_attn/st_attn_h100.cu +++ b/csrc/sliding_tile_attention/st_attn/st_attn_h100.cu @@ -161,7 +161,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { tma::load_async(k_smem[(kv_idx+1)%K::stages], g.k, kv_tile_idx, k_smem_arrived[(kv_idx+1)%K::stages]); tma::expect_bytes(v_smem_arrived[(kv_idx+1)%K::stages], sizeof(v_tile)); tma::load_async(v_smem[(kv_idx+1)%K::stages], g.v, kv_tile_idx, v_smem_arrived[(kv_idx+1)%K::stages]); - wait(compute_done[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + kittens::wait(compute_done[(kv_idx)%K::stages], (kv_idx/K::stages)%2); } } else { int qt = seq_idx / 6 / (CH * CW); @@ -188,7 +188,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { tma::load_async(k_smem[count%K::stages], g.k, kv_tile_idx, k_smem_arrived[count%K::stages]); tma::expect_bytes(v_smem_arrived[count%K::stages], sizeof(v_tile)); tma::load_async(v_smem[count%K::stages], g.v, kv_tile_idx, v_smem_arrived[count%K::stages]); - wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2); + kittens::wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2); count += 1; } else { count += 1; @@ -204,7 +204,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { tma::load_async(k_smem[count%K::stages], g.k, kv_tile_idx, k_smem_arrived[count%K::stages]); tma::expect_bytes(v_smem_arrived[count%K::stages], sizeof(v_tile)); tma::load_async(v_smem[count%K::stages], g.v, kv_tile_idx, v_smem_arrived[count%K::stages]); - wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2); + kittens::wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2); count += 1; } } @@ -237,10 +237,10 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { kv_iters = CLAMP(DT*2+1, 1, CT) * CLAMP(DH*2+1, 1, CH) * CLAMP(DW*2+1, 1, CW) * 3 - 1 ; } - wait(qsmem_semaphore, 0); + kittens::wait(qsmem_semaphore, 0); for (auto kv_idx = 0; kv_idx <= kv_iters; kv_idx++) { - wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + kittens::wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[(kv_idx)%K::stages]); copy(max_vec_last_scaled, max_vec); @@ -270,7 +270,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { copy(att_block_mma, att_block); mul_row(o_reg, o_reg, max_vec_last_scaled); - wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + kittens::wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]); warpgroup::mma_async_wait(); @@ -281,7 +281,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { if constexpr(text_kv) { for (auto kv_idx = kv_iters + 1; kv_idx <= kv_iters + 3; kv_idx++) { - wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + kittens::wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[(kv_idx)%K::stages]); copy(max_vec_last_scaled, max_vec); @@ -316,7 +316,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { copy(att_block_mma, att_block); mul_row(o_reg, o_reg, max_vec_last_scaled); - wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + kittens::wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]); warpgroup::mma_async_wait(); @@ -482,6 +482,14 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, ); fwd_attend_ker<128, false, false, true, 1, 2, 2, 5, 6, 10><<>>(g); + } else if (kernel_t_size ==5 && kernel_h_size == 6 && kernel_w_size == 1){ + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, true, 2, 3, 0, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, true, 2, 3, 0, 5, 6, 10><<>>(g); + } else if (kernel_t_size ==5 && kernel_h_size == 3 && kernel_w_size == 5){ cudaFuncSetAttribute( fwd_attend_ker<128, false, false, true, 2, 1, 2, 5, 6, 10>, diff --git a/scripts/inference/inference_hunyuan_STA.sh b/scripts/inference/inference_hunyuan_STA.sh index 4fac6bebf..0ea063ca1 100644 --- a/scripts/inference/inference_hunyuan_STA.sh +++ b/scripts/inference/inference_hunyuan_STA.sh @@ -27,7 +27,7 @@ CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_p # Inference with STA only num_gpus=1 -mask_strategy_file_path=assets/mask_strategy.json +mask_strategy_file_path=assets/mask_strategy_hunyuan.json export MODEL_BASE=data/hunyuan CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \ fastvideo/sample/sample_t2v_hunyuan_STA.py \