[feat] add sliding_tile attention triton kernel and ROCM support (#916)

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.
 

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",

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]
+    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.")

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

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}."
             )

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"
 

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 \