diff --git a/csrc/attn/sliding_tile_attn/README.md b/csrc/attn/sliding_tile_attn/README.md index 4c6bfcb47..b543d35b7 100644 --- a/csrc/attn/sliding_tile_attn/README.md +++ b/csrc/attn/sliding_tile_attn/README.md @@ -2,12 +2,12 @@ # Attention Kernel Used in FastVideo ## Sliding Tile Attention (STA) -We only support H100 for STA. +We support H100 (via TK) and any other GPU (via triton) for STA. ### Installation ```bash pip install st_attn -``` +``` Install from source: @@ -16,6 +16,14 @@ git submodule update --init --recursive python setup.py install ``` +If you want to skip the compilation of the TK kernel and only use the Triton version, try below: + +```bash +SKIP_SM90_EXT=1 python setup.py install +or +SKIP_SM90_EXT=1 pip install --no-build-isolation . +``` + If you encounter error during installation, try below: Install C++20 for ThunderKittens: ```bash @@ -30,7 +38,7 @@ sudo apt install clang-11 (If you use CUDA12.8) ```bash export CUDA_HOME=/usr/local/cuda-12.8 -export PATH=${CUDA_HOME}/bin:${PATH} +export PATH=${CUDA_HOME}/bin:${PATH} export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH ``` @@ -43,7 +51,7 @@ bash scripts/inference/v1_inference_wan_STA.sh If you want to use sliding tile attention in your custom model: ```python from st_attn import sliding_tile_attention -# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. +# 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. @@ -58,7 +66,6 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False) ### Test ```bash python ../tests/test_sta.py # test STA -python ../tests/test_vsa.py # test VSA ``` ### Benchmark ```bash @@ -67,7 +74,7 @@ python ../benchmarks/bench_sta.py ### How Does STA Work? -We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. +We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747 @@ -82,7 +89,7 @@ Here is a diagram of how the window is configured and passed through the FastVid ## Why is STA Fast? -2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. +2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. STA removes mixed blocks. diff --git a/csrc/attn/sliding_tile_attn/setup.py b/csrc/attn/sliding_tile_attn/setup.py index a33352147..fc55f38d7 100644 --- a/csrc/attn/sliding_tile_attn/setup.py +++ b/csrc/attn/sliding_tile_attn/setup.py @@ -51,21 +51,28 @@ for k in kernels: source_files.append(sources[k]['source_files'][target]) cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}') +ext_modules = [] + +if os.environ.get("SKIP_SM90_EXT", "0") != "1": + ext_modules.append( + CUDAExtension('st_attn_cuda', + sources=source_files, + extra_compile_args={ + 'cxx': cpp_flags, + 'nvcc': cuda_flags + }, + libraries=['cuda']) + ) +else: + print("ENV SKIP_SM90_EXT=1, skip st_attn_cuda compile") + setup(name=PACKAGE_NAME, version=VERSION, author=AUTHOR, description=DESCRIPTION, url=URL, packages=find_packages(), - ext_modules=[ - CUDAExtension('st_attn_cuda', - sources=source_files, - extra_compile_args={ - 'cxx': cpp_flags, - 'nvcc': cuda_flags - }, - libraries=['cuda']) - ], + ext_modules=ext_modules, cmdclass={'build_ext': BuildExtension}, classifiers=[ "Programming Language :: Python :: 3", diff --git a/csrc/attn/sliding_tile_attn/st_attn/__init__.py b/csrc/attn/sliding_tile_attn/st_attn/__init__.py index 9be8627a2..89fb94bf8 100644 --- a/csrc/attn/sliding_tile_attn/st_attn/__init__.py +++ b/csrc/attn/sliding_tile_attn/st_attn/__init__.py @@ -7,12 +7,17 @@ try: except ImportError: sta_fwd = None -def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, dit_seq_shape='30x48x80'): +try: + from st_attn.st_attn_triton import sliding_tile_attention_triton +except ImportError: + sliding_tile_attention_triton = None + +def sliding_tile_attention_SM90(q_all, k_all, v_all, window_size, text_length, has_text=True, dit_seq_shape='30x48x80'): seq_length = q_all.shape[2] dit_seq_shape_mapping = { '30x48x80':1, '36x48x48':2, - '18x48x80':3, + '18x48x80':3, } if has_text: assert q_all.shape[ @@ -46,4 +51,13 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text, kernel_aspect_ratio_flag) if has_text: _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True, kernel_aspect_ratio_flag) - return hidden_states[:, :, :seq_length] \ No newline at end of file + return hidden_states[:, :, :seq_length] + +def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, dit_seq_shape='30x48x80'): + major, minor = torch.cuda.get_device_capability(q_all.device) + if major == 9 and minor == 0 and sta_fwd is not None: + return sliding_tile_attention_SM90(q_all, k_all, v_all, window_size, text_length, has_text, dit_seq_shape) + elif sliding_tile_attention_triton is not None: + return sliding_tile_attention_triton(q_all, k_all, v_all, window_size, text_length, has_text, dit_seq_shape) + else: + raise ImportError("No suitable sliding tile attention implementation found.") diff --git a/csrc/attn/sliding_tile_attn/st_attn/st_attn_triton.py b/csrc/attn/sliding_tile_attn/st_attn/st_attn_triton.py new file mode 100644 index 000000000..ab4357a0f --- /dev/null +++ b/csrc/attn/sliding_tile_attn/st_attn/st_attn_triton.py @@ -0,0 +1,327 @@ +import math +import torch +import triton +import triton.language as tl + +def is_cuda(): + return triton.runtime.driver.active.get_current_target().backend == "cuda" + + +def is_hip(): + target = triton.runtime.driver.active.get_current_target() + return target.backend == 'hip' + +def get_common_autotune_config(): + configs = [ + triton.Config({'BLOCK_Q': BLOCK_Q, 'BLOCK_KV': BLOCK_KV}, num_stages=s, num_warps=w) \ + for BLOCK_Q in [32, 64, 128]\ + for BLOCK_KV in [32, 64, 128]\ + for s in [1, 2, 3, 4]\ + for w in [4, 8]\ + ] + return configs + + +def get_cuda_autotune_config(): + # cuda and hip can use differnt autotune configs + return get_common_autotune_config() + + +def get_hip_autotune_config(): + # cuda and hip can use differnt autotune configs + return get_common_autotune_config() + + +def get_autotune_config(): + if is_cuda(): + return get_cuda_autotune_config() + else: + return get_hip_autotune_config() + + +@triton.jit +def clamp_int(value, min_val, max_val): + ret = tl.where(value > max_val, max_val, value) + ret = tl.where(ret < min_val, min_val, ret) + return ret + + +@triton.jit +def _attn_fwd_loop( + q, k, v, kv_mask, m, l, acc, sm_scale, + MASK_KV: tl.constexpr, +): + scores = tl.dot(q, k.T) #[BLOCK_Q, BLOCK_KV] + scores = scores * sm_scale + if MASK_KV: + scores = tl.where(kv_mask[None, :], scores, -float('inf')) + + current_m = tl.max(scores, axis=1) + new_m = tl.maximum(m, current_m) + exp_scores = tl.math.exp2(scores - new_m[:, None]) + current_l = tl.sum(exp_scores, axis=1) + + # Update L <- L * exp(M - M') + L1, M <- M' + alpha = tl.math.exp2(m - new_m) + l = l * alpha + current_l + m = new_m + + # Update O <- O * exp(M - M') + P @ V + acc = (acc * alpha[:, None] + tl.dot(exp_scores.to(v.type.element_ty), v)) + + return m, l, acc + + +@triton.autotune( + configs=get_autotune_config(), + key=['head_dim'], +) +@triton.jit +def triton_sta_kernel( + Q, K, V, output, + batch_size: int, num_heads: int, seq_len: int, head_dim: int, + img_seq_len: int, + text_length: int, + canvas_t: int, canvas_h: int, canvas_w: int, + kernel_t: int, kernel_h: int, kernel_w: int, + tile_t: int, tile_h: int, tile_w: int, + scale: float, + has_text: tl.constexpr, + text_q: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_KV: tl.constexpr, + BLOCK_DIM: tl.constexpr, +): + total_tile_size = tile_t * tile_h * tile_w + q_block_per_tile = (total_tile_size + BLOCK_Q - 1) // BLOCK_Q + + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + if text_q: + q_block_idx = tl.program_id(2) + else: + q_tile_flat = tl.program_id(2) // q_block_per_tile + q_block_idx = tl.program_id(2) % q_block_per_tile + + m = tl.full((BLOCK_Q,), -float('inf'), dtype=tl.float32) + l = tl.zeros((BLOCK_Q,), dtype=tl.float32) + acc = tl.zeros((BLOCK_Q, BLOCK_DIM), dtype=tl.float32) + + q_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim + if text_q: + q_base_idx = img_seq_len + q_block_idx * BLOCK_Q + else: + q_base_idx = q_tile_flat * total_tile_size + q_block_idx * BLOCK_Q + + q_offset_in_tile = tl.arange(0, BLOCK_Q) + q_idx = q_base_idx + q_offset_in_tile + q_mask = (q_block_idx * BLOCK_Q + tl.arange(0, BLOCK_Q)) < total_tile_size + + q = tl.load( + Q + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=q_mask[:, None], + other=0.0 + ) # [BLOCK_Q, BLOCK_DIM] + + # Scale sm_scale by log_2(e) and use 2^x instead of exp + sm_scale = scale * 1.4426950408889634 + + num_tiles_t = canvas_t // tile_t + num_tiles_h = canvas_h // tile_h + num_tiles_w = canvas_w // tile_w + tiles_per_hw = num_tiles_h * num_tiles_w + + if text_q: + kv_tile_start_t = 0 + kv_tile_end_t = num_tiles_t + + kv_tile_start_h = 0 + kv_tile_end_h = num_tiles_h + + kv_tile_start_w = 0 + kv_tile_end_w = num_tiles_w + + else: + q_tile_t = q_tile_flat // tiles_per_hw + remaining = q_tile_flat % tiles_per_hw + q_tile_h = remaining // num_tiles_w + q_tile_w = remaining % num_tiles_w + + kernel_center_t = clamp_int(q_tile_t, kernel_t // 2, (num_tiles_t - 1) - kernel_t // 2) + kernel_center_h = clamp_int(q_tile_h, kernel_h // 2, (num_tiles_h - 1) - kernel_h // 2) + kernel_center_w = clamp_int(q_tile_w, kernel_w // 2, (num_tiles_w - 1) - kernel_w // 2) + + kv_tile_start_t = kernel_center_t - kernel_t // 2 + kv_tile_end_t = kernel_center_t + kernel_t // 2 + 1 + kv_tile_end_t = tl.where(kv_tile_end_t > num_tiles_t, num_tiles_t, kv_tile_end_t) + + kv_tile_start_h = kernel_center_h - kernel_h // 2 + kv_tile_end_h = kernel_center_h + kernel_h // 2 + 1 + kv_tile_end_h = tl.where(kv_tile_end_h > num_tiles_h, num_tiles_h, kv_tile_end_h) + + kv_tile_start_w = kernel_center_w - kernel_w // 2 + kv_tile_end_w = kernel_center_w + kernel_w // 2 + 1 + kv_tile_end_w = tl.where(kv_tile_end_w > num_tiles_w, num_tiles_w, kv_tile_end_w) + + # for kv_img + for kv_tile_t in tl.range(kv_tile_start_t, kv_tile_end_t): + for kv_tile_h in tl.range(kv_tile_start_h, kv_tile_end_h): + for kv_tile_w in tl.range(kv_tile_start_w, kv_tile_end_w): + kv_base_idx = (kv_tile_t * num_tiles_h * num_tiles_w + kv_tile_h * num_tiles_w + kv_tile_w) * total_tile_size + + for kv_block_idx in tl.range(0, total_tile_size, BLOCK_KV): + kv_offset_in_block = tl.arange(0, BLOCK_KV) + kv_idx = kv_base_idx + kv_block_idx + kv_offset_in_block + kv_mask = (kv_block_idx + tl.arange(0, BLOCK_KV)) < total_tile_size + + kv_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim + + k = tl.load( + K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=kv_mask[:, None], + other=0.0 + ) # [BLOCK_KV, BLOCK_DIM] + v = tl.load( + V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=kv_mask[:, None], + other=0.0 + ) # [BLOCK_KV, BLOCK_DIM] + + m, l, acc = _attn_fwd_loop(q, k, v, kv_mask, m, l, acc, sm_scale, False) + + + # for kv_text + if has_text: + kv_base_idx = img_seq_len + for kv_block_idx in tl.range(0, total_tile_size, BLOCK_KV): + kv_offset_in_block = tl.arange(0, BLOCK_KV) + kv_idx = kv_base_idx + kv_block_idx + kv_offset_in_block + kv_mask = (kv_block_idx + tl.arange(0, BLOCK_KV)) < text_length + + kv_offset = (batch_idx * num_heads + head_idx) * seq_len * head_dim + + k = tl.load( + K + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=kv_mask[:, None], + other=0.0 + ) # [BLOCK_KV, BLOCK_DIM] + v = tl.load( + V + kv_offset + kv_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + mask=kv_mask[:, None], + other=0.0 + ) # [BLOCK_KV, BLOCK_DIM] + + m, l, acc = _attn_fwd_loop(q, k, v, kv_mask, m, l, acc, sm_scale, True) + + + output_acc = acc / l[:, None] + tl.store( + output + q_offset + q_idx[:, None] * head_dim + tl.arange(0, BLOCK_DIM)[None, :], + output_acc, + mask=q_mask[:, None] + ) # [BLOCK_Q, BLOCK_DIM] + + +def sliding_tile_attention_triton( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + window_size, text_length: int, + has_text=True, dit_seq_shape='30x48x80') -> torch.Tensor: + seq_length = q.shape[2] + if has_text: + assert q.shape[2] >= 115200 and q.shape[2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '30x48x80' for HunyuanVideo" + target_size = math.ceil(seq_length / 384) * 384 + pad_size = target_size - seq_length + if pad_size > 0: + q = torch.cat([q, q[:, :, -pad_size:]], dim=2) + k = torch.cat([k, k[:, :, -pad_size:]], dim=2) + v = torch.cat([v, v[:, :, -pad_size:]], dim=2) + else: + if dit_seq_shape == '36x48x48': # Stepvideo + assert q.shape[2] == 82944 + elif dit_seq_shape == '18x48x80': # Wan + assert q.shape[2] == 69120 + else: + raise ValueError(f"Unsupported {dit_seq_shape}, current shape is {q.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan") + assert q.shape[1] == len(window_size), "Number of heads must match the number of window sizes" + + batch_size, num_heads, seq_len, head_dim = q.shape + if dit_seq_shape == '30x48x80': # Hunyuan + canvas_t, canvas_h, canvas_w = 30, 48, 80 + tile_t, tile_h, tile_w = 6, 8, 8 + elif dit_seq_shape == '36x48x48': # Stepvideo + canvas_t, canvas_h, canvas_w = 36, 48, 48 + tile_t, tile_h, tile_w = 6, 8, 8 + elif dit_seq_shape == '18x48x80': # Wan + canvas_t, canvas_h, canvas_w = 18, 48, 80 + tile_t, tile_h, tile_w = 6, 8, 8 + + img_seq_len = canvas_t * canvas_h * canvas_w + + num_tiles_t = canvas_t // tile_t + num_tiles_h = canvas_h // tile_h + num_tiles_w = canvas_w // tile_w + num_tiles = num_tiles_t * num_tiles_h * num_tiles_w + + total_tile_size = tile_t * tile_h * tile_w + + # BLOCK_Q=128 + # BLOCK_KV=128 + BLOCK_DIM = head_dim + + output = torch.empty_like(q) + + # for q_img + # kernel_size maybe different for different head + # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs + for head_index, (kernel_t, kernel_h, kernel_w) in enumerate(window_size): + for batch in range(batch_size): + q_head, k_head, v_head, o_head = (q[batch:batch + 1, head_index:head_index + 1], + k[batch:batch + 1, head_index:head_index + 1], + v[batch:batch + 1, head_index:head_index + 1], + output[batch:batch + 1, head_index:head_index + 1]) + + # triton_sta_kernel[(1, 1, num_tiles * triton.cdiv(total_tile_size, BLOCK_Q))]( + grid = lambda META: (1, 1, num_tiles * triton.cdiv(total_tile_size, META['BLOCK_Q'])) + triton_sta_kernel[grid]( + q_head, k_head, v_head, o_head, + 1, 1, seq_len, head_dim, + img_seq_len, + text_length, + canvas_t, canvas_h, canvas_w, + kernel_t, kernel_h, kernel_w, + tile_t, tile_h, tile_w, + scale=1.0 / (head_dim ** 0.5), + has_text=has_text, + text_q=False, + # BLOCK_Q=BLOCK_Q, + # BLOCK_KV=BLOCK_KV, + BLOCK_DIM=BLOCK_DIM, + ) + + # for q_text + # kernel_t, kernel_h, kernel_w is not used, set to (3, 3, 3) + if has_text: + # triton_sta_kernel[(batch_size, num_heads, triton.cdiv(total_tile_size, BLOCK_Q))]( + grid = lambda META: (batch_size, num_heads, triton.cdiv(total_tile_size, META['BLOCK_Q'])) + triton_sta_kernel[grid]( + q, k, v, output, + batch_size, num_heads, seq_len, head_dim, + img_seq_len, + text_length, + canvas_t, canvas_h, canvas_w, + 3, 3, 3, + #kernel_t, kernel_h, kernel_w, + tile_t, tile_h, tile_w, + scale=1.0 / (head_dim ** 0.5), + has_text=has_text, + text_q=True, + # BLOCK_Q=BLOCK_Q, + # BLOCK_KV=BLOCK_KV, + BLOCK_DIM=BLOCK_DIM, + ) + + if has_text: + if pad_size > 0: + output = output[:, :, :seq_length] + + return output \ No newline at end of file diff --git a/fastvideo/platforms/rocm.py b/fastvideo/platforms/rocm.py index 684dfb767..2b1445d71 100644 --- a/fastvideo/platforms/rocm.py +++ b/fastvideo/platforms/rocm.py @@ -58,6 +58,13 @@ class RocmPlatform(Platform): torch.cuda.reset_peak_memory_stats(device) return float(torch.cuda.max_memory_allocated(device)) + @classmethod + def get_torch_device(cls): + """ + Return torch.cuda + """ + return torch.cuda + @classmethod def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype) -> str: @@ -71,8 +78,23 @@ class RocmPlatform(Platform): elif selected_backend in (AttentionBackendEnum.FLASH_ATTN, None): pass - elif selected_backend in (AttentionBackendEnum.SLIDING_TILE_ATTN, - AttentionBackendEnum.SAGE_ATTN): + elif selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN: + try: + from st_attn import sliding_tile_attention # noqa: F401 + + from fastvideo.attention.backends.sliding_tile_attn import ( # noqa: F401 + SlidingTileAttentionBackend) + logger.info("Using Sliding Tile Attention backend.") + + return "fastvideo.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend" + except ImportError as e: + logger.error( + "Failed to import Sliding Tile Attention backend: %s", + str(e)) + raise ImportError( + "Sliding Tile Attention backend is not installed. ") from e + + elif selected_backend in (AttentionBackendEnum.SAGE_ATTN): raise ValueError( f"{selected_backend.name} is not supported on {cls.device_name}." ) diff --git a/pyproject.toml b/pyproject.toml index 314f174f7..16d950a3b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "timm==1.0.11", "peft>=0.15.0", "diffusers>=0.33.1", - "torch==2.9.0", + "torch>=2.9.0", "torchvision", # Acceleration & Optimization @@ -101,14 +101,14 @@ explicit = true [project.optional-dependencies] -# flash-attn: pip install flash-attn==2.8.1 --no-cache-dir --no-build-isolation +# flash-attn: pip install flash-attn==2.8.1 --no-cache-dir --no-build-isolation lint = [ "pre-commit==4.0.1", ] -test = [ +test = [ "av==14.3.0", "pytorch-msssim==1.0.0", "pytest", @@ -116,6 +116,10 @@ test = [ dev = [ "fastvideo[lint]", "fastvideo[test]", ] +rocm = [ + "amdsmi", +] + [project.scripts] fastvideo = "fastvideo.entrypoints.cli.main:main" diff --git a/scripts/inference/v1_inference_hunyuan_STA.sh b/scripts/inference/v1_inference_hunyuan_STA.sh index 1d2186c3f..0e42ad0c6 100644 --- a/scripts/inference/v1_inference_hunyuan_STA.sh +++ b/scripts/inference/v1_inference_hunyuan_STA.sh @@ -9,6 +9,7 @@ fastvideo generate \ --model-path $MODEL_BASE \ --sp-size ${num_gpus} \ --tp-size 1 \ + --num-gpus ${num_gpus} \ --height 768 \ --width 1280 \ --num-frames 117 \