diff --git a/.github/mergify.yml b/.github/mergify.yml
index a577719ef..71ad38590 100644
--- a/.github/mergify.yml
+++ b/.github/mergify.yml
@@ -197,6 +197,30 @@ pull_request_rules:
 
           For future commits, `pre-commit` will run automatically on changed files before each commit.
 
+  # ============================================================
+  # Buildkite CI failure notification
+  # ============================================================
+
+  - name: comment on Buildkite CI failure
+    conditions:
+      - check-failure~=buildkite/ci
+      - -closed
+    actions:
+      comment:
+        message: |
+          ## Buildkite CI tests failed
+
+          Hi @{{author}}, some Buildkite CI tests have failed. Check the build for details:
+          [View Buildkite build →](https://buildkite.com/fastvideo/ci)
+
+          Common causes:
+          - **Test failures**: Check the failing step's output for assertion errors or tracebacks
+          - **Import errors**: Make sure new dependencies are added to `pyproject.toml`
+          - **GPU memory**: Some tests require specific GPU types (L40S, H100 NVL)
+          - **Kernel build**: If you changed `fastvideo-kernel/`, the build may have failed
+
+          If the failure is unrelated to your changes, leave a comment explaining why.
+
   # ============================================================
   # Merge conflict detection
   # ============================================================
diff --git a/.github/scripts/runpod_api.py b/.github/scripts/runpod_api.py
deleted file mode 100644
index 280332a0c..000000000
--- a/.github/scripts/runpod_api.py
+++ /dev/null
@@ -1,249 +0,0 @@
-import argparse
-import json
-import os
-import subprocess
-import sys
-import time
-
-import requests
-
-
-def parse_arguments():
-    """Parse command line arguments"""
-    parser = argparse.ArgumentParser(description='Run tests on RunPod GPU')
-    parser.add_argument('--gpu-type', type=str, help='GPU type to use')
-    parser.add_argument('--gpu-count',
-                        type=int,
-                        help='Number of GPUs to use',
-                        default=1)
-    parser.add_argument('--test-command', type=str, help='Test command to run')
-    parser.add_argument('--disk-size',
-                        type=int,
-                        default=20,
-                        help='Container disk size in GB (default: 20)')
-    parser.add_argument('--volume-size',
-                        type=int,
-                        default=20,
-                        help='Persistent volume size in GB (default: 20)')
-    parser.add_argument(
-        '--image',
-        type=str,
-        required=True,
-        help='Docker image to use')
-    return parser.parse_args()
-
-
-args = parse_arguments()
-API_KEY = os.environ['RUNPOD_API_KEY']
-RUN_ID = os.environ['GITHUB_RUN_ID']
-JOB_ID = os.environ['JOB_ID']
-PODS_API = "https://rest.runpod.io/v1/pods"
-HEADERS = {
-    "Content-Type": "application/json",
-    "Authorization": f"Bearer {API_KEY}"
-}
-
-
-def create_pod():
-    """Create a RunPod instance"""
-    # Ensure image name is lowercase (Docker requirement)
-    image_name = args.image.lower()
-    print(f"Using specified image: {image_name}")
-
-    docker_start_cmd = [
-        "bash", 
-        "-c", 
-        "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-server -y;mkdir -p ~/.ssh;cd $_;chmod 700 ~/.ssh;echo \"$PUBLIC_KEY\" >> authorized_keys;chmod 700 authorized_keys;service ssh start;sleep infinity"
-    ]
-    
-    print(f"Creating RunPod instance with GPU: {args.gpu_type}...")
-    payload = {
-        "name": f"fastvideo-{JOB_ID}-{RUN_ID}",
-        "containerDiskInGb": args.disk_size,
-        "volumeInGb": args.volume_size,
-        "gpuTypeIds": [args.gpu_type],
-        "gpuCount": args.gpu_count,
-        "imageName": image_name,
-        "allowedCudaVersions": ["12.4"],
-        "dockerStartCmd": docker_start_cmd
-    }
-
-    response = requests.post(PODS_API, headers=HEADERS, json=payload)
-    response_data = response.json()
-    print(f"Response: {json.dumps(response_data, indent=2)}")
-
-    return response_data["id"]
-
-
-def wait_for_pod(pod_id):
-    """Wait for pod to be in RUNNING state and fully ready with SSH access"""
-    print("Waiting for RunPod to be ready...")
-
-    # First wait for RUNNING status
-    max_attempts = 10
-    attempts = 0
-    while attempts < max_attempts:
-        response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
-        pod_data = response.json()
-        status = pod_data["desiredStatus"]
-
-        if status == "RUNNING":
-            print("RunPod is running! Now waiting for ports to be assigned...")
-            break
-
-        print(
-            f"Current status: {status}, waiting... (attempt {attempts+1}/{max_attempts})"
-        )
-        time.sleep(2)
-        attempts += 1
-
-    if attempts >= max_attempts:
-        raise TimeoutError(
-            "Timed out waiting for RunPod to reach RUNNING state")
-
-    # Wait for ports to be assigned
-    max_attempts = 50
-    attempts = 0
-    while attempts < max_attempts:
-        response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
-        pod_data = response.json()
-        port_mappings = pod_data.get("portMappings")
-
-        if (port_mappings is not None and "22" in port_mappings
-                and pod_data.get("publicIp", "") != ""):
-            print("RunPod is ready with SSH access!")
-            print(f"SSH IP: {pod_data['publicIp']}")
-            print(f"SSH Port: {port_mappings['22']}")
-            break
-
-        print(
-            f"Waiting for SSH port and public IP to be available... (attempt {attempts+1}/{max_attempts})"
-        )
-        time.sleep(20)
-        attempts += 1
-
-    if attempts >= max_attempts:
-        raise TimeoutError("Timed out waiting for RunPod SSH access")
-
-
-def execute_command(pod_id):
-    """Execute command on the pod via SSH using system SSH client"""
-    print(f"Running command: {args.test_command}")
-
-    response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
-    pod_data = response.json()
-    ssh_ip = pod_data["publicIp"]
-    ssh_port = pod_data["portMappings"]["22"]
-
-    # Copy the repository to the pod using scp
-    repo_dir = os.path.abspath(os.getcwd())
-    repo_name = os.path.basename(repo_dir)
-
-    print(f"Copying repository from {repo_dir} to RunPod...")
-
-    tar_command = [
-        "tar", "-czf", "/tmp/repo.tar.gz", "-C",
-        os.path.dirname(repo_dir), repo_name
-    ]
-    subprocess.run(tar_command, check=True)
-
-    # Copy the tarball to the pod
-    scp_command = [
-        "scp", "-o", "StrictHostKeyChecking=no", "-o",
-        "UserKnownHostsFile=/dev/null", "-o", "ServerAliveInterval=60", "-o",
-        "ServerAliveCountMax=10", "-P",
-        str(ssh_port), "/tmp/repo.tar.gz", f"root@{ssh_ip}:/tmp/"
-    ]
-    subprocess.run(scp_command, check=True)
-
-    # For custom image, we can use the pre-configured environment
-    setup_steps = [
-        "tar -xzf /tmp/repo.tar.gz --no-same-owner -C /workspace/",
-        f"cd /workspace/{repo_name}",
-        "source $HOME/.local/bin/env && source /opt/venv/bin/activate",
-        args.test_command
-    ]
-    
-    remote_command = " && ".join(setup_steps)
-
-    ssh_command = [
-        "ssh", "-o", "StrictHostKeyChecking=no", "-o",
-        "UserKnownHostsFile=/dev/null", "-o", "ServerAliveInterval=60", "-o",
-        "ServerAliveCountMax=10", "-p",
-        str(ssh_port), f"root@{ssh_ip}", remote_command
-    ]
-
-    print(f"Connecting to {ssh_ip}:{ssh_port}...")
-
-    try:
-        process = subprocess.Popen(ssh_command,
-                                   stdout=subprocess.PIPE,
-                                   stderr=subprocess.STDOUT,
-                                   universal_newlines=True,
-                                   bufsize=0)
-
-        stdout_lines = []
-
-        print("Command output:")
-
-        for line in iter(process.stdout.readline, ''):
-            print(line.strip())
-            stdout_lines.append(line)
-
-        process.wait()
-
-        return_code = process.returncode
-        success = return_code == 0
-
-        stdout_str = "".join(stdout_lines)
-
-        if success:
-            print("Command executed successfully")
-        else:
-            print(f"Command failed with exit code {return_code}")
-
-        result = {
-            "success": success,
-            "return_code": return_code,
-            "stdout": stdout_str,
-            "stderr": ""
-        }
-        return result
-
-    except Exception as e:
-        print(f"Error executing SSH command: {str(e)}")
-        result = {"success": False, "error": str(e), "stdout": "", "stderr": ""}
-        return result
-
-
-def terminate_pod(pod_id):
-    """Terminate the pod"""
-    print("Terminating RunPod...")
-    requests.delete(f"{PODS_API}/{pod_id}", headers=HEADERS)
-    print(f"Terminated pod {pod_id}")
-
-
-def main():
-    pod_id = None
-    try:
-        pod_id = create_pod()
-        wait_for_pod(pod_id)
-        result = execute_command(pod_id)
-
-        if result.get("error") is not None:
-            print(f"Error executing command: {result['error']}")
-            sys.exit(1)
-
-        if not result.get("success", False):
-            print(
-                "Tests failed - check the output above for details on which tests failed"
-            )
-            sys.exit(1)
-
-    finally:
-        if pod_id:
-            terminate_pod(pod_id)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/.github/scripts/runpod_cleanup.py b/.github/scripts/runpod_cleanup.py
deleted file mode 100644
index eef991af7..000000000
--- a/.github/scripts/runpod_cleanup.py
+++ /dev/null
@@ -1,90 +0,0 @@
-import json
-import os
-import sys
-import uuid
-
-import requests
-
-API_KEY = os.environ['RUNPOD_API_KEY']
-RUN_ID = os.environ.get('GITHUB_RUN_ID', str(uuid.uuid4()))
-PODS_API = "https://rest.runpod.io/v1/pods"
-HEADERS = {
-    "Content-Type": "application/json",
-    "Authorization": f"Bearer {API_KEY}"
-}
-
-
-def get_job_ids():
-    """Parse job IDs from environment variable"""
-    job_ids_str = os.environ.get('JOB_IDS')
-    try:
-        job_ids = json.loads(job_ids_str)
-        if not isinstance(job_ids, list):
-            print("Error: JOB_IDS is not a list.")
-            sys.exit(1)
-        return job_ids
-    except json.JSONDecodeError as e:
-        print(f"Error parsing JOB_IDS: {e}")
-        sys.exit(1)
-
-
-def cleanup_pods():
-    """Find and terminate RunPod instances"""
-    print(f"Run ID: {RUN_ID}")
-
-    single_job_id = os.environ.get('JOB_ID')
-    
-    if single_job_id:
-        job_ids = [single_job_id]
-        print(f"Job ID: {single_job_id}")
-    else:
-        job_ids = get_job_ids()
-        print(f"Job IDs: {job_ids}")
-
-    # Get all pods associated with RunPod API_KEY
-    try:
-        response = requests.get(PODS_API, headers=HEADERS)
-        response.raise_for_status()
-        pods = response.json()
-    except requests.exceptions.RequestException as e:
-        print(f"Error getting pods: {e}")
-        sys.exit(1)
-
-    # Find and terminate pods created by this workflow run
-    terminated_pods = []
-    for pod in pods:
-        pod_name = pod.get("name", "")
-        pod_id = pod.get("id")
-
-        # Check if this pod was created by one of our jobs
-        if any(f"{job_id}-{RUN_ID}" in pod_name for job_id in job_ids):
-            print(f"Found pod: {pod_id} ({pod_name})")
-            try:
-                print(f"Terminating pod {pod_id}...")
-                term_response = requests.delete(f"{PODS_API}/{pod_id}",
-                                                headers=HEADERS)
-                term_response.raise_for_status()
-                terminated_pods.append(pod_id)
-                print(f"Successfully terminated pod {pod_id}")
-            except requests.exceptions.RequestException as e:
-                print(f"Error terminating pod {pod_id}: {e}")
-                sys.exit(1)
-    
-    if terminated_pods:
-        if single_job_id:
-            print(f"Terminated pod: {terminated_pods[0]}")
-        else:
-            print(f"Terminated {len(terminated_pods)} pods: {terminated_pods}")
-    else:
-        if single_job_id:
-            print(f"No pod found matching pattern: {single_job_id}-{RUN_ID}")
-        else:
-            print("No pods found to terminate.")
-
-
-def main():
-    cleanup_pods()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/.github/workflows/build-image-template.yml b/.github/workflows/_template-build-image.yml
similarity index 100%
rename from .github/workflows/build-image-template.yml
rename to .github/workflows/_template-build-image.yml
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/ci-precommit.yml
similarity index 100%
rename from .github/workflows/pre-commit.yml
rename to .github/workflows/ci-precommit.yml
diff --git a/.github/workflows/trigger-full-suite.yml b/.github/workflows/ci-trigger-full-suite.yml
similarity index 100%
rename from .github/workflows/trigger-full-suite.yml
rename to .github/workflows/ci-trigger-full-suite.yml
diff --git a/.github/workflows/issue-labeler.yml b/.github/workflows/community-issue-labeler.yml
similarity index 100%
rename from .github/workflows/issue-labeler.yml
rename to .github/workflows/community-issue-labeler.yml
diff --git a/.github/workflows/stale.yml b/.github/workflows/community-stale.yml
similarity index 100%
rename from .github/workflows/stale.yml
rename to .github/workflows/community-stale.yml
diff --git a/.github/workflows/welcome.yml b/.github/workflows/community-welcome.yml
similarity index 100%
rename from .github/workflows/welcome.yml
rename to .github/workflows/community-welcome.yml
diff --git a/.github/workflows/build-image.yml b/.github/workflows/infra-build-image.yml
similarity index 86%
rename from .github/workflows/build-image.yml
rename to .github/workflows/infra-build-image.yml
index 664f376cd..145048fb9 100644
--- a/.github/workflows/build-image.yml
+++ b/.github/workflows/infra-build-image.yml
@@ -32,7 +32,7 @@ permissions:
 jobs:
   build-python-3-10:
     if: ${{ github.event.inputs.python_3_10 == 'true' }}
-    uses: ./.github/workflows/build-image-template.yml
+    uses: ./.github/workflows/_template-build-image.yml
     with:
       python_version: '3.10'
       dockerfile_path: docker/Dockerfile.python3.10
@@ -41,7 +41,7 @@ jobs:
 
   build-python-3-11:
     if: ${{ github.event.inputs.python_3_11 == 'true' }}
-    uses: ./.github/workflows/build-image-template.yml
+    uses: ./.github/workflows/_template-build-image.yml
     with:
       python_version: '3.11'
       dockerfile_path: docker/Dockerfile.python3.11
@@ -50,7 +50,7 @@ jobs:
 
   build-python-3-12:
     if: ${{ github.event.inputs.python_3_12 == 'true' }}
-    uses: ./.github/workflows/build-image-template.yml
+    uses: ./.github/workflows/_template-build-image.yml
     with:
       python_version: '3.12'
       dockerfile_path: docker/Dockerfile.python3.12
@@ -59,9 +59,9 @@ jobs:
 
   build-python-3-12-cuda-12-9:
     if: ${{ github.event.inputs.python_3_12_cuda_12_9 == 'true' }}
-    uses: ./.github/workflows/build-image-template.yml
+    uses: ./.github/workflows/_template-build-image.yml
     with:
       python_version: '3.12'
       dockerfile_path: docker/Dockerfile.python3.12.cuda12.9.1
       tag_suffix: py3.12-cuda12.9.1
-    secrets: inherit
\ No newline at end of file
+    secrets: inherit
diff --git a/.github/workflows/docs.yml b/.github/workflows/infra-docs.yml
similarity index 94%
rename from .github/workflows/docs.yml
rename to .github/workflows/infra-docs.yml
index 693abcfaa..654404146 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/infra-docs.yml
@@ -7,14 +7,14 @@ on:
       - 'docs/**'
       - 'mkdocs.yml'
       - 'requirements-mkdocs.txt'
-      - '.github/workflows/docs.yml'
+      - '.github/workflows/infra-docs.yml'
   pull_request:
     branches: [ main ]
     paths:
       - 'docs/**'
       - 'mkdocs.yml'
       - 'requirements-mkdocs.txt'
-      - '.github/workflows/docs.yml'
+      - '.github/workflows/infra-docs.yml'
 
 permissions:
   contents: read
diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml
deleted file mode 100644
index 585978e8f..000000000
--- a/.github/workflows/pr-test.yml
+++ /dev/null
@@ -1,334 +0,0 @@
-name: PR Test
-
-on:
-  push:
-    branches: [main]
-    paths:
-      - "fastvideo/**/*.py"
-      - ".github/workflows/pr-test.yml"
-  pull_request:
-    branches: [main]
-    types: [opened, ready_for_review, synchronize, reopened]
-    paths:
-      - "fastvideo/**/*.py"
-      - ".github/workflows/pr-test.yml"
-      - "pyproject.toml"
-      - "docker/Dockerfile.python3.12"
-      - "csrc/**"
-  workflow_dispatch:
-    inputs:
-      run_encoder_test:
-        description: "Run encoder-test"
-        required: false
-        default: false
-        type: boolean
-      run_vae_test:
-        description: "Run vae-test"
-        required: false
-        default: false
-        type: boolean
-      run_transformer_test:
-        description: "Run transformer-test"
-        required: false
-        default: false
-        type: boolean
-      run_ssim_test:
-        description: "Run ssim-test"
-        required: false
-        default: false
-        type: boolean
-      run_training_test:
-        description: "Run training-test"
-        required: false
-        default: false
-        type: boolean
-      run_training_test_VSA:
-        description: "Run training-test-VSA"
-        required: false
-        default: false
-        type: boolean
-      run_precision_test_VSA:
-        description: "Run precision-test-VSA"
-        required: false
-        default: false
-        type: boolean
-      run_unit_test:
-        description: "Run unit-test"
-        required: false
-        default: false
-        type: boolean
-
-env:
-  PYTHONUNBUFFERED: "1"
-
-
-concurrency:
-  group: pr-test-${{ github.ref }}
-  cancel-in-progress: true
-
-jobs:
-  change-filter:
-    runs-on: ubuntu-latest
-    if: ${{ github.event.pull_request.draft == false || github.event_name == 'workflow_dispatch' }}
-    outputs:
-      encoder-test: ${{ steps.filter.outputs.encoder-test }}
-      vae-test: ${{ steps.filter.outputs.vae-test }}
-      transformer-test: ${{ steps.filter.outputs.transformer-test }}
-      training-test: ${{ steps.filter.outputs.training-test }}
-      training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }}
-      precision-test-VSA: ${{ steps.filter.outputs.precision-test-VSA }}
-      unit-test: ${{ steps.filter.outputs.unit-test }}
-    steps:
-      - uses: actions/checkout@v4
-      - uses: dorny/paths-filter@v3
-        id: filter
-        with:
-          filters: |
-            # Define reusable path patterns
-            common-paths: &common-paths
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.10'
-              - 'docker/Dockerfile.python3.11'
-              - 'docker/Dockerfile.python3.12'
-            vsa-kernel-paths: &vsa-kernel-paths
-              - 'csrc/attn/video_sparse_attn/**'
-              - 'csrc/attn/video_sparse_attn/tk/**'
-              - 'csrc/attn/video_sparse_attn/setup.py'
-              - 'csrc/attn/video_sparse_attn/config_vsa.py'
-              - 'csrc/attn/video_sparse_attn/vsa.cpp'
-            vsa-paths: &vsa-paths
-              - 'fastvideo/**'
-              - *common-paths
-              - *vsa-kernel-paths
-            
-            # Actual tests
-            encoder-test:
-              - 'fastvideo/models/encoders/**'
-              - 'fastvideo/models/loader/**'
-              - 'fastvideo/tests/encoders/**'
-              - *common-paths
-            vae-test:
-              - 'fastvideo/models/vaes/**'
-              - 'fastvideo/models/loader/**'
-              - 'fastvideo/tests/vaes/**'
-              - *common-paths
-            transformer-test:
-              - 'fastvideo/models/dits/**'
-              - 'fastvideo/models/loader/**'
-              - 'fastvideo/tests/transformers/**'
-              - 'fastvideo/layers/**'
-              - 'fastvideo/attention/**'
-              - *common-paths
-            training-test:
-              - 'fastvideo/**'
-              - *common-paths
-            training-test-VSA:
-              - 'fastvideo/**'
-              - *common-paths
-              - *vsa-kernel-paths
-            precision-test-VSA:
-              - *common-paths
-              - *vsa-kernel-paths
-            unit-test:
-              - 'fastvideo/**'
-              - *common-paths
-
-  encoder-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.encoder-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_encoder_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "encoder-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 1
-      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/encoders -s"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  vae-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.vae-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_vae_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "vae-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 1
-      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/vaes -s"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  transformer-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.transformer-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_transformer_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "transformer-test"
-      gpu_type: "NVIDIA L40S"
-      gpu_count: 1
-      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/transformers -s"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-      
-  ssim-test:
-    needs: change-filter
-    if: >-
-      github.event_name != 'workflow_dispatch' || (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
-    strategy:
-      fail-fast: false
-      matrix:
-        python-version: [
-          # {version: "3.10", tag: "latest"},
-          # {version: "3.11", tag: "py3.11-latest"},
-          {version: "3.12", tag: "py3.12-latest"}
-        ]
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "ssim-test-py${{ matrix.python-version.version }}"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 2
-      volume_size: 200
-      disk_size: 200
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/ssim -vs"
-      timeout_minutes: 60
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  training-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "training-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 4
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/Vanilla -srP"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-
-  training-test-VSA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test-VSA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test_VSA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "training-test-VSA"
-      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: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/VSA -srP"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-
-  precision-test-VSA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-VSA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_VSA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "precision-test-VSA"
-      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_vsa.py"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  unit-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.unit-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_unit_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "unit-test"
-      gpu_type: "NVIDIA L40S"
-      gpu_count: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/dataset/ -vs && pytest ./fastvideo/workflow/ -vs && pytest ./fastvideo/entrypoints/ -vs"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  # nightly-test:
-  #   if: >-
-  #     (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
-  #   uses: ./.github/workflows/runpod-test.yml
-  #   with:
-  #     job_id: "nightly-test"
-  #     gpu_type: "NVIDIA A40"
-  #     gpu_count: 4
-  #     volume_size: 100
-  #     disk_size: 100
-  #     image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-  #     test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/nightly/test_e2e_overfit_single_sample.py -vs"
-  #     timeout_minutes: 30
-  #   secrets:
-  #     RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-  #     RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-  #     WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-
-  runpod-cleanup:
-    # Add other jobs to this list as you create them
-    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:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - 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", "precision-test-VSA"]'
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        run: python .github/scripts/runpod_cleanup.py
diff --git a/.github/workflows/fastvideo-publish.yml b/.github/workflows/publish-fastvideo.yml
similarity index 100%
rename from .github/workflows/fastvideo-publish.yml
rename to .github/workflows/publish-fastvideo.yml
diff --git a/.github/workflows/fastvideo-kernel-publish.yml b/.github/workflows/publish-kernel.yml
similarity index 100%
rename from .github/workflows/fastvideo-kernel-publish.yml
rename to .github/workflows/publish-kernel.yml
diff --git a/.github/workflows/runpod-test.yml b/.github/workflows/runpod-test.yml
deleted file mode 100644
index e4018c12a..000000000
--- a/.github/workflows/runpod-test.yml
+++ /dev/null
@@ -1,94 +0,0 @@
-name: RunPod Test
-
-on:
-  workflow_call:
-    inputs:
-      job_id:
-        required: true
-        type: string
-        description: "Unique identifier for this test job"
-      gpu_type:
-        required: true
-        type: string
-        description: "GPU type to use (e.g. NVIDIA A40, NVIDIA L40S)"
-      gpu_count:
-        required: true
-        type: number
-        description: "Number of GPUs to use"
-      volume_size:
-        required: false
-        type: number
-        default: 20
-        description: "Volume size in GB"
-      disk_size:
-        required: false
-        type: number
-        default: 20
-        description: "Disk size in GB"
-      image:
-        required: true
-        type: string
-        description: "Docker image to use"
-      test_command:
-        required: true
-        type: string
-        description: "Command to run tests"
-      timeout_minutes:
-        required: false
-        type: number
-        default: 30
-        description: "Timeout in minutes"
-    secrets:
-      RUNPOD_API_KEY:
-        required: true
-      RUNPOD_PRIVATE_KEY:
-        required: true
-      WANDB_API_KEY:
-        required: false
-
-jobs:
-  run-test:
-    runs-on: ubuntu-latest
-    environment: runpod-runners
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.12"
-
-      - name: Set up SSH key
-        run: |
-          mkdir -p ~/.ssh
-          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
-          chmod 600 ~/.ssh/id_rsa
-          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - name: Run tests on RunPod
-        env:
-          JOB_ID: ${{ inputs.job_id }}
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-        timeout-minutes: ${{ inputs.timeout_minutes }}
-        run: >-
-          python .github/scripts/runpod_api.py 
-          --gpu-type "${{ inputs.gpu_type }}"
-          --gpu-count ${{ inputs.gpu_count }}
-          --volume-size ${{ inputs.volume_size }}
-          --disk-size ${{ inputs.disk_size }}
-          --image "${{ inputs.image }}"
-          --test-command "${{ inputs.test_command }}"
-
-      - name: Terminate RunPod Instances
-        if: ${{ always() }}
-        env:
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          JOB_ID: ${{ inputs.job_id }}
-        run: python .github/scripts/runpod_cleanup.py 
\ No newline at end of file
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
deleted file mode 100644
index 9030e3f39..000000000
--- a/.github/workflows/test.yml
+++ /dev/null
@@ -1,31 +0,0 @@
-name: Run Tests
-
-on:
-  push:
-    branches: [ main ]
-  pull_request:
-    branches: [ main ]
-
-jobs:
-  test:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check out repository
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.12'  # or any version you need
-
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip setuptools wheel
-          pip install torch
-          pip install packaging ninja
-          pip install -e .
-          pip install pytest
-
-      - name: Run Pytest
-        run: |
-          pytest --ignore csrc/attn/test
diff --git a/.github/workflows/vsa-publish.yml b/.github/workflows/vsa-publish.yml
deleted file mode 100644
index c9f257a14..000000000
--- a/.github/workflows/vsa-publish.yml
+++ /dev/null
@@ -1,257 +0,0 @@
-name: Publish Video Sparse Attention Kernel to PyPI on Version Change
-
-on:
-  push:
-    branches:
-      - main
-    paths:
-      - "csrc/attn/video_sparse_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/video_sparse_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']
-          # For version reference https://pytorch.org/get-started/previous-versions/
-          torch-cuda:
-            - torch-version: '2.5.1'
-              cuda-version: '12.4.1'
-              torch-cuda-short: 'cu124'
-            - torch-version: '2.6.0'
-              cuda-version: '12.6.3'
-              torch-cuda-short: 'cu126'
-            - torch-version: '2.7.1'
-              cuda-version: '12.8.0'
-              torch-cuda-short: 'cu128'
-
-    steps:
-      - name: Free up disk space
-        run: |
-          echo "Initial disk space:"
-          df -h
-          
-          # Remove large directories
-          sudo rm -rf /usr/share/dotnet
-          sudo rm -rf /usr/local/lib/android
-          sudo rm -rf /opt/ghc
-          sudo rm -rf /usr/local/share/boost
-          sudo rm -rf /usr/share/swift
-          sudo rm -rf /usr/local/lib/node_modules
-          sudo rm -rf /usr/local/share/powershell
-          sudo rm -rf /usr/share/rust
-          sudo rm -rf /usr/local/.ghcup
-          
-          # Remove cached files
-          sudo rm -rf /var/lib/apt/lists/*
-          sudo rm -rf /var/cache/apt/archives/*
-          
-          echo "Disk space after cleanup:"
-          df -h
-
-      - name: Checkout
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install CUDA ${{ matrix.torch-cuda.cuda-version }}
-        uses: Jimver/cuda-toolkit@v0.2.21
-        id: cuda-toolkit
-        with:
-          cuda: ${{ matrix.torch-cuda.cuda-version }}
-          linux-local-args: '["--toolkit"]'
-          method: 'network'
-
-      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
-        run: |
-          sudo apt update
-          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
-          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-          # Allow Git to Access Safe Directory
-          git config --global --add safe.directory /__w/FastVideo/FastVideo
-
-          # Set CUDA environment variables
-          export CUDA_HOME=/usr/local/cuda-${{ matrix.torch-cuda.cuda-version }}
-          export PATH=${CUDA_HOME}/bin:${PATH}
-          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-          
-          # Verify installation
-          gcc --version
-          g++ --version
-          clang-11 --version
-          nvcc --version
-
-      - name: Install PyTorch ${{ matrix.torch-cuda.torch-version }}+cu${{ matrix.torch-cuda.cuda-version }}
-        run: |
-          pip install --upgrade pip
-          # 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
-          pip install --no-cache-dir torch==${{ matrix.torch-cuda.torch-version }} --index-url https://download.pytorch.org/whl/${{matrix.torch-cuda.torch-cuda-short}}
-          nvcc --version
-          python --version
-          python -c "import torch; print('PyTorch:', torch.__version__)"
-          python -c "import torch; print('CUDA:', torch.version.cuda)"
-          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
-
-      - name: Build wheel
-        run: |
-          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
-          
-          # 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/video_sparse_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/video_sparse_attn
-
-          CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
-          TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2)
-          # 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/video_sparse_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/video_sparse_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/video_sparse_attn/dist/
diff --git a/.gitignore b/.gitignore
index 5ed80a76b..dadc1b67e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -88,5 +88,6 @@ preprocess_output_text/
 
 .claude/
 .codex/
+.sisyphus/
 openspec/
 fastvideo/tests/ssim/reference_videos/**
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 08ccadbd9..31805cdb2 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -18,10 +18,8 @@ exclude: |
     fastvideo/train\.py|
     fastvideo/utils/.*|
     examples/.*|
-    .github/workflows/fastvideo-publish.yml|
-    .github/workflows/sta-publish.yml|
-    .github/workflows/vsa-publish.yml|
-    .github/workflows/build-image-template.yml|
+    .github/workflows/publish-fastvideo.yml|
+    .github/workflows/_template-build-image.yml|
     docs/source/inference/support_matrix.md
   )
 repos:
diff --git a/docs/README.md b/docs/README.md
index 1234fb599..2a0132700 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -36,4 +36,4 @@ Then open your browser to: http://localhost:8000
 
 ## Automatic Deployment
 
-Documentation is automatically built and deployed to GitHub Pages when changes are pushed to the `main` branch via the `.github/workflows/docs.yml` workflow.
+Documentation is automatically built and deployed to GitHub Pages when changes are pushed to the `main` branch via the `.github/workflows/infra-docs.yml` workflow.
diff --git a/docs/contributing/developer_env/runpod.md b/docs/contributing/developer_env/runpod.md
index b14b683d1..f3e7c5d77 100644
--- a/docs/contributing/developer_env/runpod.md
+++ b/docs/contributing/developer_env/runpod.md
@@ -1,65 +1,195 @@
+# RunPod Development Environment
 
-# 📦 Developing FastVideo on RunPod
+RunPod gives you on-demand cloud GPUs for FastVideo development. It's useful when you need a beefy GPU to test training runs, benchmark inference, or reproduce results without waiting for shared cluster time.
 
-You can easily use the FastVideo Pod Template on [RunPod](https://www.runpod.io) for development or experimentation.
+## Prerequisites
 
-## Creating a new pod
+- A [RunPod](https://www.runpod.io) account with billing configured
+- An SSH key pair. If you don't have one, generate it with `ssh-keygen -t ed25519`
+- Your public key (`~/.ssh/id_ed25519.pub`) ready to paste into RunPod
+
+## Step 1: Create a Pod
+
+**1. Verify your account**
+
+Make sure you're logged into the right RunPod account before spending credits.
 
-- Make sure you are using the correct RunPod account.
 ![RunPod Account Selection](../../assets/images/runpod_account.png)
 
-- Use "Additional Filters" to select CUDA 12.8.
+**2. Filter by CUDA version**
+
+Use "Additional Filters" to select CUDA 12.8.
+
 ![RunPod CUDA selection](../../assets/images/runpod_cuda.png)
 
-- Click "Deploy" and Pick a single A40 or RTX 4090 GPU.
+**3. Select a GPU**
+
+Click "Deploy" and pick a GPU. See [GPU Recommendations](#gpu-recommendations) below for guidance on which GPU to choose.
+
 ![RunPod GPU Selection](../../assets/images/runpod_deploy.png)
 
-- Select the "FastVideo" or "fastvideo-dev" Pod Template.
+**4. Pick the FastVideo template**
+
+Select the "FastVideo" or "fastvideo-dev" Pod Template. This pulls the pre-built image that includes all dependencies, Flash Attention, and a ready-to-use `uv` environment.
+
 ![RunPod Pod Template Selection](../../assets/images/runpod_create.png)
 
-- Set the Pod name to "`<name>-<FastVideo>-<date>`".
+**5. Name your pod**
+
+Use a memorable name like `yourname-fastvideo-2026-03-28`. This helps if you have multiple pods running.
+
+**6. Add a persistent volume (recommended)**
+
+Attach a network volume to `/root/.cache` or `/models` for storing downloaded model weights. Models can be 10-50 GB each, and re-downloading them every session wastes time and bandwidth.
+
+**7. Deploy**
+
+Click Deploy. The pod takes a few minutes to start while the image pulls. You'll see it transition to "Running" in your dashboard.
+
+## Step 2: Connect via SSH
+
+Once the pod is running, find the "SSH over exposed TCP" connection string in the pod dashboard.
 
-- Finally, once the pod is deployed (will take a few minutes as the image is being pulled), you can SSH into it using "SSH exposed over TCP". You'll need to use the matching private ssh key you provided.
 ![RunPod SSH](../../assets/images/runpod_ssh.png)
 
-## Working with the pod
-
-After SSH'ing into your pod, you'll find the correct `uv` environment already activated and you should be in /FastVideo/ directory. Make sure to use /FastVideo/ for all your work.
-
-To pull in the latest changes from the GitHub repo:
+Connect with:
 
 ```bash
-cd /FastVideo
-git pull
+ssh root@<pod-ip> -p <port> -i ~/.ssh/id_ed25519
 ```
 
-Run your development workflows as usual:
+RunPod also supports VS Code Remote SSH if you prefer an IDE.
 
-```bash
-# Run linters
-pre-commit run --all-files
+### Custom template (advanced)
 
-# Run tests
-pytest tests/
-```
-
-Make sure to push your changes back to the GitHub repo as nothing will be saved to the pod when it is terminated.
-
-After you are done with your work, you can terminate the pod by clicking the "Terminate" and "Delete" buttons. Remember if the pod is not completely deleted, Runpod will keep charging you for it.
-
-## Extra Information:
-If you need to customize the pod template this section has some useful information. For the most part you can leave the defaults of the FastVideo Pod Template.
-
-When creating your pod template, use this image:
+If you're setting up a pod from scratch instead of the FastVideo template, use this image:
 
 ```
 ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest
 ```
 
-Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):
+And paste this as the Container Start Command to enable SSH ([RunPod docs](https://docs.runpod.io/pods/configuration/use-ssh)):
 
 ```bash
 bash -c "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-server -y;mkdir -p ~/.ssh;cd $_;chmod 700 ~/.ssh;echo \"$PUBLIC_KEY\" >> authorized_keys;chmod 700 authorized_keys;service ssh start;sleep infinity"
 ```
 
 ![RunPod template configuration](../../assets/images/runpod_template.png)
+
+## Step 3: Set Up FastVideo
+
+After SSH'ing in, the `uv` virtual environment at `/opt/venv` is already activated (configured in `.bashrc` and `.profile`). You land in the `/FastVideo` directory.
+
+**Clone or pull the repo**
+
+If the pod already has the FastVideo repo:
+
+```bash
+cd /FastVideo
+git pull
+```
+
+If starting from a blank pod:
+
+```bash
+git clone https://github.com/hao-ai-lab/FastVideo.git /FastVideo
+cd /FastVideo
+```
+
+**Install the package**
+
+```bash
+uv pip install -e .[dev]
+```
+
+The Docker image already includes Flash Attention and most heavy dependencies, so this is fast.
+
+**Build the custom kernels (optional)**
+
+VSA and STA attention kernels aren't in the Docker image by default. Build them if you're working on attention backends or need maximum inference performance:
+
+```bash
+cd /FastVideo/fastvideo-kernel
+./build.sh
+```
+
+The build script detects your GPU architecture automatically. An A100 or H100 takes about 5-10 minutes.
+
+**Verify the setup**
+
+```bash
+cd /FastVideo
+python -c "import fastvideo; print('OK')"
+pytest tests/ -q --no-header
+```
+
+## Development Workflow
+
+### Editing code on RunPod
+
+Two common approaches:
+
+**Option A: Edit on RunPod directly**
+
+Use VS Code Remote SSH or `vim`/`nano` on the pod. Commit and push when you're ready:
+
+```bash
+cd /FastVideo
+git add .
+git commit -m "your change"
+git push
+```
+
+**Option B: Edit locally, sync to RunPod**
+
+Work in your local repo, then pull on the pod:
+
+```bash
+# On RunPod:
+cd /FastVideo
+git pull
+```
+
+This keeps your local tools (editor, linters) intact while running GPU workloads on the pod.
+
+### Running linters and tests
+
+```bash
+# Lint
+pre-commit run --all-files
+
+# Full test suite
+pytest tests/
+
+# Just package tests
+pytest fastvideo/tests/ -v
+```
+
+### Storing models
+
+If you attached a persistent volume, point your model downloads there:
+
+```bash
+export HF_HOME=/models/huggingface
+export TRANSFORMERS_CACHE=/models/huggingface
+```
+
+Add these to `/root/.bashrc` so they persist across SSH sessions. The volume survives pod termination, so you only download models once.
+
+### Terminating the pod
+
+When you're done, push any commits you want to keep. RunPod does not save pod storage after termination.
+
+Go to your RunPod dashboard, click "Terminate", then "Delete". A pod that's stopped but not deleted still charges you for storage. Fully delete it to stop all charges.
+
+## GPU Recommendations
+
+| GPU | VRAM | Good for |
+|-----|------|----------|
+| RTX 4090 | 24 GB | Inference testing, small model fine-tuning, quick iteration |
+| A40 | 48 GB | Mid-size training runs, 480p video generation |
+| A100 (40 GB) | 40 GB | Multi-GPU inference, training with sequence parallelism |
+| A100 (80 GB) | 80 GB | Large model training, 720p+ video generation |
+| H100 | 80 GB | Heavy training, benchmarking, kernel development |
+
+For most development work, a single RTX 4090 or A40 is sufficient and cost-effective. Use A100/H100 when you need to reproduce training results at scale or test multi-GPU features.
diff --git a/fastvideo/tests/encoders/test_clip_encoder.py b/fastvideo/tests/encoders/test_clip_encoder.py
index cdf4162bd..266a9ef5d 100644
--- a/fastvideo/tests/encoders/test_clip_encoder.py
+++ b/fastvideo/tests/encoders/test_clip_encoder.py
@@ -21,9 +21,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder_2")
 TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer_2")
 
@@ -34,13 +32,12 @@ def test_clip_encoder():
     Tests compatibility between two different implementations for loading text encoders:
     1. load_text_encoder from fastvideo.models.hunyuan.text_encoder
     2. TextEncoderLoader from fastvideo.models.loader
-    
+
     The test verifies that both implementations:
     - Load models with the same weights and parameters
     - Produce nearly identical outputs for the same input prompts
     """
-    args = FastVideoArgs(model_path="openai/clip-vit-large-patch14",
-                         pipeline_config=HunyuanConfig())
+    args = FastVideoArgs(model_path="openai/clip-vit-large-patch14", pipeline_config=HunyuanConfig())
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
     logger.info("Loading models from %s", args.model_path)
@@ -55,6 +52,7 @@ def test_clip_encoder():
     model1 = CLIPTextModel.from_pretrained(TEXT_ENCODER_PATH).to(torch.float16).to(device).eval()
 
     from fastvideo.models.loader.component_loader import TextEncoderLoader
+
     loader = TextEncoderLoader()
     model2 = loader.load(TEXT_ENCODER_PATH, args)
 
@@ -90,10 +88,7 @@ def test_clip_encoder():
     tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
 
     # Test with some sample prompts
-    prompts = [
-        "a photo of a cat", "a beautiful landscape with mountains",
-        "an astronaut riding a horse on the moon"
-    ]
+    prompts = ["a photo of a cat", "a beautiful landscape with mountains", "an astronaut riding a horse on the moon"]
 
     logger.info("Testing CLIP text encoder with sample prompts")
 
@@ -102,14 +97,11 @@ def test_clip_encoder():
             logger.info("Testing prompt: '%s'", prompt)
 
             # Tokenize the prompt
-            tokens = tokenizer(prompt,
-                               padding="max_length",
-                               max_length=77,
-                               truncation=True,
-                               return_tensors="pt").to(device)
+            tokens = tokenizer(prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt").to(
+                device
+            )
             # Get embeddings from our implementation
-            outputs1 = model1(input_ids=tokens.input_ids,
-                              output_hidden_states=True)
+            outputs1 = model1(input_ids=tokens.input_ids, output_hidden_states=True)
 
             logger.info("Testing model2")
             print("--------------------------------")
@@ -118,24 +110,25 @@ def test_clip_encoder():
                 outputs2 = model2(
                     input_ids=tokens.input_ids,
                     # attention_mask=tokens.attention_mask,
-                    output_hidden_states=True)
+                    output_hidden_states=True,
+                )
 
             # Compare last hidden states
-            last_hidden_state1 = outputs1.last_hidden_state[
-                tokens.attention_mask == 1]
-            last_hidden_state2 = outputs2.last_hidden_state[
-                tokens.attention_mask == 1]
+            last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask == 1]
+            last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask == 1]
             # print("last_hidden_state1", last_hidden_state1)
             # print("last_hidden_state2", last_hidden_state2)
 
-            assert last_hidden_state1.shape == last_hidden_state2.shape, \
+            assert last_hidden_state1.shape == last_hidden_state2.shape, (
                 f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}"
+            )
             # Compare pooler outputs
             pooler_output1 = outputs1.pooler_output
             pooler_output2 = outputs2.pooler_output
 
-            assert pooler_output1.shape == pooler_output2.shape, \
+            assert pooler_output1.shape == pooler_output2.shape, (
                 f"Pooler output shapes don't match: {pooler_output1.shape} vs {pooler_output2.shape}"
+            )
 
             assert_close(pooler_output1, pooler_output2, atol=1e-2, rtol=1e-3)
             assert_close(last_hidden_state1, last_hidden_state2, atol=1e-2, rtol=1e-3)
diff --git a/fastvideo/tests/encoders/test_llama_encoder.py b/fastvideo/tests/encoders/test_llama_encoder.py
index 4c4fa5c22..e04ac5ba4 100644
--- a/fastvideo/tests/encoders/test_llama_encoder.py
+++ b/fastvideo/tests/encoders/test_llama_encoder.py
@@ -15,15 +15,14 @@ from fastvideo.utils import maybe_download_model
 from fastvideo.configs.models.encoders import LlamaConfig
 from torch.distributed.tensor import DTensor
 from torch.testing import assert_close
+
 logger = init_logger(__name__)
 
 os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder")
 TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer")
 
@@ -34,13 +33,12 @@ def test_llama_encoder():
     Tests compatibility between two different implementations for loading text encoders:
     1. load_text_encoder from fastvideo.models.hunyuan.text_encoder
     2. TextEncoderLoader from fastvideo.models.loader
-    
+
     The test verifies that both implementations:
     - Load models with the same weights and parameters
     - Produce nearly identical outputs for the same input prompts
     """
-    args = FastVideoArgs(model_path="meta-llama/Llama-2-7b-hf",
-                         pipeline_config=HunyuanConfig())
+    args = FastVideoArgs(model_path="meta-llama/Llama-2-7b-hf", pipeline_config=HunyuanConfig())
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -68,15 +66,15 @@ def test_llama_encoder():
     logger.info("Model1 has %d parameters", len(params1))
     logger.info("Model2 has %d parameters", len(params2))
 
-
     # check if embed_tokens are the same
     device = model1.embed_tokens.weight.device
-    assert torch.allclose(model1.embed_tokens.weight,
-                          model2.embed_tokens.weight.to_local().to(device) if isinstance(model2.embed_tokens.weight, DTensor) else model2.embed_tokens.weight.to(device))
-    weights = [
-        "layers.{}.input_layernorm.weight",
-        "layers.{}.post_attention_layernorm.weight"
-    ]
+    assert torch.allclose(
+        model1.embed_tokens.weight,
+        model2.embed_tokens.weight.to_local().to(device)
+        if isinstance(model2.embed_tokens.weight, DTensor)
+        else model2.embed_tokens.weight.to(device),
+    )
+    weights = ["layers.{}.input_layernorm.weight", "layers.{}.post_attention_layernorm.weight"]
     for layer_idx in range(hf_config.num_hidden_layers):
         for w in weights:
             name1 = w.format(layer_idx)
@@ -121,33 +119,29 @@ def test_llama_encoder():
             logger.info("Testing prompt: '%s'", prompt)
 
             # Tokenize the prompt
-            tokens = tokenizer(prompt,
-                               padding="max_length",
-                               max_length=512,
-                               truncation=True,
-                               return_tensors="pt").to(device)
+            tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True, return_tensors="pt").to(
+                device
+            )
 
             # Get outputs from our implementation
             # filter out padding input_ids
             # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1]
             # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1]
-            outputs1 = model1(input_ids=tokens.input_ids,
-                              output_hidden_states=True)
+            outputs1 = model1(input_ids=tokens.input_ids, output_hidden_states=True)
             print("--------------------------------")
             logger.info("Testing model2")
 
             # Get outputs from HuggingFace implementation
             with set_forward_context(current_timestep=0, attn_metadata=None):
-                outputs2 = model2(input_ids=tokens.input_ids,
-                                  attention_mask=tokens.attention_mask,
-                                  output_hidden_states=True)
+                outputs2 = model2(
+                    input_ids=tokens.input_ids, attention_mask=tokens.attention_mask, output_hidden_states=True
+                )
 
             # Compare last hidden states
-            last_hidden_state1 = outputs1.last_hidden_state[
-                tokens.attention_mask == 1]
-            last_hidden_state2 = outputs2.last_hidden_state[
-                tokens.attention_mask == 1]
+            last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask == 1]
+            last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask == 1]
 
-            assert last_hidden_state1.shape == last_hidden_state2.shape, \
+            assert last_hidden_state1.shape == last_hidden_state2.shape, (
                 f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}"
+            )
             assert_close(last_hidden_state1, last_hidden_state2, atol=1e-1, rtol=1e-4)
diff --git a/fastvideo/tests/transformers/test_hunyuanvideo.py b/fastvideo/tests/transformers/test_hunyuanvideo.py
index edd737530..62af0285d 100644
--- a/fastvideo/tests/transformers/test_hunyuanvideo.py
+++ b/fastvideo/tests/transformers/test_hunyuanvideo.py
@@ -7,9 +7,7 @@ import pytest
 import torch
 
 from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.distributed.parallel_state import (
-    get_sp_parallel_rank,
-    get_sp_world_size)
+from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
 from fastvideo.logger import init_logger
 from fastvideo.models.loader.component_loader import TransformerLoader
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -24,9 +22,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
 CONFIG_PATH = os.path.join(TRANSFORMER_PATH, "config.json")
 
@@ -41,9 +37,7 @@ REFERENCE_LATENT = 93.23263549804688
 @pytest.mark.usefixtures("distributed_setup")
 @pytest.mark.skip(reason="Skipping test_hunyuanvideo_distributed")
 def test_hunyuanvideo_distributed():
-    logger.info(
-        f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}"
-    )
+    logger.info(f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}")
 
     torch.cuda.set_device(f"cuda:{LOCAL_RANK}")
 
@@ -51,9 +45,7 @@ def test_hunyuanvideo_distributed():
     sp_rank = get_sp_parallel_rank()
     sp_world_size = get_sp_world_size()
 
-    logger.info(
-        f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}"
-    )
+    logger.info(f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}")
 
     config = json.load(open(CONFIG_PATH))
     # remove   "_class_name": "HunyuanVideoTransformer3DModel",   "_diffusers_version": "0.32.0.dev0",
@@ -62,9 +54,11 @@ def test_hunyuanvideo_distributed():
     config.pop("_diffusers_version")
 
     precision_str = "bf16"
-    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         dit_cpu_offload=True,
-                         pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str))
+    args = FastVideoArgs(
+        model_path=TRANSFORMER_PATH,
+        dit_cpu_offload=True,
+        pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str),
+    )
     args.device = torch.device(f"cuda:{LOCAL_RANK}")
 
     loader = TransformerLoader()
@@ -78,23 +72,12 @@ def test_hunyuanvideo_distributed():
     device = torch.device(f"cuda:{LOCAL_RANK}")
 
     # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(batch_size,
-                                16,
-                                8,
-                                16,
-                                16,
-                                device=device,
-                                dtype=torch.bfloat16)
+    hidden_states = torch.randn(batch_size, 16, 8, 16, 16, device=device, dtype=torch.bfloat16)
     chunk_per_rank = hidden_states.shape[2] // sp_world_size
-    hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) *
-                                  chunk_per_rank]
+    hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank : (sp_rank + 1) * chunk_per_rank]
 
     # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(batch_size,
-                                        seq_len + 1,
-                                        4096,
-                                        device=device,
-                                        dtype=torch.bfloat16)
+    encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 4096, device=device, dtype=torch.bfloat16)
 
     # Timestep
     timestep = torch.tensor([500], device=device, dtype=torch.bfloat16)
@@ -117,7 +100,5 @@ def test_hunyuanvideo_distributed():
 
     # Check if latents are similar
     diff_output_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(
-        f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}"
-    )
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
     assert diff_output_latents < 1e-4, f"Output latents differ significantly: max diff = {diff_output_latents}"
diff --git a/fastvideo/tests/transformers/test_wanvideo.py b/fastvideo/tests/transformers/test_wanvideo.py
index 1028138b5..9f93783f7 100644
--- a/fastvideo/tests/transformers/test_wanvideo.py
+++ b/fastvideo/tests/transformers/test_wanvideo.py
@@ -23,9 +23,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
 
 
@@ -34,22 +32,25 @@ def test_wan_transformer():
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     precision = torch.bfloat16
     precision_str = "bf16"
-    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         dit_cpu_offload=True,
-                         pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
+    args = FastVideoArgs(
+        model_path=TRANSFORMER_PATH,
+        dit_cpu_offload=True,
+        pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str),
+    )
     args.device = device
 
     loader = TransformerLoader()
     model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision)
 
-    model1 = WanTransformer3DModel.from_pretrained(
-        TRANSFORMER_PATH, device=device,
-        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+    model1 = (
+        WanTransformer3DModel.from_pretrained(TRANSFORMER_PATH, device=device, torch_dtype=precision)
+        .to(device, dtype=precision)
+        .requires_grad_(False)
+    )
 
     total_params = sum(p.numel() for p in model1.parameters())
     # Calculate weight sum for model1 (converting to float64 to avoid overflow)
-    weight_sum_model1 = sum(
-        p.to(torch.float64).sum().item() for p in model1.parameters())
+    weight_sum_model1 = sum(p.to(torch.float64).sum().item() for p in model1.parameters())
     # Also calculate mean for more stable comparison
     weight_mean_model1 = weight_sum_model1 / total_params
     logger.info("Model 1 weight sum: %s", weight_sum_model1)
@@ -57,8 +58,7 @@ def test_wan_transformer():
 
     # Calculate weight sum for model2 (converting to float64 to avoid overflow)
     total_params_model2 = sum(p.numel() for p in model2.parameters())
-    weight_sum_model2 = sum(
-        p.to(torch.float64).sum().item() for p in model2.parameters())
+    weight_sum_model2 = sum(p.to(torch.float64).sum().item() for p in model2.parameters())
     # Also calculate mean for more stable comparison
     weight_mean_model2 = weight_sum_model2 / total_params_model2
     logger.info("Model 2 weight sum: %s", weight_sum_model2)
@@ -78,20 +78,10 @@ def test_wan_transformer():
     seq_len = 30
 
     # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(batch_size,
-                                16,
-                                21,
-                                160,
-                                90,
-                                device=device,
-                                dtype=precision)
+    hidden_states = torch.randn(batch_size, 16, 21, 160, 90, device=device, dtype=precision)
 
     # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(batch_size,
-                                        seq_len + 1,
-                                        4096,
-                                        device=device,
-                                        dtype=precision)
+    encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 4096, device=device, dtype=precision)
 
     # Timestep
     timestep = torch.tensor([500], device=device, dtype=precision)
@@ -100,7 +90,7 @@ def test_wan_transformer():
         data_type="dummy",
     )
 
-    with torch.amp.autocast('cuda', dtype=precision):
+    with torch.amp.autocast("cuda", dtype=precision):
         output1 = model1(
             hidden_states=hidden_states,
             encoder_hidden_states=encoder_hidden_states,
@@ -108,13 +98,13 @@ def test_wan_transformer():
             return_dict=False,
         )[0]
         with set_forward_context(
-                current_timestep=0,
-                attn_metadata=None,
-                forward_batch=forward_batch,
+            current_timestep=0,
+            attn_metadata=None,
+            forward_batch=forward_batch,
         ):
-            output2 = model2(hidden_states=hidden_states,
-                             encoder_hidden_states=encoder_hidden_states,
-                             timestep=timestep)
+            output2 = model2(
+                hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, timestep=timestep
+            )
 
     # Check if outputs have the same shape
     assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
diff --git a/fastvideo/tests/vaes/test_hunyuan_vae.py b/fastvideo/tests/vaes/test_hunyuan_vae.py
index 6a5c422d9..c7e3408b5 100644
--- a/fastvideo/tests/vaes/test_hunyuan_vae.py
+++ b/fastvideo/tests/vaes/test_hunyuan_vae.py
@@ -9,6 +9,7 @@ from safetensors.torch import load_file
 
 from fastvideo.configs.pipelines import PipelineConfig
 from fastvideo.logger import init_logger
+
 # from fastvideo.models.vaes.hunyuanvae import (
 #     AutoencoderKLHunyuanVideo as MyHunyuanVAE)
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -22,9 +23,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 VAE_PATH = os.path.join(MODEL_PATH, "vae")
 CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
 
@@ -42,30 +41,28 @@ def test_hunyuan_vae():
     device = torch.device("cuda:0")
     precision = torch.bfloat16
     precision_str = "bf16"
-    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=HunyuanVAEConfig(), vae_precision=precision_str))
+    args = FastVideoArgs(
+        model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=HunyuanVAEConfig(), vae_precision=precision_str)
+    )
     args.device = device
     args.vae_cpu_offload = False
 
     loader = VAELoader()
     model = loader.load(VAE_PATH, args)
 
-    model.enable_tiling(tile_sample_min_height=32,
-                         tile_sample_min_width=32,
-                         tile_sample_min_num_frames=8,
-                         tile_sample_stride_height=16,
-                         tile_sample_stride_width=16,
-                         tile_sample_stride_num_frames=4)
+    model.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_min_num_frames=8,
+        tile_sample_stride_height=16,
+        tile_sample_stride_width=16,
+        tile_sample_stride_num_frames=4,
+    )
 
     batch_size = 1
 
     # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size,
-                               3,
-                               21,
-                               64,
-                               64,
-                               device=device,
-                               dtype=torch.bfloat16)
+    input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16)
 
     # Disable gradients for inference
     with torch.no_grad():
@@ -81,8 +78,5 @@ def test_hunyuan_vae():
         raise ValueError(f"Unknown device: {device_name}")
 
     diff_encoded_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(
-        f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}"
-    )
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
     assert diff_encoded_latents < 1e-4, f"Encoded latents differ significantly: max diff = {diff_encoded_latents}"
-
diff --git a/fastvideo/tests/vaes/test_wan_vae.py b/fastvideo/tests/vaes/test_wan_vae.py
index f72417b45..31a1175b8 100644
--- a/fastvideo/tests/vaes/test_wan_vae.py
+++ b/fastvideo/tests/vaes/test_wan_vae.py
@@ -20,9 +20,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 VAE_PATH = os.path.join(MODEL_PATH, "vae")
 
 
@@ -31,28 +29,23 @@ def test_wan_vae():
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     precision = torch.float32
     precision_str = "fp32"
-    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str))
+    args = FastVideoArgs(
+        model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str)
+    )
     args.device = device
     args.vae_cpu_offload = False
 
     loader = VAELoader()
     model2 = loader.load(VAE_PATH, args)
-    assert model2.use_feature_cache # Default to use the original WanVAE algorithm
+    assert model2.use_feature_cache  # Default to use the original WanVAE algorithm
 
-    model1 = AutoencoderKLWan.from_pretrained(
-        VAE_PATH, torch_dtype=precision).to(device).eval()
+    model1 = AutoencoderKLWan.from_pretrained(VAE_PATH, torch_dtype=precision).to(device).eval()
 
     # Create identical inputs for both models
     batch_size = 1
 
     # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size,
-                               3,
-                               81,
-                               32,
-                               32,
-                               device=device,
-                               dtype=precision)
+    input_tensor = torch.randn(batch_size, 3, 81, 32, 32, device=device, dtype=precision)
     # latent_tensor = torch.randn(batch_size,
     #                             16,
     #                             21,
@@ -69,18 +62,22 @@ def test_wan_vae():
         print("--------------------------------")
         latent2 = model2.encode(input_tensor)
         # Check if latents have the same shape
-        assert latent1.mean.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
+        assert latent1.mean.shape == latent2.mean.shape, (
+            f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
+        )
         # Check if latents are similar
         assert_close(latent1.mean, latent2.mean, atol=1e-4, rtol=1e-4)
         # Test decoding
         logger.info("Testing decoding...")
         latent1_tensor = latent1.mode()
-        mean1 = (torch.tensor(model1.config.latents_mean).view(
-            1, model1.config.z_dim, 1, 1, 1).to(input_tensor.device,
-                                                input_tensor.dtype))
-        std1 = (1.0 / torch.tensor(model1.config.latents_std).view(
-            1, model1.config.z_dim, 1, 1, 1)).to(input_tensor.device,
-                                                input_tensor.dtype)
+        mean1 = (
+            torch.tensor(model1.config.latents_mean)
+            .view(1, model1.config.z_dim, 1, 1, 1)
+            .to(input_tensor.device, input_tensor.dtype)
+        )
+        std1 = (1.0 / torch.tensor(model1.config.latents_std).view(1, model1.config.z_dim, 1, 1, 1)).to(
+            input_tensor.device, input_tensor.dtype
+        )
         latent1_tensor = latent1_tensor / std1 + mean1
         output1 = model1.decode(latent1_tensor).sample