nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update STA README.md (#206)

Browse Source
This commit is contained in:
Zhang Peiyuan
2025-02-20 22:26:26 -08:00
committed by GitHub
parent e768b5ec5b
commit 09ab452610
5 changed files with 41 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+4 -4
View File
@@ -6,7 +6,7 @@ FastVideo is a lightweight framework for accelerating large video diffusion mode
<p align="center">
🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan" target="_blank">FastHunyuan</a> | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank">FastMochi</a> | 🎮 <a href="https://discord.gg/REBzDQTWWt" target="_blank"> Discord </a> | 🕹️ <a href="https://replicate.com/lucataco/fast-hunyuan-video" target="_blank"> Replicate </a>
🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan" target="_blank">FastHunyuan</a> | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank">FastMochi</a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> Slack </a>
</p>
@@ -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.
csrc/sliding_tile_attention/README.md
+19 -5
View File
@@ -2,9 +2,6 @@
# Sliding Tile Atteniton Kernel
<div align="center">
<img src=../../assets/sliding_tile_attn_map.png width="80%"/>
</div>
## 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.
<div align="center">
<img src=../../assets/sliding_tile_attn_map.png width="80%"/>
</div>
## Acknowledgement
We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.
We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.
csrc/sliding_tile_attention/st_attn/__init__.py
+1 -1
View File
@@ -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]
csrc/sliding_tile_attention/st_attn/st_attn_h100.cu
+16 -8
View File
@@ -161,7 +161,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> 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<D> 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<D> 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<D> 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<D> 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<D> 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<D> 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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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>,
scripts/inference/inference_hunyuan_STA.sh
+1 -1
View File
@@ -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 \