diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index dc80e5288..3b379873f 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -27,4 +27,4 @@ jobs: pip install pytest - name: Run Pytest run: | - pytest \ No newline at end of file + pytest --ignore csrc/sliding_tile_attention/test \ No newline at end of file diff --git a/.gitignore b/.gitignore index 87df90f58..c2dc3ede3 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,3 @@ -ucf101_stride4x4x4 __pycache__ *.mp4 .ipynb_checkpoints @@ -8,10 +7,8 @@ results/ build/ fastvideo.egg-info/ wandb/ -.idea *.ipynb *.jpg -*.mp3 *.safetensors *.mp4 *.png @@ -20,28 +17,6 @@ wandb/ *.pt cache_dir/ wandb/ -sample_video* -sample_image* -512* -720* -1024* -debug* -private* -caption* -*deepspeed* -revised* -129f* -all* -read* -YSH* -*pick* -*ysh* -hw* -257f* -513f* -taming* -221hw* -65x512x512 runs/ samples/ *validation/ @@ -51,3 +26,9 @@ outputs_video sbatch.sh *.out env +dist/ +*.o +**/build/ +**.egg-info +**.pyc +**.egg \ No newline at end of file diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 000000000..d284208d0 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "csrc/sliding_tile_attention/tk"] + path = csrc/sliding_tile_attention/tk + url = https://github.com/HazyResearch/ThunderKittens.git diff --git a/README.md b/README.md index 1d9557207..6a84dbf3a 100644 --- a/README.md +++ b/README.md @@ -4,8 +4,6 @@ FastVideo is a lightweight framework for accelerating large video diffusion models. -https://github.com/user-attachments/assets/064ac1d2-11ed-4a0c-955b-4d412a96ef30 -

🤗 FastHunyuan | 🤗 FastMochi | 🎮 Discord | 🕹ī¸ Replicate @@ -14,6 +12,7 @@ https://github.com/user-attachments/assets/064ac1d2-11ed-4a0c-955b-4d412a96ef30 FastVideo currently offers: (with more to come) +- [NEW!] [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/) that reduce HunyuanVideo's inference cost from 16 minutes to 5 minutes. - FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup. - First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model). - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan. @@ -22,21 +21,10 @@ FastVideo currently offers: (with more to come) Dev in progress and highly experimental. -## đŸŽĨ More Demos -Fast-Mochi comparison with original Mochi, achieving an 8X diffusion speed boost with the FastVideo framework. - -https://github.com/user-attachments/assets/5fbc4596-56d6-43aa-98e0-da472cf8e26c - -Comparison between OpenAI Sora, original Hunyuan and FastHunyuan - -https://github.com/user-attachments/assets/d323b712-3f68-42b2-952b-94f6a49c4836 - -Comparison between original FastHunyuan, LLM-INT8 quantized FastHunyuan and NF4 quantized FastHunyuan - -https://github.com/user-attachments/assets/cf89efb5-5f68-4949-a085-f41c1ef26c94 ## Change Log +- ```2025/02/14```: Release the inference code and kernel for Sliding Tile Attention(https://hao-ai-lab.github.io/blogs/sta/). - ```2025/01/13```: Support Lora finetuning for HunyuanVideo. - ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment. - ```2024/12/17```: `FastVideo` v1.0 is released. @@ -47,9 +35,18 @@ The code is tested on Python 3.10.0, CUDA 12.1 and H100. ``` ./env_setup.sh fastvideo ``` +To try sliding tile attention (optional), please follow the instruction in [csrc/sliding_tile_attention/README.md](csrc/sliding_tile_attention/README.md) to install STA. + ## 🚀 Inference +### Inference HunyuanVideo with Sliding Tile Attention +Coming soon... Merging commits now. +```bash +python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model +# Inference with FA2 full attn +``` + ### Inference FastHunyuan on single RTX4090 We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM. ```bash @@ -176,7 +173,7 @@ For Image-Video Mixture Fine-tuning, make sure to enable the `--group_frame` opt ## 🤝 Contributing -We welcome all contributions. Please run `bash format.sh` before submitting a pull request. +We welcome all contributions. Please run `bash format.sh --all` before submitting a pull request. ## 🔧 Testing Run `pytest` to verify the data preprocessing, checkpoint saving, and sequence parallel pipelines. We recommend adding corresponding test cases in the `test` folder to support your contribution. diff --git a/assets/sliding_tile_attn_map.png b/assets/sliding_tile_attn_map.png new file mode 100644 index 000000000..90a3bd565 Binary files /dev/null and b/assets/sliding_tile_attn_map.png differ diff --git a/csrc/sliding_tile_attention/README.md b/csrc/sliding_tile_attention/README.md new file mode 100644 index 000000000..1e65f548b --- /dev/null +++ b/csrc/sliding_tile_attention/README.md @@ -0,0 +1,41 @@ + + +# Sliding Tile Atteniton Kernel + +

+ +
+ +## Installation +We test our code on Pytorch 2.5.1 and CUDA 12.4. Currently we only have implementation on H100. +```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 +python setup.py install +``` + +## Usage + +```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 +out = sliding_tile_attention(q, k, v, window_size, text_length) +``` + + +## Test +```bash +python test/test_sta.py +``` + +## Use STA with FastVideo + + +## Acknowledgement + +We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens. \ No newline at end of file diff --git a/csrc/sliding_tile_attention/config.py b/csrc/sliding_tile_attention/config.py new file mode 100644 index 000000000..a5523ae48 --- /dev/null +++ b/csrc/sliding_tile_attention/config.py @@ -0,0 +1,16 @@ +### ADD TO THIS TO REGISTER NEW KERNELS +sources = { + 'attn': { + 'source_files': { + 'h100': + 'st_attn/st_attn_h100.cu' # define these source files for each GPU target desired. + } + } +} + +### WHICH KERNELS DO WE WANT TO BUILD? +# (oftentimes during development work you don't need to redefine them all.) +kernels = ['attn'] + +### WHICH GPU TARGET DO WE WANT TO BUILD FOR? +target = 'h100' diff --git a/csrc/sliding_tile_attention/setup.py b/csrc/sliding_tile_attention/setup.py new file mode 100644 index 000000000..b78804265 --- /dev/null +++ b/csrc/sliding_tile_attention/setup.py @@ -0,0 +1,64 @@ +import os +import subprocess + +from config import kernels, sources, target +from setuptools import find_packages, setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + +target = target.lower() + +# Set environment variables +tk_root = os.getenv('THUNDERKITTENS_ROOT', + os.path.abspath(os.path.join(os.getcwd(), 'tk/'))) +python_include = subprocess.check_output([ + 'python', '-c', "import sysconfig; print(sysconfig.get_path('include'))" +]).decode().strip() +torch_include = subprocess.check_output([ + 'python', '-c', + "import torch; from torch.utils.cpp_extension import include_paths; print(' '.join(['-I' + p for p in include_paths()]))" +]).decode().strip() +print('st_attn root:', tk_root) +print('Python include:', python_include) +print('Torch include directories:', torch_include) + +# CUDA flags +cuda_flags = [ + '-DNDEBUG', '-Xcompiler=-Wno-psabi', '-Xcompiler=-fno-strict-aliasing', + '--expt-extended-lambda', '--expt-relaxed-constexpr', + '-forward-unknown-to-host-compiler', '--use_fast_math', '-std=c++20', + '-O3', '-Xnvlink=--verbose', '-Xptxas=--verbose', + '-Xptxas=--warn-on-spills', f'-I{tk_root}/include', + f'-I{tk_root}/prototype', f'-I{python_include}', '-DTORCH_COMPILE' +] + torch_include.split() +cpp_flags = ['-std=c++20', '-O3'] + +if target == 'h100': + cuda_flags.append('-DKITTENS_HOPPER') + cuda_flags.append('-arch=sm_90a') +else: + raise ValueError(f'Target {target} not supported') + +source_files = ['st_attn.cpp'] +for k in kernels: + if target not in sources[k]['source_files']: + raise KeyError( + f'Target {target} not found in source files for kernel {k}') + if isinstance(sources[k]['source_files'][target], list): + source_files.extend(sources[k]['source_files'][target]) + else: + source_files.append(sources[k]['source_files'][target]) + cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}') + +setup(name='st_attn', + version="0.0.0", + packages=find_packages(), + ext_modules=[ + CUDAExtension('st_attn_cuda', + sources=source_files, + extra_compile_args={ + 'cxx': cpp_flags, + 'nvcc': cuda_flags + }, + libraries=['cuda']) + ], + cmdclass={'build_ext': BuildExtension}) diff --git a/csrc/sliding_tile_attention/st_attn.cpp b/csrc/sliding_tile_attention/st_attn.cpp new file mode 100644 index 000000000..7d108ff11 --- /dev/null +++ b/csrc/sliding_tile_attention/st_attn.cpp @@ -0,0 +1,24 @@ +#include +#include + +#include +#include +#include + +#include + + +#ifdef TK_COMPILE_ATTN +extern torch::Tensor sta_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text +); +#endif + + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.doc() = "Sliding Block Attention Kernels"; // optional module docstring + +#ifdef TK_COMPILE_ATTN + m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention, assuming tile size is (6,8,8)"); +#endif +} \ No newline at end of file diff --git a/csrc/sliding_tile_attention/st_attn/__init__.py b/csrc/sliding_tile_attention/st_attn/__init__.py new file mode 100644 index 000000000..aba12afe2 --- /dev/null +++ b/csrc/sliding_tile_attention/st_attn/__init__.py @@ -0,0 +1,37 @@ +import math + +import torch +from st_attn_cuda import sta_fwd + + +def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length): + seq_length = q_all.shape[2] + assert q_all.shape[ + 2] == 115456, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels" + assert q_all.shape[1] == len( + window_size), "Number of heads must match the number of window sizes" + target_size = math.ceil(seq_length / 384) * 384 + pad_size = target_size - seq_length + if pad_size > 0: + q_all = torch.cat([q_all, q_all[:, :, -pad_size:]], dim=2) + k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2) + v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2) + + hidden_states = torch.empty_like(q_all) + for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size): + q_head, k_head, v_head, o_head = q_all[:, head_index:head_index + + 1], k_all[:, head_index: + head_index + + 1], v_all[:, + head_index: + head_index + + 1], hidden_states[:, + head_index: + head_index + + + 1] + _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, + w_kernel, text_length, False) + + _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, 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 new file mode 100644 index 000000000..aaad99678 --- /dev/null +++ b/csrc/sliding_tile_attention/st_attn/st_attn_h100.cu @@ -0,0 +1,551 @@ +// # Define TORCH_COMPILE macro + +#include "kittens.cuh" +#include +#include +#include + +#define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value))) +#define ABS(x) ((x) < 0 ? -(x) : (x)) + +constexpr int CONSUMER_WARPGROUPS = (3); +constexpr int PRODUCER_WARPGROUPS = (1); +constexpr int NUM_WARPGROUPS = (CONSUMER_WARPGROUPS+PRODUCER_WARPGROUPS); +constexpr int NUM_WORKERS = (NUM_WARPGROUPS*kittens::WARPGROUP_WARPS); + +using namespace kittens; +namespace cg = cooperative_groups; + +template struct fwd_attend_ker_tile_dims {}; +template<> struct fwd_attend_ker_tile_dims<64> { + constexpr static int tile_width = (64); + constexpr static int qo_height = (4*16); + constexpr static int kv_height = (8*16); + constexpr static int stages = (4); +}; +template<> struct fwd_attend_ker_tile_dims<128> { + constexpr static int tile_width = (128); + constexpr static int qo_height = (4*16); + constexpr static int kv_height = (8*16); + constexpr static int stages = (2); +}; + +template struct fwd_globals { + using q_tile = st_bf::qo_height, fwd_attend_ker_tile_dims::tile_width>; + using k_tile = st_bf::kv_height, fwd_attend_ker_tile_dims::tile_width>; + using v_tile = st_bf::kv_height, fwd_attend_ker_tile_dims::tile_width>; + using l_col_vec = col_vec::qo_height, fwd_attend_ker_tile_dims::tile_width>>; + using o_tile = st_bf::qo_height, fwd_attend_ker_tile_dims::tile_width>; + + using q_gl = gl; + using k_gl = gl; + using v_gl = gl; + using l_gl = gl; + using o_gl = gl; + + q_gl q; + k_gl k; + v_gl v; + l_gl l; + o_gl o; + + const int N; + const int text_L; + const int hr; +}; + + +template +__global__ __launch_bounds__((NUM_WORKERS)*kittens::WARP_THREADS, 1) +void fwd_attend_ker(const __grid_constant__ fwd_globals g) { + extern __shared__ int __shm[]; + tma_swizzle_allocator al((int*)&__shm[0]); + int warpid = kittens::warpid(), warpgroupid = warpid/kittens::WARPGROUP_WARPS; + + using K = fwd_attend_ker_tile_dims; + + using q_tile = st_bf; + using k_tile = st_bf; + using v_tile = st_bf; + using l_col_vec = col_vec>; + using o_tile = st_bf; + + q_tile (&q_smem)[CONSUMER_WARPGROUPS] = al.allocate(); + k_tile (&k_smem)[K::stages] = al.allocate(); + v_tile (&v_smem)[K::stages] = al.allocate(); + l_col_vec (&l_smem)[CONSUMER_WARPGROUPS] = al.allocate(); + auto (*o_smem) = reinterpret_cast(q_smem); + + int kv_blocks = g.N / (K::kv_height); + int kv_head_idx = blockIdx.y / g.hr; + int seq_idx; + if constexpr (text_q) { + seq_idx = CT * CH * CW * 6.0 + blockIdx.x * CONSUMER_WARPGROUPS; + } else { + seq_idx = blockIdx.x * CONSUMER_WARPGROUPS; + } + __shared__ kittens::semaphore qsmem_semaphore, k_smem_arrived[K::stages], v_smem_arrived[K::stages], compute_done[K::stages]; + if (threadIdx.x == 0) { + init_semaphore(qsmem_semaphore, 0, 1); + for(int j = 0; j < K::stages; j++) { + init_semaphore(k_smem_arrived[j], 0, 1); + init_semaphore(v_smem_arrived[j], 0, 1); + init_semaphore(compute_done[j], CONSUMER_WARPGROUPS, 0); + } + + tma::expect_bytes(qsmem_semaphore, sizeof(q_smem)); + + for (int wg = 0; wg < CONSUMER_WARPGROUPS; wg++) { + coord q_tile_idx = {blockIdx.z, blockIdx.y, (seq_idx) + wg, 0}; + tma::load_async(q_smem[wg], g.q, q_tile_idx, qsmem_semaphore); + } + + if constexpr (text_q){ + for (int j = 0; j < K::stages - 1; j++) { + coord kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0}; + tma::expect_bytes(k_smem_arrived[j], sizeof(k_tile)); + tma::load_async(k_smem[j], g.k, kv_tile_idx, k_smem_arrived[j]); + tma::expect_bytes(v_smem_arrived[j], sizeof(v_tile)); + tma::load_async(v_smem[j], g.v, kv_tile_idx, v_smem_arrived[j]); + } + } else { + int qt = seq_idx / 6 / (CH * CW); + int qh = (seq_idx / 6) % (CH * CW) / CW; + int qw = (seq_idx / 6) % CW; + qt = CLAMP(qt, DT, CT-DT-1); + qh = CLAMP(qh, DH, CH-DH-1); + qw = CLAMP(qw, DW, CW-DW-1); + int count = 0; + int j = 0; + while (count < K::stages - 1) { + int kt = j / 3 / (CH * CW); + int kh = (j / 3) % (CH * CW) / CW; + int kw = (j / 3) % CW; + bool mask = (ABS(qt - kt) <= DT) && (ABS(qh - kh) <= DH) && (ABS(qw - kw) <= DW); + if (mask){ + coord kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0}; + tma::expect_bytes(k_smem_arrived[count], sizeof(k_tile)); + tma::load_async(k_smem[count], g.k, kv_tile_idx, k_smem_arrived[count]); + tma::expect_bytes(v_smem_arrived[count], sizeof(v_tile)); + tma::load_async(v_smem[count], g.v, kv_tile_idx, v_smem_arrived[count]); + count += 1; + } + j += 1; + } + } + } + __syncthreads(); + + int pipe_idx = K::stages - 1; + + if(warpgroupid == NUM_WARPGROUPS-1) { + warpgroup::decrease_registers<32>(); + + int kv_iters; + if constexpr (is_causal) { + kv_iters = (seq_idx * (K::qo_height/kittens::TILE_ROW_DIM)) - 1 + (CONSUMER_WARPGROUPS * (K::qo_height/kittens::TILE_ROW_DIM)); + kv_iters = ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM)) == 0) ? (0) : ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM)) - 1); + } + else { kv_iters = kv_blocks-2;} + + if(warpid == NUM_WORKERS-4) { + if constexpr (text_q){ + for (auto kv_idx = pipe_idx - 1; kv_idx <= kv_iters; kv_idx++) { + coord kv_tile_idx = {blockIdx.z, kv_head_idx, kv_idx + 1, 0}; + tma::expect_bytes(k_smem_arrived[(kv_idx+1)%K::stages], sizeof(k_tile)); + 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); + } + } else { + int qt = seq_idx / 6 / (CH * CW); + int qh = (seq_idx / 6) % (CH * CW) / CW; + int qw = (seq_idx / 6) % CW; + qt = CLAMP(qt, DT, CT-DT-1); + qh = CLAMP(qh, DH, CH-DH-1); + qw = CLAMP(qw, DW, CW-DW-1); + int k_t_min = CLAMP(qt-DT, 0, CT-1); + int k_t_max = CLAMP(qt+DT, 0, CT-1); + int k_h_min = CLAMP(qh-DH, 0, CH-1); + int k_h_max = CLAMP(qh+DH, 0, CH-1); + int k_w_min = CLAMP(qw-DW, 0, CW-1); + int k_w_max = CLAMP(qw+DW, 0, CW-1); + int count = 0; + for (int kt = k_t_min; kt <= k_t_max; kt++) { + for (int kh = k_h_min; kh <= k_h_max; kh++) { + for (int kw = k_w_min; kw <= k_w_max; kw++) { + for (int j = 0; j <= 2; j++){ + if (count >= K::stages - 1) { + int index = ((kt * (CH * CW)) + (kh * CW) + kw) * 3 + j; + coord kv_tile_idx = {blockIdx.z, kv_head_idx, index, 0}; + tma::expect_bytes(k_smem_arrived[count%K::stages], sizeof(k_tile)); + 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); + count += 1; + } else { + count += 1; + } + } + } + } + } + // for text + for (int index = kv_blocks - 3; index < kv_blocks; index++) { + coord kv_tile_idx = {blockIdx.z, kv_head_idx, index, 0}; + tma::expect_bytes(k_smem_arrived[count%K::stages], sizeof(k_tile)); + 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); + count += 1; + } + } + + + } + } + else { + warpgroup::increase_registers<160>(); + + rt_fl<16, K::kv_height> att_block; + rt_bf<16, K::kv_height> att_block_mma; + rt_fl<16, K::tile_width> o_reg; + + col_vec> max_vec, norm_vec, max_vec_last_scaled, max_vec_scaled; + + neg_infty(max_vec); + zero(norm_vec); + zero(o_reg); + + int kv_iters; + if constexpr (is_causal) { + kv_iters = (seq_idx * 4) - 1 + (CONSUMER_WARPGROUPS * 4); + kv_iters = (kv_iters/8); + } + else if constexpr (text_q){ + // the last three kv blocks are for text, we process them separately + kv_iters = kv_blocks - 1 - 3; + } else { + 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); + for (auto kv_idx = 0; kv_idx <= kv_iters; kv_idx++) { + + 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); + if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); } + else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); } + + warpgroup::mma_async_wait(); + + row_max(max_vec, att_block, max_vec); + + if constexpr (D == 64) { + mul(att_block, att_block, 1.44269504089f*0.125f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f); + } + else { + mul(att_block, att_block, 1.44269504089f*0.08838834764f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f); + } + + sub_row(att_block, att_block, max_vec_scaled); + exp2(att_block, att_block); + sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled); + exp2(max_vec_last_scaled, max_vec_last_scaled); + mul(norm_vec, norm_vec, max_vec_last_scaled); + row_sum(norm_vec, att_block, norm_vec); + add(att_block, att_block, 0.f); + 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); + + warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]); + warpgroup::mma_async_wait(); + + if(warpgroup::laneid() == 0) arrive(compute_done[(kv_idx)%K::stages], 1); + } + // the last three kv blocks are for text, we process them separately + 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); + warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[(kv_idx)%K::stages]); + + copy(max_vec_last_scaled, max_vec); + if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); } + else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); } + + warpgroup::mma_async_wait(); + // apply non-pad mask + int offset = g.text_L - (kv_idx - (kv_iters + 1)) * K::kv_height; + // printf("k_idx_start: %d, k_idx_end: %d, text_end: %d, offset: %d\n", k_idx_start, k_idx_end, text_end, offset); + right_fill(att_block, att_block, offset, base_types::constants::neg_infty()); + + + row_max(max_vec, att_block, max_vec); + + if constexpr (D == 64) { + mul(att_block, att_block, 1.44269504089f*0.125f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f); + } + else { + mul(att_block, att_block, 1.44269504089f*0.08838834764f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f); + } + + sub_row(att_block, att_block, max_vec_scaled); + exp2(att_block, att_block); + sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled); + exp2(max_vec_last_scaled, max_vec_last_scaled); + mul(norm_vec, norm_vec, max_vec_last_scaled); + row_sum(norm_vec, att_block, norm_vec); + add(att_block, att_block, 0.f); + 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); + + warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]); + warpgroup::mma_async_wait(); + + if(warpgroup::laneid() == 0) arrive(compute_done[(kv_idx)%K::stages], 1); + } + + div_row(o_reg, o_reg, norm_vec); + warpgroup::store(o_smem[warpgroupid], o_reg); + warpgroup::sync(warpgroupid+4); + + if (warpid % 4 == 0) { + coord o_tile_idx = {blockIdx.z, blockIdx.y, (seq_idx) + warpgroupid, 0}; + tma::store_async(g.o, o_smem[warpgroupid], o_tile_idx); + } + + mul(max_vec_scaled, max_vec_scaled, 0.69314718056f); + log(norm_vec, norm_vec); + add(norm_vec, norm_vec, max_vec_scaled); + + if constexpr (D == 64) { mul(norm_vec, norm_vec, -8.0f); } + else { mul(norm_vec, norm_vec, -11.313708499f); } + + warpgroup::store(l_smem[warpgroupid], norm_vec); + warpgroup::sync(warpgroupid+4); + + if (warpid % 4 == 0) { + coord tile_idx = {blockIdx.z, blockIdx.y, 0, (seq_idx) + warpgroupid}; + tma::store_async(g.l, l_smem[warpgroupid], tile_idx); + } + tma::store_async_wait(); + } +} + + + +#include "pyutils/torch_helpers.cuh" +#include +#include + +torch::Tensor +sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text) +{ + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + + auto batch = q.size(0); + auto seq_len = q.size(2); + auto head_dim = q.size(3); + auto qo_heads = q.size(1); + auto kv_heads = k.size(1); + + // check to see that these dimensions match for all inputs + TORCH_CHECK(q.size(0) == batch, "Q batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(k.size(0) == batch, "K batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(v.size(0) == batch, "V batch dimension - idx 0 - must match for all inputs"); + + TORCH_CHECK(q.size(2) == seq_len, "Q sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(k.size(2) == seq_len, "K sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(v.size(2) == seq_len, "V sequence length dimension - idx 2 - must match for all inputs"); + + TORCH_CHECK(q.size(3) == head_dim, "Q head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(k.size(3) == head_dim, "K head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(v.size(3) == head_dim, "V head dimension - idx 3 - must match for all non-vector inputs"); + + TORCH_CHECK(qo_heads >= kv_heads, "QO heads must be greater than or equal to KV heads"); + TORCH_CHECK(qo_heads % kv_heads == 0, "QO heads must be divisible by KV heads"); + TORCH_CHECK(q.size(1) == qo_heads, "QO head dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(k.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(v.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs"); + + auto hr = qo_heads / kv_heads; + + c10::BFloat16* q_ptr = q.data_ptr(); + c10::BFloat16* k_ptr = k.data_ptr(); + c10::BFloat16* v_ptr = v.data_ptr(); + + bf16* d_q = reinterpret_cast(q_ptr); + bf16* d_k = reinterpret_cast(k_ptr); + bf16* d_v = reinterpret_cast(v_ptr); + + + + torch::Tensor l_vec = torch::empty({static_cast(batch), + static_cast(qo_heads), + static_cast(seq_len), + static_cast(1)}, + torch::TensorOptions().dtype(torch::kFloat).device(q.device()).memory_format(at::MemoryFormat::Contiguous)); + + + bf16* o_ptr = reinterpret_cast(o.data_ptr()); + bf16* d_o = reinterpret_cast(o_ptr); + + float* l_ptr = reinterpret_cast(l_vec.data_ptr()); + float* d_l = reinterpret_cast(l_ptr); + + cudaDeviceSynchronize(); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + + if (head_dim == 128) { + using q_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>; + using k_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>; + using v_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>; + using l_col_vec = col_vec::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>>; + using o_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>; + + using q_global = gl; + using k_global = gl; + using v_global = gl; + using l_global = gl; + using o_global = gl; + + using globals = fwd_globals<128>; + + q_global qg_arg{d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; + k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; + v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; + l_global lg_arg{d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + o_global og_arg{d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; + + globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast(seq_len), static_cast(text_length), static_cast(hr)}; + + auto mem_size = kittens::MAX_SHARED_MEMORY; + auto threads = NUM_WORKERS * kittens::WARP_THREADS; + + // TORCH_CHECK(seq_len % (CONSUMER_WARPGROUPS*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 192"); + dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM*4-2), qo_heads, batch); + dim3 grid_text(2, qo_heads, batch); + if (!process_text) { + if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) { + + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 1, 1, 1, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 1, 1, 1, 5, 6, 10><<>>(g); + + } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) { + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 1, 1, 2, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 1, 1, 2, 5, 6, 10><<>>(g); + + } else if (kernel_t_size == 5 && kernel_h_size == 3 && kernel_w_size == 3) { + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 2, 1, 1, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 2, 1, 1, 5, 6, 10><<>>(g); + + }else if (kernel_t_size ==3 && kernel_h_size == 5 && kernel_w_size == 5){ + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 1, 2, 2, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 1, 2, 2, 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, 2, 1, 2, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 2, 1, 2, 5, 6, 10><<>>(g); + + } else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 5){ + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 2, 2, 2, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 2, 2, 2, 5, 6, 10><<>>(g); + + } else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 7){ + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 2, 2, 3, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 2, 2, 3, 5, 6, 10><<>>(g); + } else if (kernel_t_size == 5 && kernel_h_size == 6 && kernel_w_size == 10){ + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 2, 3, 5, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 2, 3, 5, 5, 6, 10><<>>(g); + } else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 1){ + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 2, 0, 0, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 2, 0, 0, 5, 6, 10><<>>(g); + } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){ + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 0, 3, 5, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 0, 3, 5, 5, 6, 10><<>>(g); + } else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 10){ + cudaFuncSetAttribute( + fwd_attend_ker<128, false, false, 2, 0, 5, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, false, 2, 0, 5, 5, 6, 10><<>>(g); + } else { + // print error + std::cout << "Invalid kernel size" << std::endl; + //print kernel size + std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl; + } + } else { + cudaFuncSetAttribute( + fwd_attend_ker<128, false, true, 1, 1, 1, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, true, 1, 1, 1, 5, 6, 10><<>>(g); + } + + + CHECK_CUDA_ERROR(cudaGetLastError()); + cudaStreamSynchronize(stream); + } + + return o; + cudaDeviceSynchronize(); +} diff --git a/csrc/sliding_tile_attention/test/flex_sta_ref.py b/csrc/sliding_tile_attention/test/flex_sta_ref.py new file mode 100644 index 000000000..e605c0219 --- /dev/null +++ b/csrc/sliding_tile_attention/test/flex_sta_ref.py @@ -0,0 +1,80 @@ +from typing import Tuple + +import torch +from torch import BoolTensor, IntTensor +from torch.nn.attention.flex_attention import create_block_mask + +# Peiyuan: This is neccesay. Dont know why. see https://github.com/pytorch/pytorch/issues/135028 +torch._inductor.config.realize_opcount_threshold = 100 + + +def generate_sta_mask(canvas_twh, kernel_twh, tile_twh, text_length): + """Generates a 3D NATTEN attention mask with a given kernel size. + + Args: + canvas_t: The time dimension of the canvas. + canvas_h: The height of the canvas. + canvas_w: The width of the canvas. + kernel_t: The time dimension of the kernel. + kernel_h: The height of the kernel. + kernel_w: The width of the kernel. + """ + canvas_t, canvas_h, canvas_w = canvas_twh + kernel_t, kernel_h, kernel_w = kernel_twh + tile_t_size, tile_h_size, tile_w_size = tile_twh + total_tile_size = tile_t_size * tile_h_size * tile_w_size + canvas_tile_t, canvas_tile_h, canvas_tile_w = canvas_t // tile_t_size, canvas_h // tile_h_size, canvas_w // tile_w_size + img_seq_len = canvas_t * canvas_h * canvas_w + + def get_tile_t_x_y( + idx: IntTensor) -> Tuple[IntTensor, IntTensor, IntTensor]: + tile_id = idx // total_tile_size + tile_t = tile_id // (canvas_tile_h * canvas_tile_w) + tile_h = (tile_id % (canvas_tile_h * canvas_tile_w)) // canvas_tile_w + tile_w = tile_id % canvas_tile_w + return tile_t, tile_h, tile_w + + def sta_mask_3d( + b: IntTensor, + h: IntTensor, + q_idx: IntTensor, + kv_idx: IntTensor, + ) -> BoolTensor: + q_t_tile, q_x_tile, q_y_tile = get_tile_t_x_y(q_idx) + kv_t_tile, kv_x_tile, kv_y_tile = get_tile_t_x_y(kv_idx) + # kernel nominally attempts to center itself on the query, but kernel center + # is clamped to a fixed distance (kernel half-length) from the canvas edge + kernel_center_t = q_t_tile.clamp(kernel_t // 2, + (canvas_tile_t - 1) - kernel_t // 2) + kernel_center_x = q_x_tile.clamp(kernel_h // 2, + (canvas_tile_h - 1) - kernel_h // 2) + kernel_center_y = q_y_tile.clamp(kernel_w // 2, + (canvas_tile_w - 1) - kernel_w // 2) + time_mask = (kernel_center_t - kv_t_tile).abs() <= kernel_t // 2 + hori_mask = (kernel_center_x - kv_x_tile).abs() <= kernel_h // 2 + vert_mask = (kernel_center_y - kv_y_tile).abs() <= kernel_w // 2 + image_mask = (q_idx < img_seq_len) & (kv_idx < img_seq_len) + image_to_text_mask = (q_idx < img_seq_len) & ( + kv_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length) + text_to_all_mask = (q_idx >= img_seq_len) & (kv_idx < + img_seq_len + text_length) + return (image_mask & time_mask & hori_mask + & vert_mask) | image_to_text_mask | text_to_all_mask + + sta_mask_3d.__name__ = f"natten_3d_c{canvas_t}x{canvas_w}x{canvas_h}_k{kernel_t}x{kernel_w}x{kernel_h}" + return sta_mask_3d + + +def get_sliding_tile_attention_mask(kernel_size, tile_size, img_size, + text_length, device): + img_seq_len = img_size[0] * img_size[1] * img_size[2] + image_mask = generate_sta_mask(img_size, kernel_size, tile_size, + text_length) + mask = create_block_mask(image_mask, + B=None, + H=None, + Q_LEN=img_seq_len + 256, + KV_LEN=img_seq_len + 256, + device=device, + _compile=True) + return mask diff --git a/csrc/sliding_tile_attention/test/test_sta.py b/csrc/sliding_tile_attention/test/test_sta.py new file mode 100644 index 000000000..659dc195b --- /dev/null +++ b/csrc/sliding_tile_attention/test/test_sta.py @@ -0,0 +1,121 @@ +import torch +from flex_sta_ref import get_sliding_tile_attention_mask +from st_attn import sliding_tile_attention +from torch.nn.attention.flex_attention import flex_attention +# from flash_attn_interface import flash_attn_func +from tqdm import tqdm + +flex_attention = torch.compile(flex_attention, dynamic=False) + + +def flex_test(Q, K, V, kernel_size): + mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), + (30, 48, 80), 39, 'cuda') + output = flex_attention(Q, K, V, block_mask=mask) + + return output + + +def h100_fwd_kernel_test(Q, K, V, kernel_size): + o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 39) + return o + + +def generate_tensor(shape, mean, std, dtype, device): + tensor = torch.randn(shape, dtype=dtype, device=device) + + magnitude = torch.norm(tensor, dim=-1, keepdim=True) + scaled_tensor = tensor * ( + torch.randn(magnitude.shape, dtype=dtype, device=device) * std + + mean) / magnitude + + return scaled_tensor.contiguous() + + +def check_correctness(b, + h, + n, + d, + causal, + mean, + std, + num_iterations=50, + error_mode='all'): + results = { + 'TK vs FLEX': { + 'sum_diff': 0, + 'sum_abs': 0, + 'max_diff': 0 + }, + } + kernel_size_ls = [(1, 6, 10), (5, 1, 1), (3, 3, 3), (5, 5, 5)] + from tqdm import tqdm + for kernel_size in tqdm(kernel_size_ls): + for _ in range(num_iterations): + torch.manual_seed(0) + + Q = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, + 'cuda') + K = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, + 'cuda') + V = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, + 'cuda') + tk_o = h100_fwd_kernel_test(Q, K, V, kernel_size) + pt_o = flex_test(Q, K, V, kernel_size) + + diff = pt_o - tk_o + abs_diff = torch.abs(diff) + results['TK vs FLEX']['sum_diff'] += torch.sum(abs_diff).item() + results['TK vs FLEX']['max_diff'] = max( + results['TK vs FLEX']['max_diff'], + torch.max(abs_diff).item()) + + torch.cuda.empty_cache() + print("kernel_size", kernel_size) + print("max_diff", torch.max(abs_diff).item()) + print( + "avg_diff", + torch.sum(abs_diff).item() / + (b * h * n * d * (1 if error_mode == 'output' else + 3 if error_mode == 'backward' else 4))) + + total_elements = b * h * n * d * num_iterations * ( + 1 if error_mode == 'output' else + 3 if error_mode == 'backward' else 4) * len(kernel_size_ls) + for name, data in results.items(): + avg_diff = data['sum_diff'] / total_elements + max_diff = data['max_diff'] + results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff} + + return results + + +def generate_error_graphs(b, h, d, causal, mean, std, error_mode='all'): + seq_lengths = [115456] + + tk_avg_errors, tk_max_errors = [], [] + + for n in tqdm(seq_lengths, desc="Generating error data"): + results = check_correctness(b, + h, + n, + d, + causal, + mean, + std, + error_mode=error_mode) + + tk_avg_errors.append(results['TK vs FLEX']['avg_diff']) + tk_max_errors.append(results['TK vs FLEX']['max_diff']) + + +# Example usage +b, h, d = 1, 24, 128 +causal = False +mean = 1e-1 +std = 10 + +for mode in ['output']: + generate_error_graphs(b, h, d, causal, mean, std, error_mode=mode) + +print("Error graphs generated and saved for all modes.") diff --git a/csrc/sliding_tile_attention/tk b/csrc/sliding_tile_attention/tk new file mode 160000 index 000000000..1719fb726 --- /dev/null +++ b/csrc/sliding_tile_attention/tk @@ -0,0 +1 @@ +Subproject commit 1719fb72641b965d26155a0515d413b007f9dc72 diff --git a/demo/qualitative_results.md b/demo/qualitative_results.md new file mode 100644 index 000000000..bd0de273d --- /dev/null +++ b/demo/qualitative_results.md @@ -0,0 +1,15 @@ +Fast-Hunyuan comparison with original Hunyuan, achieving an 8X diffusion speed boost with the FastVideo framework. + +https://github.com/user-attachments/assets/064ac1d2-11ed-4a0c-955b-4d412a96ef30 + +Fast-Mochi comparison with original Mochi, achieving an 8X diffusion speed boost with the FastVideo framework. + +https://github.com/user-attachments/assets/5fbc4596-56d6-43aa-98e0-da472cf8e26c + +Comparison between OpenAI Sora, original Hunyuan and FastHunyuan + +https://github.com/user-attachments/assets/d323b712-3f68-42b2-952b-94f6a49c4836 + +Comparison between original FastHunyuan, LLM-INT8 quantized FastHunyuan and NF4 quantized FastHunyuan + +https://github.com/user-attachments/assets/cf89efb5-5f68-4949-a085-f41c1ef26c94 \ No newline at end of file diff --git a/fastvideo/utils/validation.py b/fastvideo/utils/validation.py index 42294e7e9..562518189 100644 --- a/fastvideo/utils/validation.py +++ b/fastvideo/utils/validation.py @@ -1,3 +1,4 @@ +# isort: skip_file import gc import os from typing import List, Optional, Union diff --git a/format.sh b/format.sh index 0f11734a5..520d74b29 100644 --- a/format.sh +++ b/format.sh @@ -124,7 +124,7 @@ spell_check() { } spell_check_all(){ - codespell --toml pyproject.toml "${CODESPELL_EXCLUDES[@]}" + codespell --toml pyproject.toml } # Spelling check of files that differ from main branch. diff --git a/pyproject.toml b/pyproject.toml index 96ef3039c..ca58216b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "fastvideo" -version = "1.2.0" +version = "0.0.1" description = "FastVideo" readme = "README.md" requires-python = ">=3.8" @@ -37,3 +37,6 @@ ignore_missing_imports = true disallow_untyped_calls = true check_untyped_defs = true no_implicit_optional = true + +[tool.codespell] +skip ="./data,./wandb,csrc/sliding_tile_attention/tk" \ No newline at end of file