training refactor done

kv cache debugging
revert
2025-09-08 03:48:01 +00:00 · 2025-09-08 02:27:19 +00:00 · 2025-09-07 23:38:02 +00:00 · 2025-09-07 23:36:43 +00:00 · 2025-09-07 08:49:36 +00:00 · 2025-09-07 02:40:19 +00:00
237 changed files with 2366 additions and 13741 deletions
@@ -104,18 +104,6 @@ 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"
@@ -129,11 +117,11 @@ steps:
                  queue: "default"
            - path:
                - "fastvideo/**"
-                - "csrc/attn/video_sparse_attn/**"
-                - "csrc/attn/video_sparse_attn/tk/**"
-                - "csrc/attn/video_sparse_attn/setup.py"
-                - "csrc/attn/video_sparse_attn/config_vsa.py"
-                - "csrc/attn/video_sparse_attn/vsa.cpp"
+                - "csrc/attn/vsa/**"
+                - "csrc/attn/tk/**"
+                - "csrc/attn/setup_vsa.py"
+                - "csrc/attn/config_vsa.py"
+                - "csrc/attn/vsa.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -145,10 +133,10 @@ steps:
                  queue: "default"
            - path:
                - "fastvideo/**"
-                - "csrc/attn/sliding_tile_attn/**"
-                - "csrc/attn/sliding_tile_attn/setup.py"
-                - "csrc/attn/sliding_tile_attn/config_sta.py"
-                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
+                - "csrc/attn/st_attn/**"
+                - "csrc/attn/setup_sta.py"
+                - "csrc/attn/config_sta.py"
+                - "csrc/attn/st_attn.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -159,10 +147,10 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "csrc/attn/sliding_tile_attn/**"
-                - "csrc/attn/sliding_tile_attn/setup.py"
-                - "csrc/attn/sliding_tile_attn/config_sta.py"
-                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
+                - "csrc/attn/st_attn/**"
+                - "csrc/attn/setup_sta.py"
+                - "csrc/attn/config_sta.py"
+                - "csrc/attn/st_attn.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -173,12 +161,12 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "csrc/attn/video_sparse_attn/**"
-                - "csrc/attn/video_sparse_attn/tk/**"
+                - "csrc/attn/vsa/**"
+                - "csrc/attn/tk/**"
                - "csrc/attn/tests/test_vsa.py"
-                - "csrc/attn/video_sparse_attn/setup.py"
-                - "csrc/attn/video_sparse_attn/config_vsa.py"
-                - "csrc/attn/video_sparse_attn/vsa.cpp"
+                - "csrc/attn/setup_vsa.py"
+                - "csrc/attn/config_vsa.py"
+                - "csrc/attn/vsa.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -188,37 +176,3 @@ 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"
@@ -109,23 +109,6 @@ 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"
-        ;;
    *)
        log "Error: Unknown test type: $TEST_TYPE"
        exit 1
@@ -62,8 +62,8 @@ on:
        required: false
        default: false
        type: boolean
-      run_unit_test:
-        description: "Run unit-test"
+      run_nightly_test:
+        description: "Run nightly-test"
        required: false
        default: false
        type: boolean
@@ -93,7 +93,6 @@ 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
@@ -103,21 +102,18 @@ 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/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'
+              - 'csrc/attn/st_attn/**'
+              - 'csrc/attn/setup_sta.py'
+              - 'csrc/attn/config_sta.py'
+              - 'csrc/attn/st_attn.cpp'
            vsa-kernel-paths: &vsa-kernel-paths
-              - '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'
+              - 'csrc/attn/vsa/**'
+              - 'csrc/attn/tk/**'
+              - 'csrc/attn/setup_vsa.py'
+              - 'csrc/attn/config_vsa.py'
+              - 'csrc/attn/vsa.cpp'
            vsa-paths: &vsa-paths
              - 'fastvideo/**'
              - *common-paths
@@ -158,9 +154,6 @@ jobs:
            precision-test-VSA:
              - *common-paths
              - *vsa-kernel-paths
-            unit-test:
-              - 'fastvideo/**'
-              - *common-paths

  encoder-test:
    needs: change-filter
@@ -241,7 +234,7 @@ jobs:
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
+  
  training-test:
    needs: change-filter
    if: >-
@@ -339,42 +332,23 @@ jobs:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

-  unit-test:
-    needs: change-filter
+  nightly-test:
    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.unit-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_unit_test == 'true')
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
-      job_id: "unit-test"
-      gpu_type: "NVIDIA L40S"
-      gpu_count: 1
+      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: "uv pip install -e .[test] && pytest ./fastvideo/dataset/ -vs && pytest ./fastvideo/workflow/ -vs && pytest ./fastvideo/entrypoints/ -vs"
+      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 }}
-
-  # 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 }}
+      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}

  runpod-cleanup:
    # Add other jobs to this list as you create them
@@ -398,4 +372,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
@@ -5,7 +5,7 @@ on:
    branches:
      - main
    paths:
-      - "csrc/attn/sliding_tile_attn/setup.py"
+      - "csrc/attn/setup_sta.py"
  workflow_dispatch:

 jobs:
@@ -23,13 +23,13 @@ jobs:
      - name: Check if version changed
        id: check-version
        run: |
-          cd csrc/attn/sliding_tile_attn
+          cd csrc/attn
          # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_sta.py)
          echo "New version: $NEW_VERSION"
          
          # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup_sta.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/sliding_tile_attn  # Move into the correct folder
+          cd csrc/attn  # Move into the correct folder
          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py bdist_wheel --dist-dir=dist
+          python setup_sta.py bdist_wheel --dist-dir=dist

      - name: Rename wheel file
        run: |
-          cd csrc/attn/sliding_tile_attn
+          cd csrc/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/sliding_tile_attn/dist/*.whl
+          path: csrc/attn/dist/*.whl
          retention-days: 90

  publish_package:
@@ -239,11 +239,11 @@ jobs:
          pip install setuptools
          pip install ninja packaging wheel

-          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
+          cd csrc/attn  # Move into the correct folder
          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py sdist --dist-dir=dist
+          python setup_sta.py sdist --dist-dir=dist

      - name: Publish release distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
-          packages-dir: csrc/attn/sliding_tile_attn/dist/
+          packages-dir: csrc/attn/dist/
@@ -5,7 +5,7 @@ on:
    branches:
      - main
    paths:
-      - "csrc/attn/video_sparse_attn/setup.py"
+      - "csrc/attn/setup_vsa.py"
  workflow_dispatch:

 jobs:
@@ -23,13 +23,13 @@ jobs:
      - name: Check if version changed
        id: check-version
        run: |
-          cd csrc/attn/video_sparse_attn
+          cd csrc/attn
          # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_vsa.py)
          echo "New version: $NEW_VERSION"
          
          # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup_vsa.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/video_sparse_attn  # Move into the correct folder
+          cd csrc/attn  # Move into the correct folder
          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py bdist_wheel --dist-dir=dist
+          python setup_vsa.py bdist_wheel --dist-dir=dist

      - name: Rename wheel file
        run: |
-          cd csrc/attn/video_sparse_attn
+          cd csrc/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/video_sparse_attn/dist/*.whl
+          path: csrc/attn/dist/*.whl
          retention-days: 90

  publish_package:
@@ -247,11 +247,11 @@ jobs:
          pip install setuptools
          pip install ninja packaging wheel

-          cd csrc/attn/video_sparse_attn  # Move into the correct folder
+          cd csrc/attn  # Move into the correct folder
          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py sdist --dist-dir=dist
+          python setup_vsa.py sdist --dist-dir=dist

      - name: Publish release distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
-          packages-dir: csrc/attn/video_sparse_attn/dist/
+          packages-dir: csrc/attn/dist/
@@ -65,5 +65,4 @@ docs/source/distillation/examples/
 !comfyui/assets/**/*.png
 !comfyui/assets/**/*.gif

-dmd_t2v_output/
-preprocess_output_text/
+dmd_t2v_output/
@@ -1,7 +1,3 @@
-[submodule "csrc/attn/video_sparse_attn/tk"]
-	path = csrc/attn/video_sparse_attn/tk
-	url = https://github.com/HazyResearch/ThunderKittens.git
-
-[submodule "csrc/attn/sliding_tile_attn/tk"]
-	path = csrc/attn/sliding_tile_attn/tk
+[submodule "csrc/attn/tk"]
+	path = csrc/attn/tk
 	url = https://github.com/HazyResearch/ThunderKittens.git
@@ -12,6 +12,9 @@ exclude: |
    scripts/.*|
    fastvideo/data_preprocess/.*|
    fastvideo/dataset/.*|
+    fastvideo/distill/.*|
+    fastvideo/distill\.py|
+    fastvideo/distill_adv\.py|
    fastvideo/models/.*|
    fastvideo/sample/.*|
    fastvideo/train\.py|
@@ -41,10 +44,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:
@@ -7,7 +7,7 @@
 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-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/tMwknPLY" 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.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> |
 </p>

 <div align="center">
@@ -155,8 +155,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 Chen, Yongqi and Huang, Haofeng 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 Huang, Haofeng and Chen, Yongqi and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
  journal={arXiv preprint arXiv:2505.13389},
  year={2025}
 }
@@ -1,18 +0,0 @@
-<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>
@@ -1,6 +0,0 @@
-<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>
@@ -1,2 +1,2 @@
 recursive-include tk *
-include config_vsa.py
+include config.py
@@ -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.8)
+(If you use CUDA12.4)
 ```bash
-export CUDA_HOME=/usr/local/cuda-12.8
+export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 ```
@@ -0,0 +1,4 @@
+    off_hz = tl.program_id(2)
+    b = off_hz // H
+    h = off_hz % H
+    meta_base = ((b * H + h) * q_tiles + q_blk)
@@ -1,7 +1,7 @@
 import os
 import subprocess

-from config_sta import kernels, sources, target
+from csrc.attn.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.6"
+VERSION = "0.0.4"
 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"
@@ -9,10 +9,10 @@ target = target.lower()

 # Package metadata
 PACKAGE_NAME = "vsa"
-VERSION = "0.0.3"
+VERSION = "0.0.1"
 AUTHOR = "Hao AI Lab"
 DESCRIPTION = "Video Sparse Attention Kernel Used in FastVideo"
-URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn/video_sparse_attn"
+URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn"

 # Set environment variables
 tk_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), 'tk/')))
@@ -1,2 +0,0 @@
-recursive-include tk *
-include config_sta.py
@@ -1,96 +0,0 @@
-
-# Attention Kernel Used in FastVideo
-
-## Sliding Tile Attention (STA)
-We only support H100 for STA.
-
-### Installation
-```bash
-pip install st_attn
-``` 
-
-Install from source:
-
-```bash
-git submodule update --init --recursive
-python setup.py install
-```
-
-If you encounter error during installation, try below:
-Install C++20 for ThunderKittens:
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-sudo apt install clang-11
-```
-(If you use CUDA12.8)
-```bash
-export CUDA_HOME=/usr/local/cuda-12.8
-export PATH=${CUDA_HOME}/bin:${PATH} 
-export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-```
-
-###  Usage
-End-2-end inference with FastVideo:
-```bash
-bash scripts/inference/v1_inference_wan_STA.sh
-```
-
-If you want to use sliding tile attention in your custom model:
-```python
-from st_attn import sliding_tile_attention
-# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
-# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256
-# a tile is a cube of size (6, 8, 8)
-# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video.
-# text_length: int ranging from 0 to 256
-# If your attention contains text token (Hunyuan)
-out = sliding_tile_attention(q, k, v, window_size, text_length)
-# If your attention does not contain text token (StepVideo)
-out = sliding_tile_attention(q, k, v, window_size, 0, False)
-```
-
-
-### Test
-```bash
-python ../tests/test_sta.py # test STA
-python ../tests/test_vsa.py # test VSA
-```
-### Benchmark
-```bash
-python ../benchmarks/bench_sta.py
-```
-
-
-### How Does STA Work?
-We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 
-
-
-https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747
-
-
-## STA Configuration Logic
-Here is a diagram of how the window is configured and passed through the FastVideo pipeline:
-
-<div align="center">
-<img src="../../../docs/source/_static/images/STA_configuration.png" width="80%"/>
-</div>
-
-
-## Why is STA Fast?
-2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. 
-
-STA removes mixed blocks.
-
-
-<div align="center">
-<img src=../../../assets/sliding_tile_attn_map.png width="80%"/>
-</div>
-
-## Acknowledgement
-
-We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.
@@ -1,61 +0,0 @@
-
-
-# Attention Kernel Used in FastVideo
-
-## Video Sparse Attention (VSA)
-
-### Installation
-We support H100 (via TK) and any other GPU (via triton) for VSA.
-
-```bash
-pip install vsa
-``` 
-
-Install from source:
-
-```bash
-git submodule update --init --recursive
-python setup.py install
-```
-
-
-If you encounter error during installation, try below:
-Install C++20 for ThunderKittens:
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-sudo apt install clang-11
-```
-(If you use CUDA12.8)
-```bash
-export CUDA_HOME=/usr/local/cuda-12.8
-export PATH=${CUDA_HOME}/bin:${PATH} 
-export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-```
-
-### Verify if you have successfully installed
-
-```bash
-# test numerical
-python ../tests/test_vsa.py
-# (For H100) test speed
-python ../benchmarks/bench_vsa_hopper.py
-```
-
-bench_vsa_hopper.py should print something like this:
-
-```bash
-Using topk=76 kv blocks per q block (out of 768 total kv blocks)
-
-=== BLOCK SPARSE ATTENTION BENCHMARK ===
-Block Sparse Forward  - TFLOPS: 5622.26
-Block Sparse Backward - TFLOPS: 3865.68
-```
-
-## Acknowledgement
-
-We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.
@@ -1,32 +0,0 @@
-# Attention Kernel Used in FastVideo
-
-## VMoBA: Mixture-of-Block Attention for Video Diffusion Models (VMoBA)
-
-### Installation
-Please ensure that you have installed FlashAttention version **2.7.1 or higher**, as some interfaces have changed in recent releases.
-
-### Usage
-
-You can use `moba_attn_varlen` in the following ways:
-
-**Install from source:**
-```bash
-python setup.py install
-```
-
-**Import after installation:**
-```python
-from vmoba import moba_attn_varlen
-```
-
-**Or import directly from the project root:**
-```python
-from csrc.attn.vmoba_attn.vmoba import moba_attn_varlen
-```
-
-### Verify if you have successfully installed
-
-```bash
-python csrc/attn/vmoba_attn/vmoba/vmoba.py
-```
-
@@ -1,26 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from setuptools import find_packages, setup
-
-PACKAGE_NAME = "vmoba"
-VERSION = "0.0.0"
-AUTHOR = "JianzongWu"
-DESCRIPTION = "VMoBA: Mixture-of-Block Attention for Video Diffusion Models"
-URL = "https://github.com/KwaiVGI/VMoBA"
-
-setup(
-    name=PACKAGE_NAME,
-    version=VERSION,
-    author=AUTHOR,
-    description=DESCRIPTION,
-    url=URL,
-    packages=find_packages(),
-    classifiers=[
-        "Programming Language :: Python :: 3",
-        "License :: OSI Approved :: Apache Software License",
-    ],
-    python_requires='>=3.12',
-    install_requires=[
-        "flash-attn >= 2.7.1",
-    ]
-)
@@ -1,97 +0,0 @@
-# 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"
@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from .vmoba import moba_attn_varlen, process_moba_input, process_moba_output
@@ -1,868 +0,0 @@
-# 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')
@@ -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/sliding_tile_attn && \
+    cd csrc/attn && \
    git submodule update --init --recursive && \
-    python setup.py install
+    python setup_sta.py install

 # Install VSA 
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
+    cd csrc/attn && \
    git submodule update --init --recursive && \
-    python setup.py install
+    python setup_vsa.py install

-EXPOSE 22
+EXPOSE 22
@@ -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/sliding_tile_attn && \
+    cd csrc/attn && \
    git submodule update --init --recursive && \
-    python setup.py install
+    python setup_sta.py install

 # Install VSA 
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
+    cd csrc/attn && \
    git submodule update --init --recursive && \
-    python setup.py install
+    python setup_vsa.py install

-EXPOSE 22
+EXPOSE 22
@@ -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/sliding_tile_attn && \
+    cd csrc/attn && \
    git submodule update --init --recursive && \
-    python setup.py install
+    python setup_sta.py install

 # Install VSA 
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
+    cd csrc/attn && \
    git submodule update --init --recursive && \
-    python setup.py install
+    python setup_vsa.py install

-EXPOSE 22
+EXPOSE 22
@@ -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/sliding_tile_attn && \
+    cd csrc/attn && \
    git submodule update --init --recursive && \
-    python setup.py install
+    python setup_sta.py install

 # Install VSA 
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
+    cd csrc/attn && \
    git submodule update --init --recursive && \
-    python setup.py install
+    python setup_vsa.py install

-EXPOSE 22
+EXPOSE 22
@@ -1,53 +0,0 @@
-# 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.
@@ -115,7 +115,6 @@ design/overview

 contributing/overview
 contributing/developer_env/index
-contributing/profiling
 :::

 :::{toctree}
@@ -1,40 +1,32 @@
 (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
@@ -42,7 +34,6 @@ os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
 ```

 #### 2. In CLI
-
 You can also set the environment variable on the command line:

 ```bash
@@ -50,7 +41,6 @@ FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
 ```

 (optimizations-flash)=
-
 ### Flash Attention

 **`FLASH_ATTN`**
@@ -67,7 +57,7 @@ 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
 ```

@@ -76,9 +66,7 @@ FastVideo will automatically detect and use `FA3` if it is installed when using
 :::

 (optimizations-sta)=
-
 ### Sliding Tile Attention
-
 **`SLIDING_TILE_ATTN`**

 ```bash
@@ -88,9 +76,7 @@ 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
@@ -101,45 +87,19 @@ 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-sage3)=
-
-### 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
-```
-
 (optimizations-teacache)=
-
 ## 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?
@@ -4,7 +4,7 @@
 You can install the Sliding Tile Attention package using

 ```
-pip install st_attn
+pip install st_attn==0.0.4
 ``` 

 # Building from Source
@@ -12,6 +12,7 @@ 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

@@ -21,20 +22,14 @@ sudo apt update
 sudo apt install clang-11
 ```

-Set up CUDA environment (if using CUDA 12.4):
+Install STA:

 ```bash
 export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-```
-
-Install STA:
-
-```bash
-cd csrc/attn/sliding_tile_attn/
 git submodule update --init --recursive
-python setup.py install
+python setup_sta.py install
 ```

 # 🧪 Test
@@ -4,7 +4,8 @@
 You can install the Video Sparse Attention package using

 ```bash
-pip install vsa
+git submodule update --init --recursive
+python setup_vsa.py install
 ```

 # Building from Source
@@ -33,9 +34,9 @@ export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 Install VSA:

 ```bash
-cd csrc/attn/video_sparse_attn/
+cd csrc/attn/
 git submodule update --init --recursive
-python setup.py install
+python setup_vsa.py install
 ```

 # 🧪 Test
@@ -1,9 +0,0 @@
-# VidProm Dataset
-
-From [Self-Forcing](https://github.com/gdhe17/Self-Forcing) repository.
-
-## Download the dataset
-
-```bash
-./download_dataset.sh
-```
@@ -1,3 +0,0 @@
-#! /bin/bash
-
-huggingface-cli download gdhe17/Self-Forcing vidprom_filtered_extended.txt --local-dir prompts
@@ -10,18 +10,31 @@
 #SBATCH --output=dmd_t2v_output/t2v_%j.out
 #SBATCH --error=dmd_t2v_output/t2v_%j.err
 #SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv

 # Basic Info
+export WANDB_MODE="online"
 export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
 # different cache dir for different processes
 export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29503
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
 export TOKENIZERS_PARALLELISM=false
-export WANDB_API_KEY=your_wandb_api_key
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"

 # Configs
 NUM_GPUS=1
@@ -31,68 +44,76 @@ GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
 REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
 FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model

-DATA_DIR=your_data_dir
-VALIDATION_DATASET_FILE=your_validation_data_dir
+DATA_DIR="data/crush-smol-single_processed_t2v/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

+# Training arguments
 training_args=(
-  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  
-  --output_dir your_output_dir
-  --max_train_steps 4000
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
+  --output_dir "checkpoints/SFwan_t2v_finetune"
+  --max_train_steps 500
  --train_batch_size 1
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
  --num_latent_t 21
  --num_height 480
  --num_width 832
+  --num_frames 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
  --enable_gradient_checkpointing_type "full"
  --log_visualization
  --simulate_generator_forward
-  --num_frames 81
  --num_frame_per_block 3  # Frame generation block size for self-forcing
  --enable_gradient_masking
  --gradient_mask_last_n_frames 21
 )

+# Parallel arguments
 parallel_args=(
-  --num_gpus $NUM_GPUS # 64
+  --num_gpus 1 # 64
  --sp_size 1
  --tp_size 1
  --hsdp_replicate_dim 1 # 64
-  --hsdp_shard_dim $NUM_GPUS
+  --hsdp_shard_dim 1
 )

+# Model arguments
 model_args=(
  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --generator_model_path $GENERATOR_MODEL_PATH
  --real_score_model_path $REAL_SCORE_MODEL_PATH
  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

+# Dataset arguments
 dataset_args=(
  --data_path "$DATA_DIR"
  --dataloader_num_workers 4
 )

+# Validation arguments
 validation_args=(
  --log_validation
  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
+  --validation_steps 10
  --validation_sampling_steps "4"
  --validation_guidance_scale "6.0" # not used for dmd inference
 )

+# Optimizer arguments
 optimizer_args=(
  --learning_rate 1e-5
  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 500
-  --weight_only_checkpointing_steps 500
+  --training_state_checkpointing_steps 50
+  --weight_only_checkpointing_steps 50
  --weight_decay 0.01
  --betas '0.0,0.999'
  --max_grad_norm 1.0
 )

+# Miscellaneous arguments
 miscellaneous_args=(
  --inference_mode False
  --checkpoints_total_limit 3
@@ -103,9 +124,10 @@ miscellaneous_args=(
  --use_ema True
  --ema_decay 0.99
  --ema_start_step 100
-  --init_weights_from_safetensors your_ode_init_weights_path
+  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
 )

+# Self-forcing DMD arguments
 dmd_args=(
  --dmd_denoising_steps '1000,750,500,250'
  --min_timestep_ratio 0.02
@@ -114,20 +136,23 @@ dmd_args=(
  --real_score_guidance_scale 3.0
  --fake_score_learning_rate 8e-6
  --fake_score_betas '0.0,0.999'
-  --warp_denoising_step
 )

+# Self-forcing specific arguments
 self_forcing_args=(
  --independent_first_frame False  # Whether to treat first frame independently
-  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --same_step_across_blocks False  # Whether to use same denoising step across all blocks
  --last_step_only False  # Whether to only use the last denoising step
  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  --validate_cache_structure False  # Set to True for debugging KV cache issues
 )

-torchrun \
--nnodes 1 \
--master_port $MASTER_PORT \
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
 --nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
@@ -137,4 +162,4 @@ torchrun \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
    "${dmd_args[@]}" \
-    "${self_forcing_args[@]}"
+    "${self_forcing_args[@]}"
@@ -1,3 +1,3 @@
 #!/bin/bash

-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -0,0 +1,43 @@
+#!/bin/bash
+
+# Download the full dataset first
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
+
+# Create a single-example dataset for debugging
+SINGLE_EXAMPLE_DIR="data/crush-smol-single"
+mkdir -p "$SINGLE_EXAMPLE_DIR/videos"
+
+# Copy the specific video that matches the validation.json style (macaron crushing)
+cp "data/crush-smol/videos/7P02AihYkCU-Scene-005.mp4" "$SINGLE_EXAMPLE_DIR/videos/"
+
+# Create a single-line videos.txt
+echo "videos/7P02AihYkCU-Scene-005.mp4" > "$SINGLE_EXAMPLE_DIR/videos.txt"
+
+# Create a single-line prompt.txt with the macaron crushing prompt
+echo "PIKA_CRUSH A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out." > "$SINGLE_EXAMPLE_DIR/prompt.txt"
+
+# Generate the JSON file and merge.txt for the single example
+python scripts/dataset_preparation/prepare_json_file.py --data_folder "$SINGLE_EXAMPLE_DIR" --output "videos2caption.json"
+
+# Create a validation.json that uses the same example for consistency
+cat > "$SINGLE_EXAMPLE_DIR/validation.json" << 'EOF'
+{
+    "data": [
+      {
+        "caption": "A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out.",
+        "image_path": null,
+        "video_path": null,
+        "num_inference_steps": 50,
+        "height": 480,
+        "width": 832,
+        "num_frames": 81
+      }
+    ]
+  }
+EOF
+
+echo "Single example dataset created at $SINGLE_EXAMPLE_DIR"
+echo "Contains:"
+echo "- 1 video: $(cat $SINGLE_EXAMPLE_DIR/videos.txt)"
+echo "- 1 prompt: $(cat $SINGLE_EXAMPLE_DIR/prompt.txt)"
+echo "- Validation file created with the same example for consistency" 
@@ -21,4 +21,4 @@ torchrun --nproc_per_node=$GPU_NUM \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
-    --preprocess_task "t2v" 
+    --preprocess_task "t2v" 
@@ -3,8 +3,8 @@
 GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
-DATA_MERGE_PATH="$(dirname "$0")/crush_smol_prompts.txt"
-OUTPUT_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
+DATA_MERGE_PATH="data/crush-smol-single/merge.txt"
+OUTPUT_DIR="data/crush-smol-single_processed_t2v/"

 torchrun --nproc_per_node=$GPU_NUM \
    fastvideo/pipelines/preprocess/v1_preprocess.py \
@@ -15,11 +15,15 @@ torchrun --nproc_per_node=$GPU_NUM \
    --max_height 480 \
    --max_width 832 \
    --num_frames 81 \
-    --flow_shift 5.0 \
    --dataloader_num_workers 0 \
    --output_dir=$OUTPUT_DIR \
    --train_fps 16 \
-    --samples_per_file 8 \
-    --flush_frequency 8 \
+    --samples_per_file 1 \
+    --flush_frequency 1 \
    --video_length_tolerance_range 5 \
-    --preprocess_task "ode_trajectory" 
+    --preprocess_task "t2v"
+
+# Copy the validation.json to the output directory for consistency
+cp "data/crush-smol-single/validation.json" "$OUTPUT_DIR/"
+
+echo "Preprocessing completed. Validation file copied to $OUTPUT_DIR/" 
@@ -0,0 +1,31 @@
+{
+    "data": [
+      {
+        "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
+        "image_path": null,
+        "video_path": null,
+        "num_inference_steps": 50,
+        "height": 480,
+        "width": 832,
+        "num_frames": 77
+      },
+      {
+        "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
+        "image_path": null,
+        "video_path": null,
+        "num_inference_steps": 50,
+        "height": 480,
+        "width": 832,
+        "num_frames": 77
+      },
+      {
+        "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
+        "image_path": null,
+        "video_path": null,
+        "num_inference_steps": 50,
+        "height": 480,
+        "width": 832,
+        "num_frames": 77
+      }
+    ]
+  }
@@ -1,157 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=t2v
-#SBATCH --partition=main
-#SBATCH --nodes=4
-#SBATCH --ntasks=4
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=dmd_t2v_output/t2v_%j.out
-#SBATCH --error=dmd_t2v_output/t2v_%j.err
-#SBATCH --exclusive
-
-# Basic Info
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-export NCCL_DEBUG_SUBSYS=INIT,NET
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export TOKENIZERS_PARALLELISM=false
-export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
-# export WANDB_API_KEY='your_wandb_api_key_here'
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-
-# Configs
-NUM_GPUS=8
-
-# Model paths for Self-Forcing DMD distillation with Wan2.2:
-# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
-# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
-# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
-GENERATOR_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Updated to Wan2.2
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
-
-# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
-DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
-# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
-# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
-# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
-VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
-
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-training_args=(
-  --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
-  --output_dir "/mnt/sharefs/users/hao.zhang/SFwan2.2_t2v_finetune"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --max_train_steps 4000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 448  # Updated to match Wan2.2 config
-  --num_width 832   # Updated to match Wan2.2 config
-  --enable_gradient_checkpointing_type "full"
-  --simulate_generator_forward
-  # --log_visualization
-  --num_frames 81
-  --num_frame_per_block 3  # Frame generation block size for self-forcing
-  --enable_gradient_masking
-  --gradient_mask_last_n_frames 21
-  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
-  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
-)
-
-parallel_args=(
-  --num_gpus 32 # 64
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1 # 64
-  --hsdp_shard_dim 32
-)
-
-model_args=(
-  --model_path $GENERATOR_MODEL_PATH
-  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
-)
-
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 4
-)
-
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 20
-  --validation_sampling_steps "4"
-  --validation_guidance_scale "6.0" # not used for dmd inference
-)
-
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 500
-  --weight_only_checkpointing_steps 500
-  --weight_decay 0.01
-  --betas '0.0,0.999'
-  --max_grad_norm 1.0
-)
-
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --flow_shift 5
-  --seed 1000
-  --use_ema True
-  --ema_decay 0.99
-  --ema_start_step 100
-)
-
-dmd_args=(
-  --dmd_denoising_steps '1000,750,500,250'
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.98
-  --dfake_gen_update_ratio 5
-  --real_score_guidance_scale 3.0
-  --fake_score_learning_rate 8e-6
-  --fake_score_betas '0.0,0.999'
-  --warp_denoising_step
-)
-
-self_forcing_args=(
-  --independent_first_frame False  # Whether to treat first frame independently
-  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
-  --last_step_only False  # Whether to only use the last denoising step
-  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
-)
-
-srun torchrun \
--nnodes $SLURM_JOB_NUM_NODES \
--nproc_per_node $NUM_GPUS \
--node_rank $SLURM_PROCID \
--rdzv_backend=c10d \
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}" \
-    "${self_forcing_args[@]}"
@@ -4,7 +4,9 @@ These are end-to-end example scripts for distilling Wan2.1 T2V 1.3B model using
 ### 0. Make sure you have installed VSA

 ```bash
-pip install vsa
+cd csrc/attn
+git submodule update --init --recursive
+python setup_vsa.py install
 ```

 ### 1. Download dataset:
@@ -39,18 +39,15 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir $OUTPUT_DIR
+  --output_dir"checkpoints/wan_t2v_finetune"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -75,8 +72,6 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -96,7 +91,7 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-6
+  --learning_rate 1e-5
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
  --weight_only_checkpointing_steps 500
@@ -139,4 +134,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -39,18 +39,15 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir "$OUTPUT_DIR"
+  --output_dir "checkpoints/wan_t2v_finetune"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -75,8 +72,6 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -96,7 +91,7 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-6
+  --learning_rate 1e-5
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
  --weight_only_checkpointing_steps 500
@@ -139,4 +134,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -39,18 +39,15 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd
-  --output_dir "$OUTPUT_DIR"
+  --output_dir "checkpoints/wan_t2v_finetune"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -75,8 +72,6 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -96,7 +91,7 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-6
+  --learning_rate 1e-5
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
  --weight_only_checkpointing_steps 500
@@ -138,4 +133,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -1,3 +1,3 @@
 #!/bin/bash

-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
@@ -4,7 +4,9 @@ These are end-to-end example scripts for distilling Wan2.2 TI2V 5B model DMD+VSA
 ### 0. Make sure you have installed VSA

 ```bash
-pip install vsa
+cd csrc/attn
+git submodule update --init --recursive
+python setup_vsa.py install
 ```

 ### Data-free Distillation
@@ -40,18 +40,15 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DIR=your_validation_path  #(example:validation_64.json)
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name Wan_distillation
-  --output_dir "$OUTPUT_DIR"
+  --output_dir "your_output_dir"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -76,8 +73,6 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -97,11 +92,11 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 4e-6
+  --learning_rate 2e-5
  --lr_scheduler "cosine_with_min_lr"
  --min_lr_ratio 0.5
  --lr_warmup_steps 100
-  --fake_score_learning_rate 2e-6
+  --fake_score_learning_rate 1e-5
  --fake_score_lr_scheduler "cosine_with_min_lr"
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
@@ -146,4 +141,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -40,8 +40,6 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DIR=your_validation_path  #(example:validation_64.json)
 # export CUDA_VISIBLE_DEVICES=4,5
@@ -75,8 +73,6 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -96,11 +92,11 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 4e-6
+  --learning_rate 2e-5
  --lr_scheduler "cosine_with_min_lr"
  --min_lr_ratio 0.5
  --lr_warmup_steps 100
-  --fake_score_learning_rate 2e-6
+  --fake_score_learning_rate 1e-5
  --fake_score_lr_scheduler "cosine_with_min_lr"
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
@@ -146,4 +142,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -4,7 +4,9 @@ These are end-to-end example scripts for distilling Wan2.2 TI2V 5B model DMD+VSA
 ### 0. Make sure you have installed VSA

 ```bash
-pip install vsa
+cd csrc/attn
+git submodule update --init --recursive
+python setup_vsa.py install
 ```

 ### 1. Download dataset:
@@ -14,29 +14,26 @@ export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
 # Configs
 NUM_GPUS=1
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" 
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir "$OUTPUT_DIR"
-  --max_train_steps 4000
-  --train_batch_size 1
+  --output_dir="checkpoints/wan_t2v_finetune"
+  --max_train_steps=4000
+  --train_batch_size=1
  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
+  --gradient_accumulation_steps=1
  --num_latent_t 31
  --num_height 704
  --num_width 1280
  --num_frames 121
  --enable_gradient_checkpointing_type "full"
-  --training_state_checkpointing_steps 500
-  --weight_only_checkpointing_steps 500
+  --training_state_checkpointing_steps=500
+  --weight_only_checkpointing_steps=500
 )

 # Parallel arguments
@@ -52,8 +49,6 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -73,8 +68,8 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-6
-  --mixed_precision "bf16"
+  --learning_rate=1e-5
+  --mixed_precision="bf16"
  --weight_decay 0.01
  --max_grad_norm 1.0
 )
@@ -103,7 +98,6 @@ dmd_args=(
 torchrun \
 --nnodes 1 \
 --nproc_per_node $NUM_GPUS \
--master_port $MASTER_PORT \
    fastvideo/training/wan_distillation_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
@@ -112,4 +106,4 @@ torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -1,116 +0,0 @@
-#!/bin/bash
-
-# Basic Info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-export MASTER_PORT=29501
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-# Configs
-NUM_GPUS=1
-MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir="checkpoints/wan_t2v_finetune"
-  --max_train_steps=4000
-  --train_batch_size=1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps=1
-  --num_latent_t 31
-  --num_height 704
-  --num_width 1280
-  --num_frames 121
-  --enable_gradient_checkpointing_type "full"
-  --training_state_checkpointing_steps=500
-  --weight_only_checkpointing_steps=500
-  --lora_rank 32
-  --lora_training True
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus 1
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 1
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 4
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 200
-  --validation_sampling_steps "3"
-  --validation_guidance_scale "6.0" # not used for dmd inference
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate=1e-4
-  --mixed_precision="bf16"
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 8
-  --seed 1000
-)
-
-# DMD arguments
-dmd_args=(
-  --dmd_denoising_steps '1000,757,522'
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.98
-  --generator_update_interval 5
-  --real_score_guidance_scale 3.5
-  --VSA_sparsity 0.8
-)
-
-torchrun \
--nnodes 1 \
--nproc_per_node $NUM_GPUS \
--master_port $MASTER_PORT \
-    fastvideo/training/wan_distillation_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
@@ -1,3 +1,3 @@
 #!/bin/bash

-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -21,4 +21,4 @@ torchrun --nproc_per_node=$GPU_NUM \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
-    --preprocess_task "t2v" 
+    --preprocess_task "t2v" 
@@ -1,44 +0,0 @@
-from fastvideo import VideoGenerator
-
-# from fastvideo.configs.sample import SamplingParam
-
-OUTPUT_PATH = "video_samples"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        # FastVideo will automatically handle distributed setup
-        num_gpus=2,
-        use_fsdp_inference=True,
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
-        distributed_executor_backend="ray",
-        # image_encoder_cpu_offload=False,
-    )
-
-    # Generate videos with the same simple API, regardless of GPU count
-    prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
-    )
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
-
-    # Generate another video with a different prompt, without reloading the
-    # model!
-    prompt2 = (
-        "A majestic lion strides across the golden savanna, its powerful frame "
-        "glistening under the warm afternoon sun. The tall grass ripples gently in "
-        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
-        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
-        "cinematic.")
-    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
-
-
-if __name__ == "__main__":
-    main()
@@ -28,4 +28,4 @@ def main():
    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)

 if __name__ == "__main__":
-    main()
+    main()
@@ -1,43 +0,0 @@
-# NOTE: This is still a work in progress, and the checkpoints are not released yet.
-
-from fastvideo import VideoGenerator
-
-# from fastvideo.configs.sample import SamplingParam
-
-OUTPUT_PATH = "video_samples_self_forcing_causal_wan2_2_14B_t2v"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    generator = VideoGenerator.from_pretrained(
-        "rand0nmr/SFWan2.2-T2V-A14B-Diffusers",
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=True,
-        dit_cpu_offload=True, # DiT need to be offloaded for MoE
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125],
-        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
-        pin_cpu_memory=True,
-        init_weights_from_safetensors="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer/",
-        init_weights_from_safetensors_2="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer/",
-        num_frame_per_block=7,
-        # image_encoder_cpu_offload=False,
-    )
-
-    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
-    # sampling_param.num_frames = 45
-    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
-    # Generate videos with the same simple API, regardless of GPU count
-    prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
-    )
-    _ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, num_frames=81)
-
-
-if __name__ == "__main__":
-    main()
@@ -1,36 +0,0 @@
-from fastvideo import VideoGenerator
-
-# from fastvideo.configs.sample import SamplingParam
-
-OUTPUT_PATH = "video_samples_wan2_1_Fun"
-OUTPUT_NAME = "wan2.1_test"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    generator = VideoGenerator.from_pretrained(
-        "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers",
-        # "alibaba-pai/Wan2.2-Fun-A14B-Control",
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=True,
-        dit_cpu_offload=True, # DiT need to be offloaded for MoE
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
-        pin_cpu_memory=True,
-        # image_encoder_cpu_offload=False,
-    )
-
-    prompt = "一位年轻女性穿着一件粉色的连衣裙，裙子上有白色的装饰和粉色的纽扣。她的头发是紫色的，头上戴着一个红色的大蝴蝶结，显得非常可爱和精致。她还戴着一个红色的领结，整体造型充满了少女感和活力。她的表情温柔，双手轻轻交叉放在身前，姿态优雅。背景是简单的灰色，没有任何多余的装饰，使得人物更加突出。她的妆容清淡自然，突显了她的清新气质。整体画面给人一种甜美、梦幻的感觉，仿佛置身于童话世界中。"
-    negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
-    # prompt                  = "A young woman with beautiful, clear eyes and blonde hair stands in the forest, wearing a white dress and a crown. Her expression is serene, reminiscent of a movie star, with fair and youthful skin. Her brown long hair flows in the wind. The video quality is very high, with a clear view. High quality, masterpiece, best quality, high resolution, ultra-fine, fantastical."
-    # negative_prompt         = "Twisted body, limb deformities, text captions, comic, static, ugly, error, messy code."
-    image_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/8.png"
-    control_video_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/pose.mp4"
-
-    video = generator.generate_video(prompt, negative_prompt=negative_prompt, image_path=image_path, video_path=control_video_path, output_path=OUTPUT_PATH, output_video_name=OUTPUT_NAME, save_video=True)
-
-if __name__ == "__main__":
-    main()
@@ -1,41 +0,0 @@
-from fastvideo import VideoGenerator
-
-OUTPUT_PATH = "video_samples_wan2_2_5B_ti2v"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    model_name = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-    generator = VideoGenerator.from_pretrained(
-        model_name,
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=True,
-        dit_cpu_offload=True,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
-        # image_encoder_cpu_offload=False,
-    )
-
-    # I2V is triggered just by passing in an image_path argument
-    prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
-    image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path)
-
-    # Generate another video with a different prompt, without reloading the
-    # model!
-
-    # T2V mode
-    prompt2 = (
-        "A majestic lion strides across the golden savanna, its powerful frame "
-        "glistening under the warm afternoon sun. The tall grass ripples gently in "
-        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
-        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
-        "cinematic.")
-    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
-
-
-if __name__ == "__main__":
-    main()
@@ -1,56 +0,0 @@
-# FastVideo Gradio Local Demo
-
-This is a Gradio-based web interface for generating videos using the FastVideo framework. The demo allows users to create videos from text prompts with various customization options.
-
-## Overview
-
-The demo uses the FastVideo framework to generate videos based on text prompts. It provides a simple web interface built with Gradio that allows users to:
-
- Enter text prompts to generate videos
- Customize video parameters (dimensions, number of frames, etc.)
- Use negative prompts to guide the generation process
- Set or randomize seeds for reproducibility
-
---
-
-## Usage
-
-Run the demo with:
-
-```bash
-python examples/inference/gradio/local/gradio_local_demo.py
-```
-
-This will start a web server at `http://0.0.0.0:7860` where you can access the interface.
-
---
-
-## Model Initialization
-
-This demo initializes a `VideoGenerator` with the minimum required arguments for inference. Users can seamlessly adjust inference options between generations, including prompts, resolution, video length, *without ever needing to reload the model*.
-
-## Video Generation
-
-The core functionality is in the `generate_video` function, which:
-1. Processes user inputs
-2. Uses the FastVideo VideoGenerator from earlier to run inference (`generator.generate_video()`)
-
-## Gradio Interface
-
-The interface is built with several components:
- A text input for the prompt
- A video display for the result
- Inference options in a collapsible accordion:
-  - Height and width sliders
-  - Number of frames slider
-  - Guidance scale slider
-  - Negative prompt options
-  - Seed controls
-
-### Inference Options
-
- **Height/Width**: Control the resolution of the generated video
- **Number of Frames**: Set how many frames to generate
- **Guidance Scale**: Control how closely the generation follows the prompt
- **Negative Prompt**: Specify what you don't want to see in the video
- **Seed**: Control randomness for reproducible results
@@ -1,656 +0,0 @@
-import argparse
-import os
-import base64
-import time
-
-import gradio as gr
-from fastvideo.entrypoints.video_generator import VideoGenerator
-from fastvideo.configs.sample.base import SamplingParam
-from copy import deepcopy
-
-
-MODEL_PATH_MAPPING = {
-    "FastWan2.1-T2V-1.3B": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
-    # "FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
-}
-
-def create_timing_display(inference_time, total_time, stage_execution_times, num_frames):
-    dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A"
-    
-    timing_html = f"""
-    <div style="margin: 10px 0;">
-        <h3 style="text-align: center; margin-bottom: 10px;">⏱️ Timing Breakdown</h3>
-        <div style="display: grid; grid-template-columns: repeat(5, 1fr); gap: 10px; margin-bottom: 10px;">
-            <div class="timing-card timing-card-highlight">
-                <div style="font-size: 20px;">🚀</div>
-                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">DiT Denoising</div>
-                <div style="font-size: 16px; color: #ffa200; font-weight: bold;">{dit_denoising_time}</div>
-            </div>
-            <div class="timing-card">
-                <div style="font-size: 20px;">🧠</div>
-                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">E2E (w. vae/text encoder)</div>
-                <div style="font-size: 16px; color: #2563eb;">{inference_time:.2f}s</div>
-            </div>
-            <div class="timing-card">
-                <div style="font-size: 20px;">🎬</div>
-                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Video Encoding</div>
-                <div style="font-size: 16px; color: #dc2626;">N/A</div>
-            </div>
-            <div class="timing-card">
-                <div style="font-size: 20px;">🌐</div>
-                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Network Transfer</div>
-                <div style="font-size: 16px; color: #059669;">N/A</div>
-            </div>
-            <div class="timing-card">
-                <div style="font-size: 20px;">📊</div>
-                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Total Processing</div>
-                <div style="font-size: 18px; color: #0277bd;">{total_time:.2f}s</div>
-            </div>
-        </div>"""
-    
-    if inference_time > 0:
-        fps = num_frames / inference_time
-        timing_html += f"""
-        <div class="performance-card" style="margin-top: 15px;">
-            <span style="font-weight: bold;">Generation Speed: </span>
-            <span style="font-size: 18px; color: #6366f1; font-weight: bold;">{fps:.1f} frames/second</span>
-        </div>"""
-    
-    return timing_html + "</div>"
-def setup_model_environment(model_path: str) -> None:
-    if "fullattn" in model_path.lower():
-        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
-    else:
-        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
-    os.environ["FASTVIDEO_STAGE_LOGGING"] = "1"
-
-def load_example_prompts():
-    def contains_chinese(text):
-        return any('\u4e00' <= char <= '\u9fff' for char in text)
-    
-    def load_from_file(filepath):
-        prompts, labels = [], []
-        try:
-            with open(filepath, "r", encoding='utf-8') as f:
-                for line in f:
-                    line = line.strip()
-                    if line and not contains_chinese(line):
-                        label = line[:100] + "..." if len(line) > 100 else line
-                        labels.append(label)
-                        prompts.append(line)
-        except Exception as e:
-            print(f"Warning: Could not read {filepath}: {e}")
-        return prompts, labels
-    
-    examples, example_labels = load_from_file("examples/inference/gradio/local/prompts_final.txt")
-    
-    if not examples:
-        examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."]
-        example_labels = ["Crowded rooftop bar at night"]
-    
-    return examples, example_labels
-
-
-def create_gradio_interface(default_params: dict[str, SamplingParam], generators: dict[str, VideoGenerator]):
-    def generate_video(
-        prompt, negative_prompt, use_negative_prompt, seed, guidance_scale,
-        num_frames, height, width, randomize_seed, model_selection, progress
-    ):
-        model_path = MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers")
-        setup_model_environment(model_path)
-        try:
-            if progress:
-                progress(0.1, desc="Loading model for local inference...")
-            
-            generator = generators[model_path]
-            params = deepcopy(default_params[model_path])
-            total_start_time = time.time()
-            if progress:
-                progress(0.2, desc="Configuring parameters...")
-
-            params.prompt = prompt
-            params.seed = int(seed)
-            params.guidance_scale = guidance_scale
-            params.num_frames = int(num_frames)
-            params.height = int(height)
-            params.width = int(width)
-
-            if randomize_seed:
-                params.seed = torch.randint(0, 1000000, (1, )).item()
-
-            if use_negative_prompt and negative_prompt:
-                params.negative_prompt = negative_prompt
-            else:
-                params.negative_prompt = default_params[model_path].negative_prompt
-
-            if progress:
-                progress(0.4, desc="Generating video locally...")
-
-            output_dir = "outputs/"
-            os.makedirs(output_dir, exist_ok=True)
-            start_time = time.time()
-            result = generator.generate_video(prompt=prompt, sampling_param=params, save_video=True, return_frames=False)
-            inference_time = time.time() - start_time
-            logging_info = result.get("logging_info", None)
-            if logging_info:
-                stage_names = logging_info.get_execution_order()
-                stage_execution_times = [
-                    logging_info.get_stage_info(stage_name).get("execution_time", 0.0) 
-                    for stage_name in stage_names
-                ]
-            else:
-                stage_names = []
-                stage_execution_times = []
-            total_time = time.time() - total_start_time
-            timing_details=create_timing_display(inference_time=inference_time, total_time=total_time, stage_execution_times=stage_execution_times, num_frames=params.num_frames)
-            safe_prompt = params.prompt[:100].replace(' ', '_').replace('/', '_').replace('\\', '_')
-            video_filename = f"{params.prompt[:100]}.mp4"
-            output_path = os.path.join(output_dir, video_filename)
-
-            if progress:
-                progress(1.0, desc="Generation complete!")
-
-            return output_path, params.seed, timing_details
-
-        except Exception as e:
-            print(f"An error occurred during local generation: {e}")
-            return None, f"Generation failed: {str(e)}", ""
-
-    examples, example_labels = load_example_prompts()
-    
-    theme = gr.themes.Base().set(
-        button_primary_background_fill="#2563eb",
-        button_primary_background_fill_hover="#1d4ed8",
-        button_primary_text_color="white",
-        slider_color="#2563eb",
-        checkbox_background_color_selected="#2563eb",
-    )
-    
-    def get_default_values(model_name):
-        model_path = MODEL_PATH_MAPPING.get(model_name)
-        if model_path and model_path in default_params:
-            params = default_params[model_path]
-            return {
-                'height': params.height,
-                'width': params.width,
-                'num_frames': params.num_frames,
-                'guidance_scale': params.guidance_scale,
-                'seed': params.seed,
-            }
-        
-        return {
-            'height': 448,
-            'width': 832,
-            'num_frames': 61,
-            'guidance_scale': 3.0,
-            'seed': 1024,
-        }
-    
-    initial_values = get_default_values("FastWan2.1-T2V-1.3B")
-    
-    with gr.Blocks(title="FastWan", theme=theme) as demo:
-        gr.Image("assets/full.svg", show_label=False, container=False, height=80)
-
-        gr.HTML("""
-        <div style="text-align: center; margin-bottom: 10px;">
-            <p style="font-size: 18px;"> Make Video Generation Go Blurrrrrrr </p>
-            <p style="font-size: 18px;"> <a href="https://github.com/hao-ai-lab/FastVideo/tree/main" target="_blank">Code</a> | <a href="https://hao-ai-lab.github.io/blogs/fastvideo_post_training/" target="_blank">Blog</a> | <a href="https://hao-ai-lab.github.io/FastVideo/" target="_blank">Docs</a>  </p>
-        </div>
-        """)
-        
-        with gr.Accordion("🎥 What Is FastVideo?", open=False):
-            gr.HTML("""
-            <div style="padding: 20px; line-height: 1.6;">
-                <p style="font-size: 16px; margin-bottom: 15px;">
-                    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.
-                </p>
-            </div>
-            """)
-        
-        with gr.Row():
-            model_selection = gr.Dropdown(
-                choices=list(MODEL_PATH_MAPPING.keys()),
-                value="FastWan2.1-T2V-1.3B",
-                label="Select Model",
-                interactive=True
-            )
-
-        with gr.Row():
-            example_dropdown = gr.Dropdown(
-                choices=example_labels,
-                label="Example Prompts",
-                value=None,
-                interactive=True,
-                allow_custom_value=False
-            )
-        
-        with gr.Row():
-            with gr.Column(scale=6):
-                prompt = gr.Text(
-                    label="Prompt",
-                    show_label=False,
-                    max_lines=3,
-                    placeholder="Describe your scene...",
-                    container=False,
-                    lines=3,
-                    autofocus=True,
-                )
-            with gr.Column(scale=1, min_width=120, elem_classes="center-button"):
-                run_button = gr.Button("Run", variant="primary", size="lg")
-        
-        with gr.Row():
-            with gr.Column():
-                error_output = gr.Text(label="Error", visible=False)
-                timing_display = gr.Markdown(label="Timing Breakdown", visible=False)
-
-        with gr.Row(equal_height=True, elem_classes="main-content-row"):
-            with gr.Column(scale=1, elem_classes="advanced-options-column"):
-                with gr.Group():
-                    gr.HTML("<div style='margin: 0 0 15px 0; text-align: center; font-size: 16px;'>Advanced Options</div>")
-                    with gr.Row():
-                        height = gr.Number(
-                            label="Height",
-                            value=initial_values['height'],
-                            interactive=False,
-                            container=True
-                        )
-                        width = gr.Number(
-                            label="Width",
-                            value=initial_values['width'],
-                            interactive=False,
-                            container=True
-                        )
-                    
-                    with gr.Row():
-                        num_frames = gr.Number(
-                            label="Number of Frames",
-                            value=initial_values['num_frames'],
-                            interactive=False,
-                            container=True
-                        )
-                        guidance_scale = gr.Slider(
-                            label="Guidance Scale",
-                            minimum=1,
-                            maximum=12,
-                            value=initial_values['guidance_scale'],
-                        )
-                    
-                    with gr.Row():
-                        use_negative_prompt = gr.Checkbox(
-                            label="Use negative prompt", value=False)
-                        negative_prompt = gr.Text(
-                            label="Negative prompt",
-                            max_lines=3,
-                            lines=3,
-                            placeholder="Enter a negative prompt",
-                            visible=False,
-                        )
-
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=1000000,
-                        step=1,
-                        value=initial_values['seed'],
-                    )
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
-                    seed_output = gr.Number(label="Used Seed")
-        
-            with gr.Column(scale=1, elem_classes="video-column"):
-                result = gr.Video(
-                    label="Generated Video", 
-                    show_label=True,
-                    height=466,
-                    width=600,
-                    container=True,
-                    elem_classes="video-component"
-                )
-        
-        gr.HTML("""
-        <style>
-        .center-button {
-            display: flex !important;
-            justify-content: center !important;
-            height: 100% !important;
-            padding-top: 1.4em !important;
-        }
-        
-        .gradio-container {
-            max-width: 1200px !important;
-            margin: 0 auto !important;
-        }
-        
-        .main {
-            max-width: 1200px !important;
-            margin: 0 auto !important;
-        }
-        
-        .gr-form, .gr-box, .gr-group {
-            max-width: 1200px !important;
-        }
-        
-        .gr-video {
-            max-width: 500px !important;
-            margin: 0 auto !important;
-        }
-        
-        .main-content-row {
-            display: flex !important;
-            align-items: flex-start !important;
-            min-height: 500px !important;
-            gap: 20px !important;
-        }
-        
-        .advanced-options-column,
-        .video-column {
-            display: flex !important;
-            flex-direction: column !important;
-            flex: 1 !important;
-            min-height: 400px !important;
-            align-items: stretch !important;
-        }
-        
-        .video-column > * {
-            margin-top: 0 !important;
-        }
-        
-        .video-column .gr-video,
-        .video-component {
-            margin-top: 0 !important;
-            padding-top: 0 !important;
-        }
-        
-        .video-column .gr-video .gr-form {
-            margin-top: 0 !important;
-        }
-        
-        .advanced-options-column .gr-group,
-        .video-column .gr-video {
-            margin-top: 0 !important;
-            vertical-align: top !important;
-        }
-        
-        .advanced-options-column > *:last-child,
-        .video-column > *:last-child {
-            flex-grow: 0 !important;
-        }
-        
-        @media (max-width: 1400px) {
-            .main-content-row {
-                min-height: 600px !important;
-            }
-            
-            .advanced-options-column,
-            .video-column {
-                min-height: 600px !important;
-            }
-        }
-        
-        @media (max-width: 1200px) {
-            .main-content-row {
-                flex-direction: column !important;
-                align-items: stretch !important;
-            }
-            
-            .advanced-options-column,
-            .video-column {
-                min-height: auto !important;
-                width: 100% !important;
-            }
-        }
-        
-        .timing-card {
-            background: var(--background-fill-secondary) !important;
-            border: 1px solid var(--border-color-primary) !important;
-            color: var(--body-text-color) !important;
-            padding: 10px;
-            border-radius: 8px;
-            text-align: center;
-            min-height: 80px;
-            display: flex;
-            flex-direction: column;
-            justify-content: center;
-        }
-        
-        .timing-card-highlight {
-            background: var(--background-fill-primary) !important;
-            border: 2px solid var(--color-accent) !important;
-        }
-        
-        .performance-card {
-            background: var(--background-fill-secondary) !important;
-            border: 1px solid var(--border-color-primary) !important;
-            color: var(--body-text-color) !important;
-            padding: 10px;
-            border-radius: 6px;
-            text-align: center;
-        }
-        
-        .gr-number input[readonly] {
-            background-color: var(--background-fill-secondary) !important;
-            border: 1px solid var(--border-color-primary) !important;
-            color: var(--body-text-color-subdued) !important;
-            cursor: default !important;
-            text-align: center !important;
-            font-weight: 500 !important;
-        }
-        </style>
-        """)
-        
-        def on_example_select(example_label):
-            if example_label and example_label in example_labels:
-                index = example_labels.index(example_label)
-                return examples[index]
-            return ""
-        
-        example_dropdown.change(
-            fn=on_example_select,
-            inputs=example_dropdown,
-            outputs=prompt,
-        )
-        
-        gr.HTML("""
-        <div style="text-align: center; margin-top: 10px; margin-bottom: 15px;">
-            <p style="font-size: 16px; margin: 0;">Note that this demo is meant to showcase FastWan's quality and that under a large number of requests, generation speed may be affected.</p>
-        </div>
-        """)
-        
-        use_negative_prompt.change(
-            fn=lambda x: gr.update(visible=x),
-            inputs=use_negative_prompt,
-            outputs=negative_prompt,
-        )
-        
-        def on_model_selection_change(selected_model):
-            if not selected_model:
-                selected_model = "FastWan2.1-T2V-1.3B"
-            
-            model_path = MODEL_PATH_MAPPING.get(selected_model)
-            
-            if model_path and model_path in default_params:
-                params = default_params[model_path]
-                return (
-                    gr.update(value=params.height),
-                    gr.update(value=params.width),
-                    gr.update(value=params.num_frames),
-                    gr.update(value=params.guidance_scale),
-                    gr.update(value=params.seed),
-                )
-            
-            return (
-                gr.update(value=448),
-                gr.update(value=832),
-                gr.update(value=61),
-                gr.update(value=3.0),
-                gr.update(value=1024),
-            )
-        
-        model_selection.change(
-            fn=on_model_selection_change,
-            inputs=model_selection,
-            outputs=[height, width, num_frames, guidance_scale, seed],
-        )
-        
-        def handle_generation(*args, progress=None, request: gr.Request = None):
-            model_selection, prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, num_frames, height, width, randomize_seed = args
-            
-            result_path, seed_or_error, timing_details = generate_video(
-                prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, 
-                num_frames, height, width, randomize_seed, model_selection, progress
-            )
-            if result_path and os.path.exists(result_path):
-                return (
-                    result_path, 
-                    seed_or_error, 
-                    gr.update(visible=False),
-                    gr.update(visible=True, value=timing_details),
-                )
-            else:
-                return (
-                    None, 
-                    seed_or_error, 
-                    gr.update(visible=True, value=seed_or_error),
-                    gr.update(visible=False),
-                )
-        
-        run_button.click(
-            fn=handle_generation,
-            inputs=[
-                model_selection,
-                prompt,
-                negative_prompt,
-                use_negative_prompt,
-                seed,
-                guidance_scale,
-                num_frames,
-                height,
-                width,
-                randomize_seed,
-            ],
-            outputs=[result, seed_output, error_output, timing_display],
-            concurrency_limit=20,
-        )
-    
-    return demo
-
-
-def main():
-    parser = argparse.ArgumentParser(description="FastVideo Gradio Local Demo")
-    parser.add_argument("--t2v_model_paths", type=str,
-                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
-                        help="Comma separated list of paths to the T2V model(s)")
-    parser.add_argument("--host", type=str, default="0.0.0.0",
-                        help="Host to bind to")
-    parser.add_argument("--port", type=int, default=7860,
-                        help="Port to bind to")
-    args = parser.parse_args()
-    generators = {}
-    default_params = {}
-    model_paths = args.t2v_model_paths.split(",")
-    for model_path in model_paths:
-        print(f"Loading model: {model_path}")
-        setup_model_environment(model_path)
-        generators[model_path] = VideoGenerator.from_pretrained(model_path)
-        default_params[model_path] = SamplingParam.from_pretrained(model_path)
-    demo = create_gradio_interface(default_params, generators)
-    print(f"Starting Gradio frontend at http://{args.host}:{args.port}")
-    print(f"T2V Models: {args.t2v_model_paths}")
-    
-    from fastapi import FastAPI, Request, HTTPException
-    from fastapi.responses import HTMLResponse, FileResponse
-    import uvicorn
-    
-    app = FastAPI()
-    
-    @app.get("/logo.png")
-    def get_logo():
-        return FileResponse(
-            "assets/full.svg",
-            media_type="image/svg+xml",
-            headers={
-                "Cache-Control": "public, max-age=3600",
-                "Access-Control-Allow-Origin": "*"
-            }
-        )
-    
-    @app.get("/favicon.ico")
-    def get_favicon():
-        favicon_path = "assets/icon-simple.svg"
-        
-        if os.path.exists(favicon_path):
-            return FileResponse(
-                favicon_path, 
-                media_type="image/svg+xml",
-                headers={
-                    "Cache-Control": "public, max-age=3600",
-                    "Access-Control-Allow-Origin": "*"
-                }
-            )
-        else:
-            raise HTTPException(status_code=404, detail="Favicon not found")
-    
-    @app.get("/", response_class=HTMLResponse)
-    def index(request: Request):
-        base_url = str(request.base_url).rstrip('/')
-        return f"""
-        <!DOCTYPE html>
-        <html lang="en">
-        <head>
-            <meta charset="UTF-8" />
-            <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-            
-            <title>FastWan</title>
-            <meta name="title" content="FastWan">
-            <meta name="description" content="Make video generation go blurrrrrrr">
-            <meta name="keywords" content="FastVideo, video generation, AI, machine learning, FastWan">
-            
-            <meta property="og:type" content="website">
-            <meta property="og:url" content="{base_url}/">
-            <meta property="og:title" content="FastWan">
-            <meta property="og:description" content="Make video generation go blurrrrrrr">
-            <meta property="og:image" content="{base_url}/logo.png">
-            <meta property="og:image:width" content="1200">
-            <meta property="og:image:height" content="630">
-            <meta property="og:site_name" content="FastWan">
-            
-            <meta property="twitter:card" content="summary_large_image">
-            <meta property="twitter:url" content="{base_url}/">
-            <meta property="twitter:title" content="FastWan">
-            <meta property="twitter:description" content="Make video generation go blurrrrrrr">
-            <meta property="twitter:image" content="{base_url}/logo.png">
-            <link rel="icon" type="image/png" sizes="32x32" href="/favicon.ico">
-            <link rel="icon" type="image/png" sizes="16x16" href="/favicon.ico">
-            <link rel="apple-touch-icon" href="/favicon.ico">
-            <style>
-                body, html {{
-                    margin: 0;
-                    padding: 0;
-                    height: 100%;
-                    overflow: hidden;
-                }}
-                iframe {{
-                    width: 100%;
-                    height: 100vh;
-                    border: none;
-                }}
-            </style>
-        </head>
-        <body>
-            <iframe src="/gradio" width="100%" height="100%" style="border: none;"></iframe>
-        </body>
-        </html>
-        """
-    
-    app = gr.mount_gradio_app(
-        app, 
-        demo, 
-        path="/gradio",
-        allowed_paths=[os.path.abspath("outputs"), os.path.abspath("fastvideo-logos")]
-    )
-    
-    uvicorn.run(app, host=args.host, port=args.port)
-
-
-if __name__ == "__main__":
-    
-    main() 
@@ -1,11 +0,0 @@
-A dynamic shot of a sleek black motorcycle accelerating down an empty highway at sunset. The bike's engine roars as it gains speed, smoke trailing from the tires. The rider, wearing a black leather jacket and helmet, leans forward with determination, gripping the handlebars tightly. The camera follows the motorcycle from a distance, capturing the dust kicked up behind it, then zooms in to show the intense focus on the rider's face. The background showcases the endless road stretching into the horizon with vibrant orange and pink hues of the setting sun. Medium shot transitioning to close-up.
-A Jedi Master Yoda, recognizable by his green skin, large ears, and wise wrinkles, is performing on a small stage, strumming a guitar with great concentration. Yoda wears a casual robe and sits on a stool, his eyes closed as he plays, fully immersed in the music. The stage is dimly lit with spotlights highlighting Yoda, creating a mystical atmosphere. The background shows a live audience watching intently. Medium close-up shot focusing on Yoda's expressive face and hands moving gracefully over the guitar strings.
-A cute, fluffy panda bear is preparing a meal in a cozy, modern kitchen. The panda is standing at a wooden countertop, wearing a white chef’s hat and apron. It skillfully stirs a pot on the stove with one hand while holding a spatula in the other. The kitchen is well-lit, with appliances and cabinets in pastel colors, creating a warm and inviting atmosphere. The panda moves gracefully, with a focused and determined expression, as steam rises from the pot. Medium shot focusing on the panda’s actions at the stove.
-In a futuristic Tokyo rooftop during a heavy rainstorm, a robotic DJ stands behind a turntable, spinning vinyl records in a cyberpunk night setting. The robot has metallic, sleek body parts with glowing blue LED lights, and it moves gracefully with the beat. Raindrops create a shimmering effect as they hit the ground and the DJ. The surrounding environment features neon signs, towering skyscrapers, and a dark, misty atmosphere. The camera starts with a wide shot of the city skyline before zooming in on the DJ performing. Sci-fi, fantasy.
-A realistic animated scene featuring a polar bear playing a guitar. The polar bear is standing upright, wearing a cozy fur vest and fingerless gloves. It holds the guitar with both hands, strumming the strings with one hand while plucking them with the other, showcasing natural, fluid motions. The polar bear's expressive face shows concentration and joy as it plays. The background is a snowy Arctic landscape with icebergs and a clear blue sky. The scene captures the bear from a mid-shot angle, focusing on its interaction with the guitar.
-The scene opens to a breathtaking view of a tranquil ocean horizon at dusk, displaying a vibrant tapestry of oranges, pinks, and purples as the sun sets. In the foreground, tall, swaying palm trees frame the scene, their silhouettes stark against the colorful sky. The ocean itself shimmers with reflections of the sunset, creating a peaceful, almost ethereal atmosphere. A small boat can be seen in the distance, centered on the horizon, adding a sense of scale and solitude to the scene. The waves gently lap the shore, creating faint patterns on the sandy beach, which stretches across the foreground. Above, the sky is dotted with scattered clouds that catch the last light of the day, enhancing the drama and beauty of the scene. The overall mood is serene and contemplative, capturing a perfect moment of nature’s grandeur.
-A large, modern semi-truck accelerating down an empty highway, gaining speed with each second. The truck's powerful engine roars as it moves forward, smoke billowing from the tires. The camera starts from a wide shot, capturing the truck in the distance, then smoothly zooms in to follow the vehicle as it speeds up. The truck's headlights illuminate the road ahead, casting a bright glow. The truck driver can be seen through the windshield, focused and determined. The background shows the vast openness of the highway stretching into the horizon under a clear blue sky. Medium to close-up shots of the truck as it accelerates.
-Soft blue light pulses from the blade’s rune-etched hilt, illuminating nearby moss-covered roots and ferns. The surrounding trees are tall and gnarled, their branches curling like claws overhead. Fog swirls gently at ground level, parting slightly as a figure in a cloak approaches from the distance. Medium shot slowly zooming toward the sword, emphasizing its mystical aura.
-The video opens with a tranquil scene in the heart of a dense forest, emphasizing two large, textured tree trunks in the foreground framing the view. Sunlight filters through the canopy above, casting intricate patterns of light and shadow on the trees and the ground. Between the tree trunks, a clear view of a calm, muddy river unfolds, its surface shimmering under the gentle sunlight. The riverbank is decorated with a variety of small bushes and vibrant foliage, subtly transitioning into the deep greens of tall, leafy plants. In the background, the dense forest looms, filled with dark, towering trees, their branches intertwining to form an intricate canopy. The scene is bathed in the soft glow of the sun, creating a serene and picturesque setting. Occasional sunbeams pierce through the foliage, adding a magical aura to the landscape. The vibrant reds and oranges of the smaller plants add contrast, bringing warmth to the earthy tones of the scenery. Overall, this harmonious blend of natural elements creates a peaceful and idyllic forest setting.
-A lone figure stands on a large, moss-covered rock, surrounded by the soft rush of a nearby stream. The figure is wearing white sneakers and shorts, with a plaid shirt that hangs loosely in the breeze. The lighting creates dramatic shadows, enhancing the textures of the rock and the subtle movement of the water below. In the background, a waterfall cascades into the stream, completing this tranquil and serene nature scene.
-In an industrial setting, a person leans casually against a railing, exuding a sense of confidence and composure. They are wearing a striking outfit, consisting of a vibrant, patterned jacket over a simple white crop top, creating a bold contrast. The atmosphere is infused with warm, ambient lighting that casts soft shadows on the concrete walls and metallic surfaces. Intricate wiring and pipes form an intricate backdrop, enhancing the urban aesthetic. Their relaxed posture and direct, engaging gaze suggest a sense of ease in this industrial environment. This scene encapsulates a blend of modern fashion and gritty, urban architecture, creating a visually compelling narrative.
@@ -1,47 +0,0 @@
-A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.
-The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object.
-The video shows a cylindrical object with a cityscape image being flattened as if it were under a hydraulic press. The object is placed on a metal platform, and a large, striped cylinder presses down on it, causing it to collapse and release a liquid inside. The background features a green wall with a yellow and red warning sign.
-A red toy car is being crushed by a large hydraulic press, which is flattening objects as if they were under a hydraulic press.
-A large, cylindrical object is seen pressing down on a small orange ball, causing it to flatten as if it were under a hydraulic press. The background features a green wall with yellow and red warning signs.
-The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is shown compressing a wooden object, which shatters into small pieces. The background features a green wall with a yellow sign displaying a lightning bolt.
-A large metal cylinder is seen descending, flattening objects as if they were under a hydraulic press. The cylinder compresses a stack of matches and boxes, causing them to crumble into small pieces. The scene is set against a green background with yellow and red signs.
-A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out.
-The video shows a metal press flattening objects as if they were under a hydraulic press. The press is pressing down on a pile of colorful gummy candies, squishing them into a pile of squiggly shapes. The press is made of metal and has a large base, and the gummy candies are of various colors, including red, green, and orange. The background is a green wall, and the press is placed on a metal surface.
-A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces.
-The video shows a stack of colorful sponges being flattened as if they were under a hydraulic press. The sponges, which are pink, white, blue, and green, are compressed into a smaller size, demonstrating the press's power. The background features a green wall with a yellow and red sign, adding context to the setting.
-A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, leaving a pile of debris around it.
-A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press.
-The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.
-The video shows a close-up of a metal cylinder pressing down on a yellow object, which is being flattened as if it were under a hydraulic press. The cylinder is positioned above the object, and the force is causing the object to compress and spread out, creating a visible deformation. The background is blurred, focusing attention on the action of the cylinder and the object being flattened.
-A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.
-The video shows a hydraulic press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing two colorful objects that resemble sandwiches. The press is yellow and black striped, and the objects being flattened are placed on a metal plate. The background is green, and the press is moving down, compressing the objects.
-The scene shows a metal press with a yellow and black striped pattern, holding a container filled with chocolate. A metal cylinder is descending, flattening the chocolate as if it were under a hydraulic press. The background is a green wall, and the press is mounted on a sturdy metal frame.
-The video shows a colorful sponge being flattened as if it were under a hydraulic press, with the sponge being compressed and eventually flattened into a thin layer.
-The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is pressing down on a stack of wooden blocks, causing them to crumble and break apart. The press is black and yellow striped, and the wooden blocks are small and rectangular. The background is green, and the press is sitting on a metal table.
-A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces.
-The video shows a stack of colorful sponges being flattened by a large, cylindrical object, which appears to be a hydraulic press. The sponges, which are pink, blue, white, and green, are compressed into a single layer, demonstrating the press's powerful force. The background features a green wall with a yellow and red sign, adding context to the industrial setting.
-A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, demonstrating the immense pressure applied by the cylinder.
-A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press. The popcorn is crushed and scattered around the base of the cylinder, creating a satisfying visual effect.
-The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is composed of a large, cylindrical metal cylinder with yellow and black stripes, and a metal base. The objects being flattened are two cylindrical blocks of cotton candy, one pink and one blue. The press is positioned on a metal table, and the background features a green wall with a yellow and red sign.
-The video shows a large orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.
-The video shows a cylindrical object being pressed down onto a flat surface, causing the objects beneath it to be flattened as if they were under a hydraulic press. The objects being flattened appear to be yellow and are being crushed into a pile of debris. The background is a greenish-gray color, and the surface on which the objects are being flattened is metallic and shiny.
-A green and blue object with a spiky texture is being flattened by a large, cylindrical metal press, demonstrating its resilience and durability.
-The video shows a stack of caramelized sugar cubes being flattened as if they were under a hydraulic press, resulting in a messy pile of broken sugar on the table.
-A large metal cylinder is seen pressing down on a pile of colorful jelly beans, flattening them as if they were under a hydraulic press.
-The video shows a machine with a yellow and black striped cylinder pressing down on a stack of colorful sponges, flattening them as if they were under a hydraulic press. The machine is situated in a green-walled room with warning signs in the background.
-The video shows a machine with a yellow and black striped cylinder, which is pressing down on two colorful objects, flattening them as if they were under a hydraulic press. The machine appears to be in a workshop or industrial setting, with a green wall in the background. The objects being flattened are green and orange, and the machine is covered in dirt and grime, indicating it has been used frequently.
-The video shows a large, industrial press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing a pile of pink objects into a pile of crumbs. The press is large and metallic, with a yellow and black striped pattern on its side. The background is a green wall with a yellow warning sign.
-The video shows a pink, sparkly ball being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press.
-A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and segments.
-The video shows a machine with a yellow and black striped cylinder, which is flattening objects as if they were under a hydraulic press. The machine is pressing down on two colorful objects, causing them to compress and flatten. The background is a green wall, and the machine appears to be in a workshop or industrial setting.
-The video shows a large, yellow and black striped cylinder flattening objects as if they were under a hydraulic press. The objects being flattened are pink and are being crushed into small pieces. The background is a green wall with a yellow sign.
-The video shows a machine with a yellow and black striped cylinder pressing down on two colorful objects, which are flattened as if they were under a hydraulic press. The machine is positioned on a metal platform, and the background is a green wall.
-A green cube is being compressed by a hydraulic press, which flattens the object as if it were under a hydraulic press. The press is shown in action, with the cube being squeezed into a smaller shape.
-A pink, sparkly ball is being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press.
-A red cabbage is being crushed by a hydraulic press, which flattens the objects as if they were under a hydraulic press. The press is shown in action, compressing the cabbage into a smaller, more compact form.
-A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and pulp.
-A large metal press is shown compressing a stack of burgers, causing them to be flattened and crushed into a pile of ground meat.
-A pizza is being crushed by a hydraulic press, causing the toppings to spread out and the crust to crumble.
-A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.
-A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.
-A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.
@@ -1,94 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=1
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_ode_init"
-  --output_dir "wan_ode_init_crush_smol"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "wan_ode_init_crush_smol"
-  --max_train_steps 6000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  --warp_denoising_step
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 1
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR" 
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "6.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 6e-6
-  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-)
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/ode_causal_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,136 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=2e6B8_16kFV_ode_vidprom
-#SBATCH --partition=main
-#SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=ode_vidprom16k/ode_vidprom8b16k_2e-6.out
-#SBATCH --error=ode_vidprom16k/ode_vidprom8b16k_2e-6.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate your-conda-env
-
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_API_KEY=your-wandb-api-key
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-
-MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="your-data-dir"
-VALIDATION_DATASET_FILE="examples/training/consistency_finetune/causal_ode_init/validation.json"
-OUTPUT_DIR="your-output-dir"
-INIT_WEIGHTS_FROM_SAFETENSORS="your-init-weights-from-safetensors" # bidirectional weights from Wan2.1-T2V-1.3B-Diffusers
-NUM_GPUS=8
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "wan_ode_init"
-  --output_dir $OUTPUT_DIR
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "vidprom_8b16k_ode_init_2e-6"
-  # --resume_from_checkpoint "ode_init_diffusers/"
-  --warp_denoising_step
-  --log_visualization
-  --max_train_steps 6001
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 480
-  --num_width 832
-  --num_frames 81
-  --dmd_denoising_steps "1000,750,500,250"
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim $NUM_GPUS
-  --hsdp_shard_dim 1
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR" 
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "6.0"
-  --init_weights_from_safetensors $INIT_WEIGHTS_FROM_SAFETENSORS
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 2e-6
-  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 500
-  --training_state_checkpointing_steps 500
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-  # --enable_gradient_checkpointing_type "full"
-)
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-srun torchrun \
-  --nnodes $SLURM_JOB_NUM_NODES \
-  --nproc_per_node $NUM_GPUS \
-  --node_rank $SLURM_PROCID \
-  --rdzv_backend=c10d \
-  --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/ode_causal_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,76 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "Elon Musk, dressed in a sleek white spacesuit with a reflective visor, walks confidently across the lunar surface. His posture is upright, and he moves steadily with purpose. The moon's rocky terrain and scattered boulders surround him, casting shadows under the dim sunlight. The background shows vast stretches of the moon's barren landscape with craters and dust clouds kicked up by his boots. The scene captures a wide shot, emphasizing the vastness and desolation of the lunar environment. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "In a dynamic action-packed sequence set in the Marvel multiverse, Spider-Man and Venom engage in an intense battle. Spider-Man, in his classic red and blue suit, swings and dodges venomous attacks from the black symbiote-covered Venom. Both characters display a range of acrobatic moves and powerful strikes. The environment is a chaotic urban landscape with crumbling buildings and neon lights, reflecting the multiversal theme. The camera captures the epic fight from various angles, including wide shots to show the scale of destruction and close-ups to highlight their fierce expressions and physical combat. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A warm, family-oriented scene depicting a father getting ready to leave the house to buy milk. The father, a middle-aged man with a kind face and a casual outfit, picks up a jacket from the coat rack. His posture is upright as he bends down slightly to put on his shoes. In the background, there are glimpses of a cozy living room with a family photograph on the wall. The camera focuses closely on the father, capturing his gentle smile and reassuring nod towards the camera before he opens the front door and steps outside. Static medium close-up shot. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "Close-up shot of a man with a prosthetic hand that functions as a rocket launcher. He looks at his new hand with a mix of amazement and concern, his facial expression showing a blend of curiosity and apprehension. The prosthetic hand is sleek and metallic, with intricate details that resemble a high-tech weapon. The background is a dimly lit laboratory with various scientific equipment and monitors displaying data. The man stands in a relaxed posture, his other hand resting on his hip, as he inspects his new limb. The scene is rendered in a realistic sci-fi style, emphasizing the futuristic technology and the man's emotional response to his new appendage. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "Realistic CCTV footage style, Kim Taehyung from the band BTS is involved in a drug deal, caught on camera. Kim Taehyung appears nervous and cautious, wearing casual clothing typical of a public space. He exchanges items discreetly with another person, who is partially obscured. Both individuals maintain a watchful demeanor, occasionally glancing around to ensure no one is watching them. The lighting is dim, with flickering fluorescent lights casting shadows on their faces. The background shows a typical urban setting with blurred figures moving in the distance. Static camera angle, medium close-up shot focusing on the interaction between Taehyung and the other individual. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "Photorealistic studio setup with professional lighting, showcasing detailed cubic dissections of experimental plastic and felt-like materials on a pristine white background. Each cube reveals intricate layers and textures of the materials, emphasizing their unique properties. The scene has a shallow depth of field initially, then slowly pulls out to reveal the full arrangement of cubes, maintaining a wide depth of field throughout the transition. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}
@@ -61,8 +61,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 2e-5
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 2000
-  --training_state_checkpointing_steps 2000
+  --checkpointing_steps 2000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
@@ -95,8 +95,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 1e-5
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
+  --checkpointing_steps 1000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
@@ -93,8 +93,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 1e-6
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
+  --checkpointing_steps 1000
  --weight_decay 0.01
  --max_grad_norm 1.0
 )
@@ -7,7 +7,9 @@ These are e2e example scripts for finetuning Wan2.1 T2V with VSA to accelerate i
 ## Make sure you have installed VSA

 ```bash
-pip install vsa
+cd csrc/attn
+git submodule update --init --recursive
+python setup_vsa.py install
 ```

 ### Download the synthetic dataset:
@@ -1,134 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=t2v
-#SBATCH --partition=main
-#SBATCH --nodes=8
-#SBATCH --ntasks=8
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=VSA_t2v_output/t2v_%j.out
-#SBATCH --error=VSA_t2v_output/t2v_%j.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate your_env
-
-# Basic Info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR=your_data_dir
-VALIDATION_DATASET_FILE=your_validation_dataset_file
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_t2v_VSA
-  --output_dir "checkpoints/wan_t2v_finetune_VSA"
-  --max_train_steps 4000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 480
-  --num_width 832
-  --num_frames 81
-  # --enable_gradient_checkpointing_type "full" # if OOM enable this
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus 64
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 64
-  --hsdp_shard_dim 1
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 4
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "5.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-6
-  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 1
-  --seed 1000
-)
-
-# VSA arguments
-vsa_args=(
-  --VSA_decay_rate 0.03 \
-  --VSA_decay_interval_steps 50 \
-  --VSA_sparsity 0.9 \
-)
-
-srun torchrun \
--nnodes $SLURM_JOB_NUM_NODES \
--nproc_per_node $NUM_GPUS \
--node_rank $SLURM_PROCID \
--rdzv_backend=c10d \
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/wan_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  
@@ -91,10 +91,9 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-6
+  --learning_rate 1e-5
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
+  --checkpointing_steps 1000
  --weight_decay 0.01
  --max_grad_norm 1.0
 )
@@ -61,8 +61,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 1e-5
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
+  --checkpointing_steps 1000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
@@ -95,8 +95,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 1e-5
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
+  --checkpointing_steps 1000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
@@ -61,8 +61,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 1e-5
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
+  --checkpointing_steps 1000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
@@ -61,8 +61,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 5e-5
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
+  --checkpointing_steps 1000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
@@ -92,8 +92,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 5e-5
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 400
-  --training_state_checkpointing_steps 400
+  --checkpointing_steps 400
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
@@ -61,8 +61,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 5e-5
  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 400
-  --training_state_checkpointing_steps 400
+  --checkpointing_steps 400
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
@@ -10,7 +10,6 @@ torchrun --nproc_per_node=$GPU_NUM \
    --model_path $MODEL_PATH \
    --mode preprocess \
    --workload_type t2v \
-    --preprocess.video_loader_type torchvision \
    --preprocess.dataset_type merged \
    --preprocess.dataset_path $DATASET_PATH \
    --preprocess.dataset_output_dir $OUTPUT_DIR \
@@ -28,4 +28,4 @@
      "num_frames": 77
    }
  ]
-}
+}
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Matthew Noto	4477dbea67	training refactor done	2025-09-08 03:48:01 +00:00
Matthew Noto	04685a3ecd	kv cache debugging	2025-09-08 02:27:19 +00:00
Matthew Noto	b8392e9b2a	revert	2025-09-07 23:38:02 +00:00
Matthew Noto	60b71c5053	new branch	2025-09-07 23:36:43 +00:00
JerryZhou54	2b7bf88a4d	new branch	2025-09-07 08:49:36 +00:00
RandNMR73	328eb611c4	training refactor	2025-09-07 02:40:19 +00:00
RandNMR73	3408e20d7a	Merge branch 'matthew/causal' of github.com:hao-ai-lab/FastVideo into matthew/causal	2025-09-07 00:17:20 +00:00
SolitaryThinker	4328fe1ebf	fix warping	2025-09-05 10:20:10 +00:00
SolitaryThinker	6d0eba5789	warp timestep	2025-09-05 09:54:35 +00:00
RandNMR73	b3e76ae7dd	Merge branch 'matthew/causal' of github.com:hao-ai-lab/FastVideo into matthew/causal	2025-09-02 00:57:20 +00:00
Matthew Noto	1ffd80ee51	training experiments in progress	2025-09-02 00:03:58 +00:00
Matthew Noto	e84fdaedde	add single example preprocessing	2025-09-01 03:49:33 +00:00
Matthew Noto	dd0fe401c9	fix kv-cache + training loop	2025-09-01 01:52:11 +00:00
Matthew Noto	60eac9f18b	clean-up + kv cache debugging	2025-08-31 08:31:19 +00:00
Matthew Noto	a75bccb75d	debugging kv cache	2025-08-31 05:30:21 +00:00
RandNMR73	c52ff91747	Merge branch 'matthew/causal' of github.com:hao-ai-lab/FastVideo into matthew/causal	2025-08-28 12:21:18 +00:00
RandNMR73	36bb0935a8	val data path fix	2025-08-28 12:18:41 +00:00
RandNMR73	d1e26abd63	Merge branch 'matthew/causal' of github.com:hao-ai-lab/FastVideo into matthew/causal	2025-08-28 11:53:46 +00:00
RandNMR73	b6f187f338	working training script	2025-08-28 11:32:53 +00:00
RandNMR73	de4938c3a7	working training script	2025-08-28 11:27:18 +00:00
RandNMR73	a9f7407228	debugging training	2025-08-28 08:06:31 +00:00
RandNMR73	258d1da0d3	training script testing	2025-08-28 05:03:12 +00:00
RandNMR73	17f6dff632	training branch	2025-08-28 04:49:51 +00:00
RandNMR73	e66f16057f	training script testing	2025-08-28 04:42:51 +00:00
JerryZhou54	8acc7c3655	small fix	2025-08-28 04:10:51 +00:00
JerryZhou54	8dd0b6536d	Fix pre-commit tests	2025-08-28 04:07:11 +00:00
JerryZhou54	1a79b30ea4	Fix inter-block issues	2025-08-28 03:56:37 +00:00
JerryZhou54	772ead0d34	Wan2.1 causal, few-step inference runnable	2025-08-28 03:56:37 +00:00
SolitaryThinker	1575102965	debugging model	2025-08-28 03:56:37 +00:00
SolitaryThinker	b198ba5607	checkpoint	2025-08-28 03:56:34 +00:00
William Lin	cf1942fd47	[dev] Will/causal infer (#764 )	2025-08-28 03:55:16 +00:00
Wei Zhou	b5519f1f91	Delete fastvideo/configs/models/dits/causal_wanvideo.py	2025-08-28 03:55:16 +00:00
JerryZhou54	a464f96b95	Fix file dir	2025-08-28 03:55:16 +00:00
JerryZhou54	363cf0d173	First version of causalwan	2025-08-28 03:55:15 +00:00
JerryZhou54	36371c5689	checkpoint	2025-08-28 03:55:15 +00:00
SolitaryThinker	8fea7c02b5	debugging model	2025-08-27 02:56:31 -07:00
SolitaryThinker	e2b6f49879	checkpoint	2025-08-27 02:12:43 -07:00
William Lin	9f24aef7cf	[dev] Will/causal infer (#764 )	2025-08-26 23:13:20 -07:00
Wei Zhou	7a489da74d	Delete fastvideo/configs/models/dits/causal_wanvideo.py	2025-08-26 16:33:25 -07:00
JerryZhou54	cf230dcccd	Fix file dir	2025-08-26 23:32:30 +00:00
JerryZhou54	c4521e8953	First version of causalwan	2025-08-26 23:31:46 +00:00
JerryZhou54	026ee8d9f4	checkpoint	2025-08-26 19:49:57 +00:00