Revert "[chore] release v0.1.7 (#955)"

This reverts commit 9cd6a86b95.

.buildkite/pipeline.yml

@@ -129,11 +129,7 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
-                - "csrc/attn/video_sparse_attn/**"
-                - "csrc/attn/video_sparse_attn/tk/**"
-                - "csrc/attn/video_sparse_attn/setup.py"
-                - "csrc/attn/video_sparse_attn/config_vsa.py"
-                - "csrc/attn/video_sparse_attn/vsa.cpp"
+                - "fastvideo-kernel/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -145,10 +141,7 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
-                - "csrc/attn/sliding_tile_attn/**"
-                - "csrc/attn/sliding_tile_attn/setup.py"
-                - "csrc/attn/sliding_tile_attn/config_sta.py"
-                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
+                - "fastvideo-kernel/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -159,10 +152,7 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "csrc/attn/sliding_tile_attn/**"
-                - "csrc/attn/sliding_tile_attn/setup.py"
-                - "csrc/attn/sliding_tile_attn/config_sta.py"
-                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
+                - "fastvideo-kernel/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -173,12 +163,7 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "csrc/attn/video_sparse_attn/**"
-                - "csrc/attn/video_sparse_attn/tk/**"
-                - "csrc/attn/tests/test_vsa.py"
-                - "csrc/attn/video_sparse_attn/setup.py"
-                - "csrc/attn/video_sparse_attn/config_vsa.py"
-                - "csrc/attn/video_sparse_attn/vsa.cpp"
+                - "fastvideo-kernel/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -189,7 +174,7 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "csrc/attn/vmoba_attn/**"
+                - "fastvideo-kernel/**"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -200,7 +185,7 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "csrc/attn/vmoba_attn/vmoba/**"
+                - "fastvideo-kernel/**"
                 - "fastvideo/attention/backends/vmoba.py"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"

.github/workflows/fastvideo-kernel-publish.yml

@@ -0,0 +1,207 @@
+name: Publish FastVideo Kernel to PyPI on Version Change
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - "fastvideo-kernel/pyproject.toml"
+  workflow_dispatch:
+
+jobs:
+  check-version-change:
+    runs-on: ubuntu-latest
+    outputs:
+      version-changed: ${{ steps.check-version.outputs.changed }}
+      new-version: ${{ steps.check-version.outputs.new-version }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 2
+
+      - name: Check if version changed
+        id: check-version
+        run: |
+          cd fastvideo-kernel
+          # Get current commit's version from pyproject.toml
+          # Use ^ to match start of line to avoid matching minimum-version
+          NEW_VERSION=$(grep -oP '^version\s*=\s*"\K[^"]+' pyproject.toml)
+          echo "New version: $NEW_VERSION"
+          
+          # Get previous version from git history
+          # Note: git show expects path relative to repo root
+          OLD_VERSION=$(git show HEAD~1:fastvideo-kernel/pyproject.toml | grep -oP '^version\s*=\s*"\K[^"]+' || echo "0.0.0")
+          echo "Old version: $OLD_VERSION"
+          
+          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
+            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
+            echo "changed=true" >> $GITHUB_OUTPUT
+            echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT
+          else
+            echo "Version did not change"
+            echo "changed=false" >> $GITHUB_OUTPUT
+          fi
+
+  build_wheels:
+    name: Build Wheel
+    needs: check-version-change
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+          os: [ubuntu-22.04]
+          python-version: ['3.10', '3.11', '3.12']
+          torch-cuda:
+            # - torch-version: '2.5.1'
+            #   cuda-version: '12.4.1'
+            #   torch-cuda-short: 'cu124'
+            # - torch-version: '2.6.0'
+            #   cuda-version: '12.6.3'
+            #   torch-cuda-short: 'cu126'
+            # - torch-version: '2.7.1'
+            #   cuda-version: '12.8.0'
+            #   torch-cuda-short: 'cu128'
+            - torch-version: '2.9.1'
+              cuda-version: '12.8.0'
+              torch-cuda-short: 'cu128'
+
+    steps:
+      - name: Free up disk space
+        run: |
+          echo "Initial disk space:"
+          df -h
+          
+          # Remove large directories
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          sudo rm -rf /usr/share/swift
+          sudo rm -rf /usr/local/lib/node_modules
+          sudo rm -rf /usr/local/share/powershell
+          sudo rm -rf /usr/share/rust
+          sudo rm -rf /usr/local/.ghcup
+          
+          # Remove cached files
+          sudo rm -rf /var/lib/apt/lists/*
+          sudo rm -rf /var/cache/apt/archives/*
+          
+          echo "Disk space after cleanup:"
+          df -h
+
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install CUDA ${{ matrix.torch-cuda.cuda-version }}
+        uses: Jimver/cuda-toolkit@v0.2.21
+        id: cuda-toolkit
+        with:
+          cuda: ${{ matrix.torch-cuda.cuda-version }}
+          linux-local-args: '["--toolkit"]'
+          method: 'network'
+
+      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
+        run: |
+          sudo apt update
+          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
+          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+          # Allow Git to Access Safe Directory
+          git config --global --add safe.directory /__w/FastVideo/FastVideo
+
+          # Set CUDA environment variables
+          export CUDA_HOME=/usr/local/cuda-${{ matrix.torch-cuda.cuda-version }}
+          export PATH=${CUDA_HOME}/bin:${PATH}
+          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+          
+          # Verify installation
+          gcc --version
+          g++ --version
+          clang-11 --version
+          nvcc --version
+
+      - name: Install PyTorch ${{ matrix.torch-cuda.torch-version }}+cu${{ matrix.torch-cuda.cuda-version }}
+        run: |
+          pip install --upgrade pip
+          pip install typing-extensions==4.12.2
+          pip install --no-cache-dir torch==${{ matrix.torch-cuda.torch-version }} --index-url https://download.pytorch.org/whl/${{matrix.torch-cuda.torch-cuda-short}}
+          nvcc --version
+          python --version
+          python -c "import torch; print('PyTorch:', torch.__version__)"
+          python -c "import torch; print('CUDA:', torch.version.cuda)"
+          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
+
+      - name: Build wheel
+        run: |
+          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
+          
+          pip install setuptools ninja packaging wheel triton scikit-build-core cmake build
+          
+          cd fastvideo-kernel
+          git submodule update --init --recursive  # Ensure ThunderKittens submodule is initialized
+          # Release builds are produced on GPU-less runners, so force-enable TK and target Hopper.
+          export TORCH_CUDA_ARCH_LIST="9.0a"
+          export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
+          
+          # Build standard wheel (no local version suffix) for PyPI
+          python -m build --wheel --outdir dist
+          
+          # Fix the wheel to be manylinux compliant
+          pip install auditwheel
+          # Target manylinux_2_35 (Ubuntu 22.04 native)
+          auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist
+          # Move fixed wheels back to dist for upload consistency
+          rm dist/*.whl
+          mv fixed_dist/*.whl dist/
+
+      - name: Upload wheel artifact
+        # Only upload if it's the "main" CUDA version we want on PyPI
+        # We upload all to artifacts for inspection/GH releases, but give them distinct artifact names
+        uses: actions/upload-artifact@v4
+        with:
+          name: fastvideo_kernel-py${{ matrix.python-version }}-${{ matrix.torch-cuda.torch-cuda-short }}-torch${{ matrix.torch-cuda.torch-version }}
+          path: fastvideo-kernel/dist/*.whl
+          retention-days: 90
+
+  publish_package:
+    name: Publish package
+    needs: [build_wheels, check-version-change]
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
+    runs-on: ubuntu-22.04
+    permissions:
+      id-token: write  # Needed for OIDC Trusted Publishing
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+
+      - name: Download PyPI wheels
+        uses: actions/download-artifact@v4
+        with:
+          path: fastvideo-kernel/dist/
+          pattern: 'fastvideo_kernel-py*'
+          merge-multiple: true
+
+      - name: Build source distribution
+        run: |
+          pip install build scikit-build-core cmake ninja
+          
+          cd fastvideo-kernel
+          # We don't need full CUDA/Torch to just package the source (sdist)
+          python -m build --sdist --outdir dist
+
+      - name: Publish release distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: fastvideo-kernel/dist/

.gitignore

@@ -30,6 +30,8 @@ env
 **/build/
 **.pyc
 **.txt
+*.log
+weights/
 
 # Distribution / packaging
 build/

.gitmodules

@@ -1,7 +1,3 @@
-[submodule "csrc/attn/video_sparse_attn/tk"]
-	path = csrc/attn/video_sparse_attn/tk
-	url = https://github.com/HazyResearch/ThunderKittens.git
-
-[submodule "csrc/attn/sliding_tile_attn/tk"]
-	path = csrc/attn/sliding_tile_attn/tk
+[submodule "fastvideo-kernel/include/tk"]
+	path = fastvideo-kernel/include/tk
 	url = https://github.com/HazyResearch/ThunderKittens.git

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ default_stages:
 exclude: |
   (?x)(
     fastvideo/third_party/.*|
-    csrc/.*|
+    fastvideo-kernel/.*|
     assets/.*|
     tests/.*|
     demo/.*|

csrc/attn/README.md

@@ -1,113 +0,0 @@
-
-
-# Attention Kernel Used in FastVideo
-
-
-
-## Video Sparse Attention (VSA)
-
-### Installation
-We support H100 (via TK) and any other GPU (via triton) for VSA.
-```bash
-git submodule update --init --recursive
-python setup_vsa.py install
-```
-
-
-If you encounter error during installation, try below:
-Install C++20 for ThunderKittens:
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-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 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-```
-
-### Verify if you have successfully installed
-
-```bash
-# test numerical
-python tests/test_vsa.py
-# (For H100) test speed
-python benchmarks/bench_vsa_hopper.py
-```
-bench_vsa_hopper.py should print something like this:
-```bash
-Using topk=76 kv blocks per q block (out of 768 total kv blocks)
-
-=== BLOCK SPARSE ATTENTION BENCHMARK ===
-Block Sparse Forward  - TFLOPS: 5622.26
-Block Sparse Backward - TFLOPS: 3865.68
-```
-
-
-## Sliding Tile Attention (STA)
-We only support H100 for STA.
-```bash
-git submodule update --init --recursive
-python setup_sta.py install
-```
-
-
-
-
-###  Usage
-End-2-end inference with FastVideo:
-```bash
-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. 
-# 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
-# If your attention contains text token (Hunyuan)
-out = sliding_tile_attention(q, k, v, window_size, text_length)
-# If your attention does not contain text token (StepVideo)
-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
-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. 
-
-
-https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747
-
-## 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. 
-
-STA removes mixed blocks.
-
-
-<div align="center">
-<img src=../../assets/sliding_tile_attn_map.png width="80%"/>
-</div>
-
-## Acknowledgement
-
-We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/attn/benchmarks/bench_sta.py

@@ -1,145 +0,0 @@
-import os
-from collections import defaultdict
-
-import matplotlib.pyplot as plt
-import numpy as np
-import torch
-from st_attn import sliding_tile_attention
-from triton.testing import do_bench
-
-
-def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
-    assert mode in ["fwd", "bwd", "fwd_bwd"]
-    f = 4 * batch * seqlen**2 * nheads * headdim // (2 if causal else 1)
-    return f if mode == "fwd" else (2.5 * f if mode == "bwd" else 3.5 * f)
-
-
-def compute_TFLOPS(flops, ms):
-    flops = flops / 1e12
-    ms = ms / 1e3
-    return flops / ms
-
-
-def benchmark_attention(configurations):
-    results = {'fwd': defaultdict(list), 'bwd': defaultdict(list)}
-
-    for B, H, N, D, causal, dit_seq_shape, window_size in configurations:
-        print("=" * 60)
-        print(f"Timing forward and backward pass for B={B}, H={H}, N={N}, D={D}, causal={causal}")
-
-        q = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
-        k = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
-        v = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
-
-        # grad_output = torch.randn_like(q, requires_grad=False).contiguous()
-        # qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
-        # kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
-        # vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()
-
-
-        # # Warmup for forward pass
-        # for _ in range(10):
-        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
-
-        # # Time the forward pass
-        # for i in range(10):
-        #     start_events_fwd[i].record()
-        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
-        #     end_events_fwd[i].record()
-        ms = do_bench(lambda: sliding_tile_attention(q, k, v, [window_size] * 24, 0, False, dit_seq_shape))
-
-        # times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
-        # time_us_fwd = np.mean(times_fwd) * 1000
-
-        tflops_fwd = compute_TFLOPS(flops(B, N, H, D, causal, 'fwd'), ms)
-        results['fwd'][(D, causal)].append((N, tflops_fwd))
-
-        print(f"Average time for forward pass (ms): {ms:.2f}")
-        print(f"Average TFLOPS: {tflops_fwd}")
-        print("-" * 60)
-
-        # torch.cuda.empty_cache()
-        # torch.cuda.synchronize()
-
-        # # Prepare for timing backward pass
-        # start_events_bwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
-        # end_events_bwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
-
-        # # Warmup for backward pass
-        # for _ in range(10):
-        #     qg, kg, vg = tk.mha_backward(q, k, v, o, l_vec, grad_output, causal)
-
-        # # Time the backward pass
-        # for i in range(10):
-        #     start_events_bwd[i].record()
-        #     qg, kg, vg = tk.mha_backward(q, k, v, o, l_vec, grad_output, causal)
-        #     end_events_bwd[i].record()
-
-        # torch.cuda.synchronize()
-        # times_bwd = [s.elapsed_time(e) for s, e in zip(start_events_bwd, end_events_bwd)]
-        # time_us_bwd = np.mean(times_bwd) * 1000
-
-        # tflops_bwd = compute_TFLOPS(flops(B, N, H, D, causal, 'bwd'), ms)
-        # results['bwd'][(D, causal)].append((N, tflops_bwd))
-
-        # print(f"Average time for backward pass(ms): {ms:.2f}")
-        # print(f"Average TFLOPS: {tflops_bwd}")
-        # print("=" * 60)
-
-    return results
-
-
-def plot_results(results):
-    os.makedirs('benchmark_results', exist_ok=True)
-    for mode in ['fwd', 'bwd']:
-        for (D, causal), values in results[mode].items():
-            seq_lens = [x[0] for x in values]
-            tflops = [x[1] for x in values]
-
-            plt.figure(figsize=(10, 6))
-            bars = plt.bar(range(len(seq_lens)), tflops, tick_label=seq_lens)
-            plt.xlabel('Sequence Length')
-            plt.ylabel('TFLOPS')
-            plt.title(f'{mode.upper()} Pass - Head Dim: {D}, Causal: {causal}')
-            plt.grid(True)
-
-            # Adding the numerical y value on top of each bar
-            for bar in bars:
-                yval = bar.get_height()
-                plt.text(bar.get_x() + bar.get_width() / 2, yval, round(yval, 2), ha='center', va='bottom')
-
-            filename = f'benchmark_results/{mode}_D{D}_causal{causal}.png'
-            plt.savefig(filename)
-            plt.close()
-
-
-# Example list of configurations to test
-configurations = [
-    (2, 24, 69120, 128, False, '18x48x80', [3, 6, 10]),
-    (2, 24, 69120, 128, True, '18x48x80', [3, 6, 10]),
-    (2, 24, 82944, 128, False, '36x48x48', [3, 3, 6]), # Stepvideo
-    (2, 24, 82944, 128, True, '36x48x48', [3, 3, 6]),
-    # (16, 16, 768*16,  128, False),
-    # (16, 16, 768*2,  128, False),
-    # (16, 16, 768*4,  128, False),
-    # (16, 16, 768*8,  128, False),
-    # (16, 16, 768*16, 128, False),
-    # (16, 16, 768,    128, True),
-    # (16, 16, 768*2,  128, True),
-    # (16, 16, 768*4,  128, True),
-    # (16, 16, 768*8,  128, True),
-    # (16, 16, 768*16, 128, True),
-    # (16, 32, 768,    64,  False),
-    # (16, 32, 768*2,  64,  False),
-    # (16, 32, 768*4,  64,  False),
-    # (16, 32, 768*8,  64,  False),
-    # (16, 32, 768*16, 64,  False),
-    # (16, 32, 768,    64,  True),
-    # (16, 32, 768*2,  64,  True),
-    # (16, 32, 768*4,  64,  True),
-    # (16, 32, 768*8,  64,  True),
-    # (16, 32, 768*16, 64,  True),
-]
-
-results = benchmark_attention(configurations)
-# plot_results(results)

csrc/attn/benchmarks/bench_vsa_hopper.py

@@ -1,224 +0,0 @@
-import torch
-import argparse
-from triton.testing import do_bench
-from vsa import block_sparse_fwd, block_sparse_bwd
-from vsa import BLOCK_M, BLOCK_N
-import  triton
-import numpy as np
-import random
-
-def set_seed(seed: int = 42):
-    # Python random module
-    random.seed(seed)
-
-    # NumPy
-    np.random.seed(seed)
-
-    # PyTorch
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)  # if using multi-GPU
-
-def parse_arguments():
-    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
-    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
-    parser.add_argument('--num_heads', type=int, default=12, help='Number of heads')
-    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
-    parser.add_argument('--topk', type=int, default=None, help='Number of kv blocks each q block attends to')
-    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[49152], help='Sequence lengths to benchmark')
-    return parser.parse_args()
-
-def create_input_tensors(batch, head, seq_len, headdim):
-    """Create random input tensors for attention."""
-    q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    return q, k, v
-
-def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"):
-    """
-    Generate a block sparse pattern where each q block attends to exactly k kv blocks.
-    
-    Args:
-        bs: batch size
-        h: number of heads
-        num_q_blocks: number of query blocks
-        num_kv_blocks: number of key-value blocks
-        k: number of kv blocks each q block attends to
-        device: device to create tensors on
-        
-    Returns:
-        q2k_block_sparse_index: [bs, h, num_q_blocks, k]
-            Contains the indices of kv blocks that each q block attends to.
-        q2k_block_sparse_num: [bs, h, num_q_blocks]
-            Contains the number of kv blocks that each q block attends to (all equal to k).
-        k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks]
-            Contains the indices of q blocks that attend to each kv block.
-        k2q_block_sparse_num: [bs, h, num_kv_blocks]
-            Contains the number of q blocks that attend to each kv block.
-        block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]
-            Binary mask where 1 indicates attention connection.
-    """
-    # Ensure k is not larger than num_kv_blocks
-    k = min(k, num_kv_blocks)
-    
-    # Create random scores for sampling
-    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device)
-    
-    # Get top-k indices for each q block
-    _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1)
-    q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32)
-    
-    # sort q2k_block_sparse_index
-    q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1)
-
-    # All q blocks attend to exactly k kv blocks
-    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device)
-    
-    # Create the corresponding mask
-    block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
-    
-    # Fill in the mask based on the indices
-    for b in range(bs):
-        for head in range(h):
-            for q_idx in range(num_q_blocks):
-                kv_indices = q2k_block_sparse_index[b, head, q_idx]
-                block_sparse_mask[b, head, q_idx, kv_indices] = True
-    
-    # Create the reverse mapping (k2q)
-    # First, initialize lists to collect q indices for each kv block
-    k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)]
-    
-    # Populate the lists based on q2k mapping
-    for b in range(bs):
-        for head in range(h):
-            flat_idx = b * h + head
-            for q_idx in range(num_q_blocks):
-                kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist()
-                for kv_idx in kv_indices:
-                    k2q_indices_list[flat_idx][kv_idx].append(q_idx)
-    
-    # Find the maximum number of q blocks that attend to any kv block
-    max_q_per_kv = 0
-    for flat_idx in range(bs * h):
-        for kv_idx in range(num_kv_blocks):
-            max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx]))
-    
-    # Create tensors for k2q mapping
-    k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, 
-                                        dtype=torch.int32, device=device)
-    k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), 
-                                       dtype=torch.int32, device=device)
-    
-    # Fill the tensors
-    for b in range(bs):
-        for head in range(h):
-            flat_idx = b * h + head
-            for kv_idx in range(num_kv_blocks):
-                q_indices = k2q_indices_list[flat_idx][kv_idx]
-                num_q = len(q_indices)
-                k2q_block_sparse_num[b, head, kv_idx] = num_q
-                if num_q > 0:
-                    k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor(
-                        q_indices, dtype=torch.int32, device=device)
-                
-    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
-
-def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops):
-    """Benchmark block sparse attention forward and backward passes."""
-    print("\n=== BLOCK SPARSE ATTENTION BENCHMARK ===")
-    
-    # Forward pass
-    # Warm-up run
-    variable_block_sizes = torch.ones(q2k_block_sparse_index.shape[2], device=q.device).int() * BLOCK_M
-    o, l_vec = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
-    torch.cuda.synchronize()
-    
-    # Benchmark forward
-    fwd_time = do_bench(
-        lambda: block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes),
-        warmup=5,
-        rep=20,
-        quantiles=None
-    )
-    
-    sparse_tflops = flops / fwd_time * 1e-12 * 1e3
-    print(f"Block Sparse Forward - TFLOPS: {sparse_tflops:.2f}")
-    
-    # Backward pass
-    grad_output = torch.randn_like(o)
-    
-    # Warm-up runs
-    for _ in range(5):
-        block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes)
-    torch.cuda.synchronize()
-    
-    # Benchmark backward
-    bwd_time = do_bench(
-        lambda: block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes),
-        warmup=5,
-        rep=20,
-        quantiles=None
-    )
-    bwd_flops = 2.5 * flops  # Approximation
-
-    sparse_bwd_tflops = bwd_flops / bwd_time * 1e-12 * 1e3
-    print(f"Block Sparse Backward - TFLOPS: {sparse_bwd_tflops:.2f}")
-    
-    return sparse_tflops, sparse_bwd_tflops
-
-def main():
-    args = parse_arguments()
-
-    set_seed(42)
-    
-    # Extract parameters
-    batch = args.batch_size
-    head = args.num_heads
-    headdim = args.head_dim
-    
-    print(f"Block Sparse Attention Benchmark")
-    print(f"batch: {batch}, head: {head}, headdim: {headdim}")
-    
-    # Test with different sequence lengths
-    for seq_len in args.seq_lengths:
-        # Skip very long sequences if they might cause OOM
-        if seq_len > 16384 and batch > 1:
-            continue
-            
-        print("="*100)
-        print(f"\nSequence length: {seq_len}")
-        
-        # Calculate theoretical FLOPs for attention
-        flops = 4 * batch * head * headdim * seq_len * seq_len
-        
-        # Create input tensors
-        q, k, v = create_input_tensors(batch, head, seq_len, headdim)
-        
-        # Setup block sparse parameters
-        num_q_blocks = seq_len // BLOCK_M
-        num_kv_blocks = seq_len // BLOCK_N
-        
-        # Determine k value (number of kv blocks per q block)
-        topk = args.topk
-        if topk is None:
-            topk = num_kv_blocks // 10  # Default to ~90% sparsity if k is not specified
-        topk = max(1, topk)       
-        print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)")
-        
-        # Generate block sparse pattern
-        q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, _ = generate_block_sparse_pattern(
-            batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda")
-        
-        # Benchmark block sparse attention
-        sparse_fwd, sparse_bwd = benchmark_block_sparse_attention(
-            q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops
-        )
-        
-        # Print results
-        print("\n=== PERFORMANCE RESULTS ===")
-        print(f"Block Sparse Forward - TFLOPS: {sparse_fwd:.2f}")
-        print(f"Block Sparse Backward - TFLOPS: {sparse_bwd:.2f}")
-
-if __name__ == "__main__":
-    main()

csrc/attn/benchmarks/bench_vsa_triton.py

@@ -1,217 +0,0 @@
-import torch
-import argparse
-import triton.testing
-from vsa import block_sparse_attn
-from vsa import BLOCK_M, BLOCK_N
-
-import numpy as np
-import random
-
-def set_seed(seed: int = 42):
-    # Python random module
-    random.seed(seed)
-
-    # NumPy
-    np.random.seed(seed)
-
-    # PyTorch
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)  # if using multi-GPU
-
-def parse_arguments():
-    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
-    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
-    parser.add_argument('--num_heads', type=int, default=12, help='Number of heads')
-    parser.add_argument('--head_dim', type=int, default=64, help='Head dimension')
-    parser.add_argument('--topk', type=int, default=None, help='Number of kv blocks each q block attends to')
-    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[49152], help='Sequence lengths to benchmark')
-    return parser.parse_args()
-
-def create_input_tensors(batch, head, seq_len, headdim):
-    """Create random input tensors for attention."""
-    q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    return q, k, v
-
-def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"):
-    """
-    Generate a block sparse pattern where each q block attends to exactly k kv blocks.
-    
-    Args:
-        bs: batch size
-        h: number of heads
-        num_q_blocks: number of query blocks
-        num_kv_blocks: number of key-value blocks
-        k: number of kv blocks each q block attends to
-        device: device to create tensors on
-        
-    Returns:
-        q2k_block_sparse_index: [bs, h, num_q_blocks, k]
-            Contains the indices of kv blocks that each q block attends to.
-        q2k_block_sparse_num: [bs, h, num_q_blocks]
-            Contains the number of kv blocks that each q block attends to (all equal to k).
-        k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks]
-            Contains the indices of q blocks that attend to each kv block.
-        k2q_block_sparse_num: [bs, h, num_kv_blocks]
-            Contains the number of q blocks that attend to each kv block.
-        block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]
-            Binary mask where 1 indicates attention connection.
-    """
-    # Ensure k is not larger than num_kv_blocks
-    k = min(k, num_kv_blocks)
-    
-    # Create random scores for sampling
-    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device)
-    
-    # Get top-k indices for each q block
-    _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1)
-    q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32)
-    
-    # sort q2k_block_sparse_index
-    q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1)
-
-    # All q blocks attend to exactly k kv blocks
-    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device)
-    
-    # Create the corresponding mask
-    block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
-    
-    # Fill in the mask based on the indices
-    for b in range(bs):
-        for head in range(h):
-            for q_idx in range(num_q_blocks):
-                kv_indices = q2k_block_sparse_index[b, head, q_idx]
-                block_sparse_mask[b, head, q_idx, kv_indices] = True
-    
-    # Create the reverse mapping (k2q)
-    # First, initialize lists to collect q indices for each kv block
-    k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)]
-    
-    # Populate the lists based on q2k mapping
-    for b in range(bs):
-        for head in range(h):
-            flat_idx = b * h + head
-            for q_idx in range(num_q_blocks):
-                kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist()
-                for kv_idx in kv_indices:
-                    k2q_indices_list[flat_idx][kv_idx].append(q_idx)
-    
-    # Find the maximum number of q blocks that attend to any kv block
-    max_q_per_kv = 0
-    for flat_idx in range(bs * h):
-        for kv_idx in range(num_kv_blocks):
-            max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx]))
-    
-    # Create tensors for k2q mapping
-    k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, 
-                                        dtype=torch.int32, device=device)
-    k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), 
-                                       dtype=torch.int32, device=device)
-    
-    # Fill the tensors
-    for b in range(bs):
-        for head in range(h):
-            flat_idx = b * h + head
-            for kv_idx in range(num_kv_blocks):
-                q_indices = k2q_indices_list[flat_idx][kv_idx]
-                num_q = len(q_indices)
-                k2q_block_sparse_num[b, head, kv_idx] = num_q
-                if num_q > 0:
-                    k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor(
-                        q_indices, dtype=torch.int32, device=device)
-                
-    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
-
-def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops):
-    """Benchmark block sparse attention forward+backward pass."""
-    print("\n=== BLOCK SPARSE ATTENTION FORWARD+BACKWARD BENCHMARK ===")
-    
-    # Combined forward+backward pass
-    # Warm-up run
-    q_fwd = q.clone().requires_grad_(True)
-    k_fwd = k.clone().requires_grad_(True)
-    v_fwd = v.clone().requires_grad_(True)
-    o = block_sparse_attn(q_fwd, k_fwd, v_fwd, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
-    grad_output = torch.randn_like(o)
-    o.backward(grad_output)
-    torch.cuda.synchronize()
-    
-    # Benchmark forward+backward
-    def forward_backward_fn():
-        q_fwd = q.clone().requires_grad_(True)
-        k_fwd = k.clone().requires_grad_(True)
-        v_fwd = v.clone().requires_grad_(True)
-        o = block_sparse_attn(q_fwd, k_fwd, v_fwd, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
-        grad_output = torch.randn_like(o)
-        o.backward(grad_output)
-    
-    total_time = triton.testing.do_bench(
-        forward_backward_fn,
-        warmup=25,
-        rep=100,
-        return_mode='mean'
-    )
-    
-    # Total flops for forward + backward (forward + 2.5x backward approximation)
-    total_flops = flops + 2.5 * flops  # 3.5x the forward flops
-    sparse_tflops = total_flops / total_time * 1e-12 * 1e3
-    print(f"Block Sparse Forward+Backward - TFLOPS: {sparse_tflops:.2f}")
-    
-    return sparse_tflops
-
-def main():
-    args = parse_arguments()
-
-    set_seed(42)
-    
-    # Extract parameters
-    batch = args.batch_size
-    head = args.num_heads
-    headdim = args.head_dim
-    
-    print(f"Block Sparse Attention Benchmark")
-    print(f"batch: {batch}, head: {head}, headdim: {headdim}")
-    
-    # Test with different sequence lengths
-    for seq_len in args.seq_lengths:
-        # Skip very long sequences if they might cause OOM
-        if seq_len > 16384 and batch > 1:
-            continue
-            
-        print("="*100)
-        print(f"\nSequence length: {seq_len}")
-        
-        # Calculate theoretical FLOPs for attention
-        flops = 4 * batch * head * headdim * seq_len * seq_len
-        
-        # Create input tensors
-        q, k, v = create_input_tensors(batch, head, seq_len, headdim)
-        
-        # Setup block sparse parameters
-        num_q_blocks = seq_len // BLOCK_M
-        num_kv_blocks = seq_len // BLOCK_N
-        
-        # Determine k value (number of kv blocks per q block)
-        topk = args.topk
-        if topk is None:
-            topk = num_kv_blocks // 10  # Default to ~90% sparsity if k is not specified
-        topk = max(1, topk)       
-        print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)")
-        
-        # Generate block sparse pattern
-        q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, _ = generate_block_sparse_pattern(
-            batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda")
-        
-        # Benchmark block sparse attention
-        sparse_fwd = benchmark_block_sparse_attention(
-            q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops
-        )
-        
-        # Print results
-        print("\n=== PERFORMANCE RESULTS ===")
-        print(f"Block Sparse Forward+Backward - TFLOPS: {sparse_fwd:.2f}")
-
-if __name__ == "__main__":
-    main()

csrc/attn/sliding_tile_attn/MANIFEST.in

@@ -1,2 +0,0 @@
-recursive-include tk *
-include config_sta.py

csrc/attn/sliding_tile_attn/README.md

@@ -1,96 +0,0 @@
-
-# Attention Kernel Used in FastVideo
-
-## Sliding Tile Attention (STA)
-We only support H100 for STA.
-
-### Installation
-```bash
-pip install st_attn
-``` 
-
-Install from source:
-
-```bash
-git submodule update --init --recursive
-python setup.py install
-```
-
-If you encounter error during installation, try below:
-Install C++20 for ThunderKittens:
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-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 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-```
-
-###  Usage
-End-2-end inference with FastVideo:
-```bash
-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. 
-# 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
-# If your attention contains text token (Hunyuan)
-out = sliding_tile_attention(q, k, v, window_size, text_length)
-# If your attention does not contain text token (StepVideo)
-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
-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. 
-
-
-https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747
-
-
-## STA Configuration Logic
-Here is a diagram of how the window is configured and passed through the FastVideo pipeline:
-
-<div align="center">
-<img src="../../../docs/assets/images/STA_configuration.png" width="80%"/>
-</div>
-
-
-## 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. 
-
-STA removes mixed blocks.
-
-
-<div align="center">
-<img src=../../../assets/sliding_tile_attn_map.png width="80%"/>
-</div>
-
-## Acknowledgement
-
-We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/attn/sliding_tile_attn/config_sta.py

@@ -1,15 +0,0 @@
-### ADD TO THIS TO REGISTER NEW KERNELS
-sources = {
-    'st_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 = ['st_attn']
-
-### WHICH GPU TARGET DO WE WANT TO BUILD FOR?
-target = 'h100'

csrc/attn/sliding_tile_attn/setup.py

@@ -1,76 +0,0 @@
-import os
-import subprocess
-
-from config_sta import kernels, sources, target
-from setuptools import find_packages, setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-target = target.lower()
-
-# Package metadata
-PACKAGE_NAME = "st_attn"
-VERSION = "0.0.6"
-AUTHOR = "Hao AI Lab"
-DESCRIPTION = "Sliding Tile Atteniton Kernel Used in FastVideo"
-URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/sliding_tile_attention"
-
-# 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=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'])
-      ],
-      cmdclass={'build_ext': BuildExtension},
-      classifiers=[
-          "Programming Language :: Python :: 3",
-          "Environment :: GPU :: NVIDIA CUDA :: 12",
-          "License :: OSI Approved :: Apache Software License",
-      ],
-      python_requires='>=3.10',
-      install_requires=["torch>=2.5.0"])

csrc/attn/sliding_tile_attn/st_attn.cpp

@@ -1,23 +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_ST_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, bool has_text, int kernel_aspect_ratio_flag
-); 
-#endif
-
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.doc() = "Sliding Block Attention Kernels"; // optional module docstring
-    
-#ifdef TK_COMPILE_ST_ATTN
-    m.def("sta_fwd",  torch::wrap_pybind_function(sta_forward), "sliding tile attention, assuming tile size is (6,8,8)");
-#endif
-}

csrc/attn/sliding_tile_attn/st_attn/__init__.py

@@ -1,49 +0,0 @@
-import math
-
-import torch
-from torch.utils.checkpoint import detach_variable
-try:
-    from st_attn_cuda import sta_fwd
-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'):
-    seq_length = q_all.shape[2]
-    dit_seq_shape_mapping = {
-        '30x48x80':1,
-        '36x48x48':2,
-        '18x48x80':3, 
-    }
-    if has_text:
-        assert q_all.shape[
-            2] >= 115200 and q_all.shape[2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '30x48x80' for HunyuanVideo"
-        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)
-    else:
-        if dit_seq_shape == '36x48x48': # Stepvideo 204x768x68
-            assert q_all.shape[2] == 82944
-        elif dit_seq_shape == '18x48x80': # Wan 69x768x1280
-            assert q_all.shape[2] == 69120
-        else:
-            raise ValueError(f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
-
-    kernel_aspect_ratio_flag = dit_seq_shape_mapping[dit_seq_shape]
-    hidden_states = torch.empty_like(q_all)
-    # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
-    for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
-        for batch in range(q_all.shape[0]):
-            q_head, k_head, v_head, o_head = (q_all[batch:batch + 1, head_index:head_index + 1],
-                                              k_all[batch:batch + 1,
-                                                    head_index:head_index + 1], v_all[batch:batch + 1,
-                                                                                      head_index:head_index + 1],
-                                              hidden_states[batch:batch + 1, head_index:head_index + 1])
-
-            _ = 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]

csrc/attn/tests/test_sta.py

@@ -1,87 +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), (18, 48, 80), 0, 'cuda', 0)
-    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, 0, False, '18x48x80')
-    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 = [(3, 3, 5), (3, 1, 10)]
-    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
-
-
-# Example usage
-b, h, d = 2, 24, 128
-n = 69120  # Sequence length
-causal = False
-mean = 1e-1
-std = 10
-
-# Run correctness check directly
-results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output')
-assert results['TK vs FLEX']['avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large"
-assert results['TK vs FLEX']['max_diff'] < 4e-2, f"Maximum difference: {results['TK vs FLEX']['max_diff']} is too large"
-print(f"Average difference: {results['TK vs FLEX']['avg_diff']}")
-print(f"Maximum difference: {results['TK vs FLEX']['max_diff']}")

csrc/attn/tests/test_vsa.py

@@ -1,245 +0,0 @@
-import torch
-import sys
-import os
-import numpy as np
-from tqdm import tqdm
-
-# Add the parent directory to the path to import block_sparse_attn
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from tests.utils import generate_block_sparse_mask_for_function, create_full_mask_from_block_mask
-from vsa import block_sparse_attn
-
-BLOCK_M = 64
-BLOCK_N = 64
-
-def pytorch_test(Q, K, V, block_sparse_mask, dO):
-    q_ = Q.clone().float().requires_grad_()
-    k_ = K.clone().float().requires_grad_()
-    v_ = V.clone().float().requires_grad_()
-
-    QK = torch.matmul(q_, k_.transpose(-2, -1))
-    QK /= (q_.size(-1) ** 0.5)
-    QK = QK.masked_fill(~block_sparse_mask.unsqueeze(0), float('-inf'))
-
-    QK = torch.nn.functional.softmax(QK, dim=-1)
-    output = torch.matmul(QK, v_)
-
-    dO_ = dO
-    output.backward(dO_)
-    return (
-        output.to(torch.bfloat16),
-        q_.grad.to(torch.bfloat16),
-        k_.grad.to(torch.bfloat16),
-        v_.grad.to(torch.bfloat16),
-    )
-
-
-def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, q_non_pad_index, kv_non_pad_index, q_num_blocks, kv_num_blocks, dO):
-    Q = Q.detach().requires_grad_()
-    K = K.detach().requires_grad_()
-    V = V.detach().requires_grad_()
-    
-    q_padded = vsa_pad(Q, q_non_pad_index, q_num_blocks, BLOCK_M)
-    k_padded = vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M)
-    v_padded = vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M)
-    output, _= block_sparse_attn(q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes)
-    output = output[:, :, q_non_pad_index, :]
-    output.backward(dO)
-    return output, Q.grad, K.grad, V.grad
-
-
-def get_non_pad_index(
-    vid_len: torch.LongTensor,
-    n_win: int,
-    win_size: int,
-):
-    device = vid_len.device
-    starts_pad = torch.arange(n_win, device=device) * win_size
-    index_pad = starts_pad[:, None] + torch.arange(win_size, device=device)[None, :]
-    index_mask = torch.arange(win_size, device=device)[None, :] < vid_len[:, None]
-
-    return index_pad[index_mask]
-
-def generate_tensor(shape, dtype, device):
-    tensor = torch.randn(shape, dtype=dtype, device=device)
-    return tensor
-
-def generate_variable_block_sizes(num_blocks, min_size=16, max_size=64, device="cuda"):
-    return torch.randint(min_size, max_size + 1, (num_blocks,), device=device, dtype=torch.int32)
-
-
-def vsa_pad(x, non_pad_index, num_blocks, block_size):
-    padded_x = torch.zeros((1, x.shape[1], num_blocks * BLOCK_M, x.shape[3]), device=x.device, dtype=x.dtype)
-    padded_x[:, :, non_pad_index, :] = x
-    return padded_x
-
-def check_correctness(h, d, num_blocks, k,  num_iterations=20, error_mode='all'):
-    results = {
-        'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-        'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-        'gK': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-        'gV': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-    }
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    variable_block_sizes = generate_variable_block_sizes(num_blocks, device=device)
-    S = int(variable_block_sizes.sum().item())
-    padded_S = num_blocks * BLOCK_M
-    non_pad_index = get_non_pad_index(variable_block_sizes, num_blocks, BLOCK_M)
-    block_mask = generate_block_sparse_mask_for_function(h, num_blocks, num_blocks, k, device)
-    full_mask = create_full_mask_from_block_mask(block_mask, variable_block_sizes, variable_block_sizes, device)
-    for _ in range(num_iterations):
-        Q = generate_tensor((1, h, S, d),  torch.bfloat16, device)
-        K = generate_tensor((1, h, S, d), torch.bfloat16, device)
-        V = generate_tensor((1, h, S, d),  torch.bfloat16, device)
-        dO = generate_tensor((1, h, S, d), torch.bfloat16, device)
-        
-        # print(Q.shape, K.shape, V.shape, dO.shape)
-
-        # dO_padded = torch.zeros_like(dO_padded)
-        # dO_padded[:, :, non_pad_index, :] = dO
-
-        pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO)
-        bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), variable_block_sizes, non_pad_index, non_pad_index, num_blocks, num_blocks, dO)
-        for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), (pt_o, bs_o)]):
-            if bs is not None:
-                diff = pt - bs
-                abs_diff = torch.abs(diff)
-                results[name]['sum_diff'] += torch.sum(abs_diff).item()
-                results[name]['sum_abs'] += torch.sum(torch.abs(pt)).item()
-                rel_max_diff = torch.max(abs_diff) / torch.mean(torch.abs(pt))
-                results[name]['max_diff'] = max(results[name]['max_diff'], rel_max_diff.item())
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    total_elements = h * S * d * num_iterations
-    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 check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, num_iterations=20, error_mode='all'):
-    results = {
-        'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-        'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-        'gK': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-        'gV': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-    }
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    q_variable_block_sizes = generate_variable_block_sizes(num_q_blocks, device=device)
-    kv_variable_block_sizes = generate_variable_block_sizes(num_kv_blocks, device=device)
-    
-    S_q = int(q_variable_block_sizes.sum().item())
-    S_kv = int(kv_variable_block_sizes.sum().item())
-    
-    q_non_pad_index = get_non_pad_index(q_variable_block_sizes, num_q_blocks, BLOCK_M)
-    kv_non_pad_index = get_non_pad_index(kv_variable_block_sizes, num_kv_blocks, BLOCK_M)
-    
-    block_mask = generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, device)
-    full_mask = create_full_mask_from_block_mask(block_mask, q_variable_block_sizes, kv_variable_block_sizes, device)
-    
-    for _ in range(num_iterations):
-        Q = generate_tensor((1, h, S_q, d),  torch.bfloat16, device)
-        K = generate_tensor((1, h, S_kv, d), torch.bfloat16, device)
-        V = generate_tensor((1, h, S_kv, d),  torch.bfloat16, device)
-        dO = generate_tensor((1, h, S_q, d), torch.bfloat16, device)
-        
-        # print(Q.shape, K.shape, V.shape, dO.shape)
-
-        pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO)
-        bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), kv_variable_block_sizes, q_non_pad_index, kv_non_pad_index, num_q_blocks, num_kv_blocks, dO)
-        
-        for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), (pt_o, bs_o)]):
-            if bs is not None:
-                diff = pt - bs
-                abs_diff = torch.abs(diff)
-                results[name]['sum_diff'] += torch.sum(abs_diff).item()
-                results[name]['sum_abs'] += torch.sum(torch.abs(pt)).item()
-                rel_max_diff = torch.max(abs_diff) / torch.mean(torch.abs(pt))
-                results[name]['max_diff'] = max(results[name]['max_diff'], rel_max_diff.item())
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    total_elements_q = h * S_q * d * num_iterations
-    total_elements_kv = h * S_kv * d * num_iterations
-    
-    for name, data in results.items():
-        total_elements = total_elements_q if name in ['gQ', 'gO'] else total_elements_kv
-        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(h, d, error_mode='all'):
-    test_configs = [
-        {"num_blocks": 16, "k": 2, "description": "Small sequence"},
-        {"num_blocks": 32, "k": 4, "description": "Medium sequence"},
-        {"num_blocks": 53, "k": 6, "description": "Large sequence"},
-    ]
-    
-    print(f"\nError Analysis for h={h}, d={d}, mode={error_mode}")
-    print("=" * 150)
-    print(f"{'Config':<20} {'Blocks':<8} {'K':<4} "
-          f"{'gQ Avg':<12} {'Rel gQ Max':<12} "
-          f"{'gK Avg':<12} {'Rel gK Max':<12} "
-          f"{'gV Avg':<12} {'Rel gV Max':<12} "
-          f"{'gO Avg':<12} {'Rel gO Max':<12}")
-    print("-" * 150)
-
-    for config in test_configs:
-        num_blocks = config["num_blocks"]
-        k = config["k"]
-        description = config["description"]
-        results = check_correctness(h, d, num_blocks, k, error_mode=error_mode)
-        print(f"{description:<20} {num_blocks:<8} {k:<4} "
-              f"{results['gQ']['avg_diff']:<12.6e} {results['gQ']['max_diff']:<12.6e} "
-              f"{results['gK']['avg_diff']:<12.6e} {results['gK']['max_diff']:<12.6e} "
-              f"{results['gV']['avg_diff']:<12.6e} {results['gV']['max_diff']:<12.6e} "
-              f"{results['gO']['avg_diff']:<12.6e} {results['gO']['max_diff']:<12.6e}")
-
-    print("-" * 150)
-
-def generate_error_graphs_qkdiff(h, d, error_mode='all'):
-    test_configs = [
-        {"num_q_blocks": 16, "num_kv_blocks": 32, "k": 2, "description": "Small Q, Med KV"},
-        {"num_q_blocks": 32, "num_kv_blocks": 16, "k": 4, "description": "Med Q, Small KV"},
-        {"num_q_blocks": 53, "num_kv_blocks": 32, "k": 6, "description": "Large Q, Med KV"},
-        {"num_q_blocks": 16, "num_kv_blocks": 48, "k": 2, "description": "Small Q, Large KV"},
-        {"num_q_blocks": 48, "num_kv_blocks": 16, "k": 2, "description": "Large Q, Small KV"},
-    ]
-    
-    print(f"\nError Analysis (QK Diff) for h={h}, d={d}, mode={error_mode}")
-    print("=" * 150)
-    print(f"{'Config':<20} {'Q Blks':<8} {'KV Blks':<8} {'K':<4} "
-          f"{'gQ Avg':<12} {'Rel gQ Max':<12} "
-          f"{'gK Avg':<12} {'Rel gK Max':<12} "
-          f"{'gV Avg':<12} {'Rel gV Max':<12} "
-          f"{'gO Avg':<12} {'Rel gO Max':<12}")
-    print("-" * 150)
-
-    for config in test_configs:
-        num_q_blocks = config["num_q_blocks"]
-        num_kv_blocks = config["num_kv_blocks"]
-        k = config["k"]
-        description = config["description"]
-        results = check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, error_mode=error_mode)
-        print(f"{description:<20} {num_q_blocks:<8} {num_kv_blocks:<8} {k:<4} "
-              f"{results['gQ']['avg_diff']:<12.6e} {results['gQ']['max_diff']:<12.6e} "
-              f"{results['gK']['avg_diff']:<12.6e} {results['gK']['max_diff']:<12.6e} "
-              f"{results['gV']['avg_diff']:<12.6e} {results['gV']['max_diff']:<12.6e} "
-              f"{results['gO']['avg_diff']:<12.6e} {results['gO']['max_diff']:<12.6e}")
-
-    print("-" * 150)
-
-if __name__ == "__main__":
-    h, d = 16, 128
-    print("Block Sparse Attention with Variable Block Sizes Analysis")
-    print("=" * 60)
-    for mode in ['backward']:
-        generate_error_graphs(h, d, error_mode=mode)
-        generate_error_graphs_qkdiff(h, d, error_mode=mode)
-    print("\nAnalysis completed for all modes.")

csrc/attn/tests/test_vsa_forward.py

@@ -1,236 +0,0 @@
-import os
-import sys
-from typing import Tuple
-
-import torch
-
-# Make sure we can import from the project root (`vsa`, `tests.utils`, etc.)
-CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
-PROJECT_ROOT = os.path.dirname(CURRENT_DIR)
-if PROJECT_ROOT not in sys.path:
-    sys.path.append(PROJECT_ROOT)
-if CURRENT_DIR not in sys.path:
-    sys.path.append(CURRENT_DIR)
-
-from tests.utils import (
-    generate_block_sparse_mask_for_function,
-    create_full_mask_from_block_mask,
-)
-from vsa import block_sparse_attn, BLOCK_M
-import test_vsa as ref  # reuse helper functions from backward test
-
-
-def pytorch_forward(
-    Q: torch.Tensor,
-    K: torch.Tensor,
-    V: torch.Tensor,
-    block_sparse_mask: torch.Tensor,
-) -> torch.Tensor:
-    """
-    Dense PyTorch reference forward:
-    - Q: [1, h, S_q, d]
-    - K,V: [1, h, S_kv, d]
-    - block_sparse_mask: [h, S_q, S_kv] bool
-    """
-    q = Q.clone().float()
-    k = K.clone().float()
-    v = V.clone().float()
-
-    attn = torch.matmul(q, k.transpose(-2, -1))  # [1, h, S_q, S_kv]
-    attn = attn / (q.size(-1) ** 0.5)
-    attn = attn.masked_fill(~block_sparse_mask.unsqueeze(0), float("-inf"))
-    attn = torch.nn.functional.softmax(attn, dim=-1)
-    out = torch.matmul(attn, v)  # [1, h, S_q, d]
-    return out.to(torch.bfloat16)
-
-
-def block_sparse_forward_test(
-    Q: torch.Tensor,
-    K: torch.Tensor,
-    V: torch.Tensor,
-    block_sparse_mask: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-    q_non_pad_index: torch.Tensor,
-    kv_non_pad_index: torch.Tensor,
-    q_num_blocks: int,
-    kv_num_blocks: int,
-) -> torch.Tensor:
-    """
-    Forward-only wrapper around `block_sparse_attn`, mirroring `block_sparse_kernel_test`
-    but without any backward / grad logic.
-    """
-    Q = Q.detach()
-    K = K.detach()
-    V = V.detach()
-
-    q_padded = ref.vsa_pad(Q, q_non_pad_index, q_num_blocks, BLOCK_M)
-    k_padded = ref.vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M)
-    v_padded = ref.vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M)
-
-    out_padded, _ = block_sparse_attn(
-        q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes
-    )
-    # Remove padding on the query side
-    out = out_padded[:, :, q_non_pad_index, :]
-    return out
-
-
-def run_forward_equal_qk(
-    h: int = 16,
-    d: int = 128,
-    num_blocks: int = 16,
-    k: int = 2,
-    num_iterations: int = 5,
-) -> Tuple[float, float]:
-    """
-    Forward-only correctness test for the case S_q == S_kv.
-    Mirrors `check_correctness` but only compares forward outputs.
-    """
-    assert torch.cuda.is_available(), "VSA kernels require CUDA"
-    device = "cuda"
-
-    variable_block_sizes = ref.generate_variable_block_sizes(
-        num_blocks, device=device
-    )
-    S = int(variable_block_sizes.sum().item())
-    non_pad_index = ref.get_non_pad_index(
-        variable_block_sizes, num_blocks, BLOCK_M
-    )
-
-    block_mask = generate_block_sparse_mask_for_function(
-        h, num_blocks, num_blocks, k, device
-    )
-    full_mask = create_full_mask_from_block_mask(
-        block_mask, variable_block_sizes, variable_block_sizes, device
-    )
-    print(f"[qkequal] h: {h}, d: {d}, num_blocks: {num_blocks}, k: {k}")
-    print(f"[qkequal] variable_block_sizes: {variable_block_sizes}, non_pad_index: {non_pad_index.shape}, block_mask: {block_mask.shape}, full_mask: {full_mask.shape}")
-    sum_diff = 0.0
-    sum_abs = 0.0
-    max_rel_diff = 0.0
-
-    for i in range(num_iterations):
-        Q = ref.generate_tensor((1, h, S, d), torch.bfloat16, device)
-        K = ref.generate_tensor((1, h, S, d), torch.bfloat16, device)
-        V = ref.generate_tensor((1, h, S, d), torch.bfloat16, device)
-
-        if i == 0: print(f"[qkequal] Q: {Q.shape}, K: {K.shape}, V: {V.shape}, full_mask: {full_mask.shape}")
-        if i == 0: print(f"[qkequal] block_mask: {block_mask.shape}")
-
-        pt_o = pytorch_forward(Q, K, V, full_mask)
-        bs_o = block_sparse_forward_test(
-            Q,
-            K,
-            V,
-            block_mask.unsqueeze(0),
-            variable_block_sizes,
-            non_pad_index,
-            non_pad_index,
-            num_blocks,
-            num_blocks,
-        )
-
-        diff = (pt_o - bs_o).abs()
-        sum_diff += diff.sum().item()
-        sum_abs += pt_o.abs().sum().item()
-        rel_max = diff.max() / (pt_o.abs().mean() + 1e-6)
-        max_rel_diff = max(max_rel_diff, rel_max.item())
-
-    total_elems = h * S * d * num_iterations
-    avg_abs_err = sum_diff / total_elems
-    return avg_abs_err, max_rel_diff
-
-
-def run_forward_qk_diff(
-    h: int = 16,
-    d: int = 128,
-    num_q_blocks: int = 16,
-    num_kv_blocks: int = 32,
-    k: int = 2,
-    num_iterations: int = 5,
-) -> Tuple[float, float]:
-    """
-    Forward-only correctness test for the case S_q != S_kv.
-
-    NOTE:
-      - The Triton backend supports different Q/KV logical lengths via padding.
-      - The SM90 (H100) CUDA backend currently assumes the same number of blocks
-        for Q and KV, so we skip this test there.
-    """
-    assert torch.cuda.is_available(), "VSA kernels require CUDA"
-
-    device = "cuda"
-
-    q_variable_block_sizes = ref.generate_variable_block_sizes(
-        num_q_blocks, device=device
-    )
-    kv_variable_block_sizes = ref.generate_variable_block_sizes(
-        num_kv_blocks, device=device
-    )
-
-    S_q = int(q_variable_block_sizes.sum().item())
-    S_kv = int(kv_variable_block_sizes.sum().item())
-
-    q_non_pad_index = ref.get_non_pad_index(
-        q_variable_block_sizes, num_q_blocks, BLOCK_M
-    )
-    kv_non_pad_index = ref.get_non_pad_index(
-        kv_variable_block_sizes, num_kv_blocks, BLOCK_M
-    )
-
-    block_mask = generate_block_sparse_mask_for_function(
-        h, num_q_blocks, num_kv_blocks, k, device
-    )
-    full_mask = create_full_mask_from_block_mask(
-        block_mask, q_variable_block_sizes, kv_variable_block_sizes, device
-    )
-
-    sum_diff = 0.0
-    sum_abs = 0.0
-    max_rel_diff = 0.0
-
-    for _ in range(num_iterations):
-        Q = ref.generate_tensor((1, h, S_q, d), torch.bfloat16, device)
-        K = ref.generate_tensor((1, h, S_kv, d), torch.bfloat16, device)
-        V = ref.generate_tensor((1, h, S_kv, d), torch.bfloat16, device)
-
-        pt_o = pytorch_forward(Q, K, V, full_mask)
-        bs_o = block_sparse_forward_test(
-            Q,
-            K,
-            V,
-            block_mask.unsqueeze(0),
-            kv_variable_block_sizes,
-            q_non_pad_index,
-            kv_non_pad_index,
-            num_q_blocks,
-            num_kv_blocks,
-        )
-
-        diff = (pt_o - bs_o).abs()
-        sum_diff += diff.sum().item()
-        sum_abs += pt_o.abs().sum().item()
-        rel_max = diff.max() / (pt_o.abs().mean() + 1e-6)
-        max_rel_diff = max(max_rel_diff, rel_max.item())
-
-    total_elems = h * S_q * d * num_iterations
-    avg_abs_err = sum_diff / total_elems
-    return avg_abs_err, max_rel_diff
-
-
-if __name__ == "__main__":
-    h, d = 16, 128
-    print("Forward Block Sparse Attention Check (QK Equal)")
-    print("=" * 80)
-    avg_err_eq, max_rel_eq = run_forward_equal_qk(h, d, num_blocks=32, k=2)
-    print(f"QK equal: avg |ΔO| = {avg_err_eq:.6e}, max rel ΔO = {max_rel_eq:.6e}")
-
-    print("\nForward Block Sparse Attention Check (QK Different)")
-    print("=" * 80)
-    avg_err_diff, max_rel_diff = run_forward_qk_diff(
-        h, d, num_q_blocks=32, num_kv_blocks=48, k=2
-    )
-    print(
-        f"QK diff:  avg |ΔO| = {avg_err_diff:.6e}, max rel ΔO = {max_rel_diff:.6e}"
-    )
-

csrc/attn/tests/utils.py

@@ -1,60 +0,0 @@
-import torch
-
-def generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, device="cuda"):
-    """
-    Generate block sparse mask of shape [h, num_q_blocks, num_kv_blocks].
-    
-    Args:
-        h: number of heads
-        num_q_blocks: number of query blocks
-        num_kv_blocks: number of key/value blocks
-        k: number of kv blocks each q block attends to
-        device: device to create tensors on
-
-    Returns:
-        block_sparse_mask: [h, num_q_blocks, num_kv_blocks] bool tensor
-    """
-    k = min(k, num_kv_blocks)
-    scores = torch.rand(h, num_q_blocks, num_kv_blocks, device=device)
-    _, indices = torch.topk(scores, k, dim=-1)
-    block_sparse_mask = torch.zeros(h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
-
-    block_sparse_mask = block_sparse_mask.scatter_(2, indices, 1).bool()
-    return block_sparse_mask
-
-
-def create_full_mask_from_block_mask(block_sparse_mask, q_variable_block_sizes, 
-                                     kv_variable_block_sizes, device="cuda"):
-    """
-    Convert block-level sparse mask to full attention mask.
-    
-    Args:
-        block_sparse_mask: [h, num_q_blocks, num_kv_blocks] bool tensor
-        q_variable_block_sizes: [num_q_blocks] tensor
-        kv_variable_block_sizes: [num_kv_blocks] tensor
-        device: device to create tensors on
-
-    Returns:
-        full_mask: [h, S_q, S_kv] bool tensor where S = total sequence length
-    """
-    h, num_q_blocks, num_kv_blocks = block_sparse_mask.shape
-    total_q_seq_len = q_variable_block_sizes.sum().item()
-    total_kv_seq_len = kv_variable_block_sizes.sum().item()
-    
-    q_cumsum = torch.cat([torch.tensor([0], device=device), q_variable_block_sizes.cumsum(dim=0)[:-1]])
-    kv_cumsum = torch.cat([torch.tensor([0], device=device), kv_variable_block_sizes.cumsum(dim=0)[:-1]])
-
-    full_mask = torch.zeros(h, total_q_seq_len, total_kv_seq_len, dtype=torch.bool, device=device)
-
-    for head in range(h):
-        for q_block in range(num_q_blocks):
-            q_start = q_cumsum[q_block]
-            q_end = q_start + q_variable_block_sizes[q_block]
-
-            for kv_block in range(num_kv_blocks):
-                if block_sparse_mask[head, q_block, kv_block]:
-                    kv_start = kv_cumsum[kv_block]
-                    kv_end = kv_start + kv_variable_block_sizes[kv_block]
-                    full_mask[head, q_start:q_end, kv_start:kv_end] = True
-
-    return full_mask

csrc/attn/video_sparse_attn/MANIFEST.in

@@ -1,2 +0,0 @@
-recursive-include tk *
-include config_vsa.py

csrc/attn/video_sparse_attn/README.md

@@ -1,61 +0,0 @@
-
-
-# Attention Kernel Used in FastVideo
-
-## Video Sparse Attention (VSA)
-
-### Installation
-We support H100 (via TK) and any other GPU (via triton) for VSA.
-
-```bash
-pip install vsa
-``` 
-
-Install from source:
-
-```bash
-git submodule update --init --recursive
-python setup.py install
-```
-
-
-If you encounter error during installation, try below:
-Install C++20 for ThunderKittens:
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-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 LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-```
-
-### Verify if you have successfully installed
-
-```bash
-# test numerical
-python ../tests/test_vsa.py
-# (For H100) test speed
-python ../benchmarks/bench_vsa_hopper.py
-```
-
-bench_vsa_hopper.py should print something like this:
-
-```bash
-Using topk=76 kv blocks per q block (out of 768 total kv blocks)
-
-=== BLOCK SPARSE ATTENTION BENCHMARK ===
-Block Sparse Forward  - TFLOPS: 5622.26
-Block Sparse Backward - TFLOPS: 3865.68
-```
-
-## Acknowledgement
-
-We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/attn/video_sparse_attn/config_vsa.py

@@ -1,15 +0,0 @@
-### ADD TO THIS TO REGISTER NEW KERNELS
-sources = {
-    'block_sparse': {
-        'source_files': {
-            'h100': 'vsa/block_sparse_h100.cu'
-        }
-    }
-}
-
-### WHICH KERNELS DO WE WANT TO BUILD?
-# (oftentimes during development work you don't need to redefine them all.)
-kernels = ['block_sparse']
-
-### WHICH GPU TARGET DO WE WANT TO BUILD FOR?
-target = 'h100'

csrc/attn/video_sparse_attn/setup.py

@@ -1,81 +0,0 @@
-import os
-import subprocess
-
-from config_vsa import kernels, sources, target
-from setuptools import find_packages, setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-target = target.lower()
-
-# Package metadata
-PACKAGE_NAME = "vsa"
-VERSION = "0.0.3"
-AUTHOR = "Hao AI Lab"
-DESCRIPTION = "Video Sparse Attention Kernel Used in FastVideo"
-URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn/video_sparse_attn"
-
-# 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('vsa 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 = ['vsa.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()}')
-
-
-ext_modules = [
-    CUDAExtension('vsa_cuda',
-        sources=source_files,
-        extra_compile_args={
-            'cxx': cpp_flags,
-            'nvcc': cuda_flags
-        },
-        libraries=['cuda'])
-]
-
-
-
-setup(name=PACKAGE_NAME,
-      version=VERSION,
-      author=AUTHOR,
-      description=DESCRIPTION,
-      url=URL,
-      packages=find_packages(),
-      ext_modules=ext_modules,
-      cmdclass={'build_ext': BuildExtension},
-      classifiers=[
-          "Programming Language :: Python :: 3",
-          "Environment :: GPU :: NVIDIA CUDA :: 12",
-          "License :: OSI Approved :: Apache Software License",
-      ],
-      python_requires='>=3.10',
-      install_requires=["torch>=2.5.0"])

csrc/attn/video_sparse_attn/vsa.cpp

@@ -1,27 +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_BLOCK_SPARSE
-extern std::vector<torch::Tensor> block_sparse_attention_forward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v,  torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num, torch::Tensor block_size
-); 
-extern std::vector<torch::Tensor> block_sparse_attention_backward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og, torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num, torch::Tensor block_size
-);
-#endif
-
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.doc() = "Video Sparse Attention Kernels"; // optional module docstring
-
-#ifdef TK_COMPILE_BLOCK_SPARSE
-    m.def("block_sparse_fwd",  torch::wrap_pybind_function(block_sparse_attention_forward), "block sparse attention");
-    m.def("block_sparse_bwd",  torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward");
-#endif
-}

csrc/attn/video_sparse_attn/vsa/__init__.py

@@ -1,80 +0,0 @@
-import torch
-from typing import Tuple
-block_sparse_attn=None
-import torch
-major, minor = torch.cuda.get_device_capability(0)
-if major == 9 and minor == 0:# check if H100
-    from vsa_cuda import block_sparse_fwd, block_sparse_bwd
-    from vsa.block_sparse_wrapper import block_sparse_attn_SM90
-    block_sparse_attn = block_sparse_attn_SM90
-else:
-    from vsa.block_sparse_wrapper import block_sparse_attn_triton
-    block_sparse_fwd = None
-    block_sparse_bwd = None
-    block_sparse_attn = block_sparse_attn_triton
-
-BLOCK_M = 64
-BLOCK_N = 64
-
-
-def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
-    QK = torch.matmul(q, k.transpose(-2, -1))
-    QK /= (q.size(-1)**0.5)
-
-    # Causal mask removed since causal is always false
-
-    QK = torch.nn.functional.softmax(QK, dim=-1)
-    output = torch.matmul(QK, v)
-    return output, QK
-
-
-def video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight=None):
-    """
-    q: [batch_size, num_heads, seq_len, head_dim]
-    k: [batch_size, num_heads, seq_len, head_dim]
-    v: [batch_size, num_heads, seq_len, head_dim]
-    topk: int
-    block_size: int or tuple of 3 ints
-    video_shape: tuple of (T, H, W)
-    compress_attn_weight: [batch_size, num_heads, seq_len, head_dim]
-    select_attn_weight: [batch_size, num_heads, seq_len, head_dim]
-    NOTE: We assume q, k, v is zero padded!!
-    V1 of sparse attention. Include compress attn and sparse attn branch, use average pooling to compress. 
-    Assume q, k, v is flattened in this way: [batch_size, num_heads, T//block_size[0], H//block_size[1], W//block_size[2], block_size[0], block_size[1], block_size[2]]
-    """
-
-    if isinstance(block_size, int):
-        block_size = (block_size, block_size, block_size)
-
-    block_elements = block_size[0] * block_size[1] * block_size[2]
-    assert block_elements == 64
-    assert q.shape[2] % block_elements == 0
-    batch_size, num_heads, seq_len, head_dim = q.shape
-    # compress attn
-    q_compress = (q.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(q.dtype)
-    k_compress = (k.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype)
-    v_compress = (v.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype)
-
-    output_compress, block_attn_score = torch_attention(q_compress, k_compress,
-                                                        v_compress)
-
-    output_compress = output_compress.view(batch_size, num_heads,
-                                           seq_len // block_elements, 1,
-                                           head_dim)
-    output_compress = output_compress.repeat(1, 1, 1, block_elements,
-                                             1).view(batch_size, num_heads,
-                                                     seq_len, head_dim)
-
-    topK_indices = torch.topk(block_attn_score, topk, dim=-1).indices
-    block_mask = torch.zeros_like(block_attn_score, dtype=torch.bool).scatter_(-1, topK_indices, True)
-    output_select, _ = block_sparse_attn(q, k, v, block_mask, variable_block_sizes)
-
-    if compress_attn_weight is not None:
-        final_output = output_compress * compress_attn_weight + output_select
-    else:
-        final_output = output_compress + output_select
-    return final_output
-

csrc/attn/video_sparse_attn/vsa/block_sparse_wrapper.py

@@ -1,185 +0,0 @@
-import torch
-try:
-    from vsa_cuda import block_sparse_fwd, block_sparse_bwd
-except ImportError:
-    block_sparse_fwd = None
-    block_sparse_bwd = None
-from vsa.block_sparse_attn_triton import triton_block_sparse_attn_forward, triton_block_sparse_attn_backward
-assert torch.__version__ >= "2.4.0", "VSA requires PyTorch 2.4.0 or higher"
-from vsa.index import map_to_index
-from typing import Tuple, Optional
-
-
-
-@torch.library.custom_op("vsa::block_sparse_attn_triton", mutates_args=(), device_types="cuda")
-def block_sparse_attn_triton(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    q = q.contiguous()
-    k = k.contiguous()
-    v = v.contiguous()
-    block_map = block_map.int()
-    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
-    o, M = triton_block_sparse_attn_forward(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
-    return o, M
-
-
-@torch.library.register_fake("vsa::block_sparse_attn_triton")
-def _block_sparse_attn_triton_fake(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    q = q.contiguous()
-    k = k.contiguous()
-    v = v.contiguous()
-    o = torch.empty_like(q)
-    M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
-    return o, M
-
-
-@torch.library.custom_op("vsa::block_sparse_attn_backward_triton", mutates_args=(), device_types="cuda")
-def block_sparse_attn_backward_triton(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    M: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    grad_output_padded = grad_output_padded.contiguous()
-    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
-    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
-    dq, dk, dv = triton_block_sparse_attn_backward(grad_output_padded, q_padded, k_padded, v_padded, o_padded, M, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes)
-    return dq, dk, dv
-
-@torch.library.register_fake("vsa::block_sparse_attn_backward_triton")
-def _block_sparse_attn_backward_triton_fake(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    M: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    grad_output_padded = grad_output_padded.contiguous()
-    dq = torch.empty_like(grad_output_padded)
-    dk = torch.empty_like(grad_output_padded)
-    dv = torch.empty_like(grad_output_padded)
-    return dq, dk, dv
-
-
-def backward_triton(ctx, grad_output1, grad_output2):
-    q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes = ctx.saved_tensors
-    dq, dk, dv = block_sparse_attn_backward_triton(grad_output1, q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes)
-    return dq, dk, dv, None, None
-
-
-def setup_context_triton(ctx, inputs, output):
-    q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
-    o_padded, M = output
-    ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes)
-    
-block_sparse_attn_triton.register_autograd(backward_triton, setup_context=setup_context_triton)
-
-
-major, minor = torch.cuda.get_device_capability(0)
-
-if major == 9 and minor == 0:# check if H100
-    @torch.library.custom_op("vsa::block_sparse_attn_SM90", mutates_args=(), device_types="cuda")
-    def block_sparse_attn_SM90(
-        q_padded: torch.Tensor,
-        k_padded: torch.Tensor, 
-        v_padded: torch.Tensor, 
-        block_map: torch.Tensor,
-        variable_block_sizes: torch.Tensor, 
-        )-> Tuple[torch.Tensor, torch.Tensor]:
-        q_padded = q_padded.contiguous()
-        k_padded = k_padded.contiguous()
-        v_padded = v_padded.contiguous()
-        q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
-        variable_block_sizes = variable_block_sizes.int()
-        o_padded, lse_padded = block_sparse_fwd(q_padded, k_padded, v_padded, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
-        return o_padded, lse_padded
-
-
-
-
-    @torch.library.register_fake("vsa::block_sparse_attn_SM90")
-    def _block_sparse_attn_SM90_fake(
-        q_padded: torch.Tensor,
-        k_padded: torch.Tensor, 
-        v_padded: torch.Tensor, 
-        block_map: torch.Tensor,
-        variable_block_sizes: torch.Tensor, 
-    ) -> Tuple[torch.Tensor,  torch.Tensor]:
-        q_padded, k_padded, v_padded = [x.contiguous() for x in (q_padded, k_padded, v_padded)]
-        B, H, S, D = q_padded.shape
-        o_padded = torch.empty_like(q_padded)
-        lse_padded = torch.empty((B, H, S, 1), device=q_padded.device, dtype=torch.float32)
-        return o_padded, lse_padded
-
-
-    @torch.library.custom_op("vsa::block_sparse_attn_backward_SM90", mutates_args=(), device_types="cuda")
-    def block_sparse_attn_backward_SM90(
-        grad_output_padded: torch.Tensor, 
-        q_padded: torch.Tensor, 
-        k_padded: torch.Tensor, 
-        v_padded: torch.Tensor, 
-        o_padded: torch.Tensor, 
-        lse_padded: torch.Tensor, 
-        block_map: torch.Tensor,
-        variable_block_sizes: torch.Tensor, 
-        )-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        grad_output_padded = grad_output_padded.contiguous()
-        k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
-        grad_q_padded, grad_k_padded, grad_v_padded = block_sparse_bwd(
-            q_padded, k_padded, v_padded, o_padded, lse_padded, grad_output_padded, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes
-        )
-        grad_q_padded = grad_q_padded.to(grad_output_padded.dtype)
-        grad_k_padded = grad_k_padded.to(grad_output_padded.dtype)
-        grad_v_padded = grad_v_padded.to(grad_output_padded.dtype)
-        return grad_q_padded, grad_k_padded, grad_v_padded
-
-    @torch.library.register_fake("vsa::block_sparse_attn_backward_SM90")
-    def _block_sparse_attn_backward_SM90_fake(
-        grad_output_padded: torch.Tensor, 
-        q_padded: torch.Tensor, 
-        k_padded: torch.Tensor, 
-        v_padded: torch.Tensor, 
-        o_padded: torch.Tensor, 
-        lse_padded: torch.Tensor, 
-        block_map: torch.Tensor,
-        variable_block_sizes: torch.Tensor, 
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        torch._check(grad_output_padded.dtype == torch.bfloat16)
-        torch._check(lse_padded.dtype == torch.float32)
-        grad_output_padded = grad_output_padded.contiguous()
-        dq = torch.empty_like(grad_output_padded)
-        dk = torch.empty_like(grad_output_padded)
-        dv = torch.empty_like(grad_output_padded)
-        return dq, dk, dv
-
-
-    def backward_SM90(ctx, grad_output1, grad_output2):
-        q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes= ctx.saved_tensors
-        dq, dk, dv = block_sparse_attn_backward_SM90(grad_output1, q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
-        return dq, dk, dv, None, None
-
-    def setup_context_SM90(ctx, inputs, output):
-        q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
-        o_padded, lse_padded = output
-        ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
-
-
-    block_sparse_attn_SM90.register_autograd(backward_SM90, setup_context=setup_context_SM90)

csrc/attn/vmoba_attn/README.md

@@ -1,32 +0,0 @@
-# Attention Kernel Used in FastVideo
-
-## VMoBA: Mixture-of-Block Attention for Video Diffusion Models (VMoBA)
-
-### Installation
-Please ensure that you have installed FlashAttention version **2.7.1 or higher**, as some interfaces have changed in recent releases.
-
-### Usage
-
-You can use `moba_attn_varlen` in the following ways:
-
-**Install from source:**
-```bash
-python setup.py install
-```
-
-**Import after installation:**
-```python
-from vmoba import moba_attn_varlen
-```
-
-**Or import directly from the project root:**
-```python
-from csrc.attn.vmoba_attn.vmoba import moba_attn_varlen
-```
-
-### Verify if you have successfully installed
-
-```bash
-python csrc/attn/vmoba_attn/vmoba/vmoba.py
-```
-

csrc/attn/vmoba_attn/setup.py

@@ -1,26 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from setuptools import find_packages, setup
-
-PACKAGE_NAME = "vmoba"
-VERSION = "0.0.0"
-AUTHOR = "JianzongWu"
-DESCRIPTION = "VMoBA: Mixture-of-Block Attention for Video Diffusion Models"
-URL = "https://github.com/KwaiVGI/VMoBA"
-
-setup(
-    name=PACKAGE_NAME,
-    version=VERSION,
-    author=AUTHOR,
-    description=DESCRIPTION,
-    url=URL,
-    packages=find_packages(),
-    classifiers=[
-        "Programming Language :: Python :: 3",
-        "License :: OSI Approved :: Apache Software License",
-    ],
-    python_requires='>=3.12',
-    install_requires=[
-        "flash-attn >= 2.7.1",
-    ]
-)

csrc/attn/vmoba_attn/vmoba/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from .vmoba import moba_attn_varlen, process_moba_input, process_moba_output

docker/Dockerfile.python3.10

@@ -55,18 +55,12 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install STA (Sliding Tile Attention)
+# Install FastVideo Unified Kernel
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/sliding_tile_attn && \
+    cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    python setup.py install
+    ./build.sh --release
 
-# Install VSA 
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
 
 EXPOSE 22

docker/Dockerfile.python3.11

@@ -55,18 +55,12 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install STA (Sliding Tile Attention)
+# Install FastVideo Unified Kernel
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/sliding_tile_attn && \
+    cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    python setup.py install
+    ./build.sh --release
 
-# Install VSA 
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
 
 EXPOSE 22

docker/Dockerfile.python3.12

@@ -55,18 +55,11 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install STA (Sliding Tile Attention)
+# Install FastVideo Unified Kernel
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/sliding_tile_attn && \
+    cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    python setup.py install
-
-# Install VSA 
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
+    ./build.sh --release
 
 EXPOSE 22

docker/Dockerfile.python3.12.cuda12.9.1

@@ -55,18 +55,12 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install STA (Sliding Tile Attention)
+# Install FastVideo Unified Kernel
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/sliding_tile_attn && \
+    cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    python setup.py install
+    ./build.sh --release
 
-# Install VSA 
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
 
 EXPOSE 22

docker/rocm7_1.Dockerfile

@@ -0,0 +1,58 @@
+FROM rocm/pytorch:rocm7.1_ubuntu22.04_py3.10_pytorch_release_2.9.1
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+SHELL ["/bin/bash", "-c"]
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    zsh \
+    vim \
+    curl \
+    gcc-11 \
+    g++-11 \
+    clang-11 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set up C++20 compilers for ThunderKittens
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+# Install uv and source its environment
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    echo 'source $HOME/.local/bin/env' >> /root/.bashrc
+
+# Copy just the pyproject.toml first to leverage Docker cache
+COPY pyproject_other.toml ./pyproject.toml
+
+# Create a dummy README to satisfy the installation
+RUN echo "# Placeholder" > README.md
+
+# Create and activate virtual environment with specific Python version and seed
+RUN source $HOME/.local/bin/env && \
+    uv venv --python 3.10 --seed /opt/venv && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir --upgrade pip 
+
+COPY . .
+
+# Install dependencies using uv and set up shell configuration
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir -e .[rocm] && \
+    git config --unset-all http.https://github.com/.extraheader || true && \
+    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
+    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+
+# Install FastVideo Unified Kernel
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd fastvideo-kernel && \
+    git submodule update --init --recursive && \
+    ./build.sh --release
+
+EXPOSE 22

docs/README.md

@@ -6,7 +6,7 @@ This directory contains the FastVideo documentation built with MkDocs.
 
 ```bash
 # Install dependencies
-pip install -r docs/requirements-mkdocs.txt
+pip install -r requirements-mkdocs.txt
 
 # Serve docs with live reload (recommended for development)
 mkdocs serve

docs/assets/custom.css

@@ -2,9 +2,11 @@
     writing-mode: sideways-lr;
     white-space: nowrap;
     max-width: 0;
-    p {
-       margin: 0;
-    }
+}
+
+/* Keep header cell paragraph content tight (avoid CSS nesting for compatibility) */
+.vertical-table-header th.head:not(.stub) p {
+    margin: 0;
 }
 
 /* Image sizing classes */

docs/attention/developer/index.md

@@ -0,0 +1,179 @@
+# Adding a New Attention Backend
+
+FastVideo allows integrating new attention mechanisms easily. This guide walks you through adding a new backend (e.g., `MyNewAttn`).
+
+## 1. Implement the Backend (Python)
+
+Create a new file in `fastvideo/attention/backends/` (e.g., `mynew_attn.py`).
+
+Your implementation should inherit from `AttentionBackend` defined in `abstract.py`.
+
+```python
+# fastvideo/attention/backends/mynew_attn.py
+import torch
+from .abstract import AttentionBackend
+# Import the context manager to access metadata (optional)
+from fastvideo.forward_context import get_forward_context
+
+# Import compiled kernel if applicable (see Section 2)
+try:
+    # Import from the top-level package
+    from fastvideo_kernel import my_compiled_attn_func
+except ImportError:
+    my_compiled_attn_func = None
+
+class MyNewAttnBackend(AttentionBackend):
+    def process_inputs(self, q, k, v, **kwargs):
+        # Pre-process inputs if necessary
+        return q, k, v
+
+    def forward(self, q, k, v, **kwargs):
+        # Optional: Access extra metadata passed via ForwardContext
+        # Only needed if your backend requires global state (e.g. window_size)
+        try:
+            context = get_forward_context()
+            metadata = context.attn_metadata
+            # Example: window_size = metadata.window_size
+        except (AssertionError, AttributeError):
+            # Handle case where context is not set (e.g. standard inference)
+            pass
+
+        if my_compiled_attn_func is not None:
+            return my_compiled_attn_func(q, k, v)
+        else:
+            # Fallback implementation (e.g., Triton or pure PyTorch)
+            return self.fallback_impl(q, k, v)
+```
+
+## 2. Passing Extra Information via ForwardContext (Optional)
+
+FastVideo uses a `ForwardContext` to pass global metadata (like current timestep, batch info, or custom attention configurations) to attention backends without changing the `forward` signature of every layer. **This is optional and only required if your backend needs dynamic per-step information.**
+
+To use this:
+1. **Set Context**: In your pipeline or generation loop, use the `set_forward_context` context manager.
+2. **Access Context**: Inside your attention backend, use `get_forward_context()`.
+
+See `docs/attention/sta/index.md` (Sliding Tile Attention) for an example of how complex configuration (window sizes) is passed this way.
+
+## 3. Adding Compiled Kernels (C++/CUDA)
+
+If your backend requires custom CUDA kernels, you need to add them to the `fastvideo-kernel` package.
+
+### A. Add Source Files
+Place your kernel implementation files in `fastvideo-kernel/csrc/attention/`.
+* `mynew_attn.cu` (CUDA implementation)
+* `mynew_attn.h` (Optional headers)
+
+### B. Register in Extension
+Update `fastvideo-kernel/csrc/common_extension.cpp` to expose your function to Python.
+
+```cpp
+// 1. Declare external function
+#ifdef COMPILE_MYNEW_ATTN
+extern torch::Tensor mynew_attn_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v);
+#endif
+
+// 2. Register in module
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    // ... other kernels ...
+
+#ifdef COMPILE_MYNEW_ATTN
+    m.def("mynew_attn_fwd", torch::wrap_pybind_function(mynew_attn_forward), "My New Attention Forward");
+#endif
+}
+```
+
+### C. Update CMakeLists.txt
+Update `fastvideo-kernel/CMakeLists.txt` to compile your new files.
+
+**Case 1: General CUDA Kernel (Runs on all GPUs)**
+Add your source file directly to `EXTENSION_SOURCES` and define the compilation flag.
+
+```cmake
+# Add to EXTENSION_SOURCES
+list(APPEND EXTENSION_SOURCES csrc/attention/mynew_attn.cu)
+
+# Add compilation definition for common_extension.cpp
+list(APPEND COMPILE_DEFS COMPILE_MYNEW_ATTN)
+```
+
+**Case 2: ThunderKittens Kernel (Hopper H100 Only)**
+If your kernel uses ThunderKittens (TK), it requires specific architecture flags (`sm_90a`). Add it inside the `ENABLE_TK_KERNELS` block.
+
+```cmake
+if(ENABLE_TK_KERNELS)
+    # Add source only if TK is enabled
+    list(APPEND EXTENSION_SOURCES csrc/attention/mynew_attn_tk.cu)
+    
+    # Add definition to guard registration
+    list(APPEND COMPILE_DEFS TK_COMPILE_MYNEW_ATTN)
+endif()
+```
+
+### D. Expose in Python Ops
+Update `fastvideo-kernel/python/fastvideo_kernel/ops.py` to make the function importable and handle fallbacks gracefully.
+
+```python
+# fastvideo-kernel/python/fastvideo_kernel/ops.py
+
+# Try to load C++ extension symbols
+try:
+    from fastvideo_kernel._C import fastvideo_kernel_ops
+    mynew_attn_fwd = getattr(fastvideo_kernel_ops, "mynew_attn_fwd", None)
+except ImportError:
+    mynew_attn_fwd = None
+
+def my_compiled_attn_func(q, k, v):
+    # Runtime check: use C++ kernel if available, else fallback
+    if mynew_attn_fwd is not None:
+        return mynew_attn_fwd(q, k, v)
+    else:
+        # Call Triton/Python fallback
+        return mynew_attn_triton(q, k, v)
+```
+
+### E. Expose in Package Init
+Update `fastvideo-kernel/python/fastvideo_kernel/__init__.py` to export the function.
+
+```python
+from fastvideo_kernel.ops import (
+    my_compiled_attn_func,
+    # ...
+)
+
+__all__ = [
+    "my_compiled_attn_func",
+    # ...
+]
+```
+
+## 4. Register the Backend
+
+Update `fastvideo/attention/backends/__init__.py` to export your new class.
+
+```python
+from .mynew_attn import MyNewAttnBackend
+```
+
+## 5. Platform Integration
+
+If your backend requires specific platform checks (e.g., checking for H100 support), handle that in `fastvideo/platforms/cuda.py` or within your backend's `__init__`.
+
+## 6. Add Documentation
+
+Create a new documentation page for your backend to explain its usage, installation (if custom kernels are needed), and features.
+
+1. **Create Directory**: `docs/attention/mynew_attn/`
+2. **Create Index**: `docs/attention/mynew_attn/index.md`
+3. **Update Navigation**: Add an entry to `mkdocs.yml` under the "Attention" tab.
+
+## Checklist
+
+* [ ] Created `fastvideo/attention/backends/mynew_attn.py`.
+* [ ] (Optional) Added CUDA kernels in `fastvideo-kernel/csrc/attention/`.
+* [ ] (Optional) Updated `common_extension.cpp` and `CMakeLists.txt`.
+* [ ] (Optional) Exposed kernel in `fastvideo-kernel/python/fastvideo_kernel/ops.py`.
+* [ ] (Optional) Exported kernel in `fastvideo-kernel/python/fastvideo_kernel/__init__.py`.
+* [ ] Implemented `forward` method respecting the standard signature.
+* [ ] Added unit tests in `tests/`.
+* [ ] Added documentation in `docs/attention/` and updated `mkdocs.yml`.

docs/attention/index.md

@@ -0,0 +1,53 @@
+# FastVideo Attention Kernels
+
+FastVideo provides highly optimized custom attention kernels to accelerate video generation.
+
+## Supported Kernels
+
+* **[Video Sparse Attention (VSA)](vsa/index.md)**: Sparse attention mechanism selecting top-k blocks.
+* **[Sliding Tile Attention (STA)](sta/index.md)**: Optimized attention for window-based video generation.
+
+## General Build Instructions
+
+These instructions apply to building the `fastvideo-kernel` package from source, which includes both STA and VSA kernels.
+
+### Prerequisites
+
+* **PyTorch**: 2.5.0+
+* **CUDA**: 12.4+ (12.8 recommended for best performance)
+* **C++ Compiler**: GCC 11+ (C++20 support required for ThunderKittens)
+
+Install system dependencies:
+
+```bash
+sudo apt update
+sudo apt install -y gcc-11 g++-11 clang-11 ninja-build
+
+# Set gcc-11 as default
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+```
+
+Set up your CUDA environment variables (adjust version as needed):
+
+```bash
+export CUDA_HOME=/usr/local/cuda-12.8
+export PATH=${CUDA_HOME}/bin:${PATH}
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+### Compile and Install
+
+Clone the repository and build the kernel:
+
+```bash
+# Clone recursively to get ThunderKittens submodule
+git clone --recursive https://github.com/hao-ai-lab/FastVideo.git
+cd FastVideo/fastvideo-kernel
+
+# Build and install
+./build.sh
+```
+
+The build script automatically detects your GPU architecture:
+* **H100 (sm_90a)**: Compiles optimized C++ ThunderKittens kernels.
+* **Other (A100, etc.)**: Skips C++ compilation; installs Python package with Triton kernels.

docs/attention/sta/index.md

@@ -0,0 +1,36 @@
+# Sliding Tile Attention (STA)
+
+Optimized attention for window-based video generation (e.g., HunyuanVideo).
+
+## Installation
+
+STA is included in the `fastvideo-kernel` package. See the [main Attention page](../index.md) for build instructions.
+
+## Usage
+
+```python
+from fastvideo_kernel import sliding_tile_attention
+
+# q, k, v: [batch_size, num_heads, seq_length, head_dim]
+# window_size: List of (t, h, w) tiles. Tile size is (6, 8, 8).
+# text_length: Number of text tokens (0-256)
+
+out = sliding_tile_attention(
+    q, k, v, 
+    window_size=[(3, 3, 3)],  # Example window
+    text_length=256
+)
+```
+
+## Citation
+
+If you use Sliding Tile Attention in your research, please cite:
+
+```bibtex
+@article{zhang2025fast,
+  title={Fast video generation with sliding tile attention},
+  author={Zhang, Peiyuan and Chen, Yongqi and Su, Runlong and Ding, Hangliang and Stoica, Ion and Liu, Zhengzhong and Zhang, Hao},
+  journal={arXiv preprint arXiv:2502.04507},
+  year={2025}
+}
+```

docs/attention/vsa/index.md

@@ -0,0 +1,36 @@
+# Video Sparse Attention (VSA)
+
+Sparse attention mechanism selecting top-k blocks.
+
+## Installation
+
+VSA is included in the `fastvideo-kernel` package. See the [main Attention page](../index.md) for build instructions.
+
+## Usage
+
+```python
+from fastvideo_kernel import video_sparse_attn
+
+# q, k, v: [batch_size, num_heads, seq_len, head_dim]
+# variable_block_sizes: Number of valid tokens per block
+# topk: Number of blocks to attend
+
+output = video_sparse_attn(
+    q, k, v, 
+    variable_block_sizes=block_sizes, 
+    topk=32
+)
+```
+
+## Citation
+
+If you use Video Sparse Attention in your research, please cite:
+
+```bibtex
+@article{zhang2025vsa,
+  title={Vsa: Faster video diffusion with trainable sparse attention},
+  author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
+  journal={arXiv preprint arXiv:2505.13389},
+  year={2025}
+}
+```

docs/contributing/overview.md

@@ -6,7 +6,7 @@ Thank you for your interest in contributing to FastVideo. We want to make the pr
 Our community is open to everyone and welcomes any contributions no matter how large or small.
 
 # Developer Environment:
-Do make sure you have CUDA 12.4 installed and supported. FastVideo currently only support Linux and CUDA GPUs, but we hope to support other platforms in the future.
+Do make sure you have CUDA 12.4 installed and supported. FastVideo currently only supports Linux and CUDA GPUs, but we hope to support other platforms in the future.
 
 We recommend using a fresh Python 3.10 Conda environment to develop FastVideo:
 

docs/contributing/profiling.md

@@ -1,7 +1,7 @@
 # Profiling FastVideo
 
 !!! warning
-    Profiling is only intended for FastVideo developers and maintainers to understand the proportion of time spent in different parts of the codebase. **FastVideo end-users should never turn on profiling** as it will significantly slow down the inference.
+    Profiling is only intended for FastVideo developers and maintainers to understand the proportion of time spent in different parts of the codebase. **FastVideo end-users should never turn on profiling** as it will significantly slow down inference.
 
 ## Profiling with PyTorch
 
@@ -49,5 +49,5 @@ Traces can be visualized using <https://ui.perfetto.dev/>.
 ### Best Practices
 
 - Keep the profiled step count small; traces can be large and slow down job shutdown while the profiler flushes data.
-- After profiling, clean up trace directories to avoid filling disks.
+- After profiling, clean up trace directories to avoid filling disk storage.
 - When adding new regions, register them in `fastvideo.profiler` and wrap the corresponding code block with `with self.profiler_controller.region("your_region"):` or the `@profile_region` decorator.

docs/contributing/testing.md

@@ -74,9 +74,11 @@ To add a new SSIM test, follow these steps:
        generator.generate_video(prompt, ...)
 
        # Compare with Reference
-       ssim_values = compute_video_ssim_torchvision(reference_path, generated_path, use_ms_ssim=True)
-        assert ssim_values[0] >= 0.98  # Threshold
-    ```
+       ssim_values = compute_video_ssim_torchvision(
+           reference_path, generated_path, use_ms_ssim=True
+       )
+       assert ssim_values[0] >= 0.98  # Threshold
+   ```
 
 4. **Reference Videos**:
    * When running the test for the first time (or when updating the reference), the test will fail because the reference video is missing. The generated video will be saved in `fastvideo/tests/ssim/generated_videos`.

docs/distillation/dmd.md

@@ -1,6 +1,6 @@
 # 🎯 Distillation
 
-We introduce a new finetuning strategy - **Sparse-distill**, which jointly integrates **[DMD](https://arxiv.org/abs/2405.14867)** and **[VSA](https://arxiv.org/abs/2505.13389)** in a single training process. This approach combines the benefits of both distillation to shorten diffusion steps and sparse attention to reduce attention computations, enabling much faster video generation.
+We introduce a new finetuning strategy - **Sparse-distill**, which jointly integrates **[DMD](https://arxiv.org/abs/2405.14867)** and **[VSA](https://arxiv.org/abs/2505.13389)** in a single training process. This approach combines the benefits of both distillation to shorten diffusion steps and sparse attention to reduce attention computation, enabling much faster video generation.
 
 ## 📊 Model Overview
 

docs/getting_started/quick_start.md

@@ -7,6 +7,11 @@ Get up and running with FastVideo in minutes!
 First, install FastVideo:
 
 ```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
 pip install fastvideo
 ```
 
@@ -15,36 +20,53 @@ pip install fastvideo
 ### Text-to-Video Generation
 
 ```python
-from fastvideo import FastVideoPipeline
+from fastvideo import VideoGenerator
 
-# Initialize the pipeline
-pipe = FastVideoPipeline.from_pretrained("wan2.1-t2v-1.3B")
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
 
-# Generate a video
-prompt = "A cat playing with a ball of yarn"
-video = pipe(prompt, num_frames=16, height=512, width=512)
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
 
-# Save the video
-video.save("output.mp4")
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
 ```
 
 ### Image-to-Video Generation
 
 ```python
-from fastvideo import FastVideoPipeline
-from PIL import Image
+from fastvideo import VideoGenerator, SamplingParam
 
-# Load an image
-image = Image.open("input.jpg")
+def main():
+    # Create the generator
+    model_name = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+    generator = VideoGenerator.from_pretrained(model_name, num_gpus=1)
 
-# Initialize the pipeline
-pipe = FastVideoPipeline.from_pretrained("wan2.1-i2v-14B-480p")
+    # Set up parameters with an initial image
+    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    sampling_param.num_frames = 107
 
-# Generate a video from the image
-video = pipe(image, num_frames=16, height=480, width=480)
+    # Generate video based on the image
+    prompt = "A photograph coming to life with gentle movement"
+    generator.generate_video(prompt, sampling_param=sampling_param,
+                             output_path="my_videos/",
+                             save_video=True)
 
-# Save the video
-video.save("output.mp4")
+if __name__ == '__main__':
+    main()
 ```
 
 ## Next Steps

docs/sliding_tile_attention/demo.md

@@ -1,16 +0,0 @@
-
-# 🔍 Demo
-This is is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
-
-<div style="text-align: center;">
-  <video controls width="800">
-    <source src="https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747" type="video/mp4">
-    Your browser does not support the video tag.
-  </video>
-</div>
-
-You can run STA using the following command:
-
-```bash
-bash scripts/inference/v1_inference_wan_STA.sh
-```

docs/sliding_tile_attention/installation.md

@@ -1,65 +0,0 @@
-
-# 🔧 Installation
-You can install the Sliding Tile Attention package using
-
-```
-pip install st_attn
-``` 
-
-# Building from Source
-We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
-First, install C++20 for ThunderKittens:
-
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-sudo apt install clang-11
-```
-
-Set up CUDA environment (if using CUDA 12.4):
-
-```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
-```
-
-Install STA:
-
-```bash
-cd csrc/attn/sliding_tile_attn/
-git submodule update --init --recursive
-python setup.py install
-```
-
-# 🧪 Test
-
-```bash
-python csrc/attn/tests/test_sta.py
-```
-
-# 📋 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
-# If your attention contains text token (Hunyuan)
-out = sliding_tile_attention(q, k, v, window_size, text_length)
-# If your attention does not contain text token (StepVideo)
-out = sliding_tile_attention(q, k, v, window_size, 0, False)
-
-```
-
-# 🚀Inference
-
-```bash
-bash scripts/inference/v1_inference_wan_STA.sh
-```

docs/training/data_preprocess.md

@@ -30,7 +30,7 @@ path_to_your_dataset_folder/
 └── prompt.txt
 ```
 
-To geranate the `videos2caption.json` and `merge.txt`, run
+To generate the `videos2caption.json` and `merge.txt`, run
 
 ``` python
 python scripts/dataset_preparation/prepare_json_file.py --data_folder mini_i2v_dataset/crush-smol_raw/ --output your_output_folder

docs/video_sparse_attention/installation.md

@@ -1,65 +0,0 @@
-
-# 🔧 Installation
-You can install the Video Sparse Attention package using
-
-```bash
-pip install vsa
-```
-
-# Building from Source
-We support H100 (via ThunderKittens) and any other GPU (via Triton) for VSA.
-
-First, install C++20 for ThunderKittens (if using H100):
-
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-sudo apt install clang-11
-```
-
-Set up CUDA environment (if using CUDA 12.8):
-
-```bash
-export CUDA_HOME=/usr/local/cuda-12.8
-export PATH=${CUDA_HOME}/bin:${PATH} 
-export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-```
-
-Install VSA:
-
-```bash
-cd csrc/attn/video_sparse_attn/
-git submodule update --init --recursive
-python setup.py install
-```
-
-# 🧪 Test
-
-```bash
-python csrc/attn/tests/test_vsa.py
-```
-
-# 📋 Usage
-
-```python
-from vsa import video_sparse_attn
-
-# q, k, v: [batch_size, num_heads, seq_len, head_dim]
-# variable_block_sizes: [num_blocks] - number of valid tokens in each block
-# topk: int - number of top-k blocks to attend to
-# block_size: int or tuple of 3 ints - size of each block (default: 64 tokens)
-# compress_attn_weight: optional weight for compressed attention branch
-
-output = video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight)
-
-``` 
-
-# 🚀Inference
-
-```bash
-bash scripts/inference/v1_inference_wan_VSA.sh
-```

fastvideo-kernel/CMakeLists.txt

@@ -0,0 +1,140 @@
+cmake_minimum_required(VERSION 3.26 FATAL_ERROR)
+project(fastvideo-kernel LANGUAGES CXX CUDA)
+
+# Import common utils if needed, but we keep it simple for now
+
+# Find Python and Torch
+find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)
+
+# Robustly find Torch include paths using Python
+execute_process(
+    COMMAND "${Python_EXECUTABLE}" -c "import torch; from torch.utils.cpp_extension import include_paths; print(';'.join(include_paths()))"
+    OUTPUT_VARIABLE TORCH_INCLUDE_PATHS
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+list(APPEND TORCH_INCLUDE_DIRS ${TORCH_INCLUDE_PATHS})
+
+# Find Torch package (still useful for libraries)
+find_package(Torch REQUIRED)
+
+# Include directories
+include_directories(
+    ${CMAKE_SOURCE_DIR}/include
+    ${CMAKE_SOURCE_DIR}/include/tk/include
+    ${CMAKE_SOURCE_DIR}/include/tk/prototype
+    ${CMAKE_SOURCE_DIR}/csrc
+    ${TORCH_INCLUDE_DIRS}
+)
+
+# ---------------------------
+# ThunderKittens (TK) toggles
+# ---------------------------
+# AUTO: enable TK only when we can confidently target Hopper (sm_90a).
+# ON:   force-enable TK kernels (intended for release wheels/images; does NOT require a GPU).
+# OFF:  never build TK kernels.
+set(FASTVIDEO_KERNEL_BUILD_TK "AUTO" CACHE STRING "Build ThunderKittens kernels: AUTO/ON/OFF")
+set_property(CACHE FASTVIDEO_KERNEL_BUILD_TK PROPERTY STRINGS AUTO ON OFF)
+
+# Prefer environment variable (used by CI) if CMake var is not explicitly set.
+if(NOT DEFINED TORCH_CUDA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST})
+    set(TORCH_CUDA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}")
+endif()
+
+message(STATUS "TORCH_CUDA_ARCH_LIST (cmake/env): ${TORCH_CUDA_ARCH_LIST}")
+message(STATUS "FASTVIDEO_KERNEL_BUILD_TK: ${FASTVIDEO_KERNEL_BUILD_TK}")
+
+set(ENABLE_TK_KERNELS OFF)
+if(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "ON")
+    set(ENABLE_TK_KERNELS ON)
+elseif(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "OFF")
+    set(ENABLE_TK_KERNELS OFF)
+else()
+    # AUTO: detect Hopper if possible.
+    if(TORCH_CUDA_ARCH_LIST)
+        # Accept common spellings: 9.0a, 90a, sm_90a.
+        string(REGEX MATCH "(^|[; ,])((9\\.0a)|(90a)|(sm_90a))([; ,]|$)" _HAS_90A "${TORCH_CUDA_ARCH_LIST}")
+        if(_HAS_90A)
+            set(ENABLE_TK_KERNELS ON)
+        endif()
+    else()
+        # Best-effort local detection (works when a CUDA device is visible).
+        execute_process(
+            COMMAND "${Python_EXECUTABLE}" -c "import torch; import sys; \nprint('1' if (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9) else '0')"
+            OUTPUT_VARIABLE _LOCAL_HAS_HOPPER
+            OUTPUT_STRIP_TRAILING_WHITESPACE
+            ERROR_QUIET
+        )
+        if(_LOCAL_HAS_HOPPER STREQUAL "1")
+            set(ENABLE_TK_KERNELS ON)
+        endif()
+    endif()
+endif()
+
+if(ENABLE_TK_KERNELS)
+    message(STATUS "ThunderKittens kernels: ENABLED")
+else()
+    message(STATUS "ThunderKittens kernels: DISABLED (will use Triton fallbacks at runtime)")
+endif()
+
+# Always try to build the extension if CUDA is available, but conditionally add sources/flags
+set(BUILD_CXX_KERNELS ON)
+
+# Compiler flags
+set(CUDA_FLAGS
+    "-DNDEBUG"
+    "-O3"
+    "-std=c++20"
+    "--use_fast_math"
+    "--expt-extended-lambda"
+    "--expt-relaxed-constexpr"
+    "-Xcompiler=-fno-strict-aliasing"
+    "-Xcompiler=-fPIC"
+    "-DTORCH_COMPILE"
+    "-Xnvlink=--verbose"
+    "-Xptxas=--verbose"
+    "-Xptxas=--warn-on-spills"
+)
+
+# If TK is enabled, ensure we target Hopper. This is required even on GPU-less builders (CI).
+if(ENABLE_TK_KERNELS)
+    if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES OR CMAKE_CUDA_ARCHITECTURES STREQUAL "")
+        set(CMAKE_CUDA_ARCHITECTURES "90a" CACHE STRING "CUDA architectures" FORCE)
+    endif()
+    list(APPEND CUDA_FLAGS "-DKITTENS_HOPPER")
+    message(STATUS "CMAKE_CUDA_ARCHITECTURES: ${CMAKE_CUDA_ARCHITECTURES}")
+endif()
+
+if(BUILD_CXX_KERNELS)
+    # Source files
+    set(EXTENSION_SOURCES csrc/common_extension.cpp)
+    
+    # Conditionally add TK kernels
+    if(ENABLE_TK_KERNELS)
+        list(APPEND EXTENSION_SOURCES 
+            csrc/attention/st_attn_h100.cu
+            csrc/attention/block_sparse_h100.cu
+        )
+    endif()
+
+    # Combined FastVideo Extension
+    # Using name 'fastvideo_kernel_ops' to distinguish from the python package namespace
+    Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI
+        ${EXTENSION_SOURCES}
+    )
+    
+    # Build compile definitions list
+    set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops)
+    if(ENABLE_TK_KERNELS)
+        list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE)
+    endif()
+
+    target_compile_definitions(fastvideo_kernel_ops PRIVATE ${COMPILE_DEFS})
+
+    target_compile_options(fastvideo_kernel_ops PRIVATE 
+        $<$<COMPILE_LANGUAGE:CUDA>:${CUDA_FLAGS}>
+    )
+    
+    # We install it to fastvideo_kernel/_C so we can load it to register the ops
+    install(TARGETS fastvideo_kernel_ops LIBRARY DESTINATION fastvideo_kernel/_C)
+endif()
+

fastvideo-kernel/LICENSE

@@ -0,0 +1,187 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.

fastvideo-kernel/MANIFEST.in

@@ -0,0 +1,6 @@
+include LICENSE
+include README.md
+include pyproject.toml
+recursive-include python/fastvideo_kernel *.py
+recursive-include csrc *.cu *.cuh *.cpp *.h
+recursive-include include/tk *.cu *.cuh *.cpp *.h *.src

fastvideo-kernel/README.md

@@ -0,0 +1,51 @@
+# FastVideo Kernel
+
+CUDA kernels for FastVideo video generation.
+
+## Installation
+
+### Standard Installation (Local Development)
+This will automatically detect your GPU architecture. If an NVIDIA Hopper (H100/sm_90a) GPU is detected, ThunderKittens kernels will be enabled. Otherwise, they will be skipped, and the package will use Triton fallbacks at runtime.
+
+```bash
+git submodule update --init --recursive
+cd fastvideo-kernel
+./build.sh
+```
+
+### Release Build (Force Enable Kernels)
+If you are building a release wheel or docker image on a machine without a GPU (e.g., CI/CD), you can force-enable the compilation of Hopper-specific ThunderKittens kernels.
+
+```bash
+cd fastvideo-kernel
+./build.sh --release
+```
+*Note: The resulting wheel will contain kernels that require an H100 GPU to run, but can be built on any machine with CUDA 12.3+ toolchain.*
+
+## Usage
+
+```python
+from fastvideo_kernel import sliding_tile_attention, video_sparse_attn, moba_attn_varlen
+
+# Example: Sliding Tile Attention
+out = sliding_tile_attention(q, k, v, window_sizes, text_len)
+
+# Example: Video Sparse Attention (with Triton fallback)
+out = video_sparse_attn(q, k, v, block_sizes, topk=5)
+
+# Example: VMoBA
+out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)
+```
+
+## Requirements
+
+- **Runtime**:
+  - NVIDIA H100 (sm_90a) for C++ optimized kernels.
+  - Any CUDA GPU for Triton-based fallbacks.
+- **Build**:
+  - CUDA Toolkit 12.3+
+  - C++20 compatible compiler (GCC 10+, Clang 11+)
+
+## Acknowledgement
+
+This package structure and build system are based on [sgl-kernel](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) from the SGLang project.

fastvideo-kernel/build.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+set -ex
+
+# Simple build script wrapping uv/pip
+# Usage:
+#   ./build.sh                # local dev build (auto-detect / skip TK kernels when not available)
+#   ./build.sh --release      # force-enable Hopper/TK kernels for release builds (no GPU required)
+
+echo "Building fastvideo-kernel..."
+
+# Ensure submodules are initialized if needed (tk)
+git submodule update --init --recursive
+
+# Install build dependencies
+pip install scikit-build-core cmake ninja
+
+RELEASE=0
+if [ "${1:-}" = "--release" ] || [ "${1:-}" = "-r" ]; then
+    RELEASE=1
+fi
+
+if [ "$RELEASE" -eq 1 ]; then
+    # Force-enable ThunderKittens kernels and compile for Hopper.
+    # Intended for producing release wheels/images on machines without a GPU.
+    export TORCH_CUDA_ARCH_LIST="9.0a"
+    export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
+fi
+
+echo "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST:-<unset>}"
+echo "CMAKE_ARGS: ${CMAKE_ARGS:-<unset>}"
+# Build and install
+# Use -v for verbose output
+pip install . -v --no-build-isolation

fastvideo-kernel/csrc/attention/block_sparse_h100.cu

@@ -639,8 +639,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
 
     // store kq and vq
 
-    // ! the following two line seems unnecessary.
-    // tma::store_async_wait(); // ensure qg is finished
+    // ensuring all writes are finished
     __syncthreads();
 
     warpgroup::store(kg_smem[0], kg_reg);
@@ -661,6 +660,145 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
     tma::store_async_wait(); 
 }
 
+
+template<int D>
+void block_sparse_attention_forward_impl(
+    bf16* d_q, bf16* d_k, bf16* d_v, float* d_l, bf16* d_o,
+    int batch, int qo_heads, int kv_heads, int seq_len, int hr,
+    int max_kv_blocks_per_q, 
+    int32_t* q2k_block_sparse_index_ptr,
+    int32_t* q2k_block_sparse_num_ptr,
+    int32_t* block_size_ptr,
+    cudaStream_t stream
+) {
+    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>;
+
+    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<D>;
+
+    q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    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), static_cast<uint32_t>(D)};
+
+    globals g{
+        qg_arg, kg_arg, vg_arg, lg_arg, og_arg, 
+        static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), 
+        q2k_block_sparse_index_ptr, q2k_block_sparse_num_ptr, block_size_ptr
+    };
+
+    // Shared memory size for the kernel
+    // 54000 bytes is calibrated for H100 shared memory constraints for these tile sizes
+    constexpr int mem_size = 54000;
+
+    dim3 grid(seq_len/(64), qo_heads, batch);
+
+    cudaFuncSetAttribute(
+        fwd_attend_ker<D>,
+        cudaFuncAttributeMaxDynamicSharedMemorySize,
+        mem_size
+    );
+
+    fwd_attend_ker<D><<<grid, (128), mem_size, stream>>>(g);
+}
+
+template<int D>
+void block_sparse_attention_backward_impl(
+    bf16* d_q, bf16* d_k, bf16* d_v, bf16* d_o, bf16* d_og, float* d_l, float* d_d, float* d_qg, float* d_kg, float* d_vg,
+    int batch, int qo_heads, int kv_heads, int seq_len, int hr, int max_q_blocks_per_kv,
+    int32_t* k2q_block_sparse_index_ptr,
+    int32_t* k2q_block_sparse_num_ptr,
+    int32_t* block_size_ptr,
+    cudaStream_t stream
+) {
+    using G = bwd_attend_ker_tile_dims<D>;
+    using og_tile = st_bf<4*16, D>;
+    using o_tile  = st_bf<4*16, D>;
+    using d_tile  = col_vec<st_fl<4*16, D>>;
+
+    using og_global = gl<bf16,  -1, -1, -1, -1, og_tile>;
+    using o_global  = gl<bf16,  -1, -1, -1, -1, o_tile>;
+    using d_global  = gl<float, -1, -1, -1, -1, d_tile>;
+
+    using prep_globals = bwd_prep_globals<D>;
+
+    constexpr int mem_size_prep = kittens::MAX_SHARED_MEMORY;
+    int threads_prep = PREP_NUM_WARPS * kittens::WARP_THREADS;
+    dim3 grid_bwd_prep(seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
+
+    og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
+
+    prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg};
+
+    cudaFuncSetAttribute(
+        bwd_attend_prep_ker<D>,
+        cudaFuncAttributeMaxDynamicSharedMemorySize,
+        mem_size_prep
+    );
+    bwd_attend_prep_ker<D><<<grid_bwd_prep, threads_prep, mem_size_prep, stream>>>(bwd_g);
+
+    using bwd_q_tile    =         st_bf<G::tile_h_qo, G::tile_width>;
+    using bwd_k_tile    =         st_bf<G::tile_h,    G::tile_width>;
+    using bwd_v_tile    =         st_bf<G::tile_h,    G::tile_width>;
+    using bwd_og_tile   =         st_bf<G::tile_h_qo, G::tile_width>;
+    using bwd_qg_tile   =         st_fl<G::tile_h_qo, G::tile_width>;
+    using bwd_kg_tile   =         st_fl<G::tile_h,    G::tile_width>;
+    using bwd_vg_tile   =         st_fl<G::tile_h,    G::tile_width>;
+    using bwd_l_tile    = row_vec<st_fl<G::tile_h_qo, G::tile_h>>;
+    using bwd_d_tile    = row_vec<st_fl<G::tile_h_qo, G::tile_h>>;
+
+    using bwd_q_global  = gl<bf16,  -1, -1, -1, -1, bwd_q_tile>;
+    using bwd_k_global  = gl<bf16,  -1, -1, -1, -1, bwd_k_tile>;
+    using bwd_v_global  = gl<bf16,  -1, -1, -1, -1, bwd_v_tile>;
+    using bwd_og_global = gl<bf16,  -1, -1, -1, -1, bwd_og_tile>;
+    using bwd_qg_global = gl<float, -1, -1, -1, -1, bwd_qg_tile>;
+    using bwd_kg_global = gl<float, -1, -1, -1, -1, bwd_kg_tile>;
+    using bwd_vg_global = gl<float, -1, -1, -1, -1, bwd_vg_tile>;
+    using bwd_l_global  = gl<float, -1, -1, -1, -1, bwd_l_tile>;
+    using bwd_d_global  = gl<float, -1, -1, -1, -1, bwd_d_tile>;
+
+    using bwd_global_args = bwd_globals<D>;
+
+    bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
+    bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
+    bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
+
+    bwd_global_args bwd_global{bwd_q_arg, bwd_k_arg, bwd_v_arg, bwd_og_arg, bwd_qg_arg, bwd_kg_arg, bwd_vg_arg, bwd_l_arg, bwd_d_arg, 
+                        static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_q_blocks_per_kv),
+                        k2q_block_sparse_index_ptr, k2q_block_sparse_num_ptr, block_size_ptr};
+
+    dim3 grid_bwd_main(seq_len/64, qo_heads, batch);
+    int threads_main = 128;
+    // Calibrated shared memory sizes for different head dimensions
+    int bwd_mem_size = (D == 64) ? 72000 : 113000;
+
+    cudaFuncSetAttribute(
+        bwd_attend_ker<D>,
+        cudaFuncAttributeMaxDynamicSharedMemorySize,
+        bwd_mem_size
+    );
+    bwd_attend_ker<D><<<grid_bwd_main, threads_main, bwd_mem_size, stream>>>(bwd_global);
+}
+
 #include "pyutils/torch_helpers.cuh"
 #include <ATen/cuda/CUDAContext.h>
 #include <iostream>
@@ -672,32 +810,23 @@ block_sparse_attention_forward(
     torch::Tensor v, 
     torch::Tensor q2k_block_sparse_index, 
     torch::Tensor q2k_block_sparse_num, 
-    torch::Tensor kv_block_size
+    torch::Tensor block_size
 )
 {
     CHECK_INPUT(q);
     CHECK_INPUT(k);
     CHECK_INPUT(v);
 
-    // q shape: (batch, qo_heads, q_seq_len, head_dim)
-    // k shape: (batch, kv_heads, kv_seq_len, head_dim)
-    // v shape: (batch, kv_heads, kv_seq_len, head_dim)
-    // q2k_block_sparse_index shape: (batch, qo_heads, num_q_blocks, max_kv_blocks_per_q)
-    // q2k_block_sparse_num shape: (batch, qo_heads, num_q_blocks)
-    // kv_block_size shape: (num_kv_blocks) This does not need other dimensions because across all batch/heads the padding is the same.
-    
     auto batch    = q.size(0);
-    auto q_seq_len  = q.size(2);
-    auto kv_seq_len = k.size(2);
+    auto seq_len  = q.size(2); 
     auto head_dim = q.size(3); 
     auto qo_heads = q.size(1);
     auto kv_heads = k.size(1);
     auto max_kv_blocks_per_q = q2k_block_sparse_index.size(3);
-    auto num_q_blocks = q2k_block_sparse_index.size(2);
-    auto num_kv_blocks = kv_block_size.size(0);
+    auto num_q_blocks = block_size.size(0);
     TORCH_CHECK(batch==1, "Batch size dim will be removed in the future, please set batch to 1");
-    TORCH_CHECK(num_q_blocks * BLOCK_M == q_seq_len, "This kernel supports variable q block size, but it assumes the input sequence is properly padded.");
-    TORCH_CHECK(num_kv_blocks * BLOCK_M == kv_seq_len, "This kernel supports variable kv block size, but it assumes the input sequence is properly padded.");
+    TORCH_CHECK(num_q_blocks * 64 == seq_len, "This kernel supports variable block size, but it assumes the input sequence is properly padded.");
+    TORCH_CHECK(num_q_blocks == q2k_block_sparse_index.size(2), "Number of Q blocks does not match between q2k_block_sparse_index and block_size");
     // 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");
@@ -705,9 +834,11 @@ block_sparse_attention_forward(
     TORCH_CHECK(q2k_block_sparse_index.size(0) == batch, "q2k_block_sparse_index batch dimension - idx 0 - must match for all inputs");
     TORCH_CHECK(q2k_block_sparse_num.size(0) == batch, "q2k_block_sparse_num batch dimension - idx 0 - must match for all inputs");
 
-    TORCH_CHECK(v.size(2) == kv_seq_len, "V sequence length dimension - idx 2 - must match K inputs");
-    TORCH_CHECK(q2k_block_sparse_num.size(2) == num_q_blocks, "q2k_block_sparse_num idx 2 - must match num_q_blocks");
-
+    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(q2k_block_sparse_index.size(2) == seq_len / BLOCK_M, "q2k_block_sparse_index idx 2 - must match seq_len / BLOCK_M");
+    TORCH_CHECK(q2k_block_sparse_num.size(2) == seq_len / BLOCK_M, "q2k_block_sparse_num idx 2 - must match seq_len / BLOCK_M");
 
     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");
@@ -733,12 +864,12 @@ block_sparse_attention_forward(
     // for the returned outputs
     torch::Tensor o     = torch::empty({static_cast<const uint>(batch), 
                                         static_cast<const uint>(qo_heads), 
-                                        static_cast<const uint>(q_seq_len), 
+                                        static_cast<const uint>(seq_len), 
                                         static_cast<const uint>(head_dim)}, v.options());
     
     torch::Tensor l_vec = torch::empty({static_cast<const uint>(batch), 
                                         static_cast<const uint>(qo_heads), 
-                                        static_cast<const uint>(q_seq_len), 
+                                        static_cast<const uint>(seq_len), 
                                         static_cast<const uint>(1)}, 
                                         torch::TensorOptions().dtype(torch::kFloat).device(q.device()).memory_format(at::MemoryFormat::Contiguous));
         
@@ -749,110 +880,32 @@ block_sparse_attention_forward(
     float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
     float* d_l   = reinterpret_cast<float*>(l_ptr);
 
-    //cudadevicesynchronize();
     const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
     const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
 
+    // Temporated implementation to avoid code duplication between head_dim=64 and 128
     if (head_dim == 64) {
-        using q_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>;
-        using k_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>;
-        using v_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>;
-        using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>>;
-        using o_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::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<64>;
-
-        q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
-        k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
-        v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
-        l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,   static_cast<unsigned int>(q_seq_len)};
-        o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
-
-        globals g{
-            qg_arg, 
-            kg_arg, 
-            vg_arg, 
-            lg_arg, 
-            og_arg, 
-            static_cast<int>(q_seq_len), 
-            static_cast<int>(hr), 
-            static_cast<int>(max_kv_blocks_per_q), 
-            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), 
+        block_sparse_attention_forward_impl<64>(
+            d_q, d_k, d_v, d_l, d_o,
+            batch, qo_heads, kv_heads, seq_len, hr,
+            max_kv_blocks_per_q,
+            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()),
             reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr()),
-            reinterpret_cast<int32_t*>(kv_block_size.data_ptr())
-        };
-
-        constexpr int mem_size = 54000;
-
-        dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch);
-
-        cudaFuncSetAttribute(
-            fwd_attend_ker<64>,
-            cudaFuncAttributeMaxDynamicSharedMemorySize,
-            mem_size
+            reinterpret_cast<int32_t*>(block_size.data_ptr()),
+            stream
         );
-
-        fwd_attend_ker<64><<<grid, (128), mem_size, stream>>>(g);
-
-        CHECK_CUDA_ERROR(cudaGetLastError());
-        // cudaStreamSynchronize(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>(q_seq_len), 128U};
-        k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
-        v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_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>(q_seq_len)};
-        o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
-
-        globals g{
-            qg_arg, 
-            kg_arg, 
-            vg_arg, 
-            lg_arg, 
-            og_arg, 
-            static_cast<int>(q_seq_len), 
-            static_cast<int>(hr), 
-            static_cast<int>(max_kv_blocks_per_q), 
-            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), 
+    } else if (head_dim == 128) {
+        block_sparse_attention_forward_impl<128>(
+            d_q, d_k, d_v, d_l, d_o,
+            batch, qo_heads, kv_heads, seq_len, hr,
+            max_kv_blocks_per_q,
+            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()),
             reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr()),
-            reinterpret_cast<int32_t*>(kv_block_size.data_ptr())
-        };
-
-        constexpr int mem_size = 54000;
-
-        dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch);
-
-        cudaFuncSetAttribute(
-            fwd_attend_ker<128>,
-            cudaFuncAttributeMaxDynamicSharedMemorySize,
-            mem_size
+            reinterpret_cast<int32_t*>(block_size.data_ptr()),
+            stream
         );
-
-        fwd_attend_ker<128><<<grid, (128), mem_size, stream>>>(g);
-
-        CHECK_CUDA_ERROR(cudaGetLastError());
-        // cudaStreamSynchronize(stream);
+    } else {
+        TORCH_CHECK(false, "Unsupported head_dim: ", head_dim, ". Only 64 and 128 are supported.");
     }
 
     return {o, l_vec};
@@ -868,7 +921,7 @@ block_sparse_attention_backward(torch::Tensor q,
                    torch::Tensor og,
                    torch::Tensor k2q_block_sparse_index,
                    torch::Tensor k2q_block_sparse_num,
-                   torch::Tensor kv_block_size)
+                   torch::Tensor block_size)
 {
     CHECK_INPUT(q);
     CHECK_INPUT(k);
@@ -877,23 +930,11 @@ block_sparse_attention_backward(torch::Tensor q,
     CHECK_INPUT(o);
     CHECK_INPUT(og);
 
-    // q: [batch, qo_heads, q_seq_len, head_dim]
-    // k: [batch, kv_heads, kv_seq_len, head_dim]
-    // v: [batch, kv_heads, kv_seq_len, head_dim]
-    // o: [batch, qo_heads, q_seq_len, head_dim]
-    // l_vec: [batch, qo_heads, q_seq_len, 1]
-    // og: [batch, qo_heads, q_seq_len, head_dim]
-    // k2q_block_sparse_index: [batch, kv_heads, num_kv_blocks, max_num_q_blocks]
-    // k2q_block_sparse_num: [batch, kv_heads, num_kv_blocks]
-    // kv_block_size: [num_kv_blocks]
-
     auto batch    = q.size(0);
-    auto q_seq_len  = q.size(2);
-    auto kv_seq_len = k.size(2);
+    auto seq_len  = q.size(2);
     auto head_dim = q.size(3);
     auto max_q_blocks_per_kv = k2q_block_sparse_index.size(3);
-    auto num_kv_blocks = kv_block_size.size(0);
-    TORCH_CHECK(k2q_block_sparse_index.size(2) == num_kv_blocks, "k2q_block_sparse_index.size(2) must match num_kv_blocks (kv_block_size.size(0))");
+    TORCH_CHECK(k2q_block_sparse_index.size(2) == block_size.size(0), "k2q_block_sparse_index.size(2) must match block_size.size(0)");
     // 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");
@@ -904,18 +945,23 @@ block_sparse_attention_backward(torch::Tensor q,
     TORCH_CHECK(k2q_block_sparse_index.size(0) == batch, "k2q_block_sparse_index batch dimension - idx 0 - must match for all inputs");
     TORCH_CHECK(k2q_block_sparse_num.size(0) == batch, "k2q_block_sparse_num batch dimension - idx 0 - must match for all inputs");
 
-    TORCH_CHECK(v.size(2)     == kv_seq_len, "V  sequence length dimension - idx 2 - must match K sequence length");
-    TORCH_CHECK(l_vec.size(2) == q_seq_len, "L  sequence length dimension - idx 2 - must match Q sequence length");
-    TORCH_CHECK(o.size(2)     == q_seq_len, "O  sequence length dimension - idx 2 - must match Q sequence length");
-    TORCH_CHECK(og.size(2)    == q_seq_len, "OG sequence length dimension - idx 2 - must match Q sequence length");
-    TORCH_CHECK(k2q_block_sparse_index.size(2) == num_kv_blocks, "k2q_block_sparse_index idx 2 - must match num_kv_blocks (kv_block_size.size(0))");
-    TORCH_CHECK(k2q_block_sparse_num.size(2) == num_kv_blocks, "k2q_block_sparse_num idx 2 - must match num_kv_blocks (kv_block_size.size(0))");
+
+    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(l_vec.size(2) == seq_len, "L  sequence length dimension - idx 2 - must match for all inputs");
+    TORCH_CHECK(o.size(2)     == seq_len, "O  sequence length dimension - idx 2 - must match for all inputs");
+    TORCH_CHECK(og.size(2)    == seq_len, "OG sequence length dimension - idx 2 - must match for all inputs");
+    TORCH_CHECK(k2q_block_sparse_index.size(2) == seq_len / BLOCK_N, "k2q_block_sparse_index idx 2 - must match seq_len / BLOCK_N");
+    TORCH_CHECK(k2q_block_sparse_num.size(2) == seq_len / BLOCK_N, "k2q_block_sparse_num idx 2 - must match seq_len / BLOCK_N");
 
     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(o.size(3)  == head_dim, "O  head dimension - idx 3 - must match for all non-vector inputs");
     TORCH_CHECK(og.size(3) == head_dim, "OG head dimension - idx 3 - must match for all non-vector inputs");
+    
+
 
     auto qo_heads = q.size(1);
     auto kv_heads = k.size(1);
@@ -942,20 +988,20 @@ block_sparse_attention_backward(torch::Tensor q,
 
     torch::Tensor qg = torch::zeros({static_cast<const uint>(batch), 
                                      static_cast<const uint>(qo_heads), 
-                                     static_cast<const uint>(q_seq_len), 
+                                     static_cast<const uint>(seq_len), 
                                      static_cast<const uint>(head_dim)},   l_vec.options());
     torch::Tensor kg = torch::zeros({static_cast<const uint>(batch), 
                                      static_cast<const uint>(kv_heads), 
-                                     static_cast<const uint>(kv_seq_len), 
+                                     static_cast<const uint>(seq_len), 
                                      static_cast<const uint>(head_dim)},   l_vec.options());
     torch::Tensor vg = torch::zeros({static_cast<const uint>(batch), 
                                      static_cast<const uint>(kv_heads), 
-                                     static_cast<const uint>(kv_seq_len), 
+                                     static_cast<const uint>(seq_len), 
                                      static_cast<const uint>(head_dim)},   l_vec.options());
     
     torch::Tensor d_vec = torch::empty({static_cast<const uint>(batch), 
                                         static_cast<const uint>(qo_heads), 
-                                        static_cast<const uint>(q_seq_len), 
+                                        static_cast<const uint>(seq_len), 
                                         static_cast<const uint>(1)},       l_vec.options());
 
     float*         qg_ptr = qg.data_ptr<float>();
@@ -984,7 +1030,7 @@ block_sparse_attention_backward(torch::Tensor q,
     //  cudaStreamSynchronize(stream);
 
     // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256");
-    dim3 grid_bwd(q_seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
+    dim3 grid_bwd(seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
 
     if (head_dim == 64)  {
         using og_tile = st_bf<4*16, 64>;
@@ -997,9 +1043,9 @@ block_sparse_attention_backward(torch::Tensor q,
 
         using bwd_prep_globals = bwd_prep_globals<64>;
 
-        og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
-        o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
-        d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
+        og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
+        o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
+        d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
 
         bwd_prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg};
 
@@ -1036,15 +1082,15 @@ block_sparse_attention_backward(torch::Tensor q,
 
         using bwd_global_args = bwd_globals<64>;
 
-        bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
-        bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
-        bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
-        bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
-        bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
-        bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
-        bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
-        bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
-        bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
+        bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
+        bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
+        bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
+        bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
+        bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
+        bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
+        bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
+        bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
+        bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
 
         bwd_global_args bwd_global{bwd_q_arg, 
                         bwd_k_arg, 
@@ -1055,14 +1101,14 @@ block_sparse_attention_backward(torch::Tensor q,
                         bwd_vg_arg, 
                         bwd_l_arg, 
                         bwd_d_arg, 
-                        static_cast<int>(kv_seq_len), // N is not used in the kernel
+                        static_cast<int>(seq_len), 
                         static_cast<int>(hr),
                         static_cast<int>(max_q_blocks_per_kv),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_index.data_ptr()),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr()),
-                        reinterpret_cast<int32_t*>(kv_block_size.data_ptr())};
+                        reinterpret_cast<int32_t*>(block_size.data_ptr())};
 
-        dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch);
+        dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
         threads = 128;
 
         //cudadevicesynchronize();
@@ -1101,9 +1147,9 @@ block_sparse_attention_backward(torch::Tensor q,
 
         using bwd_prep_globals = bwd_prep_globals<128>;
 
-        og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
-        o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
-        d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
+        og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+        o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+        d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
 
         bwd_prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg};
 
@@ -1140,15 +1186,15 @@ block_sparse_attention_backward(torch::Tensor q,
 
         using bwd_global_args = bwd_globals<128>;
 
-        bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
-        bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
-        bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
-        bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
-        bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
-        bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
-        bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
-        bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
-        bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
+        bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+        bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+        bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+        bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
+        bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
 
         bwd_global_args bwd_global{bwd_q_arg, 
                         bwd_k_arg, 
@@ -1159,14 +1205,14 @@ block_sparse_attention_backward(torch::Tensor q,
                         bwd_vg_arg, 
                         bwd_l_arg, 
                         bwd_d_arg, 
-                        static_cast<int>(kv_seq_len), // N is not used in the kernel
+                        static_cast<int>(seq_len), 
                         static_cast<int>(hr),
                         static_cast<int>(max_q_blocks_per_kv),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_index.data_ptr()),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr()),
-                        reinterpret_cast<int32_t*>(kv_block_size.data_ptr())};
+                        reinterpret_cast<int32_t*>(block_size.data_ptr())};
 
-        dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch);
+        dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
         threads = 128;
 
         //cudadevicesynchronize();

fastvideo-kernel/csrc/attention/st_attn_h100.cu

@@ -451,6 +451,10 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
 
         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)};
 
+        // Shared memory size for the kernel.
+        // We use the maximum available shared memory (kittens::MAX_SHARED_MEMORY) 
+        // which is approximately 227KB on H100, necessary for the high-performance 
+        // TMA-based attention tiles with multiple stages.
         constexpr int mem_size = kittens::MAX_SHARED_MEMORY;
         int threads = NUM_WORKERS * kittens::WARP_THREADS;
         if (has_text) {
@@ -458,104 +462,31 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
             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) {
+#define LAUNCH_IMAGE_KER(DT_VAL, DH_VAL, DW_VAL) \
+                    cudaFuncSetAttribute( \
+                        fwd_attend_ker<128, false, false, true, DT_VAL, DH_VAL, DW_VAL, 5, 6, 10>, \
+                        cudaFuncAttributeMaxDynamicSharedMemorySize, \
+                        mem_size \
+                    ); \
+                    fwd_attend_ker<128, false, false, true, DT_VAL, DH_VAL, DW_VAL, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, true, 1, 1, 1, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 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, true, 1, 1, 2, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true,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, true, 2, 1, 1, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 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, true, 1, 2, 2, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 1, 2, 2, 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, true, 2, 3, 0, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 2, 3, 0, 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, true, 2, 1, 2, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 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, true, 2, 2, 2, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 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, true, 2, 2, 3, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 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, true, 2, 3, 5, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 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, true, 2, 0, 0, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 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, true, 0, 3, 5, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true, 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, true, 2, 0, 5, 5, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, true,2, 0, 5, 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;
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3)      { LAUNCH_IMAGE_KER(1, 1, 1); }
+                else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(1, 1, 2); }
+                else if (kernel_t_size == 5 && kernel_h_size == 3 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(2, 1, 1); }
+                else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(1, 2, 2); }
+                else if (kernel_t_size == 5 && kernel_h_size == 6 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(2, 3, 0); }
+                else if (kernel_t_size == 5 && kernel_h_size == 3 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(2, 1, 2); }
+                else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(2, 2, 2); }
+                else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 7) { LAUNCH_IMAGE_KER(2, 2, 3); }
+                else if (kernel_t_size == 5 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(2, 3, 5); }
+                else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(1, 3, 5); }
+                else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(2, 0, 0); }
+                else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(0, 3, 5); }
+                else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(2, 0, 5); }
+                else {
+                    TORCH_CHECK(false, "Invalid kernel size: ", kernel_t_size, "x", kernel_h_size, "x", kernel_w_size);
                 }
+#undef LAUNCH_IMAGE_KER
             } else {
                 cudaFuncSetAttribute(
                     fwd_attend_ker<128, false, true, true, 1, 1, 1, 5, 6, 10>,
@@ -568,266 +499,67 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
         } else {
             dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
             if (kernel_aspect_ratio_flag == 2){
-                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+#define LAUNCH_IMAGE_KER(DT_VAL, DH_VAL, DW_VAL) \
+                    cudaFuncSetAttribute( \
+                        fwd_attend_ker<128, false, false, false, DT_VAL, DH_VAL, DW_VAL, 6, 6, 6>, \
+                        cudaFuncAttributeMaxDynamicSharedMemorySize, \
+                        mem_size \
+                    ); \
+                    fwd_attend_ker<128, false, false, false, DT_VAL, DH_VAL, DW_VAL, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) {
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) {
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                }  else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                }  else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<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;
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3)      { LAUNCH_IMAGE_KER(1, 1, 1); }
+                else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(1, 1, 3); }
+                else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(3, 1, 1); }
+                else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(1, 3, 3); }
+                else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(1, 3, 1); }
+                else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(3, 1, 3); }
+                else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(3, 3, 3); }
+                else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(3, 0, 0); }
+                else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(3, 0, 3); }
+                else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(3, 3, 0); }
+                else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(0, 3, 3); }
+                else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(0, 0, 3); }
+                else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(0, 3, 0); }
+                else {
+                    TORCH_CHECK(false, "Invalid kernel size: ", kernel_t_size, "x", kernel_h_size, "x", kernel_w_size);
                 }
+#undef LAUNCH_IMAGE_KER
             }
             else if (kernel_aspect_ratio_flag == 3) {
-                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+#define LAUNCH_IMAGE_KER(DT_VAL, DH_VAL, DW_VAL) \
+                    cudaFuncSetAttribute( \
+                        fwd_attend_ker<128, false, false, false, DT_VAL, DH_VAL, DW_VAL, 3, 6, 10>, \
+                        cudaFuncAttributeMaxDynamicSharedMemorySize, \
+                        mem_size \
+                    ); \
+                    fwd_attend_ker<128, false, false, false, DT_VAL, DH_VAL, DW_VAL, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 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, false, 1, 1, 2, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,1, 1, 2, 3, 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, false, 1, 2, 2, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-
-                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 7){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 9){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 1){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 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, false, 0, 3, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 7){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 7){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 9){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 0, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,1, 0, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 1, 1, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,1, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 3 && kernel_w_size == 10){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 1, 5, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,0, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 5){
-                    cudaFuncSetAttribute(
-                        fwd_attend_ker<128, false, false, false, 0, 3, 2, 3, 6, 10>,
-                        cudaFuncAttributeMaxDynamicSharedMemorySize,
-                        mem_size
-                    );
-                    fwd_attend_ker<128, false, false, false,0, 3, 2, 3, 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;
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3)      { LAUNCH_IMAGE_KER(1, 1, 1); }
+                else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(1, 1, 2); }
+                else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(1, 2, 2); }
+                else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(1, 3, 0); }
+                else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 7) { LAUNCH_IMAGE_KER(1, 2, 3); }
+                else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 9) { LAUNCH_IMAGE_KER(1, 2, 4); }
+                else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(1, 3, 5); }
+                else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(1, 3, 1); }
+                else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(1, 0, 0); }
+                else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(0, 3, 5); }
+                else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(0, 2, 5); }
+                else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 7) { LAUNCH_IMAGE_KER(0, 3, 3); }
+                else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 7) { LAUNCH_IMAGE_KER(0, 2, 3); }
+                else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 9) { LAUNCH_IMAGE_KER(0, 2, 4); }
+                else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(1, 0, 5); }
+                else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(1, 1, 5); }
+                else if (kernel_t_size == 1 && kernel_h_size == 3 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(0, 1, 5); }
+                else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(0, 3, 2); }
+                else {
+                    TORCH_CHECK(false, "Invalid kernel size: ", kernel_t_size, "x", kernel_h_size, "x", kernel_w_size);
                 }
+#undef LAUNCH_IMAGE_KER
             }
 
             else {
-                std::cout << "Unsupported kernel_aspect_ratio_flag: " << kernel_aspect_ratio_flag << std::endl;
+                TORCH_CHECK(false, "Unsupported kernel_aspect_ratio_flag: ", kernel_aspect_ratio_flag);
             }
 
         }

fastvideo-kernel/csrc/common_extension.cpp

@@ -0,0 +1,36 @@
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+#include <vector>
+
+// Forward declarations
+#ifdef TK_COMPILE_ST_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, bool has_text, int kernel_aspect_ratio_flag
+); 
+#endif
+
+#ifdef TK_COMPILE_BLOCK_SPARSE
+extern std::vector<torch::Tensor> block_sparse_attention_forward(
+    torch::Tensor q, torch::Tensor k, torch::Tensor v,  
+    torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num, torch::Tensor block_size
+); 
+extern std::vector<torch::Tensor> block_sparse_attention_backward(
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og, 
+    torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num, torch::Tensor block_size
+);
+#endif
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "FastVideo CUDA Kernels";
+
+#ifdef TK_COMPILE_ST_ATTN
+    m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention (Hopper)");
+#endif
+
+#ifdef TK_COMPILE_BLOCK_SPARSE
+    m.def("block_sparse_fwd", torch::wrap_pybind_function(block_sparse_attention_forward), "block sparse attention forward (Hopper)");
+    m.def("block_sparse_bwd", torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward (Hopper)");
+#endif
+}

fastvideo-kernel/pyproject.toml

@@ -0,0 +1,35 @@
+[build-system]
+requires = [
+    "scikit-build-core>=0.10",
+    "torch>=2.5.0",
+    "setuptools>=61.0.0",
+    "wheel"
+]
+build-backend = "scikit_build_core.build"
+
+[project]
+name = "fastvideo-kernel"
+version = "0.2.1"
+description = "Unified CUDA kernels for FastVideo"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { file = "LICENSE" }
+authors = [
+    { name = "Hao AI Lab", email = "contact@haoailab.com" }
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Environment :: GPU :: NVIDIA CUDA",
+]
+dependencies = [
+    "torch>=2.5.0",
+    "triton>=2.0.0",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/hao-ai-lab/FastVideo"
+
+[tool.scikit-build]
+cmake.build-type = "Release"
+minimum-version = "build-system.requires"
+wheel.packages = ["python/fastvideo_kernel"]

fastvideo-kernel/python/fastvideo_kernel/__init__.py

@@ -0,0 +1,21 @@
+from .version import __version__
+
+from fastvideo_kernel.ops import (
+    sliding_tile_attention,
+    video_sparse_attn,
+)
+
+from fastvideo_kernel.vmoba import (
+    moba_attn_varlen,
+    process_moba_input,
+    process_moba_output,
+)
+
+__all__ = [
+    "sliding_tile_attention",
+    "video_sparse_attn",
+    "moba_attn_varlen",
+    "process_moba_input",
+    "process_moba_output",
+    "__version__",
+]

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -0,0 +1,112 @@
+import math
+import torch
+from .triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward
+from .triton_kernels.st_attn_triton import sliding_tile_attention_triton
+from .triton_kernels.index import map_to_index
+
+# Try to load the C++ extension
+try:
+    from fastvideo_kernel._C import fastvideo_kernel_ops
+    sta_fwd = getattr(fastvideo_kernel_ops, "sta_fwd", None)
+    block_sparse_fwd = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None)
+    block_sparse_bwd = getattr(fastvideo_kernel_ops, "block_sparse_bwd", None)
+except ImportError:
+    sta_fwd = None
+    block_sparse_fwd = None
+    block_sparse_bwd = None
+
+def sliding_tile_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    window_size: list,
+    text_length: int,
+    has_text: bool = True,
+    seq_shape: str = "30x48x80",
+) -> torch.Tensor:
+    # Check if the specific op is available
+    if sta_fwd is None:
+        return sliding_tile_attention_triton(
+            q, k, v, window_size, text_length, has_text, seq_shape
+        )
+
+    seq_length = q.shape[2]
+    shape_map = {"30x48x80": 1, "36x48x48": 2, "18x48x80": 3}
+
+    if has_text:
+        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)
+
+    output = torch.empty_like(q)
+    flag = shape_map[seq_shape]
+
+    for head_idx, (t, h, w) in enumerate(window_size):
+        sta_fwd(
+            q[:, head_idx:head_idx + 1], k[:, head_idx:head_idx + 1],
+            v[:, head_idx:head_idx + 1], output[:, head_idx:head_idx + 1],
+            t, h, w, text_length, False, has_text, flag
+        )
+
+    if has_text:
+        sta_fwd(q, k, v, output, 3, 3, 3, text_length, True, True, flag)
+
+    return output[:, :, :seq_length]
+
+
+def video_sparse_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    topk: int,
+    block_size: int | tuple = 64,
+    compress_attn_weight: torch.Tensor = None,
+) -> torch.Tensor:
+    if isinstance(block_size, int):
+        block_size = (block_size, block_size, block_size)
+
+    block_elements = block_size[0] * block_size[1] * block_size[2]
+    batch, heads, seq_len, dim = q.shape
+
+    # Compression branch
+    q_c = q.view(batch, heads, seq_len // block_elements, block_elements, dim)
+    k_c = k.view(batch, heads, seq_len // block_elements, block_elements, dim)
+    v_c = v.view(batch, heads, seq_len // block_elements, block_elements, dim)
+
+    q_c = (q_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
+        q.dtype)
+    k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
+        k.dtype)
+    v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
+        v.dtype)
+
+    scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5)
+    attn = torch.softmax(scores, dim=-1)
+    out_c = torch.matmul(attn, v_c)
+
+    out_c = out_c.view(batch, heads, seq_len // block_elements, 1, dim)
+    out_c = out_c.repeat(1, 1, 1, block_elements,
+                         1).view(batch, heads, seq_len, dim)
+
+    # Sparse branch
+    topk_idx = torch.topk(scores, topk, dim=-1).indices
+    mask = torch.zeros_like(scores,
+                            dtype=torch.bool).scatter_(-1, topk_idx, True)
+
+    idx, num = map_to_index(mask)
+    
+    if block_sparse_fwd is not None:
+        out_s = block_sparse_fwd(
+            q, k, v, idx, num, variable_block_sizes.int()
+        )[0] # block_sparse_fwd returns vector<Tensor>
+    else:
+        out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num,
+                                                    variable_block_sizes)
+
+    if compress_attn_weight is not None:
+        return out_c * compress_attn_weight + out_s
+    return out_c + out_s

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py

@@ -8,7 +8,6 @@ This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
 Credits: OpenAI kernel team
 """
 
-import pytest
 import torch
 import triton
 import triton.language as tl
@@ -17,7 +16,6 @@ import triton.language as tl
 import math  # small utility needed by the sparse wrapper
 # ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
 
-
 # We don't run auto-tuning every time to keep the tutorial fast. Keeping
 # the code below and commenting out the equivalent parameters is convenient for
 # re-tuning.
@@ -29,65 +27,92 @@ configs = [
     for w in [4, 8]\
 ]
 
+
 # ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
 @triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
 @triton.jit
-def _attn_fwd_sparse(Q, K, V, sm_scale,                         #
-                     q2k_index, q2k_num, max_kv_blks,           #
-                     variable_block_sizes,
-                     M, Out,                                    #
-                     stride_qz, stride_qh, stride_qm, stride_qk,
-                     stride_kz, stride_kh, stride_kn, stride_kk,
-                     stride_vz, stride_vh, stride_vk, stride_vn,
-                     stride_oz, stride_oh, stride_om, stride_on,
-                     Z, H, N_CTX,                               #
-                     HEAD_DIM: tl.constexpr,                    #
-                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
-                     STAGE: tl.constexpr):
+def _attn_fwd_sparse(
+        Q,
+        K,
+        V,
+        sm_scale,  #
+        q2k_index,
+        q2k_num,
+        max_kv_blks,  #
+        variable_block_sizes,
+        M,
+        Out,  #
+        stride_qz,
+        stride_qh,
+        stride_qm,
+        stride_qk,
+        stride_kz,
+        stride_kh,
+        stride_kn,
+        stride_kk,
+        stride_vz,
+        stride_vh,
+        stride_vk,
+        stride_vn,
+        stride_oz,
+        stride_oh,
+        stride_om,
+        stride_on,
+        Z,
+        H,
+        N_CTX,  #
+        HEAD_DIM: tl.constexpr,  #
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        STAGE: tl.constexpr):
     """
     64×64 **block-sparse** forward kernel. Back-prop kernels remain dense
     (32×64 and 64×32) – memory footprint unchanged.
     """
 
     # ----- program-id mapping -----
-    q_blk   = tl.program_id(0)          # Q-tile index
-    off_hz  = tl.program_id(1)          # fused (batch, head)
-    b       = off_hz // H
-    h       = off_hz %  H
+    q_blk = tl.program_id(0)  # Q-tile index
+    off_hz = tl.program_id(1)  # fused (batch, head)
+    b = off_hz // H
+    h = off_hz % H
     q_tiles = N_CTX // BLOCK_M
     meta_base = ((b * H + h) * q_tiles + q_blk)
 
-    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
-    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
+    kv_blocks = tl.load(q2k_num + meta_base)  # int32
+    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
 
     # ----- base pointers -----
-    qvk_off = (b.to(tl.int64) * stride_qz +
-               h.to(tl.int64) * stride_qh)
+    qvk_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
 
-    Q_ptr = tl.make_block_ptr(
-        base=Q + qvk_off, shape=(N_CTX, HEAD_DIM),
-        strides=(stride_qm, stride_qk),
-        offsets=(q_blk * BLOCK_M, 0),
-        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
+    Q_ptr = tl.make_block_ptr(base=Q + qvk_off,
+                              shape=(N_CTX, HEAD_DIM),
+                              strides=(stride_qm, stride_qk),
+                              offsets=(q_blk * BLOCK_M, 0),
+                              block_shape=(BLOCK_M, HEAD_DIM),
+                              order=(1, 0))
 
-    K_base = tl.make_block_ptr(
-        base=K + qvk_off, shape=(HEAD_DIM, N_CTX),
-        strides=(stride_kk, stride_kn),
-        offsets=(0, 0),
-        block_shape=(HEAD_DIM, BLOCK_N), order=(0, 1))
+    K_base = tl.make_block_ptr(base=K + qvk_off,
+                               shape=(HEAD_DIM, N_CTX),
+                               strides=(stride_kk, stride_kn),
+                               offsets=(0, 0),
+                               block_shape=(HEAD_DIM, BLOCK_N),
+                               order=(0, 1))
 
-    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
-    V_base = tl.make_block_ptr(
-        base=V + qvk_off, shape=(N_CTX, HEAD_DIM),
-        strides=(stride_vk, stride_vn),
-        offsets=(0, 0),
-        block_shape=(BLOCK_N, HEAD_DIM), order=v_order)
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1,
+                                                                              0)
+    V_base = tl.make_block_ptr(base=V + qvk_off,
+                               shape=(N_CTX, HEAD_DIM),
+                               strides=(stride_vk, stride_vn),
+                               offsets=(0, 0),
+                               block_shape=(BLOCK_N, HEAD_DIM),
+                               order=v_order)
 
-    O_ptr = tl.make_block_ptr(
-        base=Out + qvk_off, shape=(N_CTX, HEAD_DIM),
-        strides=(stride_om, stride_on),
-        offsets=(q_blk * BLOCK_M, 0),
-        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
+    O_ptr = tl.make_block_ptr(base=Out + qvk_off,
+                              shape=(N_CTX, HEAD_DIM),
+                              strides=(stride_om, stride_on),
+                              offsets=(q_blk * BLOCK_M, 0),
+                              block_shape=(BLOCK_M, HEAD_DIM),
+                              order=(1, 0))
 
     # ----- accumulators -----
     offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M)
@@ -127,23 +152,30 @@ def _attn_fwd_sparse(Q, K, V, sm_scale,                         #
     acc = acc / l_i[:, None]
     tl.store(M + off_hz * N_CTX + offs_m, m_i)
     tl.store(O_ptr, acc.to(Out.type.element_ty))
+
+
 # ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
 
 
-
-
 @triton.jit
-def _attn_bwd_preprocess(O, DO,  #
-                         Delta,  #
-                         Z, H, N_CTX,  #
-                         BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr  #
-                         ):
+def _attn_bwd_preprocess(
+        O,
+        DO,  #
+        Delta,  #
+        Z,
+        H,
+        N_CTX,  #
+        BLOCK_M: tl.constexpr,
+        HEAD_DIM: tl.constexpr  #
+):
     off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
     off_hz = tl.program_id(1)
     off_n = tl.arange(0, HEAD_DIM)
     # load
-    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
-    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
+                off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
+                 off_n[None, :]).to(tl.float32)
     delta = tl.sum(o * do, axis=1)
     # write-back
     tl.store(Delta + off_hz * N_CTX + off_m, delta)
@@ -151,19 +183,32 @@ def _attn_bwd_preprocess(O, DO,  #
 
 # The main inner-loop logic for computing dK and dV.
 @triton.jit
-def _attn_bwd_dkdv(dk, dv,  #
-                   Q, k, v, sm_scale,  #
-                   DO,  #
-                   M, D,  #
-                   k2q_index, k2q_num, max_q_blks,
-                   variable_block_sizes,
-                   # shared by Q/K/V/DO.
-                   stride_tok, stride_d,  #
-                   H, N_CTX, BLOCK_M1: tl.constexpr,  #
-                   BLOCK_N1: tl.constexpr,  #
-                   HEAD_DIM: tl.constexpr,  #
-                   # Filled in by the wrapper.
-                   start_n, start_m, num_steps):
+def _attn_bwd_dkdv(
+        dk,
+        dv,  #
+        Q,
+        k,
+        v,
+        sm_scale,  #
+        DO,  #
+        M,
+        D,  #
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        # shared by Q/K/V/DO.
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,
+        BLOCK_M1: tl.constexpr,  #
+        BLOCK_N1: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr,  #
+        # Filled in by the wrapper.
+    start_n,
+        start_m,
+        num_steps):
     offs_m = start_m + tl.arange(0, BLOCK_M1)
     offs_n = start_n + tl.arange(0, BLOCK_N1)
     offs_k = tl.arange(0, HEAD_DIM)
@@ -172,21 +217,20 @@ def _attn_bwd_dkdv(dk, dv,  #
     # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
     tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
     step_m = BLOCK_M1
-    kv_blk   = tl.program_id(0)          # Q-tile index
-    off_hz  = tl.program_id(2)          # fused (batch, head)
-    b       = off_hz // H
-    h       = off_hz %  H
+    kv_blk = tl.program_id(0)  # Q-tile index
+    off_hz = tl.program_id(2)  # fused (batch, head)
+    b = off_hz // H
+    h = off_hz % H
     q_tiles = N_CTX // BLOCK_N1
     meta_base = ((b * H + h) * q_tiles + kv_blk)
 
-    q_blocks = tl.load(k2q_num  + meta_base)                 # int32
-    q_ptr    = k2q_index + meta_base * max_q_blks           # ptr to list
+    q_blocks = tl.load(k2q_num + meta_base)  # int32
+    q_ptr = k2q_index + meta_base * max_q_blks  # ptr to list
     block_size = tl.load(variable_block_sizes + kv_blk)
-    
-    
-        
-    for blk_idx in range(q_blocks*2):
-        block_sparse_offset = (tl.load(q_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_m
+
+    for blk_idx in range(q_blocks * 2):
+        block_sparse_offset = (tl.load(q_ptr + blk_idx // 2).to(tl.int32) * 2 +
+                               blk_idx % 2) * step_m
         qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok)
         # Load m before computing qk to reduce pipeline stall.
         offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
@@ -212,21 +256,32 @@ def _attn_bwd_dkdv(dk, dv,  #
     return dk, dv
 
 
-
 # the main inner-loop logic for computing dQ
 @triton.jit
-def _attn_bwd_dq(dq, q, K, V,  #
-                 do, m, D,
-                 # shared by Q/K/V/DO.
-                 q2k_index, q2k_num, max_kv_blks,
-                 variable_block_sizes,
-                 stride_tok, stride_d,  #
-                 H, N_CTX,  #
-                 BLOCK_M2: tl.constexpr,  #
-                 BLOCK_N2: tl.constexpr,  #
-                 HEAD_DIM: tl.constexpr,
-                 # Filled in by the wrapper.
-                 start_m, start_n, num_steps):
+def _attn_bwd_dq(
+        dq,
+        q,
+        K,
+        V,  #
+        do,
+        m,
+        D,
+        # shared by Q/K/V/DO.
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M2: tl.constexpr,  #
+        BLOCK_N2: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr,
+        # Filled in by the wrapper.
+        start_m,
+        start_n,
+        num_steps):
     offs_m = start_m + tl.arange(0, BLOCK_M2)
     offs_n = start_n + tl.arange(0, BLOCK_N2)
     offs_k = tl.arange(0, HEAD_DIM)
@@ -237,28 +292,27 @@ def _attn_bwd_dq(dq, q, K, V,  #
     # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
     tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
     step_n = BLOCK_N2
-    
-    q_blk   = tl.program_id(0)          # Q-tile index
-    off_hz  = tl.program_id(2)          # fused (batch, head)
-    b       = off_hz // H
-    h       = off_hz %  H
+
+    q_blk = tl.program_id(0)  # Q-tile index
+    off_hz = tl.program_id(2)  # fused (batch, head)
+    b = off_hz // H
+    h = off_hz % H
     q_tiles = N_CTX // BLOCK_M2
     meta_base = ((b * H + h) * q_tiles + q_blk)
 
-    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
-    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
-    
-    
-    for blk_idx in range(kv_blocks*2):
-        kv_idx = tl.load(kv_ptr + blk_idx//2).to(tl.int32)
-        block_size = tl.load(variable_block_sizes + kv_idx) -  (blk_idx % 2) * step_n
-        block_sparse_offset = (kv_idx*2 + blk_idx%2) * step_n * stride_tok
+    kv_blocks = tl.load(q2k_num + meta_base)  # int32
+    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+    block_size = tl.load(variable_block_sizes + q_blk)
+
+    for blk_idx in range(kv_blocks * 2):
+        block_sparse_offset = (tl.load(kv_ptr + blk_idx // 2).to(tl.int32) * 2 +
+                               blk_idx % 2) * step_n * stride_tok
         kT = tl.load(kT_ptrs + block_sparse_offset)
         vT = tl.load(vT_ptrs + block_sparse_offset)
         qk = tl.dot(q, kT)
         p = tl.math.exp2(qk - m)
         mask = tl.arange(0, BLOCK_N2) < block_size.to(tl.int32)
-        p = tl.where(mask[None, :], p , 0.0)
+        p = tl.where(mask[None, :], p, 0.0)
         # Compute dP and dS.
         dp = tl.dot(do, vT).to(tl.float32)
         ds = p * (dp - Di[:, None])
@@ -270,23 +324,37 @@ def _attn_bwd_dq(dq, q, K, V,  #
     return dq
 
 
-
 @triton.jit
-def _attn_bwd(Q, K, V, sm_scale,  #
-              DO,  #
-              DQ, DK, DV,  #
-              M, D,
-              q2k_index, q2k_num, max_kv_blks,
-              k2q_index, k2q_num, max_q_blks,
-              variable_block_sizes,
-              # shared by Q/K/V/DO.
-              stride_z, stride_h, stride_tok, stride_d,  #
-              H, N_CTX,  #
-              BLOCK_M1: tl.constexpr,  #
-              BLOCK_N1: tl.constexpr,  #
-              BLOCK_M2: tl.constexpr,  #
-              BLOCK_N2: tl.constexpr,  #
-              HEAD_DIM: tl.constexpr):
+def _attn_bwd(
+        Q,
+        K,
+        V,
+        sm_scale,  #
+        DO,  #
+        DQ,
+        DK,
+        DV,  #
+        M,
+        D,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        # shared by Q/K/V/DO.
+        stride_z,
+        stride_h,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M1: tl.constexpr,  #
+        BLOCK_N1: tl.constexpr,  #
+        BLOCK_M2: tl.constexpr,  #
+        BLOCK_N2: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr):
     LN2 = 0.6931471824645996  # = ln(2)
 
     bhid = tl.program_id(2)
@@ -320,20 +388,32 @@ def _attn_bwd(Q, K, V, sm_scale,  #
     k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
     v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
 
-
     num_steps = N_CTX // BLOCK_M1
 
     dk, dv = _attn_bwd_dkdv(  #
-        dk, dv,  #
-        Q, k, v, sm_scale,  #
+        dk,
+        dv,  #
+        Q,
+        k,
+        v,
+        sm_scale,  #
         DO,  #
-        M, D,  #
-        k2q_index, k2q_num, max_q_blks,
+        M,
+        D,  #
+        k2q_index,
+        k2q_num,
+        max_q_blks,
         variable_block_sizes,
-        stride_tok, stride_d,  #
-        H, N_CTX,  #
-        BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
-        start_n, start_m, num_steps  #
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M1,
+        BLOCK_N1,
+        HEAD_DIM,  #
+        start_n,
+        start_m,
+        num_steps  #
     )
 
     dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
@@ -358,54 +438,88 @@ def _attn_bwd(Q, K, V, sm_scale,  #
     m = m[:, None]
 
     num_steps = N_CTX // BLOCK_N2
-    dq = _attn_bwd_dq(dq, q, K, V,  #
-                      do, m, D,  #
-                      q2k_index, q2k_num, max_kv_blks,
-                      variable_block_sizes,
-                      stride_tok, stride_d,  #
-                      H, N_CTX,  #
-                      BLOCK_M2, BLOCK_N2, HEAD_DIM,  #
-                      start_m, end_n, num_steps  #
-                      )
+    dq = _attn_bwd_dq(
+        dq,
+        q,
+        K,
+        V,  #
+        do,
+        m,
+        D,  #
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M2,
+        BLOCK_N2,
+        HEAD_DIM,  #
+        start_m,
+        end_n,
+        num_steps  #
+    )
     # Write back dQ.
     dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
     dq *= LN2
     tl.store(dq_ptrs, dq)
-    
-    
-
 
 
 # ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
-def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes):
+def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
+                                     variable_block_sizes):
     B, H, T, D = q.shape
     sm_scale = 1.0 / math.sqrt(D)
     max_kv_blks = q2k_index.shape[-1]
     assert T % 64 == 0, f"T must be a multiple of 64, but got {T}"
-    assert T // 64 == q2k_num.shape[-1], f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
+    assert q2k_num.shape[
+        -1] == T // 64, f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
     o = torch.empty_like(q)
     M = torch.empty((B, H, T), dtype=torch.float32, device=q.device)
 
     grid = lambda _: (triton.cdiv(T, 64), B * H, 1)
-    _attn_fwd_sparse[grid](
-        q, k, v, sm_scale,
-        q2k_index, q2k_num, max_kv_blks,
-        variable_block_sizes,
-        M, o,
-        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
-        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
-        v.stride(0), v.stride(1), v.stride(2), v.stride(3),
-        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
-        B, H, T,
-        HEAD_DIM=D, STAGE=3
-    )
+    _attn_fwd_sparse[grid](q,
+                           k,
+                           v,
+                           sm_scale,
+                           q2k_index,
+                           q2k_num,
+                           max_kv_blks,
+                           variable_block_sizes,
+                           M,
+                           o,
+                           q.stride(0),
+                           q.stride(1),
+                           q.stride(2),
+                           q.stride(3),
+                           k.stride(0),
+                           k.stride(1),
+                           k.stride(2),
+                           k.stride(3),
+                           v.stride(0),
+                           v.stride(1),
+                           v.stride(2),
+                           v.stride(3),
+                           o.stride(0),
+                           o.stride(1),
+                           o.stride(2),
+                           o.stride(3),
+                           B,
+                           H,
+                           T,
+                           HEAD_DIM=D,
+                           STAGE=3)
 
     return o, M
 
-def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num, variable_block_sizes):
+
+def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
+                                      k2q_index, k2q_num, variable_block_sizes):
     assert do.is_contiguous()
     assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
-    
+
     B, H, T, D = q.shape
     sm_scale = 1.0 / math.sqrt(D)
     dq = torch.empty_like(q)
@@ -421,30 +535,49 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q
     pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
     delta = torch.empty_like(M)
     _attn_bwd_preprocess[pre_grid](
-        o, do,  #
+        o,
+        do,  #
         delta,  #
-        BATCH, N_HEAD, N_CTX,  #
-        BLOCK_M=PRE_BLOCK, HEAD_DIM=D  #
+        BATCH,
+        N_HEAD,
+        N_CTX,  #
+        BLOCK_M=PRE_BLOCK,
+        HEAD_DIM=D  #
     )
-    
-    
+
     max_q_blks = k2q_index.shape[-1]
     max_kv_blks = q2k_index.shape[-1]
-    
+
     grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
     _attn_bwd[grid](
-        q, arg_k, v, sm_scale, do, dq, dk, dv,  #
-        M, delta,  #
-        q2k_index, q2k_num, max_kv_blks,
-        k2q_index, k2q_num, max_q_blks,
+        q,
+        arg_k,
+        v,
+        sm_scale,
+        do,
+        dq,
+        dk,
+        dv,  #
+        M,
+        delta,  #
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
         variable_block_sizes,
-        q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
-        N_HEAD, N_CTX,  #
-        BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1,  #
-        BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2,  #
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        q.stride(3),  #
+        N_HEAD,
+        N_CTX,  #
+        BLOCK_M1=BLOCK_M1,
+        BLOCK_N1=BLOCK_N1,  #
+        BLOCK_M2=BLOCK_M2,
+        BLOCK_N2=BLOCK_N2,  #
         HEAD_DIM=D  #
     )
 
     return dq, dk, dv
-
-

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py

@@ -1,9 +1,9 @@
-
 ## pytorch sdpa version of block sparse ##
 import triton
 import triton.language as tl
 import torch
 
+
 @triton.jit
 def topk_index_to_map_kernel(
     map_ptr,
@@ -26,6 +26,7 @@ def topk_index_to_map_kernel(
         index = tl.load(index_ptr_base + i * index_kv_stride)
         tl.store(map_ptr_base + index * map_kv_stride, 1.0)
 
+
 @triton.jit
 def map_to_index_kernel(
     map_ptr,
@@ -59,6 +60,7 @@ def map_to_index_kernel(
         index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride +
         q * index_num_q_stride, num)
 
+
 def topk_index_to_map(index: torch.Tensor,
                       num_kv_blocks: int,
                       transpose_map: bool = False):
@@ -106,6 +108,7 @@ def topk_index_to_map(index: torch.Tensor,
 
     return block_map
 
+
 def map_to_index(block_map: torch.Tensor):
     """
     Convert a block map to indices and counts.

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/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-kernel/python/fastvideo_kernel/version.py

@@ -0,0 +1 @@
+__version__ = "0.2.1"

fastvideo-kernel/tests/test_correctness.py

@@ -0,0 +1,63 @@
+import torch
+import sys
+import os
+from tqdm import tqdm
+
+# Local support import
+from .support_flex_sta import get_sliding_tile_attention_mask
+
+# USE OUR NEW PACKAGE!
+from fastvideo_kernel import sliding_tile_attention
+from torch.nn.attention.flex_attention import flex_attention
+
+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), (18, 48, 80), 0, 'cuda', 0)
+    output = flex_attention(Q, K, V, block_mask=mask)
+    return output
+
+def h100_fwd_kernel_test(Q, K, V, kernel_size):
+    # Using the same parameters as the original test
+    o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 0, False, '18x48x80')
+    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=2):
+    print(f"Running correctness check: batch={b}, heads={h}, seq_len={n}, dim={d}")
+    kernel_size_ls = [(3, 3, 5), (3, 1, 10)]
+    
+    for kernel_size in kernel_size_ls:
+        print(f"Testing kernel_size: {kernel_size}")
+        for xi in tqdm(range(num_iterations)):
+            torch.manual_seed(xi)
+            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)
+            max_d = torch.max(abs_diff).item()
+            avg_d = torch.sum(abs_diff).item() / (b * h * n * d)
+            
+            if max_d > 0.1:
+                print(f"Warning: Large diff detected! max={max_d}, avg={avg_d}")
+
+    print("\n✅ TEST COMPLETE: New package matches FlexAttention behavior.")
+
+if __name__ == "__main__":
+    b, h, d = 2, 24, 128
+    n = 69120 
+    causal = False
+    mean = 1e-1
+    std = 10
+    
+    check_correctness(b, h, n, d, causal, mean, std, num_iterations=2)

fastvideo-kernel/tests/test_vmoba_correctness.py

@@ -3,7 +3,7 @@
 import torch
 import pytest
 import random
-from csrc.attn.vmoba_attn.vmoba  import moba_attn_varlen
+from fastvideo_kernel import moba_attn_varlen
 
 def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda"):
     """
@@ -51,7 +51,7 @@ def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, dev
 @pytest.mark.parametrize("moba_topk", [2, 4])
 @pytest.mark.parametrize("select_mode", ["topk", "threshold"])
 @pytest.mark.parametrize("threshold_type", ["query_head", "head_global", "overall"])
-@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_moba_attn_varlen_forward(
     batch_size, total_seqlen, num_heads, head_dim, moba_chunk_size, moba_topk, select_mode, threshold_type, dtype
 ):

fastvideo/attention/backends/sdpa.py

@@ -81,7 +81,9 @@ class SDPAImpl(AttentionImpl):
         query = query.transpose(1, 2)
         key = key.transpose(1, 2)
         value = value.transpose(1, 2)
-        attn_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+
+        if attn_metadata is not None:
+            attn_mask = getattr(attn_metadata, "attn_mask", None)
         attn_kwargs = {
             "attn_mask": attn_mask,
             "dropout_p": self.dropout,

fastvideo/attention/backends/sliding_tile_attn.py

@@ -5,7 +5,7 @@ from typing import Any
 
 import torch
 from einops import rearrange
-from st_attn import sliding_tile_attention
+from fastvideo_kernel import sliding_tile_attention
 
 import fastvideo.envs as envs
 from fastvideo.attention.backends.abstract import (AttentionBackend,

fastvideo/attention/backends/video_sparse_attn.py

@@ -6,7 +6,7 @@ from dataclasses import dataclass
 import torch
 
 try:
-    from vsa import video_sparse_attn
+    from fastvideo_kernel import video_sparse_attn
 except ImportError:
     video_sparse_attn = None
 

fastvideo/attention/backends/vmoba.py

@@ -6,8 +6,8 @@ from dataclasses import dataclass
 import torch
 from einops import rearrange
 
-from csrc.attn.vmoba_attn.vmoba import (moba_attn_varlen, process_moba_input,
-                                        process_moba_output)
+from fastvideo_kernel import (moba_attn_varlen, process_moba_input,
+                              process_moba_output)
 from fastvideo.attention.backends.abstract import (AttentionBackend,
                                                    AttentionImpl,
                                                    AttentionMetadata,

fastvideo/configs/models/dits/__init__.py

@@ -2,10 +2,12 @@ from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
+from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
 from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig",
-    "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig"
+    "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
+    "LongCatVideoConfig"
 ]

fastvideo/configs/models/dits/longcat.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat Video DiT configuration for native FastVideo implementation.
+"""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def is_longcat_blocks(n: str, m) -> bool:
+    """FSDP shard condition for LongCat transformer blocks."""
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class LongCatVideoArchConfig(DiTArchConfig):
+    """Architecture configuration for native LongCat Video DiT."""
+
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_longcat_blocks])
+
+    # Enable torch.compile for transformer blocks (major speedup!)
+    _compile_conditions: list = field(
+        default_factory=lambda: [is_longcat_blocks])
+
+    # Parameter name mapping for weight conversion
+    # Maps original LongCat third_party names -> native FastVideo names
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # Embedders
+            r"^x_embedder\.(.*)$": r"patch_embed.\1",
+            r"^t_embedder\.mlp\.0\.(.*)$": r"time_embedder.linear_1.\1",
+            r"^t_embedder\.mlp\.2\.(.*)$": r"time_embedder.linear_2.\1",
+            r"^y_embedder\.y_proj\.0\.(.*)$": r"caption_embedder.linear_1.\1",
+            r"^y_embedder\.y_proj\.2\.(.*)$": r"caption_embedder.linear_2.\1",
+
+            # Transformer blocks - AdaLN modulation
+            r"^blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$":
+            r"blocks.\1.adaln_linear_1.\2",
+
+            # Transformer blocks - Normalization
+            r"^blocks\.(\d+)\.mod_norm_attn\.(.*)$": r"blocks.\1.norm_attn.\2",
+            r"^blocks\.(\d+)\.mod_norm_ffn\.(.*)$": r"blocks.\1.norm_ffn.\2",
+            r"^blocks\.(\d+)\.pre_crs_attn_norm\.(.*)$":
+            r"blocks.\1.norm_cross.\2",
+
+            # Self-attention: QKV fused -> separate (will need splitting in converter)
+            # Original has attn.qkv.weight -> need to split into to_q, to_k, to_v
+            r"^blocks\.(\d+)\.attn\.qkv\.(.*)$":
+            r"blocks.\1.self_attn.qkv_fused.\2",  # Marker for splitting
+            r"^blocks\.(\d+)\.attn\.proj\.(.*)$":
+            r"blocks.\1.self_attn.to_out.\2",
+            r"^blocks\.(\d+)\.attn\.q_norm\.(.*)$":
+            r"blocks.\1.self_attn.q_norm.\2",
+            r"^blocks\.(\d+)\.attn\.k_norm\.(.*)$":
+            r"blocks.\1.self_attn.k_norm.\2",
+
+            # Cross-attention
+            r"^blocks\.(\d+)\.cross_attn\.q_linear\.(.*)$":
+            r"blocks.\1.cross_attn.to_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.kv_linear\.(.*)$":
+            r"blocks.\1.cross_attn.kv_fused.\2",  # Marker for splitting
+            r"^blocks\.(\d+)\.cross_attn\.proj\.(.*)$":
+            r"blocks.\1.cross_attn.to_out.\2",
+            r"^blocks\.(\d+)\.cross_attn\.q_norm\.(.*)$":
+            r"blocks.\1.cross_attn.q_norm.\2",
+            r"^blocks\.(\d+)\.cross_attn\.k_norm\.(.*)$":
+            r"blocks.\1.cross_attn.k_norm.\2",
+
+            # FFN (SwiGLU)
+            r"^blocks\.(\d+)\.ffn\.w1\.(.*)$": r"blocks.\1.ffn.w1.\2",  # gate
+            r"^blocks\.(\d+)\.ffn\.w2\.(.*)$": r"blocks.\1.ffn.w2.\2",  # down
+            r"^blocks\.(\d+)\.ffn\.w3\.(.*)$": r"blocks.\1.ffn.w3.\2",  # up
+
+            # Final layer
+            r"^final_layer\.adaLN_modulation\.1\.(.*)$":
+            r"final_layer.adaln_linear.\1",
+            r"^final_layer\.norm_final\.(.*)$": r"final_layer.norm.\1",
+            r"^final_layer\.linear\.(.*)$": r"final_layer.proj.\1",
+        })
+
+    # Reverse mapping for saving checkpoints: custom -> hf
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    # LoRA parameter name mapping
+    lora_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    # Model architecture parameters
+    hidden_size: int = 4096
+    depth: int = 48  # Number of transformer blocks
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128  # hidden_size / num_attention_heads
+
+    in_channels: int = 16  # Latent space channels
+    out_channels: int = 16
+    num_channels_latents: int = 16
+
+    # Patch embedding
+    patch_size: tuple[int, int,
+                      int] = (1, 2, 2)  # [T, H, W] - no temporal compression
+
+    # Text/caption embedding
+    caption_channels: int = 4096  # UMT5 d_model
+
+    # Timestep embedding
+    adaln_tembed_dim: int = 512
+    frequency_embedding_size: int = 256
+
+    # FFN
+    mlp_ratio: int = 4
+
+    # Attention backend support
+    _supported_attention_backends: tuple = field(default_factory=lambda: (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    ))
+
+    # Text padding behavior
+    text_tokens_zero_pad: bool = True
+
+    # Block Sparse Attention (BSA)
+    enable_bsa: bool = False
+    bsa_params: dict | None = field(
+        default_factory=lambda: {
+            "sparsity": 0.9375,
+            "cdf_threshold": None,
+            "chunk_3d_shape_q": [4, 4, 4],
+            "chunk_3d_shape_k": [4, 4, 4],
+        })
+
+    # LoRA exclusions
+    exclude_lora_layers: list[str] = field(default_factory=lambda: [])
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        # Ensure attention_head_dim matches
+        self.attention_head_dim = self.hidden_size // self.num_attention_heads
+
+
+@dataclass
+class LongCatVideoConfig(DiTConfig):
+    """Main configuration for LongCat Video DiT."""
+
+    arch_config: DiTArchConfig = field(default_factory=LongCatVideoArchConfig)
+
+    prefix: str = "longcat"

fastvideo/configs/pipelines/longcat.py

@@ -0,0 +1,355 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+import html
+
+import ftfy
+import regex as re
+import torch
+
+from fastvideo.configs.models import DiTConfig, VAEConfig
+from fastvideo.configs.models.dits.base import DiTArchConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class LongCatDiTArchConfig(DiTArchConfig):
+    """Extended DiTArchConfig with LongCat-specific fields.
+    
+    NOTE: This is for Phase 1 wrapper compatibility. For native model (Phase 2),
+    use LongCatVideoConfig from fastvideo.configs.models.dits.longcat instead.
+    """
+    # LongCat-specific architecture parameters
+    adaln_tembed_dim: int = 512
+    caption_channels: int = 4096
+    depth: int = 48
+    enable_bsa: bool = False
+    enable_flashattn3: bool = False
+    enable_flashattn2: bool = True
+    enable_xformers: bool = False
+    frequency_embedding_size: int = 256
+    in_channels: int = 16
+    mlp_ratio: int = 4
+    num_heads: int = 32
+    out_channels: int = 16
+    text_tokens_zero_pad: bool = True
+    patch_size: list[int] = field(default_factory=lambda: [1, 2, 2])
+    cp_split_hw: list[int] | None = None
+    bsa_params: dict | None = None
+
+
+def longcat_preprocess_text(prompt: str) -> str:
+    """Clean and preprocess text like original LongCat implementation.
+    
+    This function applies the same text cleaning pipeline as the original
+    LongCat-Video implementation to ensure identical tokenization results.
+    
+    Steps:
+    1. basic_clean: Fix unicode issues and unescape HTML entities
+    2. whitespace_clean: Normalize whitespace to single spaces
+    
+    Args:
+        prompt: Raw input text prompt
+        
+    Returns:
+        Cleaned and normalized text prompt
+    """
+    # basic_clean: fix unicode and HTML entities
+    text = ftfy.fix_text(prompt)
+    text = html.unescape(html.unescape(text))
+    text = text.strip()
+
+    # whitespace_clean: normalize whitespace
+    text = re.sub(r"\s+", " ", text)
+    text = text.strip()
+
+    return text
+
+
+def umt5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    """
+    Postprocess UMT5/T5 encoder outputs to fixed length 512 embeddings.
+    """
+    mask: torch.Tensor = outputs.attention_mask
+    hidden_state: torch.Tensor = outputs.last_hidden_state
+    seq_lens = mask.gt(0).sum(dim=1).long()
+    assert torch.isnan(hidden_state).sum() == 0
+    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
+    prompt_embeds_tensor: torch.Tensor = torch.stack([
+        torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
+        for u in prompt_embeds
+    ],
+                                                     dim=0)
+    return prompt_embeds_tensor
+
+
+@dataclass
+class LongCatT2V480PConfig(PipelineConfig):
+    """Configuration for LongCat pipeline (480p) aligned to LongCat-Video modules.
+
+    Components expected by loaders:
+      - tokenizer: AutoTokenizer
+      - text_encoder: UMT5EncoderModel
+      - transformer: LongCatVideoTransformer3DModel (Phase 1 wrapper)
+                  OR LongCatTransformer3DModel (Phase 2 native)
+      - vae: AutoencoderKLWan (Wan VAE, 4x8 compression)
+      - scheduler: FlowMatchEulerDiscreteScheduler
+    """
+
+    # DiT config with LongCat-specific arch_config
+    # NOTE: For Phase 1 wrapper, uses LongCatDiTArchConfig
+    # For Phase 2 native model, can use LongCatVideoConfig directly
+    dit_config: DiTConfig = field(
+        default_factory=lambda: DiTConfig(arch_config=LongCatDiTArchConfig()))
+
+    # VAE config: Wan VAE with encoder+decoder enabled
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Precision defaults
+    dit_precision: str = "bf16"
+    vae_precision: str = "bf16"
+    text_encoder_precisions: tuple[str, ...] = field(
+        default_factory=lambda: ("bf16", ))
+
+    # Text encoding (UMT5 uses T5-like config; postprocess to fixed 512)
+    text_encoder_configs: tuple[T5Config, ...] = field(
+        default_factory=lambda: (T5Config(), ))
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (longcat_preprocess_text, ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda:
+                                               (umt5_postprocess_text, ))
+
+    # LongCat-specific runtime toggles (consumed by pipeline/stages)
+    enable_kv_cache: bool = True
+    offload_kv_cache: bool = False
+    enable_bsa: bool = False
+    use_distill: bool = False
+    enhance_hf: bool = False
+    # Optional BSA parameter dict (kept for backward/phase-1 compatibility).
+    # `LongCatPipeline.initialize_pipeline()` uses this as a base and then applies
+    # CLI overrides (bsa_sparsity / bsa_chunk_{q,k} / bsa_cdf_threshold).
+    bsa_params: dict | None = None
+    # BSA runtime overrides (preferred over bsa_params if provided via CLI)
+    bsa_sparsity: float | None = None
+    bsa_cdf_threshold: float | None = None
+    bsa_chunk_q: list[int] | None = None
+    bsa_chunk_k: list[int] | None = None
+    t_thresh: float | None = None  # refine stage default controlled by sampling args
+
+    # LongCat does not need flow_shift
+    flow_shift: float | None = None
+    dmd_denoising_steps: list[int] | None = None
+
+    def __post_init__(self):
+        # LongCat inference requires vae encoder and decoder
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class LongCatT2V704PConfig(LongCatT2V480PConfig):
+    """Configuration for LongCat pipeline (704p) with BSA enabled by default.
+    
+    Uses the same resolution and BSA parameters as original LongCat refinement stage.
+    BSA parameters configured in transformer config.json with chunk_3d_shape=[4,4,4]:
+    - Input: 704×1280×96
+    - VAE (8x): 88×160×96  
+    - Patch [1,2,2]: 44×80×96
+    - chunk [4,4,4]: 96%4=0, 44%4=0, 80%4=0 ✅
+    
+    This configuration matches the original LongCat refinement stage parameters.
+    """
+
+    # Enable BSA by default for 704p
+    enable_bsa: bool = True
+
+
+ASPECT_RATIO_627 = {
+    '0.26': ([320, 1216], 1),
+    '0.31': ([352, 1120], 1),
+    '0.38': ([384, 1024], 1),
+    '0.43': ([416, 960], 1),
+    '0.52': ([448, 864], 1),
+    '0.58': ([480, 832], 1),
+    '0.67': ([512, 768], 1),
+    '0.74': ([544, 736], 1),
+    '0.86': ([576, 672], 1),
+    '0.95': ([608, 640], 1),
+    '1.05': ([640, 608], 1),
+    '1.17': ([672, 576], 1),
+    '1.29': ([704, 544], 1),
+    '1.35': ([736, 544], 1),
+    '1.50': ([768, 512], 1),
+    '1.67': ([800, 480], 1),
+    '1.73': ([832, 480], 1),
+    '2.00': ([896, 448], 1),
+    '2.31': ([960, 416], 1),
+    '2.58': ([992, 384], 1),
+    '2.75': ([1056, 384], 1),
+    '3.09': ([1088, 352], 1),
+    '3.70': ([1184, 320], 1),
+    '3.80': ([1216, 320], 1),
+    '3.90': ([1248, 320], 1),
+    '4.00': ([1280, 320], 1)
+}
+
+ASPECT_RATIO_627_F64 = {
+    '0.26': ([320, 1216], 1),
+    '0.38': ([384, 1024], 1),
+    '0.50': ([448, 896], 1),
+    '0.67': ([512, 768], 1),
+    '0.82': ([576, 704], 1),
+    '1.00': ([640, 640], 1),
+    '1.22': ([704, 576], 1),
+    '1.50': ([768, 512], 1),
+    '1.86': ([832, 448], 1),
+    '2.00': ([896, 448], 1),
+    '2.50': ([960, 384], 1),
+    '2.83': ([1088, 384], 1),
+    '3.60': ([1152, 320], 1),
+    '3.80': ([1216, 320], 1),
+    '4.00': ([1280, 320], 1)
+}
+
+ASPECT_RATIO_627_F128 = {
+    '0.25': ([256, 1024], 1),
+    '0.38': ([384, 1024], 1),
+    '0.43': ([384, 896], 1),
+    '0.57': ([512, 896], 1),
+    '0.67': ([512, 768], 1),
+    '1.00': ([640, 640], 1),
+    '1.50': ([768, 512], 1),
+    '1.75': ([896, 512], 1),
+    '2.33': ([896, 384], 1),
+    '2.67': ([1024, 384], 1),
+    '4.00': ([1024, 256], 1),
+}
+
+ASPECT_RATIO_627_F256 = {
+    '0.25': ([256, 1024], 1),
+    '0.33': ([256, 768], 1),
+    '0.50': ([256, 512], 1),
+    '0.67': ([512, 768], 1),
+    '1.00': ([512, 512], 1),
+    '1.50': ([768, 512], 1),
+    '2.00': ([512, 256], 1),
+    '3.00': ([768, 256], 1),
+    '4.00': ([1024, 256], 1),
+}
+
+ASPECT_RATIO_960 = {
+    '0.25': ([480, 1920], 1),
+    '0.29': ([512, 1792], 1),
+    '0.32': ([544, 1696], 1),
+    '0.36': ([576, 1600], 1),
+    '0.40': ([608, 1504], 1),
+    '0.49': ([672, 1376], 1),
+    '0.54': ([704, 1312], 1),
+    '0.59': ([736, 1248], 1),
+    '0.69': ([800, 1152], 1),
+    '0.74': ([832, 1120], 1),
+    '0.82': ([864, 1056], 1),
+    '0.88': ([896, 1024], 1),
+    '0.94': ([928, 992], 1),
+    '1.00': ([960, 960], 1),
+    '1.07': ([992, 928], 1),
+    '1.14': ([1024, 896], 1),
+    '1.22': ([1056, 864], 1),
+    '1.31': ([1088, 832], 1),
+    '1.35': ([1120, 832], 1),
+    '1.44': ([1152, 800], 1),
+    '1.70': ([1248, 736], 1),
+    '2.00': ([1344, 672], 1),
+    '2.05': ([1376, 672], 1),
+    '2.47': ([1504, 608], 1),
+    '2.53': ([1536, 608], 1),
+    '2.83': ([1632, 576], 1),
+    '3.06': ([1664, 544], 1),
+    '3.12': ([1696, 544], 1),
+    '3.62': ([1856, 512], 1),
+    '3.93': ([1888, 480], 1),
+    '4.00': ([1920, 480], 1)
+}
+
+ASPECT_RATIO_960_F64 = {
+    '0.22': ([448, 2048], 1),
+    '0.29': ([512, 1792], 1),
+    '0.36': ([576, 1600], 1),
+    '0.45': ([640, 1408], 1),
+    '0.55': ([704, 1280], 1),
+    '0.63': ([768, 1216], 1),
+    '0.76': ([832, 1088], 1),
+    '0.88': ([896, 1024], 1),
+    '1.00': ([960, 960], 1),
+    '1.14': ([1024, 896], 1),
+    '1.31': ([1088, 832], 1),
+    '1.50': ([1152, 768], 1),
+    '1.58': ([1216, 768], 1),
+    '1.82': ([1280, 704], 1),
+    '1.91': ([1344, 704], 1),
+    '2.20': ([1408, 640], 1),
+    '2.30': ([1472, 640], 1),
+    '2.67': ([1536, 576], 1),
+    '2.89': ([1664, 576], 1),
+    '3.62': ([1856, 512], 1),
+    '3.75': ([1920, 512], 1)
+}
+
+ASPECT_RATIO_960_F128 = {
+    '0.20': ([384, 1920], 1),
+    '0.27': ([512, 1920], 1),
+    '0.33': ([512, 1536], 1),
+    '0.42': ([640, 1536], 1),
+    '0.50': ([640, 1280], 1),
+    '0.60': ([768, 1280], 1),
+    '0.67': ([768, 1152], 1),
+    '0.78': ([896, 1152], 1),
+    '1.00': ([1024, 1024], 1),
+    '1.29': ([1152, 896], 1),
+    '1.50': ([1152, 768], 1),
+    '1.67': ([1280, 768], 1),
+    '2.00': ([1280, 640], 1),
+    '2.40': ([1536, 640], 1),
+    '3.00': ([1536, 512], 1),
+    '3.75': ([1920, 512], 1),
+    '5.00': ([1920, 384], 1),
+}
+
+ASPECT_RATIO_960_F256 = {
+    '0.33': ([512, 1536], 1),
+    '0.60': ([768, 1280], 1),
+    '1.00': ([1024, 1024], 1),
+    '1.67': ([1280, 768], 1),
+    '3.00': ([1536, 512], 1),
+}
+
+
+def get_bucket_config(resolution, scale_factor_spatial):
+    if resolution == '480p':
+        if scale_factor_spatial == 16 or scale_factor_spatial == 32:
+            return ASPECT_RATIO_627
+        elif scale_factor_spatial == 64:
+            return ASPECT_RATIO_627_F64
+        elif scale_factor_spatial == 128:
+            return ASPECT_RATIO_627_F128
+        elif scale_factor_spatial == 256:
+            return ASPECT_RATIO_627_F256
+    elif resolution == '720p':
+        if scale_factor_spatial == 16 or scale_factor_spatial == 32:
+            return ASPECT_RATIO_960
+        elif scale_factor_spatial == 64:
+            return ASPECT_RATIO_960_F64
+        elif scale_factor_spatial == 128:
+            return ASPECT_RATIO_960_F128
+        elif scale_factor_spatial == 256:
+            return ASPECT_RATIO_960_F256
+
+    raise ValueError(
+        f"Unsupported resolution '{resolution}' or scale_factor_spatial '{scale_factor_spatial}'"
+    )

fastvideo/configs/pipelines/registry.py

@@ -9,6 +9,7 @@ from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
+from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 
 # isort: off
 from fastvideo.configs.pipelines.wan import (
@@ -77,11 +78,14 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "stepvideo" in id.lower(),
     "cosmos":
     lambda id: "cosmos" in id.lower(),
+    "longcat":
+    lambda id: "longcat" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
 # Fallback configs when exact match isn't found but architecture is detected
 PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
+    "longcat": LongCatT2V480PConfig,
     "hunyuan":
     HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
     "matrixgame": MatrixGameI2V480PConfig,

fastvideo/configs/sample/base.py

@@ -3,6 +3,7 @@ from dataclasses import dataclass
 from typing import Any
 
 from fastvideo.logger import init_logger
+from fastvideo.utils import StoreBoolean
 
 logger = init_logger(__name__)
 
@@ -27,6 +28,17 @@ class SamplingParam:
     keyboard_cond: Any | None = None  # Shape: (B, T, K)
     grid_sizes: Any | None = None  # Shape: (3,) [F,H,W]
 
+    # Refine inputs (LongCat 480p->720p upscaling)
+    # Path-based refine (load stage1 video from disk, e.g. MP4)
+    refine_from: str | None = None  # Path to stage1 video (480p output from distill)
+    t_thresh: float = 0.5  # Threshold for timestep scheduling in refinement
+    spatial_refine_only: bool = False  # If True, only spatial (no temporal doubling)
+    num_cond_frames: int = 0  # Number of conditioning frames
+    # In-memory refine input (for two-stage pipeline where stage1 frames are already in memory)
+    # This mirrors LongCat's demo where a list of frames (e.g. np.ndarray or PIL.Image)
+    # is passed directly to the refinement pipeline instead of reloading from disk.
+    stage1_video: Any | None = None
+
     # Text inputs
     prompt: str | list[str] | None = None
     negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
@@ -216,6 +228,31 @@ class SamplingParam:
             default=SamplingParam.video_path,
             help="Path to input video for video-to-video generation",
         )
+        parser.add_argument(
+            "--refine-from",
+            type=str,
+            default=SamplingParam.refine_from,
+            help="Path to stage1 video for refinement (LongCat 480p->720p)",
+        )
+        parser.add_argument(
+            "--t-thresh",
+            type=float,
+            default=SamplingParam.t_thresh,
+            help=
+            "Threshold for timestep scheduling in refinement (default: 0.5)",
+        )
+        parser.add_argument(
+            "--spatial-refine-only",
+            action=StoreBoolean,
+            default=SamplingParam.spatial_refine_only,
+            help="Only perform spatial super-resolution (no temporal doubling)",
+        )
+        parser.add_argument(
+            "--num-cond-frames",
+            type=int,
+            default=SamplingParam.num_cond_frames,
+            help="Number of conditioning frames for refinement",
+        )
         parser.add_argument(
             "--moba-config-path",
             type=str,

fastvideo/fastvideo_args.py

@@ -319,6 +319,62 @@ class FastVideoArgs:
             "Path to a text file containing prompts (one per line) for batch processing",
         )
 
+        # LoRA parameters (inference-time adapter loading)
+        parser.add_argument(
+            "--lora-path",
+            type=str,
+            default=FastVideoArgs.lora_path,
+            help=
+            "Path to a LoRA adapter (directory or HF repo id). If set, LoRA will be applied at inference.",
+        )
+        parser.add_argument(
+            "--lora-nickname",
+            type=str,
+            default=FastVideoArgs.lora_nickname,
+            help=
+            "Nickname to refer to the loaded LoRA adapter (useful for swapping).",
+        )
+        parser.add_argument(
+            "--lora-target-modules",
+            nargs="+",
+            type=str,
+            default=FastVideoArgs.lora_target_modules,
+            help=
+            "Optional list of module name substrings to restrict LoRA injection (e.g. q_proj k_proj v_proj).",
+        )
+
+        # BSA runtime control (LongCat)
+        parser.add_argument(
+            "--enable-bsa",
+            action=StoreBoolean,
+            help=
+            "Enable Block Sparse Attention (BSA) at runtime (overrides config).",
+        )
+        parser.add_argument(
+            "--bsa-sparsity",
+            type=float,
+            help="BSA sparsity (e.g., 0.9375).",
+        )
+        parser.add_argument(
+            "--bsa-cdf-threshold",
+            type=float,
+            help="BSA CDF threshold (optional).",
+        )
+        parser.add_argument(
+            "--bsa-chunk-q",
+            nargs=3,
+            type=int,
+            metavar=("T", "H", "W"),
+            help="BSA chunk_3d_shape_q as three ints, e.g., 4 4 4.",
+        )
+        parser.add_argument(
+            "--bsa-chunk-k",
+            nargs=3,
+            type=int,
+            metavar=("T", "H", "W"),
+            help="BSA chunk_3d_shape_k as three ints, e.g., 4 4 4.",
+        )
+
         # STA (Sliding Tile Attention) parameters
         parser.add_argument(
             "--STA-mode",

fastvideo/layers/rotary_embedding_3d.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+3D Rotary Position Embedding (RoPE) for video transformers.
+
+Reference: https://arxiv.org/pdf/2104.09864.pdf
+"""
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+
+
+def broadcast(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor:
+    """
+    Broadcast and concatenate tensors along a dimension.
+    """
+    num_tensors = len(tensors)
+    shape_lens = set(len(t.shape) for t in tensors)
+    assert len(
+        shape_lens) == 1, "tensors must all have the same number of dimensions"
+
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+
+    dims = list(zip(*[list(t.shape) for t in tensors], strict=False))
+    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+
+    assert all(
+        len(set(t[1])) <= 2 for t in
+        expandable_dims), "invalid dimensions for broadcastable concatenation"
+
+    max_dims = [(t[0], max(t[1])) for t in expandable_dims]
+    expanded_dims = [(t[0], (t[1], ) * num_tensors) for t in max_dims]
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*[t[1] for t in expanded_dims], strict=False))
+    tensors = [
+        t[0].expand(*t[1])
+        for t in zip(tensors, expandable_shapes, strict=False)
+    ]
+
+    return torch.cat(tensors, dim=dim)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """
+    Rotate half the hidden dims of the input.
+    """
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
+
+class RotaryPositionalEmbedding3D(nn.Module):
+    """
+    3D Rotary Positional Embedding for video transformers.
+    
+    Splits the head dimension across temporal, height, and width dimensions,
+    computing separate rotary embeddings for each and concatenating them.
+    """
+
+    def __init__(
+        self,
+        head_dim: int,
+        base: float = 10000.0,
+    ):
+        """
+        Args:
+            head_dim: Dimension of each attention head
+            base: Base value for exponential frequency
+        """
+        super().__init__()
+        self.head_dim = head_dim
+        assert self.head_dim % 8 == 0, "head_dim must be a multiple of 8 for 3D RoPE"
+        self.base = base
+
+        # Cache for precomputed frequencies
+        self.freqs_dict: dict[tuple, torch.Tensor] = {}
+
+    def register_grid_size(self, grid_size: tuple[int, int, int]) -> None:
+        """
+        Precompute and register frequencies for a given grid size.
+        
+        Args:
+            grid_size: (T, H, W) tuple of grid dimensions
+        """
+        if grid_size not in self.freqs_dict:
+            self.freqs_dict[grid_size] = self.precompute_freqs_3d(grid_size)
+
+    def precompute_freqs_3d(self, grid_size: tuple[int, int,
+                                                   int]) -> torch.Tensor:
+        """
+        Precompute 3D rotary frequencies.
+        
+        Args:
+            grid_size: (num_frames, height, width)
+            
+        Returns:
+            freqs: [T*H*W, head_dim] tensor of frequencies
+        """
+        num_frames, height, width = grid_size
+
+        # Split head_dim across 3 dimensions
+        # Temporal gets the remainder to ensure exact division
+        dim_t = self.head_dim - 4 * (self.head_dim // 6)
+        dim_h = 2 * (self.head_dim // 6)
+        dim_w = 2 * (self.head_dim // 6)
+
+        # Compute frequency bands for each dimension
+        freqs_t = 1.0 / (self.base**(
+            torch.arange(0, dim_t, 2)[:(dim_t // 2)].float() / dim_t))
+        freqs_h = 1.0 / (self.base**(
+            torch.arange(0, dim_h, 2)[:(dim_h // 2)].float() / dim_h))
+        freqs_w = 1.0 / (self.base**(
+            torch.arange(0, dim_w, 2)[:(dim_w // 2)].float() / dim_w))
+
+        # Create position grids
+        grid_t = torch.arange(num_frames, dtype=torch.float32)
+        grid_h = torch.arange(height, dtype=torch.float32)
+        grid_w = torch.arange(width, dtype=torch.float32)
+
+        # Compute frequencies for each position
+        freqs_t = torch.einsum("..., f -> ... f", grid_t, freqs_t)
+        freqs_h = torch.einsum("..., f -> ... f", grid_h, freqs_h)
+        freqs_w = torch.einsum("..., f -> ... f", grid_w, freqs_w)
+
+        # Duplicate for complex pair representation
+        freqs_t = repeat(freqs_t, "... n -> ... (n r)", r=2)
+        freqs_h = repeat(freqs_h, "... n -> ... (n r)", r=2)
+        freqs_w = repeat(freqs_w, "... n -> ... (n r)", r=2)
+
+        # Broadcast and concatenate across all 3 dimensions
+        freqs = broadcast(
+            [
+                freqs_t[:, None, None, :],  # [T, 1, 1, dim_t]
+                freqs_h[None, :, None, :],  # [1, H, 1, dim_h]
+                freqs_w[None, None, :, :],  # [1, 1, W, dim_w]
+            ],
+            dim=-1,
+        )
+
+        # Flatten spatial dimensions: [T, H, W, head_dim] -> [T*H*W, head_dim]
+        freqs = rearrange(freqs, "T H W D -> (T H W) D")
+
+        return freqs
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        grid_size: tuple[int, int, int],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Apply 3D rotary positional embedding to queries and keys.
+        
+        Args:
+            q: Query tensor [B, num_heads, seq_len, head_dim]
+            k: Key tensor [B, num_heads, seq_len, head_dim]
+            grid_size: (T, H, W) tuple of grid dimensions
+            
+        Returns:
+            (q_rotated, k_rotated): Rotated query and key tensors
+        """
+        # Register grid size if not cached
+        if grid_size not in self.freqs_dict:
+            self.register_grid_size(grid_size)
+
+        # Get cached frequencies
+        freqs_cis = self.freqs_dict[grid_size].to(q.device)
+
+        # Cast to float32 for precision
+        q_, k_ = q.float(), k.float()
+        freqs_cis = freqs_cis.float()
+
+        # Compute cos and sin
+        cos = freqs_cis.cos()
+        sin = freqs_cis.sin()
+
+        # Reshape for broadcasting: [1, 1, seq_len, head_dim]
+        cos = rearrange(cos, "n d -> 1 1 n d")
+        sin = rearrange(sin, "n d -> 1 1 n d")
+
+        # Apply rotation
+        q_ = (q_ * cos) + (rotate_half(q_) * sin)
+        k_ = (k_ * cos) + (rotate_half(k_) * sin)
+
+        # Cast back to original dtype
+        return q_.type_as(q), k_.type_as(k)
+
+
+def apply_rotary_emb_3d(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    rope_module: RotaryPositionalEmbedding3D,
+    grid_size: tuple[int, int, int],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Convenience function to apply 3D RoPE.
+    
+    Args:
+        q: Query tensor [B, num_heads, seq_len, head_dim]
+        k: Key tensor [B, num_heads, seq_len, head_dim]
+        rope_module: RotaryPositionalEmbedding3D module
+        grid_size: (T, H, W) grid dimensions
+        
+    Returns:
+        (q_rotated, k_rotated): Rotated tensors
+    """
+    return rope_module(q, k, grid_size)

fastvideo/models/dits/longcat.py

@@ -0,0 +1,866 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Native LongCat Video DiT implementation using FastVideo conventions.
+
+This is a Phase 2 reimplementation that replaces the third_party wrapper
+with native FastVideo layers for better performance and integration.
+"""
+
+from typing import Any
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from fastvideo.configs.models.dits import LongCatVideoConfig
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.layernorm import RMSNorm, FP32LayerNorm
+from fastvideo.layers.activation import get_act_fn
+from fastvideo.layers.rotary_embedding_3d import RotaryPositionalEmbedding3D
+from fastvideo.attention.layer import DistributedAttention, LocalAttention
+from fastvideo.models.dits.base import CachableDiT
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.third_party.longcat_video.block_sparse_attention.bsa_interface import flash_attn_bsa_3d
+
+
+# ============================================================================
+# Embeddings
+# ============================================================================
+
+class PatchEmbed3D(nn.Module):
+    """
+    3D patch embedding using Conv3d.
+    """
+
+    def __init__(
+        self,
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        in_channels: int = 16,
+        embed_dim: int = 4096,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv3d(
+            in_channels,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, C, T, H, W]
+        Returns:
+            [B, N, C] where N = (T/pt) * (H/ph) * (W/pw)
+        """
+        # Padding if needed
+        _, _, T, H, W = x.shape
+        if W % self.patch_size[2] != 0:
+            x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
+        if H % self.patch_size[1] != 0:
+            x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
+        if T % self.patch_size[0] != 0:
+            x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - T % self.patch_size[0]))
+
+        x = self.proj(x)  # [B, C, T', H', W']
+        x = x.flatten(2).transpose(1, 2)  # [B, N, C]
+        return x
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Sinusoidal timestep embedding + MLP projection.
+    """
+
+    def __init__(
+        self,
+        frequency_embedding_size: int = 256,
+        adaln_tembed_dim: int = 512,
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+        self.frequency_embedding_size = frequency_embedding_size
+
+        # Use FastVideo's ReplicatedLinear
+        self.linear_1 = ReplicatedLinear(
+            frequency_embedding_size,
+            adaln_tembed_dim,
+            bias=True,
+            params_dtype=dtype,
+        )
+        self.act = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(
+            adaln_tembed_dim,
+            adaln_tembed_dim,
+            bias=True,
+            params_dtype=dtype,
+        )
+
+    @staticmethod
+    def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
+        """
+        Create sinusoidal timestep embeddings.
+        """
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
+        )
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t: torch.Tensor, latent_shape: tuple | None = None) -> torch.Tensor:
+        """
+        Args:
+            t: [B] or [B, T] timesteps
+            latent_shape: (T, H, W) for temporal expansion
+        Returns:
+            [B, T, C]
+        """
+        # Sinusoidal embedding in FP32
+        t_freq = self.timestep_embedding(t.flatten(), self.frequency_embedding_size)
+        
+        # Cast to model dtype before MLP
+        # Handle LoRA wrapper if present
+        linear_layer = self.linear_1.base_layer if hasattr(self.linear_1, 'base_layer') else self.linear_1
+        target_dtype = linear_layer.weight.dtype
+        if t_freq.dtype != target_dtype:
+            t_freq = t_freq.to(target_dtype)
+        
+        # MLP projection
+        t_emb, _ = self.linear_1(t_freq)
+        t_emb = self.act(t_emb)
+        t_emb, _ = self.linear_2(t_emb)
+        
+        # Reshape if needed
+        if latent_shape is not None and len(t.shape) > 1:
+            B = t.shape[0]
+            T = latent_shape[0]
+            t_emb = t_emb.reshape(B, T, -1)
+        
+        return t_emb
+
+
+class CaptionEmbedder(nn.Module):
+    """
+    Caption embedding with MLP projection and optional text compaction.
+    """
+
+    def __init__(
+        self,
+        caption_channels: int = 4096,
+        hidden_size: int = 4096,
+        text_tokens_zero_pad: bool = True,
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+        self.text_tokens_zero_pad = text_tokens_zero_pad
+
+        # Two-layer MLP using ReplicatedLinear
+        self.linear_1 = ReplicatedLinear(
+            caption_channels,
+            hidden_size,
+            bias=True,
+            params_dtype=dtype,
+        )
+        self.act = nn.SiLU()
+        self.linear_2 = ReplicatedLinear(
+            hidden_size,
+            hidden_size,
+            bias=True,
+            params_dtype=dtype,
+        )
+
+    def forward(
+        self,
+        encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            encoder_hidden_states: [B, N_text, C_text] or [B, 1, N_text, C_text]
+            encoder_attention_mask: [B, N_text] or [B, 1, 1, N_text]
+        Returns:
+            y: [B, N_text, C] - standard padded representation (like other models)
+        """
+        # Handle extra dimension from wrapper
+        if len(encoder_hidden_states.shape) == 4:
+            encoder_hidden_states = encoder_hidden_states.squeeze(1)
+
+        # Project
+        y, _ = self.linear_1(encoder_hidden_states)
+        y = self.act(y)
+        y, _ = self.linear_2(y)  # [B, N_text, C]
+        
+        # Handle attention masking - just zero out padded tokens if requested
+        if encoder_attention_mask is not None:
+            # Remove extra dimensions
+            if len(encoder_attention_mask.shape) == 4:
+                encoder_attention_mask = encoder_attention_mask.squeeze(1).squeeze(1)
+            elif len(encoder_attention_mask.shape) == 3:
+                encoder_attention_mask = encoder_attention_mask.squeeze(1)
+            
+            # Zero out padded tokens if requested
+            if self.text_tokens_zero_pad:
+                y = y * encoder_attention_mask.unsqueeze(-1)
+        
+        # Return standard format [B, N_text, C] - no compaction!
+        return y
+
+
+# ============================================================================
+# Attention Modules (Placeholders for now)
+# ============================================================================
+
+class LongCatSelfAttention(nn.Module):
+    """
+    Self-attention with 3D RoPE support and optional BSA.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        config: LongCatVideoConfig,
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        # Separate Q/K/V projections (not fused like original)
+        self.to_q = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
+        self.to_k = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
+        self.to_v = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
+
+        # Per-head RMS normalization
+        self.q_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32)
+        self.k_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32)
+
+        # Output projection
+        self.to_out = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
+
+        # 3D RoPE
+        self.rope_3d = RotaryPositionalEmbedding3D(head_dim=self.head_dim)
+
+        # BSA configuration
+        self.enable_bsa = getattr(config, 'enable_bsa', False)
+        self.bsa_params = getattr(config, 'bsa_params', None)
+
+        # FastVideo attention backend (used when BSA is disabled)
+        self.attn = DistributedAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            supported_attention_backends=config._supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,              # [B, N, C]
+        latent_shape: tuple,          # (T, H, W)
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Forward pass with 3D RoPE and optional BSA.
+        """
+        B, N, C = x.shape
+        T, H, W = latent_shape
+
+        # Project to Q/K/V
+        q, _ = self.to_q(x)
+        k, _ = self.to_k(x)
+        v, _ = self.to_v(x)
+
+        # Reshape to heads: [B, N, num_heads, head_dim]
+        q = q.view(B, N, self.num_heads, self.head_dim)
+        k = k.view(B, N, self.num_heads, self.head_dim)
+        v = v.view(B, N, self.num_heads, self.head_dim)
+
+        # Per-head RMS normalization
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # For RoPE: need [B, num_heads, N, head_dim]
+        q_rope = q.transpose(1, 2)
+        k_rope = k.transpose(1, 2)
+
+        # Apply 3D RoPE
+        q_rope, k_rope = self.rope_3d(q_rope, k_rope, grid_size=latent_shape)
+
+        # Transpose back: [B, N, num_heads, head_dim] or [B, H, N, D] for BSA
+        q = q_rope.transpose(1, 2)
+        k = k_rope.transpose(1, 2)
+
+        # === Attention: BSA or standard ===
+        if self.enable_bsa and T > 1:  # Only use BSA for multi-frame videos
+            # BSA expects [B, H, S, D] format
+            q_bsa = q.transpose(1, 2).contiguous()  # [B, num_heads, N, head_dim]
+            k_bsa = k.transpose(1, 2).contiguous()
+            v_bsa = v.transpose(1, 2).contiguous()
+            
+            # Handle SP split: BSA operates on per-rank spatial dimensions
+            # Replicate LongCat's cp_split_hw logic exactly
+            from fastvideo.distributed.parallel_state import get_sp_world_size
+            sp_size = get_sp_world_size()
+            if sp_size > 1:
+                # Calculate optimal 2D split (same as LongCat's get_optimal_split)
+                factors = []
+                for i in range(1, int(sp_size**0.5) + 1):
+                    if sp_size % i == 0:
+                        factors.append([i, sp_size // i])
+                cp_split_hw = min(factors, key=lambda x: abs(x[0] - x[1]))
+                
+                # Split H and W dimensions by their respective factors
+                T_bsa, H_bsa, W_bsa = latent_shape
+                assert H_bsa % cp_split_hw[0] == 0 and W_bsa % cp_split_hw[1] == 0, \
+                    f"H {H_bsa} must be divisible by {cp_split_hw[0]}, W {W_bsa} must be divisible by {cp_split_hw[1]}"
+                H_bsa = H_bsa // cp_split_hw[0]
+                W_bsa = W_bsa // cp_split_hw[1]
+                latent_shape_bsa = (T_bsa, H_bsa, W_bsa)
+            else:
+                latent_shape_bsa = latent_shape
+            
+            # Call BSA with per-rank latent shape
+            out = flash_attn_bsa_3d(
+                q_bsa, k_bsa, v_bsa,
+                latent_shape_q=latent_shape_bsa,
+                latent_shape_k=latent_shape_bsa,
+                **self.bsa_params
+            )  # [B, num_heads, N, head_dim]
+            
+            # Transpose back: [B, N, num_heads, head_dim]
+            out = out.transpose(1, 2)
+        else:
+            # Standard attention: [B, N, num_heads, head_dim]
+            out, _ = self.attn(q, k, v)
+
+        # Reshape and project out
+        out = out.reshape(B, N, C)
+        out, _ = self.to_out(out)
+
+        return out
+
+
+class LongCatCrossAttention(nn.Module):
+    """
+    Cross-attention for text conditioning (standard implementation like other models).
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        config: LongCatVideoConfig,
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        # Separate Q/K/V projections
+        self.to_q = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
+        self.to_k = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
+        self.to_v = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
+
+        # Per-head RMS normalization
+        self.q_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32)
+        self.k_norm = RMSNorm(self.head_dim, eps=1e-6, dtype=dtype or torch.float32)
+
+        # Output projection
+        self.to_out = ReplicatedLinear(dim, dim, bias=True, params_dtype=dtype)
+
+        # Cross-attention uses LocalAttention (FastVideo standard)
+        self.attn = LocalAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=config.arch_config._supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,              # [B, N_img, C]
+        context: torch.Tensor,        # [B, N_text, C]
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Forward pass for cross-attention (standard implementation).
+        
+        Args:
+            x: Image tokens [B, N_img, C]
+            context: Text tokens [B, N_text, C] (standard padded format)
+        """
+        B, N_img, C = x.shape
+
+        # Project Q, K, V (standard cross-attention like WanVideo/StepVideo/Cosmos)
+        q, _ = self.to_q(x)
+        k, _ = self.to_k(context)
+        v, _ = self.to_v(context)
+        
+        N_text = context.shape[1]
+
+        # Reshape to heads
+        q = q.view(B, N_img, self.num_heads, self.head_dim)
+        k = k.view(B, N_text, self.num_heads, self.head_dim)
+        v = v.view(B, N_text, self.num_heads, self.head_dim)
+
+        # Per-head RMS normalization
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Run cross-attention using FastVideo's LocalAttention
+        # LocalAttention handles different q and k/v sequence lengths automatically
+        out = self.attn(q, k, v)  # [B, N_img, num_heads, head_dim]
+
+        # Reshape and project out
+        out = out.reshape(B, N_img, C)
+        out, _ = self.to_out(out)
+
+        return out
+
+
+# ============================================================================
+# Feed-Forward Network
+# ============================================================================
+
+class LongCatSwiGLUFFN(nn.Module):
+    """
+    SwiGLU feed-forward network using FastVideo's ReplicatedLinear.
+    
+    FFN(x) = down(gate(x) * SiLU(up(x)))
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+
+        # Three projections for SwiGLU (no bias as per original)
+        self.w1 = ReplicatedLinear(dim, hidden_dim, bias=False, params_dtype=dtype)  # gate
+        self.w3 = ReplicatedLinear(dim, hidden_dim, bias=False, params_dtype=dtype)  # up
+        self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False, params_dtype=dtype)  # down
+
+        self.act = nn.SiLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass: SiLU(w1(x)) * w3(x) -> w2 (matching original LongCat)
+        """
+        w1_out, _ = self.w1(x)
+        w3_out, _ = self.w3(x)
+        combined = self.act(w1_out) * w3_out
+        out, _ = self.w2(combined)
+        return out
+
+
+# ============================================================================
+# Modulation Utilities
+# ============================================================================
+
+def modulate_fp32(norm: nn.Module, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """
+    Apply modulation in FP32 for numerical stability (matching original LongCat).
+    
+    shift and scale should already be FP32 from torch.amp.autocast context.
+    """
+    # Ensure modulation params are FP32 (should be from autocast)
+    assert shift.dtype == torch.float32 and scale.dtype == torch.float32, \
+        f"shift and scale must be FP32, got {shift.dtype} and {scale.dtype}"
+    
+    orig_dtype = x.dtype
+    
+    # Normalize and modulate in FP32
+    x_norm = norm(x.to(torch.float32))
+    x_mod = x_norm * (scale + 1) + shift
+    
+    return x_mod.to(orig_dtype)
+
+
+# ============================================================================
+# Transformer Block
+# ============================================================================
+
+class LongCatTransformerBlock(nn.Module):
+    """
+    Single-stream transformer block with:
+    - AdaLN modulation (FP32)
+    - Self-attention
+    - Cross-attention
+    - SwiGLU FFN
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: int,
+        adaln_tembed_dim: int,
+        config: LongCatVideoConfig,
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+
+        # AdaLN modulation (6 parameters: scale/shift for attn & ffn, gate for residual)
+        self.adaln_linear_1 = ReplicatedLinear(
+            adaln_tembed_dim,
+            6 * hidden_size,
+            bias=True,
+            params_dtype=dtype,
+        )
+        self.adaln_act = nn.SiLU()
+
+        # Normalization layers (CRITICAL: Use LayerNorm not RMSNorm like original!)
+        # Original LongCat uses LayerNorm_FP32 with elementwise_affine=False
+        self.norm_attn = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
+        self.norm_ffn = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
+        # Cross-attention norm has elementwise_affine=True (has weight and bias)
+        self.norm_cross = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=True)
+
+        # Self-attention
+        self.self_attn = LongCatSelfAttention(
+            dim=hidden_size,
+            num_heads=num_heads,
+            config=config,
+            dtype=dtype,
+        )
+
+        # Cross-attention
+        self.cross_attn = LongCatCrossAttention(
+            dim=hidden_size,
+            num_heads=num_heads,
+            config=config,
+            dtype=dtype,
+        )
+
+        # SwiGLU FFN
+        ffn_hidden_dim = int(hidden_size * mlp_ratio * 2 / 3)
+        # Round up to nearest multiple of 256
+        ffn_hidden_dim = 256 * ((ffn_hidden_dim + 255) // 256)
+
+        self.ffn = LongCatSwiGLUFFN(
+            dim=hidden_size,
+            hidden_dim=ffn_hidden_dim,
+            dtype=dtype,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,              # [B, N, C]
+        context: torch.Tensor,        # [B, N_text, C]
+        t: torch.Tensor,              # [B, T, C_t]
+        latent_shape: tuple,          # (T, H, W)
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Forward pass with AdaLN modulation.
+        """
+        B, N, C = x.shape
+        T, H, W = latent_shape
+        x_orig_dtype = x.dtype  # Save for later casting
+
+        # === AdaLN Modulation (CRITICAL: FP32 for stability like original) ===
+        # Use autocast to compute modulation params in FP32
+        with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
+            t_mod = self.adaln_act(t)
+            mod_params, _ = self.adaln_linear_1(t_mod)
+            # Ensure FP32 output (needed when LoRA is applied)
+            if mod_params.dtype != torch.float32:
+                mod_params = mod_params.float()
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
+                mod_params.unsqueeze(2).chunk(6, dim=-1)  # [B, T, 1, C]
+
+        # === Self-Attention ===
+        x_norm = modulate_fp32(self.norm_attn, x.view(B, T, -1, C), shift_msa, scale_msa)
+        x_norm = x_norm.view(B, N, C)
+
+        attn_out = self.self_attn(x_norm, latent_shape=latent_shape)
+
+        # Residual with gating (CRITICAL: FP32 like original, then cast back)
+        with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
+            x = x + (gate_msa * attn_out.view(B, T, -1, C)).view(B, N, C)
+        x = x.to(x_orig_dtype)
+
+        # === Cross-Attention ===
+        x_norm_cross = self.norm_cross(x)
+        cross_out = self.cross_attn(x_norm_cross, context)
+        x = x + cross_out
+
+        # === FFN ===
+        x_norm_ffn = modulate_fp32(self.norm_ffn, x.view(B, T, -1, C), shift_mlp, scale_mlp)
+        x_norm_ffn = x_norm_ffn.view(B, N, C)
+
+        ffn_out = self.ffn(x_norm_ffn)
+
+        # Residual with gating (CRITICAL: FP32 like original, then cast back)
+        with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
+            x = x + (gate_mlp * ffn_out.view(B, T, -1, C)).view(B, N, C)
+        x = x.to(x_orig_dtype)
+
+        return x
+
+
+# ============================================================================
+# Final Layer
+# ============================================================================
+
+class FinalLayer(nn.Module):
+    """
+    Final output projection with AdaLN modulation.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        out_channels: int,
+        adaln_tembed_dim: int,
+        patch_size: tuple[int, int, int],
+        dtype: torch.dtype | None = None,
+    ):
+        super().__init__()
+
+        # AdaLN for final layer (2 parameters: scale and shift)
+        self.adaln_linear = ReplicatedLinear(
+            adaln_tembed_dim,
+            2 * hidden_size,
+            bias=True,
+            params_dtype=dtype,
+        )
+        self.adaln_act = nn.SiLU()
+
+        # CRITICAL: Use LayerNorm not RMSNorm! (matches original)
+        self.norm = FP32LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
+
+        # Output projection
+        num_patch = patch_size[0] * patch_size[1] * patch_size[2]
+        self.proj = ReplicatedLinear(
+            hidden_size,
+            num_patch * out_channels,
+            bias=True,
+            params_dtype=dtype,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,  # [B, N, C]
+        t: torch.Tensor,  # [B, T, C_t]
+        latent_shape: tuple,
+    ) -> torch.Tensor:
+        """
+        Returns: [B, N, out_channels * patch_size^3]
+        """
+        B, N, C = x.shape
+        T, _, _ = latent_shape
+
+        # AdaLN modulation (FP32 for stability like original)
+        with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
+            t_mod = self.adaln_act(t)
+            mod_params, _ = self.adaln_linear(t_mod)
+            # Ensure FP32 output (needed when LoRA is applied)
+            if mod_params.dtype != torch.float32:
+                mod_params = mod_params.float()
+            shift, scale = mod_params.unsqueeze(2).chunk(2, dim=-1)
+
+        # Modulate
+        x = modulate_fp32(self.norm, x.view(B, T, -1, C), shift, scale)
+        x = x.reshape(B, N, C)
+
+        # Project
+        x, _ = self.proj(x)
+
+        return x
+
+
+# ============================================================================
+# Main Model
+# ============================================================================
+
+class LongCatTransformer3DModel(CachableDiT):
+    """
+    Native LongCat Video Transformer using FastVideo layers.
+    
+    This is a Phase 2 implementation that replaces third_party dependencies.
+    """
+
+    # FSDP sharding: shard at each transformer block
+    _fsdp_shard_conditions = [
+        lambda n, m: "blocks" in n and n.split(".")[-1].isdigit(),
+    ]
+    # torch.compile optimization: compile each transformer block for speedup
+    _compile_conditions = [
+        lambda n, m: "blocks" in n and n.split(".")[-1].isdigit(),
+    ]
+
+    # Parameter name mapping (for weight conversion)
+    param_names_mapping = {}  # Will be defined in config
+    reverse_param_names_mapping = {}
+    lora_param_names_mapping = {}
+
+    # Supported attention backends
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    def __init__(self, config: LongCatVideoConfig, hf_config: dict[str, Any]):
+        super().__init__(config=config, hf_config=hf_config)
+
+        # Extract architecture parameters
+        self.hidden_size = config.hidden_size  # 4096
+        self.num_attention_heads = config.num_attention_heads  # 32
+        self.depth = config.depth  # 48
+        self.mlp_ratio = config.mlp_ratio  # 4
+        self.in_channels = config.in_channels  # 16
+        self.out_channels = config.out_channels  # 16
+        self.num_channels_latents = config.in_channels
+        self.patch_size = config.patch_size  # [1, 2, 2]
+
+        # Embeddings
+        self.patch_embed = PatchEmbed3D(
+            patch_size=self.patch_size,
+            in_channels=self.in_channels,
+            embed_dim=self.hidden_size,
+        )
+
+        self.time_embedder = TimestepEmbedder(
+            frequency_embedding_size=config.frequency_embedding_size,
+            adaln_tembed_dim=config.adaln_tembed_dim,
+        )
+
+        self.caption_embedder = CaptionEmbedder(
+            caption_channels=config.caption_channels,
+            hidden_size=self.hidden_size,
+            text_tokens_zero_pad=getattr(config, 'text_tokens_zero_pad', True),
+        )
+
+        # Transformer blocks (48 blocks)
+        self.blocks = nn.ModuleList([
+            LongCatTransformerBlock(
+                hidden_size=self.hidden_size,
+                num_heads=self.num_attention_heads,
+                mlp_ratio=self.mlp_ratio,
+                adaln_tembed_dim=config.adaln_tembed_dim,
+                config=config,
+            )
+            for _ in range(self.depth)
+        ])
+
+        # Output projection
+        self.final_layer = FinalLayer(
+            hidden_size=self.hidden_size,
+            out_channels=self.out_channels,
+            adaln_tembed_dim=config.adaln_tembed_dim,
+            patch_size=self.patch_size,
+        )
+
+    def enable_bsa(self):
+        """Enable BSA for all self-attention layers."""
+        for block in self.blocks:
+            block.self_attn.enable_bsa = True
+    
+    def disable_bsa(self):
+        """Disable BSA for all self-attention layers."""
+        for block in self.blocks:
+            block.self_attn.enable_bsa = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,           # [B, C, T, H, W]
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],   # [B, N_text, C_text]
+        timestep: torch.LongTensor,            # [B] or [B, T]
+        encoder_attention_mask: torch.Tensor | None = None,  # [B, N_text]
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+        guidance: float | None = None,         # Unused, for API compatibility
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Forward pass with FastVideo parameter ordering.
+        
+        NOTE: This follows FastVideo convention:
+              (hidden_states, encoder_hidden_states, timestep)
+        """
+        B, _, T, H, W = hidden_states.shape
+
+        N_t = T // self.patch_size[0]
+        N_h = H // self.patch_size[1]
+        N_w = W // self.patch_size[2]
+
+        # Handle list of encoder outputs (take first one)
+        if isinstance(encoder_hidden_states, list):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        # 1. Patch embedding
+        x = self.patch_embed(hidden_states)  # [B, N, C]
+
+        # 2. Timestep embedding
+        # Expand timestep from [B] to [B, T] if needed
+        if timestep.ndim == 1:
+            timestep = timestep.unsqueeze(1).expand(-1, N_t)  # [B, T]
+        
+        t = self.time_embedder(timestep.flatten(), latent_shape=(N_t, N_h, N_w))
+        if t.ndim == 2:
+            t = t.reshape(B, N_t, -1)  # [B, T, C_t]
+
+        # 3. Caption embedding (standard format, no compaction)
+        context = self.caption_embedder(
+            encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask
+        )  # [B, N_text, C]
+
+        # 4. Transformer blocks
+        for i, block in enumerate(self.blocks):
+            x = block(
+                x, context, t,
+                latent_shape=(N_t, N_h, N_w)
+            )
+
+        # 5. Output projection
+        output = self.final_layer(x, t, latent_shape=(N_t, N_h, N_w))
+
+        # Reshape to [B, C_out, T, H, W]
+        output = self.unpatchify(output, N_t, N_h, N_w)
+
+        # Cast to float32 for better accuracy (as per original)
+        output = output.to(torch.float32)
+
+        return output
+
+    def unpatchify(self, x: torch.Tensor, N_t: int, N_h: int, N_w: int) -> torch.Tensor:
+        """
+        Args:
+            x: [B, N, C] where C = T_p * H_p * W_p * C_out
+        Returns:
+            [B, C_out, T, H, W]
+        """
+        T_p, H_p, W_p = self.patch_size
+        x = rearrange(
+            x,
+            "B (N_t N_h N_w) (T_p H_p W_p C_out) -> B C_out (N_t T_p) (N_h H_p) (N_w W_p)",
+            N_t=N_t,
+            N_h=N_h,
+            N_w=N_w,
+            T_p=T_p,
+            H_p=H_p,
+            W_p=W_p,
+            C_out=self.out_channels,
+        )
+        return x
+

fastvideo/models/loader/component_loader.py

@@ -16,6 +16,7 @@ import torch.nn as nn
 from safetensors.torch import load_file as safetensors_load_file
 from torch.distributed import init_device_mesh
 from transformers import AutoImageProcessor, AutoTokenizer
+from transformers import UMT5EncoderModel
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
 
 from fastvideo.configs.models import EncoderConfig
@@ -405,7 +406,7 @@ class VAELoader(ComponentLoader):
             target_device = get_local_torch_device()
 
         with set_default_torch_dtype(PRECISION_TO_TYPE[
-                fastvideo_args.pipeline_config.vae_precision]):
+                fastvideo_args.pipeline_config.vae_precision] if fastvideo_args.pipeline_config.vae_precision else torch.bfloat16):
             vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
             vae = vae_cls(vae_config).to(target_device)
 

fastvideo/models/registry.py

@@ -29,7 +29,9 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"),
-    "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel")
+    "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
+    "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
+    "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),  # Native (Phase 2)
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {

fastvideo/pipelines/basic/longcat/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LongCat pipeline module."""
+
+from fastvideo.pipelines.basic.longcat.longcat_pipeline import LongCatPipeline
+
+__all__ = ["LongCatPipeline"]

fastvideo/pipelines/basic/longcat/longcat_pipeline.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat video diffusion pipeline implementation (Phase 1: Wrapper).
+
+This module contains a wrapper implementation of the LongCat video diffusion pipeline
+using FastVideo's modular pipeline architecture with the original LongCat modules.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.stages import (
+    DecodingStage,
+    InputValidationStage,
+    LatentPreparationStage,
+    TextEncodingStage,
+    TimestepPreparationStage,
+)
+from fastvideo.pipelines.stages.longcat_denoising import LongCatDenoisingStage
+from fastvideo.pipelines.stages.longcat_refine_init import LongCatRefineInitStage
+from fastvideo.pipelines.stages.longcat_refine_timestep import LongCatRefineTimestepStage
+
+logger = init_logger(__name__)
+
+
+class LongCatPipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    LongCat video diffusion pipeline with LoRA support.
+    
+    Phase 1 implementation using wrapper modules from third_party/longcat_video.
+    This validates the pipeline infrastructure before full FastVideo integration.
+    """
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        """Initialize LongCat-specific components."""
+
+        # Enable BSA (Block Sparse Attention) if configured
+        pipeline_config = fastvideo_args.pipeline_config
+        transformer = self.get_module("transformer", None)
+        if transformer is None:
+            raise RuntimeError(
+                "Transformer module not found during initializing LongCat pipeline."
+            )
+        # If user toggles BSA via CLI/config
+        if pipeline_config.enable_bsa:
+            # Build effective BSA params:
+            # 1) from explicit CLI overrides if provided
+            # 2) else from pipeline_config.bsa_params
+            # 3) else fall back to reasonable defaults
+            bsa_params_cfg = pipeline_config.bsa_params
+            sparsity = pipeline_config.bsa_sparsity
+            cdf_threshold = pipeline_config.bsa_cdf_threshold
+            chunk_q = pipeline_config.bsa_chunk_q
+            chunk_k = pipeline_config.bsa_chunk_k
+
+            effective_bsa_params = dict(bsa_params_cfg) if isinstance(
+                bsa_params_cfg, dict) else {}
+            if sparsity is not None:
+                effective_bsa_params['sparsity'] = sparsity
+            if cdf_threshold is not None:
+                effective_bsa_params['cdf_threshold'] = cdf_threshold
+            if chunk_q is not None:
+                effective_bsa_params['chunk_3d_shape_q'] = chunk_q
+            if chunk_k is not None:
+                effective_bsa_params['chunk_3d_shape_k'] = chunk_k
+            # Provide defaults if still missing
+            effective_bsa_params.setdefault('sparsity', 0.9375)
+            effective_bsa_params.setdefault('chunk_3d_shape_q', [4, 4, 4])
+            effective_bsa_params.setdefault('chunk_3d_shape_k', [4, 4, 4])
+
+            if hasattr(transformer, 'enable_bsa'):
+                logger.info(
+                    "Enabling Block Sparse Attention (BSA) for LongCat transformer"
+                )
+                transformer.enable_bsa()
+                # Propagate params to all attention modules
+                if hasattr(transformer, 'blocks'):
+                    try:
+                        for blk in transformer.blocks:
+                            if hasattr(blk, 'self_attn'):
+                                blk.self_attn.bsa_params = effective_bsa_params
+                    except Exception as e:
+                        logger.warning(
+                            "Failed to set BSA params on all blocks: %s", e)
+                logger.info("BSA parameters in effect: %s",
+                            effective_bsa_params)
+            else:
+                logger.warning(
+                    "BSA is enabled in config but transformer does not support it"
+                )
+        else:
+            # Explicitly disable if present
+            if hasattr(transformer, 'disable_bsa'):
+                transformer.disable_bsa()
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+        # Add refine initialization stage (will be skipped if not refining)
+        self.add_stage(stage_name="longcat_refine_init_stage",
+                       stage=LongCatRefineInitStage(vae=self.get_module("vae")))
+
+        # First prepare generic timesteps (for non-refine paths)
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        # Then override timesteps for refinement (will be a no-op if not refining),
+        # matching LongCat's generate_refine schedule.
+        self.add_stage(stage_name="longcat_refine_timestep_stage",
+                       stage=LongCatRefineTimestepStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer", None)))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=LongCatDenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           transformer_2=self.get_module("transformer_2", None),
+                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae"),
+                           pipeline=self))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae"),
+                                           pipeline=self))
+
+
+EntryClass = LongCatPipeline

fastvideo/pipelines/composed_pipeline_base.py

@@ -287,6 +287,7 @@ class ComposedPipelineBase(ABC):
         # remove keys that are not pipeline modules
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
+        model_index.pop("workload_type", None)
         if "boundary_ratio" in model_index and model_index[
                 "boundary_ratio"] is not None:
             logger.info(

fastvideo/pipelines/pipeline_batch_info.py

@@ -91,6 +91,14 @@ class ForwardBatch:
     video_path: str | None = None
     video_latent: torch.Tensor | None = None
 
+    # Refine inputs (LongCat)
+    refine_from: str | None = None
+    t_thresh: float = 0.5
+    spatial_refine_only: bool = False
+    num_cond_frames: int = 0
+    stage1_video: list[
+        PIL.Image.Image] | None = None  # Loaded frames from refine_from
+
     # Primary encoder embeddings
     prompt_embeds: list[torch.Tensor] = field(default_factory=list)
     negative_prompt_embeds: list[torch.Tensor] | None = None

fastvideo/pipelines/pipeline_registry.py

@@ -29,6 +29,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "Cosmos2VideoToWorldPipeline": "cosmos",
     "MatrixGamePipeline": "matrixgame",
     "MatrixGameCausalDMDPipeline": "matrixgame",
+    "LongCatPipeline": "longcat",
 }
 
 _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = {

fastvideo/pipelines/stages/decoding.py

@@ -187,6 +187,21 @@ class DecodingStage(PipelineStage):
         # Convert to CPU float32 for compatibility
         frames = frames.cpu().float()
 
+        # Crop padding if this is a LongCat refinement
+        if hasattr(batch, 'num_cond_frames_added') and hasattr(
+                batch, 'new_frame_size_before_padding'):
+            num_cond_frames_added = batch.num_cond_frames_added
+            new_frame_size = batch.new_frame_size_before_padding
+            if num_cond_frames_added > 0 or frames.shape[2] != new_frame_size:
+                # frames is [B, C, T, H, W], crop temporal dimension
+                frames = frames[:, :,
+                                num_cond_frames_added:num_cond_frames_added +
+                                new_frame_size, :, :]
+                logger.info(
+                    "Cropped LongCat refinement padding: %s:%s, final shape: %s",
+                    num_cond_frames_added,
+                    num_cond_frames_added + new_frame_size, frames.shape)
+
         # Update batch with decoded image
         batch.output = frames
 

fastvideo/pipelines/stages/latent_preparation.py

@@ -136,16 +136,26 @@ class LatentPreparationStage(PipelineStage):
             )
         # Generate or use provided latents
         if latents is None:
-            latents = randn_tensor(shape,
-                                   generator=generator,
-                                   device=device,
-                                   dtype=dtype)
+            latents = randn_tensor(
+                shape,
+                generator=generator,
+                device=device,
+                dtype=dtype,
+            )
+            if hasattr(self.scheduler, "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
         else:
+            # Pre-initialized latents:
+            # - For LongCat refine (refine_from or stage1_video present), we should not re-scale by init_noise_sigma.
+            # - For other models, keep the original behavior.
             latents = latents.to(device)
+            is_longcat_refine = (batch.refine_from
+                                 is not None) or (batch.stage1_video
+                                                  is not None)
+            if (not is_longcat_refine) and hasattr(self.scheduler,
+                                                   "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
 
-        # Scale the initial noise if needed
-        if hasattr(self.scheduler, "init_noise_sigma"):
-            latents = latents * self.scheduler.init_noise_sigma
         # Update batch with prepared latents
         batch.latents = latents
         batch.raw_latent_shape = bcthw_shape

fastvideo/pipelines/stages/longcat_denoising.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat-specific denoising stage implementing CFG-zero optimized guidance.
+"""
+
+import torch
+from tqdm import tqdm
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising import DenoisingStage
+from fastvideo.forward_context import set_forward_context
+
+logger = init_logger(__name__)
+
+
+class LongCatDenoisingStage(DenoisingStage):
+    """
+    LongCat denoising stage with CFG-zero optimized guidance scale.
+    
+    Implements:
+    1. Optimized CFG scale from CFG-zero paper
+    2. Negation of noise prediction before scheduler step (flow matching convention)
+    3. Batched CFG computation (unlike standard FastVideo separate passes)
+    """
+
+    def optimized_scale(self, positive_flat, negative_flat) -> torch.Tensor:
+        """
+        Calculate optimized scale from CFG-zero paper.
+        
+        st_star = (v_cond^T * v_uncond) / ||v_uncond||^2
+        
+        Args:
+            positive_flat: Conditional prediction, flattened [B, -1]
+            negative_flat: Unconditional prediction, flattened [B, -1]
+        
+        Returns:
+            st_star: Optimized scale [B, 1]
+        """
+        # Calculate dot product
+        dot_product = torch.sum(positive_flat * negative_flat,
+                                dim=1,
+                                keepdim=True)
+        # Squared norm of uncondition
+        squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
+        # st_star = v_cond^T * v_uncond / ||v_uncond||^2
+        st_star = dot_product / squared_norm
+        return st_star
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Run LongCat denoising loop with optimized CFG.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with denoised latents.
+        """
+        if not fastvideo_args.model_loaded["transformer"]:
+            from fastvideo.models.model_loader import TransformerLoader
+            loader = TransformerLoader()
+            self.transformer = loader.load(
+                fastvideo_args.model_paths["transformer"], fastvideo_args)
+            pipeline = self.pipeline() if self.pipeline else None
+            if pipeline:
+                pipeline.add_module("transformer", self.transformer)
+            fastvideo_args.model_loaded["transformer"] = True
+
+        # Get transformer dtype
+        if hasattr(self.transformer, 'module'):
+            transformer_dtype = next(self.transformer.module.parameters()).dtype
+        else:
+            transformer_dtype = next(self.transformer.parameters()).dtype
+
+        target_dtype = transformer_dtype
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        # Extract batch parameters
+        latents = batch.latents
+        timesteps = batch.timesteps
+        prompt_embeds = batch.prompt_embeds[0]  # LongCat uses single encoder
+        prompt_attention_mask = batch.prompt_attention_mask[
+            0] if batch.prompt_attention_mask else None
+        guidance_scale = batch.guidance_scale
+        do_classifier_free_guidance = batch.do_classifier_free_guidance
+
+        # Get negative prompts if doing CFG
+        if do_classifier_free_guidance:
+            negative_prompt_embeds = batch.negative_prompt_embeds[0]
+            negative_prompt_attention_mask = (batch.negative_attention_mask[0]
+                                              if batch.negative_attention_mask
+                                              else None)
+            # Concatenate for batched processing
+            prompt_embeds_combined = torch.cat(
+                [negative_prompt_embeds, prompt_embeds], dim=0)
+            if prompt_attention_mask is not None:
+                prompt_attention_mask_combined = torch.cat(
+                    [negative_prompt_attention_mask, prompt_attention_mask],
+                    dim=0)
+            else:
+                prompt_attention_mask_combined = None
+        else:
+            prompt_embeds_combined = prompt_embeds
+            prompt_attention_mask_combined = prompt_attention_mask
+
+        # Denoising loop
+        num_inference_steps = len(timesteps)
+        with tqdm(total=num_inference_steps,
+                  desc="LongCat Denoising") as progress_bar:
+            for i, t in enumerate(timesteps):
+                # Expand latents for CFG
+                if do_classifier_free_guidance:
+                    latent_model_input = torch.cat([latents] * 2)
+                else:
+                    latent_model_input = latents
+
+                latent_model_input = latent_model_input.to(target_dtype)
+
+                # Expand timestep to match batch size
+                timestep = t.expand(
+                    latent_model_input.shape[0]).to(target_dtype)
+
+                # Run transformer with context
+                batch.is_cfg_negative = False
+                with set_forward_context(
+                        current_timestep=i,
+                        attn_metadata=None,
+                        forward_batch=batch,
+                ), torch.autocast(device_type='cuda',
+                                  dtype=target_dtype,
+                                  enabled=autocast_enabled):
+                    noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        encoder_hidden_states=prompt_embeds_combined,
+                        timestep=timestep,
+                        encoder_attention_mask=prompt_attention_mask_combined,
+                    )
+
+                # Apply CFG with optimized scale
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+
+                    B = noise_pred_cond.shape[0]
+                    positive = noise_pred_cond.reshape(B, -1)
+                    negative = noise_pred_uncond.reshape(B, -1)
+
+                    # Calculate optimized scale (CFG-zero)
+                    st_star = self.optimized_scale(positive, negative)
+
+                    # Reshape for broadcasting
+                    st_star = st_star.view(B, 1, 1, 1, 1)
+
+                    # Apply optimized CFG formula
+                    noise_pred = (
+                        noise_pred_uncond * st_star + guidance_scale *
+                        (noise_pred_cond - noise_pred_uncond * st_star))
+
+                # CRITICAL: Negate noise prediction for flow matching scheduler
+                noise_pred = -noise_pred
+
+                # Compute previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred,
+                                              t,
+                                              latents,
+                                              return_dict=False)[0]
+
+                progress_bar.update()
+
+        # Update batch with denoised latents
+        batch.latents = latents
+        return batch

fastvideo/pipelines/stages/longcat_refine_init.py

@@ -0,0 +1,310 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat refinement initialization stage.
+
+This stage prepares the latent variables for LongCat's 480p->720p refinement by:
+1. Loading the stage1 (480p) video
+2. Upsampling it to 720p resolution
+3. Encoding it with VAE
+4. Mixing with noise according to t_thresh
+"""
+
+import math
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.vision_utils import load_video
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.configs.pipelines.longcat import get_bucket_config
+
+logger = init_logger(__name__)
+
+
+class LongCatRefineInitStage(PipelineStage):
+    """
+    Stage for initializing LongCat refinement from a stage1 (480p) video.
+    
+    This replicates the logic from LongCatVideoPipeline.generate_refine():
+    - Load stage1_video frames
+    - Upsample spatially and temporally
+    - VAE encode and normalize
+    - Mix with noise according to t_thresh
+    """
+
+    def __init__(self, vae) -> None:
+        super().__init__()
+        self.vae = vae
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Initialize latents for refinement.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with initialized latents for refinement.
+        """
+        refine_from = batch.refine_from
+        in_memory_stage1 = batch.stage1_video
+
+        # Only run for refinement tasks: either a path (refine_from) or in-memory video is provided
+        if refine_from is None and in_memory_stage1 is None:
+            # Not a refinement task, skip
+            return batch
+
+        # ------------------------------------------------------------------
+        # 1. Obtain stage1 frames (either from disk or from in-memory input)
+        # ------------------------------------------------------------------
+        if in_memory_stage1 is not None:
+            # User provided stage1 frames directly (e.g., from distilled stage output)
+            if len(in_memory_stage1) == 0:
+                raise ValueError(
+                    "stage1_video is empty; expected a non-empty list of frames"
+                )
+
+            if isinstance(in_memory_stage1[0], Image.Image):
+                pil_images = in_memory_stage1
+            else:
+                # Assume numpy arrays or torch tensors with shape [H, W, C]
+                pil_images = [
+                    Image.fromarray(np.array(frame))
+                    for frame in in_memory_stage1
+                ]
+
+            logger.info(
+                "Initializing LongCat refinement from in-memory stage1_video (%s frames)",
+                len(pil_images))
+        else:
+            # Path-based refine: load video from disk (original design)
+            logger.info("Initializing LongCat refinement from file: %s",
+                        refine_from)
+            stage1_video_path = Path(refine_from)
+            if not stage1_video_path.exists():
+                raise FileNotFoundError(
+                    f"Stage1 video not found: {refine_from}")
+
+            # Load video frames as PIL Images
+            pil_images, original_fps = load_video(str(stage1_video_path),
+                                                  return_fps=True)
+            logger.info("Loaded stage1 video: %s frames @ %s fps",
+                        len(pil_images), original_fps)
+
+        # Store in batch for reference (use PIL images, same as official demo)
+        batch.stage1_video = pil_images
+
+        # Get parameters from batch
+        num_frames = len(pil_images)
+        spatial_refine_only = batch.spatial_refine_only
+        t_thresh = batch.t_thresh
+        num_cond_frames = batch.num_cond_frames if hasattr(
+            batch, 'num_cond_frames') else 0
+
+        # Calculate new frame count (temporal upsampling if not spatial_refine_only)
+        new_num_frames = num_frames if spatial_refine_only else 2 * num_frames
+        logger.info(
+            "Refine mode: %s",
+            'spatial only' if spatial_refine_only else 'spatial + temporal')
+
+        # Update batch.num_frames to reflect the upsampled count
+        batch.num_frames = new_num_frames
+
+        # Use bucket system to select resolution (exactly like LongCat)
+        # Calculate scale_factor_spatial considering SP split
+        sp_size = fastvideo_args.sp_size if fastvideo_args.sp_size > 0 else 1
+        vae_scale_factor_spatial = 8  # VAE spatial downsampling
+        patch_size_spatial = 2  # LongCat patch size
+        bsa_latent_granularity = 4
+        scale_factor_spatial = vae_scale_factor_spatial * patch_size_spatial * bsa_latent_granularity  # 64
+
+        # Calculate optimal split like LongCat (cp_split_hw logic)
+        # For sp_size=1: [1,1], max=1
+        # For sp_size=2: [1,2], max=2
+        # For sp_size=4: [2,2], max=2
+        # For sp_size=8: [2,4], max=4
+        if sp_size > 1:
+            # Get optimal 2D split factors (mimic context_parallel_util.get_optimal_split)
+            factors = []
+            for i in range(1, int(sp_size**0.5) + 1):
+                if sp_size % i == 0:
+                    factors.append([i, sp_size // i])
+            cp_split_hw = min(factors, key=lambda x: abs(x[0] - x[1]))
+            scale_factor_spatial *= max(cp_split_hw)
+            logger.info("SP split: sp_size=%s, cp_split_hw=%s, max_split=%s",
+                        sp_size, cp_split_hw, max(cp_split_hw))
+        else:
+            cp_split_hw = [1, 1]
+
+        # Get bucket config and find closest bucket for the input aspect ratio
+        bucket_config = get_bucket_config('720p', scale_factor_spatial)
+
+        # Get input aspect ratio from stage1 video
+        input_height, input_width = pil_images[0].height, pil_images[0].width
+        input_ratio = input_height / input_width
+
+        # Find closest bucket
+        closest_ratio = min(bucket_config.keys(),
+                            key=lambda x: abs(float(x) - input_ratio))
+        height, width = bucket_config[closest_ratio][0]
+
+        logger.info("Input aspect ratio: %.2f (%sx%s)", input_ratio,
+                    input_width, input_height)
+        logger.info("Matched bucket ratio: %s -> resolution: %sx%s",
+                    closest_ratio, width, height)
+        logger.info("Target: %sx%s @ %s frames (sp_size=%s, scale_factor=%s)",
+                    width, height, new_num_frames, sp_size,
+                    scale_factor_spatial)
+
+        # Override batch height/width with bucket-selected resolution
+        batch.height = height
+        batch.width = width
+
+        # Convert PIL images to tensor [T, C, H, W]
+        stage1_video_tensor = torch.stack([
+            torch.from_numpy(np.array(img)).permute(2, 0, 1)  # HWC -> CHW
+            for img in pil_images
+        ]).float()  # [T, C, H, W]
+
+        device = batch.prompt_embeds[0].device
+        dtype = batch.prompt_embeds[0].dtype
+        stage1_video_tensor = stage1_video_tensor.to(device=device, dtype=dtype)
+
+        # Replicate LongCat's exact preprocessing (lines 1227-1235 in pipeline_longcat_video.py)
+        # First: spatial interpolation to target (height, width) on [T, C, H, W]
+        video_down = F.interpolate(stage1_video_tensor,
+                                   size=(height, width),
+                                   mode='bilinear',
+                                   align_corners=True)
+
+        # Rearrange to [C, T, H, W] and add batch dimension -> [1, C, T, H, W]
+        video_down = video_down.permute(1, 0, 2,
+                                        3).unsqueeze(0)  # [1, C, T, H, W]
+        video_down = video_down / 255.0  # Normalize to [0, 1]
+
+        # Then: temporal+spatial interpolation to (new_num_frames, height, width)
+        video_up = F.interpolate(video_down,
+                                 size=(new_num_frames, height, width),
+                                 mode='trilinear',
+                                 align_corners=True)
+
+        # Rescale to [-1, 1] for VAE
+        video_up = video_up * 2.0 - 1.0
+
+        logger.info("Upsampled video shape: %s", video_up.shape)
+
+        # Padding logic (exactly like LongCat lines 1237-1255)
+        # Only pad temporal dimension to ensure BSA compatibility
+        vae_scale_factor_temporal = 4
+        num_noise_frames = video_up.shape[2] - num_cond_frames
+
+        num_cond_latents = 0
+        num_cond_frames_added = 0
+        if num_cond_frames > 0:
+            num_cond_latents = 1 + math.ceil(
+                (num_cond_frames - 1) / vae_scale_factor_temporal)
+            num_cond_latents = math.ceil(
+                num_cond_latents /
+                bsa_latent_granularity) * bsa_latent_granularity
+            num_cond_frames_added = 1 + (
+                num_cond_latents -
+                1) * vae_scale_factor_temporal - num_cond_frames
+            num_cond_frames = num_cond_frames + num_cond_frames_added
+
+        num_noise_latents = math.ceil(num_noise_frames /
+                                      vae_scale_factor_temporal)
+        num_noise_latents = math.ceil(
+            num_noise_latents / bsa_latent_granularity) * bsa_latent_granularity
+        num_noise_frames_added = num_noise_latents * vae_scale_factor_temporal - num_noise_frames
+
+        if num_cond_frames_added > 0 or num_noise_frames_added > 0:
+            logger.info(
+                "Padding temporal dimension for BSA: cond_frames+=%s, noise_frames+=%s",
+                num_cond_frames_added, num_noise_frames_added)
+            pad_front = video_up[:, :, 0:1].repeat(1, 1, num_cond_frames_added,
+                                                   1, 1)
+            pad_back = video_up[:, :, -1:].repeat(1, 1, num_noise_frames_added,
+                                                  1, 1)
+            video_up = torch.cat([pad_front, video_up, pad_back], dim=2)
+            logger.info("Padded video shape: %s", video_up.shape)
+
+        # Update batch with actual frame count after padding
+        batch.num_frames = video_up.shape[2]
+
+        # Store padding info for later cropping (CRITICAL for correct output!)
+        batch.num_cond_frames_added = num_cond_frames_added
+        batch.num_noise_frames_added = num_noise_frames_added
+        batch.new_frame_size_before_padding = new_num_frames
+
+        # Store num_cond_latents for denoising stage
+        if num_cond_latents > 0:
+            batch.num_cond_latents = num_cond_latents
+            logger.info("Will use num_cond_latents=%s during denoising",
+                        num_cond_latents)
+
+        logger.info("Padding info: cond+=%s, noise+=%s, original=%s",
+                    num_cond_frames_added, num_noise_frames_added,
+                    new_num_frames)
+
+        # VAE encode
+        logger.info("Encoding stage1 video with VAE...")
+        vae_dtype = next(self.vae.parameters()).dtype
+        vae_device = next(self.vae.parameters()).device
+        video_up = video_up.to(dtype=vae_dtype, device=vae_device)
+
+        with torch.no_grad():
+            latent_dist = self.vae.encode(video_up)
+            # Extract tensor from latent distribution
+            if hasattr(latent_dist, 'latent_dist'):
+                # Nested distribution wrapper
+                latent_up = latent_dist.latent_dist.sample()
+            elif hasattr(latent_dist, 'sample'):
+                # DiagonalGaussianDistribution or similar
+                latent_up = latent_dist.sample()
+            elif hasattr(latent_dist, 'latents'):
+                # Direct latents tensor
+                latent_up = latent_dist.latents
+            else:
+                # Assume it's already a tensor
+                latent_up = latent_dist
+
+        # Normalize latents using VAE config (exactly like LongCat)
+        if hasattr(self.vae.config, 'latents_mean') and hasattr(
+                self.vae.config, 'latents_std'):
+            latents_mean = torch.tensor(self.vae.config.latents_mean).view(
+                1, self.vae.config.z_dim, 1, 1, 1).to(latent_up.device,
+                                                      latent_up.dtype)
+            # LongCat uses: 1.0 / latents_std (equivalent to dividing by latents_std)
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
+                1, self.vae.config.z_dim, 1, 1, 1).to(latent_up.device,
+                                                      latent_up.dtype)
+            # LongCat: (latents - mean) * (1/std)
+            latent_up = (latent_up - latents_mean) * latents_std
+
+        logger.info("Encoded latent shape: %s", latent_up.shape)
+
+        # Mix with noise according to t_thresh
+        # latent_up = (1 - t_thresh) * latent_up + t_thresh * noise
+        noise = torch.randn_like(latent_up).contiguous()
+        latent_up = (1 - t_thresh) * latent_up + t_thresh * noise
+
+        logger.info("Applied t_thresh=%s noise mixing", t_thresh)
+
+        # Store in batch
+        batch.latents = latent_up.to(dtype)
+        batch.raw_latent_shape = latent_up.shape
+
+        logger.info("LongCat refinement initialization complete")
+
+        return batch

fastvideo/pipelines/stages/longcat_refine_timestep.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LongCat refinement timestep preparation stage.
+
+This stage prepares special timesteps for LongCat refinement that start from t_thresh.
+"""
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+
+logger = init_logger(__name__)
+
+
+class LongCatRefineTimestepStage(PipelineStage):
+    """
+    Stage for preparing timesteps specific to LongCat refinement.
+    
+    For refinement, we need to start from t_thresh instead of t=1.0, so we:
+    1. Generate normal timesteps for num_inference_steps
+    2. Filter to only keep timesteps < t_thresh * 1000
+    3. Prepend t_thresh * 1000 as the first timestep
+    """
+
+    def __init__(self, scheduler) -> None:
+        super().__init__()
+        self.scheduler = scheduler
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Prepare refinement-specific timesteps.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with refinement timesteps.
+        """
+        # Only apply if this is a refinement task
+        # Trigger when either a refine_from path or in-memory stage1_video is provided
+        if batch.refine_from is None and batch.stage1_video is None:
+            return batch
+
+        device = get_local_torch_device()
+        num_inference_steps = batch.num_inference_steps
+        t_thresh = batch.t_thresh
+
+        logger.info("Preparing LongCat refinement timesteps (t_thresh=%s)",
+                    t_thresh)
+
+        # ------------------------------------------------------------------
+        # 1) Match LongCatVideoPipeline.get_timesteps_sigmas (non-distill):
+        #    sigmas = linspace(1, 0.001, num_inference_steps) on CPU
+        # ------------------------------------------------------------------
+        base_sigmas = torch.linspace(
+            1.0,
+            0.001,
+            num_inference_steps,
+            dtype=torch.float32,
+            device=
+            "cpu",  # scheduler.set_timesteps expects CPU-convertible sigmas
+        )
+        # Let the scheduler build its internal timestep schedule from sigmas
+        self.scheduler.set_timesteps(num_inference_steps,
+                                     sigmas=base_sigmas,
+                                     device=device)
+        base_timesteps = self.scheduler.timesteps
+
+        # ------------------------------------------------------------------
+        # 2) Apply t_thresh cropping exactly like generate_refine:
+        #    timesteps = [t_thresh*1000] + [t for t in base_timesteps if t < t_thresh*1000]
+        #    sigmas = timesteps / 1000  (with trailing zero)
+        # ------------------------------------------------------------------
+        t_thresh_value = t_thresh * 1000.0
+        t_thresh_tensor = torch.tensor(t_thresh_value,
+                                       dtype=base_timesteps.dtype,
+                                       device=device)
+        filtered_timesteps = base_timesteps[base_timesteps < t_thresh_tensor]
+
+        timesteps = torch.cat(
+            [t_thresh_tensor.unsqueeze(0), filtered_timesteps])
+
+        # Update scheduler with these custom timesteps and corresponding sigmas
+        self.scheduler.timesteps = timesteps
+        sigmas = torch.cat([timesteps / 1000.0, torch.zeros(1, device=device)])
+        self.scheduler.sigmas = sigmas
+
+        logger.info("Refinement timesteps: %s steps starting from t=%s",
+                    len(timesteps), t_thresh)
+        logger.info("First few timesteps: %s", timesteps[:5].tolist())
+
+        # Store in batch so downstream stages (denoising) use the same schedule
+        batch.timesteps = timesteps
+
+        return batch

fastvideo/pipelines/stages/timestep_preparation.py

@@ -7,6 +7,8 @@ This module contains implementations of timestep preparation stages for diffusio
 
 import inspect
 
+import torch
+
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
@@ -71,7 +73,12 @@ class TimestepPreparationStage(PipelineStage):
                     f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                     f" timestep schedules. Please check whether you are using the correct scheduler."
                 )
-            scheduler.set_timesteps(timesteps=timesteps,
+            # Convert timesteps to CPU if it's a tensor (for numpy conversion in scheduler)
+            if isinstance(timesteps, torch.Tensor):
+                timesteps_for_scheduler = timesteps.cpu()
+            else:
+                timesteps_for_scheduler = timesteps
+            scheduler.set_timesteps(timesteps=timesteps_for_scheduler,
                                     device=device,
                                     **extra_set_timesteps_kwargs)
             timesteps = scheduler.timesteps

fastvideo/platforms/cuda.py

@@ -126,7 +126,7 @@ class CudaPlatformBase(Platform):
         logger.info("Selected backend: %s", selected_backend)
         if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN:
             try:
-                from st_attn import sliding_tile_attention  # noqa: F401
+                from fastvideo_kernel import sliding_tile_attention  # noqa: F401
 
                 from fastvideo.attention.backends.sliding_tile_attn import (  # noqa: F401
                     SlidingTileAttentionBackend)
@@ -169,7 +169,7 @@ class CudaPlatformBase(Platform):
                 )
         elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
             try:
-                from vsa import block_sparse_attn  # noqa: F401
+                from fastvideo_kernel import video_sparse_attn  # noqa: F401
 
                 from fastvideo.attention.backends.video_sparse_attn import (  # noqa: F401
                     VideoSparseAttentionBackend)
@@ -188,8 +188,7 @@ class CudaPlatformBase(Platform):
 
         elif selected_backend == AttentionBackendEnum.VMOBA_ATTN:
             try:
-                from csrc.attn.vmoba_attn.vmoba import (  # noqa: F401
-                    moba_attn_varlen)
+                from fastvideo_kernel import moba_attn_varlen  # noqa: F401
                 from fastvideo.attention.backends.vmoba import (  # noqa: F401
                     VMOBAAttentionBackend)
                 logger.info("Using Video MOBA Attention backend.")

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

fastvideo/tests/modal/pr_test.py

@@ -4,6 +4,7 @@ app = modal.App()
 
 import os
 
+model_vol = modal.Volume.from_name("hf-model-weights")
 image_version = os.getenv("IMAGE_VERSION")
 image_tag = f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{image_version}"
 print(f"Using image: {image_tag}")
@@ -52,6 +53,7 @@ def run_test(pytest_command: str):
     git clone {git_repo} /FastVideo &&
     cd /FastVideo &&
     {checkout_command} &&
+    uv pip install -e fastvideo-kernel &&
     uv pip install -e .[test] &&
     {pytest_command}
     """
@@ -74,9 +76,15 @@ def run_vae_tests():
 def run_transformer_tests():
     run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/transformers -vs")
 
-@app.function(gpu="L40S:2", image=image, timeout=2700, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})])
+@app.function(
+    gpu="L40S:4", 
+    image=image, 
+    timeout=2700, 
+    secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})],
+    volumes={"/root/data": model_vol} 
+)
 def run_ssim_tests():
-    run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/ssim -vs")
+    run_test("export MODEL_PATH='/root/data/weights' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/ssim -vs")
 
 @app.function(gpu="L40S:4", image=image, timeout=900, secrets=[modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})])
 def run_training_tests():
@@ -96,15 +104,16 @@ def run_inference_tests_STA():
 
 @app.function(gpu="H100:1", image=image, timeout=900)
 def run_precision_tests_STA():
-    run_test("python csrc/attn/tests/test_sta.py")
+    run_test("pytest fastvideo-kernel/tests/test_correctness.py")
 
 @app.function(gpu="H100:1", image=image, timeout=900)
 def run_precision_tests_VSA():
-    run_test("python csrc/attn/tests/test_vsa.py")
+    # VSA correctness is covered by the same file now
+    run_test("pytest fastvideo-kernel/tests/test_correctness.py")
 
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_precision_tests_vmoba():
-    run_test("pytest csrc/attn/vmoba_attn/tests/test_vmoba_attn.py")
+    run_test("pytest fastvideo-kernel/tests/test_vmoba_correctness.py")
 
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_inference_tests_vmoba():

fastvideo/third_party/longcat_video/block_sparse_attention/__init__.py

@@ -0,0 +1,3 @@
+"""Block-sparse attention kernels for LongCat."""
+
+