[Misc] Remove STA (#1124)

Co-authored-by: Will Lin <wlsaidhi@gmail.com>
2026-02-23 15:14:42 -08:00
parent 37cab18780
commit fc80d95d7e
48 changed files with 199 additions and 761610 deletions
@@ -139,18 +139,6 @@ steps:
                  - TEST_TYPE=training_vsa
                agents:
                  queue: "default"
-            - path:
-                - "fastvideo/**"
-                - "fastvideo-kernel/**"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Inference Tests STA"
-                env:
-                  - TEST_TYPE=inference_sta
-                agents:
-                  queue: "default"
            - path:
                - "fastvideo-kernel/**"
                - "pyproject.toml"
@@ -89,10 +89,6 @@ case "$TEST_TYPE" in
        log "Running training VSA tests..."
        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests_VSA"
        ;;
-    "inference_sta")
-        log "Running inference STA tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_STA"
-        ;;
    "kernel_tests")
        log "Running kernel tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_kernel_tests"
@@ -47,16 +47,6 @@ on:
        required: false
        default: false
        type: boolean
-      run_inference_test_STA:
-        description: "Run inference-test-STA"
-        required: false
-        default: false
-        type: boolean
-      run_precision_test_STA:
-        description: "Run precision-test-STA"
-        required: false
-        default: false
-        type: boolean
      run_precision_test_VSA:
        description: "Run precision-test-VSA"
        required: false
@@ -90,8 +80,6 @@ jobs:
      transformer-test: ${{ steps.filter.outputs.transformer-test }}
      training-test: ${{ steps.filter.outputs.training-test }}
      training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }}
-      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:
@@ -106,12 +94,6 @@ jobs:
              - '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'
            vsa-kernel-paths: &vsa-kernel-paths
              - 'csrc/attn/video_sparse_attn/**'
              - 'csrc/attn/video_sparse_attn/tk/**'
@@ -148,13 +130,6 @@ jobs:
              - 'fastvideo/**'
              - *common-paths
              - *vsa-kernel-paths
-            inference-test-STA:
-              - 'fastvideo/**'
-              - *common-paths
-              - *sta-kernel-paths
-            precision-test-STA:
-              - *common-paths
-              - *sta-kernel-paths
            precision-test-VSA:
              - *common-paths
              - *vsa-kernel-paths
@@ -282,44 +257,6 @@ jobs:
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}

-  inference-test-STA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.inference-test-STA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_inference_test_STA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "inference-test-STA"
-      gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 2
-      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/tests/inference/STA -srP"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  precision-test-STA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-STA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_STA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "precision-test-STA"
-      gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_sta.py"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
  precision-test-VSA:
    needs: change-filter
    if: >-
@@ -378,7 +315,7 @@ jobs:

  runpod-cleanup:
    # Add other jobs to this list as you create them
-    needs: [encoder-test, vae-test, transformer-test, ssim-test, training-test, training-test-VSA, inference-test-STA, precision-test-STA, precision-test-VSA] 
+    needs: [encoder-test, vae-test, transformer-test, ssim-test, training-test, training-test-VSA, precision-test-VSA] 
    if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
    runs-on: ubuntu-latest
    steps:
@@ -395,7 +332,7 @@ jobs:

      - name: Cleanup all RunPod instances
        env:
-          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"]'
+          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", "precision-test-VSA"]'
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        run: python .github/scripts/runpod_cleanup.py
@@ -1,249 +0,0 @@
-name: Publish Sliding Tile Attention Kernel to PyPI on Version Change
-
-on:
-  push:
-    branches:
-      - main
-    paths:
-      - "csrc/attn/sliding_tile_attn/setup.py"
-  workflow_dispatch:
-
-jobs:
-  check-version-change:
-    runs-on: ubuntu-latest
-    outputs:
-      version-changed: ${{ steps.check-version.outputs.changed }}
-      new-version: ${{ steps.check-version.outputs.new-version }}
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 2
-
-      - name: Check if version changed
-        id: check-version
-        run: |
-          cd csrc/attn/sliding_tile_attn
-          # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
-          echo "New version: $NEW_VERSION"
-          
-          # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
-          echo "Old version: $OLD_VERSION"
-          
-          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
-            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
-            echo "changed=true" >> $GITHUB_OUTPUT
-            echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT
-          else
-            echo "Version did not change"
-            echo "changed=false" >> $GITHUB_OUTPUT
-          fi
-
-  build_wheels:
-    name: Build Wheel
-    needs: check-version-change
-    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
-    runs-on: ${{ matrix.os }}
-
-    strategy:
-      fail-fast: false
-      matrix:
-          # Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the
-          # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
-          os: [ubuntu-22.04]
-          python-version: ['3.10', '3.11', '3.12', '3.13']
-          torch-version: ['2.5.1', '2.6.0']
-          cuda-version: ['12.4.1', '12.5.1', '12.6.3']
-
-    steps:
-      - name: Free up disk space
-        run: |
-          echo "Initial disk space:"
-          df -h
-          
-          # Remove large directories
-          sudo rm -rf /usr/share/dotnet
-          sudo rm -rf /usr/local/lib/android
-          sudo rm -rf /opt/ghc
-          sudo rm -rf /usr/local/share/boost
-          sudo rm -rf /usr/share/swift
-          sudo rm -rf /usr/local/lib/node_modules
-          sudo rm -rf /usr/local/share/powershell
-          sudo rm -rf /usr/share/rust
-          sudo rm -rf /usr/local/.ghcup
-          
-          # Remove cached files
-          sudo rm -rf /var/lib/apt/lists/*
-          sudo rm -rf /var/cache/apt/archives/*
-          
-          echo "Disk space after cleanup:"
-          df -h
-
-      - name: Checkout
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install CUDA ${{ matrix.cuda-version }}
-        uses: Jimver/cuda-toolkit@v0.2.21
-        id: cuda-toolkit
-        with:
-          cuda: ${{ matrix.cuda-version }}
-          linux-local-args: '["--toolkit"]'
-          method: 'network'
-
-      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
-        run: |
-          sudo apt update
-          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
-          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-          # Allow Git to Access Safe Directory
-          git config --global --add safe.directory /__w/FastVideo/FastVideo
-
-          # Set CUDA environment variables
-          export CUDA_HOME=/usr/local/cuda-${{ matrix.cuda-version }}
-          export PATH=${CUDA_HOME}/bin:${PATH}
-          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-          
-          # Verify installation
-          gcc --version
-          g++ --version
-          clang-11 --version
-          nvcc --version
-
-      - name: Install PyTorch ${{ matrix.torch-version }}+cu${{ matrix.cuda-version }}
-        run: |
-          pip install --upgrade pip
-          # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
-          # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
-          pip install typing-extensions==4.12.2
-          # We want to figure out the CUDA version to download pytorch
-          # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
-          # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
-          export TORCH_CUDA_VERSION=124
-          pip install --no-cache-dir torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
-          nvcc --version
-          python --version
-          python -c "import torch; print('PyTorch:', torch.__version__)"
-          python -c "import torch; print('CUDA:', torch.version.cuda)"
-          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
-
-      - name: Build wheel
-        run: |
-          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
-          
-          # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
-          # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
-          # However this still fails so I'm using a newer version of setuptools
-          pip install setuptools
-          pip install ninja packaging wheel
-
-          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
-          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py bdist_wheel --dist-dir=dist
-
-      - name: Rename wheel file
-        run: |
-          cd csrc/attn/sliding_tile_attn
-
-          CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
-          TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2)
-          # Get the correct version format
-          tmpname=cu${CUDA_SHORT_VERSION}torch${TORCH_SHORT_VERSION}
-          wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
-          # Rename with version information
-          ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
-          echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
-
-      - name: Upload wheel artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: ${{ env.wheel_name }}
-          path: csrc/attn/sliding_tile_attn/dist/*.whl
-          retention-days: 90
-
-  publish_package:
-    name: Publish package
-    needs: [build_wheels, check-version-change]
-    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
-    runs-on: ubuntu-22.04
-    permissions:
-      id-token: write  # Needed for OIDC Trusted Publishing
-
-    steps:
-      - uses: actions/checkout@v4
-
-      - uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-      
-      - name: Install CUDA 12.4.1
-        uses: Jimver/cuda-toolkit@v0.2.21
-        id: cuda-toolkit
-        with:
-          cuda: 12.4.1
-          linux-local-args: '["--toolkit"]'
-          method: 'network'
-          sub-packages: '["nvcc"]'
-
-      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
-        run: |
-          sudo apt update
-          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
-          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-          # Allow Git to Access Safe Directory
-          git config --global --add safe.directory /__w/FastVideo/FastVideo
-
-          # Set CUDA environment variables
-          export CUDA_HOME=/usr/local/cuda-12.4.1
-          export PATH=${CUDA_HOME}/bin:${PATH}
-          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-          
-          # Verify installation
-          gcc --version
-          g++ --version
-          clang-11 --version
-          nvcc --version
-
-      - name: Install PyTorch 2.5.1+cu12.4.1
-        run: |
-          pip install --upgrade pip
-          # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
-          # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
-          pip install typing-extensions==4.12.2
-          # We want to figure out the CUDA version to download pytorch
-          # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
-          # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
-          export TORCH_CUDA_VERSION=124
-          pip install --no-cache-dir torch==2.5.1 --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
-          nvcc --version
-          python --version
-          python -c "import torch; print('PyTorch:', torch.__version__)"
-          python -c "import torch; print('CUDA:', torch.version.cuda)"
-          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
-
-      - name: Build source distribution
-        run: |
-          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
-          
-          # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
-          # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
-          # However this still fails so I'm using a newer version of setuptools
-          pip install setuptools
-          pip install ninja packaging wheel
-
-          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
-          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py sdist --dist-dir=dist
-
-      - name: Publish release distributions to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          packages-dir: csrc/attn/sliding_tile_attn/dist/
@@ -54,7 +54,8 @@ To use this:
 1. **Set Context**: In your pipeline or generation loop, use the `set_forward_context` context manager.
 2. **Access Context**: Inside your attention backend, use `get_forward_context()`.

-See [`docs/attention/sta/index.md`](../sta/index.md) (Sliding Tile Attention) for an example of how complex configuration (window sizes) is passed this way.
+See [`docs/attention/sta/index.md`](../sta/index.md) for a legacy STA example
+of passing complex configuration (window sizes) through `ForwardContext`.

 ## 3. Adding Compiled Kernels (C++/CUDA)

@@ -5,13 +5,16 @@ FastVideo provides highly optimized custom attention kernels to accelerate video
 ## Supported Kernels

 * **[Video Sparse Attention (VSA)](vsa/index.md)**: Sparse attention mechanism selecting top-k blocks.
-* **[Sliding Tile Attention (STA)](sta/index.md)**: Optimized attention for window-based video generation.
+* **[Sliding Tile Attention (STA)](sta/index.md)**: STA kernel support is kept in
+  `fastvideo-kernel`; full FastVideo STA pipeline workflow is archived in
+  `sta_do_not_delete`.
 * **Backend development guide**: See the developer guide at
  [Attention Backend Development](../contributing/attention_backend.md).

 ## General Build Instructions

-These instructions apply to building the `fastvideo-kernel` package from source, which includes both STA and VSA kernels.
+These instructions apply to building the `fastvideo-kernel` package from
+source, which includes both STA and VSA kernels.

 ### Prerequisites

@@ -1,27 +1,84 @@
 # Sliding Tile Attention (STA)

-Optimized attention for window-based video generation (e.g., HunyuanVideo).
+STA inference integration is archived from `main`.

-## Installation
+The full STA pipeline code (including mask search and STA inference wiring in
+`fastvideo/`) is preserved in:

-STA is included in the `fastvideo-kernel` package. See the [main Attention page](../index.md) for build instructions.
+- https://github.com/hao-ai-lab/FastVideo/tree/sta_do_not_delete

-## Usage
+In this branch, STA kernels in `fastvideo-kernel` are still kept.

-```python
-from fastvideo_kernel import sliding_tile_attention
+## Why STA is not in `main`

-# q, k, v: [batch_size, num_heads, seq_length, head_dim]
-# window_size: List of (t, h, w) tiles. Tile size is (6, 8, 8).
-# text_length: Number of text tokens (0-256)
+We do not keep STA pipeline integration in `main` because we believe Video
+Sparse Attention (VSA) is strictly better than STA for the actively maintained
+FastVideo inference path.

-out = sliding_tile_attention(
-    q, k, v, 
-    window_size=[(3, 3, 3)],  # Example window
-    text_length=256
-)
+## What to checkout for STA workflows
+
+To run the full STA workflow, switch to the archived branch:
+
+```bash
+git fetch origin
+git checkout sta_do_not_delete
 ```

+## Mask Search (archive branch)
+
+The reference script is:
+
+- `examples/inference/sta_mask_search/inference_wan_sta.sh`
+
+It runs two stages:
+
+1. `STA_searching` (full search), output at
+   `inference_results/sta/mask_search_full`
+2. `STA_tuning` (sparse tuning), output at
+   `inference_results/sta/mask_search_sparse`
+
+Run:
+
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
+export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
+
+bash examples/inference/sta_mask_search/inference_wan_sta.sh
+```
+
+## STA Inference (archive branch)
+
+With a selected mask strategy, run inference with:
+
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
+export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
+
+fastvideo generate \
+  --model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \
+  --num-gpus 2 \
+  --tp-size 2 \
+  --sp-size 2 \
+  --height 768 \
+  --width 1280 \
+  --num-frames 69 \
+  --num-inference-steps 50 \
+  --prompt "A cinematic wildlife shot of a lion walking in golden grasslands." \
+  --output-path outputs_video/STA/
+```
+
+Python usage on the archive branch can also set `STA_mode` in
+`VideoGenerator.from_pretrained(...)`:
+
+- `STA_searching`
+- `STA_tuning`
+- `STA_inference`
+
+## Kernel-level API (current branch)
+
+STA kernels remain available from `fastvideo-kernel`. See
+[Attention overview](../index.md) for build instructions.
+
 ## Citation

 If you use Sliding Tile Attention in your research, please cite:
@@ -8,7 +8,7 @@ This guide explains how to add and run tests in FastVideo. The testing suite is
 * **Component Tests**: Located in `fastvideo/tests/encoders`, `fastvideo/tests/transformers`, and `fastvideo/tests/vaes`. These verify the loading and basic functionality of model components.
 * **SSIM Tests**: Located in `fastvideo/tests/ssim`. These are regression tests that compare generated videos against reference videos using the Structural Similarity Index Measure (SSIM) to detect quality degradation.
 * **Training Tests**: Located in `fastvideo/tests/training`. These validate training loops, loss calculations, and specific training techniques like LoRA, Distillation, and VSA.
-* **Inference Tests**: Located in `fastvideo/tests/inference`. These test specialized inference pipelines and optimizations (e.g., STA, V-MoBA).
+* **Inference Tests**: Located in `fastvideo/tests/inference`. These test specialized inference pipelines and optimizations (e.g., VSA, V-MoBA).

 For now, we will focus on **SSIM Tests**.

@@ -26,7 +26,7 @@ FastVideo is an inference and post-training framework for diffusion models. It f
 FastVideo has the following features:

 - State-of-the-art performance optimizations for inference
-  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [Sliding Tile Attention](attention/sta/index.md)
  - [Sage Attention](https://arxiv.org/abs/2410.02367)
 - E2E post-training support
  - Data preprocessing pipeline for video data
@@ -44,7 +44,7 @@ Use the navigation menu on the left to explore different sections:
 - **Inference**: Learn how to use FastVideo for video generation
 - **Training**: Data preprocessing and fine-tuning workflows
 - **Distillation**: Post-training optimization techniques
- **Sliding Tile Attention**: Advanced attention mechanisms
+- **Sliding Tile Attention**: Legacy workflow docs and kernel notes
 - **Video Sparse Attention**: Efficient attention for video models
 - **Design**: Framework architecture and design principles
 - **Developer Guide**: Contributing and development setup
@@ -123,7 +123,7 @@ self.attn = DistributedAttention(
    softmax_scale=None,
    causal=False,
    supported_attention_backends=(
-        AttentionBackendEnum.SLIDING_TILE_ATTN,
+        AttentionBackendEnum.VIDEO_SPARSE_ATTN,
        AttentionBackendEnum.FLASH_ATTN,
        AttentionBackendEnum.TORCH_SDPA,
    )
@@ -8,7 +8,7 @@ This page describes the various options for speeding up generation times in Fast
 - Optimized Attention Backends

  - [Flash Attention](#flash-attention)
-  - [Sliding Tile Attention](#sliding-tile-attention)
+  - [Sliding Tile Attention (Archived)](#sliding-tile-attention-archived)
  - [Sage Attention](#sage-attention)
  - [Sage Attention 3](#sage-attention-3)

@@ -18,13 +18,14 @@ This page describes the various options for speeding up generation times in Fast

 - 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`
 - Video MoBA Attention: `FASTVIDEO_ATTENTION_BACKEND=VMOBA_ATTN`
 - Sparse Linear Attention: `FASTVIDEO_ATTENTION_BACKEND=SLA_ATTN`
 - SageSLA Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_SLA_ATTN`
+- Sliding Tile Attention (archived branch only):
+  `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`

 ### Configuring Backends

@@ -35,7 +36,7 @@ There are two ways to configure the attention backend in FastVideo.
 In python, set the `FASTVIDEO_ATTENTION_BACKEND` environment variable before instantiating `VideoGenerator` like this:

 ```python
-os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
 ```

 #### 2. In CLI
@@ -66,12 +67,20 @@ pip install ninja
 python setup.py install
 ```

-### Sliding Tile Attention
+### Sliding Tile Attention (Archived)

 **`SLIDING_TILE_ATTN`**

-Sliding Tile Attention is provided by `fastvideo-kernel`.
-See [STA docs](../attention/sta/index.md) for installation details.
+The full STA integration in `fastvideo/` is archived from `main` and preserved
+at:
+
+- https://github.com/hao-ai-lab/FastVideo/tree/sta_do_not_delete
+
+We keep STA off `main` because we believe VSA is strictly better than STA for
+the actively maintained FastVideo path.
+
+Kernel code in `fastvideo-kernel` is still retained. For mask search and STA
+inference workflow, see [STA docs](../attention/sta/index.md).

 ### Video Sparse Attention

@@ -6,6 +6,13 @@ For the canonical, code-level list of model IDs recognized by
 `VideoGenerator.from_pretrained(...)`, see the registrations in
 `fastvideo/registry.py` (`register_configs(...)` entries).

+!!! note
+    The full STA integration in `fastvideo/` is archived from `main` and kept
+    in `sta_do_not_delete`:
+    https://github.com/hao-ai-lab/FastVideo/tree/sta_do_not_delete
+    We do this because we believe VSA is strictly better than STA for the
+    actively maintained `main` inference path.
+
 The symbols used have the following meanings:

 - ✅ = Full compatibility
@@ -46,29 +53,32 @@ pipeline initialization and sampling.
  }
 </style>

-| Model Name | HuggingFace Model ID | Resolutions | Sliding Tile Attn | Sage Attn | VSA | BSA |
-|------------|---------------------|-------------|-------------------|-----------|-----|-----|
-| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ✅ | ⭕ |
-| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ |
-| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ✅ | ⭕ | ⭕ |
-| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ✅ | ⭕ | ⭕ |
-| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ✅ | ⭕ | ⭕ |
-| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ✅ | ✅ | ⭕ | ⭕ |
-| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p<br>544px960p | ✅ | ✅ | ⭕ | ⭕ |
-| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅* | ✅ | ⭕ | ⭕ |
-| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅* | ✅ | ⭕ | ⭕ |
-| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅* | ✅ | ⭕ | ⭕ |
-| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720P | ✅ | ✅ | ⭕ | ⭕ |
-| TurboWan2.1 T2V 1.3B | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ |
-| TurboWan2.1 T2V 14B | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | 480P, 720P | ⭕ | ⭕ | ⭕ | ⭕ |
-| TurboWan2.2 I2V A14B | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ⭕ | ⭕ | ⭕ | ⭕ |
-| LongCat T2V 13.6B | See note** | 480P<br>720P | ❌ | ❌ | ⭕ | ✅ |
-| Matrix Game 2.0 Base | `FastVideo/Matrix-Game-2.0-Base-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ |
-| Matrix Game 2.0 GTA | `FastVideo/Matrix-Game-2.0-GTA-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ |
-| Matrix Game 2.0 TempleRun | `FastVideo/Matrix-Game-2.0-TempleRun-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ |
+| Model Name | HuggingFace Model ID | Resolutions | TeaCache | Sliding Tile Attn (Legacy Branch) | Sage Attn | VSA | BSA |
+|------------|---------------------|-------------|----------|-------------------|-----------|-----|-----|
+| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
+| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
+| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ |
+| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
+| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
+| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
+| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
+| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ | ⭕ |
+| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅ | ✅* | ✅ | ⭕ | ⭕ |
+| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ | ⭕ |
+| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
+| TurboWan2.1 T2V 1.3B | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
+| TurboWan2.1 T2V 14B | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | 480P, 720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
+| TurboWan2.2 I2V A14B | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
+| LongCat T2V 13.6B | See note** | 480P<br>720P | ❌ | ❌ | ❌ | ⭕ | ✅ |
+| Matrix Game 2.0 Base | `FastVideo/Matrix-Game-2.0-Base-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
+| Matrix Game 2.0 GTA | `FastVideo/Matrix-Game-2.0-GTA-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
+| Matrix Game 2.0 TempleRun | `FastVideo/Matrix-Game-2.0-TempleRun-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |

 **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.

+`Sliding Tile Attn (Legacy Branch)` entries refer to the archived
+`sta_do_not_delete` branch workflow, not active `main` inference wiring.
+
 ## Canonical Supported IDs

 The authoritative source for model-ID recognition is
@@ -78,7 +88,9 @@ resolve default pipeline and sampling configuration for it.
 ## Special requirements

 ### Sliding Tile Attention
- Currently only Hopper GPUs (H100s) are supported.
+- Full STA pipeline usage is on the archived branch:
+  https://github.com/hao-ai-lab/FastVideo/tree/sta_do_not_delete
+- STA currently requires Hopper GPUs (H100s).

 ### TurboWan2.1 (TurboDiffusion)
 - Uses TurboDiffusionPipeline with RCM scheduler for 1-4 step generation
@@ -49,7 +49,9 @@ combinations.
 If forcing a backend fails, verify optional dependencies are installed:

 - `FLASH_ATTN`: `flash-attn`
- `SLIDING_TILE_ATTN` and `VIDEO_SPARSE_ATTN`: `fastvideo-kernel`
+- `VIDEO_SPARSE_ATTN`: `fastvideo-kernel`
+- `SLIDING_TILE_ATTN`: STA legacy workflow in
+  `sta_do_not_delete` + `fastvideo-kernel`
 - `SAGE_ATTN` / `SAGE_ATTN_THREE`: SageAttention packages

 As a fallback, use:
@@ -1,5 +1,44 @@
-# STA Mask Search Examples
+# STA Mask Search (Archived Workflow)
+
+The full STA integration is preserved in:
+
+- https://github.com/hao-ai-lab/FastVideo/tree/sta_do_not_delete
+
+Switch to the branch before running mask search:

 ```bash
+git fetch origin
+git checkout sta_do_not_delete
+```
+
+Run mask search + tuning:
+
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
+export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
+
 bash examples/inference/sta_mask_search/inference_wan_sta.sh
-```
+```
+
+This script runs:
+
+- `STA_searching` and writes to `inference_results/sta/mask_search_full`
+- `STA_tuning` and writes to `inference_results/sta/mask_search_sparse`
+
+Example STA inference (same archived branch):
+
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
+export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
+
+fastvideo generate \
+  --model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \
+  --num-gpus 2 \
+  --tp-size 2 \
+  --sp-size 2 \
+  --height 768 \
+  --width 1280 \
+  --num-frames 69 \
+  --num-inference-steps 50 \
+  --prompt "A cinematic wildlife shot of a lion walking in golden grasslands."
+```
@@ -1,39 +0,0 @@
-#!/bin/bash
-
-export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
-export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
-export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers
-
-base_port=29503
-num_gpu=1
-gpu_ids=$(seq 0 $((num_gpu-1)))
-skip_time_steps=12
-
-output_path="inference_results/sta/mask_search_full"
-STA_mode="STA_searching"
-for i in $gpu_ids; do
-    port=$((base_port+i))
-    CUDA_VISIBLE_DEVICES=$i MASTER_PORT=$port python examples/inference/sta_mask_search/wan_example.py \
-        --prompt_path ./assets/prompt_${i}.txt \
-        --output_path $output_path \
-        --STA_mode $STA_mode &
-    sleep 1
-done
-wait
-echo "STA searching completed"
-
-output_path="inference_results/sta/mask_search_sparse"
-STA_mode="STA_tuning"
-for i in $gpu_ids; do
-    port=$((base_port+i))
-    CUDA_VISIBLE_DEVICES=$i MASTER_PORT=$port python examples/inference/sta_mask_search/wan_example.py \
-        --prompt_path ./assets/prompt_${i}.txt \
-        --output_path $output_path \
-        --STA_mode $STA_mode \
-        --skip_time_steps $skip_time_steps &
-    sleep 1
-done
-wait
-echo "STA tuning completed"
-
-echo "All jobs completed"
@@ -1,63 +0,0 @@
-import os
-import argparse
-from fastvideo import VideoGenerator, SamplingParam
-
-def main(args):
-    os.makedirs(args.output_path, exist_ok=True)
-    # Create a video generator with a pre-trained model
-    generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-14B-Diffusers",
-        num_gpus=args.num_gpus,  # Adjust based on your hardware
-        STA_mode=args.STA_mode,
-        skip_time_steps=args.skip_time_steps
-    )
-
-    # Prompts for your video
-    prompt = args.prompt
-    prompt_path = args.prompt_path
-    negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
-
-    if prompt_path is not None:
-        with open(prompt_path, "r") as f:
-            prompts = f.readlines()
-    else:
-        prompts = [prompt]
-
-    params = SamplingParam(
-        height=args.height,
-        width=args.width,
-        num_frames=args.num_frames,
-        num_inference_steps=args.num_inference_steps,
-        fps=args.fps,
-        guidance_scale=args.guidance_scale,
-        seed=args.seed,
-        return_frames=True,  # Also return frames from this call (defaults to False)
-        output_path=args.output_path,  # Controls where videos are saved
-        save_video=True,
-        negative_prompt=negative_prompt
-    )
-
-    # Generate the video
-    for prompt in prompts:
-        video = generator.generate_video(
-            prompt,
-            sampling_param=params,
-        )
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--prompt", type=str, default="A man is dancing.")
-    parser.add_argument("--prompt_path", type=str, default=None)
-    parser.add_argument("--height", type=int, default=768)
-    parser.add_argument("--width", type=int, default=1280)
-    parser.add_argument("--num_frames", type=int, default=69)
-    parser.add_argument("--num_inference_steps", type=int, default=50)
-    parser.add_argument("--fps", type=int, default=16)
-    parser.add_argument("--guidance_scale", type=float, default=5.0)
-    parser.add_argument("--seed", type=int, default=12345)
-    parser.add_argument("--output_path", type=str, default="my_videos/")
-    parser.add_argument("--num_gpus", type=int, default=1)
-    parser.add_argument("--STA_mode", type=str, default="STA_searching")
-    parser.add_argument("--skip_time_steps", type=int, default=12)
-    args = parser.parse_args()
-    main(args)
@@ -1,405 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import json
-import os
-from collections import defaultdict
-from typing import Any
-
-import numpy as np
-
-from fastvideo.utils import dict_to_3d_list
-
-
-def configure_sta(mode: str = 'STA_searching',
-                  layer_num: int = 40,
-                  time_step_num: int = 50,
-                  head_num: int = 40,
-                  **kwargs) -> list[list[list[Any]]]:
-    """
-    Configure Sliding Tile Attention (STA) parameters based on the specified mode.
-    
-    Parameters:
-    ----------
-    mode : str
-        The STA mode to use. Options are:
-        - 'STA_searching': Generate a set of mask candidates for initial search
-        - 'STA_tuning': Select best mask strategy based on previously saved results
-        - 'STA_inference': Load and use a previously tuned mask strategy
-    layer_num: int, number of layers
-    time_step_num: int, number of timesteps
-    head_num: int, number of heads
-    
-    **kwargs : dict
-        Mode-specific parameters:
-        
-        For 'STA_searching':
-        - mask_candidates: list of str, optional, mask candidates to use
-        - mask_selected: list of int, optional, indices of selected masks
-        
-        For 'STA_tuning':
-        - mask_search_files_path: str, required, path to mask search results
-        - mask_candidates: list of str, optional, mask candidates to use
-        - mask_selected: list of int, optional, indices of selected masks
-        - skip_time_steps: int, optional, number of time steps to use full attention (default 12)
-        - save_dir: str, optional, directory to save mask strategy (default "mask_candidates")
-        
-        For 'STA_inference':
-        - load_path: str, optional, path to load mask strategy (default "mask_candidates/mask_strategy.json")
-    """
-    valid_modes = [
-        'STA_searching', 'STA_tuning', 'STA_inference', 'STA_tuning_cfg'
-    ]
-    if mode not in valid_modes:
-        raise ValueError(f"Mode must be one of {valid_modes}, got {mode}")
-
-    if mode == 'STA_searching':
-        # Get parameters with defaults
-        mask_candidates: list[str] | None = kwargs.get('mask_candidates')
-        if mask_candidates is None:
-            raise ValueError(
-                "mask_candidates is required for STA_searching mode")
-        mask_selected: list[int] = kwargs.get('mask_selected',
-                                              list(range(len(mask_candidates))))
-
-        # Parse selected masks
-        selected_masks: list[list[int]] = []
-        for index in mask_selected:
-            mask = mask_candidates[index]
-            masks_list = [int(x) for x in mask.split(',')]
-            selected_masks.append(masks_list)
-
-        # Create 3D mask structure with fixed dimensions (t=50, l=60)
-        masks_3d: list[list[list[list[int]]]] = []
-        for i in range(time_step_num):  # Fixed t dimension = 50
-            row = []
-            for j in range(layer_num):  # Fixed l dimension = 60
-                row.append(selected_masks)  # Add all masks at each position
-            masks_3d.append(row)
-
-        return masks_3d
-
-    elif mode == 'STA_tuning':
-        # Get required parameters
-        mask_search_files_path: str | None = kwargs.get(
-            'mask_search_files_path')
-        if not mask_search_files_path:
-            raise ValueError(
-                "mask_search_files_path is required for STA_tuning mode")
-
-        # Get optional parameters with defaults
-        mask_candidates_tuning: list[str] | None = kwargs.get('mask_candidates')
-        if mask_candidates_tuning is None:
-            raise ValueError("mask_candidates is required for STA_tuning mode")
-        mask_selected_tuning: list[int] = kwargs.get(
-            'mask_selected', list(range(len(mask_candidates_tuning))))
-        skip_time_steps_tuning: int | None = kwargs.get('skip_time_steps')
-        save_dir_tuning: str | None = kwargs.get('save_dir', "mask_candidates")
-
-        # Parse selected masks
-        selected_masks_tuning: list[list[int]] = []
-        for index in mask_selected_tuning:
-            mask = mask_candidates_tuning[index]
-            masks_list = [int(x) for x in mask.split(',')]
-            selected_masks_tuning.append(masks_list)
-
-        # Read JSON results
-        results = read_specific_json_files(mask_search_files_path)
-        averaged_results = average_head_losses(results, selected_masks_tuning)
-
-        # Add full attention mask for specific cases
-        full_attention_mask_tuning: list[int] | None = kwargs.get(
-            'full_attention_mask')
-        if full_attention_mask_tuning is not None:
-            selected_masks_tuning.append(full_attention_mask_tuning)
-
-        # Select best mask strategy
-        timesteps_tuning: int = kwargs.get('timesteps', time_step_num)
-        if skip_time_steps_tuning is None:
-            skip_time_steps_tuning = 12
-        mask_strategy, sparsity, strategy_counts = select_best_mask_strategy(
-            averaged_results, selected_masks_tuning, skip_time_steps_tuning,
-            timesteps_tuning, head_num)
-
-        # Save mask strategy
-        if save_dir_tuning is not None:
-            os.makedirs(save_dir_tuning, exist_ok=True)
-            file_path = os.path.join(
-                save_dir_tuning,
-                f'mask_strategy_s{skip_time_steps_tuning}.json')
-            with open(file_path, 'w') as f:
-                json.dump(mask_strategy, f, indent=4)
-            print(f"Successfully saved mask_strategy to {file_path}")
-
-        # Print sparsity and strategy counts for information
-        print(f"Overall sparsity: {sparsity:.4f}")
-        print("\nStrategy usage counts:")
-        total_heads = time_step_num * layer_num * head_num  # Fixed dimensions
-        for strategy, count in strategy_counts.items():
-            print(
-                f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)"
-            )
-
-        # Convert dictionary to 3D list with fixed dimensions
-        mask_strategy_3d = dict_to_3d_list(mask_strategy,
-                                           t_max=time_step_num,
-                                           l_max=layer_num,
-                                           h_max=head_num)
-
-        return mask_strategy_3d
-    elif mode == 'STA_tuning_cfg':
-        # Get required parameters for both positive and negative paths
-        mask_search_files_path_pos: str | None = kwargs.get(
-            'mask_search_files_path_pos')
-        mask_search_files_path_neg: str | None = kwargs.get(
-            'mask_search_files_path_neg')
-        save_dir_cfg: str | None = kwargs.get('save_dir')
-
-        if not mask_search_files_path_pos or not mask_search_files_path_neg or not save_dir_cfg:
-            raise ValueError(
-                "mask_search_files_path_pos, mask_search_files_path_neg, and save_dir are required for STA_tuning_cfg mode"
-            )
-
-        # Get optional parameters with defaults
-        mask_candidates_cfg: list[str] | None = kwargs.get('mask_candidates')
-        if mask_candidates_cfg is None:
-            raise ValueError(
-                "mask_candidates is required for STA_tuning_cfg mode")
-        mask_selected_cfg: list[int] = kwargs.get(
-            'mask_selected', list(range(len(mask_candidates_cfg))))
-        skip_time_steps_cfg: int | None = kwargs.get('skip_time_steps')
-
-        # Parse selected masks
-        selected_masks_cfg: list[list[int]] = []
-        for index in mask_selected_cfg:
-            mask = mask_candidates_cfg[index]
-            masks_list = [int(x) for x in mask.split(',')]
-            selected_masks_cfg.append(masks_list)
-
-        # Read JSON results for both positive and negative paths
-        pos_results = read_specific_json_files(mask_search_files_path_pos)
-        neg_results = read_specific_json_files(mask_search_files_path_neg)
-        # Combine positive and negative results into one list
-        combined_results = pos_results + neg_results
-
-        # Average the combined results
-        averaged_results = average_head_losses(combined_results,
-                                               selected_masks_cfg)
-
-        # Add full attention mask for specific cases
-        full_attention_mask_cfg: list[int] | None = kwargs.get(
-            'full_attention_mask')
-        if full_attention_mask_cfg is not None:
-            selected_masks_cfg.append(full_attention_mask_cfg)
-
-        timesteps_cfg: int = kwargs.get('timesteps', time_step_num)
-        if skip_time_steps_cfg is None:
-            skip_time_steps_cfg = 12
-        # Select best mask strategy using combined results
-        mask_strategy, sparsity, strategy_counts = select_best_mask_strategy(
-            averaged_results, selected_masks_cfg, skip_time_steps_cfg,
-            timesteps_cfg, head_num)
-
-        # Save mask strategy
-        os.makedirs(save_dir_cfg, exist_ok=True)
-        file_path = os.path.join(save_dir_cfg,
-                                 f'mask_strategy_s{skip_time_steps_cfg}.json')
-        with open(file_path, 'w') as f:
-            json.dump(mask_strategy, f, indent=4)
-        print(f"Successfully saved mask_strategy to {file_path}")
-
-        # Print sparsity and strategy counts for information
-        print(f"Overall sparsity: {sparsity:.4f}")
-        print("\nStrategy usage counts:")
-        total_heads = time_step_num * layer_num * head_num  # Fixed dimensions
-        for strategy, count in strategy_counts.items():
-            print(
-                f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)"
-            )
-
-        # Convert dictionary to 3D list with fixed dimensions
-        mask_strategy_3d = dict_to_3d_list(mask_strategy,
-                                           t_max=time_step_num,
-                                           l_max=layer_num,
-                                           h_max=head_num)
-
-        return mask_strategy_3d
-
-    else:  # STA_inference
-        # Get parameters with defaults
-        load_path: str | None = kwargs.get(
-            'load_path', "mask_candidates/mask_strategy.json")
-        if load_path is None:
-            raise ValueError("load_path is required for STA_inference mode")
-
-        # Load previously saved mask strategy
-        with open(load_path) as f:
-            mask_strategy = json.load(f)
-
-        # Convert dictionary to 3D list with fixed dimensions
-        mask_strategy_3d = dict_to_3d_list(mask_strategy,
-                                           t_max=time_step_num,
-                                           l_max=layer_num,
-                                           h_max=head_num)
-
-        return mask_strategy_3d
-
-
-# Helper functions
-
-
-def read_specific_json_files(folder_path: str) -> list[dict[str, Any]]:
-    """Read and parse JSON files containing mask search results."""
-    json_contents: list[dict[str, Any]] = []
-
-    # List files only in the current directory (no walk)
-    files = os.listdir(folder_path)
-    # Filter files
-    matching_files = [f for f in files if 'mask' in f and f.endswith('.json')]
-    print(f"Found {len(matching_files)} matching files: {matching_files}")
-
-    for file_name in matching_files:
-        file_path = os.path.join(folder_path, file_name)
-        with open(file_path) as file:
-            data = json.load(file)
-            json_contents.append(data)
-
-    return json_contents
-
-
-def average_head_losses(
-        results: list[dict[str, Any]],
-        selected_masks: list[list[int]]) -> dict[str, dict[str, np.ndarray]]:
-    """Average losses across all prompts for each mask strategy."""
-    # Initialize a dictionary to store the averaged results
-    averaged_losses: dict[str, dict[str, np.ndarray]] = {}
-    loss_type = 'L2_loss'
-    # Get all loss types (e.g., 'L2_loss')
-    averaged_losses[loss_type] = {}
-
-    for mask in selected_masks:
-        mask_str = str(mask)
-        data_shape = np.array(results[0][loss_type][mask_str]).shape
-        accumulated_data = np.zeros(data_shape)
-
-        # Sum across all prompts
-        for prompt_result in results:
-            accumulated_data += np.array(prompt_result[loss_type][mask_str])
-
-        # Average by dividing by number of prompts
-        averaged_data = accumulated_data / len(results)
-        averaged_losses[loss_type][mask_str] = averaged_data
-
-    return averaged_losses
-
-
-def select_best_mask_strategy(
-        averaged_results: dict[str, dict[str, np.ndarray]],
-        selected_masks: list[list[int]],
-        skip_time_steps: int = 12,
-        timesteps: int = 50,
-        head_num: int = 40
-) -> tuple[dict[str, list[int]], float, dict[str, int]]:
-    """Select the best mask strategy for each head based on loss minimization."""
-    best_mask_strategy: dict[str, list[int]] = {}
-    loss_type = 'L2_loss'
-    # Get the shape of time steps and layers
-    layers = len(averaged_results[loss_type][str(selected_masks[0])][0])
-
-    # Counter for sparsity calculation
-    total_tokens = 0  # total number of masked tokens
-    total_length = 0  # total sequence length
-
-    strategy_counts: dict[str, int] = {
-        str(strategy): 0
-        for strategy in selected_masks
-    }
-    full_attn_strategy = selected_masks[-1]  # Last strategy is full attention
-    print(f"Strategy {full_attn_strategy}, skip first {skip_time_steps} steps ")
-
-    for t in range(timesteps):
-        for layer_idx in range(layers):
-            for h in range(head_num):
-                if t < skip_time_steps:  # First steps use full attention
-                    strategy = full_attn_strategy
-                else:
-                    # Get losses for this head across all strategies
-                    head_losses = []
-                    for strategy in selected_masks[:
-                                                   -1]:  # Exclude full attention
-                        head_losses.append(averaged_results[loss_type][str(
-                            strategy)][t][layer_idx][h])
-
-                    # Find which strategy gives minimum loss
-                    best_strategy_idx = np.argmin(head_losses)
-                    strategy = selected_masks[best_strategy_idx]
-
-                best_mask_strategy[f'{t}_{layer_idx}_{h}'] = strategy
-
-                # Calculate sparsity
-                nums = strategy  # strategy is already a list of numbers
-                total_tokens += nums[0] * nums[1] * nums[
-                    2]  # masked tokens for chosen strategy
-                total_length += full_attn_strategy[0] * full_attn_strategy[
-                    1] * full_attn_strategy[2]
-
-                # Count strategy usage
-                strategy_counts[str(strategy)] += 1
-
-    overall_sparsity = 1 - total_tokens / total_length
-
-    return best_mask_strategy, overall_sparsity, strategy_counts
-
-
-def save_mask_search_results(
-        mask_search_final_result: list[dict[str, list[float]]],
-        prompt: str,
-        mask_strategies: list[str],
-        output_dir: str = 'output/mask_search_result/') -> str | None:
-    if not mask_search_final_result:
-        print("No mask search results to save")
-        return None
-
-    # Create result dictionary with defaultdict for nested lists
-    mask_search_dict: dict[str, dict[str, list[list[float]]]] = {
-        "L2_loss": defaultdict(list),
-        "L1_loss": defaultdict(list)
-    }
-
-    mask_selected = list(range(len(mask_strategies)))
-    selected_masks: list[list[int]] = []
-    for index in mask_selected:
-        mask = mask_strategies[index]
-        masks_list = [int(x) for x in mask.split(',')]
-        selected_masks.append(masks_list)
-
-    # Process each mask strategy
-    for i, mask_strategy in enumerate(selected_masks):
-        mask_strategy_str = str(mask_strategy)
-        # Process L2 loss
-        step_results: list[list[float]] = []
-        for step_data in mask_search_final_result:
-            if isinstance(step_data, dict) and "L2_loss" in step_data:
-                layer_losses = [float(loss) for loss in step_data["L2_loss"]]
-                step_results.append(layer_losses)
-        mask_search_dict["L2_loss"][mask_strategy_str] = step_results
-
-        step_results = []
-        for step_data in mask_search_final_result:
-            if isinstance(step_data, dict) and "L1_loss" in step_data:
-                layer_losses = [float(loss) for loss in step_data["L1_loss"]]
-                step_results.append(layer_losses)
-        mask_search_dict["L1_loss"][mask_strategy_str] = step_results
-
-    # Create the output directory if it doesn't exist
-    os.makedirs(output_dir, exist_ok=True)
-
-    # Create a filename based on the first 20 characters of the prompt
-    filename = prompt[:50].replace(" ", "_")
-    filepath = os.path.join(output_dir, f'mask_search_{filename}.json')
-
-    # Save the results to a JSON file
-    with open(filepath, 'w') as f:
-        json.dump(mask_search_dict, f, indent=4)
-
-    print(f"Successfully saved mask research results to {filepath}")
-
-    return filepath
@@ -1,263 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import json
-from dataclasses import dataclass
-from typing import Any
-
-import torch
-from einops import rearrange
-from fastvideo_kernel import sliding_tile_attention
-
-import fastvideo.envs as envs
-from fastvideo.attention.backends.abstract import (AttentionBackend,
-                                                   AttentionImpl,
-                                                   AttentionMetadata,
-                                                   AttentionMetadataBuilder)
-from fastvideo.distributed import get_sp_group
-from fastvideo.forward_context import ForwardContext, get_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.utils import dict_to_3d_list
-
-logger = init_logger(__name__)
-
-
-class RangeDict(dict):
-
-    def __getitem__(self, item: int) -> str:
-        for key in self.keys():
-            if isinstance(key, tuple):
-                low, high = key
-                if low <= item <= high:
-                    return str(super().__getitem__(key))
-            elif key == item:
-                return str(super().__getitem__(key))
-        raise KeyError(f"seq_len {item} not supported for STA")
-
-
-class SlidingTileAttentionBackend(AttentionBackend):
-
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        # TODO(will-refactor): check this
-        return [32, 64, 96, 128, 160, 192, 224, 256]
-
-    @staticmethod
-    def get_name() -> str:
-        return "SLIDING_TILE_ATTN"
-
-    @staticmethod
-    def get_impl_cls() -> type["SlidingTileAttentionImpl"]:
-        return SlidingTileAttentionImpl
-
-    @staticmethod
-    def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]:
-        return SlidingTileAttentionMetadata
-
-    @staticmethod
-    def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]:
-        return SlidingTileAttentionMetadataBuilder
-
-
-@dataclass
-class SlidingTileAttentionMetadata(AttentionMetadata):
-    current_timestep: int
-    STA_param: list[list[
-        Any]]  # each timestep with one metadata, shape [num_layers, num_heads]
-
-
-class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder):
-
-    def __init__(self):
-        pass
-
-    def prepare(self):
-        pass
-
-    def build(  # type: ignore
-        self,
-        STA_param: list[list[Any]],
-        current_timestep: int,
-        **kwargs: dict[str, Any],
-    ) -> SlidingTileAttentionMetadata:
-        param = STA_param
-        if param is None:
-            return SlidingTileAttentionMetadata(
-                current_timestep=current_timestep, STA_param=[])
-        return SlidingTileAttentionMetadata(current_timestep=current_timestep,
-                                            STA_param=param[current_timestep])
-
-
-class SlidingTileAttentionImpl(AttentionImpl):
-
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        causal: bool,
-        softmax_scale: float,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        # TODO(will-refactor): for now this is the mask strategy, but maybe we should
-        # have a more general config for STA?
-        config_file = envs.FASTVIDEO_ATTENTION_CONFIG
-        if config_file is None:
-            raise ValueError("FASTVIDEO_ATTENTION_CONFIG is not set")
-
-        # TODO(kevin): get mask strategy for different STA modes
-        with open(config_file) as f:
-            mask_strategy = json.load(f)
-        self.mask_strategy = dict_to_3d_list(mask_strategy)
-
-        self.prefix = prefix
-        sp_group = get_sp_group()
-        self.sp_size = sp_group.world_size
-        # STA config
-        self.STA_base_tile_size = [6, 8, 8]
-        self.dit_seq_shape_mapping = RangeDict({
-            (115200, 115456): '30x48x80',
-            82944: '36x48x48',
-            69120: '18x48x80',
-        })
-        self.full_window_mapping = {
-            '30x48x80': [5, 6, 10],
-            '36x48x48': [6, 6, 6],
-            '18x48x80': [3, 6, 10]
-        }
-
-    def tile(self, x: torch.Tensor) -> torch.Tensor:
-        return rearrange(
-            x,
-            "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-            n_t=self.full_window_size[0],
-            n_h=self.full_window_size[1],
-            n_w=self.full_window_size[2],
-            ts_t=self.STA_base_tile_size[0],
-            ts_h=self.STA_base_tile_size[1],
-            ts_w=self.STA_base_tile_size[2])
-
-    def untile(self, x: torch.Tensor) -> torch.Tensor:
-        x = rearrange(
-            x,
-            "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-            n_t=self.full_window_size[0],
-            n_h=self.full_window_size[1],
-            n_w=self.full_window_size[2],
-            ts_t=self.STA_base_tile_size[0],
-            ts_h=self.STA_base_tile_size[1],
-            ts_w=self.STA_base_tile_size[2])
-        return x
-
-    def preprocess_qkv(
-        self,
-        qkv: torch.Tensor,
-        attn_metadata: AttentionMetadata,
-    ) -> torch.Tensor:
-        img_sequence_length = qkv.shape[1]
-        self.dit_seq_shape_str = self.dit_seq_shape_mapping[img_sequence_length]
-        self.full_window_size = self.full_window_mapping[self.dit_seq_shape_str]
-        self.dit_seq_shape_int = list(
-            map(int, self.dit_seq_shape_str.split('x')))
-        self.img_seq_length = self.dit_seq_shape_int[
-            0] * self.dit_seq_shape_int[1] * self.dit_seq_shape_int[2]
-        return self.tile(qkv)
-
-    def postprocess_output(
-        self,
-        output: torch.Tensor,
-        attn_metadata: SlidingTileAttentionMetadata,
-    ) -> torch.Tensor:
-        return self.untile(output)
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        attn_metadata: SlidingTileAttentionMetadata,
-    ) -> torch.Tensor:
-        if self.mask_strategy is None:
-            raise ValueError(
-                "mask_strategy cannot be None for SlidingTileAttention")
-        if self.mask_strategy[0] is None:
-            raise ValueError(
-                "mask_strategy[0] cannot be None for SlidingTileAttention")
-
-        timestep = attn_metadata.current_timestep
-        forward_context: ForwardContext = get_forward_context()
-        forward_batch = forward_context.forward_batch
-        if forward_batch is None:
-            raise ValueError("forward_batch cannot be None")
-        # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl'
-        layer_idx = int(self.prefix.split('.')[-3])
-        if attn_metadata.STA_param is None or len(
-                attn_metadata.STA_param) <= layer_idx:
-            raise ValueError("Invalid STA_param")
-        STA_param = attn_metadata.STA_param[layer_idx]
-
-        text_length = q.shape[1] - self.img_seq_length
-        has_text = text_length > 0
-
-        query = q.transpose(1, 2).contiguous()
-        key = k.transpose(1, 2).contiguous()
-        value = v.transpose(1, 2).contiguous()
-
-        head_num = query.size(1)
-        sp_group = get_sp_group()
-        current_rank = sp_group.rank_in_group
-        start_head = current_rank * head_num
-
-        # searching or tuning mode
-        if len(STA_param) < head_num * sp_group.world_size:
-            sparse_attn_hidden_states_all = []
-            full_mask_window = STA_param[-1]
-            for window_size in STA_param[:-1]:
-                sparse_hidden_states = sliding_tile_attention(
-                    query, key, value, [window_size] * head_num, text_length,
-                    has_text, self.dit_seq_shape_str).transpose(1, 2)
-                sparse_attn_hidden_states_all.append(sparse_hidden_states)
-
-            hidden_states = sliding_tile_attention(
-                query, key, value, [full_mask_window] * head_num, text_length,
-                has_text, self.dit_seq_shape_str).transpose(1, 2)
-
-            attn_L2_loss = []
-            attn_L1_loss = []
-            # average loss across all heads
-            for sparse_attn_hidden_states in sparse_attn_hidden_states_all:
-                # L2 loss
-                attn_L2_loss_ = torch.mean((sparse_attn_hidden_states.float() -
-                                            hidden_states.float())**2,
-                                           dim=[0, 1, 3]).cpu().numpy()
-                attn_L2_loss_ = [round(float(x), 6) for x in attn_L2_loss_]
-                attn_L2_loss.append(attn_L2_loss_)
-                # L1 loss
-                attn_L1_loss_ = torch.mean(
-                    torch.abs(sparse_attn_hidden_states.float() -
-                              hidden_states.float()),
-                    dim=[0, 1, 3]).cpu().numpy()
-                attn_L1_loss_ = [round(float(x), 6) for x in attn_L1_loss_]
-                attn_L1_loss.append(attn_L1_loss_)
-
-            layer_loss_save = {"L2_loss": attn_L2_loss, "L1_loss": attn_L1_loss}
-
-            if forward_batch.is_cfg_negative:
-                if forward_batch.mask_search_final_result_neg is not None:
-                    forward_batch.mask_search_final_result_neg[timestep].append(
-                        layer_loss_save)
-            else:
-                if forward_batch.mask_search_final_result_pos is not None:
-                    forward_batch.mask_search_final_result_pos[timestep].append(
-                        layer_loss_save)
-        else:
-            windows = [
-                STA_param[head_idx + start_head] for head_idx in range(head_num)
-            ]
-
-            hidden_states = sliding_tile_attention(
-                query, key, value, windows, text_length, has_text,
-                self.dit_seq_shape_str).transpose(1, 2)
-
-        return hidden_states
@@ -43,6 +43,5 @@
      "require_post_norm": null
    }
  ],
-  "mask_strategy_file_path": null,
  "enable_torch_compile": false
-}
+}
@@ -24,7 +24,6 @@ class ModelConfig:
    arch_config: ArchConfig = field(default_factory=ArchConfig)

    # FastVideo-specific parameters here
-    # i.e. STA, quantization

    def __getattr__(self, name):
        # Only called if 'name' is not found in ModelConfig directly
@@ -15,11 +15,10 @@ class DiTArchConfig(ArchConfig):
    reverse_param_names_mapping: dict = field(default_factory=dict)
    lora_param_names_mapping: dict = field(default_factory=dict)
    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
-        AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
-        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
-        AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN,
-        AttentionBackendEnum.SAGE_ATTN_THREE, AttentionBackendEnum.SLA_ATTN,
-        AttentionBackendEnum.SAGE_SLA_ATTN)
+        AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+        AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE,
+        AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN)

    hidden_size: int = 0
    num_attention_heads: int = 0
@@ -1,5 +1,4 @@
-from fastvideo.configs.pipelines.base import (PipelineConfig,
-                                              SlidingTileAttnConfig)
+from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
 from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
@@ -15,8 +14,8 @@ from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig,
 __all__ = [
    "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig",
    "PipelineConfig", "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig",
-    "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
-    "WanT2V720PConfig", "WanI2V720PConfig", "SelfForcingWanT2V480PConfig",
-    "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
+    "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig",
+    "WanI2V720PConfig", "SelfForcingWanT2V480PConfig", "CosmosConfig",
+    "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
    "get_pipeline_config_cls_from_name"
 ]
@@ -2,7 +2,6 @@
 import json
 from collections.abc import Callable
 from dataclasses import asdict, dataclass, field, fields
-from enum import Enum
 from typing import Any, cast

 import torch
@@ -17,15 +16,6 @@ from fastvideo.utils import FlexibleArgumentParser, StoreBoolean, shallow_asdict
 logger = init_logger(__name__)


-class STA_Mode(str, Enum):
-    """STA (Sliding Tile Attention) modes."""
-    STA_INFERENCE = "STA_inference"
-    STA_SEARCHING = "STA_searching"
-    STA_TUNING = "STA_tuning"
-    STA_TUNING_CFG = "STA_tuning_cfg"
-    NONE = None
-
-
 def preprocess_text(prompt: str) -> str:
    return prompt

@@ -76,11 +66,6 @@ class PipelineConfig:
                                  ...] = field(default_factory=lambda:
                                               (postprocess_text, ))

-    # STA (Sliding Tile Attention) parameters
-    mask_strategy_file_path: str | None = None
-    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
-    skip_time_steps: int = 15
-
    # DMD parameters
    dmd_denoising_steps: list[int] | None = field(default=None)

@@ -191,24 +176,6 @@ class PipelineConfig:
            "Comma-separated list of denoising steps (e.g., '1000,757,522')",
        )

-        # STA (Sliding Tile Attention) parameters
-        parser.add_argument(
-            f"--{prefix_with_dot}STA-mode",
-            type=str,
-            dest=f"{prefix_with_dot.replace('-', '_')}STA_mode",
-            default=PipelineConfig.STA_mode.value,
-            choices=[mode.value for mode in STA_Mode],
-            help=
-            "STA mode: STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None",
-        )
-        parser.add_argument(
-            f"--{prefix_with_dot}skip-time-steps",
-            type=int,
-            dest=f"{prefix_with_dot.replace('-', '_')}skip_time_steps",
-            default=PipelineConfig.skip_time_steps,
-            help="Number of time steps to warmup (full attention) for STA",
-        )
-
        # Add VAE configuration arguments
        from fastvideo.configs.models.vaes.base import VAEConfig
        VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
@@ -290,11 +257,6 @@ class PipelineConfig:
        kwargs[prefix_with_dot + 'model_path'] = model_path
        pipeline_config.update_config_from_dict(kwargs, config_cli_prefix)

-        # Convert STA_mode string to enum if necessary
-        if isinstance(pipeline_config.STA_mode, str) and not isinstance(
-                pipeline_config.STA_mode, STA_Mode):
-            pipeline_config.STA_mode = STA_Mode(pipeline_config.STA_mode)
-
        return pipeline_config

    def check_pipeline_config(self) -> None:
@@ -378,24 +340,6 @@ class PipelineConfig:
            self.__post_init__()


-@dataclass
-class SlidingTileAttnConfig(PipelineConfig):
-    """Configuration for sliding tile attention."""
-
-    # Override any BaseConfig defaults as needed
-    # Add sliding tile specific parameters
-    window_size: int = 16
-    stride: int = 8
-
-    # You can provide custom defaults for inherited fields
-    height: int = 576
-    width: int = 1024
-
-    # Additional configuration specific to sliding tile attention
-    pad_to_square: bool = False
-    use_overlap_optimization: bool = True
-
-
 def parse_int_list(value: str) -> list[int]:
    """Parse a comma-separated string of integers into a list."""
    if not value:
@@ -10,7 +10,7 @@ from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig
 from fastvideo.configs.models.encoders import BaseEncoderOutput
 from fastvideo.configs.models.encoders.reason1 import Reason1Config, Reason1ArchConfig
 from fastvideo.configs.models.vaes import Cosmos25VAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig, STA_Mode
+from fastvideo.configs.pipelines.base import PipelineConfig


 def _identity_preprocess_text(prompt: str) -> str:
@@ -80,9 +80,6 @@ class Cosmos25Config(PipelineConfig):
    vae_tiling: bool = False
    vae_sp: bool = False

-    STA_mode: STA_Mode = STA_Mode.NONE
-    skip_time_steps: int = 0
-
    def __post_init__(self):
        self.vae_config.load_encoder = True
        self.vae_config.load_decoder = True
@@ -36,6 +36,5 @@
      "lora_config": null
    }
  ],
-  "mask_strategy_file_path": null,
  "enable_torch_compile": false
-}
+}
@@ -36,7 +36,6 @@
      "lora_config": null
    }
  ],
-  "mask_strategy_file_path": null,
  "enable_torch_compile": false,
  "image_encoder_config": {
    "prefix": "clip",
@@ -46,4 +45,4 @@
    "require_post_norm": null
  },
  "image_encoder_precision": "fp32"
-}
+}
@@ -20,7 +20,6 @@ if TYPE_CHECKING:
    FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None
    FASTVIDEO_TRACE_FUNCTION: int = 0
    FASTVIDEO_ATTENTION_BACKEND: str | None = None
-    FASTVIDEO_ATTENTION_CONFIG: str | None = None
    FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "spawn"
    FASTVIDEO_TARGET_DEVICE: str = "cuda"
    MAX_JOBS: str | None = None
@@ -201,19 +200,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
    # Available options:
    # - "TORCH_SDPA": use torch.nn.MultiheadAttention
    # - "FLASH_ATTN": use FlashAttention
-    # - "SLIDING_TILE_ATTN" : use Sliding Tile Attention
    # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention
    # - "SAGE_ATTN": use Sage Attention
    # - "SAGE_ATTN_THREE": use Sage Attention 3
    "FASTVIDEO_ATTENTION_BACKEND":
    lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),

-    # Path to the attention configuration file. Only used for sliding tile
-    # attention for now.
-    "FASTVIDEO_ATTENTION_CONFIG":
-    lambda: (None if os.getenv("FASTVIDEO_ATTENTION_CONFIG", None) is None else
-             os.path.expanduser(os.getenv("FASTVIDEO_ATTENTION_CONFIG", "."))),
-
    # Use dedicated multiprocess context for workers.
    "FASTVIDEO_WORKER_MULTIPROC_METHOD":
    lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "spawn"),
@@ -139,9 +139,6 @@ class FastVideoArgs:
    vae_cpu_offload: bool = True
    pin_cpu_memory: bool = True

-    # STA (Sliding Tile Attention) parameters
-    mask_strategy_file_path: str | None = None
-
    # Compilation
    enable_torch_compile: bool = False
    torch_compile_kwargs: dict[str, Any] = field(default_factory=dict)
@@ -457,12 +454,6 @@ class FastVideoArgs:
            help="BSA chunk_3d_shape_k as three ints, e.g., 4 4 4.",
        )

-        # STA (Sliding Tile Attention) parameters
-        parser.add_argument(
-            "--mask-strategy-file-path",
-            type=str,
-            help="Path to mask strategy JSON file for STA",
-        )
        parser.add_argument(
            "--enable-torch-compile",
            action=StoreBoolean,
@@ -753,7 +753,7 @@ class IndividualTokenRefinerBlock(nn.Module):
        self.attn = LocalAttention(
            num_heads=num_attention_heads,
            head_size=hidden_size // num_attention_heads,
-            # TODO: remove hardcode; remove STA
+            # TODO: remove hardcode
            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
                                          AttentionBackendEnum.TORCH_SDPA),
        )
@@ -770,7 +770,7 @@ class IndividualTokenRefinerBlock(nn.Module):
        self.attn = LocalAttention(
            num_heads=num_attention_heads,
            head_size=hidden_size // num_attention_heads,
-            # TODO: remove hardcode; remove STA
+            # TODO: remove hardcode
            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
                                          AttentionBackendEnum.TORCH_SDPA),
        )
@@ -209,11 +209,7 @@ class ForwardBatch:
    save_video: bool = True
    return_frames: bool = False

-    # STA parameters
-    STA_param: list | None = None
    is_cfg_negative: bool = False
-    mask_search_final_result_pos: list[list] | None = None
-    mask_search_final_result_neg: list[list] | None = None

    # VSA parameters
    VSA_sparsity: float = 0.0
@@ -10,14 +10,6 @@ from fastvideo.pipelines.stages.denoising import DenoisingStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult

-try:
-    from fastvideo.attention.backends.sliding_tile_attn import (
-        SlidingTileAttentionBackend)
-    st_attn_available = True
-except ImportError:
-    st_attn_available = False
-    SlidingTileAttentionBackend = None  # type: ignore
-
 try:
    from fastvideo.attention.backends.video_sparse_attn import (
        VideoSparseAttentionBackend)
@@ -101,10 +93,6 @@ class CausalDMDDenosingStage(DenoisingStage):
            },
        )

-        # STA
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
        # Latents and prompts
        assert batch.latents is not None, "latents must be provided"
        latents = batch.latents  # [B, C, T, H, W]
@@ -261,7 +249,6 @@ class CausalDMDDenosingStage(DenoisingStage):
                                                  w),  # type: ignore
                                patch_size=fastvideo_args.pipeline_config.
                                dit_config.patch_size,  # type: ignore
-                                STA_param=batch.STA_param,  # type: ignore
                                VSA_sparsity=fastvideo_args.
                                VSA_sparsity,  # type: ignore
                                device=get_local_torch_device(),  # type: ignore
@@ -12,7 +12,6 @@ import torch
 from tqdm.auto import tqdm

 from fastvideo.attention import get_attn_backend
-from fastvideo.configs.pipelines.base import STA_Mode
 from fastvideo.distributed import (get_local_torch_device, get_world_group)
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
@@ -28,13 +27,6 @@ from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.utils import dict_to_3d_list, masks_like

-try:
-    from fastvideo.attention.backends.sliding_tile_attn import (
-        SlidingTileAttentionBackend)
-    st_attn_available = True
-except ImportError:
-    st_attn_available = False
-
 try:
    from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend
    from fastvideo.utils import is_vmoba_available
@@ -77,7 +69,6 @@ class DenoisingStage(PipelineStage):
            head_size=attn_head_size,
            dtype=torch.float16,  # TODO(will): hack
            supported_attention_backends=(
-                AttentionBackendEnum.SLIDING_TILE_ATTN,
                AttentionBackendEnum.VIDEO_SPARSE_ATTN,
                AttentionBackendEnum.VMOBA_ATTN,
                AttentionBackendEnum.FLASH_ATTN,
@@ -184,10 +175,6 @@ class DenoisingStage(PipelineStage):
            },
        )

-        # Prepare STA parameters
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
        # Get latents and embeddings
        latents = batch.latents
        prompt_embeds = batch.prompt_embeds
@@ -362,10 +349,8 @@ class DenoisingStage(PipelineStage):
                with torch.autocast(device_type="cuda",
                                    dtype=target_dtype,
                                    enabled=autocast_enabled):
-                    if (st_attn_available
-                            and self.attn_backend == SlidingTileAttentionBackend
-                        ) or (vsa_available and self.attn_backend
-                              == VideoSparseAttentionBackend):
+                    if (vsa_available and self.attn_backend
+                            == VideoSparseAttentionBackend):
                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
                        )

@@ -380,7 +365,6 @@ class DenoisingStage(PipelineStage):
                                patch_size=fastvideo_args.
                                pipeline_config.  # type: ignore
                                dit_config.patch_size,  # type: ignore
-                                STA_param=batch.STA_param,  # type: ignore
                                VSA_sparsity=fastvideo_args.
                                VSA_sparsity,  # type: ignore
                                device=get_local_torch_device(),
@@ -519,10 +503,6 @@ class DenoisingStage(PipelineStage):
                if mgr2 is not None and getattr(mgr2, "enabled", False):
                    mgr2.release_all()

-        # Save STA mask search results if needed
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.pipeline_config.STA_mode == STA_Mode.STA_SEARCHING:
-            self.save_sta_search_results(batch)
-
        # deallocate transformer if on mps
        if torch.backends.mps.is_available():
            logger.info("Memory before deallocating transformer: %s",
@@ -602,145 +582,6 @@ class DenoisingStage(PipelineStage):
                     (1 - guidance_rescale) * noise_cfg)
        return noise_cfg

-    def prepare_sta_param(self, batch: ForwardBatch,
-                          fastvideo_args: FastVideoArgs):
-        """
-        Prepare Sliding Tile Attention (STA) parameters and settings.
-        
-        Args:
-            batch: The current batch information.
-            fastvideo_args: The inference arguments.
-        """
-        # TODO(kevin): STA mask search, currently only support Wan2.1 with 69x768x1280
-        from fastvideo.attention.backends.STA_configuration import configure_sta
-        STA_mode = fastvideo_args.pipeline_config.STA_mode
-        skip_time_steps = fastvideo_args.pipeline_config.skip_time_steps
-        if batch.timesteps is None:
-            raise ValueError("Timesteps must be provided")
-        timesteps_num = batch.timesteps.shape[0]
-
-        logger.info("STA_mode: %s", STA_mode)
-        if (batch.num_frames, batch.height,
-                batch.width) != (69, 768, 1280) and STA_mode != "STA_inference":
-            raise NotImplementedError(
-                "STA mask search/tuning is not supported for this resolution")
-
-        if STA_mode == STA_Mode.STA_SEARCHING or STA_mode == STA_Mode.STA_TUNING or STA_mode == STA_Mode.STA_TUNING_CFG:
-            size = (batch.width, batch.height)
-            if size == (1280, 768):
-                # TODO: make it configurable
-                sparse_mask_candidates_searching = [
-                    "3, 1, 10", "1, 5, 7", "3, 3, 3", "1, 6, 5", "1, 3, 10",
-                    "3, 6, 1"
-                ]
-                sparse_mask_candidates_tuning = [
-                    "3, 1, 10", "1, 5, 7", "3, 3, 3", "1, 6, 5", "1, 3, 10",
-                    "3, 6, 1"
-                ]
-                full_mask = ["3,6,10"]
-            else:
-                raise NotImplementedError(
-                    "STA mask search is not supported for this resolution")
-        layer_num = self.transformer.config.num_layers
-        # specific for HunyuanVideo
-        if hasattr(self.transformer.config, "num_single_layers"):
-            layer_num += self.transformer.config.num_single_layers
-        head_num = self.transformer.config.num_attention_heads
-
-        if STA_mode == STA_Mode.STA_SEARCHING:
-            STA_param = configure_sta(
-                mode=STA_Mode.STA_SEARCHING,
-                layer_num=layer_num,
-                head_num=head_num,
-                time_step_num=timesteps_num,
-                mask_candidates=sparse_mask_candidates_searching +
-                full_mask,  # last is full mask; Can add more sparse masks while keep last one as full mask
-            )
-        elif STA_mode == STA_Mode.STA_TUNING:
-            STA_param = configure_sta(
-                mode=STA_Mode.STA_TUNING,
-                layer_num=layer_num,
-                head_num=head_num,
-                time_step_num=timesteps_num,
-                mask_search_files_path=
-                f'output/mask_search_result_pos_{size[0]}x{size[1]}/',
-                mask_candidates=sparse_mask_candidates_tuning,
-                full_attention_mask=[int(x) for x in full_mask[0].split(',')],
-                skip_time_steps=
-                skip_time_steps,  # Use full attention for first 12 steps
-                save_dir=
-                f'output/mask_search_strategy_{size[0]}x{size[1]}/',  # Custom save directory
-                timesteps=timesteps_num)
-        elif STA_mode == STA_Mode.STA_TUNING_CFG:
-            STA_param = configure_sta(
-                mode=STA_Mode.STA_TUNING_CFG,
-                layer_num=layer_num,
-                head_num=head_num,
-                time_step_num=timesteps_num,
-                mask_search_files_path_pos=
-                f'output/mask_search_result_pos_{size[0]}x{size[1]}/',
-                mask_search_files_path_neg=
-                f'output/mask_search_result_neg_{size[0]}x{size[1]}/',
-                mask_candidates=sparse_mask_candidates_tuning,
-                full_attention_mask=[int(x) for x in full_mask[0].split(',')],
-                skip_time_steps=skip_time_steps,
-                save_dir=f'output/mask_search_strategy_{size[0]}x{size[1]}/',
-                timesteps=timesteps_num)
-        elif STA_mode == STA_Mode.STA_INFERENCE:
-            import fastvideo.envs as envs
-            config_file = envs.FASTVIDEO_ATTENTION_CONFIG
-            if config_file is None:
-                raise ValueError("FASTVIDEO_ATTENTION_CONFIG is not set")
-            STA_param = configure_sta(mode=STA_Mode.STA_INFERENCE,
-                                      layer_num=layer_num,
-                                      head_num=head_num,
-                                      time_step_num=timesteps_num,
-                                      load_path=config_file)
-
-        batch.STA_param = STA_param
-        batch.mask_search_final_result_pos = [[] for _ in range(timesteps_num)]
-        batch.mask_search_final_result_neg = [[] for _ in range(timesteps_num)]
-
-    def save_sta_search_results(self, batch: ForwardBatch):
-        """
-        Save the STA mask search results.
-        
-        Args:
-            batch: The current batch information.
-        """
-        size = (batch.width, batch.height)
-        if size == (1280, 768):
-            # TODO: make it configurable
-            sparse_mask_candidates_searching = [
-                "3, 1, 10", "1, 5, 7", "3, 3, 3", "1, 6, 5", "1, 3, 10",
-                "3, 6, 1"
-            ]
-        else:
-            raise NotImplementedError(
-                "STA mask search is not supported for this resolution")
-
-        from fastvideo.attention.backends.STA_configuration import save_mask_search_results
-        if batch.mask_search_final_result_pos is not None and batch.prompt is not None:
-            save_mask_search_results(
-                [
-                    dict(layer_data)
-                    for layer_data in batch.mask_search_final_result_pos
-                ],
-                prompt=str(batch.prompt),
-                mask_strategies=sparse_mask_candidates_searching,
-                output_dir=f'output/mask_search_result_pos_{size[0]}x{size[1]}/'
-            )
-        if batch.mask_search_final_result_neg is not None and batch.prompt is not None:
-            save_mask_search_results(
-                [
-                    dict(layer_data)
-                    for layer_data in batch.mask_search_final_result_neg
-                ],
-                prompt=str(batch.prompt),
-                mask_strategies=sparse_mask_candidates_searching,
-                output_dir=f'output/mask_search_result_neg_{size[0]}x{size[1]}/'
-            )
-
    def verify_input(self, batch: ForwardBatch,
                     fastvideo_args: FastVideoArgs) -> VerificationResult:
        """Verify denoising stage inputs."""
@@ -1396,10 +1237,6 @@ class DmdDenoisingStage(DenoisingStage):
            },
        )

-        # Prepare STA parameters
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
        # Get latents and embeddings
        assert batch.latents is not None, "latents must be provided"
        latents = batch.latents
@@ -1464,7 +1301,6 @@ class DmdDenoisingStage(DenoisingStage):
                                patch_size=fastvideo_args.
                                pipeline_config.  # type: ignore
                                dit_config.patch_size,  # type: ignore
-                                STA_param=batch.STA_param,  # type: ignore
                                VSA_sparsity=fastvideo_args.
                                VSA_sparsity,  # type: ignore
                                device=get_local_torch_device(),  # type: ignore
@@ -15,14 +15,6 @@ from fastvideo.pipelines.stages.denoising import DenoisingStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult

-try:
-    from fastvideo.attention.backends.sliding_tile_attn import (
-        SlidingTileAttentionBackend)
-    st_attn_available = True
-except ImportError:
-    st_attn_available = False
-    SlidingTileAttentionBackend = None  # type: ignore
-
 try:
    from fastvideo.attention.backends.video_sparse_attn import (
        VideoSparseAttentionBackend)
@@ -168,9 +160,6 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
        image_kwargs = {"encoder_hidden_states_image": image_embeds}
        pos_cond_kwargs: dict[str, Any] = {}

-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
        assert batch.latents is not None, "latents must be provided"
        latents = batch.latents
        b, c, t, h, w = latents.shape
@@ -480,7 +469,6 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
                        raw_latent_shape=(current_num_frames, h, w),
                        patch_size=ctx.fastvideo_args.pipeline_config.
                        dit_config.patch_size,
-                        STA_param=batch.STA_param,
                        VSA_sparsity=ctx.fastvideo_args.VSA_sparsity,
                        device=get_local_torch_device(),
                    )
@@ -691,9 +679,6 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
        image_kwargs = {"encoder_hidden_states_image": image_embeds}
        pos_cond_kwargs: dict[str, Any] = {}

-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
        assert batch.latents is not None, "latents must be provided"
        latents = batch.latents
        b, c, t, h, w = latents.shape
@@ -15,7 +15,6 @@ import numpy as np
 from tqdm.auto import tqdm

 from fastvideo.attention import get_attn_backend
-from fastvideo.configs.pipelines.base import STA_Mode
 from fastvideo.distributed import (get_local_torch_device, get_world_group)
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
@@ -27,13 +26,6 @@ from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.platforms import AttentionBackendEnum

-try:
-    from fastvideo.attention.backends.sliding_tile_attn import (
-        SlidingTileAttentionBackend)
-    st_attn_available = True
-except ImportError:
-    st_attn_available = False
-
 try:
    from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend
    from fastvideo.utils import is_vmoba_available
@@ -76,7 +68,6 @@ class SRDenoisingStage(PipelineStage):
            head_size=attn_head_size,
            dtype=torch.float16,  # TODO(will): hack
            supported_attention_backends=(
-                AttentionBackendEnum.SLIDING_TILE_ATTN,
                AttentionBackendEnum.VIDEO_SPARSE_ATTN,
                AttentionBackendEnum.VMOBA_ATTN,
                AttentionBackendEnum.FLASH_ATTN,
@@ -171,10 +162,6 @@ class SRDenoisingStage(PipelineStage):
            },
        )

-        # Prepare STA parameters
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
-            self.prepare_sta_param(batch, fastvideo_args)
-
        # Get latents and embeddings
        prompt_embeds = batch.prompt_embeds
        assert not torch.isnan(
@@ -253,10 +240,8 @@ class SRDenoisingStage(PipelineStage):
                with torch.autocast(device_type="cuda",
                                    dtype=target_dtype,
                                    enabled=autocast_enabled):
-                    if (st_attn_available
-                            and self.attn_backend == SlidingTileAttentionBackend
-                        ) or (vsa_available and self.attn_backend
-                              == VideoSparseAttentionBackend):
+                    if (vsa_available and self.attn_backend
+                            == VideoSparseAttentionBackend):
                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
                        )

@@ -271,7 +256,6 @@ class SRDenoisingStage(PipelineStage):
                                patch_size=fastvideo_args.
                                pipeline_config.  # type: ignore
                                dit_config.patch_size,  # type: ignore
-                                STA_param=batch.STA_param,  # type: ignore
                                VSA_sparsity=fastvideo_args.
                                VSA_sparsity,  # type: ignore
                                device=get_local_torch_device(),
@@ -342,10 +326,6 @@ class SRDenoisingStage(PipelineStage):
        # Update batch with final latents
        batch.latents = latents

-        # Save STA mask search results if needed
-        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.pipeline_config.STA_mode == STA_Mode.STA_SEARCHING:
-            self.save_sta_search_results(batch)
-
        # deallocate transformer if on mps
        if torch.backends.mps.is_available():
            logger.info("Memory before deallocating transformer: %s",
@@ -124,22 +124,7 @@ class CudaPlatformBase(Platform):
        logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s",
                    envs.FASTVIDEO_ATTENTION_BACKEND)
        logger.info("Selected backend: %s", selected_backend)
-        if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN:
-            try:
-                from fastvideo_kernel import sliding_tile_attention  # noqa: F401
-
-                from fastvideo.attention.backends.sliding_tile_attn import (  # noqa: F401
-                    SlidingTileAttentionBackend)
-                logger.info("Using Sliding Tile Attention backend.")
-
-                return "fastvideo.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend"
-            except ImportError as e:
-                logger.error(
-                    "Failed to import Sliding Tile Attention backend: %s",
-                    str(e))
-                raise ImportError(
-                    "Sliding Tile Attention backend is not installed. ") from e
-        elif selected_backend == AttentionBackendEnum.SAGE_ATTN:
+        if selected_backend == AttentionBackendEnum.SAGE_ATTN:
            try:
                from sageattention import sageattn  # noqa: F401

@@ -12,7 +12,6 @@ logger = init_logger(__name__)

 class AttentionBackendEnum(enum.Enum):
    FLASH_ATTN = enum.auto()
-    SLIDING_TILE_ATTN = enum.auto()
    TORCH_SDPA = enum.auto()
    SAGE_ATTN = enum.auto()
    SAGE_ATTN_THREE = enum.auto()
@@ -78,22 +78,6 @@ class RocmPlatform(Platform):
        elif selected_backend in (AttentionBackendEnum.FLASH_ATTN, None):
            pass

-        elif selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN:
-            try:
-                from fastvideo_kernel import sliding_tile_attention  # noqa: F401
-
-                from fastvideo.attention.backends.sliding_tile_attn import (  # noqa: F401
-                    SlidingTileAttentionBackend)
-                logger.info("Using Sliding Tile Attention backend.")
-
-                return "fastvideo.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend"
-            except ImportError as e:
-                logger.error(
-                    "Failed to import Sliding Tile Attention backend: %s",
-                    str(e))
-                raise ImportError(
-                    "Sliding Tile Attention backend is not installed. ") from e
-
        elif selected_backend in (AttentionBackendEnum.SAGE_ATTN):
            raise ValueError(
                f"{selected_backend.name} is not supported on {cls.device_name}."
@@ -1,52 +0,0 @@
-import os
-import sys
-import subprocess
-from pathlib import Path
-import pytest
-
-NUM_NODES = "1"
-NUM_GPUS_PER_NODE = "2"
-
-# Set environment variables
-os.environ["FASTVIDEO_ATTENTION_CONFIG"] = "assets/mask_strategy_wan.json"
-os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
-
-def test_inference():
-    """Test the inference functionality"""
-    # Create command as in wan_14B-STA.sh
-    cmd = [
-        "fastvideo", "generate",
-        "--model-path", "Wan-AI/Wan2.1-T2V-14B-Diffusers",
-        "--sp-size", "2",
-        "--tp-size", "2",
-        "--num-gpus", "2",
-        "--height", "768",
-        "--width", "1280",
-        "--num-frames", "69",
-        "--num-inference-steps", "2",
-        "--fps", "16",
-        "--guidance-scale", "5.0",
-        "--flow-shift", "5.0",
-        "--prompt", "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.",
-        "--negative-prompt", "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
-        "--seed", "1024",
-        "--output-path", "outputs_video/STA_1024/",
-    ]
-    
-    # Run the command
-    subprocess.run(cmd, check=True)
-    
-    # Verify output directory exists
-    output_dir = Path("outputs_video/STA_1024/")
-    assert output_dir.exists(), f"Output directory {output_dir} does not exist"
-    
-    # Verify that video files were generated
-    video_files = list(output_dir.glob("*.mp4"))
-    assert len(video_files) > 0, "No video files were generated"
-    
-    # Verify the video file properties
-    for video_file in video_files:
-        assert video_file.stat().st_size > 0, f"Video file {video_file} is empty"
-
-if __name__ == "__main__":
-    test_inference()
@@ -115,10 +115,6 @@ def run_training_lora_tests():
 def run_training_tests_VSA():
    run_test("wandb login $WANDB_API_KEY && pytest ./fastvideo/tests/training/VSA -srP")

-@app.function(gpu="H100:2", image=image, timeout=900)
-def run_inference_tests_STA():
-    run_test("pytest ./fastvideo/tests/inference/STA -srP")
-
@app.function(gpu="H100:1", image=image, timeout=900)
 def run_kernel_tests():
    run_test("pytest fastvideo-kernel/tests/ -vs")
@@ -56,7 +56,6 @@ PRECISION_TO_TYPE = {
 }

 STR_BACKEND_ENV_VAR: str = "FASTVIDEO_ATTENTION_BACKEND"
-STR_ATTN_CONFIG_ENV_VAR: str = "FASTVIDEO_ATTENTION_CONFIG"


 def find_nccl_library() -> str:
@@ -134,7 +134,7 @@ nav:
    - Examples:
      - Basic: inference/examples/basic.md
      - Optimizations: inference/examples/optimizations.md
-      - STA Mask Search: inference/examples/sta_mask_search.md
+      - STA Mask Search (Archived): inference/examples/sta_mask_search.md
  - Training:
    - Overview: training/overview.md
    - Data Preprocessing: training/data_preprocess.md
@@ -152,7 +152,7 @@ nav:
  - Attention:
    - Overview: attention/index.md
    - Video Sparse Attention: attention/vsa/index.md
-    - Sliding Tile Attention: attention/sta/index.md
+    - Sliding Tile Attention (Archived): attention/sta/index.md
    - Backend Development: contributing/attention_backend.md
  - Utilities:
    - LoRA: utilities/lora.md
@@ -143,7 +143,7 @@ check_untyped_defs = true
 follow_imports = "silent"

 [tool.codespell]
-skip ="./data,./wandb,./csrc/sliding_tile_attention/tk"
+skip ="./data,./wandb"

 [tool.ruff]
 # Allow lines to be as long as 80.
@@ -179,7 +179,6 @@ ignore = [
 ]

 [tool.ruff.lint.per-file-ignores]
-"csrc/sliding_tile_attention/test/bench.py" = ["F841"]
 # Ignore all files that end in `_test.py`.
 "fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py" = ["E741"]

@@ -121,7 +121,7 @@ check_untyped_defs = true
 follow_imports = "silent"

 [tool.codespell]
-skip ="./data,./wandb,./csrc/sliding_tile_attention/tk"
+skip ="./data,./wandb"

 [tool.ruff]
 # Allow lines to be as long as 80.
@@ -157,7 +157,6 @@ ignore = [
 ]

 [tool.ruff.lint.per-file-ignores]
-"csrc/sliding_tile_attention/test/bench.py" = ["F841"]
 # Ignore all files that end in `_test.py`.
 "fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py" = ["E741"]

@@ -1,22 +0,0 @@
-#!/bin/bash
-
-num_gpus=2
-export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_hunyuan.json
-export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
-export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
-# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo
-fastvideo generate \
-    --model-path $MODEL_BASE \
-    --sp-size ${num_gpus} \
-    --tp-size 1 \
-    --num-gpus ${num_gpus} \
-    --height 768 \
-    --width 1280 \
-    --num-frames 117 \
-    --num-inference-steps 50 \
-    --guidance-scale 1 \
-    --embedded-cfg-scale 6 \
-    --flow-shift 7 \
-    --prompt "A beautiful woman in a red dress walking down a street" \
-    --seed 1024 \
-    --output-path outputs_video/
@@ -1,21 +0,0 @@
-#!/bin/bash
-
-num_gpus=2
-export FASTVIDEO_ATTENTION_CONFIG=assets/mask_strategy_wan.json
-export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
-export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers
-fastvideo generate \
-    --model-path $MODEL_BASE \
-    --sp-size $num_gpus \
-    --tp-size 1 \
-    --num-gpus $num_gpus \
-    --height 768 \
-    --width 1280 \
-    --num-frames 69 \
-    --num-inference-steps 50 \
-    --fps 16 \
-    --guidance-scale 5.0 \
-    --prompt "A beautiful woman in a red dress walking down a street" \
-    --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
-    --seed 1024 \
-    --output-path outputs_video/