add benchmark for ops (#176)

.github/workflows/test.yml

@@ -27,4 +27,4 @@ jobs:
           pip install pytest
       - name: Run Pytest
         run: |
-          pytest --ignore csrc/sliding_tile_attention/test
+          pytest

.gitignore

@@ -1,3 +1,4 @@
+ucf101_stride4x4x4
 __pycache__
 *.mp4
 .ipynb_checkpoints
@@ -7,8 +8,10 @@ results/
 build/
 fastvideo.egg-info/
 wandb/
+.idea
 *.ipynb
 *.jpg
+*.mp3
 *.safetensors
 *.mp4
 *.png
@@ -17,6 +20,28 @@ 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/
@@ -26,9 +51,3 @@ outputs_video
 sbatch.sh
 *.out
 env
-dist/
-*.o
-**/build/
-**.egg-info
-**.pyc
-**.egg

.gitmodules

@@ -1,3 +0,0 @@
-[submodule "csrc/sliding_tile_attention/tk"]
-	path = csrc/sliding_tile_attention/tk
-	url = https://github.com/HazyResearch/ThunderKittens.git

README.md

@@ -4,6 +4,8 @@
 
 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> 
@@ -12,7 +14,6 @@ FastVideo is a lightweight framework for accelerating large video diffusion mode
 
 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.
@@ -21,10 +22,21 @@ 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.
@@ -35,18 +47,9 @@ 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
@@ -173,7 +176,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 --all` before submitting a pull request.
+We welcome all contributions. Please run `bash format.sh` 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.

assets/prompt.txt

@@ -1,8 +1,8 @@
-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.
+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.

csrc/sliding_tile_attention/README.md

@@ -1,41 +0,0 @@
-
-
-# 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.

csrc/sliding_tile_attention/config.py

@@ -1,16 +0,0 @@
-### 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'

csrc/sliding_tile_attention/setup.py

@@ -1,64 +0,0 @@
-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})

csrc/sliding_tile_attention/st_attn.cpp

@@ -1,24 +0,0 @@
-#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
-}

csrc/sliding_tile_attention/st_attn/__init__.py

@@ -1,37 +0,0 @@
-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]

csrc/sliding_tile_attention/st_attn/st_attn_h100.cu

@@ -1,559 +0,0 @@
-// # 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();
-}

csrc/sliding_tile_attention/test/flex_sta_ref.py

@@ -1,80 +0,0 @@
-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

csrc/sliding_tile_attention/test/test_sta.py

@@ -1,121 +0,0 @@
-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.")

demo/qualitative_results.md

@@ -1,15 +0,0 @@
-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

fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py

@@ -36,8 +36,10 @@ 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
-import imageio
-import torchvision
+
+from fastvideo.utils.communications import all_gather
+from fastvideo.utils.parallel_states import (get_sequence_parallel_state,
+                                             nccl_info)
 
 from ...constants import PRECISION_TO_TYPE
 from ...modules import HYVideoDiffusionTransformer
@@ -604,7 +606,6 @@ 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"""
@@ -841,6 +842,13 @@ 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,
@@ -861,25 +869,10 @@ 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
 
@@ -912,11 +905,10 @@ 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,
+                        latent_model_input,  # [2, 16, 33, 24, 42]
                         encoder_hidden_states,
-                        t_expand,
-                        prompt_mask,
-                        mask_param=mask_param,
+                        t_expand,  # [2]
+                        prompt_mask,  # [2, 256]fpdb
                         guidance=guidance_expand,
                         return_dict=False,
                     )[0]
@@ -964,64 +956,9 @@ 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

fastvideo/models/hunyuan/inference.py

@@ -369,7 +369,6 @@ class HunyuanVideoSampler(Inference):
         embedded_guidance_scale=None,
         batch_size=1,
         num_videos_per_prompt=1,
-        mask_strategy=None,
         **kwargs,
     ):
         """
@@ -525,7 +524,6 @@ 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

fastvideo/models/hunyuan/modules/attenion.py

@@ -1,7 +1,5 @@
 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
@@ -34,54 +32,12 @@ def attention(
     return out
 
 
-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):
+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)
     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)
@@ -101,50 +57,28 @@ def parallel_attention(q,
     sequence_length = query.size(1)
     encoder_sequence_length = encoder_query.size(1)
 
-    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)
+    # 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)
 
-        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)
+    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)
 

fastvideo/models/hunyuan/modules/models.py

@@ -138,7 +138,6 @@ 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,
@@ -215,7 +214,6 @@ 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
@@ -243,6 +241,7 @@ class MMDoubleStreamBlock(nn.Module):
                          scale=txt_mod2_scale)),
             gate=txt_mod2_gate,
         )
+
         return img, txt
 
 
@@ -319,7 +318,6 @@ 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)
@@ -341,10 +339,8 @@ 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:, :, :]
@@ -365,7 +361,6 @@ 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
@@ -597,7 +592,6 @@ 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,
@@ -654,30 +648,25 @@ 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 index, block in enumerate(self.single_blocks):
-                mask_param.append(index + len(self.double_blocks))
+            for _, block in enumerate(self.single_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, ...])

fastvideo/ops/modulate/k_modulate.py

@@ -0,0 +1,102 @@
+import triton
+import triton.language as tl
+
+CONFIG_LIST = [
+    triton.Config({"BLOCK_M": 256, "BLOCK_N": 32}, num_stages=2, num_warps=4),
+    triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_stages=2, num_warps=4),
+    triton.Config({"BLOCK_M": 128, "BLOCK_N": 32}, num_stages=2, num_warps=4),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_stages=2, num_warps=4),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_stages=2, num_warps=4),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 32}, num_stages=2, num_warps=4),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 64}, num_stages=2, num_warps=4),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128}, num_stages=2, num_warps=4),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 256}, num_stages=2, num_warps=4),
+]
+
+
+@triton.autotune(
+    configs=CONFIG_LIST,
+    key=["M", "N"],
+)
+@triton.jit
+def _modulate_fwd(
+        x_ptr,  # *Pointer* to first input vector.
+        output_ptr,  # *Pointer* to output vector.
+        scale_ptr,
+        shift_ptr,
+        m_stride,
+        s_stride,
+        M,
+        N,
+        seq_len,
+        BLOCK_M: tl.constexpr,  # Number of elements each program should process.
+        BLOCK_N: tl.constexpr,
+        # NOTE: `constexpr` so it can be used as a shape value.
+):
+    row_id = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
+    rows = row_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    s_rows = (row_id // seq_len) * BLOCK_M
+    col_id = tl.program_id(axis=1)
+    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    x_ptrs = x_ptr + rows[:, None] * m_stride + cols[None, :]
+    scale_ptrs = scale_ptr + s_rows * s_stride + cols[None, :]
+    shift_ptrs = shift_ptr + s_rows * s_stride + cols[None, :]
+
+    col_mask = cols[None, :] < N
+    block_mask = (rows[:, None] < M) & col_mask
+    s_block_mask = col_mask
+    x = tl.load(x_ptrs, mask=block_mask, other=0.0)
+    scale = tl.load(scale_ptrs, mask=s_block_mask, other=0.0)
+    shift = tl.load(shift_ptrs, mask=s_block_mask, other=0.0)
+
+    output = x * (1 + scale) + shift
+    # Write x + y back to DRAM.
+    tl.store(output_ptr + rows[:, None] * m_stride + cols[None, :], output, mask=block_mask)
+
+
+@triton.autotune(
+    configs=CONFIG_LIST,
+    key=["M", "N"],
+)
+@triton.jit
+def _modulate_bwd(
+        dx_ptr,  # *Pointer* to first input vector.
+        x_ptr,
+        dy_ptr,  # *Pointer* to output vector.
+        scale_ptr,
+        dscale_ptr,
+        m_stride,
+        s_stride,
+        M,
+        N,
+        seq_len,
+        BLOCK_M: tl.constexpr,  # Number of elements each program should process.
+        BLOCK_N: tl.constexpr,
+        # NOTE: `constexpr` so it can be used as a shape value.
+):
+    row_id = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
+    rows = row_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    s_rows = (row_id // seq_len) * BLOCK_M
+    col_id = tl.program_id(axis=1)
+    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    x_ptrs = x_ptr + rows[:, None] * m_stride + cols[None, :]
+    dy_ptrs = dy_ptr + rows[:, None] * m_stride + cols[None, :]
+    dx_ptrs = dx_ptr + rows[:, None] * m_stride + cols[None, :]
+    dscale_ptrs = dscale_ptr + rows[:, None] * m_stride + cols[None, :]
+
+    scale_ptrs = scale_ptr + s_rows * s_stride + cols[None, :]
+
+    col_mask = cols[None, :] < N
+    block_mask = (rows[:, None] < M) & col_mask
+    s_block_mask = col_mask
+    x = tl.load(x_ptrs, mask=block_mask, other=0.0)
+    dy = tl.load(dy_ptrs, mask=block_mask, other=0.0)
+    scale = tl.load(scale_ptrs, mask=s_block_mask, other=0.0)
+
+    dx = dy * (1 + scale)
+    dscale = dy * x
+    # Write x + y back to DRAM.
+    tl.store(dx_ptrs, dx, mask=block_mask)
+    tl.store(dscale_ptrs, dscale, mask=block_mask)

fastvideo/ops/modulate/modulate.py

@@ -0,0 +1,63 @@
+import torch
+import triton
+
+from fastvideo.ops.modulate.k_modulate import _modulate_fwd, _modulate_bwd
+
+
+class _FusedModulate(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, scale, shift):
+        y = torch.empty_like(x)
+        batch, seq_len, dim = x.shape
+        M = batch * seq_len
+        N = dim
+        x = x.view(-1, dim).contiguous()
+        scale = scale.view(-1, dim).contiguous()
+        shift = shift.view(-1, dim).contiguous()
+
+        def grid(meta):
+            return (
+                triton.cdiv(batch * seq_len, meta["BLOCK_M"]),
+                triton.cdiv(dim, meta["BLOCK_N"]),
+            )
+
+        _modulate_fwd[grid](x, y, scale, shift, x.stride(0), scale.stride(0), M, N, seq_len)
+
+        ctx.save_for_backward(x, scale)
+        ctx.batch = batch
+        ctx.seq_len = seq_len
+        ctx.dim = dim
+        return y
+
+    @staticmethod
+    def backward(ctx, dy):  # pragma: no cover  # this is covered, but called directly from C++
+        x, scale = ctx.saved_tensors
+
+        batch, seq_len, dim = ctx.batch, ctx.seq_len, ctx.dim
+        M = batch * seq_len
+        N = dim
+
+        # allocate output
+        dy = dy.contiguous()
+        dx = torch.empty_like(dy)
+        dscale = torch.empty_like(dy)
+        dshift = torch.sum(dy, dim=1)
+
+        def grid(meta):
+            return (
+                triton.cdiv(batch * seq_len, meta["BLOCK_M"]),
+                triton.cdiv(dim, meta["BLOCK_N"]),
+            )
+
+        _modulate_bwd[grid](dx, x, dy, scale, dscale, x.stride(0), scale.stride(0), M, N, seq_len)
+
+        dscale = torch.sum(dscale, dim=1)
+        return dx, dscale, dshift
+
+
+def fused_modulate(
+        x: torch.Tensor,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
+) -> torch.Tensor:
+    return _FusedModulate.apply(x, scale, shift)

fastvideo/sample/sample_t2v_cogvideo.py

@@ -0,0 +1,306 @@
+"""
+This script demonstrates how to generate a video using the CogVideoX model with the Hugging Face `diffusers` pipeline.
+The script supports different types of video generation, including text-to-video (t2v), image-to-video (i2v),
+and video-to-video (v2v), depending on the input data and different weight.
+
+- text-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
+- video-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
+- image-to-video: THUDM/CogVideoX-5b-I2V or THUDM/CogVideoX1.5-5b-I2V
+
+Running the Script:
+To run the script, use the following command with appropriate arguments:
+
+```bash
+$ python cli_demo.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX1.5-5b --generate_type "t2v"
+```
+
+Additional options are available to specify the model path, guidance scale, number of inference steps, video generation type, and output paths.
+"""
+
+import argparse
+import logging
+import os
+from typing import Literal, Optional
+
+import torch
+from diffusers import (
+    CogVideoXDPMScheduler,
+    CogVideoXImageToVideoPipeline,
+    CogVideoXPipeline,
+    CogVideoXVideoToVideoPipeline,
+)
+from diffusers.utils import export_to_video, load_image, load_video
+
+
+logging.basicConfig(level=logging.INFO)
+
+# Recommended resolution for each model (width, height)
+RESOLUTION_MAP = {
+    # cogvideox1.5-*
+    "cogvideox1.5-5b-i2v": (1360, 768),
+    "cogvideox1.5-5b": (1360, 768),
+    # cogvideox-*
+    "cogvideox-5b-i2v": (720, 480),
+    "cogvideox-5b": (720, 480),
+    "cogvideox-2b": (720, 480),
+}
+
+
+def generate_video(
+    prompt: str,
+    model_path: str,
+    lora_path: str = None,
+    lora_rank: int = 128,
+    num_frames: int = 81,
+    width: Optional[int] = None,
+    height: Optional[int] = None,
+    output_path: str = "./output.mp4",
+    image_or_video_path: str = "",
+    num_inference_steps: int = 50,
+    guidance_scale: float = 6.0,
+    num_videos_per_prompt: int = 1,
+    dtype: torch.dtype = torch.bfloat16,
+    generate_type: str = Literal[
+        "t2v", "i2v", "v2v"
+    ],  # i2v: image to video, v2v: video to video
+    seed: int = 42,
+    fps: int = 16,
+):
+    """
+    Generates a video based on the given prompt and saves it to the specified path.
+
+    Parameters:
+    - prompt (str): The description of the video to be generated.
+    - model_path (str): The path of the pre-trained model to be used.
+    - lora_path (str): The path of the LoRA weights to be used.
+    - lora_rank (int): The rank of the LoRA weights.
+    - output_path (str): The path where the generated video will be saved.
+    - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality.
+    - num_frames (int): Number of frames to generate. CogVideoX1.0 generates 49 frames for 6 seconds at 8 fps, while CogVideoX1.5 produces either 81 or 161 frames, corresponding to 5 seconds or 10 seconds at 16 fps.
+    - width (int): The width of the generated video, applicable only for CogVideoX1.5-5B-I2V
+    - height (int): The height of the generated video, applicable only for CogVideoX1.5-5B-I2V
+    - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt.
+    - num_videos_per_prompt (int): Number of videos to generate per prompt.
+    - dtype (torch.dtype): The data type for computation (default is torch.bfloat16).
+    - generate_type (str): The type of video generation (e.g., 't2v', 'i2v', 'v2v').·
+    - seed (int): The seed for reproducibility.
+    - fps (int): The frames per second for the generated video.
+    """
+
+    # 1.  Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16).
+    # add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload()
+    # function to use Multi GPUs.
+
+    image = None
+    video = None
+
+    model_name = model_path.split("/")[-1].lower()
+    desired_resolution = RESOLUTION_MAP[model_name]
+    if width is None or height is None:
+        width, height = desired_resolution
+        logging.info(
+            f"\033[1mUsing default resolution {desired_resolution} for {model_name}\033[0m"
+        )
+    elif (width, height) != desired_resolution:
+        if generate_type == "i2v":
+            # For i2v models, use user-defined width and height
+            logging.warning(
+                f"\033[1;31mThe width({width}) and height({height}) are not recommended for {model_name}. The best resolution is {desired_resolution}.\033[0m"
+            )
+        else:
+            # Otherwise, use the recommended width and height
+            logging.warning(
+                f"\033[1;31m{model_name} is not supported for custom resolution. Setting back to default resolution {desired_resolution}.\033[0m"
+            )
+            width, height = desired_resolution
+
+    if generate_type == "i2v":
+        pipe = CogVideoXImageToVideoPipeline.from_pretrained(
+            model_path, torch_dtype=dtype
+        )
+        image = load_image(image=image_or_video_path)
+    elif generate_type == "t2v":
+        pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype)
+    else:
+        pipe = CogVideoXVideoToVideoPipeline.from_pretrained(
+            model_path, torch_dtype=dtype
+        )
+        video = load_video(image_or_video_path)
+
+    # If you're using with lora, add this code
+    if lora_path:
+        pipe.load_lora_weights(
+            lora_path,
+            weight_name="pytorch_lora_weights.safetensors",
+            adapter_name="test_1",
+        )
+        pipe.fuse_lora(lora_scale=1 / lora_rank)
+
+    # 2. Set Scheduler.
+    # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`.
+    # We recommend using `CogVideoXDDIMScheduler` for CogVideoX-2B.
+    # using `CogVideoXDPMScheduler` for CogVideoX-5B / CogVideoX-5B-I2V.
+
+    # pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+    pipe.scheduler = CogVideoXDPMScheduler.from_config(
+        pipe.scheduler.config, timestep_spacing="trailing"
+    )
+
+    # 3. Enable CPU offload for the model.
+    # turn off if you have multiple GPUs or enough GPU memory(such as H100) and it will cost less time in inference
+    # and enable to("cuda")
+
+    # pipe.to("cuda")
+    pipe.enable_sequential_cpu_offload()
+    pipe.vae.enable_slicing()
+    pipe.vae.enable_tiling()
+
+    # 4. Generate the video frames based on the prompt.
+    # `num_frames` is the Number of frames to generate.
+    if generate_type == "i2v":
+        video_generate = pipe(
+            height=height,
+            width=width,
+            prompt=prompt,
+            image=image,
+            # The path of the image, the resolution of video will be the same as the image for CogVideoX1.5-5B-I2V, otherwise it will be 720 * 480
+            num_videos_per_prompt=num_videos_per_prompt,  # Number of videos to generate per prompt
+            num_inference_steps=num_inference_steps,  # Number of inference steps
+            num_frames=num_frames,  # Number of frames to generate
+            use_dynamic_cfg=True,  # This id used for DPM scheduler, for DDIM scheduler, it should be False
+            guidance_scale=guidance_scale,
+            generator=torch.Generator().manual_seed(
+                seed
+            ),  # Set the seed for reproducibility
+        ).frames[0]
+    elif generate_type == "t2v":
+        video_generate = pipe(
+            height=height,
+            width=width,
+            prompt=prompt,
+            num_videos_per_prompt=num_videos_per_prompt,
+            num_inference_steps=num_inference_steps,
+            num_frames=num_frames,
+            use_dynamic_cfg=True,
+            guidance_scale=guidance_scale,
+            generator=torch.Generator().manual_seed(seed),
+        ).frames[0]
+    else:
+        video_generate = pipe(
+            height=height,
+            width=width,
+            prompt=prompt,
+            video=video,  # The path of the video to be used as the background of the video
+            num_videos_per_prompt=num_videos_per_prompt,
+            num_inference_steps=num_inference_steps,
+            num_frames=num_frames,
+            use_dynamic_cfg=True,
+            guidance_scale=guidance_scale,
+            generator=torch.Generator().manual_seed(
+                seed
+            ),  # Set the seed for reproducibility
+        ).frames[0]
+    export_to_video(video_generate, output_path, fps=fps)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Generate a video from a text prompt using CogVideoX"
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="The description of the video to be generated",
+    )
+    parser.add_argument(
+        "--image_or_video_path",
+        type=str,
+        default=None,
+        help="The path of the image to be used as the background of the video",
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default="THUDM/CogVideoX1.5-5B",
+        help="Path of the pre-trained model use",
+    )
+    parser.add_argument(
+        "--lora_path",
+        type=str,
+        default=None,
+        help="The path of the LoRA weights to be used",
+    )
+    parser.add_argument(
+        "--lora_rank", type=int, default=128, help="The rank of the LoRA weights"
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default="./output.mp4",
+        help="The path save generated video",
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=6.0,
+        help="The scale for classifier-free guidance",
+    )
+    parser.add_argument(
+        "--num_inference_steps", type=int, default=50, help="Inference steps"
+    )
+    parser.add_argument(
+        "--num_frames",
+        type=int,
+        default=81,
+        help="Number of steps for the inference process",
+    )
+    parser.add_argument(
+        "--width", type=int, default=None, help="The width of the generated video"
+    )
+    parser.add_argument(
+        "--height", type=int, default=None, help="The height of the generated video"
+    )
+    parser.add_argument(
+        "--fps",
+        type=int,
+        default=16,
+        help="The frames per second for the generated video",
+    )
+    parser.add_argument(
+        "--num_videos_per_prompt",
+        type=int,
+        default=1,
+        help="Number of videos to generate per prompt",
+    )
+    parser.add_argument(
+        "--generate_type", type=str, default="t2v", help="The type of video generation"
+    )
+    parser.add_argument(
+        "--dtype", type=str, default="bfloat16", help="The data type for computation"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, help="The seed for reproducibility"
+    )
+
+    args = parser.parse_args()
+    dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16
+    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
+    generate_video(
+        prompt=args.prompt,
+        model_path=args.model_path,
+        lora_path=args.lora_path,
+        lora_rank=args.lora_rank,
+        output_path=args.output_path,
+        num_frames=args.num_frames,
+        width=args.width,
+        height=args.height,
+        image_or_video_path=args.image_or_video_path,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        num_videos_per_prompt=args.num_videos_per_prompt,
+        dtype=dtype,
+        generate_type=args.generate_type,
+        seed=args.seed,
+        fps=args.fps,
+    )

fastvideo/sample/sample_t2v_hunyuan_STA.py

@@ -1,464 +0,0 @@
-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)

fastvideo/utils/validation.py

@@ -1,4 +1,3 @@
-# isort: skip_file
 import gc
 import os
 from typing import List, Optional, Union

format.sh

@@ -124,7 +124,7 @@ spell_check() {
 }
 
 spell_check_all(){
-  codespell --toml pyproject.toml
+  codespell --toml pyproject.toml "${CODESPELL_EXCLUDES[@]}"
 }
 
 # Spelling check of files that differ from main branch.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fastvideo"
-version = "0.0.1"
+version = "1.2.0"
 description = "FastVideo"
 readme = "README.md"
 requires-python = ">=3.8"
@@ -37,6 +37,3 @@ 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"

scripts/inference/inference_hunyuan_STA.sh

@@ -1,25 +0,0 @@
-#!/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 

tests/test_kernels/benchmark.py

@@ -0,0 +1,297 @@
+# Copyright (c) 2023, Tri Dao.
+"""Useful functions for writing test code."""
+
+import torch
+import torch.utils.benchmark as benchmark
+
+
+def get_abs_err(x, y):
+    return (x - y).flatten().abs().max().item()
+
+
+def get_err_ratio(x, y):
+    err = (x - y).flatten().square().mean().sqrt().item()
+    base = x.flatten().square().mean().sqrt().item()
+    return err / base
+
+
+def assert_close(prefix, ref, tri, ratio):
+    msg = f"{prefix} diff: {get_abs_err(ref, tri):.6f} ratio: {get_err_ratio(ref, tri):.6f}"
+    print(msg)
+    assert get_err_ratio(ref, tri) < ratio, msg
+
+
+def benchmark_forward(
+    fn,
+    *inputs,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Forward pass")
+
+    def amp_wrapper(*inputs, **kwinputs):
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            fn(*inputs, **kwinputs)
+
+    t = benchmark.Timer(
+        stmt="fn_amp(*inputs, **kwinputs)",
+        globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_backward(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the backward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Backward pass")
+    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+        y = fn(*inputs, **kwinputs)
+        if type(y) is tuple:
+            y = y[0]
+    if grad is None:
+        grad = torch.randn_like(y)
+    else:
+        if grad.shape != y.shape:
+            raise RuntimeError("Grad shape does not match output shape")
+
+    def f(*inputs, y, grad):
+        # Set .grad to None to avoid extra operation of gradient accumulation
+        for x in inputs:
+            if isinstance(x, torch.Tensor):
+                x.grad = None
+        y.backward(grad, retain_graph=True)
+
+    t = benchmark.Timer(
+        stmt="f(*inputs, y=y, grad=grad)",
+        globals={"f": f, "inputs": inputs, "y": y, "grad": grad},
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_combined(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    if verbose:
+        print(desc, "- Forward + Backward pass")
+    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+        y = fn(*inputs, **kwinputs)
+        if type(y) is tuple:
+            y = y[0]
+    if grad is None:
+        grad = torch.randn_like(y)
+    else:
+        if grad.shape != y.shape:
+            raise RuntimeError("Grad shape does not match output shape")
+
+    def f(grad, *inputs, **kwinputs):
+        for x in inputs:
+            if isinstance(x, torch.Tensor):
+                x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            y = fn(*inputs, **kwinputs)
+            if type(y) is tuple:
+                y = y[0]
+        y.backward(grad, retain_graph=True)
+
+    t = benchmark.Timer(
+        stmt="f(grad, *inputs, **kwinputs)",
+        globals={
+            "f": f,
+            "fn": fn,
+            "inputs": inputs,
+            "grad": grad,
+            "kwinputs": kwinputs,
+        },
+        num_threads=torch.get_num_threads(),
+    )
+    m = t.timeit(repeats)
+    if verbose:
+        print(m)
+    return t, m
+
+
+def benchmark_fwd_bwd(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    return (
+        benchmark_forward(
+            fn,
+            *inputs,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_backward(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+    )
+
+
+def benchmark_all(
+    fn,
+    *inputs,
+    grad=None,
+    repeats=10,
+    desc="",
+    verbose=True,
+    amp=False,
+    amp_dtype=torch.float16,
+    **kwinputs,
+):
+    """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function."""
+    return (
+        benchmark_forward(
+            fn,
+            *inputs,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_backward(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+        benchmark_combined(
+            fn,
+            *inputs,
+            grad=grad,
+            repeats=repeats,
+            desc=desc,
+            verbose=verbose,
+            amp=amp,
+            amp_dtype=amp_dtype,
+            **kwinputs,
+        ),
+    )
+
+
+def pytorch_profiler(
+    fn,
+    *inputs,
+    trace_filename=None,
+    backward=False,
+    amp=False,
+    amp_dtype=torch.float16,
+    cpu=False,
+    verbose=True,
+    **kwinputs,
+):
+    """Wrap benchmark functions in Pytorch profiler to see CUDA information."""
+    if backward:
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+            g = torch.randn_like(out)
+    for _ in range(30):  # Warm up
+        if backward:
+            for x in inputs:
+                if isinstance(x, torch.Tensor):
+                    x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+        # Backward should be done outside autocast
+        if backward:
+            out.backward(g, retain_graph=True)
+    activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [
+        torch.profiler.ProfilerActivity.CUDA
+    ]
+    with torch.profiler.profile(
+        activities=activities,
+        record_shapes=True,
+        # profile_memory=True,
+        with_stack=True,
+    ) as prof:
+        if backward:
+            for x in inputs:
+                if isinstance(x, torch.Tensor):
+                    x.grad = None
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp):
+            out = fn(*inputs, **kwinputs)
+            if type(out) is tuple:
+                out = out[0]
+        if backward:
+            out.backward(g, retain_graph=True)
+    if verbose:
+        # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50))
+        print(prof.key_averages().table(row_limit=50))
+    if trace_filename is not None:
+        prof.export_chrome_trace(trace_filename)
+
+
+def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs):
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    fn(*inputs, **kwinputs)
+    torch.cuda.synchronize()
+    mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000)
+    if verbose:
+        print(f"{desc} max memory: {mem}GB")
+    torch.cuda.empty_cache()
+    return mem

tests/test_kernels/benchmark_modulate.py

@@ -0,0 +1,83 @@
+import torch
+from torch import nn
+
+from benchmark import (
+    benchmark_combined,
+    benchmark_forward,
+    benchmark_backward,
+    assert_close,
+)
+from fastvideo.ops.modulate.modulate import fused_modulate
+
+
+def torch_modulate(x, scale, shift):
+    return x * (scale.unsqueeze(1) + 1) + shift.unsqueeze(1)
+
+
+# from flash_attn import flash_attn_func
+
+
+def time_fwd(func, *args, **kwargs):
+    time_fb = benchmark_forward(func, *args, **kwargs)
+    return time_fb[1].mean
+
+
+def time_fwd_bwd(func, *args, **kwargs):
+    time_fb = benchmark_combined(func, *args, **kwargs)
+    return time_fb[1].mean
+
+
+def time_bwd(func, *args, **kwargs):
+    time_fb = benchmark_backward(func, *args, **kwargs)
+    return time_fb[1].mean
+
+
+device = "cuda"
+dtype = torch.bfloat16
+
+
+batch_sizes = [1, 8, 32]
+seq_lengths = [128, 512, 1024]
+hidden_dims = [768, 1024, 2048]
+
+# methods = (["torch", "triton", "thunderkitten"])
+methods = ["torch", "triton"]
+time_f = {}
+time_b = {}
+time_f_b = {}
+speed_f = {}
+speed_b = {}
+speed_f_b = {}
+for B in batch_sizes:
+    for T in seq_lengths:
+        for D in hidden_dims:
+            config = (B, T, D)
+
+            norm_func = nn.LayerNorm(D, elementwise_affine=False, eps=1e-6)
+            x = torch.randn(B, T, D, device="cuda", requires_grad=True, dtype=dtype)
+            x = norm_func(x.to(torch.float32)).to(dtype)
+            shift = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype)
+            scale = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype)
+            # test torch
+            o_ref = torch_modulate(x, scale, shift)
+            o_ref.sum().backward(retain_graph=True)
+            f_b = time_fwd_bwd(torch_modulate, x, scale, shift, verbose=False)
+            time_f_b[config, "torch"] = f_b
+            # test triton
+            o2 = fused_modulate(x, scale, shift)
+            o2.sum().backward(retain_graph=True)
+            f_b = time_fwd_bwd(fused_modulate, x, scale, shift, verbose=False)
+            time_f_b[config, "triton"] = f_b
+            # test if the results are close
+            assert_close("  o", o_ref, o2, 0.005)
+            # time_f_b[config, "thunderkitten"] = f_b
+
+            print(f"### batch size={B}, seq length={T}, B={B}, hidden dim={D} ###")
+            for method in methods:
+                # time_f_b[config, method] = time_f[config, method] + time_b[config, method]
+                print(
+                    f"{method:>50} fwd + bwd:\t {time_f_b[config, method]*1000:>6.4f} ms "
+                )
+
+# with open('flash2_attn_time.plk', 'wb') as fp:
+#     pickle.dump((speed_f, speed_b, speed_f_b), fp, protocol=pickle.HIGHEST_PROTOCOL)