add readme

Co-authored-by: RandNMR73 <notomatthew31@gmail.com>
Co-authored-by: JerryZhou54 <zhouw.jerry2017@outlook.com>
Co-authored-by: Davids048 <jundasu@ucsd.edu>
Co-authored-by: Peiyuan Zhang <a1286225768@slurm-h200-204-239.slurm-compute.tenant-slurm.svc.cluster.local>

.buildkite/pipeline.yml

@@ -22,7 +22,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                command: "timeout 20m .buildkite/scripts/pr_test.sh"
                 label: "Encoder Tests"
                 env:
                   - TEST_TYPE=encoder
@@ -35,7 +35,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                command: "timeout 20m .buildkite/scripts/pr_test.sh"
                 label: "VAE Tests"
                 env:
                   - TEST_TYPE=vae
@@ -61,7 +61,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 45m .buildkite/scripts/pr_test.sh"
+                command: "timeout 90m .buildkite/scripts/pr_test.sh"
                 label: "SSIM Tests"
                 env:
                   - TEST_TYPE=ssim
@@ -76,7 +76,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                command: "timeout 20m .buildkite/scripts/pr_test.sh"
                 label: "LoRA Inference Tests"
                 env:
                   - TEST_TYPE=inference_lora
@@ -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,48 +152,16 @@ 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:
                 command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Precision Tests STA"
+                label: "Kernel Tests"
                 env:
-                  - TEST_TYPE=precision_sta
-                agents:
-                  queue: "default"
+                  - TEST_TYPE=kernel_tests
             - 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"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Precision Tests VSA"
-                env:
-                  - TEST_TYPE=precision_vsa
-                agents:
-                  queue: "default"
-            - path:
-                - "csrc/attn/vmoba_attn/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Precision Tests VMoBA"
-                env:
-                  - TEST_TYPE=precision_vmoba
-                agents:
-                  queue: "default"
-            - path:
-                - "csrc/attn/vmoba_attn/vmoba/**"
+                - "fastvideo-kernel/**"
                 - "fastvideo/attention/backends/vmoba.py"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
@@ -222,3 +183,14 @@ steps:
                   - TEST_TYPE=unit_test
                 agents:
                   queue: "default"
+            # - path:
+            #     - "scripts/lora_extraction/**"
+            #     - "pyproject.toml"
+            #     - "docker/Dockerfile.python3.12"
+            #   config:
+            #     command: "timeout 90m .buildkite/scripts/pr_test.sh"
+            #     label: "LoRA Extraction Tests"
+            #     env:
+            #       - TEST_TYPE=lora_extraction
+            #     agents:
+            #       queue: "default"

.buildkite/scripts/pr_test.sh

@@ -31,9 +31,9 @@ log "Setting up Modal authentication from Buildkite secrets..."
 MODAL_TOKEN_ID=$(buildkite-agent secret get modal_token_id)
 MODAL_TOKEN_SECRET=$(buildkite-agent secret get modal_token_secret)
 
+# Retrieve other secrets
 WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
-
-WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
+HF_API_KEY=$(buildkite-agent secret get hf_api_key)
 
 if [ -n "$MODAL_TOKEN_ID" ] && [ -n "$MODAL_TOKEN_SECRET" ]; then
     log "Retrieved Modal credentials from Buildkite secrets"
@@ -63,19 +63,19 @@ MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUI
 case "$TEST_TYPE" in
     "encoder")
         log "Running encoder tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
         ;;
     "vae")
         log "Running VAE tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
         ;;
     "transformer")
         log "Running transformer tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
         ;;
     "ssim")
         log "Running SSIM tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
         ;;
     "training")
         log "Running training tests..."
@@ -93,13 +93,9 @@ case "$TEST_TYPE" in
         log "Running inference STA tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_STA"
         ;;
-    "precision_sta")
-        log "Running precision STA tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_STA"
-        ;;
-    "precision_vsa")
-        log "Running precision VSA tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_VSA"
+    "kernel_tests")
+        log "Running kernel tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_kernel_tests"
         ;;
     "inference_lora")
         log "Running LoRA tests..."
@@ -118,14 +114,14 @@ case "$TEST_TYPE" in
         log "Running V-MoBA inference tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_vmoba"
         ;;
-    "precision_vmoba")
-        log "Running V-MoBA precision tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_vmoba"
-        ;;
     "unit_test")
         log "Running unit tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_unit_test"
         ;;
+    "lora_extraction")
+        log "Running LoRA extraction tests..."
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_lora_extraction_tests"
+        ;;
     *)
         log "Error: Unknown test type: $TEST_TYPE"
         exit 1

.github/workflows/docs.yml

@@ -1,83 +1,72 @@
-# Sample workflow for building and deploying a Hugo site to GitHub Pages
-name: Deploy FastVideo Docs to Pages
+name: Deploy Documentation
 
 on:
-  # Runs on pushes targeting the default branch
   push:
-    branches:
-      - main
+    branches: [ main ]
     paths:
-      - "docs/**/*.md"
-      - "fastvideo/examples/**/*.py"
+      - 'docs/**'
+      - 'mkdocs.yml'
+      - 'requirements-mkdocs.txt'
+      - '.github/workflows/docs.yml'
   pull_request:
-    branches: 
-      - main
-    types: [opened, ready_for_review, synchronize, reopened]
+    branches: [ main ]
     paths:
-      - "docs/**/*.md"
-      - "fastvideo/examples/**/*.py"
+      - 'docs/**'
+      - 'mkdocs.yml'
+      - 'requirements-mkdocs.txt'
+      - '.github/workflows/docs.yml'
 
-  # Allows you to run this workflow manually from the Actions tab
-  workflow_dispatch:
-
-# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
 permissions:
   contents: read
   pages: write
   id-token: write
 
-# Allow only one concurrent deployment, skipping runs queued between the run in-progress and latest queued.
-# However, do NOT cancel in-progress runs as we want to allow these production deployments to complete.
 concurrency:
   group: "pages"
   cancel-in-progress: false
 
-# Default to bash
-defaults:
-  run:
-    shell: bash
-
 jobs:
-  pre-commit:
-    uses: ./.github/workflows/pre-commit.yml
-
-  # Build job
   build:
     runs-on: ubuntu-latest
-    needs: pre-commit
     steps:
       - name: Checkout
         uses: actions/checkout@v4
-      - name: Setup Pages
-        id: pages
-        uses: actions/configure-pages@v5
-      - name: Set up Python
+
+      - name: Setup Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: '3.12'
+
       - name: Install dependencies
         run: |
-          cd docs
-          pip install -r requirements-docs.txt
-      - name: Build docs
-        run: |
-          cd docs
-          make clean
-          make html
+          python -m pip install --upgrade pip
+          pip install -r requirements-mkdocs.txt
+
+      - name: Setup Pages
+        uses: actions/configure-pages@v4
+
+      - name: Generate docs examples
+        run: python docs/generate_examples.py
+
+      - name: Check docs links
+        run: python scripts/check_docs_links.py
+
+      - name: Build documentation
+        run: mkdocs build
+
       - name: Upload artifact
         uses: actions/upload-pages-artifact@v3
         with:
-          path: ./docs/build/html
+          path: ./site
 
-  # Deployment job
   deploy:
     environment:
       name: github-pages
       url: ${{ steps.deployment.outputs.page_url }}
-    if: ${{ github.event_name == 'push' }}
     runs-on: ubuntu-latest
     needs: build
+    if: github.ref == 'refs/heads/main'
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v4

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

@@ -0,0 +1,225 @@
+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'
+            - torch-version: '2.10.0'
+              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
+          # Point auditwheel at torch libs, but do not vendor them into the wheel.
+          TORCH_LIB_DIR=$(python - <<'PY'
+          import os
+          import torch
+
+          print(os.path.join(os.path.dirname(torch.__file__), "lib"))
+          PY
+          )
+          export LD_LIBRARY_PATH="${TORCH_LIB_DIR}:${LD_LIBRARY_PATH}"
+          # Target manylinux_2_35 (Ubuntu 22.04 native)
+          auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist \
+            --exclude libtorch_cuda.so \
+            --exclude libtorch_cpu.so \
+            --exclude libtorch.so \
+            --exclude libc10.so \
+            --exclude libc10_cuda.so \
+            --exclude libtorch_python.so
+          # 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

@@ -14,8 +14,11 @@ wandb/
 *.pt
 cache_dir/
 wandb/
+venv/
+.venv/
 runs/
 samples/
+Miniconda3-latest-Linux-x86_64.sh
 *validation/
 data/
 outputs/
@@ -28,6 +31,13 @@ env
 **/build/
 **.pyc
 **.txt
+*.log
+weights/
+
+# SSIM test outputs
+fastvideo/tests/ssim/generated_videos/
+**/.cache/**
+
 
 # Distribution / packaging
 build/
@@ -37,12 +47,13 @@ dist/
 eggs/
 .eggs/
 
-# Sphinx documentation
-docs/_build/
-docs/source/getting_started/examples/
-docs/source/inference/examples/
-docs/source/training/examples/
-docs/source/distillation/examples/
+# MkDocs documentation
+site/
+docs/getting_started/examples/
+docs/inference/examples/
+docs/training/examples/
+docs/distillation/examples/
+!requirements-mkdocs.txt
 
 # VSCode
 .vscode/
@@ -61,9 +72,14 @@ docs/source/distillation/examples/
 !fastvideo/tests/ssim/reference_videos/**/*.mp4
 
 # Static images
-!docs/source/_static/images/**/*.png
+!docs/assets/images/**/*.png
 !comfyui/assets/**/*.png
 !comfyui/assets/**/*.gif
+!assets/images/**/*.png
+!assets/images/**/*.jpg
+!assets/images/**/*.jpeg
+!assets/images/**/*.gif
+!assets/videos/**/*.mp4
 
 dmd_t2v_output/
 preprocess_output_text/

.gitmodules

@@ -1,7 +1,6 @@
-[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
+[submodule "fastvideo-kernel/include/cutlass"]
+	path = fastvideo-kernel/include/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

.pre-commit-config.yaml

@@ -4,12 +4,13 @@ default_stages:
 exclude: |
   (?x)(
     fastvideo/third_party/.*|
-    csrc/.*|
+    fastvideo-kernel/.*|
     assets/.*|
     tests/.*|
     demo/.*|
     predict\.py|
     scripts/.*|
+    assets/prompts/.*|
     fastvideo/data_preprocess/.*|
     fastvideo/dataset/.*|
     fastvideo/models/.*|
@@ -67,7 +68,7 @@ repos:
     entry: bash
     args:
       - -c
-      - 'git ls-files | grep -v "^fastvideo/tests/ssim/" | grep -v "^fastvideo/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
+      - 'git ls-files | grep -v "^\"*fastvideo/tests/ssim/" | grep -v "^\"*fastvideo/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
     language: system
     always_run: true
     pass_filenames: false

AGENTS.md

@@ -0,0 +1,42 @@
+# Repository Guidelines
+
+## Project Structure & Module Organization
+- Core Python package: `fastvideo/` (models, pipelines, training, distributed runtime, CLI entrypoints).
+- CUDA/custom kernels: `fastvideo-kernel/` (separate build/test flow).
+- Tests:
+  - `fastvideo/tests/` for package-level tests (dataset, encoders, inference, training, SSIM, workflow).
+  - `tests/local_tests/` for additional local/component checks.
+- Docs and guides: `docs/` (MkDocs source), with contributor docs in `docs/contributing/`.
+- Runnable examples and scripts: `examples/` and `scripts/`.
+- Static assets: `assets/` (including `assets/images/`, `assets/videos/`, and `assets/prompts/`) and `comfyui/assets/`.
+
+## Build, Test, and Development Commands
+- `uv pip install -e .[dev]`: editable install with lint/test extras.
+- `pre-commit install --hook-type pre-commit --hook-type commit-msg`: enable local hooks.
+- `pre-commit run --all-files`: run formatter/lint/type/spelling checks.
+- `pytest tests/`: run top-level test suite.
+- `pytest fastvideo/tests/ -v`: run package tests.
+- `pytest fastvideo/tests/ssim/ -vs`: run SSIM regression tests (GPU-heavy).
+- `cd fastvideo-kernel && ./build.sh`: build kernel extensions.
+
+## Coding Style & Naming Conventions
+- Python 3.10+; 4-space indentation; keep code and imports readable and explicit.
+- Style tools are configured in `pyproject.toml` and `.pre-commit-config.yaml`:
+  - `yapf` (format), `ruff` (lint, auto-fix), `mypy` (typing), `codespell`.
+- Target line length is 80.
+- Naming: `snake_case` for functions/files, `PascalCase` for classes, `UPPER_SNAKE_CASE` for constants.
+
+## Testing Guidelines
+- Use `pytest` and place tests near relevant domains (e.g., `fastvideo/tests/encoders/`).
+- Prefer descriptive names like `test_<feature>_<expected_behavior>.py`.
+- For new pipelines/backends, include at least one regression-oriented test; add SSIM coverage when output quality must be preserved.
+- Document GPU assumptions in tests that require specific hardware.
+
+## Commit & Pull Request Guidelines
+- Follow existing commit style: short subject with optional tag prefix, e.g. `[bugfix]: ...`, `[feat]: ...`, `[misc]: ...`, and include PR reference like `(#1234)` when applicable.
+- Keep commits focused by concern (feature, refactor, fix).
+- PRs should include:
+  - clear problem/solution summary,
+  - test evidence (`pytest`/SSIM outputs or rationale if skipped),
+  - linked issue/PR context,
+  - screenshots or sample outputs for UI/demo/docs changes.

CLAUDE.md

@@ -0,0 +1 @@
+@AGENTS.md

README.md

@@ -1,166 +1,101 @@
-<div align="center">
-<img src=assets/logos/logo.svg width="30%"/>
-</div>
+# Sliding Tile Attention (STA) Branch
 
-**FastVideo is a unified post-training and inference framework for accelerated video generation.**
+This branch is a stash/testing branch for people who want to try
+Sliding Tile Attention (STA). The top-level README is intentionally
+STA-only.
 
-FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
+## What is STA
 
-<p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/tMwknPLY" target="_blank"> <b> WeChat </b> </a> |
-</p>
+Sliding Tile Attention is an optimized attention backend for
+window-based video generation.
 
-<div align="center">
-<img src=assets/fastwan.png width="90%"/>
-</div>
+- Blog: https://hao-ai-lab.github.io/blogs/sta/
+- Paper: https://arxiv.org/abs/2502.04507
+- In-repo STA docs: `docs/attention/sta/index.md`
 
-## NEWS
-- ```2025/08/04```: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
-- ```2025/06/14```: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
-- ```2025/04/24```: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
-- ```2025/02/18```: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
+## Setup
 
-## Key Features
-
-FastVideo has the following features:
-- End-to-end post-training support:
-  - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) for Wan2.1 and Wan2.2 to achineve >50x denoising speedup
-  - Data preprocessing pipeline for video data
-  - Support full finetuning and LoRA finetuning for state-of-the-art open video DiTs
-  - Scalable training with FSDP2, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs
-- State-of-the-art performance optimizations for inference
-  - [Video Sparse Attention](https://arxiv.org/pdf/2505.13389)
-  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
-  - [TeaCache](https://arxiv.org/pdf/2411.19108)
-  - [Sage Attention](https://arxiv.org/abs/2410.02367)
-- Diverse hardware and OS support
-  - Support H100, A100, 4090
-  - Support Linux, Windows, MacOS
-
-## Getting Started
-We recommend using an environment manager such as `Conda` to create a clean environment:
+Install FastVideo from source:
 
 ```bash
-# Create and activate a new conda environment
-conda create -n fastvideo python=3.12
-conda activate fastvideo
-
-# Install FastVideo
-pip install fastvideo
+uv pip install -e .
 ```
 
-Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) for more detailed installation instructions.
-
-## Sparse Distillation
-For our sparse distillation techniques, please see our [distillation docs](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html) and check out our [blog](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
-
-See below for recipes and datasets:
-
-|                                            Model                                              |                                               Sparse Distillation                                                 |                                                  Dataset                                                  |
-|:-------------------------------------------------------------------------------------------:  |:---------------------------------------------------------------------------------------------------------------:  |:--------------------------------------------------------------------------------------------------------: |
-| [FastWan2.1-T2V-1.3B](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)         |    [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P)      | [FastVideo Synthetic Wan2.1 480P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)      |
-| [FastWan2.1-T2V-14B-Preview](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-Diffusers)   |                                                   Coming soon!                                                    |   [FastVideo Synthetic Wan2.1 720P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x768x1280_250k)   |
-| [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers)           | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free)   | [FastVideo Synthetic Wan2.2 720P](https://huggingface.co/datasets/FastVideo/Wan2.2-Syn-121x704x1280_32k)  |
-
-## Inference
-### Generating Your First Video
-Here's a minimal example to generate a video using the default settings. Make sure VSA kernels are [installed](https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation.html). Create a file called `example.py` with the following code:
-
-```python
-import os
-from fastvideo import VideoGenerator
-
-def main():
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
-
-    # Create a video generator with a pre-trained model
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
-        num_gpus=1,  # Adjust based on your hardware
-    )
-
-    # Define a prompt for your video
-    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
-
-    # 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()
-```
-
-Run the script with:
+Build the STA kernel package:
 
 ```bash
-python example.py
+cd fastvideo-kernel
+./build.sh
+cd ..
 ```
 
-For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html).
+## Run STA Inference
 
-### Other docs:
+STA backend:
 
-- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview.html)
-- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
+```
 
-## Distillation and Finetuning
-- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html)
-<!-- - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html) -->
+Ready-to-run examples:
 
-## 📑 Development Plan
-<!-- - More distillation methods -->
-  <!-- - [ ] Add Distribution Matching Distillation -->
-More FastWan Models Coming Soon!
-- [ ] Add FastWan2.1-T2V-14B
-- [ ] Add FastWan2.2-T2V-14B
-- [ ] Add FastWan2.2-I2V-14B
-<!-- - Optimization features
-- Code updates -->
-  <!-- - [ ] fp8 support -->
-  <!-- - [ ] faster load model and save model support -->
+- HunyuanVideo: `scripts/inference/v1_inference_hunyuan_STA.sh`
+- Wan2.1-T2V-14B: `scripts/inference/v1_inference_wan_STA.sh`
 
-See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/issues/468).
+Run:
 
-## 🤝 Contributing
+```bash
+bash scripts/inference/v1_inference_hunyuan_STA.sh
+# or
+bash scripts/inference/v1_inference_wan_STA.sh
+```
 
-We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview.html)
+Both scripts already set STA-related env vars:
 
-## Acknowledgement
-We learned and reused code from the following projects:
-- [Wan-Video](https://github.com/Wan-Video)
-- [ThunderKittens](https://github.com/HazyResearch/ThunderKittens)
-- [Triton](https://github.com/triton-lang/triton)
-- [DMD2](https://github.com/tianweiy/DMD2)
-- [diffusers](https://github.com/huggingface/diffusers)
-- [xDiT](https://github.com/xdit-project/xDiT)
-- [vLLM](https://github.com/vllm-project/vllm)
-- [SGLang](https://github.com/sgl-project/sglang)
+- `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
+- `FASTVIDEO_ATTENTION_CONFIG` to an STA mask strategy JSON
 
-We thank [MBZUAI](https://ifm.mbzuai.ac.ae/), [Anyscale](https://www.anyscale.com/), and [GMI Cloud](https://www.gmicloud.ai/) for their support throughout this project.
+## STA Mask Strategy Files
+
+- HunyuanVideo config: `assets/mask_strategy_hunyuan.json`
+- Wan config: `assets/mask_strategy_wan.json`
+
+## STA Mask Search (Wan2.1-T2V-14B)
+
+Run mask search + tuning from repo root:
+
+```bash
+bash examples/inference/sta_mask_search/inference_wan_sta.sh
+```
+
+What this script does:
+
+- Runs `STA_searching` first.
+- Runs `STA_tuning` next (`skip_time_steps=12` by default).
+- Uses prompt shards from `assets/prompt_0.txt` to `assets/prompt_7.txt`.
+
+Important notes:
+
+- Default script is set to 8 GPUs (`num_gpu=8`). If needed, edit
+  `examples/inference/sta_mask_search/inference_wan_sta.sh`.
+- STA searching/tuning currently supports `69x768x1280` (Wan setting).
+
+Generated files:
+
+- Search results: `output/mask_search_result_pos_1280x768/`
+- Tuned strategy: `output/mask_search_strategy_1280x768/mask_strategy_s12.json`
+
+Use the tuned mask for STA inference:
+
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
+export FASTVIDEO_ATTENTION_CONFIG=output/mask_search_strategy_1280x768/mask_strategy_s12.json
+python examples/inference/sta_mask_search/wan_example.py --STA_mode STA_inference --num_gpus 1
+```
 
 ## Citation
-If you find FastVideo useful, please considering citing our work:
 
 ```bibtex
-@software{fastvideo2024,
-  title        = {FastVideo: A Unified Framework for Accelerated Video Generation},
-  author       = {The FastVideo Team},
-  url          = {https://github.com/hao-ai-lab/FastVideo},
-  month        = apr,
-  year         = {2024},
-}
-
-@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}
-}
-
 @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},

assets/prompt_0.txt

@@ -0,0 +1 @@
+Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.

assets/prompt_1.txt

@@ -0,0 +1 @@
+A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature.

assets/prompt_2.txt

@@ -0,0 +1 @@
+A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.

assets/prompt_3.txt

@@ -0,0 +1 @@
+A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones.

assets/prompt_4.txt

@@ -0,0 +1 @@
+A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robots immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1

assets/prompt_5.txt

@@ -0,0 +1 @@
+fox in the forest close-up quickly turned its head to the left

assets/prompt_6.txt

@@ -0,0 +1 @@
+Man walking his dog in the woods on a hot sunny day

assets/prompt_7.txt

@@ -0,0 +1 @@
+A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic.

assets/prompts/mixkit_i2v.jsonl

@@ -0,0 +1,110 @@
+[
+    {
+        "prompt": "Young man skating with a skateboard on the ramps with graffiti of a park with trees, on a sunny day.", 
+        "image_path": "assets/images/mixkit-boy-skating-with-a-skateboard-in-a-park-with-ramps-34389.png"
+    },
+    {
+        "prompt": "In the midst of the joyous New Year's Eve celebration, the cheerful group of friends, their spirits lifted by the festivities, decides to immortalize the moment with a vibrant snapshot", 
+        "image_path": "assets/images/mixkit-a-cheerful-group-of-friends-celebrate-new-years-eve-and-51525.png"
+    },
+    {
+        "prompt": "A man and a woman playing in a field with grass, during a bright afternoon, while cars pass by in the distance.", 
+        "image_path": "assets/images/mixkit-a-cute-couple-playing-on-the-grass-4688.png"
+    },
+    {
+        "prompt": "Aerial view of a rocky mountain in the forest at a sunny day drone flight footage", 
+        "image_path": "assets/images/mixkit-aerial-view-of-a-rocky-mountain-in-the-forest-50589.png"
+    },
+    {
+        "prompt": "A little girl wearing a pink security helmet and denim overall discovers the art of cycling amidst the serene park, as the camera captures her graceful progress.", 
+        "image_path": "assets/images/mixkit-a-little-girl-cruises-through-the-forest-path-on-her-50088.png"
+    },
+    {
+        "prompt": "Aerial shot of a beach shore with sea waves. Big rocks on the sand at an alone beach.", 
+        "image_path": "assets/images/mixkit-aerial-shot-of-a-beach-with-sea-waves-1087.png"
+    },
+    {
+        "prompt": "Young woman cleaning her house decorated with plants and decorations, while dancing happily to music in her headphones.", 
+        "image_path": "assets/images/mixkit-woman-cleaning-her-house-dancing-happy-43379.png"
+    },
+    {
+        "prompt": "Aerial tour in a meadow surrounded by hills on the horizon, while some birds fly low over a lake.", 
+        "image_path": "assets/images/mixkit-birds-flying-low-over-a-lake-in-a-meadow-41417.png"
+    },
+    {
+        "prompt": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.", 
+        "image_path": "assets/images/mixkit-a-rancher-riding-a-horse-at-sunset-1143.png"
+    },
+    {
+        "prompt": "In the video, a martial artist dressed in a traditional white uniform with a black belt demonstrates a series of precise movements against a stark black background. The individual gracefully transitions between stances, embodying a sense of focused discipline and control. Each motion is executed with a deliberate pace, showcasing the fluidity of martial arts techniques. The soft lighting creates subtle highlights on the uniform, adding depth to the figure as it moves. The practitioner begins with an open-hand pose, feet firmly grounded, gradually shifting to a powerful forward punch. The fluidity of the sequence displays a mastery of balance and poise. Every trajectory of the limbs is precise and deliberate, capturing the elegance and strength of martial arts. The serene, isolated setting enhances the intensity and concentration of the practitioner. This visual presentation is an elegant interplay of motion and stillness, displaying the art form's discipline and grace.", 
+        "image_path": "assets/images/mixkit-a-young-man-practicing-his-karate-moves-49635.png"
+    },
+    {
+        "prompt": "In a serene and softly lit yoga studio, three individuals engage in a yoga session, each performing an upward-facing stretch. The central figure is a woman with shoulder-length brown hair, dressed in a light cropped top and green leggings, her posture reflecting grace and concentration. To her right, another participant, a woman in a purple outfit, mirrors the pose with equal poise. On her left, a person with a bun focuses intently, supported slightly by yoga blocks beneath their hands. The warm-colored wooden floor contrasts soothingly with the soft pastel mural on the back wall, featuring an abstract design and partial visage of a serene face. Natural light floods the space from a large window on the right, where lush greens peek through, adding an element of tranquility. In the corner of the room, a collection of meditation instruments, including a gong and a Buddha statue, subtly frame the peaceful setting. The mood is calm yet focused, as all three participants are deeply engaged in their practice. The scene combines elements of balance, harmony, and a shared journey towards mindfulness. This depiction captures the essence of a yoga session that blends personal growth with collective experience.", 
+        "image_path": "assets/images/mixkit-small-group-of-people-doing-yoga-together-43730.png"
+    },
+    {
+        "prompt": "In the deep blue expanse of the ocean, two dolphins glide effortlessly, their sleek bodies reflecting the sunlight filtering through the water. The prominent shadows and caustics create a shimmering effect on their skin, capturing the beauty of their natural habitat. Each dolphin moves with a fluid grace, occasionally interacting with gentle nudges, showcasing their playful and social nature. The scene is vibrant and dynamic, with the clear blue background accentuating the dolphins' movements, making it an ideal subject for AI recreation.", 
+        "image_path": "assets/images/mixkit-dolphins-underwater-4133.png"
+    },
+    {
+        "prompt": "A bustling ski slope comes alive with skiers descending a pristine, snow-covered hill, surrounded by towering, snow-draped evergreens. Several figures stand atop the slope, silhouetted against a clear blue sky, preparing to embark on their ski run. The chair lift on the right continuously drops off eager adventurers, adding to the excitement at the hilltop. Each skier, clad in colorful winter gear, carves distinct paths into the textured snow as they weave their way down. The interplay of sunlight and shadows accentuates the myriad tracks etched into the slope, creating a dynamic visual rhythm. The scene captures a vibrant winter wonderland, full of action and the thrill of a perfect ski day.", 
+        "image_path": "assets/images/mixkit-skiers-on-a-snowy-slope-3327.png"
+    },
+    {
+        "prompt": "A determined climber is scaling a massive rock face, showcasing exceptional strength and skill. The person, clad in a teal shirt and dark pants, climbs with precision, their movements measured and deliberate. They are secured by climbing gear, which includes ropes and a harness, emphasizing their commitment to safety. The rugged texture of the sandy-colored rock provides an imposing backdrop, adding drama and scale to the climb. In the distance, other large rock formations and sparse vegetation can be seen under a bright, overcast sky, contributing to the natural and adventurous atmosphere. The scene captures a moment of focus and challenge, highlighting the climber's tenacity and the breathtaking environment.",
+        "image_path": "assets/images/mixkit-alpinist-climbing-a-huge-rock-in-a-desert-43306.png"
+    },
+    {
+        "prompt": "A silver SUV drives along a winding, snow-covered mountain road, with dense pine trees blanketed in snow lining both sides. The scene is serene, with the vehicle moving smoothly, possibly on a winter journey or vacation. As the SUV disappears around the bend, another, darker SUV follows, creating a sense of motion and perspective on the snow-dusted asphalt. The towering, snow-laden rock formation to the right contrasts with the dark green of the pines, highlighting the peacefulness of the wintry landscape.",
+        "image_path": "assets/images/mixkit-curve-on-a-snowy-forest-road-3317.png"
+    },
+    {
+        "prompt": "A solitary boat glides across the expansive, tranquil expanse of a serene lake. The vessel leaves a gentle wake behind, creating delicate ripples across the mirror-like surface. The water appears a rich shade of teal, seamlessly blending with the sky at the horizon. Silhouettes of distant trees are faintly visible, creating a picturesque backdrop that enhances the solitary journey of the boat. The sky is a calm gradient, shifting from soft oranges near the shore to the pale blues above. In the distance, a few slender poles emerge from the water, remnants of an old structure or natural formation. The mood of the scene is one of peace and solitude, with the boat journeying steadily through the quiet landscape. There is a sense of endless possibilities as the boat moves toward the unseen beyond the frame. The simplicity and stillness of the scene invite contemplation and reflection, encapsulating a perfect moment of quietude on the water.",
+        "image_path": "assets/images/mixkit-motorboat-on-a-large-lake-with-turquoise-blue-waters-4996.png"
+    },
+    {
+        "prompt": "A man wearing grey shorts jumps rope in a gym, weights and gym equipment in the background.",
+        "image_path": "assets/images/gray_short_man.jpg"
+    },
+    {
+        "prompt": "Flying over a peninsula covered in bushy trees, while discovering the sea around it, painted a beautiful turquoise blue, on a sunny day.",
+        "image_path": "assets/images/peninsula.jpg"
+    },
+    {
+        "prompt": "Skillful cyclist doing a wheelie on a bike while riding through a forest, on a dirt road, surrounded by many trees, in the morning.",
+        "image_path": "assets/images/cyclist.jpg"
+    },
+    {
+        "prompt": "Some friends dancing and having fun together in circles, at a party surrounded by colored lights at a party, in a fancy old place, in a view from below them.",
+        "image_path": "assets/images/friends.jpg"
+    },
+    {
+        "prompt": "A saxophonist wearing a blazer dances while playing a song in a park.",
+        "image_path": "assets/images/saxophonist.jpg"
+    },
+    {
+        "prompt": "Romantic couple embracing and looking at each other in the middle of a forest, during a break on a road trip through nature.",
+        "image_path": "assets/images/romance.jpg"
+    },
+    {
+        "prompt": "Man dressed in 80's style dances very happily in his kitchen while listening to music on his radio and drinking wine.",
+        "image_path": "assets/images/80s_dance.jpg"
+    },
+    {
+        "prompt": "Pair of jazz musicians performing a song with their saxophone and trombone on an abandoned train.",
+        "image_path": "assets/images/jazz.jpg"
+    },
+    {
+        "prompt": "A young woman with short hair wearing pink sunglasses chews gum and makes a bubble gum with the city in the background.",
+        "image_path": "assets/images/pink.jpg"
+    },
+    {
+        "prompt": "Natural aerial landscape with a relief covered with abundant trees and vegetation and a thick layer of mist.",
+        "image_path": "assets/images/natural.jpg"
+    },
+    {
+        "prompt": "Loving couple sitting on a log on the shore of a lake outside, sharing an affectionate hug.",
+        "image_path": "assets/images/couple.jpg"
+    }
+]

assets/videos/README.md

@@ -0,0 +1,7 @@
+# FastVideo/assets/videos
+
+This folder is used to store **video assets for examples**, primarily **input videos** consumed by scripts under `FastVideo/examples/`.
+
+- **Typical contents**: short input clips for demos (e.g., video2world / image2video examples).
+- **Non-critical**: these assets are for convenience and are not required to use the FastVideo library.
+- **Large files**: avoid committing large videos to git; prefer shared storage or download-on-demand.

benchmarks/fvd/README.md

@@ -0,0 +1,106 @@
+# FVD (Fréchet Video Distance) Benchmark
+
+Evaluate generated video quality using FVD with the I3D feature extractor.
+
+## Quick Start
+
+**Run the benchmark:**
+
+```bash
+bash benchmarks/scripts/run.sh
+```
+
+That's it! The script auto-installs dependencies and runs the benchmark.
+
+**To customize:** Edit `benchmarks/fvd/run_fvd.py` to change:
+- Video paths (`real_dir`, `gen_dir`)
+- Number of videos, frames, sampling strategy
+- Device, batch size, caching, etc.
+
+## Advanced Usage (CLI)
+
+For more control without editing Python files, use the CLI.
+
+**First-time setup** (one-time per pod/environment):
+
+```bash
+bash benchmarks/scripts/setup_fvd.sh
+```
+
+Then run any configuration you want:
+
+```bash
+# Custom configuration
+python -m benchmarks.fvd.cli \
+    --real-path data/real/ \
+    --gen-path outputs/gen/ \
+    --num-videos 1024 \
+    --num-frames 32 \
+    --clip-strategy random \
+    --batch-size 32 \
+    --seed 42 \
+    --extractor clip
+```
+
+**Standard protocols:**
+
+```bash
+# Use predefined protocols
+python -m benchmarks.fvd.cli \
+    --real-path data/real/ \
+    --gen-path outputs/gen/ \
+    --protocol fvd2048_16f  # or fvd2048_128f, quick_test, etc.
+```
+
+This would use i3d model by default as the feature extractor
+
+**Feature caching** (speed up repeated evaluations):
+
+```bash
+python -m benchmarks.fvd.cli \
+    --real-path data/real/ \
+    --gen-path outputs/gen/ \
+    --protocol fvd2048_16f \
+    --cache-real-features fvd-cache/extractor_name  # Directory path (will save/load fvd-cache/extractor_name/extractor-name_real_features.pkl)
+```
+
+Run `python -m benchmarks.fvd.cli --help` for all options.
+
+## Available Protocols
+
+- `fvd2048_16f` - Standard (2048 videos, 16 frames)
+- `fvd2048_128f` - Long videos (128 frames)
+- `fvd2048_128f_subsample8` - Subsampled long videos
+- `quick_test` - Fast testing (10 videos)
+
+## Configuration Options
+
+Key options in `FVDConfig`:
+
+```python
+num_videos=2048,              # Videos to evaluate
+num_frames_per_clip=16,       # Frames per clip
+clip_strategy='beginning',    # beginning|random|uniform|middle|sliding
+frame_stride=1,               # Frame subsampling
+batch_size=32,                # GPU batch size
+device='cuda',                # cuda|cpu
+cache_real_features=None,     # Cache path for speed
+seed=42,                      # Reproducibility
+extractor='i3d',              # i3d|clip|videomae
+```
+
+## Programmatic Usage
+
+```python
+from benchmarks.fvd import compute_fvd_with_config, FVDConfig
+
+config = FVDConfig.fvd2048_16f()  # or custom config
+results = compute_fvd_with_config('data/real/', 'outputs/gen/', config)
+print(f"FVD: {results['fvd']:.2f}")
+```
+
+## Notes
+
+- Requires minimum 10 frames per clip
+- Supports both video files (.mp4, .avi, etc.) and frame directories
+- `--cache-real-features` expects a **directory path** (e.g., `cache/real`), it will automatically create/load `real_features.pkl` inside that directory

benchmarks/fvd/__init__.py

@@ -0,0 +1,38 @@
+"""
+FastVideo Frechet Video Distance (FVD) Benchmark Module.
+    >>> from fastvideo.benchmarks.fvd import compute_fvd_with_config, FVDConfig
+    >>> config = FVDConfig.fvd2048_16f()  # Standard protocol
+    >>> results = compute_fvd_with_config('data/real/', 'outputs/gen/', config)
+    >>> print(f"FVD: {results['fvd']:.2f}")
+"""
+
+from .fvd import (
+    compute_fvd,
+    compute_fvd_with_config,
+    compute_frechet_distance,
+    compute_statistics,
+    FVDConfig,
+)
+from .feature_extractors import (BaseFeatureExtractor, I3DFeatureExtractor,
+                                 load_extractor)
+from .video_utils import (
+    load_video_auto,
+    sample_clips_from_video,
+    load_video_clips_streaming,
+    ClipSamplingStrategy,
+)
+
+__all__ = [
+    'compute_fvd',
+    'compute_fvd_with_config',
+    'compute_frechet_distance',
+    'compute_statistics',
+    'FVDConfig',
+    'BaseFeatureExtractor',
+    'I3DFeatureExtractor',
+    'load_extractor',
+    'load_video_auto',
+    'sample_clips_from_video',
+    'load_video_clips_streaming',
+    'ClipSamplingStrategy',
+]

benchmarks/fvd/cli.py

@@ -0,0 +1,107 @@
+import argparse
+import sys
+import traceback
+from .fvd import compute_fvd_with_config, FVDConfig
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description='Compute Fréchet Video Distance (FVD)')
+
+    # Required arguments
+    parser.add_argument('--real-path',
+                        type=str,
+                        required=True,
+                        help='Path to real videos')
+    parser.add_argument('--gen-path',
+                        type=str,
+                        required=True,
+                        help='Path to generated videos')
+
+    # Extractor selection
+    parser.add_argument('--extractor',
+                        type=str,
+                        default='i3d',
+                        choices=['i3d', 'clip', 'videomae'],
+                        help='Feature extractor model to use (default: i3d)')
+
+    # Standard args
+    parser.add_argument('--seed',
+                        type=int,
+                        default=None,
+                        help='Random seed for reproducibility')
+    parser.add_argument('--protocol',
+                        type=str,
+                        default=None,
+                        choices=['fvd2048_16f', 'fvd2048_128f', 'quick_test'],
+                        help='Use standard protocol (overrides other settings)')
+    parser.add_argument('--num-videos',
+                        type=int,
+                        default=2048,
+                        help='Number of videos to use')
+    parser.add_argument('--num-frames',
+                        type=int,
+                        default=16,
+                        help='Number of frames per clip')
+    parser.add_argument('--clip-strategy',
+                        type=str,
+                        default='beginning',
+                        help='Clip sampling strategy')
+    parser.add_argument('--batch-size',
+                        type=int,
+                        default=32,
+                        help='Batch size for feature extraction')
+    parser.add_argument('--device',
+                        type=str,
+                        default='cuda',
+                        help='Device to use (cuda or cpu)')
+    parser.add_argument('--cache-real-features',
+                        type=str,
+                        default=None,
+                        help='Path to cache real video features')
+    parser.add_argument('--quiet',
+                        action='store_true',
+                        help='Suppress progress output')
+
+    args = parser.parse_args()
+
+    # Create config
+    if args.protocol:
+        protocol_map = {
+            'fvd2048_16f': FVDConfig.fvd2048_16f,
+            'fvd2048_128f': FVDConfig.fvd2048_128f,
+            'quick_test': FVDConfig.quick_test,
+        }
+        config = protocol_map[args.protocol]()
+        # Apply overrides
+        config.device = args.device
+        config.cache_real_features = args.cache_real_features
+        config.extractor_model = args.extractor  # Apply extractor arg
+    else:
+        config = FVDConfig(
+            num_videos=args.num_videos,
+            num_frames_per_clip=args.num_frames,
+            extractor_model=args.extractor,  # Apply extractor arg
+            clip_strategy=args.clip_strategy,
+            batch_size=args.batch_size,
+            device=args.device,
+            cache_real_features=args.cache_real_features,
+            seed=args.seed)
+
+    try:
+        _ = compute_fvd_with_config(
+            args.real_path,  # noqa: F841
+            args.gen_path,
+            config,
+            verbose=not args.quiet)
+
+        return 0
+
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        traceback.print_exc(file=sys.stderr)
+        return 1
+
+
+if __name__ == '__main__':
+    sys.exit(main())

benchmarks/fvd/feature_extractors.py

@@ -0,0 +1,264 @@
+"""
+Pluggable Feature Extractors for FVD Computation.
+Supports I3D (standard), CLIP, and VideoMAE via a common interface.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from abc import ABC, abstractmethod
+from huggingface_hub import hf_hub_download
+from tqdm import tqdm
+
+try:
+    from transformers import CLIPModel, CLIPProcessor, VideoMAEModel
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+
+
+class BaseFeatureExtractor(ABC, nn.Module):
+    """Abstract base class for all video feature extractors."""
+
+    def __init__(self, device: str = 'cuda'):
+        super().__init__()
+        self.device = torch.device(
+            device if torch.cuda.is_available() else 'cpu')
+
+    @property
+    @abstractmethod
+    def feature_dim(self) -> int:
+        """Dimension of the output feature vector."""
+        pass
+
+    @abstractmethod
+    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            videos: [B, T, C, H, W] in [0, 255] range.
+        Returns:
+            Preprocessed tensor ready for the model.
+        """
+        pass
+
+    @abstractmethod
+    def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
+        """
+        Extract features for a single batch.
+        Args:
+            videos: [B, T, C, H, W] (raw input)
+        Returns:
+            Features: [B, feature_dim]
+        """
+        pass
+
+    @torch.no_grad()
+    def extract_features(self,
+                         videos: torch.Tensor,
+                         batch_size: int = 32,
+                         verbose: bool = True) -> torch.Tensor:
+        """
+        Extract features for a large tensor of videos by batching.
+        """
+        N = len(videos)
+        all_features = []
+
+        iterator = range(0, N, batch_size)
+        if verbose:
+            iterator = tqdm(
+                iterator,
+                desc=f"Extracting features ({self.__class__.__name__})")
+
+        for i in iterator:
+            batch = videos[i:i + batch_size].to(self.device)
+            features = self.extract_features_batch(batch)
+            all_features.append(features.cpu())
+
+        return torch.cat(all_features, dim=0)
+
+
+# 1. I3D Extractor (The Standard FVD Metric)
+class I3DFeatureExtractor(BaseFeatureExtractor):
+    REPO_ID = 'flateon/FVD-I3D-torchscript'
+    MODEL_FILENAME = 'i3d_torchscript.pt'
+
+    def __init__(self, device: str = 'cuda', cache_dir: str | None = None):
+        super().__init__(device)
+        self.cache_dir = cache_dir
+        self.model = self._load_model()
+        self.model.eval()
+        self.model.to(self.device)
+
+    @property
+    def feature_dim(self) -> int:
+        return 400
+
+    def _load_model(self) -> torch.nn.Module:
+        try:
+            model_path = hf_hub_download(repo_id=self.REPO_ID,
+                                         filename=self.MODEL_FILENAME,
+                                         cache_dir=self.cache_dir)
+            return torch.jit.load(model_path, map_location=self.device)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load I3D model: {e}") from e
+
+    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
+        """Standard I3D preprocessing: Resize to 224, Norm to [-1, 1]."""
+        B, T, C, H, W = videos.shape
+
+        if T < 10:
+            raise ValueError(f"I3D requires at least 10 frames, got {T}")
+
+        # Normalize to [0, 1]
+        if videos.max() > 1.0:
+            videos = videos / 255.0
+
+        # Scale to [-1, 1]
+        videos = videos * 2.0 - 1.0
+
+        # Resize to 224x224
+        if H != 224 or W != 224:
+            videos = videos.reshape(B * T, C, H, W)
+            videos = F.interpolate(videos,
+                                   size=(224, 224),
+                                   mode='bilinear',
+                                   align_corners=False)
+            videos = videos.reshape(B, T, C, 224, 224)
+
+        # [B, T, C, H, W] -> [B, C, T, H, W]
+        return videos.permute(0, 2, 1, 3, 4).contiguous()
+
+    def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
+        batch = self.preprocess(videos)
+        # TorchScript I3D returns raw logits when return_features=True
+        return self.model(batch,
+                          rescale=False,
+                          resize=False,
+                          return_features=True)
+
+
+# 2. CLIP Extractor (Semantic/Content Quality)
+class CLIPFeatureExtractor(BaseFeatureExtractor):
+
+    def __init__(self,
+                 device: str = 'cuda',
+                 model_name: str = "openai/clip-vit-base-patch32"):
+        if not TRANSFORMERS_AVAILABLE:
+            raise ImportError(
+                "Please install transformers: pip install transformers")
+        super().__init__(device)
+        self.processor = CLIPProcessor.from_pretrained(model_name)
+        self.model = CLIPModel.from_pretrained(model_name).to(self.device)
+        self.model.eval()
+        self._feature_dim = self.model.config.projection_dim
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
+        # Ensure values are [0, 255]
+        if videos.max() <= 1.0:
+            videos = videos * 255.0
+
+        return videos.to(torch.uint8)
+
+    def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
+        # Input: [B, T, C, H, W]
+        B, T, C, H, W = videos.shape
+        videos = self.preprocess(videos)
+
+        # Flatten B*T to treat frames as images
+        images = videos.view(B * T, C, H, W)
+
+        # HF Processor
+        inputs = self.processor(images=images,
+                                return_tensors="pt",
+                                padding=True)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        # Extract features [B*T, Dim]
+        outputs = self.model.get_image_features(**inputs)
+
+        # Reshape [B, T, Dim] and Average Pooling over time
+        outputs = outputs.view(B, T, -1)
+        return outputs.mean(dim=1)
+
+
+# 3. VideoMAE Extractor (Structure/Motion Quality)
+class VideoMAEFeatureExtractor(BaseFeatureExtractor):
+
+    def __init__(self,
+                 device: str = 'cuda',
+                 model_name: str = "MCG-NJU/videomae-base"):
+        if not TRANSFORMERS_AVAILABLE:
+            raise ImportError(
+                "Please install transformers: pip install transformers")
+        super().__init__(device)
+        self.model = VideoMAEModel.from_pretrained(model_name).to(self.device)
+        self.model.eval()
+
+        self.register_buffer(
+            'mean',
+            torch.tensor([0.485, 0.456, 0.406],
+                         device=self.device).view(1, 1, 3, 1, 1))
+        self.register_buffer(
+            'std',
+            torch.tensor([0.229, 0.224, 0.225],
+                         device=self.device).view(1, 1, 3, 1, 1))
+
+    @property
+    def feature_dim(self) -> int:
+        return self.model.config.hidden_size
+
+    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
+        """
+        Efficient GPU-based preprocessing.
+        Input: [B, T, C, H, W] in range [0, 255]
+        """
+        B, T, C, H, W = videos.shape
+
+        # 1. Resize to 224x224
+        if H != 224 or W != 224:
+            videos = videos.view(B * T, C, H, W)
+            videos = F.interpolate(videos,
+                                   size=(224, 224),
+                                   mode='bilinear',
+                                   align_corners=False)
+            videos = videos.view(B, T, C, 224, 224)
+
+        # 2. Normalize to [0, 1]
+        if videos.dtype != torch.float32:
+            videos = videos.float()
+
+        if videos.max() > 1.0:
+            videos = videos / 255.0
+
+        # 3. Apply ImageNet Mean/Std
+        return (videos - self.mean) / self.std
+
+    def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
+        # Input: [B, T, C, H, W]
+
+        # Fast GPU Preprocessing
+        pixel_values = self.preprocess(videos)
+
+        # Forward pass
+        outputs = self.model(pixel_values)
+
+        # Global Average Pooling of last hidden state [B, T_patches, 768] -> [B, 768]
+        return outputs.last_hidden_state.mean(dim=1)
+
+
+# Factory
+def load_extractor(name: str, device: str = 'cuda') -> BaseFeatureExtractor:
+    name = name.lower()
+    if name == 'i3d':
+        return I3DFeatureExtractor(device)
+    elif name == 'clip':
+        return CLIPFeatureExtractor(device)
+    elif name == 'videomae':
+        return VideoMAEFeatureExtractor(device)
+    else:
+        raise ValueError(
+            f"Unknown extractor: {name}. Options: i3d, clip, videomae")

benchmarks/fvd/fvd.py

@@ -0,0 +1,405 @@
+import numpy as np
+import scipy.linalg
+import torch
+from pathlib import Path
+from collections.abc import Iterator
+import pickle
+from dataclasses import dataclass, field
+from .feature_extractors import BaseFeatureExtractor, load_extractor
+from .video_utils import ClipSamplingStrategy, load_video_clips_streaming
+
+
+def compute_statistics(features: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Compute mean and covariance."""
+    mu = np.mean(features, axis=0)
+    sigma = np.cov(features, rowvar=False)
+    return mu, sigma
+
+
+def compute_frechet_distance(mu1: np.ndarray,
+                             sigma1: np.ndarray,
+                             mu2: np.ndarray,
+                             sigma2: np.ndarray,
+                             eps: float = 1e-6) -> float:
+    """
+    Compute Fréchet distance between two Gaussians.
+    """
+    sigma1 = sigma1 + eps * np.eye(sigma1.shape[0])
+    sigma2 = sigma2 + eps * np.eye(sigma2.shape[0])
+
+    diff = mu1 - mu2
+    mean_distance = np.sum(diff**2)
+
+    trace_sum = np.trace(sigma1 + sigma2)
+
+    covmean = scipy.linalg.sqrtm(sigma1 @ sigma2)
+
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            print(
+                f"Warning: Imaginary component: {np.max(np.abs(covmean.imag))}")
+        covmean = covmean.real
+
+    trace_product = np.trace(covmean)
+
+    fvd = mean_distance + trace_sum - 2 * trace_product
+
+    return float(fvd)
+
+
+@dataclass
+class FVDConfig:
+    # default configuration for FVD computation:
+
+    # Video selection
+    num_videos: int = 2048
+
+    # Feature Extractor Selection
+    extractor_model: str = 'i3d'  # Options: 'i3d', 'clip', 'videomae'
+
+    # Clip sampling
+    num_frames_per_clip: int = 16
+    num_clips_per_video: int = 1
+    clip_strategy: str | ClipSamplingStrategy = 'beginning'
+
+    # Temporal subsampling
+    frame_stride: int = 1  # 1=no subsampling, 2=every 2nd, 8=every 8th
+    temporal_stride: int = 1  # For sliding window clips
+
+    # Data processing
+    video_extensions: list[str] = field(
+        default_factory=lambda: ['.mp4', '.avi', '.mov', '.mkv'])
+    support_frame_dirs: bool = True
+
+    # Computation
+    batch_size: int = 32
+    device: str = 'cuda'
+
+    use_streaming: bool = True
+    resize_before_extraction: bool = True
+
+    # Caching
+    cache_real_features: str | None = None
+    i3d_model_path: str | None = None
+
+    # Reproducibility
+    seed: int | None = None
+
+    @classmethod
+    def fvd2048_16f(cls) -> 'FVDConfig':
+        """Standard FVD protocol: 2048 videos, 16 frames, beginning clip."""
+        return cls(num_videos=2048,
+                   num_frames_per_clip=16,
+                   clip_strategy='beginning',
+                   use_streaming=True)
+
+    @classmethod
+    def fvd2048_128f(cls) -> 'FVDConfig':
+        """Long video protocol: 2048 videos, 128 frames."""
+        return cls(num_videos=2048,
+                   num_frames_per_clip=128,
+                   clip_strategy='beginning',
+                   use_streaming=True)
+
+    @classmethod
+    def quick_test(cls) -> 'FVDConfig':
+        """Quick test config: 100 videos, 16 frames."""
+        return cls(num_videos=100,
+                   num_frames_per_clip=16,
+                   clip_strategy='beginning')
+
+    def to_dict(self) -> dict:
+        """Export config to dict for logging"""
+        d = self.__dict__.copy()
+        d['clip_strategy'] = str(self.clip_strategy)
+        return d
+
+    def __str__(self) -> str:
+        """Human-readable protocol name"""
+        desc = f"FVD_{self.extractor_model.upper()}_{self.num_videos}_{self.num_frames_per_clip}f"
+        if self.frame_stride > 1:
+            desc += f"_subsample{self.frame_stride}"
+        if self.num_clips_per_video > 1:
+            desc += f"_{self.num_clips_per_video}clips"
+        if self.clip_strategy != 'beginning':
+            desc += f"_{self.clip_strategy}"
+        return desc
+
+
+def extract_features_streaming(video_generator: Iterator[torch.Tensor],
+                               extractor: BaseFeatureExtractor,
+                               batch_size: int = 32,
+                               max_clips: int | None = None,
+                               verbose: bool = True) -> np.ndarray:
+    """
+    Extract features from a video clip generator using streaming.
+    """
+    all_features = []
+    batch = []
+
+    if verbose:
+        print(f"Extracting features with batch_size={batch_size}...")
+
+    with torch.no_grad():
+        for clip_count, clip in enumerate(video_generator):
+            batch.append(clip)
+
+            # Process batch when full
+            if len(batch) == batch_size:
+                batch_tensor = torch.stack(batch).to(extractor.device)
+                features = extractor.extract_features_batch(batch_tensor)
+
+                all_features.append(features.detach().cpu().numpy())
+                batch = []
+
+                if verbose and clip_count % (batch_size * 10) == 0:
+                    print(f"Processed {clip_count} clips...")
+
+            if max_clips is not None and clip_count >= max_clips:
+                break
+
+        # Process remaining clips
+        if len(batch) > 0:
+            batch_tensor = torch.stack(batch).to(extractor.device)
+            features = extractor.extract_features_batch(batch_tensor)
+            all_features.append(features.detach().cpu().numpy())
+
+    if len(all_features) == 0:
+        raise RuntimeError("No features extracted - check video loading")
+
+    features = np.concatenate(all_features, axis=0)
+
+    if verbose:
+        print(f"Extracted {len(features)} feature vectors")
+
+    return features
+
+
+def load_or_compute_features(videos: str | Path | torch.Tensor,
+                             extractor: BaseFeatureExtractor,
+                             config: FVDConfig,
+                             cache_path: str | None = None,
+                             cache_name: str = "real_features") -> np.ndarray:
+    """Load features from cache or compute (with streaming support)"""
+
+    if cache_path is not None:
+        script_dir = Path(__file__).parent
+        cache_dir = script_dir / cache_path
+        cache_file = cache_dir / f"{config.extractor_model}_{cache_name}.pkl"
+
+        if cache_file.exists():
+            print(f"Loading cached features from {cache_file}")
+            with open(cache_file, 'rb') as f:
+                features = pickle.load(f)
+
+            # Validate and limit based on config
+            max_features = config.num_videos * config.num_clips_per_video
+
+            if len(features) < max_features:
+                print(
+                    f"WARNING: Cache has {len(features)} features but need {max_features}"
+                )
+                print("Cached features insufficient - will recompute...")
+            elif len(features) > max_features:
+                print(
+                    f"Using {max_features} features from cache (truncated from {len(features)})"
+                )
+                features = features[:max_features]
+                return features
+            else:
+                print(f"Using all {len(features)} cached features")
+                return features
+
+    print("Computing features from scratch...")
+
+    if isinstance(videos, (str | Path)):
+        target_size = (224, 224) if config.resize_before_extraction else None
+
+        video_generator = load_video_clips_streaming(
+            videos,
+            num_frames=config.num_frames_per_clip,
+            max_videos=config.num_videos,
+            clip_strategy=config.clip_strategy,
+            frame_stride=config.frame_stride,
+            num_clips_per_video=config.num_clips_per_video,
+            video_extensions=config.video_extensions,
+            support_frame_dirs=config.support_frame_dirs,
+            target_size=target_size,
+            verbose=True)
+
+        max_clips = config.num_videos * config.num_clips_per_video
+        features = extract_features_streaming(video_generator,
+                                              extractor,
+                                              batch_size=config.batch_size,
+                                              max_clips=max_clips,
+                                              verbose=True)
+    else:
+        print(f"Extracting features from {len(videos)} video tensors...")
+        features = extractor.extract_features(videos,
+                                              batch_size=config.batch_size,
+                                              verbose=True)
+        features = features.numpy()
+
+    # Validate feature count
+    expected_count = config.num_videos * config.num_clips_per_video
+    if len(features) < expected_count:
+        raise ValueError(
+            f"ERROR: Only extracted {len(features)} features, but need {expected_count}!\n"
+            f"Found fewer videos than expected. Check your video directory.")
+    elif len(features) > expected_count:
+        print(f"Truncating {len(features)} features to {expected_count}")
+        features = features[:expected_count]
+
+    # Cache features if requested
+    if cache_path is not None:
+        script_dir = Path(__file__).parent
+        cache_dir = script_dir / cache_path
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_file = cache_dir / f"{config.extractor_model}_{cache_name}.pkl"
+        print(f"Caching features to {cache_file}")
+        with open(cache_file, 'wb') as f:
+            pickle.dump(features, f)
+
+    return features
+
+
+def compute_fvd_with_config(real_videos: str | Path | torch.Tensor,
+                            gen_videos: str | Path | torch.Tensor,
+                            config: FVDConfig,
+                            verbose: bool = True) -> dict:
+    """
+        Compute FVD using a standardized configuration.
+
+        This is the recommended way to compute FVD for reproducibility.
+
+        Args:
+            real_videos: Path or tensors
+            gen_videos: Path or tensors
+            config: FVDConfig specifying protocol
+            verbose: Print progress
+
+        Returns:
+            results: Dictionary with:
+                - 'fvd': FVD score (float)
+                - 'protocol': Protocol name (str)
+                - 'model': Feature extractor model name (str)
+                - 'config': Configuration dict
+
+        Example:
+            >>> config = FVDConfig.fvd2048_16f()
+            >>> results = compute_fvd_with_config('data/real/', 'outputs/gen/', config)
+            >>> print(f"FVD: {results['fvd']:.2f}")
+    """
+
+    # Seed for reproducibility
+    if config.seed is not None:
+        import random as _rnd
+        _rnd.seed(config.seed)
+        np.random.seed(config.seed)
+        torch.manual_seed(config.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(config.seed)
+
+    if verbose:
+        print("=" * 70)
+        print(f"Computing FVD with protocol: {config}")
+        print(f"Model: {config.extractor_model.upper()}")
+        print("=" * 70)
+        print("\nConfiguration:")
+        for key, value in config.to_dict().items():
+            print(f"  {key}: {value}")
+        print()
+
+    # Initialize Extractor using Factory
+    if verbose:
+        print(
+            f"\nInitializing {config.extractor_model.upper()} model on {config.device}..."
+        )
+
+    extractor = load_extractor(config.extractor_model, device=config.device)
+
+    # Extract features
+    if verbose:
+        print(f"\n{'='*70}")
+        print("Extracting REAL video features...")
+        print(f"{'='*70}")
+
+    real_features = load_or_compute_features(
+        videos=real_videos,
+        extractor=extractor,
+        config=config,
+        cache_path=config.cache_real_features,
+        cache_name="real_features")
+
+    if verbose:
+        print(f"\n{'='*70}")
+        print("Extracting GENERATED video features...")
+        print(f"{'='*70}")
+
+    gen_features = load_or_compute_features(videos=gen_videos,
+                                            extractor=extractor,
+                                            config=config,
+                                            cache_path=None,
+                                            cache_name="gen_features")
+
+    if verbose:
+        print(f"\nReal videos/clips: {len(real_features)}")
+        print(f"Generated videos/clips: {len(gen_features)}")
+        print(f"\n{'='*70}")
+        print("Computing statistics...")
+        print(f"{'='*70}")
+
+    mu_real, sigma_real = compute_statistics(real_features)
+    mu_gen, sigma_gen = compute_statistics(gen_features)
+
+    if verbose:
+        print(f"\n{'='*70}")
+        print("Computing Fréchet distance...")
+        print(f"{'='*70}")
+
+    fvd = compute_frechet_distance(mu_real, sigma_real, mu_gen, sigma_gen)
+
+    if verbose:
+        print(f"\n{'='*70}")
+        print(f"FVD Score ({config.extractor_model.upper()}): {fvd:.4f}")
+        print(f"Protocol: {config}")
+        print(f"{'='*70}\n")
+
+    results = {
+        'fvd': fvd,
+        'protocol': str(config),
+        'model': config.extractor_model,
+        'config': config.to_dict(),
+    }
+
+    return results
+
+
+def compute_fvd(real_videos: str | Path | torch.Tensor,
+                gen_videos: str | Path | torch.Tensor,
+                num_frames: int = 16,
+                batch_size: int = 32,
+                device: str = 'cuda',
+                num_videos: int | None = 2048,
+                cache_real_features: str | None = None,
+                i3d_model_path: str | None = None,
+                seed: int | None = None,
+                verbose: bool = True) -> float:
+    """
+    Backward compatibility wrapper for computing FVD (defaults to I3D).
+    """
+    num_videos = num_videos if num_videos is not None else 2048
+
+    config = FVDConfig(
+        num_videos=num_videos,
+        num_frames_per_clip=num_frames,
+        extractor_model='i3d',  # Default to I3D
+        batch_size=batch_size,
+        device=device,
+        cache_real_features=cache_real_features,
+        i3d_model_path=i3d_model_path,
+        seed=seed,
+    )
+
+    result = compute_fvd_with_config(real_videos, gen_videos, config, verbose)
+    return result['fvd']

benchmarks/fvd/i3d_model.py

@@ -0,0 +1,142 @@
+"""I3D Feature Extractor for FVD Computation"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+from tqdm import tqdm
+from contextlib import suppress
+
+
+class I3DFeatureExtractor(nn.Module):
+    """
+    I3D feature extractor for FVD computation.
+    Extracts 400-dimensional features from videos using I3D model
+    trained on Kinetics-400.
+    """
+
+    REPO_ID = 'flateon/FVD-I3D-torchscript'
+    MODEL_FILENAME = 'i3d_torchscript.pt'
+
+    def __init__(self,
+                 device: str = 'cuda',
+                 cache_dir: str | Path | None = None):
+        super().__init__()
+
+        self.device_str = device
+        if device == 'cuda' and not torch.cuda.is_available():
+            print(
+                "Warning: CUDA requested but not available – falling back to CPU"
+            )
+            self.device = torch.device('cpu')
+        else:
+            self.device = torch.device(device)
+
+        self.cache_dir: str | None
+        if cache_dir is not None:
+            self.cache_dir = str(Path(cache_dir).resolve())
+        else:
+            self.cache_dir = None  # Use HF default cache
+
+        self.model = self._load_model()
+        self.model.eval()
+
+        with suppress(Exception):
+            self.model.to(self.device)
+
+    def _load_model(self) -> torch.nn.Module:
+        """Download and load I3D TorchScript model from Hugging Face Hub."""
+        print(f"Loading I3D model from Hugging Face Hub ({self.REPO_ID})...")
+
+        try:
+            # Download model from Hugging Face Hub
+            model_path = hf_hub_download(repo_id=self.REPO_ID,
+                                         filename=self.MODEL_FILENAME,
+                                         cache_dir=self.cache_dir)
+
+            # Load directly to chosen device
+            model = torch.jit.load(model_path, map_location=self.device)
+            print("I3D model loaded successfully")
+            return model
+
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to load I3D model from Hugging Face Hub. Error: {e}\n"
+                f"Ensure you have internet connection and huggingface_hub installed:\n"
+                f"pip install huggingface_hub") from e
+
+    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
+        """
+        Preprocess videos for I3D.
+        
+        Args:
+            videos: [B, T, C, H, W], values in [0, 255]
+        
+        Returns:
+            Preprocessed videos [B, C, T, 224, 224] (normalized and resized)
+        """
+        B, T, C, H, W = videos.shape
+
+        if T < 10:
+            raise ValueError(f"I3D requires at least 10 frames, got {T}")
+
+        # Normalize to [0, 1] if needed
+        if videos.max() > 1.0:
+            videos = videos / 255.0
+
+        # Resize to 224x224 if needed
+        if H != 224 or W != 224:
+            videos = videos.reshape(B * T, C, H, W)
+            videos = F.interpolate(videos,
+                                   size=(224, 224),
+                                   mode='bilinear',
+                                   align_corners=False)
+            videos = videos.reshape(B, T, C, 224, 224)
+
+        # Convert to [B, C, T, H, W] format
+        videos = videos.permute(0, 2, 1, 3, 4).contiguous()
+
+        return videos
+
+    @torch.no_grad()
+    def extract_features(self,
+                         videos: torch.Tensor,
+                         batch_size: int = 32,
+                         verbose: bool = True) -> torch.Tensor:
+        """
+        Extract I3D features
+        
+        Args:
+            videos: [N, T, C, H, W], values in [0, 255]
+            batch_size: Batch size for processing
+            verbose: Show progress bar
+        
+        Returns:
+            Features [N, 400]
+        """
+        N = len(videos)
+        all_features = []
+
+        iterator = range(0, N, batch_size)
+        if verbose:
+            iterator = tqdm(iterator, desc="Extracting I3D features")
+
+        for i in iterator:
+            batch = videos[i:i + batch_size].to(self.device)
+            batch = self.preprocess(batch)  # Now returns [B, C, T, H, W]
+
+            # Use the HF model without rescale/resize (we handle it in preprocess)
+            features = self.model(batch,
+                                  rescale=False,
+                                  resize=False,
+                                  return_features=True)
+
+            all_features.append(features.cpu())
+
+        return torch.cat(all_features, dim=0)
+
+    def __call__(self,
+                 videos: torch.Tensor,
+                 batch_size: int = 32) -> torch.Tensor:
+        return self.extract_features(videos, batch_size=batch_size)

benchmarks/fvd/run_fvd.py

@@ -0,0 +1,54 @@
+import sys
+from pathlib import Path
+
+root_dir = Path(__file__).parent.parent.parent
+sys.path.insert(0, str(root_dir))
+
+from benchmarks.fvd.fvd import FVDConfig, compute_fvd_with_config  # noqa: E402
+
+
+def main() -> None:
+    script_dir = Path(__file__).parent.resolve()
+
+    # Define directories
+    real_dir = "benchmarks/data/real_videos"
+    gen_dir = "benchmarks/data/generated_videos"
+
+    # Compare all 3 models
+    models_to_test = ['i3d', 'clip', 'videomae']
+
+    print(f"\n{'='*60}")
+    print("STARTING COMPARISON BENCHMARK")
+    print(f"{'='*60}")
+
+    for model_name in models_to_test:
+        print(f"\n>>> Running evaluation with {model_name.upper()}...")
+
+        try:
+            cfg = FVDConfig(
+                num_videos=650,
+                num_frames_per_clip=16,
+                extractor_model=model_name,
+                clip_strategy='beginning',
+                device='cuda',
+                seed=42,
+                # Use separate cache folders for each model to avoid conflicts
+                cache_real_features=str(script_dir / f'fvd-cache/{model_name}'),
+            )
+
+            results = compute_fvd_with_config(real_dir,
+                                              gen_dir,
+                                              cfg,
+                                              verbose=False)
+            print(f"FVD: {results['fvd']}\nModel: {results['model']}")
+
+        except Exception as e:
+            print(f"{model_name.upper()} Failed: {e}")
+
+    print(f"\n{'='*60}")
+    print("BENCHMARK COMPLETE")
+    print(f"{'='*60}")
+
+
+if __name__ == '__main__':
+    main()

benchmarks/fvd/validate_fvd.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+import sys
+from pathlib import Path
+import shutil
+import random
+from fvd import compute_fvd_with_config, FVDConfig
+
+script_path = Path(__file__).resolve()
+fastvideo_root = script_path.parent.parent.parent
+sys.path.insert(0, str(fastvideo_root))
+
+
+def split_videos(video_dir: Path, n_per_subset: int = 128, seed: int = 42):
+    subset_a = video_dir.parent / 'bair_full_subset_A'
+    subset_b = video_dir.parent / 'bair_full_subset_B'
+
+    if subset_a.exists():
+        shutil.rmtree(subset_a)
+    if subset_b.exists():
+        shutil.rmtree(subset_b)
+
+    subset_a.mkdir(parents=True)
+    subset_b.mkdir(parents=True)
+
+    videos = sorted(video_dir.glob('*.mp4'))
+
+    random.seed(seed)
+    shuffled = list(videos)
+    random.shuffle(shuffled)
+
+    needed = n_per_subset * 2
+    if len(shuffled) > needed:
+        shuffled = shuffled[:needed]
+
+    mid = len(shuffled) // 2
+
+    print(f"\nSplitting {len(shuffled)} BAIR FULL videos:")
+    print(f"  Subset A: {mid} videos")
+    print(f"  Subset B: {len(shuffled) - mid} videos")
+
+    for v in shuffled[:mid]:
+        shutil.copy2(v, subset_a / v.name)
+
+    for v in shuffled[mid:]:
+        shutil.copy2(v, subset_b / v.name)
+
+    return subset_a, subset_b, mid
+
+
+def validate_fvd(subset_a: Path, subset_b: Path, num_videos: int):
+    config = FVDConfig(num_videos=num_videos,
+                       num_frames_per_clip=16,
+                       clip_strategy='beginning',
+                       batch_size=8,
+                       device='cuda',
+                       seed=42)
+
+    print("\n" + "=" * 70)
+    print("TEST 1: Identity Test")
+    print("=" * 70)
+
+    result1 = compute_fvd_with_config(real_videos=str(subset_a),
+                                      gen_videos=str(subset_a),
+                                      config=config,
+                                      verbose=False)
+    fvd_identity = result1['fvd']
+    print(f"\nIdentity FVD: {fvd_identity:.2f}")
+
+    print("\n" + "=" * 70)
+    print("TEST 2: Real vs Real")
+    print("=" * 70)
+
+    result2 = compute_fvd_with_config(real_videos=str(subset_a),
+                                      gen_videos=str(subset_b),
+                                      config=config,
+                                      verbose=False)
+    fvd_real = result2['fvd']
+    print(f"\nReal vs Real FVD: {fvd_real:.2f}")
+
+    print("\n" + "=" * 70)
+    print("RESULTS")
+    print("=" * 70)
+    print(f"Identity:        {fvd_identity:.2f}")
+    print(f"Real vs Real:    {fvd_real:.2f}")
+
+
+def main() -> None:
+    bair_dir = Path('benchmarks/data/bair_full_videos')
+
+    subset_a, subset_b, count = split_videos(bair_dir,
+                                             n_per_subset=128,
+                                             seed=42)
+    validate_fvd(subset_a, subset_b, count)
+
+
+if __name__ == '__main__':
+    main()

benchmarks/fvd/video_utils.py

@@ -0,0 +1,490 @@
+import torch
+import cv2
+import numpy as np
+from pathlib import Path
+from collections.abc import Iterator
+from tqdm import tqdm
+from enum import Enum
+
+
+class ClipSamplingStrategy(Enum):
+    """Clip sampling strategies for FVD evaluation."""
+    BEGINNING = 'beginning'  # Take first N frames (most common)
+    RANDOM = 'random'  # Random N consecutive frames
+    UNIFORM = 'uniform'  # Uniformly spaced frames across video
+    MIDDLE = 'middle'  # Middle N frames
+    SLIDING = 'sliding'  # Multiple sliding windows
+    ALL = 'all'  # All possible clips
+
+
+def _load_video_cv2(video_path: str | Path,
+                    num_frames: int | None = 16,
+                    sample_strategy: str = 'uniform') -> torch.Tensor:
+    """
+    Load video from video file using OpenCV.
+    
+    Args:
+        video_path: Path to video file (MP4, AVI, MOV, MKV)
+        num_frames: Number of frames to extract
+        sample_strategy: 'uniform' or 'random'
+    
+    Returns:
+        video: [T, C, H, W]
+    """
+    video_path = str(video_path)
+    cap = cv2.VideoCapture(video_path)
+
+    if not cap.isOpened():
+        raise RuntimeError(f"Cannot open video: {video_path}")
+
+    frames = []
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    if num_frames is None:
+        # Read all available frames
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frames.append(frame)
+
+        cap.release()
+        if len(frames) == 0:
+            raise RuntimeError(f"Video has 0 frames: {video_path}")
+
+        frames = np.stack(frames)  # [T, H, W, C]
+        frames = torch.from_numpy(frames).permute(0, 3, 1,
+                                                  2).float()  # [T, C, H, W]
+        return frames
+
+    if total_frames == 0:
+        raise RuntimeError(f"Video has 0 frames: {video_path}")
+
+    # Determine frame indices for sampling
+    if total_frames < num_frames:
+        frame_indices = list(range(
+            total_frames)) + [total_frames - 1] * (num_frames - total_frames)
+    elif sample_strategy == 'uniform':
+        frame_indices = np.linspace(0, total_frames - 1, num_frames,
+                                    dtype=int).tolist()
+    elif sample_strategy == 'random':
+        frame_indices = sorted(
+            np.random.choice(total_frames, num_frames, replace=False))
+    else:
+        raise ValueError(f"Unknown sample_strategy: {sample_strategy}")
+
+    # Extract frames
+    for idx in frame_indices:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
+        ret, frame = cap.read()
+
+        if not ret:
+            if len(frames) > 0:
+                frames.append(frames[-1].copy())
+            else:
+                h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+                w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+                frames.append(np.zeros((h, w, 3), dtype=np.uint8))
+            continue
+
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames.append(frame)
+
+    cap.release()
+
+    frames = np.stack(frames)  # [T, H, W, C]
+    frames = torch.from_numpy(frames).permute(0, 3, 1,
+                                              2).float()  # [T, C, H, W]
+
+    return frames
+
+
+def _load_video_from_frames(
+        frame_dir: str | Path,
+        num_frames: int | None = 16,
+        sample_strategy: str = 'uniform',
+        frame_extensions: list[str] | None = None) -> torch.Tensor:
+    """
+    Load video from directory of frame images.
+    
+    Args:
+        frame_dir: Directory containing frames
+        num_frames: Number of frames to sample
+        sample_strategy: 'uniform' or 'random'
+        frame_extensions: Image file extensions to look for
+    
+    Returns:
+        video: [T, C, H, W]
+    """
+    if frame_extensions is None:
+        frame_extensions = ['.jpg', '.png', '.jpeg', '.bmp']
+
+    frame_dir = Path(frame_dir)
+
+    if not frame_dir.exists():
+        raise FileNotFoundError(f"Frame directory not found: {frame_dir}")
+
+    # Find all frames
+    frame_files: list[Path] = []
+    for ext in frame_extensions:
+        frame_files.extend(frame_dir.glob(f"*{ext}"))
+
+    if len(frame_files) == 0:
+        raise ValueError(
+            f"No frames found in {frame_dir} with extensions {frame_extensions}"
+        )
+
+    frame_files = sorted(frame_files, key=lambda x: x.name)
+    total_frames = len(frame_files)
+
+    # Determine frame indices
+    if num_frames is None:
+        frame_indices = list(range(total_frames))
+    else:
+        if total_frames < num_frames:
+            frame_indices = list(range(total_frames)) + [total_frames - 1] * (
+                num_frames - total_frames)
+        elif sample_strategy == 'uniform':
+            frame_indices = np.linspace(0,
+                                        total_frames - 1,
+                                        num_frames,
+                                        dtype=int).tolist()
+        elif sample_strategy == 'random':
+            frame_indices = sorted(
+                np.random.choice(total_frames, num_frames, replace=False))
+        else:
+            raise ValueError(f"Unknown sample_strategy: {sample_strategy}")
+
+    # Load frames
+    frames = []
+    for idx in frame_indices:
+        frame_path = frame_files[idx]
+        frame = cv2.imread(str(frame_path))
+
+        if frame is None:
+            raise RuntimeError(f"Failed to load frame: {frame_path}")
+
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames.append(frame)
+
+    # Stack and convert to tensor
+    frames = np.stack(frames)  # [T, H, W, C]
+    frames = torch.from_numpy(frames).permute(0, 3, 1,
+                                              2).float()  # [T, C, H, W]
+
+    return frames
+
+
+def _detect_video_format(path: str | Path) -> str:
+    """
+    Detect if path is a video file or frame directory.
+    
+    Returns:
+        'video_file', 'frame_directory', or 'unknown'
+    """
+    path = Path(path)
+
+    if path.is_file():
+        return 'video_file'
+    elif path.is_dir():
+        # Check if contains image files
+        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
+        for ext in image_extensions:
+            if list(path.glob(f"*{ext}")):
+                return 'frame_directory'
+        return 'unknown'
+    else:
+        raise ValueError(f"Path does not exist: {path}")
+
+
+def load_video_auto(video_path: str | Path,
+                    num_frames: int | None = 16,
+                    sample_strategy: str = 'uniform') -> torch.Tensor:
+    """
+    Automatically detect format and load video.
+    
+    Supports:
+    - Video files (MP4, AVI, MOV, MKV)
+    - Frame directories (JPG, PNG)
+    
+    Args:
+        video_path: Path to video file or frame directory
+        num_frames: Number of frames to extract
+        sample_strategy: 'uniform' or 'random'
+    
+    Returns:
+        video: [T, C, H, W]
+    """
+    format_type = _detect_video_format(video_path)
+
+    if format_type == 'video_file':
+        return _load_video_cv2(video_path, num_frames, sample_strategy)
+    elif format_type == 'frame_directory':
+        return _load_video_from_frames(video_path, num_frames, sample_strategy)
+    else:
+        raise ValueError(f"Unknown video format at {video_path}")
+
+
+def sample_clips_from_video(
+        video: torch.Tensor,
+        num_frames_per_clip: int = 16,
+        num_clips: int = 1,
+        strategy: str | ClipSamplingStrategy = ClipSamplingStrategy.BEGINNING,
+        frame_stride: int = 1,
+        temporal_stride: int = 1) -> list[torch.Tensor]:
+    """
+    Sample clips from a video with various strategies.
+        
+    Args:
+        video: [T, C, H, W] full video
+        num_frames_per_clip: Frames per clip
+        num_clips: Number of clips to extract
+        strategy: ClipSamplingStrategy or string ('beginning', 'random', etc.)
+        frame_stride: Skip frames (FPS control: 1=all, 2=every 2nd, 8=every 8th)
+        temporal_stride: Stride between clips for sliding window
+    
+    Returns:
+        List of clips, each [num_frames_per_clip, C, H, W]
+    
+    Examples:
+        >>> # Beginning clip (most common for FVD)
+        >>> clips = sample_clips_from_video(video, 16, strategy='beginning')
+        
+        >>> # Multiple random clips
+        >>> clips = sample_clips_from_video(video, 16, num_clips=4, strategy='random')
+        
+        >>> # Subsample FPS by 2x (every 2nd frame)
+        >>> clips = sample_clips_from_video(video, 16, frame_stride=2)
+        
+        >>> # Sliding window with overlap
+        >>> clips = sample_clips_from_video(video, 16, strategy='sliding', temporal_stride=8)
+    """
+    # Convert string to enum if needed
+    if isinstance(strategy, str):
+        strategy = ClipSamplingStrategy(strategy)
+
+    T, C, H, W = video.shape
+
+    # Apply frame stride (FPS subsampling)
+    if frame_stride > 1:
+        video = video[::frame_stride]
+        T = len(video)
+
+    effective_clip_length = num_frames_per_clip
+
+    # Handle videos shorter than clip length
+    if effective_clip_length > T:
+        pad_length = effective_clip_length - T
+        last_frame = video[-1:].repeat(pad_length, 1, 1, 1)
+        video = torch.cat([video, last_frame], dim=0)
+        T = len(video)
+
+    clips = []
+
+    if strategy == ClipSamplingStrategy.BEGINNING:
+        # Take first clip (most common for FVD evaluation)
+        clip = video[:effective_clip_length]
+        clips.append(clip)
+
+    elif strategy == ClipSamplingStrategy.MIDDLE:
+        # Take middle clip
+        start = (T - effective_clip_length) // 2
+        clip = video[start:start + effective_clip_length]
+        clips.append(clip)
+
+    elif strategy == ClipSamplingStrategy.RANDOM:
+        # Sample N random clips
+        for _ in range(num_clips):
+            if effective_clip_length == T:
+                start = 0
+            else:
+                start = np.random.randint(0, T - effective_clip_length + 1)
+            clip = video[start:start + effective_clip_length]
+            clips.append(clip)
+
+    elif strategy == ClipSamplingStrategy.UNIFORM:
+        # Uniformly spaced clips
+        if num_clips == 1:
+            # Single clip from middle
+            start = (T - effective_clip_length) // 2
+            clip = video[start:start + effective_clip_length]
+            clips.append(clip)
+        else:
+            # Multiple uniformly spaced clips
+            step = (T - effective_clip_length) / (num_clips -
+                                                  1) if num_clips > 1 else 0
+            for i in range(num_clips):
+                start = int(i * step)
+                start = min(start, T - effective_clip_length)
+                clip = video[start:start + effective_clip_length]
+                clips.append(clip)
+
+    elif strategy == ClipSamplingStrategy.SLIDING:
+        # Sliding window with stride
+        for start in range(0, T - effective_clip_length + 1, temporal_stride):
+            clip = video[start:start + effective_clip_length]
+            clips.append(clip)
+            if len(clips) >= num_clips:
+                break
+
+    elif strategy == ClipSamplingStrategy.ALL:
+        # All possible clips (overlapping)
+        for start in range(T - effective_clip_length + 1):
+            clip = video[start:start + effective_clip_length]
+            clips.append(clip)
+
+    else:
+        raise ValueError(f"Unknown strategy: {strategy}")
+
+    return clips
+
+
+def load_video_clips_streaming(directory: str | Path,
+                               num_frames: int = 16,
+                               max_videos: int | None = None,
+                               clip_strategy: str
+                               | ClipSamplingStrategy = 'beginning',
+                               frame_stride: int = 1,
+                               num_clips_per_video: int = 1,
+                               video_extensions: list[str] | None = None,
+                               support_frame_dirs: bool = True,
+                               target_size: tuple[int, int] | None = (224, 224),
+                               verbose: bool = True) -> Iterator[torch.Tensor]:
+    """
+    This generator yields clips one-by-one instead of loading all videos into RAM.
+    Perfect for large datasets where memory is limited.
+    
+    Args:
+        directory: Path to directory with videos
+        num_frames: Frames per clip
+        max_videos: Max videos to load
+        clip_strategy: 'beginning', 'random', 'uniform', etc.
+        frame_stride: Frame skip (1=all, 2=every 2nd, 8=every 8th)
+        num_clips_per_video: Number of clips per video
+        video_extensions: Video file extensions
+        support_frame_dirs: Also load frame directories
+        target_size: Resize clips to (H, W). If None, keep original size.
+        verbose: Show progress
+    
+    Yields:
+        clip: [T, C, H, W] individual clips
+    
+    Example:
+        >>> for clip in load_video_clips_streaming('data/videos/', num_frames=16):
+        >>>     features = model.extract_features(clip.unsqueeze(0))
+        >>>     # Process one clip at a time - low memory usage!
+    """
+    if video_extensions is None:
+        video_extensions = ['.mp4', '.avi', '.mov', '.mkv']
+
+    directory = Path(directory)
+
+    if not directory.exists():
+        raise FileNotFoundError(f"Directory not found: {directory}")
+
+    # Find video paths
+    video_paths: list[Path] = []
+
+    # Find video files
+    for ext in video_extensions:
+        video_paths.extend(directory.glob(f"**/*{ext}"))
+
+    # Find frame directories if enabled
+    if support_frame_dirs:
+        for subdir in directory.iterdir():
+            if subdir.is_dir():
+                # Check if it contains frames
+                image_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
+                for ext in image_extensions:
+                    if list(subdir.glob(f"*{ext}")):
+                        video_paths.append(subdir)
+                        break
+
+    if len(video_paths) == 0:
+        raise ValueError(f"No videos found in {directory}")
+
+    video_paths = sorted(video_paths)
+
+    if max_videos is not None:
+        video_paths = video_paths[:max_videos]
+
+    if verbose:
+        print(f"Found {len(video_paths)} videos in {directory}")
+        if num_clips_per_video > 1:
+            print(f"Extracting {num_clips_per_video} clips per video...")
+        if frame_stride > 1:
+            print(f"Subsampling frames with stride {frame_stride}...")
+        if target_size:
+            print(f"Resizing clips to {target_size}...")
+
+    # Track statistics
+    failed_count = 0
+    total_clips = 0
+
+    iterator = tqdm(video_paths,
+                    desc="Loading videos") if verbose else video_paths
+
+    for video_path in iterator:
+        try:
+            # Load full video
+            video = load_video_auto(video_path,
+                                    num_frames=None,
+                                    sample_strategy='uniform')
+
+            # Sample clips from video
+            clips = sample_clips_from_video(video,
+                                            num_frames_per_clip=num_frames,
+                                            num_clips=num_clips_per_video,
+                                            strategy=clip_strategy,
+                                            frame_stride=frame_stride)
+
+            if target_size is not None:
+                resized_clips = []
+                for clip in clips:
+                    T, C, H, W = clip.shape
+                    if target_size != (H, W):
+                        # Resize to target size
+                        clip = clip.contiguous(
+                        )  # Fix non-contiguous tensors first
+                        clip_flat = clip.view(T * C, H,
+                                              W).unsqueeze(0)  # [1, T*C, H, W]
+                        clip_resized = torch.nn.functional.interpolate(
+                            clip_flat,
+                            size=target_size,
+                            mode='bilinear',
+                            align_corners=False)
+                        clip = clip_resized.squeeze(0).view(
+                            T, C, target_size[0],
+                            target_size[1])  # Back to [T, C, H, W]
+                    resized_clips.append(clip)
+                clips = resized_clips
+
+            # Yield clips one by one
+            for clip in clips:
+                yield clip
+                total_clips += 1
+
+            # Free memory
+            del video, clips
+
+        except Exception as e:
+            failed_count += 1
+            if verbose:
+                print(f"\nWarning: Failed to load {video_path}: {e}")
+            continue
+
+    # Validate
+    if total_clips == 0:
+        raise RuntimeError(f"Failed to load any videos from {directory}")
+
+    failure_rate = failed_count / len(video_paths)
+    if failure_rate > 0.1:  # More than 10% failed
+        print(
+            f"\nWARNING: {failure_rate:.1%} of videos failed to load ({failed_count}/{len(video_paths)})"
+        )
+
+    if verbose:
+        print(
+            f"\nSuccessfully loaded {total_clips} clips from {len(video_paths) - failed_count} videos"
+        )

benchmarks/scripts/run.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+# 1. Install missing dependency
+pip install -q opencv-python-headless transformers huggingface_hub
+
+# 2. Run FVD script
+python benchmarks/fvd/run_fvd.py

benchmarks/scripts/setup_fvd.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+# 1. Install missing dependency
+pip install -q opencv-python-headless

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/source/_static/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_vsa.py

@@ -1,156 +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, non_pad_index, dO):
-    Q = Q.detach().requires_grad_()
-    K = K.detach().requires_grad_()
-    V = V.detach().requires_grad_()
-    
-    q_padded = vsa_pad(Q, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
-    k_padded = vsa_pad(K, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
-    v_padded = vsa_pad(V, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
-    output, _= block_sparse_attn(q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes)
-    output = output[:, :, 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=32, 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, k, device)
-    full_mask = create_full_mask_from_block_mask(block_mask, 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)
-
-        # 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, 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 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)
-
-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)
-    print("\nAnalysis completed for all modes.")

csrc/attn/tests/utils.py

@@ -1,54 +0,0 @@
-import torch
-
-def generate_block_sparse_mask_for_function(h, num_blocks, k, device="cuda"):
-    """
-    Generate block sparse mask of shape [h, num_blocks, num_blocks].
-    
-    Args:
-        h: number of heads
-        num_blocks: number of blocks
-        k: number of kv blocks each q block attends to
-        device: device to create tensors on
-
-    Returns:
-        block_sparse_mask: [h, num_blocks, num_blocks] bool tensor
-    """
-    k = min(k, num_blocks)
-    scores = torch.rand(h, num_blocks, num_blocks, device=device)
-    _, indices = torch.topk(scores, k, dim=-1)
-    block_sparse_mask = torch.zeros(h, num_blocks, num_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, variable_block_sizes, device="cuda"):
-    """
-    Convert block-level sparse mask to full attention mask.
-    
-    Args:
-        block_sparse_mask: [h, num_blocks, num_blocks] bool tensor
-        variable_block_sizes: [num_blocks] tensor
-        device: device to create tensors on
-
-    Returns:
-        full_mask: [h, S, S] bool tensor where S = total sequence length
-    """
-    h, num_blocks, _ = block_sparse_mask.shape
-    total_seq_len = variable_block_sizes.sum().item()
-    cumsum = torch.cat([torch.tensor([0], device=device), variable_block_sizes.cumsum(dim=0)[:-1]])
-
-    full_mask = torch.zeros(h, total_seq_len, total_seq_len, dtype=torch.bool, device=device)
-
-    for head in range(h):
-        for q_block in range(num_blocks):
-            q_start = cumsum[q_block]
-            q_end = q_start + variable_block_sizes[q_block]
-
-            for kv_block in range(num_blocks):
-                if block_sparse_mask[head, q_block, kv_block]:
-                    kv_start = cumsum[kv_block]
-                    kv_end = kv_start + 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_attn_triton.py

@@ -1,449 +0,0 @@
-"""
-Fused Attention
-===============
-
-This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
-(https://tridao.me/publications/flash2/flash2.pdf)
-
-Credits: OpenAI kernel team
-"""
-
-import pytest
-import torch
-import triton
-import triton.language as tl
-
-# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
-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.
-configs = [
-    triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
-    for BM in [64]\
-    for BN in [64]\
-    for s in [3, 4, 7]\
-    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):
-    """
-    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_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
-
-    # ----- base pointers -----
-    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))
-
-    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)
-
-    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)
-    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
-    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
-    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
-    qk_scale = sm_scale * 1.44269504  # 1/ln2
-    q = tl.load(Q_ptr)
-
-    # ----- sparse loop over valid K/V tiles -----
-    for i in range(0, kv_blocks):
-        kv_idx = tl.load(kv_ptr + i).to(tl.int32)
-        block_size = tl.load(variable_block_sizes + kv_idx)
-        K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N))
-        V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0))
-
-        k = tl.load(K_ptr)
-        qk = tl.dot(q, k)
-        # mask out invalid columns
-        mask = tl.arange(0, BLOCK_N) < block_size
-        qk = tl.where(mask[None, :], qk, -float("inf"))
-
-        m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
-        p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
-        l_ij = tl.sum(p, 1)
-
-        alpha = tl.math.exp2(m_i - m_ij)
-        l_i = l_i * alpha + l_ij
-        acc = acc * alpha[:, None]
-
-        v = tl.load(V_ptr)
-        acc = tl.dot(p.to(tl.bfloat16), v, acc)
-        m_i = m_ij
-
-    # ----- epilogue -----
-    m_i += tl.math.log2(l_i)
-    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  #
-                         ):
-    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)
-    delta = tl.sum(o * do, axis=1)
-    # write-back
-    tl.store(Delta + off_hz * N_CTX + off_m, delta)
-
-
-# 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):
-    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)
-    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
-    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    # 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
-    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
-    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
-        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)
-        m = tl.load(M + offs_m)
-        qkT = tl.dot(k, qT)
-        pT = tl.math.exp2(qkT - m[None, :])
-        mask = tl.arange(0, BLOCK_N1) < block_size
-        pT = tl.where(mask[:, None], pT, 0.0)
-
-        do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
-        # Compute dV.
-        ppT = pT
-        ppT = ppT.to(tl.bfloat16)
-        dv += tl.dot(ppT, do)
-        # D (= delta) is pre-divided by ds_scale.
-        Di = tl.load(D + offs_m)
-        # Compute dP and dS.
-        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
-        dsT = pT * (dpT - Di[None, :])
-        dsT = dsT.to(tl.bfloat16)
-        dk += tl.dot(dsT, tl.trans(qT))
-        # Increment pointers.
-    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):
-    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)
-    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
-    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
-    # D (= delta) is pre-divided by ds_scale.
-    Di = tl.load(D + offs_m)
-    # 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_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):
-        block_sparse_offset = (tl.load(kv_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_n * stride_tok
-        block_size = tl.load(variable_block_sizes + blk_idx//2) - (blk_idx%2) * step_n
-        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)
-        # Compute dP and dS.
-        dp = tl.dot(do, vT).to(tl.float32)
-        ds = p * (dp - Di[:, None])
-        ds = ds.to(tl.bfloat16)
-        # Compute dQ.
-        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
-        dq += tl.dot(ds, tl.trans(kT))
-        # Increment pointers.
-    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):
-    LN2 = 0.6931471824645996  # = ln(2)
-
-    bhid = tl.program_id(2)
-    off_chz = (bhid * N_CTX).to(tl.int64)
-    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
-    pid = tl.program_id(0)
-
-    # offset pointers for batch/head
-    Q += adj
-    K += adj
-    V += adj
-    DO += adj
-    DQ += adj
-    DK += adj
-    DV += adj
-    M += off_chz
-    D += off_chz
-
-    # load scales
-    offs_k = tl.arange(0, HEAD_DIM)
-
-    start_n = pid * BLOCK_N1
-    start_m = 0
-
-    offs_n = start_n + tl.arange(0, BLOCK_N1)
-
-    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
-    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
-
-    # load K and V: they stay in SRAM throughout the inner loop.
-    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,  #
-        DO,  #
-        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  #
-    )
-
-    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
-    tl.store(dv_ptrs, dv)
-
-    # Write back dK.
-    dk *= sm_scale
-    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
-    tl.store(dk_ptrs, dk)
-
-    # THIS BLOCK DOES DQ:
-    start_m = pid * BLOCK_M2
-    end_n = 0
-
-    offs_m = start_m + tl.arange(0, BLOCK_M2)
-
-    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
-    dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
-    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
-
-    m = tl.load(M + offs_m)
-    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  #
-                      )
-    # 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):
-    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]}"
-    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
-    )
-
-    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):
-    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)
-    dk = torch.empty_like(k)
-    dv = torch.empty_like(v)
-    BATCH, N_HEAD, N_CTX = q.shape[:3]
-    BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
-    RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
-    arg_k = k
-    arg_k = arg_k * (sm_scale * RCP_LN2)
-    PRE_BLOCK = 64
-    assert N_CTX % PRE_BLOCK == 0
-    pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
-    delta = torch.empty_like(M)
-    _attn_bwd_preprocess[pre_grid](
-        o, do,  #
-        delta,  #
-        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,
-        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,  #
-        HEAD_DIM=D  #
-    )
-
-    return dq, dk, dv
-
-

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

demo/gradio_web_demo.py

@@ -1,195 +0,0 @@
-import argparse
-import os
-import tempfile
-
-import gradio as gr
-import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import export_to_video
-
-from fastvideo.distill.solver import PCMFMScheduler
-from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
-from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
-
-
-def init_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--prompts", nargs="+", default=[])
-    parser.add_argument("--num_frames", type=int, default=25)
-    parser.add_argument("--height", type=int, default=480)
-    parser.add_argument("--width", type=int, default=848)
-    parser.add_argument("--num_inference_steps", type=int, default=8)
-    parser.add_argument("--guidance_scale", type=float, default=4.5)
-    parser.add_argument("--model_path", type=str, default="data/mochi")
-    parser.add_argument("--seed", type=int, default=12345)
-    parser.add_argument("--transformer_path", type=str, default=None)
-    parser.add_argument("--scheduler_type", type=str, default="pcm_linear_quadratic")
-    parser.add_argument("--lora_checkpoint_dir", type=str, default=None)
-    parser.add_argument("--shift", type=float, default=8.0)
-    parser.add_argument("--num_euler_timesteps", type=int, default=50)
-    parser.add_argument("--linear_threshold", type=float, default=0.1)
-    parser.add_argument("--linear_range", type=float, default=0.75)
-    parser.add_argument("--cpu_offload", action="store_true")
-    return parser.parse_args()
-
-
-def load_model(args):
-    if args.scheduler_type == "euler":
-        scheduler = FlowMatchEulerDiscreteScheduler()
-    else:
-        linear_quadratic = True if "linear_quadratic" in args.scheduler_type else False
-        scheduler = PCMFMScheduler(
-            1000,
-            args.shift,
-            args.num_euler_timesteps,
-            linear_quadratic,
-            args.linear_threshold,
-            args.linear_range,
-        )
-
-    if args.transformer_path:
-        transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path)
-    else:
-        transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder="transformer/")
-
-    pipe = MochiPipeline.from_pretrained(args.model_path, transformer=transformer, scheduler=scheduler)
-    pipe.enable_vae_tiling()
-    # pipe.to(device)
-    # if args.cpu_offload:
-    pipe.enable_sequential_cpu_offload()
-    return pipe
-
-
-def generate_video(
-    prompt,
-    negative_prompt,
-    use_negative_prompt,
-    seed,
-    guidance_scale,
-    num_frames,
-    height,
-    width,
-    num_inference_steps,
-    randomize_seed=False,
-):
-    if randomize_seed:
-        seed = torch.randint(0, 1000000, (1, )).item()
-
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-
-    if not use_negative_prompt:
-        negative_prompt = None
-
-    with torch.autocast("cuda", dtype=torch.bfloat16):
-        output = pipe(
-            prompt=[prompt],
-            negative_prompt=negative_prompt,
-            height=height,
-            width=width,
-            num_frames=num_frames,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            generator=generator,
-        ).frames[0]
-
-    output_path = os.path.join(tempfile.mkdtemp(), "output.mp4")
-    export_to_video(output, output_path, fps=30)
-    return output_path, seed
-
-
-examples = [
-    "A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand’s movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough.",
-    "A vintage train snakes through the mountains, its plume of white steam rising dramatically against the jagged peaks. The cars glint in the late afternoon sun, their deep crimson and gold accents lending a touch of elegance. The tracks carve a precarious path along the cliffside, revealing glimpses of a roaring river far below. Inside, passengers peer out the large windows, their faces lit with awe as the landscape unfolds.",
-    "A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background. Strings of fairy lights hang above, casting a warm, golden glow over the scene. Groups of people gather around high tables, their laughter blending with the soft rhythm of live jazz. The aroma of freshly mixed cocktails and charred appetizers wafts through the air, mingling with the cool night breeze.",
-]
-
-args = init_args()
-pipe = load_model(args)
-print("load model successfully")
-with gr.Blocks() as demo:
-    gr.Markdown("# Fastvideo Mochi Video Generation Demo")
-
-    with gr.Group():
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            run_button = gr.Button("Run", scale=0)
-        result = gr.Video(label="Result", show_label=False)
-
-    with gr.Accordion("Advanced options", open=False):
-        with gr.Group():
-            with gr.Row():
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=1024,
-                    step=32,
-                    value=args.height,
-                )
-                width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=args.width)
-
-            with gr.Row():
-                num_frames = gr.Slider(
-                    label="Number of Frames",
-                    minimum=21,
-                    maximum=163,
-                    value=args.num_frames,
-                )
-                guidance_scale = gr.Slider(
-                    label="Guidance Scale",
-                    minimum=1,
-                    maximum=12,
-                    value=args.guidance_scale,
-                )
-                num_inference_steps = gr.Slider(
-                    label="Inference Steps",
-                    minimum=4,
-                    maximum=100,
-                    value=args.num_inference_steps,
-                )
-
-            with gr.Row():
-                use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
-            seed = gr.Slider(label="Seed", minimum=0, maximum=1000000, step=1, value=args.seed)
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            seed_output = gr.Number(label="Used Seed")
-
-    gr.Examples(examples=examples, inputs=prompt)
-
-    use_negative_prompt.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_negative_prompt,
-        outputs=negative_prompt,
-    )
-
-    run_button.click(
-        fn=generate_video,
-        inputs=[
-            prompt,
-            negative_prompt,
-            use_negative_prompt,
-            seed,
-            guidance_scale,
-            num_frames,
-            height,
-            width,
-            num_inference_steps,
-            randomize_seed,
-        ],
-        outputs=[result, seed_output],
-    )
-
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)

demo/qualitative_results.md

@@ -1,15 +0,0 @@
-Fast-Hunyuan comparison with original Hunyuan, achieving an 8X diffusion speed boost with the FastVideo framework.
-
-https://github.com/user-attachments/assets/064ac1d2-11ed-4a0c-955b-4d412a96ef30
-
-Fast-Mochi comparison with original Mochi, achieving an 8X diffusion speed boost with the FastVideo framework.
-
-https://github.com/user-attachments/assets/5fbc4596-56d6-43aa-98e0-da472cf8e26c
-
-Comparison between OpenAI Sora, original Hunyuan and FastHunyuan
-
-https://github.com/user-attachments/assets/d323b712-3f68-42b2-952b-94f6a49c4836
-
-Comparison between original FastHunyuan, LLM-INT8 quantized FastHunyuan and NF4 quantized FastHunyuan
-
-https://github.com/user-attachments/assets/cf89efb5-5f68-4949-a085-f41c1ef26c94

docker/Dockerfile.python3.10

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.8.0-devel-ubuntu22.04
+FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     uv pip install --no-cache-dir --upgrade pip && \
     uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp310-cp310-linux_x86_64.whl
 
 COPY . .
 
@@ -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
 
-# 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

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.8.0-devel-ubuntu22.04
+FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     uv pip install --no-cache-dir --upgrade pip && \
     uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp311-cp311-linux_x86_64.whl
 
 COPY . .
 
@@ -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
 
-# 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

@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     uv pip install --no-cache-dir --upgrade pip && \
     uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp312-cp312-linux_x86_64.whl
 
 COPY . .
 
@@ -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
 
 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
 
-# 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 --rocm
+
+EXPOSE 22

docs/Makefile

@@ -1,26 +0,0 @@
-# Minimal makefile for Sphinx documentation
-#
-
-# You can set these variables from the command line, and also
-# from the environment for the first two.
-SPHINXOPTS    ?=
-SPHINXBUILD   ?= sphinx-build
-SOURCEDIR     = source
-BUILDDIR      = build
-
-# Put it first so that "make" without argument is like "make help".
-help:
-	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
-
-.PHONY: help Makefile
-
-# Catch-all target: route all unknown targets to Sphinx using the new
-# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
-%: Makefile
-	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
-
-clean:
-	@$(SPHINXBUILD) -M clean "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
-	rm -rf "$(SOURCEDIR)/getting_started/examples"
-	rm -rf "$(SOURCEDIR)/inference/examples"
-	rm -rf "$(SOURCEDIR)/training/examples"

docs/README.md

@@ -1,20 +1,39 @@
-# FastVideo documents
+# FastVideo Documentation
 
-## Build the docs
+This directory contains the FastVideo documentation built with MkDocs.
+
+## Build the docs locally
 
 ```bash
-# Install dependencies.
-pip install -r requirements-docs.txt
+# Install dependencies
+pip install -r requirements-mkdocs.txt
 
-# Build the docs.
-make clean
-make html
+# Serve docs with live reload (recommended for development)
+mkdocs serve
+
+# Or build static site
+mkdocs build
 ```
 
-## Open the docs with your browser
+## View the docs
+
+### Development server (with live reload)
 
 ```bash
-python -m http.server -d build/html/
+mkdocs serve
 ```
 
-Launch your browser and open localhost:8000.
+Then open your browser to: http://127.0.0.1:8000
+
+### Static build
+
+```bash
+mkdocs build
+python -m http.server -d site/
+```
+
+Then open your browser to: http://localhost:8000
+
+## Automatic Deployment
+
+Documentation is automatically built and deployed to GitHub Pages when changes are pushed to the `main` branch via the `.github/workflows/docs.yml` workflow.

docs/api/fastvideo.md

@@ -0,0 +1,257 @@
+# FastVideo API Reference
+
+This page contains the complete API reference for the FastVideo library.
+
+## fastvideo
+
+### Modules
+
+| Name | Description |
+|------|-------------|
+| [attention](#fastvideoattention) | Attention mechanisms and backends for video generation |
+| [configs](#fastvideoconfigs) | Configuration classes for pipelines, models, and sampling |
+| [distributed](#fastvideodistributed) | Distributed execution and communication utilities |
+| [entrypoints](#fastvideoentrypoints) | Main API entry points for video generation |
+| [models](#fastvideomodels) | Model implementations (transformers, VAEs, schedulers) |
+| [pipelines](#fastvideopipelines) | Core pipeline classes for video diffusion |
+| [training](#fastvideotraining) | Training utilities and helpers |
+| [workflow](#fastvideoworkflow) | Workflow management and orchestration |
+| [dataset](#fastvideodataset) | Dataset handling and preprocessing |
+| [layers](#fastvideolayers) | Custom neural network layers |
+| [platforms](#fastvideoplatforms) | Platform-specific implementations |
+| [utils](#fastvideoutils) | Utility functions and helpers |
+| [worker](#fastvideoworker) | Execution workers for video generation |
+
+## fastvideo.attention
+
+::: fastvideo.attention
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.configs
+
+::: fastvideo.configs
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+### Submodules
+
+#### fastvideo.configs.pipelines
+
+::: fastvideo.configs.pipelines
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.configs.models
+
+::: fastvideo.configs.models
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.configs.sample
+
+::: fastvideo.configs.sample
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+## fastvideo.distributed
+
+::: fastvideo.distributed
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.entrypoints
+
+::: fastvideo.entrypoints
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.models
+
+::: fastvideo.models
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+### Submodules
+
+#### fastvideo.models.registry
+
+::: fastvideo.models.registry
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.models.loader
+
+::: fastvideo.models.loader
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+## fastvideo.registry
+
+::: fastvideo.registry
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 3
+
+## fastvideo.pipelines
+
+::: fastvideo.pipelines
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+### Submodules
+
+#### fastvideo.pipelines.composed_pipeline_base
+
+::: fastvideo.pipelines.composed_pipeline_base
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.pipelines.lora_pipeline
+
+::: fastvideo.pipelines.lora_pipeline
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.pipelines.pipeline_batch_info
+
+::: fastvideo.pipelines.pipeline_batch_info
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.pipelines.pipeline_registry
+
+::: fastvideo.pipelines.pipeline_registry
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.pipelines.stages
+
+::: fastvideo.pipelines.stages
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+## fastvideo.training
+
+::: fastvideo.training
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.workflow
+
+::: fastvideo.workflow
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.dataset
+
+::: fastvideo.dataset
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.layers
+
+::: fastvideo.layers
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.platforms
+
+::: fastvideo.platforms
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.utils
+
+::: fastvideo.utils
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 3
+
+## fastvideo.worker
+
+::: fastvideo.worker
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3

docs/api/summary.md

@@ -0,0 +1,27 @@
+# API Summary
+
+This page provides a quick overview of the main FastVideo API components.
+
+## Video Generator
+
+::: fastvideo.VideoGenerator
+    options:
+      show_root_heading: false
+      show_source: false
+      heading_level: 3
+
+## Initialization Configuration
+
+::: fastvideo.PipelineConfig
+    options:
+      show_root_heading: false
+      show_source: false
+      heading_level: 3
+
+## Sampling Configuration
+
+::: fastvideo.SamplingParam
+    options:
+      show_root_heading: false
+      show_source: false
+      heading_level: 3

docs/assets/custom.css

@@ -0,0 +1,43 @@
+.vertical-table-header th.head:not(.stub) {
+    writing-mode: sideways-lr;
+    white-space: nowrap;
+    max-width: 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 */
+.image-small {
+    max-width: 200px;
+    height: auto;
+}
+
+.image-medium {
+    max-width: 400px;
+    height: auto;
+}
+
+.image-large {
+    max-width: 600px;
+    height: auto;
+}
+
+.image-full {
+    max-width: 100%;
+    height: auto;
+}
+
+/* Responsive images */
+img {
+    max-width: 100%;
+    height: auto;
+}
+
+/* Center images */
+.image-center {
+    display: block;
+    margin: 0 auto;
+}