fix

Co-authored-by: KyleS1016 <kyle.s@gmicloud.ai>

.buildkite/pipeline.yml

@@ -104,6 +104,18 @@ steps:
                   - TEST_TYPE=distillation_dmd
                 agents:
                   queue: "default"
+            - path:
+                - "fastvideo/training/*self_forcing_distillation_pipeline.py"
+                - "fastvideo/tests/training/self-forcing/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Self-Forcing Tests"
+                env:
+                  - TEST_TYPE=self_forcing
+                agents:
+                  queue: "default"
             - path:
                 - "fastvideo/**"
                 - "pyproject.toml"
@@ -117,11 +129,11 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
-                - "csrc/attn/vsa/**"
-                - "csrc/attn/tk/**"
-                - "csrc/attn/setup_vsa.py"
-                - "csrc/attn/config_vsa.py"
-                - "csrc/attn/vsa.cpp"
+                - "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"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -133,10 +145,10 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
-                - "csrc/attn/st_attn/**"
-                - "csrc/attn/setup_sta.py"
-                - "csrc/attn/config_sta.py"
-                - "csrc/attn/st_attn.cpp"
+                - "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"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -147,10 +159,10 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "csrc/attn/st_attn/**"
-                - "csrc/attn/setup_sta.py"
-                - "csrc/attn/config_sta.py"
-                - "csrc/attn/st_attn.cpp"
+                - "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"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -161,12 +173,12 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "csrc/attn/vsa/**"
-                - "csrc/attn/tk/**"
+                - "csrc/attn/video_sparse_attn/**"
+                - "csrc/attn/video_sparse_attn/tk/**"
                 - "csrc/attn/tests/test_vsa.py"
-                - "csrc/attn/setup_vsa.py"
-                - "csrc/attn/config_vsa.py"
-                - "csrc/attn/vsa.cpp"
+                - "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:
@@ -176,3 +188,48 @@ steps:
                   - 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/attention/backends/vmoba.py"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Inference Tests VMoBA"
+                env:  
+                  - TEST_TYPE=inference_vmoba
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Unit Tests"
+                env:
+                  - 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..."
@@ -109,6 +109,27 @@ case "$TEST_TYPE" in
         log "Running distillation DMD tests..."
         MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_distill_dmd_tests"
         ;;
+        # run_inference_tests_vmoba
+    "self_forcing")
+        log "Running self-forcing tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_self_forcing_tests"
+        ;;
+    "inference_vmoba")
+        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,82 +1,65 @@
-# 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: 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

.github/workflows/pr-test.yml

@@ -62,8 +62,8 @@ on:
         required: false
         default: false
         type: boolean
-      run_nightly_test:
-        description: "Run nightly-test"
+      run_unit_test:
+        description: "Run unit-test"
         required: false
         default: false
         type: boolean
@@ -93,6 +93,7 @@ jobs:
       inference-test-STA: ${{ steps.filter.outputs.inference-test-STA }}
       precision-test-STA: ${{ steps.filter.outputs.precision-test-STA }}
       precision-test-VSA: ${{ steps.filter.outputs.precision-test-VSA }}
+      unit-test: ${{ steps.filter.outputs.unit-test }}
     steps:
       - uses: actions/checkout@v4
       - uses: dorny/paths-filter@v3
@@ -102,18 +103,21 @@ jobs:
             # Define reusable path patterns
             common-paths: &common-paths
               - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.10'
+              - 'docker/Dockerfile.python3.11'
               - 'docker/Dockerfile.python3.12'
             sta-kernel-paths: &sta-kernel-paths
-              - 'csrc/attn/st_attn/**'
-              - 'csrc/attn/setup_sta.py'
-              - 'csrc/attn/config_sta.py'
-              - 'csrc/attn/st_attn.cpp'
+              - 'csrc/attn/sliding_tile_attn/**'
+              - 'csrc/attn/sliding_tile_attn/tk/**'
+              - 'csrc/attn/sliding_tile_attn/setup.py'
+              - 'csrc/attn/sliding_tile_attn/config_sta.py'
+              - 'csrc/attn/sliding_tile_attn/st_attn.cpp'
             vsa-kernel-paths: &vsa-kernel-paths
-              - 'csrc/attn/vsa/**'
-              - 'csrc/attn/tk/**'
-              - 'csrc/attn/setup_vsa.py'
-              - 'csrc/attn/config_vsa.py'
-              - 'csrc/attn/vsa.cpp'
+              - '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'
             vsa-paths: &vsa-paths
               - 'fastvideo/**'
               - *common-paths
@@ -154,6 +158,9 @@ jobs:
             precision-test-VSA:
               - *common-paths
               - *vsa-kernel-paths
+            unit-test:
+              - 'fastvideo/**'
+              - *common-paths
 
   encoder-test:
     needs: change-filter
@@ -234,7 +241,7 @@ jobs:
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-  
+
   training-test:
     needs: change-filter
     if: >-
@@ -332,23 +339,42 @@ jobs:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
 
-  nightly-test:
+  unit-test:
+    needs: change-filter
     if: >-
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
+      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.unit-test == 'true') || 
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_unit_test == 'true')
     uses: ./.github/workflows/runpod-test.yml
     with:
-      job_id: "nightly-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 4
+      job_id: "unit-test"
+      gpu_type: "NVIDIA L40S"
+      gpu_count: 1
       volume_size: 100
       disk_size: 100
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/nightly/test_e2e_overfit_single_sample.py -vs"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/dataset/ -vs && pytest ./fastvideo/workflow/ -vs && pytest ./fastvideo/entrypoints/ -vs"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+
+  # nightly-test:
+  #   if: >-
+  #     (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
+  #   uses: ./.github/workflows/runpod-test.yml
+  #   with:
+  #     job_id: "nightly-test"
+  #     gpu_type: "NVIDIA A40"
+  #     gpu_count: 4
+  #     volume_size: 100
+  #     disk_size: 100
+  #     image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+  #     test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/nightly/test_e2e_overfit_single_sample.py -vs"
+  #     timeout_minutes: 30
+  #   secrets:
+  #     RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+  #     RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
+  #     WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
 
   runpod-cleanup:
     # Add other jobs to this list as you create them
@@ -372,4 +398,4 @@ jobs:
           JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12", "training-test", "training-test-VSA", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]'
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
-        run: python .github/scripts/runpod_cleanup.py
+        run: python .github/scripts/runpod_cleanup.py

.github/workflows/sta-publish.yml

@@ -5,7 +5,7 @@ on:
     branches:
       - main
     paths:
-      - "csrc/attn/setup_sta.py"
+      - "csrc/attn/sliding_tile_attn/setup.py"
   workflow_dispatch:
 
 jobs:
@@ -23,13 +23,13 @@ jobs:
       - name: Check if version changed
         id: check-version
         run: |
-          cd csrc/attn
+          cd csrc/attn/sliding_tile_attn
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_sta.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup_sta.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -144,13 +144,13 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_sta.py bdist_wheel --dist-dir=dist
+          python setup.py bdist_wheel --dist-dir=dist
 
       - name: Rename wheel file
         run: |
-          cd csrc/attn
+          cd csrc/attn/sliding_tile_attn
 
           CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
           TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2)
@@ -165,7 +165,7 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: ${{ env.wheel_name }}
-          path: csrc/attn/dist/*.whl
+          path: csrc/attn/sliding_tile_attn/dist/*.whl
           retention-days: 90
 
   publish_package:
@@ -239,11 +239,11 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_sta.py sdist --dist-dir=dist
+          python setup.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: csrc/attn/dist/
+          packages-dir: csrc/attn/sliding_tile_attn/dist/

.github/workflows/vsa-publish.yml

@@ -5,7 +5,7 @@ on:
     branches:
       - main
     paths:
-      - "csrc/attn/setup_vsa.py"
+      - "csrc/attn/video_sparse_attn/setup.py"
   workflow_dispatch:
 
 jobs:
@@ -23,13 +23,13 @@ jobs:
       - name: Check if version changed
         id: check-version
         run: |
-          cd csrc/attn
+          cd csrc/attn/video_sparse_attn
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_vsa.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup_vsa.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -152,13 +152,13 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/video_sparse_attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_vsa.py bdist_wheel --dist-dir=dist
+          python setup.py bdist_wheel --dist-dir=dist
 
       - name: Rename wheel file
         run: |
-          cd csrc/attn
+          cd csrc/attn/video_sparse_attn
 
           CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
           TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2)
@@ -173,7 +173,7 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: ${{ env.wheel_name }}
-          path: csrc/attn/dist/*.whl
+          path: csrc/attn/video_sparse_attn/dist/*.whl
           retention-days: 90
 
   publish_package:
@@ -247,11 +247,11 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/video_sparse_attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_vsa.py sdist --dist-dir=dist
+          python setup.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: csrc/attn/dist/
+          packages-dir: csrc/attn/video_sparse_attn/dist/

.gitignore

@@ -14,6 +14,8 @@ wandb/
 *.pt
 cache_dir/
 wandb/
+venv/
+.venv/
 runs/
 samples/
 *validation/
@@ -37,12 +39,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,6 +64,9 @@ 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
+
+dmd_t2v_output/
+preprocess_output_text/

.gitmodules

@@ -1,3 +1,7 @@
-[submodule "csrc/attn/tk"]
-	path = csrc/attn/tk
+[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
 	url = https://github.com/HazyResearch/ThunderKittens.git

.pre-commit-config.yaml

@@ -10,11 +10,9 @@ exclude: |
     demo/.*|
     predict\.py|
     scripts/.*|
+    prompts/.*|
     fastvideo/data_preprocess/.*|
     fastvideo/dataset/.*|
-    fastvideo/distill/.*|
-    fastvideo/distill\.py|
-    fastvideo/distill_adv\.py|
     fastvideo/models/.*|
     fastvideo/sample/.*|
     fastvideo/train\.py|
@@ -44,10 +42,10 @@ repos:
   - id: codespell
     additional_dependencies: ['tomli']
     args: ['--toml', 'pyproject.toml']
-- repo: https://github.com/PyCQA/isort
-  rev: 6.0.1
-  hooks:
-  - id: isort
+# - repo: https://github.com/PyCQA/isort
+#   rev: 6.0.1
+#   hooks:
+#   - id: isort
 - repo: https://github.com/jackdewinter/pymarkdown
   rev: v0.9.30
   hooks:

README.md

@@ -1,13 +1,12 @@
 <div align="center">
 <img src=assets/logos/logo.svg width="30%"/>
 </div>
-
 **FastVideo is a unified post-training and inference framework for accelerated video generation.**
 
 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.
 
 <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-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/rG0QpZdw" target="_blank"> <b> WeChat </b> </a> |
+    | 🕹️ <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/"><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-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/c7g1qdD" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 <div align="center">
@@ -15,6 +14,7 @@ FastVideo features an end-to-end unified pipeline for accelerating diffusion mod
 </div>
 
 ## NEWS
+- ```2025/11/19```: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py)
 - ```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!
@@ -49,10 +49,10 @@ conda activate fastvideo
 pip install fastvideo
 ```
 
-Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) for more detailed installation instructions.
+Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) 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/).
+For our sparse distillation techniques, please see our [distillation docs](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/) and check out our [blog](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
 
 See below for recipes and datasets:
 
@@ -64,7 +64,7 @@ See below for recipes and datasets:
 
 ## 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:
+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/). Create a file called `example.py` with the following code:
 
 ```python
 import os
@@ -100,35 +100,32 @@ Run the script with:
 python example.py
 ```
 
-For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html).
+For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/).
 
 ### Other docs:
 
-- [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)
+- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview/)
+- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/)
 
 ## Distillation and Finetuning
-- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html)
+- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/)
 <!-- - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html) -->
 
-## 📑 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 -->
+## Awesome work using FastVideo or our research projects
 
-See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/issues/468).
+- [SGLang](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen): SGLang's diffusion inference functionality is based  on a fork of FastVideo on Sept. 24, 2025. [![Star](https://img.shields.io/github/stars/sgl-project/sglang.svg?style=social&label=Star)](https://github.com/sgl-project/sglang)
+
+- [DanceGRPO](https://github.com/XueZeyue/DanceGRPO): A  unified framework to adapt Group Relative Policy Optimization (GRPO) to visual generation paradigms. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/XueZeyue/DanceGRPO.svg?style=social&label=Star)](https://github.com/XueZeyue/DanceGRPO)
+- [SRPO](https://github.com/Tencent-Hunyuan/SRPO): A method to directly align the full diffusion trajectory with fine-grained human preference. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/Tencent-Hunyuan/SRPO.svg?style=social&label=Star)](https://github.com/Tencent-Hunyuan/SRPO)
+- [DCM](https://github.com/Vchitect/DCM): Dual-expert consistency model for efficient and high-quality video generation. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/Vchitect/DCM.svg?style=social&label=Star)](https://github.com/Vchitect/DCM)
+- [Hunyuan Video 1.5](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5): A leading lightweight video generation model, where they proposed SSTA based on Sliding Tile Attention. [![Star](https://img.shields.io/github/stars/Tencent-Hunyuan/HunyuanVideo-1.5.svg?style=social&label=Star)](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5)
+- [Kandinsky-5.0](https://github.com/kandinskylab/kandinsky-5): A family of diffusion models for video & image generation, where their NABLA attention includes a Sliding Tile Attention branch. [![Star](https://img.shields.io/github/stars/kandinskylab/kandinsky-5.svg?style=social&label=Star)](https://github.com/kandinskylab/kandinsky-5)
+- [LongCat Video](https://github.com/meituan-longcat/LongCat-Video): A foundational video generation model with 13.6B parameters with block-sparse attention similar to Video Sparse Attention. [![Star](https://img.shields.io/github/stars/meituan-longcat/LongCat-Video.svg?style=social&label=Star)](https://github.com/meituan-longcat/LongCat-Video)
 
 ## 🤝 Contributing
 
-We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview.html)
-
+We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview/).
+See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/issues/468).
 ## Acknowledgement
 We learned and reused code from the following projects:
 - [Wan-Video](https://github.com/Wan-Video)
@@ -155,8 +152,8 @@ If you find FastVideo useful, please considering citing our work:
 }
 
 @article{zhang2025vsa,
-  title={VSA: Faster Video Diffusion with Trainable Sparse Attention},
-  author={Zhang, Peiyuan and Huang, Haofeng and Chen, Yongqi and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
+  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}
 }

assets/full.svg

@@ -0,0 +1,18 @@
+<svg width="252" height="105" viewBox="0 0 252 105" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM100.768 13.1217L87.7028 29.4852H103.143L100.768 13.1217Z" fill="#356CFF"/>
+<path d="M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM109.081 90.697L116.802 65.8487C116.802 65.8487 120.959 65.8487 132.242 65.8487C143.525 65.8487 137.586 78.5759 135.211 84.0304C133.307 88.4021 127.491 90.697 122.74 90.697C117.989 90.697 109.081 90.697 109.081 90.697Z" fill="#356CFF"/>
+<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944H159.747C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273L125.188 48.273L124 37.97L147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852H131.836C120.142 29.4852 125.897 1.00043 141.337 1.00043L173.188 1.00056Z" fill="#356CFF"/>
+<path d="M179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056Z" fill="#356CFF"/>
+<path d="M161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM237.948 77.9692C239.984 70.6965 240.917 65.242 228.446 65.242C215.975 65.242 211.818 71.9087 210.037 77.9692C208.255 84.0298 208.255 91.3025 219.538 91.3025C230.821 91.3025 235.911 85.2419 237.948 77.9692Z" fill="#356CFF"/>
+<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944M173.188 1.00056C173.188 1.00056 156.777 1.00043 141.337 1.00043M173.188 1.00056L141.337 1.00043M141.337 20.3944C146.088 20.3944 150.839 20.3944 159.747 20.3944M141.337 20.3944H159.747M159.747 20.3944C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273M148.463 48.273C139.556 48.273 125.188 48.273 125.188 48.273M148.463 48.273L125.188 48.273M125.188 48.273L124 37.97M124 37.97C124 37.97 141.931 37.97 147.87 37.97M124 37.97L147.87 37.97M147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852M151.433 29.4852C146.682 29.4852 138.962 29.4852 131.836 29.4852M151.433 29.4852H131.836M131.836 29.4852C120.142 29.4852 125.897 1.00043 141.337 1.00043M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057ZM96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM87.7028 29.4852L100.768 13.1217L103.143 29.4852H87.7028ZM89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457ZM108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM116.802 65.8487L109.081 90.697C109.081 90.697 117.989 90.697 122.74 90.697C127.491 90.697 133.307 88.4021 135.211 84.0304C137.586 78.5759 143.525 65.8487 132.242 65.8487C120.959 65.8487 116.802 65.8487 116.802 65.8487ZM179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056ZM161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457ZM230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM228.446 65.242C240.917 65.242 239.984 70.6965 237.948 77.9692C235.911 85.2419 230.821 91.3025 219.538 91.3025C208.255 91.3025 208.255 84.0298 210.037 77.9692C211.818 71.9087 215.975 65.242 228.446 65.242Z" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M15.2524 55.5451L21.191 100.999L24.7541 100.999L18.8156 55.5451L15.2524 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M8.12646 55.5451L14.065 100.999L15.2527 100.999L9.31417 55.5451L8.12646 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M1 55.5451L6.93853 100.999L7.53239 100.999L1.59385 55.5451L1 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
+<path d="M15.2524 48.2724L30.0988 1H33.6619L18.8156 48.2724H15.2524Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M8.12646 48.2724L22.9728 1H24.1605L9.31417 48.2724H8.12646Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M1 48.2724L15.8463 1H16.4402L1.59385 48.2724H1Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
+<path d="M85.3271 55.5457H67.5116L87 12.7363L44.3513 68.2729H58.6038L43.1636 101L85.3271 55.5457Z" fill="#FDC717" stroke="#FDC717" stroke-width="1.18771" stroke-miterlimit="16"/>
+</svg>

assets/icon-simple.svg

@@ -0,0 +1,6 @@
+<svg width="160" height="93" viewBox="0 0 160 93" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M28.8511 91.66L57.6319 1.86368H64.5394L35.7585 91.66H28.8511Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+<path d="M15.0376 91.66L43.8185 1.86368H46.1209L17.3401 91.66H15.0376Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+<path d="M1.22217 91.66L30.003 1.86366H31.1543L2.3734 91.66H1.22217Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.15122"/>
+<path d="M71.4465 1.86483L42.666 91.6599H69.144L78.3538 58.2746H123.251L129.007 39.855H84.1099L89.866 22.5868H152.032L157.788 1.86483H71.4465Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+</svg>

benchmarks/fvd/README.md

@@ -0,0 +1,103 @@
+# 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
+```
+
+**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.
+```
+
+**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 cache/real  # Directory path (will save/load cache/real/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
+```
+
+## 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
+
+- I3D model auto-downloads from Hugging Face on first run
+- 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,35 @@
+"""
+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 .i3d_model import I3DFeatureExtractor
+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',
+    'I3DFeatureExtractor',
+    'load_video_auto',
+    'sample_clips_from_video',
+    'load_video_clips_streaming',
+    'ClipSamplingStrategy',
+]

benchmarks/fvd/cli.py

@@ -0,0 +1,185 @@
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from .fvd import compute_fvd_with_config, FVDConfig
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description='Compute Fréchet Video Distance (FVD)',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Standard FVD2048_16f protocol
+  python -m fastvideo.benchmarks.fvd.cli \\
+      --real-path data/real/ \\
+      --gen-path outputs/gen/ \\
+      --protocol fvd2048_16f
+
+  # Custom configuration
+  python -m fastvideo.benchmarks.fvd.cli \\
+      --real-path data/real/ \\
+      --gen-path outputs/gen/ \\
+      --num-videos 1024 \\
+      --num-frames 32 \\
+      --clip-strategy random \\
+      --frame-stride 2
+        """)
+
+    # Required arguments
+    parser.add_argument('--real-path',
+                        type=str,
+                        required=True,
+                        help='Path to real videos directory')
+    parser.add_argument('--gen-path',
+                        type=str,
+                        required=True,
+                        help='Path to generated videos directory')
+
+    # Reproducibility
+    parser.add_argument(
+        '--seed',
+        type=int,
+        default=None,
+        help='Random seed for reproducibility (np.random, random, torch)')
+
+    # Protocol presets
+    parser.add_argument('--protocol',
+                        type=str,
+                        default=None,
+                        choices=[
+                            'fvd2048_16f', 'fvd2048_128f',
+                            'fvd2048_128f_subsample8', 'quick_test'
+                        ],
+                        help='Use standard protocol (overrides other settings)')
+
+    # Video selection
+    parser.add_argument('--num-videos',
+                        type=int,
+                        default=2048,
+                        help='Number of videos to use (default: 2048)')
+
+    # Clip sampling
+    parser.add_argument('--num-frames',
+                        type=int,
+                        default=16,
+                        help='Number of frames per clip (default: 16)')
+    parser.add_argument('--num-clips',
+                        type=int,
+                        default=1,
+                        help='Number of clips per video (default: 1)')
+    parser.add_argument(
+        '--clip-strategy',
+        type=str,
+        default='beginning',
+        choices=['beginning', 'random', 'uniform', 'middle', 'sliding', 'all'],
+        help='Clip sampling strategy (default: beginning)')
+    parser.add_argument(
+        '--frame-stride',
+        type=int,
+        default=1,
+        help='Frame stride for FPS subsampling (default: 1, no subsampling)')
+    parser.add_argument('--temporal-stride',
+                        type=int,
+                        default=1,
+                        help='Temporal stride for sliding window (default: 1)')
+
+    # Data processing
+    parser.add_argument('--no-frame-dirs',
+                        action='store_true',
+                        help='Disable frame directory support')
+
+    # Computation
+    parser.add_argument('--batch-size',
+                        type=int,
+                        default=32,
+                        help='Batch size for feature extraction (default: 32)')
+    parser.add_argument('--device',
+                        type=str,
+                        default='cuda',
+                        choices=['cuda', 'cpu'],
+                        help='Device to use (default: cuda)')
+
+    # Caching
+    parser.add_argument('--cache-real-features',
+                        type=str,
+                        default=None,
+                        help='Path to cache real video features')
+    parser.add_argument('--i3d-model-path',
+                        type=str,
+                        default=None,
+                        help='Custom cache path for I3D model')
+
+    # Output
+    parser.add_argument('--output',
+                        type=str,
+                        default='fvd_results.json',
+                        help='Output JSON file (default: fvd_results.json)')
+    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,
+            'fvd2048_128f_subsample8': FVDConfig.fvd2048_128f_subsample8,
+            'quick_test': FVDConfig.quick_test,
+        }
+        config = protocol_map[args.protocol]()
+
+        # Override device and caching from args
+        config.device = args.device
+        config.cache_real_features = args.cache_real_features
+        config.i3d_model_path = args.i3d_model_path
+        config.batch_size = args.batch_size
+        config.seed = args.seed
+    else:
+        # Custom config from args
+        config = FVDConfig(num_videos=args.num_videos,
+                           num_frames_per_clip=args.num_frames,
+                           num_clips_per_video=args.num_clips,
+                           clip_strategy=args.clip_strategy,
+                           frame_stride=args.frame_stride,
+                           temporal_stride=args.temporal_stride,
+                           support_frame_dirs=not args.no_frame_dirs,
+                           batch_size=args.batch_size,
+                           device=args.device,
+                           cache_real_features=args.cache_real_features,
+                           i3d_model_path=args.i3d_model_path,
+                           seed=args.seed)
+
+    # Compute FVD
+    try:
+        results = compute_fvd_with_config(real_videos=args.real_path,
+                                          gen_videos=args.gen_path,
+                                          config=config,
+                                          verbose=not args.quiet)
+
+        # Save results
+        output_path = Path(args.output)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(output_path, 'w') as f:
+            json.dump(results, f, indent=2)
+
+        print(f"\nResults saved to {output_path}")
+        print(f"FVD: {results['fvd']:.2f}")
+        print(f"Protocol: {results['protocol']}")
+
+        return 0
+
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        import traceback
+        traceback.print_exc()
+        return 1
+
+
+if __name__ == '__main__':
+    sys.exit(main())

benchmarks/fvd/fvd.py

@@ -0,0 +1,447 @@
+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 .i3d_model import I3DFeatureExtractor
+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
+
+    # 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.
+        
+        most common FVD configuration used in papers
+        """
+        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 fvd2048_128f_subsample8(cls) -> 'FVDConfig':
+        """
+        Long video with FPS subsampling: 2048 videos, 128 frames (every 8th).
+        Used for very long videos - samples every 8th frame
+        """
+        return cls(num_videos=2048,
+                   num_frames_per_clip=16,
+                   frame_stride=8,
+                   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"""
+        return {
+            'num_videos': self.num_videos,
+            'num_frames_per_clip': self.num_frames_per_clip,
+            'num_clips_per_video': self.num_clips_per_video,
+            'clip_strategy': str(self.clip_strategy),
+            'frame_stride': self.frame_stride,
+            'temporal_stride': self.temporal_stride,
+            'batch_size': self.batch_size,
+            'device': self.device,
+            'seed': self.seed,
+            'use_streaming': self.use_streaming,
+        }
+
+    def __str__(self) -> str:
+        """Human-readable protocol name"""
+        desc = f"FVD{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: I3DFeatureExtractor,
+                               batch_size: int = 32,
+                               max_clips: int | None = None,
+                               verbose: bool = True) -> np.ndarray:
+    """
+    Extract features from a video clip generator using streaming.
+
+    Args:
+        video_generator: Iterator yielding clips [T, C, H, W]
+        extractor: I3D feature extractor
+        batch_size: Batch size for processing
+        max_clips: Maximum clips to process (for validation)
+        verbose: Show progress
+    
+    Returns:
+        features: [N, 400] numpy array
+    """
+    all_features = []
+    batch = []
+    clip_count = 0
+
+    if verbose:
+        print(f"Extracting features with batch_size={batch_size}...")
+
+    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_tensor,
+                                                  batch_size=batch_size,
+                                                  verbose=False)
+            all_features.append(features.cpu().numpy())
+
+            batch = []  # Clear batch
+
+            if verbose and clip_count % (batch_size * 10) == 0:
+                print(f"Processed {clip_count} clips...")
+
+        # Stop if we've reached max_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_tensor,
+                                              batch_size=len(batch),
+                                              verbose=False)
+        all_features.append(features.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: I3DFeatureExtractor,
+                             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:
+        cache_file = Path(cache_path) / f"{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("Recomputing features...")
+            elif len(features) > max_features:
+                features = features[:max_features]
+                return features
+            else:
+                return features
+
+    # Compute features
+    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:
+        # Already a tensor
+        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:
+        cache_dir = Path(cache_path)
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_file = cache_dir / f"{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(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:
+    """
+    Compute Fréchet Video Distance (FVD)
+    
+    For advanced control, use compute_fvd_with_config() instead.
+    
+    Args:
+        real_videos: Path to real videos or tensor [N, T, C, H, W]
+        gen_videos: Path to generated videos or tensor [N, T, C, H, W]
+        num_frames: Frames per video (default: 16)
+        batch_size: Batch size (default: 32)
+        device: 'cuda' or 'cpu' (default: 'cuda')
+        num_videos: Max videos (default: 2048)
+        cache_real_features: Cache path for real features
+        i3d_model_path: Custom I3D model cache path
+        seed: Random seed for reproducibility
+        verbose: Print progress
+    
+    Returns:
+        FVD score (float). Lower is better.
+    """
+    num_videos = num_videos if num_videos is not None else 2048
+
+    config = FVDConfig(
+        num_videos=num_videos,
+        num_frames_per_clip=num_frames,
+        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']
+
+
+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)
+            - 'config': Configuration dict
+    
+    Example:
+        >>> config = FVDConfig.fvd2048_16f()
+        >>> results = compute_fvd_with_config('data/real/', 'outputs/gen/', config)
+        >>> print(f"FVD: {results['fvd']:.2f}")
+        >>> print(f"Protocol: {results['protocol']}")  # "FVD2048_16f"
+    """
+    # 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("=" * 70)
+        print("\nConfiguration:")
+        for key, value in config.to_dict().items():
+            print(f"  {key}: {value}")
+        print()
+
+    # Initialize I3D
+    if verbose:
+        print(f"\nInitializing I3D model on {config.device}...")
+
+    extractor = I3DFeatureExtractor(device=config.device,
+                                    cache_dir=config.i3d_model_path)
+
+    # 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: {fvd:.4f}")
+        print(f"Protocol: {config}")
+        print(f"{'='*70}\n")
+
+    results = {
+        'fvd': fvd,
+        'protocol': str(config),
+        'config': config.to_dict(),
+    }
+
+    return results

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,34 @@
+import sys
+from pathlib import Path
+from benchmarks.fvd.fvd import FVDConfig, compute_fvd_with_config
+
+root_dir = Path(__file__).parent.parent.parent
+sys.path.insert(0, str(root_dir))
+
+
+def main() -> None:
+    # Get script directory
+    script_dir = Path(__file__).parent.resolve()
+
+    clip_strategy = 'beginning'  # Options: 'uniform', 'random', 'beginning', 'end', 'all'
+    cfg = FVDConfig(
+        num_videos=650,
+        num_frames_per_clip=16,
+        num_clips_per_video=1,
+        clip_strategy=clip_strategy,
+        frame_stride=1,
+        batch_size=32,
+        device='cuda',
+        seed=42,
+        cache_real_features=str(script_dir / f'fvd-cache/{clip_strategy}'),
+    )
+
+    real_dir = "benchmarks/data/real_videos"
+    gen_dir = "benchmarks/data/generated_videos"
+
+    results = compute_fvd_with_config(real_dir, gen_dir, cfg, verbose=True)
+    print(f"FVD = {results['fvd']:.2f}")
+
+
+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
+
+# 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

@@ -25,9 +25,9 @@ sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave
 sudo apt update
 sudo apt install clang-11
 ```
-(If you use CUDA12.4)
+(If you use CUDA12.8)
 ```bash
-export CUDA_HOME=/usr/local/cuda-12.4
+export CUDA_HOME=/usr/local/cuda-12.8
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 ```

csrc/attn/scratch.sh

@@ -1,4 +0,0 @@
-    off_hz = tl.program_id(2)
-    b = off_hz // H
-    h = off_hz % H
-    meta_base = ((b * H + h) * q_tiles + q_blk)

csrc/attn/sliding_tile_attn/MANIFEST.in

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

csrc/attn/sliding_tile_attn/README.md

@@ -0,0 +1,96 @@
+
+# Attention Kernel Used in FastVideo
+
+## Sliding Tile Attention (STA)
+We only support H100 for STA.
+
+### Installation
+```bash
+pip install st_attn
+``` 
+
+Install from source:
+
+```bash
+git submodule update --init --recursive
+python setup.py install
+```
+
+If you encounter error during installation, try below:
+Install C++20 for ThunderKittens:
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+(If you use CUDA12.8)
+```bash
+export CUDA_HOME=/usr/local/cuda-12.8
+export PATH=${CUDA_HOME}/bin:${PATH} 
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+###  Usage
+End-2-end inference with FastVideo:
+```bash
+bash scripts/inference/v1_inference_wan_STA.sh
+```
+
+If you want to use sliding tile attention in your custom model:
+```python
+from st_attn import sliding_tile_attention
+# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
+# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256
+# a tile is a cube of size (6, 8, 8)
+# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video.
+# text_length: int ranging from 0 to 256
+# If your attention contains text token (Hunyuan)
+out = sliding_tile_attention(q, k, v, window_size, text_length)
+# If your attention does not contain text token (StepVideo)
+out = sliding_tile_attention(q, k, v, window_size, 0, False)
+```
+
+
+### Test
+```bash
+python ../tests/test_sta.py # test STA
+python ../tests/test_vsa.py # test VSA
+```
+### Benchmark
+```bash
+python ../benchmarks/bench_sta.py
+```
+
+
+### How Does STA Work?
+We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 
+
+
+https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747
+
+
+## STA Configuration Logic
+Here is a diagram of how the window is configured and passed through the FastVideo pipeline:
+
+<div align="center">
+<img src="../../../docs/assets/images/STA_configuration.png" width="80%"/>
+</div>
+
+
+## Why is STA Fast?
+2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. 
+
+STA removes mixed blocks.
+
+
+<div align="center">
+<img src=../../../assets/sliding_tile_attn_map.png width="80%"/>
+</div>
+
+## Acknowledgement
+
+We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/attn/sliding_tile_attn/setup.py

@@ -1,7 +1,7 @@
 import os
 import subprocess
 
-from csrc.attn.config_sta import kernels, sources, target
+from config_sta import kernels, sources, target
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
@@ -9,7 +9,7 @@ target = target.lower()
 
 # Package metadata
 PACKAGE_NAME = "st_attn"
-VERSION = "0.0.4"
+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"

csrc/attn/video_sparse_attn/MANIFEST.in

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

csrc/attn/video_sparse_attn/README.md

@@ -0,0 +1,61 @@
+
+
+# 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/setup.py

@@ -9,10 +9,10 @@ target = target.lower()
 
 # Package metadata
 PACKAGE_NAME = "vsa"
-VERSION = "0.0.1"
+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"
+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/')))

csrc/attn/video_sparse_attn/vsa/block_sparse_attn_triton.py

@@ -250,8 +250,9 @@ def _attn_bwd_dq(dq, q, K, V,  #
     
     
     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
+        kv_idx = tl.load(kv_ptr + blk_idx//2).to(tl.int32)
+        block_size = tl.load(variable_block_sizes + kv_idx) -  (blk_idx % 2) * step_n
+        block_sparse_offset = (kv_idx*2 + blk_idx%2) * step_n * stride_tok
         kT = tl.load(kT_ptrs + block_sparse_offset)
         vT = tl.load(vT_ptrs + block_sparse_offset)
         qk = tl.dot(q, kT)

csrc/attn/vmoba_attn/README.md

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

@@ -0,0 +1,26 @@
+# 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/tests/test_vmoba_attn.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import pytest
+import random
+from csrc.attn.vmoba_attn.vmoba  import moba_attn_varlen
+
+def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda"):
+    """
+    Generates random data for testing the variable-length attention function.
+    """
+    torch.manual_seed(42)
+    random.seed(42)
+    torch.cuda.manual_seed_all(42)
+
+    # Generate sequence lengths for each item in the batch
+    if batch_size > 1:
+        # Ensure sequence lengths are reasonably distributed
+        avg_seqlen = total_seqlen // batch_size
+        seqlens = [random.randint(avg_seqlen // 2, avg_seqlen + avg_seqlen // 2) for _ in range(batch_size - 1)]
+        remaining_len = total_seqlen - sum(seqlens)
+        if remaining_len > 0:
+            seqlens.append(remaining_len)
+        else: # Adjust if sum exceeds total_seqlen
+            seqlens.append(avg_seqlen)
+            current_sum = sum(seqlens)
+            seqlens[-1] -= (current_sum - total_seqlen)
+        # Ensure all lengths are positive
+        seqlens = [max(1, s) for s in seqlens]
+        # Final adjustment to match total_seqlen
+        seqlens[-1] += total_seqlen - sum(seqlens)
+
+    else:
+        seqlens = [total_seqlen]
+
+    cu_seqlens = torch.tensor([0] + list(torch.cumsum(torch.tensor(seqlens), 0)), device=device, dtype=torch.int32)
+    max_seqlen = max(seqlens) if seqlens else 0
+
+    q = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False)
+    k = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False)
+    v = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False)
+
+    return q, k, v, cu_seqlens, max_seqlen
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("total_seqlen", [512, 1024])
+@pytest.mark.parametrize("num_heads", [8])
+@pytest.mark.parametrize("head_dim", [64])
+@pytest.mark.parametrize("moba_chunk_size", [64])
+@pytest.mark.parametrize("moba_topk", [2, 4])
+@pytest.mark.parametrize("select_mode", ["topk", "threshold"])
+@pytest.mark.parametrize("threshold_type", ["query_head", "head_global", "overall"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_moba_attn_varlen_forward(
+    batch_size, total_seqlen, num_heads, head_dim, moba_chunk_size, moba_topk, select_mode, threshold_type, dtype
+):
+    """
+    Tests the forward pass of moba_attn_varlen for basic correctness.
+    It checks output shape, dtype, and for the presence of NaNs/Infs.
+    """
+    if dtype == torch.float32:
+        pytest.skip("float32 is not supported in flash attention")
+
+    q, k, v, cu_seqlens, max_seqlen = generate_test_data(
+        batch_size, total_seqlen, num_heads, head_dim, dtype
+    )
+
+    # Ensure chunk size is not larger than the smallest sequence length
+    min_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).min().item()
+    if moba_chunk_size > min_seqlen:
+        pytest.skip("moba_chunk_size is larger than the minimum sequence length in the batch")
+
+    try:
+        output = moba_attn_varlen(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            moba_chunk_size=moba_chunk_size,
+            moba_topk=moba_topk,
+            select_mode=select_mode,
+            threshold_type=threshold_type,
+            simsum_threshold=0.5, # A reasonable default for threshold mode
+        )
+    except Exception as e:
+        pytest.fail(f"moba_attn_varlen forward pass failed with exception: {e}")
+
+    # 1. Check output shape
+    assert output.shape == q.shape, f"Expected output shape {q.shape}, but got {output.shape}"
+
+    # 2. Check output dtype
+    assert output.dtype == q.dtype, f"Expected output dtype {q.dtype}, but got {output.dtype}"
+
+    # 3. Check for NaNs or Infs in the output
+    assert torch.all(torch.isfinite(output)), "Output contains NaN or Inf values"

csrc/attn/vmoba_attn/vmoba/__init__.py

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

csrc/attn/vmoba_attn/vmoba/vmoba.py

@@ -0,0 +1,868 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapt from https://github.com/KwaiVGI/VMoBA/blob/main/src/vmoba.py
+
+import random
+import time
+import os
+import torch
+from typing import Tuple
+try:
+    from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
+    from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
+except ImportError:
+    def _unsupported(*args, **kwargs):
+        raise ImportError("flash-attn is not installed. Please install it, e.g., `pip install flash-attn`.")
+    _flash_attn_varlen_forward = _unsupported
+    _flash_attn_varlen_backward = _unsupported
+    flash_attn_varlen_func = _unsupported
+    
+from functools import lru_cache
+from einops import rearrange
+
+
+@lru_cache(maxsize=16)
+def calc_chunks(cu_seqlen, moba_chunk_size):
+    """
+    Calculate chunk boundaries.
+
+    For vision tasks we include all chunks (even the last one which might be shorter)
+    so that every chunk can be selected.
+    """
+    batch_sizes = cu_seqlen[1:] - cu_seqlen[:-1]
+    batch_num_chunk = (batch_sizes + (moba_chunk_size - 1)) // moba_chunk_size
+    cu_num_chunk = torch.ones(
+        batch_num_chunk.numel() + 1,
+        device=cu_seqlen.device,
+        dtype=batch_num_chunk.dtype,
+    )
+    cu_num_chunk[1:] = batch_num_chunk.cumsum(dim=0)
+    num_chunk = cu_num_chunk[-1]
+    chunk_sizes = torch.full(
+        (num_chunk + 1,), moba_chunk_size, dtype=torch.int32, device=cu_seqlen.device
+    )
+    chunk_sizes[0] = 0
+    batch_last_chunk_size = batch_sizes - (batch_num_chunk - 1) * moba_chunk_size
+    chunk_sizes[cu_num_chunk[1:]] = batch_last_chunk_size
+    cu_chunk = chunk_sizes.cumsum(dim=-1, dtype=torch.int32)
+    chunk_to_batch = torch.zeros(
+        (num_chunk,), dtype=torch.int32, device=cu_seqlen.device
+    )
+    chunk_to_batch[cu_num_chunk[1:-1]] = 1
+    chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32)
+
+    # Do not filter out any chunk
+    filtered_chunk_indices = torch.arange(
+        num_chunk, device=cu_seqlen.device, dtype=torch.int32
+    )
+    num_filtered_chunk = num_chunk
+
+    return cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch
+
+
+# --- Threshold Selection Helper Functions ---
+
+def _select_threshold_query_head(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects chunks for each <query, head> pair based on threshold.
+    Normalization and sorting happen along the chunk dimension (dim=0).
+    """
+    C, H, S = gate.shape
+    eps = 1e-6
+
+    # LSE‐style normalization per <head, query> (across chunks)
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
+
+    row_min = gate_min_val.amin(dim=0)        # (H, S)
+    row_max = gate_masked.amax(dim=0)         # (H, S)
+    denom = row_max - row_min
+    denom = torch.where(denom <= eps, torch.ones_like(denom), denom) # avoid divide‑by‑zero
+
+    gate_norm = (gate - row_min.unsqueeze(0)) / denom.unsqueeze(0)
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
+
+    # 1) pull out the self‐chunk’s normalized weight for each <head,seq>
+    self_norm = (gate_norm * gate_self_chunk_mask).sum(dim=0)  # (H, S)
+
+    # 2) compute how much more normalized weight we need beyond self
+    total_norm_sum = gate_norm.sum(dim=0) # (H, S)
+    remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps)  # (H, S)
+    remain_ratio = torch.clamp(remain_ratio, min=0.0) # if already ≥ thresh, no extra needed
+
+    # 3) zero out the self‐chunk in a copy, so we only sort “others”
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0
+
+    # 4) sort the other chunks by descending norm, per <head,seq>
+    sorted_norm, sorted_idx = torch.sort(others_norm, descending=True, dim=0) # (C, H, S)
+
+    # 5) cumulative‑sum the sorted norms per <head,seq>
+    cumsum_others = sorted_norm.cumsum(dim=0) # (C, H, S)
+
+    # 6) for each <head,seq>, find the smallest k where cumsum_ratio ≥ remain_ratio
+    ratio = cumsum_others / (total_norm_sum.unsqueeze(0) + eps) # (C, H, S)
+    cond = ratio >= remain_ratio.unsqueeze(0) # (C, H, S) boolean mask
+    any_cond = cond.any(dim=0) # (H, S)
+    # Find the index of the first True value along dim 0. If none, use C-1.
+    cutoff = torch.where(any_cond, cond.float().argmax(dim=0), torch.full_like(any_cond, fill_value=C - 1)) # (H, S)
+
+    # 7) build a mask in sorted order up to that cutoff
+    idx_range = torch.arange(C, device=gate.device).view(-1, 1, 1) # (C, 1, 1)
+    sorted_mask = idx_range <= cutoff.unsqueeze(0) # (C, H, S)
+
+    # 8) scatter it back to original chunk order
+    others_mask = torch.zeros_like(gate, dtype=torch.bool)
+    others_mask.scatter_(0, sorted_idx, sorted_mask)
+
+    # 9) finally, include every self‐chunk plus all selected others
+    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)
+
+    return final_gate_mask
+
+
+def _select_threshold_block(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects <query, head> pairs for each block based on threshold.
+    Normalization and sorting happen across the head and sequence dimensions (dim=1, 2).
+    """
+    C, H, S = gate.shape
+    HS = H * S
+    eps = 1e-6
+
+    # LSE‐style normalization per block (across heads and queries)
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
+
+    block_max = gate_masked.amax(dim=(1, 2), keepdim=True) # (C, 1, 1)
+    block_min = gate_min_val.amin(dim=(1, 2), keepdim=True) # (C, 1, 1)
+    block_denom = block_max - block_min
+    block_denom = torch.where(block_denom <= eps, torch.ones_like(block_denom), block_denom) # (C, 1, 1)
+
+    gate_norm = (gate - block_min) / block_denom # (C, H, S)
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
+
+    # 1) identify normalized weights of entries that *are* self-chunks (from query perspective)
+    self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S)
+    # Sum these weights *per block*
+    self_norm_sum_per_block = self_norm_entries.sum(dim=(1, 2)) # (C,)
+
+    # 2) compute how much more normalized weight each block needs beyond its self-chunk contributions
+    total_norm_sum_per_block = gate_norm.sum(dim=(1, 2)) # (C,)
+    remain_ratio = simsum_threshold - self_norm_sum_per_block / (total_norm_sum_per_block + eps) # (C,)
+    remain_ratio = torch.clamp(remain_ratio, min=0.0) # (C,)
+
+    # 3) zero out the self‐chunk entries in a copy, so we only sort “others”
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries
+
+    # 4) sort the other <head, seq> pairs by descending norm, per block
+    others_flat = others_norm.contiguous().view(C, HS) # (C, H*S)
+    sorted_others_flat, sorted_indices_flat = torch.sort(others_flat, dim=1, descending=True) # (C, H*S)
+
+    # 5) cumulative‑sum the sorted norms per block
+    cumsum_others_flat = sorted_others_flat.cumsum(dim=1) # (C, H*S)
+
+    # 6) for each block, find the smallest k where cumsum_ratio ≥ remain_ratio
+    ratio_flat = cumsum_others_flat / (total_norm_sum_per_block.unsqueeze(1) + eps) # (C, H*S)
+    cond_flat = ratio_flat >= remain_ratio.unsqueeze(1) # (C, H*S) boolean mask
+    any_cond = cond_flat.any(dim=1) # (C,)
+    # Find the index of the first True value along dim 1. If none, use HS-1.
+    cutoff_flat = torch.where(any_cond, cond_flat.float().argmax(dim=1), torch.full_like(any_cond, fill_value=HS - 1)) # (C,)
+
+    # 7) build a mask in sorted order up to that cutoff per block
+    idx_range_flat = torch.arange(HS, device=gate.device).unsqueeze(0) # (1, H*S)
+    sorted_mask_flat = idx_range_flat <= cutoff_flat.unsqueeze(1) # (C, H*S)
+
+    # 8) scatter it back to original <head, seq> order per block
+    others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C, H*S)
+    others_mask_flat.scatter_(1, sorted_indices_flat, sorted_mask_flat)
+    others_mask = others_mask_flat.view(C, H, S) # (C, H, S)
+
+    # 9) finally, include every self‐chunk entry plus all selected others
+    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)
+
+    return final_gate_mask
+
+
+def _select_threshold_overall(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects <chunk, query, head> triplets globally based on threshold.
+    Normalization and sorting happen across all valid entries.
+    """
+    C, H, S = gate.shape
+    CHS = C * H * S
+    eps = 1e-6
+
+    # LSE‐style normalization globally across all valid entries
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
+
+    overall_max = gate_masked.max() # scalar
+    overall_min = gate_min_val.min() # scalar
+    overall_denom = overall_max - overall_min
+    overall_denom = torch.where(overall_denom <= eps, torch.tensor(1.0, device=gate.device, dtype=gate.dtype), overall_denom)
+
+    gate_norm = (gate - overall_min) / overall_denom # (C, H, S)
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
+
+    # 1) identify normalized weights of entries that *are* self-chunks
+    self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S)
+    # Sum these weights globally
+    self_norm_sum_overall = self_norm_entries.sum() # scalar
+
+    # 2) compute how much more normalized weight is needed globally beyond self-chunk contributions
+    total_norm_sum_overall = gate_norm.sum() # scalar
+    remain_ratio = simsum_threshold - self_norm_sum_overall / (total_norm_sum_overall + eps) # scalar
+    remain_ratio = torch.clamp(remain_ratio, min=0.0) # scalar
+
+    # 3) zero out the self‐chunk entries in a copy, so we only sort “others”
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries
+
+    # 4) sort all other entries by descending norm, globally
+    others_flat = others_norm.flatten() # (C*H*S,)
+    valid_others_mask_flat = valid_gate_mask.flatten() & ~gate_self_chunk_mask.flatten() # Mask for valid, non-self entries
+    
+    # Only sort the valid 'other' entries
+    valid_others_indices = torch.where(valid_others_mask_flat)[0]
+    valid_others_values = others_flat[valid_others_indices]
+
+    sorted_others_values, sort_perm = torch.sort(valid_others_values, descending=True) # (N_valid_others,)
+    sorted_original_indices = valid_others_indices[sort_perm] # Original indices in C*H*S space, sorted by value
+
+    # 5) cumulative‑sum the sorted valid 'other' norms globally
+    cumsum_others_values = sorted_others_values.cumsum(dim=0) # (N_valid_others,)
+
+    # 6) find the smallest k where cumsum_ratio ≥ remain_ratio globally
+    ratio_values = cumsum_others_values / (total_norm_sum_overall + eps) # (N_valid_others,)
+    cond_values = ratio_values >= remain_ratio # (N_valid_others,) boolean mask
+    any_cond = cond_values.any() # scalar
+    
+    # Find the index of the first True value in the *sorted* list. If none, use all valid others.
+    cutoff_idx_in_sorted = torch.where(
+        any_cond, 
+        cond_values.float().argmax(dim=0), 
+        torch.tensor(len(sorted_others_values) - 1, device=gate.device, dtype=torch.long)
+    )
+
+    # 7) build a mask selecting the top-k others based on the cutoff
+    # Select the original indices corresponding to the top entries in the sorted list
+    selected_other_indices = sorted_original_indices[:cutoff_idx_in_sorted + 1]
+
+    # 8) create the mask in the original flat shape
+    others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C*H*S,)
+    if selected_other_indices.numel() > 0: # Check if any 'other' indices were selected
+        others_mask_flat[selected_other_indices] = True
+    others_mask = others_mask_flat.view(C, H, S) # (C, H, S)
+
+    # 9) finally, include every self‐chunk entry plus all selected others
+    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)
+
+    return final_gate_mask
+
+
+def _select_threshold_head_global(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects <chunk, query> globally for each head based on threshold.
+    """
+    C, H, S = gate.shape
+    eps = 1e-6
+
+    # 1) LSE‐style normalization per head (across chunks and sequence dims)
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)
+
+    max_per_head = gate_masked.amax(dim=(0, 2), keepdim=True)  # (1, H, 1)
+    min_per_head = gate_min_val.amin(dim=(0, 2), keepdim=True)  # (1, H, 1)
+    denom = max_per_head - min_per_head
+    denom = torch.where(denom <= eps, torch.ones_like(denom), denom)
+
+    gate_norm = (gate - min_per_head) / denom
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0)  # (C, H, S)
+
+    # 2) sum normalized self‐chunk contributions per head
+    self_norm_sum = (gate_norm * gate_self_chunk_mask).sum(dim=(0, 2))  # (H,)
+
+    # 3) total normalized sum per head
+    total_norm_sum = gate_norm.sum(dim=(0, 2))  # (H,)
+
+    # 4) how much more normalized weight needed per head
+    remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps)  # (H,)
+    remain_ratio = torch.clamp(remain_ratio, min=0.0)
+
+    # 5) zero out self‐chunk entries to focus on "others"
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0  # (C, H, S)
+
+    # 6) flatten chunk and sequence dims, per head
+    CS = C * S
+    others_flat = others_norm.permute(1, 0, 2).reshape(H, CS)  # (H, C*S)
+    valid_flat  = (valid_gate_mask & ~gate_self_chunk_mask) \
+                   .permute(1, 0, 2).reshape(H, CS)           # (H, C*S)
+
+    # 7) vectorized selection of “others” per head
+    masked_flat = torch.where(valid_flat, others_flat, torch.zeros_like(others_flat))
+    sorted_vals, sorted_idx = torch.sort(masked_flat, dim=1, descending=True)  # (H, C*S)
+
+    cumsum_vals = sorted_vals.cumsum(dim=1)                                    # (H, C*S)
+    ratio_vals  = cumsum_vals / (total_norm_sum.unsqueeze(1) + eps)            # (H, C*S)
+    cond        = ratio_vals >= remain_ratio.unsqueeze(1)                      # (H, C*S)
+
+    has_cutoff = cond.any(dim=1)                                               # (H,)
+    default    = torch.full((H,), CS - 1, device=gate.device, dtype=torch.long)
+    cutoff     = torch.where(has_cutoff, cond.float().argmax(dim=1), default)  # (H,)
+
+    idx_range   = torch.arange(CS, device=gate.device).unsqueeze(0)            # (1, C*S)
+    sorted_mask = idx_range <= cutoff.unsqueeze(1)                             # (H, C*S)
+
+    selected_flat = torch.zeros_like(valid_flat)                               # (H, C*S)
+    selected_flat.scatter_(1, sorted_idx, sorted_mask)                         # (H, C*S)
+
+    # 8) reshape selection mask back to (C, H, S)
+    others_mask = selected_flat.reshape(H, C, S).permute(1, 0, 2)               # (C, H, S)
+
+    # 9) include self‐chunks plus selected others, and obey valid mask
+    final_gate_mask = valid_gate_mask & (gate_self_chunk_mask | others_mask)
+
+    return final_gate_mask
+
+
+class MixedAttention(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        self_attn_cu_seqlen,
+        moba_q,
+        moba_kv,
+        moba_cu_seqlen_q,
+        moba_cu_seqlen_kv,
+        max_seqlen,
+        moba_chunk_size,
+        moba_q_sh_indices,
+    ):
+        ctx.max_seqlen = max_seqlen
+        ctx.moba_chunk_size = moba_chunk_size
+        ctx.softmax_scale = softmax_scale = q.shape[-1] ** (-0.5)
+
+        # Non-causal self-attention branch
+        # return out, softmax_lse, S_dmask, rng_state
+        self_attn_out_sh, self_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=self_attn_cu_seqlen,
+            cu_seqlens_k=self_attn_cu_seqlen,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+        )
+        # MOBA attention branch (non-causal)
+        moba_attn_out, moba_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
+            q=moba_q,
+            k=moba_kv[:, 0],
+            v=moba_kv[:, 1],
+            cu_seqlens_q=moba_cu_seqlen_q,
+            cu_seqlens_k=moba_cu_seqlen_kv,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=moba_chunk_size,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+        )
+
+        self_attn_lse_sh = self_attn_lse_hs.t().contiguous()
+        moba_attn_lse = moba_attn_lse_hs.t().contiguous()
+
+        output = torch.zeros((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+        output_2d = output.view(-1, q.shape[2])
+
+        max_lse_1d = self_attn_lse_sh.view(-1)
+        max_lse_1d = max_lse_1d.index_reduce(
+            0, moba_q_sh_indices, moba_attn_lse.view(-1), "amax"
+        )
+        self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh)
+        moba_attn_lse = (
+            moba_attn_lse.view(-1)
+            .sub(max_lse_1d.index_select(0, moba_q_sh_indices))
+            .reshape_as(moba_attn_lse)
+        )
+
+        mixed_attn_se_sh = self_attn_lse_sh.exp()
+        moba_attn_se = moba_attn_lse.exp()
+
+        mixed_attn_se_sh.view(-1).index_add_(
+            0, moba_q_sh_indices, moba_attn_se.view(-1)
+        )
+        mixed_attn_lse_sh = mixed_attn_se_sh.log()
+
+        # Combine self-attention output
+        factor = (self_attn_lse_sh - mixed_attn_lse_sh).exp()  # [S, H]
+        self_attn_out_sh = self_attn_out_sh * factor.unsqueeze(-1)
+        output_2d += self_attn_out_sh.reshape_as(output_2d)
+
+        # Combine MOBA attention output
+        mixed_attn_lse = (
+            mixed_attn_lse_sh.view(-1)
+            .index_select(0, moba_q_sh_indices)
+            .view_as(moba_attn_lse)
+        )
+        factor = (moba_attn_lse - mixed_attn_lse).exp()  # [S, H]
+        moba_attn_out = moba_attn_out * factor.unsqueeze(-1)
+        raw_attn_out = moba_attn_out.view(-1, moba_attn_out.shape[-1])
+        output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out)
+        output = output.to(q.dtype)
+        mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(mixed_attn_se_sh)
+        ctx.save_for_backward(
+            output,
+            mixed_attn_lse_sh,
+            q,
+            k,
+            v,
+            self_attn_cu_seqlen,
+            moba_q,
+            moba_kv,
+            moba_cu_seqlen_q,
+            moba_cu_seqlen_kv,
+            moba_q_sh_indices,
+        )
+
+        return output
+
+    @staticmethod
+    def backward(ctx, d_output):
+
+        max_seqlen = ctx.max_seqlen
+        moba_chunk_size = ctx.moba_chunk_size
+        softmax_scale = ctx.softmax_scale
+
+        (
+            output,
+            mixed_attn_vlse_sh,
+            q,
+            k,
+            v,
+            self_attn_cu_seqlen,
+            moba_q,
+            moba_kv,
+            moba_cu_seqlen_q,
+            moba_cu_seqlen_kv,
+            moba_q_sh_indices,
+        ) = ctx.saved_tensors
+
+        d_output = d_output.contiguous()
+
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        _ = _flash_attn_varlen_backward(
+            dout=d_output,
+            q=q,
+            k=k,
+            v=v,
+            out=output,
+            softmax_lse=mixed_attn_vlse_sh.t().contiguous(),
+            dq=dq,
+            dk=dk,
+            dv=dv,
+            cu_seqlens_q=self_attn_cu_seqlen,
+            cu_seqlens_k=self_attn_cu_seqlen,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=True,
+            window_size_left=-1,
+            window_size_right=-1
+        )
+
+        headdim = q.shape[-1]
+        d_moba_output = (
+            d_output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)
+        )
+        moba_output = (
+            output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)
+        )
+
+        mixed_attn_vlse = (
+            mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1)
+        )
+
+        dmq = torch.empty_like(moba_q)
+        dmkv = torch.empty_like(moba_kv)
+        _ = _flash_attn_varlen_backward(
+            dout=d_moba_output,
+            q=moba_q,
+            k=moba_kv[:, 0],
+            v=moba_kv[:, 1],
+            out=moba_output,
+            softmax_lse=mixed_attn_vlse,
+            dq=dmq,
+            dk=dmkv[:,0],
+            dv=dmkv[:,1],
+            cu_seqlens_q=moba_cu_seqlen_q,
+            cu_seqlens_k=moba_cu_seqlen_kv,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=moba_chunk_size,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=True,
+            window_size_left=-1,
+            window_size_right=-1
+        )
+
+        return dq, dk, dv, None, dmq, dmkv, None, None, None, None, None
+
+
+def moba_attn_varlen(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    max_seqlen: int,
+    moba_chunk_size: int,
+    moba_topk: int,
+    select_mode: str = 'threshold', # "topk" or "threshold"
+    simsum_threshold: float = 0.25,
+    threshold_type: str = 'query_head',
+) -> torch.Tensor:
+    """
+    Accelerated MOBA attention for vision tasks with proper LSE normalization.
+
+    This version:
+      - Splits KV into chunks.
+      - For each query head, selects the top-k relevant KV chunks (including the self chunk)
+        by amplifying the diagonal (self-chunk) logits.
+      - Aggregates the attention outputs from the selected chunks using a log-sum-exp
+        reduction so that attending to each query over the selected chunks is equivalent
+        to the original algorithm.
+    """
+    # Stack keys and values.
+    kv = torch.stack((k, v), dim=1)
+    seqlen, num_head, head_dim = q.shape
+
+    # Compute chunk boundaries.
+    cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(
+        cu_seqlens, moba_chunk_size
+    )
+
+    self_attn_cu_seqlen = cu_chunk
+
+    # Update top-k selection to include the self chunk.
+    moba_topk = min(moba_topk, num_filtered_chunk)
+
+    # --- Build filtered KV from chunks ---
+    chunk_starts = cu_chunk[filtered_chunk_indices]          # [num_filtered_chunk]
+    chunk_ends = cu_chunk[filtered_chunk_indices + 1]          # [num_filtered_chunk]
+    chunk_lengths = chunk_ends - chunk_starts                  # [num_filtered_chunk]
+    max_chunk_len = int(chunk_lengths.max().item())
+
+    range_tensor = torch.arange(max_chunk_len, device=kv.device, dtype=chunk_starts.dtype).unsqueeze(0)
+    indices = chunk_starts.unsqueeze(1) + range_tensor
+    indices = torch.clamp(indices, max=kv.shape[0] - 1)
+    valid_mask = range_tensor < chunk_lengths.unsqueeze(1)
+    gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, *kv.shape[1:])
+    gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(gathered)
+
+    # Compute key_gate_weight over valid tokens.
+    key_values = gathered[:, :, 0].float()  # [num_filtered_chunk, max_chunk_len, num_head, head_dim]
+    valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1)
+    key_sum = (key_values * valid_mask_exp).sum(dim=1)
+    divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1)
+    key_gate_weight = key_sum / divisor  # [num_filtered_chunk, num_head, head_dim]
+
+    # Compute gate logits between key_gate_weight and queries.
+    q_float = q.float()
+    # gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float)  # [num_filtered_chunk, num_head, seqlen]
+    gate = torch.bmm(key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2)
+
+    # Amplify the diagonal (self chunk) contributions.
+    gate_seq_idx = torch.arange(seqlen, device=q.device, dtype=torch.int32).unsqueeze(0).expand(num_filtered_chunk, seqlen)
+    chunk_start = cu_chunk[filtered_chunk_indices]  # [num_filtered_chunk]
+    chunk_end = cu_chunk[filtered_chunk_indices + 1]  # [num_filtered_chunk]
+    gate_self_chunk_mask = ((gate_seq_idx >= chunk_start.unsqueeze(1)) &
+                            (gate_seq_idx < chunk_end.unsqueeze(1))).unsqueeze(1).expand(-1, num_head, -1)
+    amplification_factor = 1e9  # Example factor; adjust as needed.
+    origin_gate = gate.clone()
+    gate = gate.clone()
+    if select_mode == "topk":
+        gate[gate_self_chunk_mask] += amplification_factor
+
+    # Exclude positions that are outside the valid batch boundaries.
+    batch_starts = cu_seqlens[chunk_to_batch[filtered_chunk_indices]]
+    batch_ends = cu_seqlens[chunk_to_batch[filtered_chunk_indices] + 1]
+    gate_batch_start_mask = gate_seq_idx < batch_starts.unsqueeze(1)
+    gate_batch_end_mask = gate_seq_idx >= batch_ends.unsqueeze(1)
+    gate_inf_mask = gate_batch_start_mask | gate_batch_end_mask
+    gate.masked_fill_(gate_inf_mask.unsqueeze(1), -float("inf"))
+
+    if select_mode == 'topk':
+        # We amplify self‐chunk in gate already, so self entries will rank highest.
+        valid_gate_mask = gate != -float("inf")
+        if threshold_type == 'query_head':
+            # === per‐<head,seq> top-k across chunks (original behavior) ===
+            # gate: (C, H, S)
+            _, gate_topk_idx = torch.topk(gate, k=moba_topk, dim=0, largest=True, sorted=False)
+            gate_idx_mask = torch.zeros_like(gate, dtype=torch.bool)
+            gate_idx_mask.scatter_(0, gate_topk_idx, True)
+            gate_mask = valid_gate_mask & gate_idx_mask
+        elif threshold_type == 'overall':
+            # === global top-k across all (chunk, head, seq) entries ===
+            C, H, S = gate.shape
+            flat_gate = gate.flatten()
+            flat_mask = valid_gate_mask.flatten()
+            flat_gate_masked = torch.where(flat_mask, flat_gate, -float("inf"))
+            # pick topk global entries
+            vals, idx = torch.topk(flat_gate_masked, k=moba_topk * H * S, largest=True, sorted=False)
+            others_mask_flat = torch.zeros_like(flat_mask, dtype=torch.bool)
+            others_mask_flat[idx] = True
+            gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view(gate.shape)
+        elif threshold_type == 'head_global':
+            # per-head top-k across all chunks and sequence positions
+            C, H, S = gate.shape
+            CS = C * S
+            flat_gate  = gate.permute(1, 0, 2).reshape(H, CS)
+            flat_valid = valid_gate_mask.permute(1, 0, 2).reshape(H, CS)
+            flat_gate_masked = torch.where(flat_valid, flat_gate, torch.full_like(flat_gate, -float('inf')))
+            # pick top-k indices per head
+            _, topk_idx = torch.topk(flat_gate_masked, k=moba_topk * S, dim=1, largest=True, sorted=False)
+            gate_idx_flat = torch.zeros_like(flat_valid, dtype=torch.bool)
+            gate_idx_flat.scatter_(1, topk_idx, True)
+            gate_mask = gate_idx_flat.reshape(H, C, S).permute(1, 0, 2)
+        else:
+            raise ValueError(
+                f"Invalid threshold_type for topk: {threshold_type}. "
+                "Choose 'query_head', 'block', or 'overall'."
+            )
+    elif select_mode == 'threshold':
+        # Delegate to the specific thresholding function
+        valid_gate_mask = gate != -float("inf")  # (num_chunk, num_head, seqlen)
+        if threshold_type == 'query_head':
+            gate_mask = _select_threshold_query_head(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        elif threshold_type == 'block':
+            gate_mask = _select_threshold_block(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        elif threshold_type == 'overall':
+            gate_mask = _select_threshold_overall(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        elif threshold_type == 'head_global':
+            gate_mask = _select_threshold_head_global(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        else:
+            raise ValueError(f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'.")
+    else:
+        raise ValueError(f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'.")
+    
+    # eliminate self_chunk in MoBA branch
+    gate_mask = gate_mask & ~gate_self_chunk_mask
+    # if gate_mask is all false, perform flash_attn instead
+    if gate_mask.sum() == 0:
+        return flash_attn_varlen_func(
+            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=False
+        )
+
+    # Determine which query positions are selected.
+    # nonzero_indices has shape [N, 3] where each row is [chunk_index, head_index, seq_index].
+    moba_q_indices = gate_mask.reshape(gate_mask.shape[0], -1).nonzero(as_tuple=True)[-1] # [(h s k)]
+    moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (moba_q_indices // seqlen)
+    moba_q = rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1)
+
+    # Build cumulative sequence lengths for the selected queries.
+    moba_seqlen_q = gate_mask.sum(dim=-1).flatten()
+    q_zero_mask = moba_seqlen_q == 0
+    valid_expert_mask = ~q_zero_mask
+    if q_zero_mask.sum() > 0:
+        moba_seqlen_q = moba_seqlen_q[valid_expert_mask]
+    moba_cu_seqlen_q = torch.cat(
+        (
+            torch.tensor([0], device=q.device, dtype=moba_seqlen_q.dtype),
+            moba_seqlen_q.cumsum(dim=0),
+        ),
+        dim=0,
+    ).to(torch.int32)
+
+    # Rearrange gathered KV for the MOBA branch.
+    experts_tensor = rearrange(gathered, "nc cl two h d -> (nc h) cl two d")
+    valid_expert_lengths = chunk_lengths.unsqueeze(1).expand(num_filtered_chunk, num_head).reshape(-1).to(torch.int32)
+    if q_zero_mask.sum() > 0:
+        experts_tensor = experts_tensor[valid_expert_mask]
+        valid_expert_lengths = valid_expert_lengths[valid_expert_mask]
+
+    seq_range = torch.arange(experts_tensor.shape[1], device=experts_tensor.device).unsqueeze(0)
+    mask = seq_range < valid_expert_lengths.unsqueeze(1)
+    moba_kv = experts_tensor[mask]  # Shape: ((nc h cl_valid) two d)
+    moba_kv = moba_kv.unsqueeze(2)   # Shape: ((nc h cl_valid) two 1 d)
+
+    moba_cu_seqlen_kv = torch.cat(
+        [torch.zeros(1, device=experts_tensor.device, dtype=torch.int32),
+         valid_expert_lengths.cumsum(dim=0)],
+        dim=0,
+    ).to(torch.int32)
+
+    assert (
+        moba_cu_seqlen_kv.shape == moba_cu_seqlen_q.shape
+    ), f"Mismatch between moba_cu_seqlen_kv.shape and moba_cu_seqlen_q.shape: {moba_cu_seqlen_kv.shape} vs {moba_cu_seqlen_q.shape}"
+
+    return MixedAttention.apply(
+        q,
+        k,
+        v,
+        self_attn_cu_seqlen,
+        moba_q,
+        moba_kv,
+        moba_cu_seqlen_q,
+        moba_cu_seqlen_kv,
+        max_seqlen,
+        moba_chunk_size,
+        moba_q_sh_indices,
+    )
+
+
+def process_moba_input(
+    x,
+    patch_resolution,
+    chunk_size,
+):
+    """
+    Process inputs for the attention function.
+
+    Args:
+        x (torch.Tensor): Input tensor with shape [batch_size, num_patches, num_heads, head_dim].
+        patch_resolution (tuple): Tuple containing the patch resolution (t, h, w).
+        chunk_size (int): Size of the chunk. (maybe tuple or int, according to chunk type)
+
+    Returns:
+        torch.Tensor: Processed input tensor.
+    """
+    if isinstance(chunk_size, float) or isinstance(chunk_size, int):
+        moba_chunk_size = int(chunk_size * patch_resolution[1] * patch_resolution[2])
+    else:
+        assert isinstance(chunk_size, (Tuple, list)), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}"
+        if len(chunk_size) == 2:
+            assert patch_resolution[1] % chunk_size[0] == 0 and patch_resolution[2] % chunk_size[1] == 0, f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}"
+            nch, ncw = patch_resolution[1] // chunk_size[0], patch_resolution[2] // chunk_size[1]
+            x = rearrange(x, "b (t nch ch ncw cw) n d -> b (nch ncw t ch cw) n d", t=patch_resolution[0], nch=nch, ncw=ncw, ch=chunk_size[0], cw=chunk_size[1])
+            moba_chunk_size = patch_resolution[0] * chunk_size[0] * chunk_size[1]
+        elif len(chunk_size) == 3:
+            assert patch_resolution[0] % chunk_size[0] == 0 and patch_resolution[1] % chunk_size[1] == 0 and patch_resolution[2] % chunk_size[2] == 0, f"patch_resolution {patch_resolution} should be divisible by 3d chunk_size {chunk_size}"
+            nct, nch, ncw = patch_resolution[0] // chunk_size[0], patch_resolution[1] // chunk_size[1], patch_resolution[2] // chunk_size[2]
+            x = rearrange(x, "b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d", nct=nct, nch=nch, ncw=ncw, ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2])
+            moba_chunk_size = chunk_size[0] * chunk_size[1] * chunk_size[2]
+        else:
+            raise ValueError(f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}")
+        
+    return x, moba_chunk_size
+
+
+def process_moba_output(
+    x,
+    patch_resolution,
+    chunk_size,
+):
+    if isinstance(chunk_size, float) or isinstance(chunk_size, int):
+        pass
+    elif len(chunk_size) == 2:
+        x = rearrange(x, "b (nch ncw t ch cw) n d -> b (t nch ch ncw cw) n d", nch=patch_resolution[1] // chunk_size[0], ncw=patch_resolution[2] // chunk_size[1], t=patch_resolution[0], ch=chunk_size[0], cw=chunk_size[1])
+    elif len(chunk_size) == 3:
+        x = rearrange(x, "b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d", nct=patch_resolution[0] // chunk_size[0], nch=patch_resolution[1] // chunk_size[1], ncw=patch_resolution[2] // chunk_size[2], ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2])
+
+    return x
+
+
+# TEST
+def generate_data(batch_size, seqlen, num_head, head_dim, dtype):
+    random.seed(0)
+    torch.manual_seed(0)
+    torch.cuda.manual_seed(0)
+    device = torch.cuda.current_device()
+
+    q = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
+    k = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
+    v = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
+    print(f"q.shape: {q.shape}, k.shape: {k.shape}, v.shape: {v.shape}")
+    cu_seqlens = torch.arange(0, q.shape[0] * q.shape[1] + 1, q.shape[1], dtype=torch.int32, device='cuda')
+    max_seqlen = q.shape[1]
+    q = rearrange(q, "b s ... -> (b s) ...")
+    k = rearrange(k, "b s ... -> (b s) ...")
+    v = rearrange(v, "b s ... -> (b s) ...")
+
+    return q, k, v, cu_seqlens, max_seqlen
+
+
+def test_attn_varlen_moba_speed(batch, head, seqlen, head_dim, moba_chunk_size, moba_topk, dtype=torch.bfloat16, select_mode='threshold', simsum_threshold=0.25, threshold_type='query_head'):
+    """Speed test comparing flash_attn vs moba_attention"""
+    # Get data
+    q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype)
+    print(f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}")
+    vo_grad = torch.randn_like(q)
+    
+    # Warmup
+    warmup_iters = 3
+    perf_test_iters = 10
+
+    # Warmup
+    for _ in range(warmup_iters):
+        o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False)
+        torch.autograd.backward(o, vo_grad)
+    
+    torch.cuda.synchronize()
+    start_flash = time.perf_counter()
+    for _ in range(perf_test_iters):
+        o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False)
+        torch.autograd.backward(o, vo_grad)
+        
+    torch.cuda.synchronize()
+    time_flash = (time.perf_counter() - start_flash) / perf_test_iters * 1000
+
+    # Warmup
+    for _ in range(warmup_iters):
+        om = moba_attn_varlen(q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type)
+        torch.autograd.backward(om, vo_grad)
+
+    torch.cuda.synchronize()
+    start_moba = time.perf_counter()
+    for _ in range(perf_test_iters):
+        om = moba_attn_varlen(q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type)
+        torch.autograd.backward(om, vo_grad)
+    
+    torch.cuda.synchronize()
+    time_moba = (time.perf_counter() - start_moba) / perf_test_iters * 1000
+    
+    print(f"Flash: {time_flash:.2f}ms, MoBA: {time_moba:.2f}ms")
+    print(f"Speedup:  {time_flash / time_moba:.2f}x")
+
+
+
+if __name__ == "__main__":
+    """
+    CUDA_VISIBLE_DEVICES=1 \
+    python -u csrc/attn/vmoba_attn/vmoba/vmoba.py
+    """
+    test_attn_varlen_moba_speed(batch=1, head=12, seqlen=32760, head_dim=128, moba_chunk_size=32760 // 3 // 6 // 4, moba_topk=3, select_mode='threshold', simsum_threshold=0.3, threshold_type='query_head')

docker/Dockerfile.python3.10

@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/sliding_tile_attn && \
     git submodule update --init --recursive && \
-    python setup_sta.py install
+    python setup.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
     git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install
 
-EXPOSE 22
+EXPOSE 22

docker/Dockerfile.python3.11

@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/sliding_tile_attn && \
     git submodule update --init --recursive && \
-    python setup_sta.py install
+    python setup.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
     git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install
 
-EXPOSE 22
+EXPOSE 22

docker/Dockerfile.python3.12

@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/sliding_tile_attn && \
     git submodule update --init --recursive && \
-    python setup_sta.py install
+    python setup.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
     git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install
 
-EXPOSE 22
+EXPOSE 22

docker/Dockerfile.python3.12.cuda12.9.1

@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/sliding_tile_attn && \
     git submodule update --init --recursive && \
-    python setup_sta.py install
+    python setup.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
     git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install
 
-EXPOSE 22
+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 docs/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,248 @@
+# 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.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,41 @@
+.vertical-table-header th.head:not(.stub) {
+    writing-mode: sideways-lr;
+    white-space: nowrap;
+    max-width: 0;
+    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;
+}

docs/assets/logos/icon_simple.svg

@@ -0,0 +1,6 @@
+<svg width="160" height="93" viewBox="0 0 160 93" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M28.8511 91.66L57.6319 1.86368H64.5394L35.7585 91.66H28.8511Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+<path d="M15.0376 91.66L43.8185 1.86368H46.1209L17.3401 91.66H15.0376Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+<path d="M1.22217 91.66L30.003 1.86366H31.1543L2.3734 91.66H1.22217Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.15122"/>
+<path d="M71.4465 1.86483L42.666 91.6599H69.144L78.3538 58.2746H123.251L129.007 39.855H84.1099L89.866 22.5868H152.032L157.788 1.86483H71.4465Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
+</svg>

docs/assets/logos/logo.svg

@@ -0,0 +1,18 @@
+<svg width="252" height="105" viewBox="0 0 252 105" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM100.768 13.1217L87.7028 29.4852H103.143L100.768 13.1217Z" fill="#356CFF"/>
+<path d="M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM109.081 90.697L116.802 65.8487C116.802 65.8487 120.959 65.8487 132.242 65.8487C143.525 65.8487 137.586 78.5759 135.211 84.0304C133.307 88.4021 127.491 90.697 122.74 90.697C117.989 90.697 109.081 90.697 109.081 90.697Z" fill="#356CFF"/>
+<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944H159.747C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273L125.188 48.273L124 37.97L147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852H131.836C120.142 29.4852 125.897 1.00043 141.337 1.00043L173.188 1.00056Z" fill="#356CFF"/>
+<path d="M179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056Z" fill="#356CFF"/>
+<path d="M161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM237.948 77.9692C239.984 70.6965 240.917 65.242 228.446 65.242C215.975 65.242 211.818 71.9087 210.037 77.9692C208.255 84.0298 208.255 91.3025 219.538 91.3025C230.821 91.3025 235.911 85.2419 237.948 77.9692Z" fill="#356CFF"/>
+<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944M173.188 1.00056C173.188 1.00056 156.777 1.00043 141.337 1.00043M173.188 1.00056L141.337 1.00043M141.337 20.3944C146.088 20.3944 150.839 20.3944 159.747 20.3944M141.337 20.3944H159.747M159.747 20.3944C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273M148.463 48.273C139.556 48.273 125.188 48.273 125.188 48.273M148.463 48.273L125.188 48.273M125.188 48.273L124 37.97M124 37.97C124 37.97 141.931 37.97 147.87 37.97M124 37.97L147.87 37.97M147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852M151.433 29.4852C146.682 29.4852 138.962 29.4852 131.836 29.4852M151.433 29.4852H131.836M131.836 29.4852C120.142 29.4852 125.897 1.00043 141.337 1.00043M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057ZM96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM87.7028 29.4852L100.768 13.1217L103.143 29.4852H87.7028ZM89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457ZM108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM116.802 65.8487L109.081 90.697C109.081 90.697 117.989 90.697 122.74 90.697C127.491 90.697 133.307 88.4021 135.211 84.0304C137.586 78.5759 143.525 65.8487 132.242 65.8487C120.959 65.8487 116.802 65.8487 116.802 65.8487ZM179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056ZM161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457ZM230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM228.446 65.242C240.917 65.242 239.984 70.6965 237.948 77.9692C235.911 85.2419 230.821 91.3025 219.538 91.3025C208.255 91.3025 208.255 84.0298 210.037 77.9692C211.818 71.9087 215.975 65.242 228.446 65.242Z" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M15.2524 55.5451L21.191 100.999L24.7541 100.999L18.8156 55.5451L15.2524 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M8.12646 55.5451L14.065 100.999L15.2527 100.999L9.31417 55.5451L8.12646 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M1 55.5451L6.93853 100.999L7.53239 100.999L1.59385 55.5451L1 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
+<path d="M15.2524 48.2724L30.0988 1H33.6619L18.8156 48.2724H15.2524Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M8.12646 48.2724L22.9728 1H24.1605L9.31417 48.2724H8.12646Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M1 48.2724L15.8463 1H16.4402L1.59385 48.2724H1Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
+<path d="M85.3271 55.5457H67.5116L87 12.7363L44.3513 68.2729H58.6038L43.1636 101L85.3271 55.5457Z" fill="#FDC717" stroke="#FDC717" stroke-width="1.18771" stroke-miterlimit="16"/>
+</svg>

docs/contributing/developer_env/docker.md

@@ -1,4 +1,4 @@
-(docker)=
+
 # 🐳 Using the FastVideo Docker Image
 
 If you prefer a containerized development environment or want to avoid managing dependencies manually, you can use our prebuilt Docker image:

docs/contributing/developer_env/index.md

@@ -3,11 +3,3 @@
 # 🧰 Developer Environment
 
 Accelerate your FastVideo development workflow by leveraging Docker images and cloud GPUs for efficient experimentation and reproducible environments.
-
-:::{toctree}
-:caption: Contents
-:maxdepth: 1
-
-docker
-runpod
-:::

docs/contributing/developer_env/runpod.md

@@ -1,4 +1,3 @@
-(runpod)=
 
 # 📦 Developing FastVideo on RunPod
 
@@ -10,7 +9,7 @@ Choose a GPU that supports CUDA 12.8
 
 Pick 1 or 2 L40S GPU(s)
 
-![RunPod CUDA selection](../../_static/images/runpod_cuda.png)
+![RunPod CUDA selection](../../assets/images/runpod_cuda.png)
 
 When creating your pod template, use this image:
 
@@ -24,11 +23,11 @@ Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.
 bash -c "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-server -y;mkdir -p ~/.ssh;cd $_;chmod 700 ~/.ssh;echo \"$PUBLIC_KEY\" >> authorized_keys;chmod 700 authorized_keys;service ssh start;sleep infinity"
 ```
 
-![RunPod template configuration](../../_static/images/runpod_template.png)
+![RunPod template configuration](../../assets/images/runpod_template.png)
 
 After deploying, the pod will take a few minutes to pull the image and start the SSH service.
 
-![RunPod ssh](../../_static/images/runpod_ssh.png)
+![RunPod ssh](../../assets/images/runpod_ssh.png)
 
 ## Working with the pod
 

docs/contributing/overview.md

@@ -1,4 +1,3 @@
-(developer-overview)=
 
 # 🛠️ Contributing to FastVideo
 
@@ -71,3 +70,7 @@ uv pip install ninja
 
 python setup.py install
 ```
+
+## Testing
+
+Please refer to the [Testing Guide](testing.md) for more information on how to add and run tests in FastVideo.

docs/contributing/profiling.md

@@ -0,0 +1,53 @@
+# Profiling FastVideo
+
+!!! warning
+    Profiling is only intended for FastVideo developers and maintainers to understand the proportion of time spent in different parts of the codebase. **FastVideo end-users should never turn on profiling** as it will significantly slow down the inference.
+
+## Profiling with PyTorch
+
+FastVideo exposes a process-wide torch profiler that you can enable via environment variables. Set `FASTVIDEO_TORCH_PROFILER_DIR` to an absolute directory path to start collecting traces, and specify the regions you want recorded with `FASTVIDEO_TORCH_PROFILE_REGIONS`:
+
+```bash
+FASTVIDEO_TORCH_PROFILER_DIR=/mnt/traces/fastvideo \
+FASTVIDEO_TORCH_PROFILE_REGIONS="profiler_region_model_loading,profiler_region_training_step"
+```
+
+All profiled regions must be registered in `fastvideo.profiler`; the current list includes:
+
+- `profiler_region_model_loading` — pipeline/module loading
+- `profiler_region_inference_pre_denoising`
+- `profiler_region_inference_denoising`
+- `profiler_region_inference_post_denoising`
+- `profiler_region_training_checkpoint_saving`
+- `profiler_region_training_dit`
+- `profiler_region_training_validation`
+- `profiler_region_training_epoch`
+- `profiler_region_training_step`
+- `profiler_region_training_backward`
+- `profiler_region_training_optimizer`
+- `profiler_region_distillation_teacher_forward`
+- `profiler_region_distillation_student_forward`
+- `profiler_region_distillation_loss`
+- `profiler_region_distillation_update`
+
+While profiling is enabled, FastVideo records additional annotations:
+
+- `fastvideo.region::<name>` spans are emitted when entering a region.
+- `fastvideo.profiler.enable_collection` / `fastvideo.profiler.disable_collection` events mark when torch profiler collection is toggled on or off.
+
+Only one profiler instance is created per process; subsequent pipelines reuse the same controller. If you set `FASTVIDEO_TORCH_PROFILE_REGIONS` incorrectly (e.g. misspelled name), FastVideo logs a warning and ignores that entry.
+
+Additional knobs:
+
+- `FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES`
+- `FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY`
+- `FASTVIDEO_TORCH_PROFILER_WITH_STACK`
+- `FASTVIDEO_TORCH_PROFILER_WITH_FLOPS`
+
+Traces can be visualized using <https://ui.perfetto.dev/>.
+
+### Best Practices
+
+- Keep the profiled step count small; traces can be large and slow down job shutdown while the profiler flushes data.
+- After profiling, clean up trace directories to avoid filling disks.
+- When adding new regions, register them in `fastvideo.profiler` and wrap the corresponding code block with `with self.profiler_controller.region("your_region"):` or the `@profile_region` decorator.

docs/contributing/testing.md

@@ -0,0 +1,129 @@
+# Testing in FastVideo
+
+This guide explains how to add and run tests in FastVideo. The testing suite is divided into several categories to ensure correctness across components, training workflows, and inference quality.
+
+## Test Types
+
+* **Unit Tests**: Located in `fastvideo/tests/dataset`, `fastvideo/tests/entrypoints`, and `fastvideo/tests/workflow`. These test individual functions and classes.
+* **Component Tests**: Located in `fastvideo/tests/encoders`, `fastvideo/tests/transformers`, and `fastvideo/tests/vaes`. These verify the loading and basic functionality of model components.
+* **SSIM Tests**: Located in `fastvideo/tests/ssim`. These are regression tests that compare generated videos against reference videos using the Structural Similarity Index Measure (SSIM) to detect quality degradation.
+* **Training Tests**: Located in `fastvideo/tests/training`. These validate training loops, loss calculations, and specific training techniques like LoRA, Distillation, and VSA.
+* **Inference Tests**: Located in `fastvideo/tests/inference`. These test specialized inference pipelines and optimizations (e.g., STA, V-MoBA).
+
+For now, we will focus on **SSIM Tests**.
+
+## SSIM Tests
+
+SSIM tests are located in `fastvideo/tests/ssim`. These tests generate videos using specific models and parameters, and compare them against reference videos to ensure that changes in the codebase do not degrade generation quality or alter the output unexpectedly.
+
+!!! note
+    If you are adding an SSIM test, this serves as a safeguard. Any future code changes that break or cause errors with the specific arguments and configurations you defined will trigger a failure. Therefore, it is important to include multiple settings and arguments that cover the core features of your new pipeline to ensure robust regression testing.
+
+### Directory Structure
+
+```
+fastvideo/tests/ssim/
+├── <GPU>_reference_videos/   # Reference videos organized by GPU type (e.g., L40S_reference_videos)
+│   ├── <Model_Name>/
+│   │   ├── <Backend>/        # e.g., FLASH_ATTN, TORCH_SDPA
+│   │   │   └── <Video_File>
+├── test_causal_similarity.py
+├── test_inference_similarity.py
+├── update_reference_videos.sh
+└── ...
+```
+
+### Adding a New SSIM Test
+
+To add a new SSIM test, follow these steps:
+
+1. **Create or Update a Test File**: You can add a new test function to an existing file (like `test_inference_similarity.py`) or create a new one if testing a distinct category of models.
+
+2. **Define Model Parameters**: Define the configuration for the model you want to test. This includes model path, dimensions, inference steps, and other generation parameters. **Note:** Consider using lower `num_inference_steps` or reduced resolution (e.g., 480p instead of 720p) to keep test execution time reasonable, provided it doesn't compromise the test's ability to detect regression.
+
+   ```python
+   MY_MODEL_PARAMS = {
+       "num_gpus": 1,
+       "model_path": "organization/model-name",
+       "height": 480,
+       "width": 832,
+       "num_frames": 45,
+       "num_inference_steps": 20,
+       # ... other parameters
+   }
+   ```
+
+3. **Implement the Test Function**:
+   * Use `pytest.mark.parametrize` to run the test with different prompts, backends, and models.
+   * Set the attention backend environment variable.
+   * Initialize the `VideoGenerator`.
+   * Generate the video.
+   * Compare the generated video with the reference video using `compute_video_ssim_torchvision`.
+
+   Example structure:
+
+   ```python
+   @pytest.mark.parametrize("prompt", TEST_PROMPTS)
+   @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
+   def test_my_model_similarity(prompt, ATTENTION_BACKEND):
+       # Setup output directories
+       # ...
+
+       # Initialize Generator
+       generator = VideoGenerator.from_pretrained(...)
+       generator.generate_video(prompt, ...)
+
+       # Compare with Reference
+       ssim_values = compute_video_ssim_torchvision(reference_path, generated_path, use_ms_ssim=True)
+        assert ssim_values[0] >= 0.98  # Threshold
+    ```
+
+4. **Reference Videos**:
+   * When running the test for the first time (or when updating the reference), the test will fail because the reference video is missing. The generated video will be saved in `fastvideo/tests/ssim/generated_videos`.
+   * Inspect the generated video to ensure it meets quality expectations.
+   * Move the generated video to the appropriate reference folder: `fastvideo/tests/ssim/<GPU>_reference_videos/<Model>/<Backend>/`.
+   * You can use the helper script `update_reference_videos.sh` to automate copying videos from `generated_videos` to `L40S_reference_videos`. Note: Check the script to ensure paths match your environment (it defaults to `L40S_reference_videos`).
+
+### Running Tests Locally
+
+To run the SSIM tests locally:
+
+```bash
+pytest fastvideo/tests/ssim/ -vs
+```
+
+Ensure you have the necessary GPUs available as defined in your test parameters.
+
+## Modal Workflow
+
+FastVideo uses [Modal](https://modal.com/) for running tests in a CI environment. The workflow scripts are located in `fastvideo/tests/modal/`.
+
+### `pr_test.py`
+
+The main entry point for CI tests is `fastvideo/tests/modal/pr_test.py`. This script defines Modal functions that execute the pytest suites on specific hardware (e.g., L40S, H100).
+
+### Updating Modal Configuration
+
+If you add a new test that requires:
+* **Different GPU Hardware**: You may need to change the `@app.function(gpu=...)` decorator.
+* **Longer Execution Time**: Increase the `timeout` parameter.
+* **New Environment Variables/Secrets**: Add them to `secrets=[...]` or the image environment. For example, if your model is gated on Hugging Face, ensure `HF_API_KEY` is passed.
+
+For SSIM tests, the `run_ssim_tests` function in `pr_test.py` currently runs:
+
+```python
+@app.function(gpu="L40S:2", image=image, timeout=2700, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})])
+def run_ssim_tests():
+    run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/ssim -vs")
+```
+
+If your new test file is inside `fastvideo/tests/ssim`, it will automatically be picked up by this command. However, ensure that the `gpu="L40S:2"` configuration is sufficient for your model. If your model requires more GPUs (e.g., 4 or 8), you might need to create a separate Modal function or update the existing one.
+
+### Workflow Scripts
+
+The shell script that triggers these tests in the CI pipeline is located at `.buildkite/scripts/pr_test.sh`. If you add a new test category (e.g., a new folder outside of `ssim`), you will need to:
+1. Add a new function in `fastvideo/tests/modal/pr_test.py`.
+2. Add a new case in `.buildkite/scripts/pr_test.sh` to handle the new test type.
+
+!!! note
+    If you are a maintainer, you'll need to finally manually update the workflow script in Buildkite. Otherwise, a maintainer will help you update.

docs/design/overview.md

@@ -29,7 +29,6 @@ FastVideo separates model components from execution logic with these principles:
 - **Custom Attention Backends**: Components can support and use different Attention implementations
 - **Pipeline Abstraction**: Consistent interface across diffusion models
 
-(design-fastvideo-args)=
 ## FastVideoArgs
 
 The `FastVideoArgs` class in `fastvideo/fastvideo_args.py` serves as the central configuration system for FastVideo. It contains all parameters needed to control model loading, inference configuration, performance optimization settings, and more.
@@ -61,7 +60,6 @@ with set_current_fastvideo_args(fastvideo_args):
     result = generate_video()
 ```
 
-(design-pipeline-system)=
 ## Pipeline System
 
 ### `ComposedPipelineBase`
@@ -108,7 +106,8 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
     return batch
 ```
 
-(design-forwardbatch)=
+![Pipeline execution and data flow](../assets/images/pipeline.png)
+
 ### ForwardBatch
 
 Defined in `fastvideo/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:
@@ -120,12 +119,10 @@ Defined in `fastvideo/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsul
 
 This structure facilitates clear state transitions between stages.
 
-(design-model-components)=
 ## Model Components
 
 The `fastvideo/models/` directory contains implementations of the core neural network models used in video diffusion:
 
-(design-transformer-models)=
 ### Transformer Models
 
 Transformer networks perform the actual denoising during diffusion:
@@ -152,7 +149,6 @@ def forward(
     return noise_pred  # Predicted noise residual
 ```
 
-(design-vae-variational-auto-encoder)=
 ### VAE (Variational Auto-Encoder)
 
 VAEs handle conversion between pixel space and latent space:
@@ -170,7 +166,6 @@ FastVideo's VAE implementations include:
 - Optional tiling for large frames
 - Distributed weight support
 
-(design-text-and-image-encoders)=
 ### Text and Image Encoders
 
 Encoders process conditioning inputs into embeddings:
@@ -188,7 +183,6 @@ FastVideo implements optimizations such as:
 - Caching for common prompts
 - Precision-tuned computation
 
-(design-schedulers)=
 ### Schedulers
 
 Schedulers manage the diffusion sampling process:
@@ -216,7 +210,10 @@ def step(
     return prev_sample
 ```
 
-(design-optimized-attention)=
+This diagram shows how models are discovered, validated, and loaded across entrypoints, executors, pipelines, and model loaders.
+
+![Model loading flow](../assets/images/load_models.png)
+
 ## Optimized Attention
 
 The `fastvideo/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:
@@ -240,17 +237,17 @@ self.attn = LocalAttention(
 # export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 ```
 
+![Attention backend selector design](../assets/images/attention_backend.png)
+
 ### Attention Patterns
 Supports various patterns with memory optimization techniques:
 - **Cross/Self/Temporal/Global-Local Attention**
 - Chunking, progressive computation, optimized masking
 
-(design-distributed-processing)=
 ## Distributed Processing
 
 The `fastvideo/distributed/` directory contains implementations for distributed model execution:
 
-(design-tensor-parallelism)=
 ### Tensor Parallelism
 
 Tensor parallelism splits model weights across devices:
@@ -307,7 +304,6 @@ Efficient communication primitives minimize distributed overhead:
 - **Tensor-Parallel AllReduce**: Combines partial results
 - **Distributed Synchronization**: Coordinates execution
 
-(design-forwardcontext)=
 ## Forward Context Management
 
 ### ForwardContext
@@ -330,7 +326,6 @@ with set_forward_context(current_timestep, attn_metadata, fastvideo_args):
     output = model(inputs)
 ```
 
-(design-executor-and-worker-abstractions)=
 ## Executor and Worker System
 
 The `fastvideo/worker/` directory contains the distributed execution framework:
@@ -357,7 +352,6 @@ Each GPU worker:
 
 This design allows FastVideo to efficiently utilize multiple GPUs while providing a simple, unified interface for model execution.
 
-(design-platforms)=
 ## Platforms
 
 The `fastvideo/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:
@@ -388,7 +382,6 @@ else:
 
 The platform system is designed to be extensible for future hardware targets.
 
-(design-logger)=
 ## Logger
 See [PR](https://github.com/hao-ai-lab/FastVideo/pull/356)
 

docs/distillation/data_preprocess.md

@@ -1,4 +1,3 @@
-(v0-data-preprocess)=
 
 # 🧱 Data Preprocess for Distillation
 

docs/examples/examples_index.md

@@ -0,0 +1,12 @@
+# 💡 Examples
+
+A collection of examples demonstrating usage of FastVideo.
+
+All documented examples are autogenerated using [generate_examples.py](https://github.com/hao-ai-lab/FastVideo/blob/main/docs/generate_examples.py) from examples found in the [examples](https://github.com/hao-ai-lab/FastVideo/tree/main/examples) directory.
+
+## Examples
+
+- [Examples Distillation Index](distillation/examples/examples_distillation_index.md)
+- [Examples Training Index](training/examples/examples_training_index.md)
+- [Examples Inference Index](inference/examples/examples_inference_index.md)
+

docs/generate_examples.py

@@ -6,10 +6,11 @@ import re
 from dataclasses import dataclass, field
 from pathlib import Path
 
-ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
-ROOT_DIR_RELATIVE = '../../../..'
+ROOT_DIR = Path(__file__).parent.parent.resolve()
+ROOT_DIR_RELATIVE = '../..'
 EXAMPLE_DIR = ROOT_DIR / "examples"
-EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
+EXAMPLE_DOC_DIR = ROOT_DIR / "docs/getting_started/examples"
+GITHUB_REPO = "hao-ai-lab/FastVideo"  # Update this to your repo
 
 
 def fix_case(text: str) -> str:
@@ -71,9 +72,16 @@ class Index:
 
     def generate(self) -> str:
         content = f"# {self.title}\n\n{self.description}\n\n"
-        content += ":::{toctree}\n"
-        content += f":caption: {self.caption}\n:maxdepth: {self.maxdepth}\n"
-        content += "\n".join(self.documents) + "\n:::\n"
+        if self.caption:
+            content += f"## {self.caption}\n\n"
+        # Generate a simple list of links for MkDocs
+        for doc in self.documents:
+            # Convert document path to proper link
+            doc_link = doc.replace("\\", "/")
+            # Get just the filename for the link text
+            doc_title = fix_case(Path(doc).stem.replace("_", " ").title())
+            content += f"- [{doc_title}]({doc_link}.md)\n"
+        content += "\n"
         return content
 
 
@@ -142,30 +150,66 @@ class Example:
         return fix_case(self.path.stem.replace("_", " ").title())
 
     def generate(self) -> str:
-        # Convert the path to a relative path from __file__
-        make_relative = lambda path: ROOT_DIR_RELATIVE / path.relative_to(
-            ROOT_DIR)
+        # Create GitHub link to source
+        github_path = str(self.path.relative_to(ROOT_DIR)).replace("\\", "/")
+        github_url = f"https://github.com/{GITHUB_REPO}/blob/main/{github_path}"
+        content = f"**Source:** [{github_path}]({github_url})\n\n"
 
-        content = f"Source <gh-file:{self.path.relative_to(ROOT_DIR)}>.\n\n"
-        include = "include" if self.main_file.suffix == ".md" else \
-            "literalinclude"
-        if include == "literalinclude":
+        # Add title for code files
+        if self.main_file.suffix != ".md":
             content += f"# {self.title}\n\n"
-        content += f":::{{{include}}} {make_relative(self.main_file)}\n"  # type: ignore[no-untyped-call]
-        if include == "literalinclude":
-            content += f":language: {self.main_file.suffix[1:]}\n"
-        content += ":::\n\n"
+
+        # Include main file content
+        if self.main_file.suffix == ".md":
+            # For markdown files, include the content directly
+            with open(self.main_file, encoding='utf-8') as f:
+                content += f.read() + "\n\n"
+        else:
+            # For code files, use code blocks
+            language = self.main_file.suffix[1:] if self.main_file.suffix else ""
+            with open(self.main_file, encoding='utf-8') as f:
+                file_content = f.read()
+            content += f"```{language}\n{file_content}\n```\n\n"
 
         if not self.other_files:
             return content
 
-        content += "## Example materials\n\n"
+        content += "## Additional Files\n\n"
+        # Define binary/non-text file extensions to skip
+        binary_extensions = {
+            '.mp4', '.avi', '.mov', '.mkv', '.gif', '.jpg', '.jpeg', '.png',
+            '.webp', '.bmp', '.pdf', '.zip', '.tar', '.gz', '.mp3', '.wav'
+        }
+
         for file in sorted(self.other_files):
-            include = "include" if file.suffix == ".md" else "literalinclude"
-            content += f":::{{admonition}} {file.relative_to(self.path)}\n"
-            content += ":class: dropdown\n\n"
-            content += f":::{{{include}}} {make_relative(file)}\n:::\n"  # type: ignore[no-untyped-call]
-            content += ":::\n\n"
+            # Skip binary files
+            if file.suffix.lower() in binary_extensions:
+                continue
+
+            file_rel_path = file.relative_to(self.path)
+            # Use collapsible admonition syntax for MkDocs
+            content += f"??? note \"{file_rel_path}\"\n\n"
+
+            try:
+                if file.suffix == ".md":
+                    # Include markdown content with indentation
+                    with open(file, encoding='utf-8') as f:
+                        for line in f:
+                            content += f"    {line}"
+                else:
+                    # Include code with proper formatting
+                    language = file.suffix[1:] if file.suffix else ""
+                    with open(file, encoding='utf-8') as f:
+                        file_content = f.read()
+                    # Indent the code block for the admonition
+                    content += f"    ```{language}\n"
+                    for line in file_content.split('\n'):
+                        content += f"    {line}\n"
+                    content += "    ```\n"
+                content += "\n"
+            except UnicodeDecodeError:
+                # Skip files that can't be decoded as UTF-8
+                continue
 
         return content
 
@@ -195,7 +239,7 @@ class NestedStructure:
 
 def create_category_indices() -> dict[str, Index]:
     """Create category indices with their respective configurations."""
-    main_index_dir = ROOT_DIR / "docs/source/examples"
+    main_index_dir = ROOT_DIR / "docs/examples"
     if not main_index_dir.exists():
         main_index_dir.mkdir(parents=True)
 
@@ -203,17 +247,16 @@ def create_category_indices() -> dict[str, Index]:
         "inference":
         Index(
             path=ROOT_DIR /
-            "docs/source/inference/examples/examples_inference_index.md",
+            "docs/inference/examples/examples_inference_index.md",
             title="🚀 Examples",
             description=
-            "Inference examples demonstrate how to use FastVideo inference. We recommend starting with <project:basic.md>.",
+            "Inference examples demonstrate how to use FastVideo inference. We recommend starting with [basic.md](basic.md).",
             caption="Examples",
             maxdepth=1,
         ),
         "training":
         Index(
-            path=ROOT_DIR /
-            "docs/source/training/examples/examples_training_index.md",
+            path=ROOT_DIR / "docs/training/examples/examples_training_index.md",
             title="🚀 Examples",
             description=
             "Training examples demonstrate how to use FastVideo training.",
@@ -223,7 +266,7 @@ def create_category_indices() -> dict[str, Index]:
         "distillation":
         Index(
             path=ROOT_DIR /
-            "docs/source/distillation/examples/examples_distillation_index.md",
+            "docs/distillation/examples/examples_distillation_index.md",
             title="🚀 Examples",
             description=
             "Distillation examples demonstrate how to use FastVideo distillation.",
@@ -246,9 +289,21 @@ def find_examples(category_indices: dict[str, Index],
     examples = []
     glob_patterns = ["*.py", "*.md", "*.sh"]
 
+    # Map category names to actual directory names
+    category_dir_mapping = {
+        "distillation": "distill",  # examples/distill/ -> distillation category
+    }
+
     # Find categorised examples
     for category in category_indices:
-        category_dir = EXAMPLE_DIR / category
+        # Use mapped directory name if available, otherwise use category name
+        dir_name = category_dir_mapping.get(category, category)
+        category_dir = EXAMPLE_DIR / dir_name
+
+        # Skip if directory doesn't exist
+        if not category_dir.exists():
+            continue
+
         globs = [category_dir.glob(pattern) for pattern in glob_patterns]
         for path in itertools.chain(*globs):
             examples.append(Example(path, category))
@@ -279,11 +334,18 @@ def create_nested_structures(
                                                 dict[str,
                                                      NestedStructure]]]] = {}
 
+    # Map category names to actual directory names
+    category_dir_mapping = {
+        "distillation": "distill",
+    }
+
     for example in examples:
         if example.category not in ["training", "distillation"]:
             continue
 
-        category_dir = EXAMPLE_DIR / example.category
+        # Use mapped directory name if available
+        dir_name = category_dir_mapping.get(example.category, example.category)
+        category_dir = EXAMPLE_DIR / dir_name
         relative_path = example.path.relative_to(category_dir)
         path_parts = relative_path.parts
 
@@ -415,7 +477,7 @@ def generate_nested_examples(nested_structures: dict[str, dict[str, dict[
             category_index.documents.append(method)
 
 
-def generate_examples(generate_main_index=False):
+def generate_examples(generate_main_index: bool = False) -> None:
     """
     Generate example documentation.
     
@@ -429,12 +491,14 @@ def generate_examples(generate_main_index=False):
     # Create the main examples index only if requested
     examples_index = None
     if generate_main_index:
-        main_index_dir = ROOT_DIR / "docs/source/examples"
+        main_index_dir = ROOT_DIR / "docs/examples"
         examples_index = Index(
             path=main_index_dir / "examples_index.md",
             title="💡 Examples",
             description=
-            "A collection of examples demonstrating usage of FastVideo.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.",
+            "A collection of examples demonstrating usage of FastVideo.\n\n"
+            f"All documented examples are autogenerated using [generate_examples.py](https://github.com/{GITHUB_REPO}/blob/main/docs/generate_examples.py) "
+            f"from examples found in the [examples](https://github.com/{GITHUB_REPO}/tree/main/examples) directory.",
             caption="Examples",
             maxdepth=2)
 
@@ -471,3 +535,19 @@ def generate_examples(generate_main_index=False):
     if generate_main_index and examples_index:
         with open(examples_index.path, "w+") as f:
             f.write(examples_index.generate())
+
+
+def on_pre_build_hook(config, **kwargs):
+    """
+    MkDocs hook to generate examples before building the documentation.
+    This function is called automatically by the mkdocs-simple-hooks plugin.
+    """
+    print("Generating example documentation...")
+    generate_examples(generate_main_index=True)
+    print("Example documentation generated successfully!")
+
+
+if __name__ == "__main__":
+    print("Generating example documentation...")
+    generate_examples(generate_main_index=True)
+    print("Example documentation generated successfully!")

docs/getting_started/installation.md

@@ -0,0 +1,41 @@
+
+# 🔧 Installation
+
+FastVideo supports the following hardware platforms:
+
+- [NVIDIA CUDA](installation/gpu.md)
+- [Apple silicon](installation/mps.md)
+
+## Quick Installation
+
+### Using pip
+
+```bash
+pip install fastvideo
+```
+
+### Using conda
+
+```bash
+conda install -c conda-forge fastvideo
+```
+
+### From source
+
+```bash
+git clone https://github.com/hao-ai-lab/FastVideo.git
+cd FastVideo
+pip install -e .
+```
+
+## Hardware Requirements
+
+- **NVIDIA GPUs**: CUDA 11.8+ with compute capability 7.0+
+- **Apple Silicon**: macOS 12.0+ with M1/M2/M3 chips
+- **CPU**: x86_64 architecture (for CPU-only inference)
+
+## Next Steps
+
+- [Quick Start Guide](quick_start.md) - Get started with your first video generation
+- [Configuration](../inference/configuration.md) - Learn about configuration options
+- [Examples](../inference/examples/) - Explore example scripts and notebooks

docs/getting_started/installation/gpu.md

@@ -30,17 +30,8 @@ conda create -n fastvideo python=3.12 -y
 conda activate fastvideo
 ```
 
-:::{note}
-[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
-:::
-
 #### uv
 
-:::{tip}
-We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
-Note that you can also use `uv` to install FastVideo in a Conda environment.
-:::
-
 Or you can create a new Python environment using [uv](https://docs.astral.sh/uv/), a very fast Python environment manager. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following command:
 
 ```console

docs/getting_started/installation/mps.md

@@ -31,17 +31,8 @@ conda create -n fastvideo python=3.12.4 -y
 conda activate fastvideo
 ```
 
-:::{note}
-[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
-:::
-
 #### uv
 
-:::{tip}
-We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
-Note that you can also use `uv` to install FastVideo in a Conda environment.
-:::
-
 Or you can create a new Python environment using [uv](https://docs.astral.sh/uv/), a very fast Python environment manager. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following command:
 
 ```console

docs/getting_started/quick_start.md

@@ -0,0 +1,55 @@
+# 🚀 Quick Start
+
+Get up and running with FastVideo in minutes!
+
+## Installation
+
+First, install FastVideo:
+
+```bash
+pip install fastvideo
+```
+
+## Basic Usage
+
+### Text-to-Video Generation
+
+```python
+from fastvideo import FastVideoPipeline
+
+# Initialize the pipeline
+pipe = FastVideoPipeline.from_pretrained("wan2.1-t2v-1.3B")
+
+# Generate a video
+prompt = "A cat playing with a ball of yarn"
+video = pipe(prompt, num_frames=16, height=512, width=512)
+
+# Save the video
+video.save("output.mp4")
+```
+
+### Image-to-Video Generation
+
+```python
+from fastvideo import FastVideoPipeline
+from PIL import Image
+
+# Load an image
+image = Image.open("input.jpg")
+
+# Initialize the pipeline
+pipe = FastVideoPipeline.from_pretrained("wan2.1-i2v-14B-480p")
+
+# Generate a video from the image
+video = pipe(image, num_frames=16, height=480, width=480)
+
+# Save the video
+video.save("output.mp4")
+```
+
+## Next Steps
+
+- [Installation Guide](installation.md) - Detailed installation instructions
+- [Configuration](../inference/configuration.md) - Learn about configuration options
+- [Examples](../inference/examples/) - Explore more examples
+- [Optimizations](../inference/optimizations.md) - Performance optimization tips

docs/getting_started/v1_api.md

@@ -0,0 +1,42 @@
+# V1 API
+
+FastVideo's V1 API provides a streamlined interface for video generation tasks with powerful customization options. This page documents the primary components of the API.
+
+## Video Generator
+
+This class will be the primary Python API for generating videos and images.
+
+::: fastvideo.entrypoints.video_generator.VideoGenerator
+    options:
+      show_root_heading: true
+      show_source: false
+      members:
+        - from_pretrained
+      heading_level: 3
+
+`VideoGenerator.from_pretrained()` should be the primary way of creating a new video generator.
+
+## Configuring FastVideo
+
+The following two classes `PipelineConfig` and `SamplingParam` are used to configure initialization and sampling parameters, respectively.
+
+### PipelineConfig
+
+::: fastvideo.configs.pipelines.base.PipelineConfig
+    options:
+      show_root_heading: true
+      show_source: false
+      members:
+        - from_pretrained
+        - dump_to_json
+      heading_level: 4
+
+### SamplingParam
+
+::: fastvideo.configs.sample.base.SamplingParam
+    options:
+      show_root_heading: true
+      show_source: false
+      members:
+        - from_pretrained
+      heading_level: 4

docs/index.md

@@ -1,131 +1,52 @@
 # Welcome to FastVideo
 
-:::{figure} ../../assets/logos/logo.svg
-:align: center
-:alt: FastVideo
-:class: no-scaled-link
-:width: 60%
-:::
+<div style="text-align: center;">
+  <img src="assets/logos/logo.svg" alt="FastVideo" style="width: 60%;" />
+</div>
 
-:::{raw} html
-<p style="text-align:center">
-<strong>FastVideo is a unified inference and post-training framework for accelerated video generation.
-</strong>
-</p>
+<div style="text-align: center;">
+<strong>FastVideo is a unified inference and post-training framework for accelerated video generation.</strong>
+</div>
 
-<p style="text-align:center">
+<div style="text-align: center;">
 <script async defer src="https://buttons.github.io/buttons.js"></script>
 <a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/" data-show-count="true" data-size="large" aria-label="Star">Star</a>
 <a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/subscription" data-icon="octicon-eye" data-size="large" aria-label="Watch">Watch</a>
 <a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/fork" data-icon="octicon-repo-forked" data-size="large" aria-label="Fork">Fork</a>
-</p>
-:::
+</div>
 
 FastVideo is an inference and post-training framework for diffusion models. It 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.
 
 <div style="text-align: center;">
-  <img src=_static/images/fastwan.png width="100%"/>
+  <img src="assets/images/fastwan.png" style="width: 100%;"/>
 </div>
 
 ## Key Features
 
 FastVideo has the following features:
+
 - State-of-the-art performance optimizations for inference
   - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
   - [TeaCache](https://arxiv.org/pdf/2411.19108)
   - [Sage Attention](https://arxiv.org/abs/2410.02367)
 - E2E post-training support
-  - Data preprocessing pipeline for video data.
+  - Data preprocessing pipeline for video data
   - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) for Wan2.1 and Wan2.2 using [Video Sparse Attention](https://arxiv.org/pdf/2505.13389) and [Distribution Matching Distillation](https://tianweiy.github.io/dmd2/)
   - 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.
 
 ## Documentation
 
-% How to start using FastVideo?
+Welcome to FastVideo! This documentation will help you get started with our unified inference and post-training framework for accelerated video generation.
 
-:::{toctree}
-:caption: Getting Started
-:maxdepth: 1
+Use the navigation menu on the left to explore different sections:
 
-getting_started/installation
-<!-- getting_started/v1_api -->
-:::
-
-:::{toctree}
-:caption: Inference
-:maxdepth: 1
-
-inference/inference_quick_start
-inference/examples/examples_inference_index
-inference/configuration
-inference/optimizations
-inference/comfyui
-inference/support_matrix
-inference/cli
-inference/add_pipeline
-:::
-
-:::{toctree}
-:caption: Training
-:maxdepth: 1
-
-training/examples/examples_training_index
-training/data_preprocess
-<!-- training/finetune -->
-:::
-
-:::{toctree}
-:caption: Distillation
-:maxdepth: 1
-
-distillation/examples/examples_distillation_index
-distillation/data_preprocess
-distillation/dmd
-:::
-
-% What is STA Kernel?
-
-:::{toctree}
-:caption: Sliding Tile Attention
-:maxdepth: 1
-
-sliding_tile_attention/installation
-sliding_tile_attention/demo
-:::
-
-% What is VSA Kernel?
-
-:::{toctree}
-:caption: Video Sparse Attention
-:maxdepth: 1
-
-video_sparse_attention/installation
-:::
-
-:::{toctree}
-:caption: Design
-:maxdepth: 1
-design/overview
-:::
-
-:::{toctree}
-:caption: Developer Guide
-:maxdepth: 2
-
-contributing/overview
-contributing/developer_env/index
-:::
-
-:::{toctree}
-:caption: API Reference
-:maxdepth: 2
-
-<!-- api/summary -->
-api/fastvideo/fastvideo
-:::
-
-## Indices and tables
-
-- {ref}`genindex`
-- {ref}`modindex`
+- **Getting Started**: Installation and quick start guides
+- **Inference**: Learn how to use FastVideo for video generation
+- **Training**: Data preprocessing and fine-tuning workflows
+- **Distillation**: Post-training optimization techniques
+- **Sliding Tile Attention**: Advanced attention mechanisms
+- **Video Sparse Attention**: Efficient attention for video models
+- **Design**: Framework architecture and design principles
+- **Developer Guide**: Contributing and development setup
+- **API Reference**: Complete API documentation

docs/inference/add_pipeline.md

@@ -1,4 +1,3 @@
-(add-pipeline)=
 
 # 🏗️ Adding a New Pipeline
 

docs/inference/configuration.md

@@ -1,4 +1,4 @@
-(inference-configuration)=
+
 # Configuration
 
 ## Multi-GPU Setup

docs/inference/optimizations.md

@@ -1,32 +1,37 @@
-(inference-optimizations)=
+
 # Optimizations
 
 This page describes the various options for speeding up generation times in FastVideo.
 
 ## Table of Contents
+
 - Optimized Attention Backends
+
   - [Flash Attention](#optimizations-flash)
   - [Sliding Tile Attention](#optimizations-sta)
   - [Sage Attention](#optimizations-sage)
+  - [Sage Attention 3](#optimizations-sage3)
 
 - Caching Techniques
   - [TeaCache](#optimizations-teacache)
 
-(optimizations-backends)=
 ## Attention Backends
 
 ### Available Backends
+
 - Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`
 - Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN`
 - Sliding Tile Attention: `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
 - Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN`
 - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
+- Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN_THREE`
 
 ### Configuring Backends
 
 There are two ways to configure the attention backend in FastVideo.
 
 #### 1. In Python
+
 In python, set the `FASTVIDEO_ATTENTION_BACKEND` environment variable before instantiating `VideoGenerator` like this:
 
 ```python
@@ -34,13 +39,13 @@ os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
 ```
 
 #### 2. In CLI
+
 You can also set the environment variable on the command line:
 
 ```bash
 FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
 ```
 
-(optimizations-flash)=
 ### Flash Attention
 
 **`FLASH_ATTN`**
@@ -57,16 +62,12 @@ And if using a Hopper+ GPU (ie H100), installing [Flash Attention 3](https://git
 git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention
 
 cd hopper
-pip install ninja 
+pip install ninja
 python setup.py install
 ```
 
-:::{note}
-FastVideo will automatically detect and use `FA3` if it is installed when using `FLASH_ATTN` backend.
-:::
-
-(optimizations-sta)=
 ### Sliding Tile Attention
+
 **`SLIDING_TILE_ATTN`**
 
 ```bash
@@ -75,8 +76,8 @@ pip install st_attn==0.0.4
 
 Please see [this page](#sta-installation) for more installation instructions.
 
-(optimizations-vsa)=
 ### Video Sparse Attention
+
 **`VIDEO_SPARSE_ATTN`**
 
 ```bash
@@ -86,20 +87,40 @@ python setup_vsa.py install
 
 Please see [this page](#vsa-installation) for more installation instructions.
 
-(optimizations-sage)=
 ### Sage Attention
+
 **`SAGE_ATTN`**
 
 To use [SageAttention](https://github.com/thu-ml/SageAttention) 2.1.1, please compile from source:
 
 ```bash
 git clone https://github.com/thu-ml/SageAttention.git
-cd sageattention 
+cd sageattention
 python setup.py install  # or pip install -e .
 ```
 
-(optimizations-teacache)=
+### Sage Attention 3
+
+**`SAGE_ATTN_THREE`**
+
+[SageAttention 3](https://huggingface.co/jt-zhang/SageAttention3) is an advanced attention mechanism that leverages FP4 quantization and Blackwell GPU Tensor Cores for significant performance improvements.
+
+#### Hardware Requirements
+
+- RTX5090
+
+#### Installation
+
+Note that Sage Attention 3 requires `python>=3.13`, `torch>=2.8.0`, `CUDA >=12.8`. If you are using `uv` and using `torch==2.8.0` make sure that `sentencepiece==0.2.1` in the pyproject.toml file.
+
+To use Sage Attention 3 in FastVideo, first get access to the SageAttention3 code, then move `sageattn/` and `setup.py` to the directory `fastvideo/attention/backends`, then install from using:
+
+```bash
+python setup.py install
+```
+
 ## Teacache
+
 TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo.
 
 ### What is TeaCache?

docs/inference/support_matrix.md

@@ -0,0 +1,66 @@
+# Compatibility Matrix
+
+The table below shows every supported model and optimizations supported for them.
+
+The symbols used have the following meanings:
+
+- ✅ = Full compatibility
+- ❌ = No compatibility
+- ⭕ = Does not apply to this model
+
+## Models x Optimization
+
+The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods and FastVideo will use the optimal default parameters when initializing and generating videos.
+
+<style>
+  /* Target tables in this section */
+  #models-x-optimization + p + table {
+    display: block;
+    overflow-x: auto;
+    width: 100%;
+    font-size: 0.85rem;
+  }
+  
+  #models-x-optimization + p + table td,
+  #models-x-optimization + p + table th {
+    text-align: center;
+    white-space: nowrap;
+    padding: 0.5em;
+  }
+  
+  /* First two columns can wrap */
+  #models-x-optimization + p + table td:nth-child(1),
+  #models-x-optimization + p + table td:nth-child(2) {
+    white-space: normal;
+    min-width: 120px;
+  }
+  
+  #models-x-optimization + p + table td:nth-child(2) code {
+    font-size: 0.75rem;
+  }
+</style>
+
+| Model Name | HuggingFace Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | VSA |
+|------------|---------------------|-------------|----------|-------------------|-----------|-----|
+| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ |
+| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ |
+| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ |
+| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ |
+| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ |
+| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ |
+| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ |
+| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ |
+| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅ | ✅* | ✅ | ⭕ |
+| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ |
+| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720P | ✅ | ✅ | ✅ | ⭕ |
+| StepVideo T2V | `FastVideo/stepvideo-t2v-diffusers` | 768px768px204f<br>544px992px204f<br>544px992px136f | ❌ | ❌ | ✅ | ⭕ |
+
+**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
+
+## Special requirements
+
+### StepVideo T2V
+- The self-attention in text-encoder (step_llm) only supports CUDA capabilities sm_80 sm_86 and sm_90
+
+### Sliding Tile Attention
+- Currently only Hopper GPUs (H100s) are supported.

docs/make.bat

@@ -1,35 +0,0 @@
-@ECHO OFF
-
-pushd %~dp0
-
-REM Command file for Sphinx documentation
-
-if "%SPHINXBUILD%" == "" (
-	set SPHINXBUILD=sphinx-build
-)
-set SOURCEDIR=source
-set BUILDDIR=build
-
-%SPHINXBUILD% >NUL 2>NUL
-if errorlevel 9009 (
-	echo.
-	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
-	echo.installed, then set the SPHINXBUILD environment variable to point
-	echo.to the full path of the 'sphinx-build' executable. Alternatively you
-	echo.may add the Sphinx directory to PATH.
-	echo.
-	echo.If you don't have Sphinx installed, grab it from
-	echo.https://www.sphinx-doc.org/
-	exit /b 1
-)
-
-if "%1" == "" goto help
-
-%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
-goto end
-
-:help
-%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
-
-:end
-popd

docs/requirements-docs.txt

@@ -1,15 +0,0 @@
-sphinx==7.4.7
-sphinx-argparse==0.5.2
-sphinx-autodoc2==0.5.0
-sphinx-book-theme==1.1.4
-sphinx-copybutton==0.5.2
-sphinx-design==0.6.1
-sphinx-togglebutton==0.3.2
-myst-parser==3.0.1
-msgspec
-commonmark # Required by sphinx-argparse when using :markdownhelp:
-
-# packages to install to build the documentation
-cachetools
-# -f https://download.pytorch.org/whl/cpu
-torch

docs/seed_parameter_behavior.md

@@ -1,51 +0,0 @@
-# Seed Parameter Behavior in vLLM
-
-## Overview
-
-The `seed` parameter in vLLM is used to control the random states for various random number generators. This parameter can affect the behavior of random operations in user code, especially when working with models in vLLM.
-
-## Default Behavior
-
-By default, the `seed` parameter is set to `None`. When the `seed` parameter is `None`, the global random states for `random`, `np.random`, and `torch.manual_seed` are not set. This means that the random operations will behave as expected, without any fixed random states.
-
-## Specifying a Seed
-
-If a specific seed value is provided, the global random states for `random`, `np.random`, and `torch.manual_seed` will be set accordingly. This can be useful for reproducibility, as it ensures that the random operations produce the same results across multiple runs.
-
-## Example Usage
-
-### Without Specifying a Seed
-
-```python
-import random
-from vllm import LLM
-
-# Initialize a vLLM model without specifying a seed
-model = LLM(model="Qwen/Qwen2.5-0.5B-Instruct")
-
-# Try generating random numbers
-print(random.randint(0, 100))  # Outputs different numbers across runs
-```
-
-### Specifying a Seed
-
-```python
-import random
-from vllm import LLM
-
-# Initialize a vLLM model with a specific seed
-model = LLM(model="Qwen/Qwen2.5-0.5B-Instruct", seed=42)
-
-# Try generating random numbers
-print(random.randint(0, 100))  # Outputs the same number across runs
-```
-
-## Important Notes
-
-- If the `seed` parameter is not specified, the behavior of global random states remains unaffected.
-- If a specific seed value is provided, the global random states for `random`, `np.random`, and `torch.manual_seed` will be set to that value.
-- This behavior can be useful for reproducibility but may lead to non-intuitive behavior if the user is not explicitly aware of it.
-
-## Conclusion
-
-Understanding the behavior of the `seed` parameter in vLLM is crucial for ensuring the expected behavior of random operations in your code. By default, the `seed` parameter is set to `None`, which means that the global random states are not affected. However, specifying a seed value can help achieve reproducibility in your experiments.

docs/sliding_tile_attention/demo.md

@@ -1,4 +1,3 @@
-(sta-demo)=
 
 # 🔍 Demo
 This is is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.

docs/sliding_tile_attention/installation.md

@@ -1,10 +1,9 @@
-(sta-installation)=
 
 # 🔧 Installation
 You can install the Sliding Tile Attention package using
 
 ```
-pip install st_attn==0.0.4
+pip install st_attn
 ``` 
 
 # Building from Source
@@ -12,7 +11,6 @@ We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have impleme
 First, install C++20 for ThunderKittens:
 
 ```bash
-cd csrc/sliding_tile_attention/
 sudo apt update
 sudo apt install gcc-11 g++-11
 
@@ -22,14 +20,20 @@ sudo apt update
 sudo apt install clang-11
 ```
 
-Install STA:
+Set up CUDA environment (if using CUDA 12.4):
 
 ```bash
 export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+Install STA:
+
+```bash
+cd csrc/attn/sliding_tile_attn/
 git submodule update --init --recursive
-python setup_sta.py install
+python setup.py install
 ```
 
 # 🧪 Test

docs/source/_static/custom.css

@@ -1,8 +0,0 @@
-.vertical-table-header th.head:not(.stub) {
-    writing-mode: sideways-lr;
-    white-space: nowrap;
-    max-width: 0;
-    p {
-       margin: 0;
-    }
-}