stage

fix kernel issue (#185 )
Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>
2025-02-18 22:12:12 +00:00 · 2025-02-16 21:35:56 -08:00 · 2025-02-15 16:22:01 -08:00 · 2025-02-15 15:44:34 -08:00
26 changed files with 361666 additions and 90 deletions
@@ -27,4 +27,4 @@ jobs:
          pip install pytest
      - name: Run Pytest
        run: |
-          pytest
+          pytest --ignore csrc/sliding_tile_attention/test
@@ -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
@@ -0,0 +1,3 @@
+[submodule "csrc/sliding_tile_attention/tk"]
+	path = csrc/sliding_tile_attention/tk
+	url = https://github.com/HazyResearch/ThunderKittens.git
@@ -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
-

 <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> 
@@ -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
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model 
+# Inference with FA2 full attn
+sh scripts/inference/inference_hunyuan_STA.sh
+```
+
 ### 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.
@@ -1,8 +1,8 @@
-Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.
-A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature.
-A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.
-A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones.
-A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robots immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1
-fox in the forest close-up quickly turned its head to the left
-Man walking his dog in the woods on a hot sunny day
-A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic.
+Panda playing the guitar
+A princess is riding a horse across a river, realistic
+A corgi wearing sunglasses walks on the beach of a tropical island
+A princess is brushing her long golden hair in the garden.
+Several giant wooly mammoths approach treading through a snowy meadow, their long wooly fur lightly blows in the wind as they walk, snow covered trees
+A litter of golden retriever puppies playing in the snow. Their heads pop out of the snow, covered in.
+A llama sits in a cozy reading nook, surrounded by plush pillows and soft blankets. Warm, golden lighting from a floor lamp creates a welcoming atmosp
+A close up view of a glass sphere that has a zen garden within it. There is a small dwarf in the sphere who is raking the zen garden and creating patterns in the sand.
@@ -0,0 +1,41 @@
+
+
+# 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.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.
@@ -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'
@@ -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})
@@ -0,0 +1,24 @@
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+
+#include <vector>
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+
+#include <cuda_runtime.h>
+
+
+#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
+}
@@ -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]
@@ -0,0 +1,559 @@
+// # Define TORCH_COMPILE macro
+
+#include "kittens.cuh"
+#include <cooperative_groups.h>
+#include <iostream>
+#include <stdio.h>
+
+#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<int D> 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<int D> struct fwd_globals {
+    using q_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::qo_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+    using k_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::kv_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+    using v_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::kv_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+    using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<D>::qo_height, fwd_attend_ker_tile_dims<D>::tile_width>>;
+    using o_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::qo_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+
+    using q_gl = gl<bf16,  -1, -1, -1, -1, q_tile>;
+    using k_gl = gl<bf16,  -1, -1, -1, -1, k_tile>;
+    using v_gl = gl<bf16,  -1, -1, -1, -1, v_tile>;
+    using l_gl = gl<float, -1, -1, -1, -1, l_col_vec>;
+    using o_gl = gl<bf16,  -1, -1, -1, -1, o_tile>;
+
+    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<int D, bool is_causal, bool text_q, int DT, int DH, int DW, int CT, int CH, int CW>
+__global__  __launch_bounds__((NUM_WORKERS)*kittens::WARP_THREADS, 1)
+void fwd_attend_ker(const __grid_constant__ fwd_globals<D> 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<D>;
+
+    using q_tile    =         st_bf<K::qo_height, K::tile_width>;
+    using k_tile    =         st_bf<K::kv_height, K::tile_width>;
+    using v_tile    =         st_bf<K::kv_height, K::tile_width>;
+    using l_col_vec = col_vec<st_fl<K::qo_height, K::tile_width>>;
+    using o_tile    =         st_bf<K::qo_height, K::tile_width>;
+    
+    q_tile    (&q_smem)[CONSUMER_WARPGROUPS] = al.allocate<q_tile, CONSUMER_WARPGROUPS>();
+    k_tile    (&k_smem)[K::stages]           = al.allocate<k_tile, K::stages          >();
+    v_tile    (&v_smem)[K::stages]           = al.allocate<v_tile, K::stages          >();
+    l_col_vec (&l_smem)[CONSUMER_WARPGROUPS] = al.allocate<l_col_vec, CONSUMER_WARPGROUPS>();
+    auto      (*o_smem)                      = reinterpret_cast<o_tile(*)>(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> 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<k_tile> 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<k_tile> 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<bf16>)) - 1 + (CONSUMER_WARPGROUPS * (K::qo_height/kittens::TILE_ROW_DIM<bf16>)); 
+            kv_iters = ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM<bf16>)) == 0) ? (0) : ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM<bf16>)) - 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<k_tile> 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<k_tile> 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<k_tile> 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<rt_fl<16, K::kv_height>> 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<float>::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> 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<l_col_vec> 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 <ATen/cuda/CUDAContext.h>
+#include <iostream>
+
+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>();
+    c10::BFloat16* k_ptr = k.data_ptr<c10::BFloat16>();
+    c10::BFloat16* v_ptr = v.data_ptr<c10::BFloat16>();
+
+    bf16*  d_q = reinterpret_cast<bf16*>(q_ptr);
+    bf16*  d_k = reinterpret_cast<bf16*>(k_ptr);
+    bf16*  d_v = reinterpret_cast<bf16*>(v_ptr);
+    
+
+    
+    torch::Tensor l_vec = torch::empty({static_cast<const uint>(batch), 
+                                        static_cast<const uint>(qo_heads), 
+                                        static_cast<const uint>(seq_len), 
+                                        static_cast<const uint>(1)}, 
+                                        torch::TensorOptions().dtype(torch::kFloat).device(q.device()).memory_format(at::MemoryFormat::Contiguous));
+        
+
+    bf16*  o_ptr = reinterpret_cast<bf16*>(o.data_ptr<c10::BFloat16>());
+    bf16*  d_o   = reinterpret_cast<bf16*>(o_ptr);
+
+    float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
+    float* d_l   = reinterpret_cast<float*>(l_ptr);
+
+    cudaDeviceSynchronize();
+    auto stream = at::cuda::getCurrentCUDAStream().stream(); 
+
+
+    if (head_dim == 128) {
+        using q_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using k_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using v_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>>;
+        using o_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+
+        using q_global = gl<bf16,  -1, -1, -1, -1, q_tile>;
+        using k_global = gl<bf16,  -1, -1, -1, -1, k_tile>;
+        using v_global = gl<bf16,  -1, -1, -1, -1, v_tile>;
+        using l_global = gl<float, -1, -1, -1, -1, l_col_vec>;
+        using o_global = gl<bf16,  -1, -1, -1, -1, o_tile>;
+
+        using globals      = fwd_globals<128>;
+
+        q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+        k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,   static_cast<unsigned int>(seq_len)};
+        o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+
+        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len),  static_cast<int>(text_length), static_cast<int>(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<bf16>*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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<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, 2, 1, 2, 5, 6, 10>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, 2, 1, 2, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<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, 2, 3, 0, 5, 6, 10>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, false, 2, 3, 0, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(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><<<grid_text, (32*NUM_WORKERS), mem_size, stream>>>(g);
+        }
+
+
+        CHECK_CUDA_ERROR(cudaGetLastError());
+        cudaStreamSynchronize(stream);
+    }
+
+    return o;
+    cudaDeviceSynchronize();
+}
@@ -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
@@ -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.")
@@ -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
@@ -36,10 +36,8 @@ from diffusers.utils import (USE_PEFT_BACKEND, BaseOutput, deprecate, logging,
                             replace_example_docstring, scale_lora_layers)
 from diffusers.utils.torch_utils import randn_tensor
 from einops import rearrange
-
-from fastvideo.utils.communications import all_gather
-from fastvideo.utils.parallel_states import (get_sequence_parallel_state,
-                                             nccl_info)
+import imageio
+import torchvision

 from ...constants import PRECISION_TO_TYPE
 from ...modules import HYVideoDiffusionTransformer
@@ -606,6 +604,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
        enable_vae_sp: bool = False,
        n_tokens: Optional[int] = None,
        embedded_guidance_scale: Optional[float] = None,
+        mask_strategy: Optional[Dict[str, list]] = None,
        **kwargs,
    ):
        r"""
@@ -842,13 +841,6 @@ class HunyuanVideoPipeline(DiffusionPipeline):
            latents,
        )

-        world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
-        if get_sequence_parallel_state():
-            latents = rearrange(latents,
-                                "b t (n s) h w -> b t n s h w",
-                                n=world_size).contiguous()
-            latents = latents[:, :, rank, :, :, :]
-
        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_func_kwargs(
            self.scheduler.step,
@@ -869,10 +861,25 @@ class HunyuanVideoPipeline(DiffusionPipeline):
        num_warmup_steps = len(
            timesteps) - num_inference_steps * self.scheduler.order
        self._num_timesteps = len(timesteps)
+        
+        import os
+        if not os.path.exists("./data/step_by_step"):
+            os.makedirs("./data/step_by_step", exist_ok=True)
+            
+        # create step by step folder for this prompt
+        prompt_folder = f"./data/step_by_step/{prompt}"
+        if not os.path.exists(prompt_folder):
+            os.makedirs(prompt_folder, exist_ok=True)

        # if is_progress_bar:
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
+                step_folder = f"{prompt_folder}/{i}"
+                if not os.path.exists(step_folder):
+                    os.makedirs(step_folder, exist_ok=True)
+                mask_param = [
+                    mask_strategy, i
+                ]  # if mask_strategy is None, STA will not be used
                if self.interrupt:
                    continue

@@ -905,10 +912,11 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                    encoder_hidden_states = torch.cat(
                        [prompt_embeds_2, prompt_embeds], dim=1)
                    noise_pred = self.transformer(  # For an input image (129, 192, 336) (1, 256, 256)
-                        latent_model_input,  # [2, 16, 33, 24, 42]
+                        latent_model_input,
                        encoder_hidden_states,
-                        t_expand,  # [2]
-                        prompt_mask,  # [2, 256]fpdb
+                        t_expand,
+                        prompt_mask,
+                        mask_param=mask_param,
                        guidance=guidance_expand,
                        return_dict=False,
                    )[0]
@@ -956,9 +964,64 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                    if callback is not None and i % callback_steps == 0:
                        step_idx = i // getattr(self.scheduler, "order", 1)
                        callback(step_idx, t, latents)
+                import copy
+                temp_latent = copy.deepcopy(latents)
+                    
+                expand_temporal_dim = False
+                if len(temp_latent.shape) == 4:
+                    if isinstance(self.vae, AutoencoderKLCausal3D):
+                        temp_latent = temp_latent.unsqueeze(2)
+                        expand_temporal_dim = True
+                elif len(temp_latent.shape) == 5:
+                    pass
+                else:
+                    raise ValueError(
+                        f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {temp_latent.shape}."
+                    )
+
+                if (
+                    hasattr(self.vae.config, "shift_factor")
+                    and self.vae.config.shift_factor
+                ):
+                    temp_latent = (
+                        temp_latent / self.vae.config.scaling_factor
+                        + self.vae.config.shift_factor
+                    )
+                else:
+                    temp_latent = temp_latent / self.vae.config.scaling_factor
+
+                with torch.autocast(
+                    device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled
+                ):
+                    if enable_tiling:
+                        self.vae.enable_tiling()
+                        image = self.vae.decode(
+                            temp_latent, return_dict=False, generator=generator
+                        )[0]
+                    else:
+                        image = self.vae.decode(
+                            temp_latent, return_dict=False, generator=generator
+                        )[0]
+
+                if expand_temporal_dim or image.shape[2] == 1:
+                    image = image.squeeze(2)
+                
+                image = (image / 2 + 0.5).clamp(0, 1)
+                # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+                image = image.cpu().float()
+                
+                # save image
+                output_path = f"{step_folder}/output.mp4"
+                
+                
+                videos = rearrange(image, "b c t h w -> t b c h w")
+                video_frames = []
+                for x in videos:
+                    x = torchvision.utils.make_grid(x, nrow=6)
+                    x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                    video_frames.append((x * 255).numpy().astype(np.uint8))
+                imageio.mimsave(output_path, video_frames, fps=24)

-        if get_sequence_parallel_state():
-            latents = all_gather(latents, dim=2)

        if not output_type == "latent":
            expand_temporal_dim = False
@@ -369,6 +369,7 @@ class HunyuanVideoSampler(Inference):
        embedded_guidance_scale=None,
        batch_size=1,
        num_videos_per_prompt=1,
+        mask_strategy=None,
        **kwargs,
    ):
        """
@@ -524,6 +525,7 @@ class HunyuanVideoSampler(Inference):
            vae_ver=self.args.vae,
            enable_tiling=self.args.vae_tiling,
            enable_vae_sp=self.args.vae_sp,
+            mask_strategy=mask_strategy,
        )[0]
        out_dict["samples"] = samples
        out_dict["prompts"] = prompt
@@ -1,5 +1,7 @@
 import torch
 import torch.nn.functional as F
+from einops import rearrange
+from st_attn import sliding_tile_attention

 from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad
 from fastvideo.utils.communications import all_gather, all_to_all_4D
@@ -32,12 +34,54 @@ def attention(
    return out


-def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask):
-    # 1GPU torch.Size([1, 11264, 24, 128]) tensor([    0, 11275, 11520], device='cuda:0', dtype=torch.int32)
-    # 2GPU torch.Size([1, 5632, 24, 128]) tensor([   0, 5643, 5888], device='cuda:0', dtype=torch.int32)
+def tile(x, sp_size):
+    x = rearrange(x,
+                  "b (sp t h w) head d -> b (t sp h w) head d",
+                  sp=sp_size,
+                  t=30,
+                  h=48 // sp_size,
+                  w=80)
+    return rearrange(
+        x,
+        "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
+        n_t=5,
+        n_h=6,
+        n_w=10,
+        ts_t=6,
+        ts_h=8,
+        ts_w=8)
+
+
+def untile(x, sp_size):
+    x = rearrange(
+        x,
+        "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
+        n_t=5,
+        n_h=6,
+        n_w=10,
+        ts_t=6,
+        ts_h=8,
+        ts_w=8)
+    return rearrange(x,
+                     "b (t sp h w) head d -> b (sp t h w) head d",
+                     sp=sp_size,
+                     t=30,
+                     h=48 // sp_size,
+                     w=80)
+
+
+def parallel_attention(q,
+                       k,
+                       v,
+                       img_q_len,
+                       img_kv_len,
+                       text_mask,
+                       mask_param=None):
    query, encoder_query = q
    key, encoder_key = k
    value, encoder_value = v
+    text_length = text_mask.sum()
+
    if get_sequence_parallel_state():
        # batch_size, seq_len, attn_heads, head_dim
        query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
@@ -57,28 +101,50 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask):
    sequence_length = query.size(1)
    encoder_sequence_length = encoder_query.size(1)

-    # Hint: please check encoder_query.shape
-    query = torch.cat([query, encoder_query], dim=1)
-    key = torch.cat([key, encoder_key], dim=1)
-    value = torch.cat([value, encoder_value], dim=1)
-    # B, S, 3, H, D
-    qkv = torch.stack([query, key, value], dim=2)
+    if mask_param[0] is not None:
+        query = torch.cat([tile(query, nccl_info.sp_size), encoder_query],
+                          dim=1).transpose(1, 2)
+        key = torch.cat([tile(key, nccl_info.sp_size), encoder_key],
+                        dim=1).transpose(1, 2)
+        value = torch.cat([tile(value, nccl_info.sp_size), encoder_value],
+                          dim=1).transpose(1, 2)

-    attn_mask = F.pad(text_mask, (sequence_length, 0), value=True)
-    hidden_states = flash_attn_no_pad(qkv,
-                                      attn_mask,
-                                      causal=False,
-                                      dropout_p=0.0,
-                                      softmax_scale=None)
+        head_num = query.size(1)
+        mask_strategy, time_step, layer_idx = mask_param
+        windows = [
+            mask_strategy[f"{time_step}_{layer_idx}_{head_idx}"]
+            for head_idx in range(head_num)
+        ]
+
+        hidden_states = sliding_tile_attention(query, key, value, windows,
+                                               text_length).transpose(1, 2)
+    else:
+        query = torch.cat([query, encoder_query], dim=1)
+        key = torch.cat([key, encoder_key], dim=1)
+        value = torch.cat([value, encoder_value], dim=1)
+        # B, S, 3, H, D
+        qkv = torch.stack([query, key, value], dim=2)
+
+        attn_mask = F.pad(text_mask, (sequence_length, 0), value=True)
+        hidden_states = flash_attn_no_pad(qkv,
+                                          attn_mask,
+                                          causal=False,
+                                          dropout_p=0.0,
+                                          softmax_scale=None)

    hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
        (sequence_length, encoder_sequence_length), dim=1)
+
+    if mask_param[0] is not None:
+        hidden_states = untile(hidden_states, nccl_info.sp_size)
+
    if get_sequence_parallel_state():
        hidden_states = all_to_all_4D(hidden_states,
                                      scatter_dim=1,
                                      gather_dim=2)
        encoder_hidden_states = all_gather(encoder_hidden_states,
                                           dim=2).contiguous()
+
    hidden_states = hidden_states.to(query.dtype)
    encoder_hidden_states = encoder_hidden_states.to(query.dtype)

@@ -138,6 +138,7 @@ class MMDoubleStreamBlock(nn.Module):
        vec: torch.Tensor,
        freqs_cis: tuple = None,
        text_mask: torch.Tensor = None,
+        mask_param=None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        (
            img_mod1_shift,
@@ -214,6 +215,7 @@ class MMDoubleStreamBlock(nn.Module):
            img_q_len=img_q.shape[1],
            img_kv_len=img_k.shape[1],
            text_mask=text_mask,
+            mask_param=mask_param,
        )

        # attention computation end
@@ -241,7 +243,6 @@ class MMDoubleStreamBlock(nn.Module):
                         scale=txt_mod2_scale)),
            gate=txt_mod2_gate,
        )
-
        return img, txt


@@ -318,6 +319,7 @@ class MMSingleStreamBlock(nn.Module):
        txt_len: int,
        freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None,
        text_mask: torch.Tensor = None,
+        mask_param=None,
    ) -> torch.Tensor:
        mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
        x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale)
@@ -339,8 +341,10 @@ class MMSingleStreamBlock(nn.Module):
            return encoder_state.narrow(
                dim, nccl_info.rank_within_group * local_heads, local_heads)

-        freqs_cis = (shrink_head(freqs_cis[0],
-                                 dim=0), shrink_head(freqs_cis[1], dim=0))
+        freqs_cis = (
+            shrink_head(freqs_cis[0], dim=0),
+            shrink_head(freqs_cis[1], dim=0),
+        )

        img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :]
        img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :]
@@ -361,6 +365,7 @@ class MMSingleStreamBlock(nn.Module):
            img_q_len=img_q.shape[1],
            img_kv_len=img_k.shape[1],
            text_mask=text_mask,
+            mask_param=mask_param,
        )

        # attention computation end
@@ -592,6 +597,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
        encoder_hidden_states: torch.Tensor,
        timestep: torch.LongTensor,
        encoder_attention_mask: torch.Tensor,
+        mask_param=None,
        output_features=False,
        output_features_stride=8,
        attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -648,25 +654,30 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):

        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
        # --------------------- Pass through DiT blocks ------------------------
-        for _, block in enumerate(self.double_blocks):
-            double_block_args = [img, txt, vec, freqs_cis, text_mask]

+        for index, block in enumerate(self.double_blocks):
+            mask_param.append(index)
+            double_block_args = [
+                img, txt, vec, freqs_cis, text_mask, mask_param
+            ]
            img, txt = block(*double_block_args)
-
+            mask_param = mask_param[:-1]
        # Merge txt and img to pass through single stream blocks.
        x = torch.cat((img, txt), 1)
        if output_features:
            features_list = []
        if len(self.single_blocks) > 0:
-            for _, block in enumerate(self.single_blocks):
+            for index, block in enumerate(self.single_blocks):
+                mask_param.append(index + len(self.double_blocks))
                single_block_args = [
                    x,
                    vec,
                    txt_seq_len,
                    (freqs_cos, freqs_sin),
                    text_mask,
+                    mask_param,
                ]
-
+                mask_param = mask_param[:-1]
                x = block(*single_block_args)
                if output_features and _ % output_features_stride == 0:
                    features_list.append(x[:, :img_seq_len, ...])
@@ -0,0 +1,464 @@
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+import imageio
+import numpy as np
+import torch
+import torch.distributed as dist
+import torchvision
+from einops import rearrange
+
+from fastvideo.models.hunyuan.inference import HunyuanVideoSampler
+from fastvideo.models.hunyuan.modules.modulate_layers import modulate
+from fastvideo.utils.parallel_states import (
+    initialize_sequence_parallel_state, nccl_info)
+
+
+def teacache_forward(
+    self,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    timestep: torch.LongTensor,
+    encoder_attention_mask: torch.Tensor,
+    mask_param=None,
+    output_features=False,
+    output_features_stride=8,
+    attention_kwargs: Optional[Dict[str, Any]] = None,
+    return_dict: bool = False,
+    guidance=None,
+) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+    if guidance is None:
+        guidance = torch.tensor([6016.0],
+                                device=hidden_states.device,
+                                dtype=torch.bfloat16)
+
+    img = x = hidden_states
+    text_mask = encoder_attention_mask
+    t = timestep
+    txt = encoder_hidden_states[:, 1:]
+    text_states_2 = encoder_hidden_states[:, 0, :self.config.text_states_dim_2]
+    _, _, ot, oh, ow = x.shape  # codespell:ignore
+    tt, th, tw = (
+        ot // self.patch_size[0],  # codespell:ignore
+        oh // self.patch_size[1],  # codespell:ignore
+        ow // self.patch_size[2],  # codespell:ignore
+    )
+    original_th = nccl_info.sp_size * th
+    freqs_cos, freqs_sin = self.get_rotary_pos_embed((tt, original_th, tw))
+    # Prepare modulation vectors.
+    vec = self.time_in(t)
+
+    # text modulation
+    vec = vec + self.vector_in(text_states_2)
+
+    # guidance modulation
+    if self.guidance_embed:
+        if guidance is None:
+            raise ValueError(
+                "Didn't get guidance strength for guidance distilled model.")
+
+        # our timestep_embedding is merged into guidance_in(TimestepEmbedder)
+        vec = vec + self.guidance_in(guidance)
+
+    # Embed image and text.
+    img = self.img_in(img)
+    if self.text_projection == "linear":
+        txt = self.txt_in(txt)
+    elif self.text_projection == "single_refiner":
+        txt = self.txt_in(txt, t,
+                          text_mask if self.use_attention_mask else None)
+    else:
+        raise NotImplementedError(
+            f"Unsupported text_projection: {self.text_projection}")
+
+    txt_seq_len = txt.shape[1]
+    img_seq_len = img.shape[1]
+
+    freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
+
+    if self.enable_teacache:
+        inp = img.clone()
+        vec_ = vec.clone()
+        (
+            img_mod1_shift,
+            img_mod1_scale,
+            img_mod1_gate,
+            img_mod2_shift,
+            img_mod2_scale,
+            img_mod2_gate,
+        ) = self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1)
+        normed_inp = self.double_blocks[0].img_norm1(inp)
+        modulated_inp = modulate(normed_inp,
+                                 shift=img_mod1_shift,
+                                 scale=img_mod1_scale)
+        if self.cnt == 0 or self.cnt == self.num_steps - 1:
+            should_calc = True
+            self.accumulated_rel_l1_distance = 0
+        else:
+            coefficients = [
+                7.33226126e+02, -4.01131952e+02, 6.75869174e+01,
+                -3.14987800e+00, 9.61237896e-02
+            ]
+            rescale_func = np.poly1d(coefficients)
+            self.accumulated_rel_l1_distance += rescale_func(
+                ((modulated_inp - self.previous_modulated_input).abs().mean() /
+                 self.previous_modulated_input.abs().mean()).cpu().item())
+            if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+                should_calc = False
+            else:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = modulated_inp
+        self.cnt += 1
+        if self.cnt == self.num_steps:
+            self.cnt = 0
+    if self.enable_teacache:
+        if not should_calc:
+            img += self.previous_residual
+        else:
+            ori_img = img.clone()
+            # --------------------- Pass through DiT blocks ------------------------
+            for index, block in enumerate(self.double_blocks):
+                mask_param.append(index)
+                double_block_args = [
+                    img, txt, vec, freqs_cis, text_mask, mask_param
+                ]
+                img, txt = block(*double_block_args)
+                mask_param = mask_param[:-1]
+
+            # Merge txt and img to pass through single stream blocks.
+            x = torch.cat((img, txt), 1)
+            if output_features:
+                features_list = []
+            if len(self.single_blocks) > 0:
+                for index, block in enumerate(self.single_blocks):
+                    mask_param.append(index + len(self.double_blocks))
+                    single_block_args = [
+                        x,
+                        vec,
+                        txt_seq_len,
+                        (freqs_cos, freqs_sin),
+                        text_mask,
+                        mask_param,
+                    ]
+                    mask_param = mask_param[:-1]
+                    x = block(*single_block_args)
+                    if output_features and _ % output_features_stride == 0:
+                        features_list.append(x[:, :img_seq_len, ...])
+
+            img = x[:, :img_seq_len, ...]
+            self.previous_residual = img - ori_img
+    else:
+        # --------------------- Pass through DiT blocks ------------------------
+        for index, block in enumerate(self.double_blocks):
+            mask_param.append(index)
+            double_block_args = [
+                img, txt, vec, freqs_cis, text_mask, mask_param
+            ]
+            img, txt = block(*double_block_args)
+            mask_param = mask_param[:-1]
+        # Merge txt and img to pass through single stream blocks.
+        x = torch.cat((img, txt), 1)
+        if output_features:
+            features_list = []
+        if len(self.single_blocks) > 0:
+            for index, block in enumerate(self.single_blocks):
+                mask_param.append(index + len(self.double_blocks))
+                single_block_args = [
+                    x,
+                    vec,
+                    txt_seq_len,
+                    (freqs_cos, freqs_sin),
+                    text_mask,
+                    mask_param,
+                ]
+                mask_param = mask_param[:-1]
+                x = block(*single_block_args)
+                if output_features and _ % output_features_stride == 0:
+                    features_list.append(x[:, :img_seq_len, ...])
+
+        img = x[:, :img_seq_len, ...]
+
+    # ---------------------------- Final layer ------------------------------
+    img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+
+    img = self.unpatchify(img, tt, th, tw)
+    assert not return_dict, "return_dict is not supported."
+    if output_features:
+        features_list = torch.stack(features_list, dim=0)
+    else:
+        features_list = None
+    return (img, features_list)
+
+
+def initialize_distributed():
+    local_rank = int(os.getenv("LOCAL_RANK", 0))  # Fetch local rank
+    # No need to remap `local_rank` because `CUDA_VISIBLE_DEVICES` already restricts visibility
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(
+        backend="nccl", init_method="env://", world_size=1, rank=local_rank
+    )
+
+
+def main(args):
+    initialize_distributed()
+    local_rank = int(os.getenv('RANK', 0))
+    world_size = int(os.getenv('WORLD_SIZE', 1))
+    print(nccl_info.sp_size)
+
+    print(args)
+    models_root_path = Path(args.model_path)
+    if not models_root_path.exists():
+        raise ValueError(f"`models_root` not exists: {models_root_path}")
+
+    # Create save folder to save the samples
+    save_path = args.output_path
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+    # Load models
+    hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(
+        models_root_path, args=args)
+
+    # Get the updated args
+    args = hunyuan_video_sampler.args
+
+    # teacache
+    hunyuan_video_sampler.pipeline.transformer.__class__.enable_teacache = args.enable_teacache
+    hunyuan_video_sampler.pipeline.transformer.__class__.cnt = 0
+    hunyuan_video_sampler.pipeline.transformer.__class__.num_steps = args.num_inference_steps
+    hunyuan_video_sampler.pipeline.transformer.__class__.rel_l1_thresh = args.rel_l1_thresh  # 0.1 for 1.6x speedup, 0.15 for 2.1x speedup
+    hunyuan_video_sampler.pipeline.transformer.__class__.accumulated_rel_l1_distance = 0
+    hunyuan_video_sampler.pipeline.transformer.__class__.previous_modulated_input = None
+    hunyuan_video_sampler.pipeline.transformer.__class__.previous_residual = None
+    hunyuan_video_sampler.pipeline.transformer.__class__.forward = teacache_forward
+
+    with open(args.mask_strategy_file_path, 'r') as f:
+        mask_strategy = json.load(f)
+    with open(args.prompt) as f:
+        prompts = [line.strip() for line in f.readlines()]
+
+    for idx, prompt in enumerate(prompts):
+
+        if idx % world_size != local_rank:
+            continue
+        outputs = hunyuan_video_sampler.predict(
+            prompt=prompt,
+            height=args.height,
+            width=args.width,
+            video_length=args.num_frames,
+            seed=args.seed,
+            negative_prompt=args.neg_prompt,
+            infer_steps=args.num_inference_steps,
+            guidance_scale=args.guidance_scale,
+            num_videos_per_prompt=args.num_videos,
+            flow_shift=args.flow_shift,
+            batch_size=args.batch_size,
+            embedded_guidance_scale=args.embedded_cfg_scale,
+            mask_strategy=mask_strategy,
+        )
+        videos = rearrange(outputs["samples"], "b c t h w -> t b c h w")
+        outputs = []
+        for x in videos:
+            x = torchvision.utils.make_grid(x, nrow=6)
+            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+            outputs.append((x * 255).numpy().astype(np.uint8))
+        os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
+        imageio.mimsave(os.path.join(args.output_path, f"{prompt[:100]}.mp4"),
+                        outputs,
+                        fps=args.fps)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # Basic parameters
+    parser.add_argument("--prompt", type=str, help="prompt file for inference")
+    parser.add_argument("--num_frames", type=int, default=16)
+    parser.add_argument("--height", type=int, default=256)
+    parser.add_argument("--width", type=int, default=256)
+    parser.add_argument("--num_inference_steps", type=int, default=50)
+    parser.add_argument("--model_path", type=str, default="data/hunyuan")
+    parser.add_argument("--output_path", type=str, default="./outputs/video")
+    parser.add_argument("--fps", type=int, default=24)
+
+    # Additional parameters
+    parser.add_argument(
+        "--sliding_block_size",
+        type=str,
+        default="8,6,10",
+        help="Sliding block size for sliding block attention.",
+    )
+    parser.add_argument(
+        "--denoise-type",
+        type=str,
+        default="flow",
+        help="Denoise type for noised inputs.",
+    )
+    parser.add_argument("--seed",
+                        type=int,
+                        default=None,
+                        help="Seed for evaluation.")
+    parser.add_argument("--neg_prompt",
+                        type=str,
+                        default=None,
+                        help="Negative prompt for sampling.")
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=1.0,
+        help="Classifier free guidance scale.",
+    )
+    parser.add_argument(
+        "--embedded_cfg_scale",
+        type=float,
+        default=6.0,
+        help="Embedded classifier free guidance scale.",
+    )
+    parser.add_argument("--flow_shift",
+                        type=int,
+                        default=7,
+                        help="Flow shift parameter.")
+    parser.add_argument("--batch_size",
+                        type=int,
+                        default=1,
+                        help="Batch size for inference.")
+    parser.add_argument(
+        "--num_videos",
+        type=int,
+        default=1,
+        help="Number of videos to generate per prompt.",
+    )
+    parser.add_argument(
+        "--load-key",
+        type=str,
+        default="module",
+        help=
+        "Key to load the model states. 'module' for the main model, 'ema' for the EMA model.",
+    )
+    parser.add_argument(
+        "--use-cpu-offload",
+        action="store_true",
+        help="Use CPU offload for the model load.",
+    )
+    parser.add_argument(
+        "--dit-weight",
+        type=str,
+        default=
+        "data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
+    )
+    parser.add_argument(
+        "--reproduce",
+        action="store_true",
+        help=
+        "Enable reproducibility by setting random seeds and deterministic algorithms.",
+    )
+    parser.add_argument(
+        "--disable-autocast",
+        action="store_true",
+        help=
+        "Disable autocast for denoising loop and vae decoding in pipeline sampling.",
+    )
+
+    # Flow Matching
+    parser.add_argument(
+        "--flow-reverse",
+        action="store_true",
+        help="If reverse, learning/sampling from t=1 -> t=0.",
+    )
+    parser.add_argument("--flow-solver",
+                        type=str,
+                        default="euler",
+                        help="Solver for flow matching.")
+    parser.add_argument(
+        "--use-linear-quadratic-schedule",
+        action="store_true",
+        help=
+        "Use linear quadratic schedule for flow matching. Following MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)",
+    )
+    parser.add_argument(
+        "--linear-schedule-end",
+        type=int,
+        default=25,
+        help="End step for linear quadratic schedule for flow matching.",
+    )
+
+    # Model parameters
+    parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill")
+    parser.add_argument("--latent-channels", type=int, default=16)
+    parser.add_argument("--precision",
+                        type=str,
+                        default="bf16",
+                        choices=["fp32", "fp16", "bf16"])
+    parser.add_argument("--rope-theta",
+                        type=int,
+                        default=256,
+                        help="Theta used in RoPE.")
+
+    parser.add_argument("--vae", type=str, default="884-16c-hy")
+    parser.add_argument("--vae-precision",
+                        type=str,
+                        default="fp16",
+                        choices=["fp32", "fp16", "bf16"])
+    parser.add_argument("--vae-tiling", action="store_true", default=True)
+    parser.add_argument("--vae-sp", action="store_true", default=False)
+
+    parser.add_argument("--text-encoder", type=str, default="llm")
+    parser.add_argument(
+        "--text-encoder-precision",
+        type=str,
+        default="fp16",
+        choices=["fp32", "fp16", "bf16"],
+    )
+    parser.add_argument("--text-states-dim", type=int, default=4096)
+    parser.add_argument("--text-len", type=int, default=256)
+    parser.add_argument("--tokenizer", type=str, default="llm")
+    parser.add_argument("--prompt-template",
+                        type=str,
+                        default="dit-llm-encode")
+    parser.add_argument("--prompt-template-video",
+                        type=str,
+                        default="dit-llm-encode-video")
+    parser.add_argument("--hidden-state-skip-layer", type=int, default=2)
+    parser.add_argument("--apply-final-norm", action="store_true")
+
+    parser.add_argument("--text-encoder-2", type=str, default="clipL")
+    parser.add_argument(
+        "--text-encoder-precision-2",
+        type=str,
+        default="fp16",
+        choices=["fp32", "fp16", "bf16"],
+    )
+    parser.add_argument("--text-states-dim-2", type=int, default=768)
+    parser.add_argument("--tokenizer-2", type=str, default="clipL")
+    parser.add_argument("--text-len-2", type=int, default=77)
+    parser.add_argument("--skip_time_steps", type=int, default=10)
+    parser.add_argument(
+        "--mask_strategy_selected",
+        type=lambda x: [int(i) for i in x.strip('[]').split(',')
+                        ],  # Convert string to list of integers
+        default=[1, 2, 6],  # Now can be directly set as a list
+        help="order of candidates")
+    parser.add_argument(
+        "--rel_l1_thresh",
+        type=float,
+        default=0.15,
+        help="0.1 for 1.6x speedup, 0.15 for 2.1x speedup",
+    )
+    parser.add_argument(
+        "--enable_teacache",
+        action="store_true",
+        help="Use teacache for speeding up inference",
+    )
+    parser.add_argument("--mask_strategy_file_path",
+                        type=str,
+                        default="assets/mask_strategy.json")
+    args = parser.parse_args()
+    # process for vae sequence parallel
+    if args.vae_sp and not args.vae_tiling:
+        raise ValueError(
+            "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True."
+        )
+    main(args)
@@ -1,3 +1,4 @@
+# isort: skip_file
 import gc
 import os
 from typing import List, Optional, Union
@@ -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.
@@ -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"
@@ -0,0 +1,25 @@
+#!/bin/bash
+
+num_gpus=8
+mask_strategy_file_path=assets/mask_strategy.json
+export MODEL_BASE=data/hunyuan
+rel_l1_thresh=0.15
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \
+    fastvideo/sample/sample_t2v_hunyuan_STA.py \
+    --height 768 \
+    --width 1280 \
+    --num_frames 117 \
+    --num_inference_steps 50 \
+    --guidance_scale 1 \
+    --embedded_cfg_scale 6 \
+    --flow_shift 7 \
+    --flow-reverse \
+    --prompt ./assets/prompt.txt \
+    --seed 12345 \
+    --output_path outputs_video/hunyuan_STA/ \
+    --model_path $MODEL_BASE \
+    --mask_strategy_file_path $mask_strategy_file_path \
+    --rel_l1_thresh $rel_l1_thresh \
+    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+    --vae-sp \
+    --enable_teacache
Author	SHA1	Message	Date
rlsu9	cadaa0646a	stage	2025-02-18 22:12:12 +00:00
Yongqi ChenandPeiyuan Zhang	ac490399c6	fix kernel issue (#185 ) Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>	2025-02-16 21:35:56 -08:00
Yongqi ChenandPeiyuan Zhang	9ea39cee57	Add STA and teacache forward (#184 ) Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>	2025-02-15 16:22:01 -08:00
Zhang Peiyuanandrlsu9	52e6e612a2	add sliding tile attn (#182 ) Co-authored-by: rlsu9 <r3su@ucsd.edu>	2025-02-15 15:44:34 -08:00