[Feature][Training]Update example fine-tuning scripts to enable gradient checkpointing (#618)

Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: Wenxuan Tan <wenxuan.tan@wisc.edu>

.buildkite/pipeline.yml

@@ -0,0 +1,138 @@
+env:
+  IMAGE_VERSION: "py3.12-latest"
+  BUILDKITE_CLEAN_CHECKOUT: true
+
+steps:
+  - label: "pre-commit"
+    command: ".buildkite/scripts/pre_commit.sh"
+    agents:
+      queue: "default"
+
+  - wait
+
+  - label: "Trigger Tests"
+    plugins:
+      - monorepo-diff#v1.4.0:
+          diff: "git diff --name-only $BUILDKITE_PULL_REQUEST_BASE_BRANCH...HEAD"
+          watch:
+            - path: 
+                - "fastvideo/v1/models/encoders/**"
+                - "fastvideo/v1/models/loader/**"
+                - "fastvideo/v1/tests/encoders/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Encoder Tests"
+                env:
+                  - TEST_TYPE=encoder
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/models/vaes/**"
+                - "fastvideo/v1/models/loader/**"
+                - "fastvideo/v1/tests/vaes/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "VAE Tests"
+                env:
+                  - TEST_TYPE=vae
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/models/dits/**"
+                - "fastvideo/v1/models/loader/**"
+                - "fastvideo/v1/tests/transformers/**"
+                - "fastvideo/v1/layers/**"
+                - "fastvideo/v1/attention/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Transformer Tests"
+                env:
+                  - TEST_TYPE=transformer
+                agents:
+                  queue: "default"
+            - path: 
+                - "fastvideo/v1/**/*.py"
+              config:
+                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                label: "SSIM Tests"
+                env:
+                  - TEST_TYPE=ssim
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Training Tests"
+                env:
+                  - TEST_TYPE=training
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "csrc/attn/vsa/**"
+                - "csrc/attn/tk/**"
+                - "csrc/attn/setup_vsa.py"
+                - "csrc/attn/config_vsa.py"
+                - "csrc/attn/vsa.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Training Tests VSA"
+                env:
+                  - TEST_TYPE=training_vsa
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/v1/**"
+                - "csrc/attn/st_attn/**"
+                - "csrc/attn/setup_sta.py"
+                - "csrc/attn/config_sta.py"
+                - "csrc/attn/st_attn.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Inference Tests STA"
+                env:
+                  - TEST_TYPE=inference_sta
+                agents:
+                  queue: "default"
+            - path:
+                - "csrc/attn/st_attn/**"
+                - "csrc/attn/setup_sta.py"
+                - "csrc/attn/config_sta.py"
+                - "csrc/attn/st_attn.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Precision Tests STA"
+                env:
+                  - TEST_TYPE=precision_sta
+                agents:
+                  queue: "default"
+            - path:
+                - "csrc/attn/vsa/**"
+                - "csrc/attn/tk/**"
+                - "csrc/attn/setup_vsa.py"
+                - "csrc/attn/config_vsa.py"
+                - "csrc/attn/vsa.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Precision Tests VSA"
+                env:
+                  - TEST_TYPE=precision_vsa
+                agents:
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -0,0 +1,117 @@
+#!/bin/bash
+set -uo pipefail
+
+log() {
+  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
+}
+
+log "=== Starting Modal test execution ==="
+
+# Change to the project directory
+cd "$(dirname "$0")/../.."
+PROJECT_ROOT=$(pwd)
+log "Project root: $PROJECT_ROOT"
+
+# Install Modal if not available
+if ! python3 -m modal --version &> /dev/null; then
+    log "Modal not found, installing..."
+    python3 -m pip install modal
+    
+    # Verify installation
+    if ! python3 -m modal --version &> /dev/null; then
+        log "Error: Failed to install modal. Please install it manually."
+        exit 1
+    fi
+fi
+
+log "modal version: $(python3 -m modal --version)"
+
+# Set up Modal authentication using Buildkite secrets
+log "Setting up Modal authentication from Buildkite secrets..."
+MODAL_TOKEN_ID=$(buildkite-agent secret get modal_token_id)
+MODAL_TOKEN_SECRET=$(buildkite-agent secret get modal_token_secret)
+
+WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
+
+WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
+
+if [ -n "$MODAL_TOKEN_ID" ] && [ -n "$MODAL_TOKEN_SECRET" ]; then
+    log "Retrieved Modal credentials from Buildkite secrets"
+    python3 -m modal token set --token-id "$MODAL_TOKEN_ID" --token-secret "$MODAL_TOKEN_SECRET" --profile buildkite-ci --activate --verify
+    if [ $? -eq 0 ]; then
+        log "Modal authentication successful"
+    else
+        log "Error: Failed to set Modal credentials"
+        exit 1
+    fi
+else
+    log "Error: Could not retrieve Modal credentials from Buildkite secrets."
+    log "Please ensure 'modal_token_id' and 'modal_token_secret' secrets are set in Buildkite."
+    exit 1
+fi
+
+MODAL_TEST_FILE="fastvideo/v1/tests/modal/pr_test.py"
+
+if [ -z "${TEST_TYPE:-}" ]; then
+    log "Error: TEST_TYPE environment variable is not set"
+    exit 1
+fi
+log "Test type: $TEST_TYPE"
+
+MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$BUILDKITE_PULL_REQUEST IMAGE_VERSION=$IMAGE_VERSION"
+
+case "$TEST_TYPE" in
+    "encoder")
+        log "Running encoder tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
+        ;;
+    "vae")
+        log "Running VAE tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
+        ;;
+    "transformer")
+        log "Running transformer tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
+        ;;
+    "ssim")
+        log "Running SSIM tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
+        ;;
+    "training")
+        log "Running training tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests"
+        ;;
+    "training_vsa")
+        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"
+        ;;
+    "precision_sta")
+        log "Running precision STA tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_STA"
+        ;;
+    "precision_vsa")
+        log "Running precision VSA tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_VSA"
+        ;;
+    *)
+        log "Error: Unknown test type: $TEST_TYPE"
+        exit 1
+        ;;
+esac
+
+log "Executing: $MODAL_COMMAND"
+eval "$MODAL_COMMAND"
+TEST_EXIT_CODE=$?
+
+if [ $TEST_EXIT_CODE -eq 0 ]; then
+    log "Modal test completed successfully"
+else
+    log "Error: Modal test failed with exit code: $TEST_EXIT_CODE"
+fi
+
+log "=== Test execution completed with exit code: $TEST_EXIT_CODE ==="
+exit $TEST_EXIT_CODE

.buildkite/scripts/pre_commit.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+set -uo pipefail
+
+log() {
+  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
+}
+
+log "=== Starting pre-commit checks ==="
+
+cd "$(dirname "$0")/../.."
+PROJECT_ROOT=$(pwd)
+log "Project root: $PROJECT_ROOT"
+
+if ! python3 -m pre_commit --version &> /dev/null; then
+    log "pre-commit not found, installing..."
+    python3 -m pip install --user pre-commit==4.0.1
+    
+    if ! python3 -m pre_commit --version &> /dev/null; then
+        log "Error: Failed to install pre-commit."
+        exit 1
+    fi
+fi
+
+log "Pre-commit version: $(python3 -m pre_commit --version)"
+
+log "Installing/updating pre-commit hooks..."
+python3 -m pre_commit install --install-hooks
+
+log "Running pre-commit checks on all files..."
+python3 -m pre_commit run --all-files
+PRE_COMMIT_EXIT_CODE=$?
+
+if [ $PRE_COMMIT_EXIT_CODE -eq 0 ]; then
+    log "Pre-commit checks completed successfully"
+else
+    log "Error: Pre-commit checks failed with exit code: $PRE_COMMIT_EXIT_CODE"
+fi
+
+log "=== Pre-commit checks completed with exit code: $PRE_COMMIT_EXIT_CODE ==="
+exit $PRE_COMMIT_EXIT_CODE 

.github/ISSUE_TEMPLATE/1-bug-report.yml

@@ -4,14 +4,6 @@ title: "[Bug] "
 labels: ['Bug']
 
 body:
-- type: textarea
-  attributes:
-    label: Environment
-    description: |
-      Please share your environment with us. You can run the command **python fastvideo/utils/env_utils.py** and copy-paste its output below.
-    placeholder: FastVideo version, platform, python version, cuda version...
-  validations:
-    required: true
 - type: textarea
   attributes:
     label: Describe the bug
@@ -25,5 +17,13 @@ body:
       What command or script did you run? Which **model** are you using?
     placeholder: |
       A placeholder for the command.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Environment
+    description: |
+      Please share your environment with us. You can run the command **python fastvideo/utils/collect_env.py** and copy-paste its output below.
+    placeholder: FastVideo version, platform, python version, cuda version...
   validations:
     required: true

.github/scripts/runpod_api.py

@@ -0,0 +1,249 @@
+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()

.github/scripts/runpod_cleanup.py

@@ -0,0 +1,90 @@
+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()

.github/workflows/build-image-template.yml

@@ -0,0 +1,106 @@
+name: Build Image Template
+
+on:
+  workflow_call:
+    inputs:
+      python_version:
+        required: true
+        type: string
+      dockerfile_path:
+        required: true
+        type: string
+      tag_suffix:
+        required: true
+        type: string
+
+jobs:
+  build-and-push:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      
+      - name: Free up disk space
+        run: |
+          # Display initial space
+          echo "Initial disk space:"
+          df -h
+          
+          # Remove large directories directly
+          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/*
+          
+          # Clean Docker
+          docker system prune -af --volumes
+          
+          # Display available space after cleanup
+          echo "Disk space after cleanup:"
+          df -h
+      
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+      
+      - name: Prepare tags
+        id: prepare-tags
+        run: |
+          SHORT_SHA=$(echo ${{ github.sha }} | cut -c1-7)
+          
+          TAGS="type=raw,value=${{ inputs.tag_suffix }}-latest"
+          TAGS="${TAGS}\ntype=raw,value=${{ inputs.tag_suffix }}-sha-${SHORT_SHA}"
+          
+          # Set Python 3.10 as the default image
+          if [[ "${{ inputs.python_version }}" == "3.10" ]]; then
+            TAGS="${TAGS}\ntype=raw,value=latest"
+          fi
+          
+          {
+            echo "tags<<EOF"
+            echo -e "$TAGS"
+            echo "EOF" 
+          } >> $GITHUB_OUTPUT
+      
+      - name: Extract metadata for Docker
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ghcr.io/${{ github.repository }}/fastvideo-dev
+          tags: ${{ steps.prepare-tags.outputs.tags }}
+      
+      - name: Build and push Docker image
+        id: build-push
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: ${{ inputs.dockerfile_path }}
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+      
+      - name: Success message
+        run: |
+          echo "✅ Python ${{ inputs.python_version }} image successfully built and pushed to ghcr.io/${{ github.repository }}/fastvideo-dev:${{ inputs.tag_suffix }}-latest"
+          echo "To run tests with this image, manually trigger the 'Run Tests' workflow." 

.github/workflows/build-image.yml

@@ -0,0 +1,52 @@
+name: Build and Push Docker Images
+
+on:
+  workflow_dispatch:
+    inputs:
+      python_3_10:
+        description: 'Build Python 3.10 image'
+        required: false
+        default: false
+        type: boolean
+      python_3_11:
+        description: 'Build Python 3.11 image'
+        required: false
+        default: false
+        type: boolean
+      python_3_12:
+        description: 'Build Python 3.12 image'
+        required: false
+        default: false
+        type: boolean
+
+permissions:
+  contents: read
+  packages: write
+
+jobs:
+  build-python-3-10:
+    if: ${{ github.event.inputs.python_3_10 == 'true' }}
+    uses: ./.github/workflows/build-image-template.yml
+    with:
+      python_version: '3.10'
+      dockerfile_path: docker/Dockerfile.python3.10
+      tag_suffix: py3.10
+    secrets: inherit
+
+  build-python-3-11:
+    if: ${{ github.event.inputs.python_3_11 == 'true' }}
+    uses: ./.github/workflows/build-image-template.yml
+    with:
+      python_version: '3.11'
+      dockerfile_path: docker/Dockerfile.python3.11
+      tag_suffix: py3.11
+    secrets: inherit
+
+  build-python-3-12:
+    if: ${{ github.event.inputs.python_3_12 == 'true' }}
+    uses: ./.github/workflows/build-image-template.yml
+    with:
+      python_version: '3.12'
+      dockerfile_path: docker/Dockerfile.python3.12
+      tag_suffix: py3.12
+    secrets: inherit

.github/workflows/codespell.yml

@@ -1,45 +0,0 @@
-name: codespell
-
-on:
-  # Trigger the workflow on push or pull request,
-  # but only for the main branch
-  push:
-    branches:
-      - main
-    paths:
-      - "**/*.py"
-      - "**/*.md"
-      - "**/*.rst"
-      - pyproject.toml
-      - requirements-lint.txt
-      - .github/workflows/codespell.yml
-  pull_request:
-    branches:
-      - main
-    paths:
-      - "**/*.py"
-      - "**/*.md"
-      - "**/*.rst"
-      - pyproject.toml
-      - requirements-lint.txt
-      - .github/workflows/codespell.yml
-
-jobs:
-  codespell:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check out repository
-        uses: actions/checkout@v3
-
-      - name: Set up Python
-        uses: actions/setup-python@v4
-        with:
-          python-version: '3.12'  # or any version you need
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -r requirements-lint.txt
-      - name: Spelling check with codespell
-        run: |
-          # Refer to the above environment variable here
-          codespell --toml pyproject.toml $CODESPELL_EXCLUDES

.github/workflows/docs.yml

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

.github/workflows/fastvideo-publish.yml

@@ -6,6 +6,7 @@ on:
       - main
     paths:
       - 'pyproject.toml'  # Trigger when pyproject.toml changes
+  workflow_dispatch:
 
 jobs:
   check-version-change:
@@ -15,7 +16,7 @@ jobs:
       new-version: ${{ steps.check-version.outputs.new-version }}
     steps:
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           fetch-depth: 2
 
@@ -23,11 +24,11 @@ jobs:
         id: check-version
         run: |
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'version\s*=\s*"\K[^"]+' pyproject.toml)
+          NEW_VERSION=$(grep -oP "version\\s*=\\s*\"\\K[^\"]+\"" pyproject.toml)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./pyproject.toml | grep -oP 'version\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./pyproject.toml | grep -oP "version\\s*=\\s*\"\\K[^\"]+\"" || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -41,17 +42,17 @@ jobs:
 
   build-publish-main:
     needs: check-version-change
-    if: needs.check-version-change.outputs.version-changed == 'true'
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
     runs-on: ubuntu-latest
     permissions:
       id-token: write  # Needed for OIDC Trusted Publishing
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
 
       - name: Set up Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: '3.10'
 

.github/workflows/matchers/actionlint.json

@@ -0,0 +1,17 @@
+{
+  "problemMatcher": [
+    {
+      "owner": "actionlint",
+      "pattern": [
+        {
+          "regexp": "^(?:\\x1b\\[\\d+m)?(.+?)(?:\\x1b\\[\\d+m)*:(?:\\x1b\\[\\d+m)*(\\d+)(?:\\x1b\\[\\d+m)*:(?:\\x1b\\[\\d+m)*(\\d+)(?:\\x1b\\[\\d+m)*: (?:\\x1b\\[\\d+m)*(.+?)(?:\\x1b\\[\\d+m)* \\[(.+?)\\]$",
+          "file": 1,
+          "line": 2,
+          "column": 3,
+          "message": 4,
+          "code": 5
+        }
+      ]
+    }
+  ]
+}

.github/workflows/matchers/mypy.json

@@ -0,0 +1,16 @@
+{
+  "problemMatcher": [
+    {
+      "owner": "mypy",
+      "pattern": [
+        {
+          "regexp": "^(.+):(\\d+):\\s(error|warning):\\s(.+)$",
+          "file": 1,
+          "line": 2,
+          "severity": 3,
+          "message": 4
+        }
+      ]
+    }
+  ]
+}

.github/workflows/pr-test.yml

@@ -0,0 +1,375 @@
+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_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
+        default: false
+        type: boolean
+      run_nightly_test:
+        description: "Run nightly-test"
+        required: false
+        default: false
+        type: boolean
+
+env:
+  PYTHONUNBUFFERED: "1"
+
+
+concurrency:
+  group: pr-test-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  pre-commit:
+    uses: ./.github/workflows/pre-commit.yml
+
+  change-filter:
+    runs-on: ubuntu-latest
+    needs: pre-commit
+    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 }}
+      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 }}
+    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.12'
+            sta-kernel-paths: &sta-kernel-paths
+              - 'csrc/attn/st_attn/**'
+              - 'csrc/attn/setup_sta.py'
+              - 'csrc/attn/config_sta.py'
+              - 'csrc/attn/st_attn.cpp'
+            vsa-kernel-paths: &vsa-kernel-paths
+              - 'csrc/attn/vsa/**'
+              - 'csrc/attn/tk/**'
+              - 'csrc/attn/setup_vsa.py'
+              - 'csrc/attn/config_vsa.py'
+              - 'csrc/attn/vsa.cpp'
+            vsa-paths: &vsa-paths
+              - 'fastvideo/v1/**'
+              - *common-paths
+              - *vsa-kernel-paths
+            
+            # Actual tests
+            encoder-test:
+              - 'fastvideo/v1/models/encoders/**'
+              - 'fastvideo/v1/models/loader/**'
+              - 'fastvideo/v1/tests/encoders/**'
+              - *common-paths
+            vae-test:
+              - 'fastvideo/v1/models/vaes/**'
+              - 'fastvideo/v1/models/loader/**'
+              - 'fastvideo/v1/tests/vaes/**'
+              - *common-paths
+            transformer-test:
+              - 'fastvideo/v1/models/dits/**'
+              - 'fastvideo/v1/models/loader/**'
+              - 'fastvideo/v1/tests/transformers/**'
+              - 'fastvideo/v1/layers/**'
+              - 'fastvideo/v1/attention/**'
+              - *common-paths
+            training-test:
+              - 'fastvideo/v1/**'
+              - *common-paths
+            training-test-VSA:
+              - 'fastvideo/v1/**'
+              - *common-paths
+              - *vsa-kernel-paths
+            inference-test-STA:
+              - 'fastvideo/v1/**'
+              - *common-paths
+              - *sta-kernel-paths
+            precision-test-STA:
+              - *common-paths
+              - *sta-kernel-paths
+            precision-test-VSA:
+              - *common-paths
+              - *vsa-kernel-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/v1/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/v1/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/v1/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/v1/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/v1/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/v1/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 }}
+
+  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/v1/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: >-
+      (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_block_sparse.py"
+      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/v1/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, inference-test-STA, precision-test-STA, 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", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]'
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        run: python .github/scripts/runpod_cleanup.py

.github/workflows/pre-commit.yml

@@ -0,0 +1,18 @@
+name: pre-commit
+
+on:
+    workflow_call:
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v5
+      with:
+        python-version: "3.12"
+    - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
+    - run: echo "::add-matcher::.github/workflows/matchers/mypy.json"
+    - uses: pre-commit/action@v3.0.1
+      with:
+        extra_args: --all-files --hook-stage manual

.github/workflows/ruff.yml

@@ -1,50 +0,0 @@
-name: ruff
-
-on:
-  # Trigger the workflow on push or pull request,
-  # but only for the main branch
-  push:
-    branches:
-      - main
-    paths:
-      - "**/*.py"
-      - pyproject.toml
-      - requirements-lint.txt
-      - .github/workflows/matchers/ruff.json
-      - .github/workflows/ruff.yml
-  pull_request:
-    branches:
-      - main
-    # This workflow is only relevant when one of the following files changes.
-    # However, we have github configured to expect and require this workflow
-    # to run and pass before github with auto-merge a pull request. Until github
-    # allows more flexible auto-merge policy, we can just run this on every PR.
-    # It doesn't take that long to run, anyway.
-    #paths:
-    #  - "**/*.py"
-    #  - pyproject.toml
-    #  - requirements-lint.txt
-    #  - .github/workflows/matchers/ruff.json
-    #  - .github/workflows/ruff.yml
-
-jobs:
-  ruff:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check out repository
-        uses: actions/checkout@v3
-
-      - name: Set up Python
-        uses: actions/setup-python@v4
-        with:
-          python-version: '3.12'  # or any version you need
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -r requirements-lint.txt
-      - name: Analysing the code with ruff
-        run: |
-          ruff check .
-      - name: Run isort
-        run: |
-          isort . --check-only

.github/workflows/runpod-test.yml

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

.github/workflows/sta-publish.yml

@@ -5,7 +5,8 @@ on:
     branches:
       - main
     paths:
-      - "csrc/sliding_tile_attention/setup.py"
+      - "csrc/attn/setup_sta.py"
+  workflow_dispatch:
 
 jobs:
   check-version-change:
@@ -15,20 +16,20 @@ jobs:
       new-version: ${{ steps.check-version.outputs.new-version }}
     steps:
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           fetch-depth: 2
 
       - name: Check if version changed
         id: check-version
         run: |
-          cd csrc/sliding_tile_attention
+          cd csrc/attn
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_sta.py)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup_sta.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -43,7 +44,7 @@ jobs:
   build_wheels:
     name: Build Wheel
     needs: check-version-change
-    if: needs.check-version-change.outputs.version-changed == 'true'
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
     runs-on: ${{ matrix.os }}
 
     strategy:
@@ -57,6 +58,29 @@ jobs:
           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
 
@@ -112,19 +136,21 @@ jobs:
 
       - 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/sliding_tile_attention  # Move into the correct folder
-          git submodule update --init --recursive tk # Ensure ThunderKittens submodule is initialized
-          python setup.py bdist_wheel --dist-dir=dist
+          cd csrc/attn  # Move into the correct folder
+          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
+          python setup_sta.py bdist_wheel --dist-dir=dist
 
       - name: Rename wheel file
         run: |
-          cd csrc/sliding_tile_attention
+          cd csrc/attn
 
           CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
           TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2)
@@ -139,13 +165,13 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: ${{ env.wheel_name }}
-          path: csrc/sliding_tile_attention/dist/*.whl
+          path: csrc/attn/dist/*.whl
           retention-days: 90
 
   publish_package:
     name: Publish package
-    needs: [build_wheels]
-    if: needs.check-version-change.outputs.version-changed == 'true'
+    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
@@ -205,17 +231,19 @@ jobs:
 
       - 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/sliding_tile_attention  # Move into the correct folder
-          git submodule update --init --recursive tk # Ensure ThunderKittens submodule is initialized
-          python setup.py sdist --dist-dir=dist
+          cd csrc/attn  # Move into the correct folder
+          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
+          python setup_sta.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: csrc/sliding_tile_attention/dist/
+          packages-dir: csrc/attn/dist/

.github/workflows/test.yml

@@ -11,10 +11,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Check out repository
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
 
       - name: Set up Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: '3.12'  # or any version you need
 
@@ -23,11 +23,9 @@ jobs:
           python -m pip install --upgrade pip setuptools wheel
           pip install torch
           pip install packaging ninja
-          # remove st-attn dependency because no cuda environment
-          sed -i '/st_attn/d' pyproject.toml
           pip install -e .
           pip install pytest
 
       - name: Run Pytest
         run: |
-          pytest --ignore csrc/sliding_tile_attention/test
+          pytest --ignore csrc/attn/test

.github/workflows/yapf.yml

@@ -1,38 +0,0 @@
-name: yapf
-
-on:
-  # Trigger the workflow on push or pull request,
-  # but only for the main branch
-  push:
-    branches:
-      - main
-    paths:
-      - "**/*.py"
-      - .github/workflows/yapf.yml
-  pull_request:
-    branches:
-      - main
-    paths:
-      - "**/*.py"
-      - .github/workflows/yapf.yml
-
-jobs:
-  yapf:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check out repository
-        uses: actions/checkout@v3
-
-      - name: Set up Python
-        uses: actions/setup-python@v4
-        with:
-          python-version: '3.12'  # or any version you need
-
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install yapf==0.32.0
-          pip install toml==0.10.2
-      - name: Running yapf
-        run: |
-          yapf --diff --recursive .

.gitignore

@@ -1,11 +1,8 @@
 __pycache__
-*.mp4
 .ipynb_checkpoints
 *.pth
 UCF-101/
 results/
-build/
-fastvideo.egg-info/
 wandb/
 *.ipynb
 *.jpg
@@ -26,11 +23,41 @@ outputs_video
 sbatch.sh
 *.out
 env
-dist/
 *.o
 **/build/
-**.egg-info
 **.pyc
-**.egg
 **.txt
-**.json
+
+# Distribution / packaging
+build/
+dist/
+*.egg-info/
+*.egg
+eggs/
+.eggs/
+
+# Sphinx documentation
+docs/_build/
+docs/source/getting_started/examples/
+docs/source/inference/examples/
+
+# VSCode
+.vscode/
+
+# DS Store
+.DS_Store
+
+# vim swap files
+*.swo
+*.swp
+
+# Python pickle files
+*.pkl
+
+# Reference videos
+!fastvideo/v1/tests/ssim/reference_videos/**/*.mp4
+
+# Static images
+!docs/source/_static/images/**/*.png
+!comfyui/assets/**/*.png
+!comfyui/assets/**/*.gif

.gitmodules

@@ -1,3 +1,3 @@
-[submodule "csrc/sliding_tile_attention/tk"]
-	path = csrc/sliding_tile_attention/tk
+[submodule "csrc/attn/tk"]
+	path = csrc/attn/tk
 	url = https://github.com/HazyResearch/ThunderKittens.git

.pre-commit-config.yaml

@@ -0,0 +1,83 @@
+default_stages:
+  - pre-commit # Run locally
+  - manual # Run in CI
+exclude: |
+  (?x)(
+    fastvideo/v1/third_party/.*|
+    csrc/.*|
+    assets/.*|
+    tests/.*|
+    demo/.*|
+    predict\.py|
+    scripts/.*|
+    fastvideo/data_preprocess/.*|
+    fastvideo/dataset/.*|
+    fastvideo/distill/.*|
+    fastvideo/distill\.py|
+    fastvideo/distill_adv\.py|
+    fastvideo/models/.*|
+    fastvideo/sample/.*|
+    fastvideo/train\.py|
+    fastvideo/utils/.*|
+    examples/.*|
+    .github/workflows/fastvideo-publish.yml|
+    .github/workflows/sta-publish.yml|
+    .github/workflows/build-image-template.yml|
+    docs/source/inference/support_matrix.md
+  )
+repos:
+- repo: https://github.com/google/yapf
+  rev: v0.43.0
+  hooks:
+  - id: yapf
+    args: [--in-place, --verbose]
+    additional_dependencies: [toml] # TODO: Remove when yapf is upgraded
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.11.12
+  hooks:
+  - id: ruff
+    args: [--output-format, github, --fix]
+- repo: https://github.com/codespell-project/codespell
+  rev: v2.4.1
+  hooks:
+  - id: codespell
+    additional_dependencies: ['tomli']
+    args: ['--toml', 'pyproject.toml']
+- repo: https://github.com/PyCQA/isort
+  rev: 6.0.1
+  hooks:
+  - id: isort
+- repo: https://github.com/jackdewinter/pymarkdown
+  rev: v0.9.30
+  hooks:
+  - id: pymarkdown
+    args: [fix]
+- repo: https://github.com/rhysd/actionlint
+  rev: v1.7.7
+  hooks:
+  - id: actionlint
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v1.15.0
+  hooks:
+  - id: mypy
+    args: [--python-version, '3.10', --follow-imports, "skip", ]
+    additional_dependencies: [types-cachetools, types-setuptools, types-PyYAML, types-requests]
+- repo: local
+  hooks:
+  - id: check-filenames
+    name: Check for spaces in all filenames
+    entry: bash
+    args:
+      - -c
+      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
+    language: system
+    always_run: true
+    pass_filenames: false
+  # Keep `suggestion` last
+  - id: suggestion
+    name: Suggestion
+    entry: bash -c 'echo "To bypass pre-commit hooks, add --no-verify to git commit."'
+    language: system
+    verbose: true
+    pass_filenames: false
+  # Insert new entries above the `suggestion` entry

README.md

@@ -2,229 +2,162 @@
 <img src=assets/logo.jpg width="30%"/>
 </div>
 
-FastVideo is a lightweight framework for accelerating large video diffusion models.
+**FastVideo is a unified framework for accelerated video generation.**
 
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 
 <p align="center">
-    🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank">FastHunyuan</a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank">FastMochi</a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> Slack </a> 
-</p> 
+    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |
+</p>
 
+<div align="center">
+<img src=assets/perf.png width="90%"/>
+</div>
 
+## NEWS
+- ```2025/06/14```: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
+- ```2025/04/24```: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
+- ```2025/02/18```: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
 
+## Key Features
 
-
-https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1
-
-
-
-FastVideo currently offers: (with more to come)
-
-- [NEW!] [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
-- First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-- Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
 - Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 - Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
 
-Dev in progress and highly experimental.
-
-
-
-## Change Log
-- ```2025/02/20```: FastVideo now supports STA on [StepVideo](https://github.com/stepfun-ai/Step-Video-T2V) with 3.4X speedup!
-- ```2025/02/18```: Release the inference code and kernel for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- ```2025/01/13```: Support Lora finetuning for HunyuanVideo.
-- ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment.
-- ```2024/12/17```: `FastVideo` v1.0 is released.
-
-
-## 🔧 Installation
-The code is tested on Python 3.10.0, CUDA 12.4 and H100.
-```
-./env_setup.sh fastvideo
-```
-To try Sliding Tile Attention (optional), please follow the instruction in [csrc/sliding_tile_attention/README.md](csrc/sliding_tile_attention/README.md) to install STA.
-
-## 🚀 Inference
-### Inference StepVideo with Sliding Tile Attention 
-First, download the model:
-```
-python scripts/huggingface/download_hf.py --repo_id=stepfun-ai/stepvideo-t2v --local_dir=data/stepvideo-t2v --repo_type=model 
-```
-Use the following scripts to run inference for StepVideo. When using STA for inference, the generated videos will have dimensions of 204×768×768 (currently, this is the only supported shape).
-```bash
-sh scripts/inference/inference_stepvideo_STA.sh # Inference stepvideo with STA
-sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo
-```
-
-### Inference HunyuanVideo with Sliding Tile Attention
-First, download the model:
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model 
-```
-We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
-```bash
-sh scripts/inference/inference_hunyuan_STA.sh
-```
-### Video Demos using STA + Teacache
-Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100.
-
-### Inference FastHunyuan on single RTX4090
-We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan_hf_quantization.sh
-```
-For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):
-
-
-| Configuration                  | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time |
-|--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------|
-| BF16 + Pipeline CPU Offload    | 23.883G                   | 33.744G                                    | 81s            | 121.5s          |
-| INT8 + Pipeline CPU Offload    | 13.911G                   | 27.979G                                    | 88s            | 116.7s          |
-| NF4 + Pipeline CPU Offload     | 9.453G                    | 19.26G                                     | 78s            | 114.5s          |
-           
-
-
-For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section:
-
-### FastHunyuan
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan.sh
-```
-You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).
-
-### FastMochi
+## Getting Started
+We recommend using an environment manager such as `Conda` to create a clean environment:
 
 ```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffusers --local_dir=data/FastMochi-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_mochi_sp.sh
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
 ```
 
+Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) for more detailed installation instructions.
+
+## Inference
+### Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+
+```python
+from fastvideo import VideoGenerator
+
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:
 
-## 🎯 Distill
-Our distillation recipe is based on [Phased Consistency Model](https://github.com/G-U-N/Phased-Consistency-Model). We did not find significant improvement using multi-phase distillation, so we keep the one phase setup similar to the original latent consistency model's recipe.
-We use the [MixKit](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/all_mixkit) dataset for distillation. To avoid running the text encoder and VAE during training, we preprocess all data to generate text embeddings and VAE latents.
-Preprocessing instructions can be found [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide preprocessed data that can be downloaded directly using the following command:
 ```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/HD-Mixkit-Finetune-Hunyuan --local_dir=data/HD-Mixkit-Finetune-Hunyuan --repo_type=dataset
+python example.py
 ```
-Next, download the original model weights with:
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model # original hunyuan
-python scripts/huggingface/download_hf.py --repo_id=genmo/mochi-1-preview --local_dir=data/mochi --repo_type=model # original mochi
-```
-To launch the distillation process, use the following commands:
-```
-bash scripts/distill/distill_hunyuan.sh # for hunyuan
-bash scripts/distill/distill_mochi.sh # for mochi
-```
-We also provide an optional script for distillation with adversarial loss, located at `fastvideo/distill_adv.py`. Although we tried adversarial loss, we did not observe significant improvements.
-## Finetune
-### ⚡ Full Finetune
-Ensure your data is prepared and preprocessed in the format specified in [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide a mochi preprocessed Black Myth Wukong data that can be downloaded directly:
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/Mochi-Black-Myth --local_dir=data/Mochi-Black-Myth --repo_type=dataset
-```
-Download the original model weights as specified in [Distill Section](#-distill):
 
-Then you can run the finetune with:
-```
-bash scripts/finetune/finetune_mochi.sh # for mochi
-```
-**Note that for finetuning, we did not tune the hyperparameters in the provided script.**
-### ⚡ Lora Finetune 
+For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html).
 
-Hunyuan supports Lora fine-tuning of videos up to 720p. Demos and prompts of Black-Myth-Wukong can be found in [here](https://huggingface.co/FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight). You can download the Lora weight through:
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight --local_dir=data/Hunyuan-Black-Myth-Wukong-lora-weight --repo_type=model
-```
-#### Minimum Hardware Requirement
-- 40 GB GPU memory each for 2 GPUs with lora.
-- 30 GB GPU memory each for 2 GPUs with CPU offload and lora.  
+### Other docs:
 
+- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview.html)
+- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
 
-Currently, both Mochi and Hunyuan models support Lora finetuning through diffusers. To generate personalized videos from your own dataset, you'll need to follow three main steps: dataset preparation, finetuning, and inference.
-
-#### Dataset Preparation
-We provide scripts to better help you get started to train on your own characters!  
-You can run this to organize your dataset to get the videos2caption.json before preprocess. Specify your video folder and corresponding caption folder (caption files should be .txt files and have the same name with its video):
-```
-python scripts/dataset_preparation/prepare_json_file.py --video_dir data/input_videos/ --prompt_dir data/captions/ --output_path data/output_folder/videos2caption.json --verbose
-```
-Also, we provide script to resize your videos:
-```
-python scripts/data_preprocess/resize_videos.py 
-```
-#### Finetuning
-After basic dataset preparation and preprocess, you can start to finetune your model using Lora:
-```
-bash scripts/finetune/finetune_hunyuan_hf_lora.sh
-```
-#### Inference
-For inference with Lora checkpoint, you can run the following scripts with additional parameter `--lora_checkpoint_dir`:
-```
-bash scripts/inference/inference_hunyuan_hf.sh 
-```
-**We also provide scripts for Mochi in the same directory.**
-
-#### Finetune with Both Image and Video
-Our codebase support finetuning with both image and video. 
-```bash
-bash scripts/finetune/finetune_hunyuan.sh
-bash scripts/finetune/finetune_mochi_lora_mix.sh
-```
-For Image-Video Mixture Fine-tuning, make sure to enable the `--group_frame` option in your script.
+## Distillation and Finetuning
+- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/training/distillation.html)
+- [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html)
 
 ## 📑 Development Plan
 
-- More distillation methods
-  - [ ] Add Distribution Matching Distillation
+<!-- - More distillation methods -->
+  <!-- - [ ] Add Distribution Matching Distillation -->
 - More models support
-  - [ ] Add CogvideoX model
-- Code update
-  - [ ] fp8 support
-  - [ ] faster load model and save model support
+  <!-- - [ ] Add CogvideoX model -->
+  - [x] Add StepVideo to V1
+- Optimization features
+  - [x] Teacache in V1
+  - [x] SageAttention in V1
+- Code updates
+  - [x] V1 Configuration API
+  - [ ] Support Training in V1
+  <!-- - [ ] fp8 support -->
+  <!-- - [ ] faster load model and save model support -->
 
 ## 🤝 Contributing
 
-We welcome all contributions. Please run `bash format.sh --all` before submitting a pull request.
-
-## 🔧 Testing
-Run `pytest` to verify the data preprocessing, checkpoint saving, and sequence parallel pipelines. We recommend adding corresponding test cases in the `test` folder to support your contribution.
+We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview.html)
 
 ## Acknowledgement
-We learned and reused code from the following projects: [PCM](https://github.com/G-U-N/Phased-Consistency-Model), [diffusers](https://github.com/huggingface/diffusers), [OpenSoraPlan](https://github.com/PKU-YuanGroup/Open-Sora-Plan), and [xDiT](https://github.com/xdit-project/xDiT).
+We learned and reused code from the following projects:
+- [PCM](https://github.com/G-U-N/Phased-Consistency-Model)
+- [diffusers](https://github.com/huggingface/diffusers)
+- [OpenSoraPlan](https://github.com/PKU-YuanGroup/Open-Sora-Plan)
+- [xDiT](https://github.com/xdit-project/xDiT)
+- [vLLM](https://github.com/vllm-project/vllm)
+- [SGLang](https://github.com/sgl-project/sglang)
 
-We thank MBZUAI and Anyscale for their support throughout this project.
+We thank MBZUAI and [Anyscale](https://www.anyscale.com/) for their support throughout this project.
 
-## Citation 
+## Citation
 If you use FastVideo for your research, please cite our paper:
 
 ```bibtex
+@misc{zhang2025vsafastervideodiffusion,
+      title={VSA: Faster Video Diffusion with Trainable Sparse Attention}, 
+      author={Peiyuan Zhang and Haofeng Huang and Yongqi Chen and Will Lin and Zhengzhong Liu and Ion Stoica and Eric Xing and Hao Zhang},
+      year={2025},
+      eprint={2505.13389},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2505.13389}, 
+}
 @misc{zhang2025fastvideogenerationsliding,
-      title={Fast Video Generation with Sliding Tile Attention}, 
+      title={Fast Video Generation with Sliding Tile Attention},
       author={Peiyuan Zhang and Yongqi Chen and Runlong Su and Hangliang Ding and Ion Stoica and Zhenghong Liu and Hao Zhang},
       year={2025},
       eprint={2502.04507},
       archivePrefix={arXiv},
       primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2502.04507}, 
+      url={https://arxiv.org/abs/2502.04507},
 }
 @misc{ding2025efficientvditefficientvideodiffusion,
-      title={Efficient-vDiT: Efficient Video Diffusion Transformers With Attention Tile}, 
+      title={Efficient-vDiT: Efficient Video Diffusion Transformers With Attention Tile},
       author={Hangliang Ding and Dacheng Li and Runlong Su and Peiyuan Zhang and Zhijie Deng and Ion Stoica and Hao Zhang},
       year={2025},
       eprint={2502.06155},
       archivePrefix={arXiv},
       primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2502.06155}, 
+      url={https://arxiv.org/abs/2502.06155},
 }
 ```

comfyui/README.md

@@ -0,0 +1,138 @@
+# ComfyUI-FastVideo
+
+A custom node suite for ComfyUI that provides accelerated video generation using [FastVideo](https://github.com/hao-ai-labs/FastVideo). See the [blog post](https://hao-ai-lab.github.io/blogs/fastvideo/) about FastVideo V1 to learn more.
+
+## Multi-GPU Parallel Inference
+
+One of the key features ComfyUI-FastVideo brings to ComfyUI is its ability to distribute the generation workload across multiple GPUs, resulting in significantly faster inference times.
+
+![Wan2.1-I2V-14B-480P-Diffusers](./assets/wani2v.gif)
+
+Example of Wan2.1-I2V-14B-480P-Diffusers model running on 4 GPUs.
+## Features
+
+- Generate high-quality videos from text prompts and images
+- Configurable video parameters (prompt, resolution, frame count, FPS)
+- Support for multiple GPUs with tensor and sequence parallelism
+- Advanced configuration options for VAE, Text Encoder, and DIT components
+- Interruption/cancellation support for long-running generations
+
+## Installation
+
+### Requirements
+
+- [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
+- CUDA-capable GPU(s) with sufficient VRAM
+
+### Install using ComfyUI Manager
+
+Coming soon!
+
+### Manual Installation
+
+#### Copy the FastVideo `comfyui` directory into your ComfyUI custom_nodes directory:
+
+```bash
+cp -r /path/to/FastVideo/comfyui /path/to/ComfyUI/custom_nodes/FastVideo
+```
+
+#### Install dependencies:
+
+Currently, the only dependency is `fastvideo`, which can be installed using pip.
+
+```bash
+pip install fastvideo
+```
+
+#### Install missing custom nodes:
+
+`ComfyUI-VideoHelperSuite`:
+
+```bash
+cd /path/to/ComfyUI/custom_nodes
+git clone https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git
+```
+
+If you're seeing `ImportError: libGL.so.1: cannot open shared object file: No such file or directory`,
+you may need to install ffmpeg
+
+```bash
+apt-get update && apt-get install ffmpeg
+```
+
+## Usage
+
+After installation, the following nodes will be available in the ComfyUI interface under the "fastvideo" category:
+
+- **Video Generator**: The main node for generating videos from prompts
+- **Inference Args**: Configure video generation parameters
+- **VAE Config**
+- **Text Encoder Config**
+- **DIT Config**
+- **Load Image Path**: Load images for potential conditioning
+
+You may have noticed many arguments on the nodes have 'auto' as the default value. This is because FastVideo will automatically detect the best values for these parameters based on the model and the hardware. However, you can also manually configure these parameters to get the best performance for your specific use case. We plan on releasing more optimized workflow files for different models and hardware configurations in the future.
+
+You can see what some of the default configurations are by looking at the FastVideo repo:
+- [Wan2.1-I2V-14B-480P-Diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json)
+- [FastHunyuan-diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/v1/configs/fasthunyuan_t2v.json)
+
+### Node Configuration
+
+#### Video Generator
+
+- **prompt**: Text description of the video to generate
+- **output_path**: Directory where generated videos will be saved
+- **num_gpus**: Number of GPUs to use for generation
+- **model_path**: Path to the FastVideo model
+- **embedded_cfg_scale**: Classifier-free guidance scale
+- **sp_size**: Sequence parallelism size (usually should match num_gpus)
+- **tp_size**: Tensor parallelism size (usually should match num_gpus)
+- **precision**: Model precision (fp16 or bf16)
+
+`model_path takes either a model id from huggingface or a local path to a model. Models by default will be downloaded to ~/.cache/huggingface/hub/ and cached for subsequent runs.`
+
+#### Inference Args
+
+- **height/width**: Resolution of the output video
+- **num_frames**: Number of frames to generate
+- **num_inference_steps**: Number of diffusion steps per frame
+- **guidance_scale**: Classifier-free guidance scale
+- **flow_shift**: Frame flow shift parameter
+- **seed**: Random seed for reproducible generation
+- **fps**: Frames per second of the output video
+- **image_path**: Optional path to input image for conditioning (for i2v models)
+
+## Memory Management
+
+Models will remain loaded in GPU memory between runs when you only change inference arguments (such as prompt, resolution, frame count, FPS, guidance scale, etc.) or the prompt text. This allows for faster subsequent generations since the model doesn't need to be reloaded.
+
+However, if you need to change the following parameters, you will need to restart the ComfyUI server:
+- **Number of GPUs** (`num_gpus`)
+- **Model path** (`model_path`)
+- **Tensor parallelism size** (`tp_size`)
+- **Sequence parallelism size** (`sp_size`)
+
+These parameters affect the model's distribution across GPUs and require a complete reinitialization of the model pipeline.
+
+## Example workflows
+
+### Text to Video
+
+FastVideo-FastHunyuan-diffusers
+
+![FastVideo-FastHunyuan-diffusers](./assets/fasthunyuan.png)
+
+- [FastHunyuan-diffusers.json](./examples/FastHunyuan-diffusers.json)
+
+### Image to Video
+
+Wan2.1-I2V-14B-480P-Diffusers
+
+![Wan2.1-I2V-14B-480P-Diffusers](./assets/wani2v.png)
+
+- [Wan2.1-I2V-14B-480P-Diffusers.json](./examples/Wan2.1-I2V-14B-480P-Diffusers.json)
+
+## License
+
+This project is licensed under Apache 2.0.

comfyui/__init__.py

@@ -0,0 +1,5 @@
+from .video_generator.nodes import (NODE_CLASS_MAPPINGS,
+                                    NODE_DISPLAY_NAME_MAPPINGS)
+
+WEB_DIRECTORY = "./web"
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']

comfyui/examples/FastHunyuan-diffusers.json

@@ -0,0 +1,645 @@
+{
+  "id": "23a1f065-bbba-4a8f-b144-944e1318fcbf",
+  "revision": 0,
+  "last_node_id": 8,
+  "last_link_id": 7,
+  "nodes": [
+    {
+      "id": 4,
+      "type": "VAEConfig",
+      "pos": [
+        374.2159423828125,
+        554.85888671875
+      ],
+      "size": [
+        334.080078125,
+        322
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "vae_config",
+          "type": "VAE_CONFIG",
+          "links": [
+            2
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VAEConfig"
+      },
+      "widgets_values": [
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "load_encoder": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "load_decoder": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "tile_sample_min_height": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 256
+        },
+        "tile_sample_min_width": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 256
+        },
+        "tile_sample_min_num_frames": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 16
+        },
+        "tile_sample_stride_height": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 192
+        },
+        "tile_sample_stride_width": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 192
+        },
+        "tile_sample_stride_num_frames": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 12
+        },
+        "blend_num_frames": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 0
+        },
+        "use_tiling": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "use_temporal_tiling": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "use_parallel_tiling": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 5,
+      "type": "TextEncoderConfig",
+      "pos": [
+        416.4937744140625,
+        953.6171875
+      ],
+      "size": [
+        270,
+        106
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "text_encoder_config",
+          "type": "TEXT_ENCODER_CONFIG",
+          "links": [
+            7
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "TextEncoderConfig"
+      },
+      "widgets_values": [
+        -99999,
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "prefix": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        },
+        "quant_config": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        },
+        "lora_config": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        }
+      }
+    },
+    {
+      "id": 6,
+      "type": "DITConfig",
+      "pos": [
+        415.1928405761719,
+        1154.1573486328125
+      ],
+      "size": [
+        270,
+        82
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "dit_config",
+          "type": "DIT_CONFIG",
+          "links": [
+            6
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "DITConfig"
+      },
+      "widgets_values": [
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "prefix": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        },
+        "quant_config": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": ""
+        }
+      }
+    },
+    {
+      "id": 1,
+      "type": "VideoGenerator",
+      "pos": [
+        818.804931640625,
+        348.9299621582031
+      ],
+      "size": [
+        400,
+        436
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "inference_args",
+          "shape": 7,
+          "type": "INFERENCE_ARGS",
+          "link": 3
+        },
+        {
+          "name": "vae_config",
+          "shape": 7,
+          "type": "VAE_CONFIG",
+          "link": 2
+        },
+        {
+          "name": "text_encoder_config",
+          "shape": 7,
+          "type": "TEXT_ENCODER_CONFIG",
+          "link": 7
+        },
+        {
+          "name": "dit_config",
+          "shape": 7,
+          "type": "DIT_CONFIG",
+          "link": 6
+        }
+      ],
+      "outputs": [
+        {
+          "name": "video_path",
+          "type": "STRING",
+          "links": [
+            4
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VideoGenerator"
+      },
+      "widgets_values": [
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones.",
+        "/workspace/ComfyUI/outputs_video/",
+        2,
+        "FastVideo/FastHunyuan-diffusers",
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "embedded_cfg_scale": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 6
+        },
+        "sp_size": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 2
+        },
+        "tp_size": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 2
+        },
+        "vae_precision": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": "fp16"
+        },
+        "vae_tiling": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "vae_sp": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        },
+        "text_encoder_precision": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": "fp16"
+        },
+        "precision": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": "fp16"
+        },
+        "use_cpu_offload": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 3,
+      "type": "VHS_LoadVideoPath",
+      "pos": [
+        1350.136962890625,
+        331.20361328125
+      ],
+      "size": [
+        231.8896484375,
+        286
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        },
+        {
+          "name": "video",
+          "type": "STRING",
+          "widget": {
+            "name": "video"
+          },
+          "link": 4
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            5
+          ]
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": null
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadVideoPath"
+      },
+      "widgets_values": {
+        "video": "",
+        "force_rate": 0,
+        "custom_width": 0,
+        "custom_height": 0,
+        "frame_load_cap": 0,
+        "skip_first_frames": 0,
+        "select_every_nth": 1,
+        "format": "Wan",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "",
+            "type": "path",
+            "format": "video/",
+            "force_rate": 0,
+            "custom_width": 0,
+            "custom_height": 0,
+            "frame_load_cap": 0,
+            "skip_first_frames": 0,
+            "select_every_nth": 1
+          }
+        }
+      }
+    },
+    {
+      "id": 2,
+      "type": "InferenceArgs",
+      "pos": [
+        411.46307373046875,
+        178.18182373046875
+      ],
+      "size": [
+        278.73828125,
+        298
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "inference_args",
+          "type": "INFERENCE_ARGS",
+          "links": [
+            3
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "InferenceArgs"
+      },
+      "widgets_values": [
+        720,
+        1280,
+        45,
+        6,
+        -99999,
+        -99999,
+        1025,
+        "fixed",
+        24,
+        -99999,
+        -99999
+      ],
+      "auto_widget_states": {
+        "height": {
+          "isAuto": false,
+          "value": 720,
+          "cachedValue": 720
+        },
+        "width": {
+          "isAuto": false,
+          "value": 1280,
+          "cachedValue": 1280
+        },
+        "num_frames": {
+          "isAuto": false,
+          "value": 45,
+          "cachedValue": 45
+        },
+        "num_inference_steps": {
+          "isAuto": false,
+          "value": 6,
+          "cachedValue": 6
+        },
+        "guidance_scale": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 1
+        },
+        "flow_shift": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": 17
+        },
+        "seed": {
+          "isAuto": false,
+          "value": 1025,
+          "cachedValue": 1024
+        },
+        "fps": {
+          "isAuto": false,
+          "value": 24,
+          "cachedValue": 24
+        },
+        "image_path": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": "X://insert/path/here.mp4"
+        },
+        "enable_teacache": {
+          "isAuto": true,
+          "value": -99999,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 8,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1668.3499755859375,
+        328.22625732421875
+      ],
+      "size": [
+        507.507080078125,
+        622.2227172851562
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 5
+        },
+        {
+          "name": "audio",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 24,
+        "loop_count": 0,
+        "filename_prefix": "",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 19,
+        "save_metadata": true,
+        "trim_to_audio": false,
+        "pingpong": false,
+        "save_output": false,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "._00003.mp4",
+            "subfolder": "",
+            "type": "temp",
+            "format": "video/h264-mp4",
+            "frame_rate": 24,
+            "workflow": "._00003.png",
+            "fullpath": "/workspace/ComfyUI/temp/._00003.mp4"
+          }
+        }
+      }
+    }
+  ],
+  "links": [
+    [
+      2,
+      4,
+      0,
+      1,
+      1,
+      "VAE_CONFIG"
+    ],
+    [
+      3,
+      2,
+      0,
+      1,
+      0,
+      "INFERENCE_ARGS"
+    ],
+    [
+      4,
+      1,
+      0,
+      3,
+      2,
+      "STRING"
+    ],
+    [
+      5,
+      3,
+      0,
+      8,
+      0,
+      "IMAGE"
+    ],
+    [
+      6,
+      6,
+      0,
+      1,
+      3,
+      "DIT_CONFIG"
+    ],
+    [
+      7,
+      5,
+      0,
+      1,
+      2,
+      "TEXT_ENCODER_CONFIG"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.9090909090909091,
+      "offset": [
+        112.86678372727341,
+        -71.45635903989245
+      ]
+    },
+    "frontendVersion": "1.20.4",
+    "VHS_latentpreview": false,
+    "VHS_latentpreviewrate": 0,
+    "VHS_MetadataImage": true,
+    "VHS_KeepIntermediate": true
+  },
+  "version": 0.4
+}

comfyui/examples/Wan2.1-I2V-14B-480P-Diffusers.json

@@ -0,0 +1,697 @@
+{
+  "id": "23a1f065-bbba-4a8f-b144-944e1318fcbf",
+  "revision": 0,
+  "last_node_id": 8,
+  "last_link_id": 7,
+  "nodes": [
+    {
+      "id": 7,
+      "type": "LoadImagePath",
+      "pos": [
+        33.15385437011719,
+        191.2037353515625
+      ],
+      "size": [
+        270,
+        334
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "image_path",
+          "type": "STRING",
+          "links": [
+            1
+          ]
+        },
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": null
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImagePath"
+      },
+      "widgets_values": [
+        "woman.jpg",
+        "image"
+      ]
+    },
+    {
+      "id": 3,
+      "type": "VHS_LoadVideoPath",
+      "pos": [
+        1350.136962890625,
+        331.20361328125
+      ],
+      "size": [
+        231.8896484375,
+        286
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        },
+        {
+          "name": "video",
+          "type": "STRING",
+          "widget": {
+            "name": "video"
+          },
+          "link": 4
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            5
+          ]
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": null
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadVideoPath"
+      },
+      "widgets_values": {
+        "video": "",
+        "force_rate": 0,
+        "custom_width": 0,
+        "custom_height": 0,
+        "frame_load_cap": 0,
+        "skip_first_frames": 0,
+        "select_every_nth": 1,
+        "format": "Wan",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "",
+            "type": "path",
+            "format": "video/",
+            "force_rate": 0,
+            "custom_width": 0,
+            "custom_height": 0,
+            "frame_load_cap": 0,
+            "skip_first_frames": 0,
+            "select_every_nth": 1
+          }
+        }
+      }
+    },
+    {
+      "id": 8,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1668.3499755859375,
+        328.22625732421875
+      ],
+      "size": [
+        214.7587890625,
+        334
+      ],
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 5
+        },
+        {
+          "name": "audio",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 24,
+        "loop_count": 0,
+        "filename_prefix": "",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 19,
+        "save_metadata": true,
+        "trim_to_audio": false,
+        "pingpong": false,
+        "save_output": false,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {}
+        }
+      }
+    },
+    {
+      "id": 4,
+      "type": "VAEConfig",
+      "pos": [
+        374.2159423828125,
+        554.85888671875
+      ],
+      "size": [
+        334.080078125,
+        322
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "vae_config",
+          "type": "VAE_CONFIG",
+          "links": [
+            2
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VAEConfig"
+      },
+      "widgets_values": [
+        true,
+        true,
+        256,
+        256,
+        16,
+        192,
+        192,
+        12,
+        0,
+        true,
+        true,
+        true
+      ],
+      "auto_widget_states": {
+        "load_encoder": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "load_decoder": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "tile_sample_min_height": {
+          "isAuto": true,
+          "value": 256,
+          "cachedValue": 256
+        },
+        "tile_sample_min_width": {
+          "isAuto": true,
+          "value": 256,
+          "cachedValue": 256
+        },
+        "tile_sample_min_num_frames": {
+          "isAuto": true,
+          "value": 16,
+          "cachedValue": 16
+        },
+        "tile_sample_stride_height": {
+          "isAuto": true,
+          "value": 192,
+          "cachedValue": 192
+        },
+        "tile_sample_stride_width": {
+          "isAuto": true,
+          "value": 192,
+          "cachedValue": 192
+        },
+        "tile_sample_stride_num_frames": {
+          "isAuto": true,
+          "value": 12,
+          "cachedValue": 12
+        },
+        "blend_num_frames": {
+          "isAuto": true,
+          "value": 0,
+          "cachedValue": 0
+        },
+        "use_tiling": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "use_temporal_tiling": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "use_parallel_tiling": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 2,
+      "type": "InferenceArgs",
+      "pos": [
+        411.46307373046875,
+        178.18182373046875
+      ],
+      "size": [
+        278.73828125,
+        298
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "image_path",
+          "shape": 7,
+          "type": "STRING",
+          "widget": {
+            "name": "image_path"
+          },
+          "link": 1
+        }
+      ],
+      "outputs": [
+        {
+          "name": "inference_args",
+          "type": "INFERENCE_ARGS",
+          "links": [
+            3
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "InferenceArgs"
+      },
+      "widgets_values": [
+        832,
+        480,
+        45,
+        20,
+        1,
+        17,
+        1024,
+        "fixed",
+        24,
+        "X://insert/path/here.mp4",
+        true
+      ],
+      "auto_widget_states": {
+        "height": {
+          "isAuto": false,
+          "value": 832,
+          "cachedValue": 720
+        },
+        "width": {
+          "isAuto": false,
+          "value": 480,
+          "cachedValue": 1280
+        },
+        "num_frames": {
+          "isAuto": false,
+          "value": 45,
+          "cachedValue": 45
+        },
+        "num_inference_steps": {
+          "isAuto": false,
+          "value": 20,
+          "cachedValue": 6
+        },
+        "guidance_scale": {
+          "isAuto": true,
+          "value": 1,
+          "cachedValue": 1
+        },
+        "flow_shift": {
+          "isAuto": true,
+          "value": 17,
+          "cachedValue": 17
+        },
+        "seed": {
+          "isAuto": false,
+          "value": 1024,
+          "cachedValue": 1024
+        },
+        "fps": {
+          "isAuto": false,
+          "value": 24,
+          "cachedValue": 24
+        },
+        "image_path": {
+          "isAuto": true,
+          "value": "X://insert/path/here.mp4",
+          "cachedValue": "X://insert/path/here.mp4"
+        },
+        "enable_teacache": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 1,
+      "type": "VideoGenerator",
+      "pos": [
+        818.804931640625,
+        348.9299621582031
+      ],
+      "size": [
+        400,
+        436
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "inference_args",
+          "shape": 7,
+          "type": "INFERENCE_ARGS",
+          "link": 3
+        },
+        {
+          "name": "vae_config",
+          "shape": 7,
+          "type": "VAE_CONFIG",
+          "link": 2
+        },
+        {
+          "name": "text_encoder_config",
+          "shape": 7,
+          "type": "TEXT_ENCODER_CONFIG",
+          "link": 7
+        },
+        {
+          "name": "dit_config",
+          "shape": 7,
+          "type": "DIT_CONFIG",
+          "link": 6
+        }
+      ],
+      "outputs": [
+        {
+          "name": "video_path",
+          "type": "STRING",
+          "links": [
+            4
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VideoGenerator"
+      },
+      "widgets_values": [
+        "A woman crying from laughter.",
+        "/workspace/ComfyUI/outputs_video/",
+        4,
+        "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
+        6,
+        2,
+        2,
+        "fp16",
+        true,
+        true,
+        "fp16",
+        "fp16",
+        true
+      ],
+      "auto_widget_states": {
+        "embedded_cfg_scale": {
+          "isAuto": true,
+          "value": 6,
+          "cachedValue": 6
+        },
+        "sp_size": {
+          "isAuto": true,
+          "value": 2,
+          "cachedValue": 2
+        },
+        "tp_size": {
+          "isAuto": true,
+          "value": 2,
+          "cachedValue": 2
+        },
+        "vae_precision": {
+          "isAuto": true,
+          "value": "fp16",
+          "cachedValue": "fp16"
+        },
+        "vae_tiling": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "vae_sp": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        },
+        "text_encoder_precision": {
+          "isAuto": true,
+          "value": "fp16",
+          "cachedValue": "fp16"
+        },
+        "precision": {
+          "isAuto": true,
+          "value": "fp16",
+          "cachedValue": "fp16"
+        },
+        "use_cpu_offload": {
+          "isAuto": true,
+          "value": true,
+          "cachedValue": true
+        }
+      }
+    },
+    {
+      "id": 5,
+      "type": "TextEncoderConfig",
+      "pos": [
+        416.4937744140625,
+        953.6171875
+      ],
+      "size": [
+        270,
+        106
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "text_encoder_config",
+          "type": "TEXT_ENCODER_CONFIG",
+          "links": [
+            7
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "TextEncoderConfig"
+      },
+      "widgets_values": [
+        "",
+        "",
+        ""
+      ],
+      "auto_widget_states": {
+        "prefix": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        },
+        "quant_config": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        },
+        "lora_config": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        }
+      }
+    },
+    {
+      "id": 6,
+      "type": "DITConfig",
+      "pos": [
+        415.1928405761719,
+        1154.1573486328125
+      ],
+      "size": [
+        270,
+        82
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "dit_config",
+          "type": "DIT_CONFIG",
+          "links": [
+            6
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "DITConfig"
+      },
+      "widgets_values": [
+        "",
+        ""
+      ],
+      "auto_widget_states": {
+        "prefix": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        },
+        "quant_config": {
+          "isAuto": true,
+          "value": "",
+          "cachedValue": ""
+        }
+      }
+    }
+  ],
+  "links": [
+    [
+      1,
+      7,
+      0,
+      2,
+      0,
+      "STRING"
+    ],
+    [
+      2,
+      4,
+      0,
+      1,
+      1,
+      "VAE_CONFIG"
+    ],
+    [
+      3,
+      2,
+      0,
+      1,
+      0,
+      "INFERENCE_ARGS"
+    ],
+    [
+      4,
+      1,
+      0,
+      3,
+      2,
+      "STRING"
+    ],
+    [
+      5,
+      3,
+      0,
+      8,
+      0,
+      "IMAGE"
+    ],
+    [
+      6,
+      6,
+      0,
+      1,
+      3,
+      "DIT_CONFIG"
+    ],
+    [
+      7,
+      5,
+      0,
+      1,
+      2,
+      "TEXT_ENCODER_CONFIG"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.8264462809917354,
+      "offset": [
+        646.7950212991898,
+        66.17259910028655
+      ]
+    },
+    "frontendVersion": "1.20.4",
+    "VHS_latentpreview": false,
+    "VHS_latentpreviewrate": 0,
+    "VHS_MetadataImage": true,
+    "VHS_KeepIntermediate": true
+  },
+  "version": 0.4
+}

comfyui/video_generator/dit_config.py

@@ -0,0 +1,31 @@
+class DITConfig:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "optional": {
+                "prefix": ("STRING", {
+                    "default": ""
+                }),
+                "quant_config": ("STRING", {
+                    "default": ""
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("DIT_CONFIG", )
+    RETURN_NAMES = ("dit_config", )
+    FUNCTION = "set_args"
+    CATEGORY = "fastvideo"
+
+    def set_args(self, prefix, quant_config):
+        raw_args = {"prefix": prefix, "quant_config": quant_config}
+
+        # Filter out keys where value is -99999
+        args = {k: v for k, v in raw_args.items() if str(int(v)) != str(-99999)}
+
+        return (args, )

comfyui/video_generator/inference_args.py

@@ -0,0 +1,89 @@
+class InferenceArgs:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "optional": {
+                "height": ("INT", {
+                    "default": 720
+                }),
+                "width": ("INT", {
+                    "default": 1280
+                }),
+                "num_frames": ("INT", {
+                    "default": 45
+                }),
+                "num_inference_steps": ("INT", {
+                    "default": 6
+                }),
+                "guidance_scale": ("FLOAT", {
+                    "default": 1.0
+                }),
+                "flow_shift": ("INT", {
+                    "default": 17
+                }),
+                "seed": ("INT", {
+                    "default": 1024
+                }),
+                "fps": ("INT", {
+                    "default": 24
+                }),
+                "image_path": ("STRING", {
+                    "default": "X://insert/path/here.mp4"
+                }),
+                "enable_teacache": ([True, False], {
+                    "default": False
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("INFERENCE_ARGS", )
+    RETURN_NAMES = ("inference_args", )
+    FUNCTION = "set_args"
+    CATEGORY = "fastvideo"
+
+    def set_args(
+        self,
+        height,
+        width,
+        num_frames,
+        num_inference_steps,
+        guidance_scale,
+        flow_shift,
+        seed,
+        fps,
+        image_path,
+        enable_teacache,
+    ):
+        raw_args = {
+            "height": height,
+            "width": width,
+            "num_frames": num_frames,
+            "num_inference_steps": num_inference_steps,
+            "guidance_scale": guidance_scale,
+            "flow_shift": flow_shift,
+            "seed": seed,
+            "fps": fps,
+            "image_path": image_path,
+            "enable_teacache": enable_teacache,
+        }
+
+        # Filter out keys where value is -99999, handling different types properly
+        args = {}
+        for k, v in raw_args.items():
+            try:
+                if isinstance(v, str):
+                    if v != "-99999":
+                        args[k] = v
+                elif v != -99999:
+                    # If it's not a string, compare directly
+                    args[k] = v
+            except (ValueError, TypeError):
+                # Include any value that causes an error in comparison
+                args[k] = v
+
+        return (args, )

comfyui/video_generator/load_image.py

@@ -0,0 +1,103 @@
+import hashlib
+import os
+
+import folder_paths
+import numpy as np
+import torch
+from PIL import Image, ImageOps, ImageSequence
+
+from .node_helpers import pillow
+
+
+class LoadImagePath:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = [
+            f for f in os.listdir(input_dir)
+            if os.path.isfile(os.path.join(input_dir, f))
+        ]
+        files = folder_paths.filter_files_content_types(files, ["image"])
+        return {
+            "required": {
+                "image": (sorted(files), {
+                    "image_upload": True
+                })
+            },
+        }
+
+    CATEGORY = "fastvideo"
+
+    RETURN_TYPES = ("STRING", "IMAGE", "MASK")
+    RETURN_NAMES = ("image_path", "IMAGE", "MASK")
+    FUNCTION = "load_image"
+
+    def load_image(self, image):
+        image_path = folder_paths.get_annotated_filepath(image)
+
+        img = pillow(Image.open, image_path)
+
+        output_images: list[torch.Tensor] = []
+        output_masks: list[torch.Tensor] = []
+        w, h = None, None
+
+        excluded_formats = ['MPO']
+
+        for i in ImageSequence.Iterator(img):
+            processed_image = pillow(ImageOps.exif_transpose, i)
+            if processed_image is None:
+                continue
+
+            if processed_image.mode == 'I':
+                processed_image = processed_image.point(lambda i: i * (1 / 255))
+            image = processed_image.convert("RGB")
+
+            if len(output_images) == 0:
+                w = image.size[0]
+                h = image.size[1]
+
+            if image.size[0] != w or image.size[1] != h:
+                continue
+
+            image = np.array(image).astype(np.float32) / 255.0
+            image = torch.from_numpy(image)[
+                None,
+            ]
+            if 'A' in processed_image.getbands():
+                mask = np.array(processed_image.getchannel('A')).astype(
+                    np.float32) / 255.0
+                mask = 1. - torch.from_numpy(mask)
+            elif processed_image.mode == 'P' and 'transparency' in processed_image.info:
+                mask = np.array(
+                    processed_image.convert('RGBA').getchannel('A')).astype(
+                        np.float32) / 255.0
+                mask = 1. - torch.from_numpy(mask)
+            else:
+                mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
+            output_images.append(image)
+            output_masks.append(mask.unsqueeze(0))
+
+        if len(output_images) > 1 and img.format not in excluded_formats:
+            output_image = torch.cat(output_images, dim=0)
+            output_mask = torch.cat(output_masks, dim=0)
+        else:
+            output_image = output_images[0]
+            output_mask = output_masks[0]
+
+        return (image_path, output_image, output_mask)
+
+    @classmethod
+    def IS_CHANGED(s, image):
+        image_path = folder_paths.get_annotated_filepath(image)
+        m = hashlib.sha256()
+        with open(image_path, 'rb') as f:
+            m.update(f.read())
+        return m.digest().hex()
+
+    @classmethod
+    def VALIDATE_INPUTS(s, image):
+        if not folder_paths.exists_annotated_filepath(image):
+            return "Invalid image file: {}".format(image)
+
+        return True

comfyui/video_generator/node_helpers.py

@@ -0,0 +1,68 @@
+import hashlib
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+import torch
+from comfy.cli_args import args
+from PIL import ImageFile, UnidentifiedImageError
+
+T = TypeVar('T')
+
+
+def conditioning_set_values(conditioning: list[Any],
+                            values: dict[str, Any] | None = None) -> list[Any]:
+    if values is None:
+        values = {}
+    c = []
+    for t in conditioning:
+        n = [t[0], t[1].copy()]
+        for k in values:
+            n[1][k] = values[k]
+        c.append(n)
+
+    return c
+
+
+def pillow(fn: Callable[[Any], T], arg: Any) -> T:
+    prev_value = None
+    try:
+        x = fn(arg)
+    except (OSError, UnidentifiedImageError, ValueError
+            ):  #PIL issues #4472 and #2445, also fixes ComfyUI issue #3416
+        prev_value = ImageFile.LOAD_TRUNCATED_IMAGES
+        ImageFile.LOAD_TRUNCATED_IMAGES = True
+        x = fn(arg)
+    finally:
+        if prev_value is not None:
+            ImageFile.LOAD_TRUNCATED_IMAGES = prev_value
+    return x
+
+
+def hasher() -> Callable[[], Any]:
+    hashfuncs = {
+        "md5": hashlib.md5,
+        "sha1": hashlib.sha1,
+        "sha256": hashlib.sha256,
+        "sha512": hashlib.sha512
+    }
+    return hashfuncs[args.default_hashing_function]
+
+
+def string_to_torch_dtype(string: str) -> torch.dtype | None:
+    if string == "fp32":
+        return torch.float32
+    if string == "fp16":
+        return torch.float16
+    if string == "bf16":
+        return torch.bfloat16
+    return None
+
+
+def image_alpha_fix(destination: torch.Tensor,
+                    source: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    if destination.shape[-1] < source.shape[-1]:
+        source = source[..., :destination.shape[-1]]
+    elif destination.shape[-1] > source.shape[-1]:
+        destination = torch.nn.functional.pad(destination, (0, 1))
+        destination[..., -1] = 1.0
+    return destination, source

comfyui/video_generator/nodes.py

@@ -0,0 +1,24 @@
+from .dit_config import DITConfig
+from .inference_args import InferenceArgs
+from .load_image import LoadImagePath
+from .text_encoder_config import TextEncoderConfig
+from .vae_config import VAEConfig
+from .video_generator import VideoGenerator
+
+NODE_CLASS_MAPPINGS = {
+    "VideoGenerator": VideoGenerator,
+    "InferenceArgs": InferenceArgs,
+    "VAEConfig": VAEConfig,
+    "TextEncoderConfig": TextEncoderConfig,
+    "DITConfig": DITConfig,
+    "LoadImagePath": LoadImagePath
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "VideoGenerator": "Video Generator",
+    "InferenceArgs": "Inference Args",
+    "VAEConfig": "VAE Config",
+    "TextEncoderConfig": "Text Encoder Config",
+    "DITConfig": "DIT Config",
+    "LoadImagePath": "Load Image Path"
+}

comfyui/video_generator/text_encoder_config.py

@@ -0,0 +1,38 @@
+class TextEncoderConfig:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "optional": {
+                "prefix": ("STRING", {
+                    "default": ""
+                }),
+                "quant_config": ("STRING", {
+                    "default": ""
+                }),
+                "lora_config": ("STRING", {
+                    "default": ""
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("TEXT_ENCODER_CONFIG", )
+    RETURN_NAMES = ("text_encoder_config", )
+    FUNCTION = "set_args"
+    CATEGORY = "fastvideo"
+
+    def set_args(self, prefix, quant_config, lora_config):
+        raw_args = {
+            "prefix": prefix,
+            "quant_config": quant_config,
+            "lora_config": lora_config
+        }
+
+        # Filter out keys where value is -99999
+        args = {k: v for k, v in raw_args.items() if str(int(v)) != str(-99999)}
+
+        return (args, )

comfyui/video_generator/vae_config.py

@@ -0,0 +1,88 @@
+class VAEConfig:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "optional": {
+                "load_encoder": ([True, False], {
+                    "default": True
+                }),
+                "load_decoder": ([True, False], {
+                    "default": True
+                }),
+                "tile_sample_min_height": ("INT", {
+                    "default": 256
+                }),
+                "tile_sample_min_width": ("INT", {
+                    "default": 256
+                }),
+                "tile_sample_min_num_frames": ("INT", {
+                    "default": 16
+                }),
+                "tile_sample_stride_height": ("INT", {
+                    "default": 192
+                }),
+                "tile_sample_stride_width": ("INT", {
+                    "default": 192
+                }),
+                "tile_sample_stride_num_frames": ("INT", {
+                    "default": 12
+                }),
+                "blend_num_frames": ("INT", {
+                    "default": 0
+                }),
+                "use_tiling": ([True, False], {
+                    "default": True
+                }),
+                "use_temporal_tiling": ([True, False], {
+                    "default": True
+                }),
+                "use_parallel_tiling": ([True, False], {
+                    "default": True
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("VAE_CONFIG", )
+    RETURN_NAMES = ("vae_config", )
+    FUNCTION = "set_args"
+    CATEGORY = "fastvideo"
+
+    def set_args(
+        self,
+        load_encoder,
+        load_decoder,
+        tile_sample_min_height,
+        tile_sample_min_width,
+        tile_sample_min_num_frames,
+        tile_sample_stride_height,
+        tile_sample_stride_width,
+        tile_sample_stride_num_frames,
+        blend_num_frames,
+        use_tiling,
+        use_temporal_tiling,
+        use_parallel_tiling,
+    ):
+        raw_args = {
+            "load_encoder": load_encoder,
+            "load_decoder": load_decoder,
+            "tile_sample_min_height": tile_sample_min_height,
+            "tile_sample_min_width": tile_sample_min_width,
+            "tile_sample_min_num_frames": tile_sample_min_num_frames,
+            "tile_sample_stride_height": tile_sample_stride_height,
+            "tile_sample_stride_width": tile_sample_stride_width,
+            "tile_sample_stride_num_frames": tile_sample_stride_num_frames,
+            "blend_num_frames": blend_num_frames,
+            "use_tiling": use_tiling,
+            "use_temporal_tiling": use_temporal_tiling,
+            "use_parallel_tiling": use_parallel_tiling,
+        }
+
+        # Filter out any value explicitly set to -99999
+        args = {k: v for k, v in raw_args.items() if str(int(v)) != str(-99999)}
+
+        return (args, )

comfyui/video_generator/video_generator.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+import glob
+import os
+import signal
+import sys
+import threading
+import time
+from typing import Any
+
+from comfy.model_management import processing_interrupted
+
+from fastvideo import PipelineConfig
+from fastvideo import VideoGenerator as FastVideoGenerator
+
+sys.path.insert(
+    0,
+    os.path.dirname(
+        os.path.dirname(
+            os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
+
+
+# Custom exception for interruption
+class GenerationInterruptedException(Exception):
+    pass
+
+
+# Custom exception for interruption that ComfyUI will recognize
+class GenerationCancelledException(Exception):
+
+    def __init__(self,
+                 message: str = "Generation was cancelled by user") -> None:
+        self.message = message
+        super().__init__(self.message)
+
+
+def update_config_from_args(config: Any, args_dict: dict[str, Any]) -> None:
+    """
+    Update configuration object from arguments dictionary.
+    
+    Args:
+        config: The configuration object to update
+        args_dict: Dictionary containing arguments
+    """
+    for key, value in args_dict.items():
+        if hasattr(config, key) and value is not None:
+            if key == "text_encoder_precisions" and isinstance(value, list):
+                setattr(config, key, tuple(value))
+            else:
+                setattr(config, key, value)
+
+
+class VideoGenerator:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": ("STRING", {
+                    "multiline":
+                    True,
+                    "default":
+                    "A ripe orange tumbles gently from a tree and lands on the head of a lounging capybara, "
+                    "who blinks slowly in response. The moment is quietly humorous and oddly serene, framed by "
+                    "lush green foliage and dappled sunlight. Mid-shot, warm and whimsical tones."
+                }),
+                "output_path": ("STRING", {
+                    "default": "/workspace/ComfyUI/outputs_video/"
+                }),
+                "num_gpus": ("INT", {
+                    "default": 2,
+                    "min": 1,
+                    "max": 16
+                }),
+                "model_path": ("STRING", {
+                    "default": "FastVideo/FastHunyuan-diffusers"
+                })
+            },
+            "optional": {
+                "inference_args": ("INFERENCE_ARGS", ),
+                "embedded_cfg_scale": ("FLOAT", {
+                    "default": 6.0
+                }),
+                "sp_size": ("INT", {
+                    "default": 2
+                }),
+                "tp_size": ("INT", {
+                    "default": 2
+                }),
+                "vae_config": ("VAE_CONFIG", ),
+                "vae_precision": (["fp16", "bf16"], {
+                    "default": "fp16"
+                }),
+                "vae_tiling": ([True, False], {
+                    "default": True
+                }),
+                "vae_sp": ([True, False], {
+                    "default": False
+                }),
+                "text_encoder_config": ("TEXT_ENCODER_CONFIG", ),
+                "text_encoder_precision": (["fp16", "bf16"], {
+                    "default": "fp16"
+                }),
+                "dit_config": ("DIT_CONFIG", ),
+                "precision": (["fp16", "bf16"], {
+                    "default": "fp16"
+                }),
+                "use_cpu_offload": ([True, False], {
+                    "default": False
+                }),
+            }
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        return True
+
+    RETURN_TYPES = ("STRING", )
+    RETURN_NAMES = ("video_path", )
+    FUNCTION = "launch_inference"
+    CATEGORY = "fastvideo"
+
+    generator: FastVideoGenerator | None = None
+    _interrupt_thread: threading.Thread | None = None
+    _generation_active: bool = False
+    _generation_interrupted: bool = False
+    _interrupt_event: threading.Event = threading.Event()
+    _generation_thread: threading.Thread | None = None
+    _generation_result: str | None = None
+    _generation_exception: Exception | None = None
+
+    def _monitor_for_interruption(self):
+        """Background thread that monitors for interruption requests"""
+        time.sleep(2)  # Give the generation thread time to send execute_forward
+
+        while self._generation_active and not self._interrupt_event.is_set():
+            if processing_interrupted():
+                print("Video generation interrupted by user")
+                self._generation_interrupted = True
+
+                # Try to send interrupt signal to worker processes
+                if self.generator is not None and hasattr(
+                        self.generator, 'executor'):
+                    try:
+                        # The MultiprocExecutor has a workers attribute
+                        if hasattr(self.generator.executor, 'workers'):
+                            for worker in self.generator.executor.workers:
+                                if worker.is_alive():
+                                    os.kill(worker.pid, signal.SIGINT)
+                            print("Interrupt signal sent to worker processes")
+                    except Exception as e:
+                        print(f"Error sending interrupt signal: {e}")
+
+                # Set the interrupt event to notify other threads
+                self._interrupt_event.set()
+                break
+            time.sleep(0.5)
+
+    def _run_generation(self, prompt: str, output_path: str,
+                        inference_args: dict[str, Any]) -> None:
+        """Thread function to run the generation"""
+        try:
+            if self.generator is not None:
+                self.generator.generate_video(prompt=prompt,
+                                              output_path=output_path,
+                                              **inference_args)
+                self._generation_result = os.path.join(output_path,
+                                                       f"{prompt[:100]}.mp4")
+            else:
+                raise RuntimeError("Generator is not initialized")
+        except Exception as e:
+            self._generation_exception = e
+            self._interrupt_event.set()
+
+    def load_output_video(self, output_dir):
+        video_extensions = ["*.mp4", "*.avi", "*.mov", "*.mkv"]
+        video_files = []
+
+        for ext in video_extensions:
+            video_files.extend(glob.glob(os.path.join(output_dir, ext)))
+
+        if not video_files:
+            print("No video files found in output directory: %s", output_dir)
+            return ""
+
+        video_files.sort()
+        return video_files[0]
+
+    def launch_inference(
+        self,
+        prompt,
+        output_path,
+        num_gpus,
+        model_path,
+        embedded_cfg_scale,
+        sp_size,
+        tp_size,
+        vae_precision,
+        vae_tiling,
+        vae_sp,
+        text_encoder_precision,
+        precision,
+        inference_args=None,
+        vae_config=None,
+        text_encoder_config=None,
+        dit_config=None,
+        use_cpu_offload=None,
+    ):
+        print('Running FastVideo inference')
+
+        # Reset interruption flag and event
+        self._generation_interrupted = False
+        self._interrupt_event.clear()
+        self._generation_result = None
+        self._generation_exception = None
+
+        # Load pipeline config from model path
+        pipeline_config = PipelineConfig.from_pretrained(model_path)
+        print('pipeline_config', pipeline_config)
+
+        # Update configs with provided config dictionaries
+        if dit_config is not None:
+            update_config_from_args(pipeline_config.dit_config, dit_config)
+
+        if vae_config is not None:
+            update_config_from_args(pipeline_config.vae_config, vae_config)
+
+        if text_encoder_config is not None:
+            update_config_from_args(pipeline_config.text_encoder_configs,
+                                    text_encoder_config)
+
+        # Update top-level pipeline config with remaining arguments
+        raw_pipeline_args = {}
+        if embedded_cfg_scale is not None:
+            raw_pipeline_args['embedded_cfg_scale'] = embedded_cfg_scale
+        if precision is not None:
+            raw_pipeline_args['precision'] = precision
+        if vae_precision is not None:
+            raw_pipeline_args['vae_precision'] = vae_precision
+        if vae_tiling is not None:
+            raw_pipeline_args['vae_tiling'] = vae_tiling
+        if vae_sp is not None:
+            raw_pipeline_args['vae_sp'] = vae_sp
+        if text_encoder_precision is not None:
+            raw_pipeline_args['text_encoder_precision'] = text_encoder_precision
+
+        # Filter out any value explicitly set to -99999 (auto values)
+        pipeline_args = {
+            k: v
+            for k, v in raw_pipeline_args.items() if str(int(v)) != str(-99999)
+        }
+
+        update_config_from_args(pipeline_config, pipeline_args)
+
+        raw_generation_args = {}
+        if num_gpus is not None:
+            raw_generation_args['num_gpus'] = num_gpus
+        if tp_size is not None:
+            raw_generation_args['tp_size'] = tp_size
+        if sp_size is not None:
+            raw_generation_args['sp_size'] = sp_size
+        if use_cpu_offload is not None:
+            raw_generation_args['use_cpu_offload'] = use_cpu_offload
+
+        generation_args = {
+            k: v
+            for k, v in raw_generation_args.items()
+            if str(int(v)) != str(-99999)
+        }
+
+        if self.generator is None:
+            print('generation_args', generation_args)
+            print('pipeline_config', pipeline_config)
+            self.generator = FastVideoGenerator.from_pretrained(
+                model_path=model_path,
+                **generation_args,
+                pipeline_config=pipeline_config)
+
+        print('inference_args', inference_args)
+
+        # Start a thread to run the generation
+        self._generation_thread = threading.Thread(target=self._run_generation,
+                                                   args=(prompt, output_path,
+                                                         inference_args),
+                                                   daemon=True)
+        self._generation_thread.start()
+
+        # Start a background thread to monitor for interruptions
+        self._generation_active = True
+        self._interrupt_thread = threading.Thread(
+            target=self._monitor_for_interruption, daemon=True)
+        self._interrupt_thread.start()
+
+        # Wait for either completion or interruption
+        while self._generation_thread.is_alive(
+        ) and not self._interrupt_event.is_set():
+            self._generation_thread.join(timeout=0.5)
+
+        self._generation_active = False
+        if self._interrupt_thread:
+            self._interrupt_thread.join(timeout=1.0)
+            self._interrupt_thread = None
+
+        if self._generation_interrupted:
+            print("Video generation was cancelled by user")
+            raise GenerationCancelledException()
+        elif self._generation_exception:
+            # Re-raise the exception from the generation thread
+            raise self._generation_exception
+        elif self._generation_result:
+            return (self._generation_result, )
+        else:
+            # This shouldn't happen, but just in case
+            print("Generation completed but no result was produced")
+            raise Exception("Generation failed to produce a result")

comfyui/web/js/FastVideo.core.js

@@ -0,0 +1,593 @@
+import { app } from '../../../scripts/app.js'
+
+function chainCallback(object, property, callback) {
+    if (object == undefined) {
+        console.error("Tried to add callback to non-existent object");
+        return;
+    }
+    if (property in object && object[property]) {
+        const callback_orig = object[property];
+        object[property] = function () {
+            const r = callback_orig.apply(this, arguments);
+            return callback.apply(this, arguments) ?? r;
+        };
+    } else {
+        object[property] = callback;
+    }
+}
+
+function drawAutoAnnotated(ctx, node, widget_width, y, H) {
+    const litegraph_base = LiteGraph;
+    const show_text = app.canvas.ds.scale >= 0.5;
+    const margin = 15;
+
+    const autoTextWidth = 30;
+    const autoTextRightMargin = 5;
+
+    ctx.textAlign = 'left';
+    ctx.strokeStyle = litegraph_base.WIDGET_OUTLINE_COLOR;
+    ctx.fillStyle = litegraph_base.WIDGET_BGCOLOR;
+
+    ctx.beginPath();
+    if (show_text && ctx.roundRect) {
+        ctx.roundRect(margin, y, widget_width - margin * 2, H, [H * 0.5]);
+    } else {
+        ctx.rect(margin, y, widget_width - margin * 2, H);
+    }
+    ctx.fill();
+
+    if (show_text) {
+        if (!this.disabled) ctx.stroke();
+        const isAuto = this.isAuto === true;
+
+        ctx.save();
+        if (isAuto) {
+            ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
+            ctx.strokeStyle = litegraph_base.WIDGET_TEXT_COLOR;
+        } else {
+            ctx.fillStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
+            ctx.strokeStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
+        }
+
+        // Position for the cog
+        const cogX = widget_width - autoTextRightMargin - autoTextWidth - 6;
+        const cogY = y + H * 0.5;
+        const cogRadius = 6;
+        const toothLength = 2;
+        const numTeeth = 8;
+        const holeRadius = 2; // Radius of the center hole
+
+        // Draw the cog
+        ctx.beginPath();
+        ctx.arc(cogX, cogY, cogRadius - toothLength, 0, Math.PI * 2);
+        ctx.fill();
+
+        // Draw the center hole (by clearing it)
+        ctx.beginPath();
+        ctx.arc(cogX, cogY, holeRadius, 0, Math.PI * 2);
+        ctx.fillStyle = litegraph_base.WIDGET_BGCOLOR;
+        ctx.fill();
+
+        // Reset fill style for the teeth
+        if (isAuto) {
+            ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
+        } else {
+            ctx.fillStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
+        }
+
+        // Draw teeth
+        ctx.beginPath();
+        for (let i = 0; i < numTeeth; i++) {
+            const angle = (i / numTeeth) * Math.PI * 2;
+            const innerX = cogX + (cogRadius - toothLength) * Math.cos(angle);
+            const innerY = cogY + (cogRadius - toothLength) * Math.sin(angle);
+            const outerX = cogX + cogRadius * Math.cos(angle);
+            const outerY = cogY + cogRadius * Math.sin(angle);
+
+            ctx.moveTo(innerX, innerY);
+            ctx.lineTo(outerX, outerY);
+        }
+        ctx.lineWidth = 2;
+        ctx.stroke();
+
+        ctx.restore();
+
+        // Draw label
+        ctx.fillStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
+        const label = this.label || this.name;
+        if (label != null) {
+            ctx.fillText(label, margin * 2 + 5, y + H * 0.7);
+        }
+
+        // Draw value
+        ctx.textAlign = 'right';
+        const text = isAuto ? "auto" : this.displayValue();
+        ctx.fillStyle = isAuto ? litegraph_base.WIDGET_SECONDARY_TEXT_COLOR : litegraph_base.WIDGET_TEXT_COLOR;
+        ctx.fillText(text, widget_width - autoTextRightMargin - autoTextWidth - 15, y + H * 0.7);
+
+        // Draw increment/decrement buttons if not in AUTO mode and not a string widget
+        if (!isAuto && !this.disabled && this.config[0] !== "FVAUTOSTRING") {
+            // Draw decrement button (left triangle)
+            ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
+            ctx.beginPath();
+            ctx.moveTo(margin + 16, y + 5);
+            ctx.lineTo(margin + 6, y + H * 0.5);
+            ctx.lineTo(margin + 16, y + H - 5);
+            ctx.fill();
+
+            // Draw increment button (right triangle)
+            ctx.beginPath();
+            ctx.moveTo(widget_width - margin - 16, y + 5);
+            ctx.lineTo(widget_width - margin - 6, y + H * 0.5);
+            ctx.lineTo(widget_width - margin - 16, y + H - 5);
+            ctx.fill();
+        }
+    }
+}
+
+function mouseAutoAnnotated(event, [x, y], node) {
+    const widget_width = node.size[0];
+    const margin = 15;
+    const H = 20; // Widget height
+
+    const autoTextWidth = 30;
+    const autoTextRightMargin = 5;
+
+    const cogRadius = 6;
+
+    if (this.isAuto) {
+        if (event.type === "pointerup" || event.type === "mouseup") {
+            const cogX = widget_width - autoTextRightMargin - autoTextWidth - 6;
+            const cogLeftEdge = cogX - cogRadius;
+            const cogRightEdge = cogX + cogRadius;
+
+            if (x > cogLeftEdge && x < cogRightEdge) {
+                this.isAuto = false;
+                this.value = this.cachedValue !== undefined ? this.cachedValue : (this.options.default || 0);
+
+                if (this.callback) {
+                    this.callback(this.value);
+                }
+                node.graph.setDirtyCanvas(true, false);
+            }
+        }
+
+        // Block ALL events in auto mode except cog clicks
+        event.preventDefault?.();
+        event.stopPropagation?.();
+        event.stopImmediatePropagation?.();
+        return true; // Always return true to indicate event was handled
+    }
+
+    // Determine if clicking on increment/decrement buttons
+    const delta = this.config[0] === "FVAUTOSTRING" ? 0 :
+        (x < 40 ? -1 : x > widget_width - 48 ? 1 : 0);
+
+    if (event.type === "pointerdown" || event.type === "mousedown") {
+        // ComfyUI appears to intercept pointerdown events, so this code path is never reached
+        console.log("pointerdown received (unexpected)");
+        return false;
+    } else if (event.type === "pointerup" || event.type === "mouseup") {
+        // Stop event propagation to prevent double handling
+        event.preventDefault?.();
+        event.stopPropagation?.();
+        event.stopImmediatePropagation?.();
+
+        const cogX = widget_width - autoTextRightMargin - autoTextWidth - 6;
+        const cogLeftEdge = widget_width - autoTextRightMargin - autoTextWidth - 6 - cogRadius;
+        const cogRightEdge = widget_width - autoTextRightMargin - autoTextWidth - 6 + cogRadius;
+
+        if (x > cogLeftEdge && x < cogRightEdge) {
+            this.isAuto = !this.isAuto;
+
+            if (this.isAuto) {
+                this.cachedValue = this.value;
+                this.value = -99999;
+            } else {
+                this.value = this.cachedValue !== undefined ? this.cachedValue : (this.options.default || 0);
+            }
+
+            if (this.callback) {
+                this.callback(this.value);
+            }
+
+            node.graph.setDirtyCanvas(true, false);
+            return true;
+        }
+
+        // If in auto mode and NOT clicking the cog, block all other interactions
+        if (this.isAuto) {
+            return true;
+        }
+
+        // Handle increment/decrement buttons if not in auto mode
+        if (delta !== 0 && !this.isAuto) {
+            if (this.config[0] === "FVAUTOCOMBO") {
+                const options = this.options.values || [];
+                if (options.length === 0) return true;
+
+                let currentIndex = -1;
+                for (let i = 0; i < options.length; i++) {
+                    const optValue = typeof options[i] === 'object' ? options[i].value : options[i];
+                    if (optValue == this.value || String(optValue) === String(this.value)) {
+                        currentIndex = i;
+                        break;
+                    }
+                }
+
+                if (currentIndex === -1) {
+                    currentIndex = 0;
+                }
+
+                let newIndex = currentIndex + delta;
+                if (newIndex < 0) {
+                    newIndex = options.length - 1;
+                } else if (newIndex >= options.length) {
+                    newIndex = 0;
+                }
+
+                const newOption = options[newIndex];
+                this.value = typeof newOption === 'object' ? newOption.value : newOption;
+
+                if (this.callback) {
+                    this.callback(this.value);
+                }
+
+                node.graph.setDirtyCanvas(true, false);
+                return true;
+            } else {
+                let v = parseFloat(this.value);
+                const increment = delta * 0.1 * (this.options.step || 1);
+
+                v += increment;
+
+                // Apply min/max constraints
+                if (this.options.min != null) {
+                    v = Math.max(this.options.min, v);
+                }
+                if (this.options.max != null) {
+                    v = Math.min(this.options.max, v);
+                }
+
+                // Round to precision or to integer
+                if (this.config[0] === "FVAUTOINT") {
+                    v = Math.round(v);
+                } else if (this.options.precision !== undefined) {
+                    const precision = Math.pow(10, this.options.precision);
+                    v = Math.round(v * precision) / precision;
+                }
+
+                this.value = v;
+
+                if (this.callback) {
+                    this.callback(this.value);
+                }
+
+                node.graph.setDirtyCanvas(true, false);
+                return true;
+            }
+        }
+
+        if (delta === 0 && !this.isAuto) {
+            if (this.config[0] === "FVAUTOCOMBO") {
+                const options = this.options.values || [];
+
+                // Create menu items
+                const menuItems = options.map(opt => {
+                    const value = typeof opt === 'object' ? opt.value : opt;
+                    const label = typeof opt === 'object' ? opt.label : opt.toString();
+
+                    return {
+                        content: label,
+                        callback: () => {
+                            this.value = value;
+
+                            if (this.callback) {
+                                this.callback(this.value);
+                            }
+
+                            node.graph.setDirtyCanvas(true, false);
+                        }
+                    };
+                });
+
+                new LiteGraph.ContextMenu(menuItems, {
+                    event: event,
+                    title: null,
+                    callback: null,
+                    extra: node
+                });
+
+                return true;
+            } else if (this.config[0] === "FVAUTOSTRING") {
+                const d_callback = (v) => {
+                    this.value = v;
+
+                    if (this.callback) {
+                        this.callback(this.value);
+                    }
+
+                    node.graph.setDirtyCanvas(true, false);
+                };
+
+                const dialog = app.canvas.prompt(
+                    'Value',
+                    this.value,
+                    d_callback,
+                    event
+                );
+
+                return true;
+            } else {
+                // For numeric widgets, show input dialog
+                const d_callback = (v) => {
+                    this.value = this.parseValue?.(v) ?? Number(v);
+
+                    // Apply min/max constraints
+                    if (this.options.min != null) {
+                        this.value = Math.max(this.options.min, this.value);
+                    }
+                    if (this.options.max != null) {
+                        this.value = Math.min(this.options.max, this.value);
+                    }
+
+                    // Round to precision or to integer
+                    if (this.config[0] === "FVAUTOINT") {
+                        this.value = Math.round(this.value);
+                    } else if (this.options.precision !== undefined) {
+                        const precision = Math.pow(10, this.options.precision);
+                        this.value = Math.round(this.value * precision) / precision;
+                    }
+
+                    if (this.callback) {
+                        this.callback(this.value);
+                    }
+
+                    node.graph.setDirtyCanvas(true, false);
+                };
+
+                const dialog = app.canvas.prompt(
+                    'Value',
+                    this.value,
+                    d_callback,
+                    event
+                );
+
+                return true;
+            }
+        }
+
+        return true;
+    }
+
+    return false;
+}
+
+function makeAutoAnnotated(widget, inputData) {
+    const original = {
+        callback: widget.callback,
+        type: widget.type,
+        value: widget.value
+    };
+
+    // Add AUTO properties to the widget
+    Object.assign(widget, {
+        type: "BOOLEAN",
+        draw: drawAutoAnnotated,
+        mouse: mouseAutoAnnotated,
+        onMouse: null, // Explicitly disable original onMouse handler
+
+        isAuto: true,
+        cachedValue: widget.value,
+        config: inputData,
+        options: Object.assign({}, inputData[1], widget.options),
+        original: original,  // Store original properties for reference
+
+        // Disable other potential mouse handlers with no-op functions
+        onClick: function () {
+            return false;
+        },
+        onPointerUp: function () {
+            return false;
+        },
+        onPointerDown: function () {
+            return false;
+        },
+        onMouseUp: function () {
+            return false;
+        },
+        onMouseDown: function () {
+            return false;
+        },
+
+        computeSize(width) {
+            return [width, 20];
+        },
+        displayValue: function () {
+            if (this.config[0] === "FVAUTOINT") {
+                return Math.round(this.value).toString();
+            }
+            if (this.config[0] === "FVAUTOCOMBO") {
+                return this.value;
+            }
+            if (this.config[0] === "FVAUTOSTRING") {
+                return this.value;
+            }
+            // For FLOAT values, check if it's actually an integer
+            if (Number.isInteger(this.value)) {
+                return this.value.toString();
+            }
+            return this.value.toFixed(this.options.precision || 2);
+        },
+        parseValue: function (v) {
+            if (this.config[0] === "FVAUTOSTRING") {
+                return v;
+            }
+            if (typeof v === "string") {
+                return parseFloat(v);
+            }
+            return v;
+        },
+        serializeValue: function () {
+            // Return special value for AUTO mode
+            return this.isAuto ? -99999 : this.value;
+        },
+        deserializeValue: function (data) {
+            if (data === -99999) {
+                this.isAuto = true;
+                this.value = -99999;
+            } else {
+                this.isAuto = false;
+                this.value = data;
+                this.cachedValue = data;
+            }
+        }
+    });
+
+    // Override callback to handle AUTO mode
+    widget.callback = function (v) {
+        if (this.isAuto) {
+            return; // Don't call the original callback in AUTO mode
+        }
+        const result = original.callback?.call(this, v);
+        return result;
+    };
+
+    // Override any potential click handlers
+    const originalOnClick = widget.onClick;
+    if (originalOnClick) {
+        widget.onClick = function (...args) {
+            if (this.isAuto) {
+                return false;
+            }
+            return originalOnClick.call(this, ...args);
+        };
+    }
+
+    return widget;
+}
+
+app.registerExtension({
+    name: "FastVideo.AutoWidgets",
+
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (nodeData?.name == "VideoGenerator" || nodeData?.name === "InferenceArgs" || nodeData?.name === "VAEConfig" ||
+            nodeData?.name === "TextEncoderConfig" || nodeData?.name === "DITConfig") {
+            // Add serialization support
+            chainCallback(nodeType.prototype, "onSerialize", function (info) {
+                if (!this.widgets) {
+                    return;
+                }
+                // Ensure widgets_values exists
+                if (!info.widgets_values) {
+                    info.widgets_values = {};
+                }
+
+                // Store AUTO widget states in a separate property
+                if (!info.auto_widget_states) {
+                    info.auto_widget_states = {};
+                }
+
+                // Handle AUTO widgets specially
+                for (const w of this.widgets) {
+                    if (w.type === "BOOLEAN" && w.isAuto !== undefined) {
+                        // Store the serialized value (for Python node)
+                        info.widgets_values[w.name] = w.serializeValue();
+
+                        // Store the full state (for UI restoration)
+                        info.auto_widget_states[w.name] = {
+                            isAuto: w.isAuto,
+                            value: w.value,
+                            cachedValue: w.cachedValue
+                        };
+                    }
+                }
+            });
+
+            // Add deserialization support
+            chainCallback(nodeType.prototype, "onConfigure", function (info) {
+                if (!this.widgets) {
+                    return;
+                }
+
+                // First, restore from widgets_values (for backward compatibility)
+                if (info.widgets_values && Array.isArray(info.widgets_values)) {
+                    for (let i = 0; i < this.widgets.length && i < info.widgets_values.length; i++) {
+                        const w = this.widgets[i];
+                        const value = info.widgets_values[i];
+
+                        if (w.type === "BOOLEAN" && w.isAuto !== undefined) {
+                            w.deserializeValue(value);
+                        }
+                    }
+                }
+
+                // Then, restore full state if available
+                if (info.auto_widget_states) {
+                    for (const w of this.widgets) {
+                        if (w.type === "BOOLEAN" && w.isAuto !== undefined && w.name in info.auto_widget_states) {
+                            const state = info.auto_widget_states[w.name];
+
+                            w.isAuto = state.isAuto;
+                            w.cachedValue = state.cachedValue;
+                            w.value = state.isAuto ? -99999 : state.value;
+
+                            w.callback?.(w.value);
+                        }
+                    }
+                }
+
+                // Force a redraw
+                this.graph?.setDirtyCanvas(true, true);
+            });
+
+
+            // Override addInput to handle AUTO widgets
+            chainCallback(nodeType.prototype, "onNodeCreated", function () {
+                // Convert any existing widgets to AUTO widgets if needed
+                let new_widgets = [];
+                const intWidgetNames = ["sp_size", "tp_size", "height", "width", "num_frames", "num_inference_steps", "flow_shift", "seed", "fps", "scale_factor",
+                    "tile_sample_min_height", "tile_sample_min_width", "tile_sample_min_num_frames", "tile_sample_stride_height", "tile_sample_stride_width",
+                    "tile_sample_stride_num_frames", "blend_num_frames"
+                ]
+                const floatWidgetNames = ["embedded_cfg_scale", "guidance_scale"]
+                const comboWidgetNames = ["vae_tiling", "vae_precision", "vae_sp", "text_encoder_precision", "precision",
+                    "load_encoder", "load_decoder", "use_tiling", "use_temporal_tiling", "use_parallel_tiling", "use_cpu_offload", "enable_teacache"
+                ]
+                const stringWidgetNames = ["prefix", "quant_config", "lora_config", "image_path"]
+
+                if (this.widgets) {
+                    for (let w of this.widgets) {
+                        if (intWidgetNames.includes(w.name)) {
+                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOINT", { "default": 0 }]));
+                        } else if (floatWidgetNames.includes(w.name)) {
+                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOFLOAT", { "default": 0 }]));
+                        } else if (comboWidgetNames.includes(w.name)) {
+                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOCOMBO", { "default": 0 }]));
+                        } else if (stringWidgetNames.includes(w.name)) {
+                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOSTRING", { "default": "" }]));
+                        } else {
+                            new_widgets.push(w);
+                        }
+                    }
+                    this.widgets = new_widgets;
+
+                    const autoWidgets = this.widgets.filter(w => w.type === "BOOLEAN" && w.isAuto !== undefined);
+                }
+
+                this.graph?.setDirtyCanvas(true, true);
+            });
+        }
+    },
+
+    async init() {
+        // Force a redraw of all nodes when the extension initializes
+        if (app.graph) {
+            setTimeout(() => {
+                app.graph.setDirtyCanvas(true, true);
+            }, 1000);
+        }
+    }
+});
+
+console.log("FastVideo.core.js loaded");

csrc/attn/README.md

@@ -4,9 +4,8 @@
 
 
 ## Installation
-We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
+We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only support H100/H200, because ThunderKittens uses TMA but doesn't support Blackwell yet.
 First, install C++20 for ThunderKittens:
-
 ```bash
 sudo apt update
 sudo apt install gcc-11 g++-11
@@ -16,17 +15,27 @@ sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave
 sudo apt update
 sudo apt install clang-11
 ```
-Install STA:
+
+## Environment Setup
+First, set up your CUDA environment:
 ```bash
 export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 git submodule update --init --recursive
-python setup.py install
+```
+
+## Install Sliding Tile Attention (STA)
+```bash
+python setup_sta.py install
+```
+
+## Install Video Sparse Attention (VSA)
+```bash
+python setup_vsa.py install
 ```
 
 ## Usage
-
 ```python
 from st_attn import sliding_tile_attention
 # assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
@@ -44,8 +53,14 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)
 
 ## Test
 ```bash
-python test/test_sta.py
+python tests/test_sta.py # test STA
+python tests/test_block_sparse.py # test VSA
 ```
+## Benchmark
+```bash
+python benchmarks/bench_sta.py
+```
+
 
 ## How Does STA Work?
 We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 

csrc/attn/benchmarks/bench_sta.py

@@ -5,6 +5,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from st_attn import sliding_tile_attention
+from triton.testing import do_bench
 
 
 def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
@@ -13,16 +14,16 @@ def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
     return f if mode == "fwd" else (2.5 * f if mode == "bwd" else 3.5 * f)
 
 
-def efficiency(flop, time):
-    flop = flop / 1e12
-    time = time / 1e6
-    return flop / time
+def compute_TFLOPS(flops, ms):
+    flops = flops / 1e12
+    ms = ms / 1e3
+    return flops / ms
 
 
 def benchmark_attention(configurations):
     results = {'fwd': defaultdict(list), 'bwd': defaultdict(list)}
 
-    for B, H, N, D, causal in configurations:
+    for B, H, N, D, causal, dit_seq_shape, window_size in configurations:
         print("=" * 60)
         print(f"Timing forward and backward pass for B={B}, H={H}, N={N}, D={D}, causal={causal}")
 
@@ -30,38 +31,31 @@ def benchmark_attention(configurations):
         k = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
         v = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
 
-        grad_output = torch.randn_like(q, requires_grad=False).contiguous()
+        # grad_output = torch.randn_like(q, requires_grad=False).contiguous()
+        # qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
+        # kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
+        # vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()
 
-        qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
-        kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
-        vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()
 
-        # Prepare for timing forward pass
-        start_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
-        end_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+        # # Warmup for forward pass
+        # for _ in range(10):
+        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
 
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+        # # Time the forward pass
+        # for i in range(10):
+        #     start_events_fwd[i].record()
+        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
+        #     end_events_fwd[i].record()
+        ms = do_bench(lambda: sliding_tile_attention(q, k, v, [window_size] * 24, 0, False, dit_seq_shape))
 
-        # Warmup for forward pass
-        for _ in range(10):
-            o = sliding_tile_attention(q, k, v, [[6, 6, 6]] * 24, 0, False)
+        # times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
+        # time_us_fwd = np.mean(times_fwd) * 1000
 
-        # Time the forward pass
-        for i in range(10):
-            start_events_fwd[i].record()
-            o = sliding_tile_attention(q, k, v, [[6, 6, 6]] * 24, 0, False)
-            end_events_fwd[i].record()
-
-        torch.cuda.synchronize()
-        times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
-        time_us_fwd = np.mean(times_fwd) * 1000
-
-        tflops_fwd = efficiency(flops(B, N, H, D, causal, 'fwd'), time_us_fwd)
+        tflops_fwd = compute_TFLOPS(flops(B, N, H, D, causal, 'fwd'), ms)
         results['fwd'][(D, causal)].append((N, tflops_fwd))
 
-        print(f"Average time for forward pass in us: {time_us_fwd:.2f}")
-        print(f"Average efficiency for forward pass in TFLOPS: {tflops_fwd}")
+        print(f"Average time for forward pass (ms): {ms:.2f}")
+        print(f"Average TFLOPS: {tflops_fwd}")
         print("-" * 60)
 
         # torch.cuda.empty_cache()
@@ -85,15 +79,14 @@ def benchmark_attention(configurations):
         # times_bwd = [s.elapsed_time(e) for s, e in zip(start_events_bwd, end_events_bwd)]
         # time_us_bwd = np.mean(times_bwd) * 1000
 
-        # tflops_bwd = efficiency(flops(B, N, H, D, causal, 'bwd'), time_us_bwd)
+        # tflops_bwd = compute_TFLOPS(flops(B, N, H, D, causal, 'bwd'), ms)
         # results['bwd'][(D, causal)].append((N, tflops_bwd))
 
-        # print(f"Average time for backward pass in us: {time_us_bwd:.2f}")
-        # print(f"Average efficiency for backward pass in TFLOPS: {tflops_bwd}")
-        print("=" * 60)
+        # print(f"Average time for backward pass(ms): {ms:.2f}")
+        # print(f"Average TFLOPS: {tflops_bwd}")
+        # print("=" * 60)
 
         torch.cuda.empty_cache()
-        torch.cuda.synchronize()
 
     return results
 
@@ -124,7 +117,10 @@ def plot_results(results):
 
 # Example list of configurations to test
 configurations = [
-    (2, 24, 82944, 128, False),
+    (2, 24, 69120, 128, False, '18x48x80', [3, 6, 10]),
+    (2, 24, 69120, 128, True, '18x48x80', [3, 6, 10]),
+    (2, 24, 82944, 128, False, '36x48x48', [3, 3, 6]), # Stepvideo
+    (2, 24, 82944, 128, True, '36x48x48', [3, 3, 6]),
     # (16, 16, 768*16,  128, False),
     # (16, 16, 768*2,  128, False),
     # (16, 16, 768*4,  128, False),

csrc/attn/benchmarks/bench_vsa.py

@@ -0,0 +1,225 @@
+import torch
+import argparse
+from flash_attn.utils.benchmark import benchmark_forward
+from vsa import block_sparse_attention_fwd, block_sparse_attention_backward
+from vsa import BLOCK_M, BLOCK_N
+
+import numpy as np
+import random
+
+def set_seed(seed: int = 42):
+    # Python random module
+    random.seed(seed)
+
+    # NumPy
+    np.random.seed(seed)
+
+    # PyTorch
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # if using multi-GPU
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
+    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
+    parser.add_argument('--num_heads', type=int, default=12, help='Number of heads')
+    parser.add_argument('--head_dim', type=int, default=64, help='Head dimension')
+    parser.add_argument('--topk', type=int, default=None, help='Number of kv blocks each q block attends to')
+    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[49152], help='Sequence lengths to benchmark')
+    return parser.parse_args()
+
+def create_input_tensors(batch, head, seq_len, headdim):
+    """Create random input tensors for attention."""
+    q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    return q, k, v
+
+def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"):
+    """
+    Generate a block sparse pattern where each q block attends to exactly k kv blocks.
+    
+    Args:
+        bs: batch size
+        h: number of heads
+        num_q_blocks: number of query blocks
+        num_kv_blocks: number of key-value blocks
+        k: number of kv blocks each q block attends to
+        device: device to create tensors on
+        
+    Returns:
+        q2k_block_sparse_index: [bs, h, num_q_blocks, k]
+            Contains the indices of kv blocks that each q block attends to.
+        q2k_block_sparse_num: [bs, h, num_q_blocks]
+            Contains the number of kv blocks that each q block attends to (all equal to k).
+        k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks]
+            Contains the indices of q blocks that attend to each kv block.
+        k2q_block_sparse_num: [bs, h, num_kv_blocks]
+            Contains the number of q blocks that attend to each kv block.
+        block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]
+            Binary mask where 1 indicates attention connection.
+    """
+    # Ensure k is not larger than num_kv_blocks
+    k = min(k, num_kv_blocks)
+    
+    # Create random scores for sampling
+    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device)
+    
+    # Get top-k indices for each q block
+    _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1)
+    q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32)
+    
+    # sort q2k_block_sparse_index
+    q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1)
+
+    # All q blocks attend to exactly k kv blocks
+    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device)
+    
+    # Create the corresponding mask
+    block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
+    
+    # Fill in the mask based on the indices
+    for b in range(bs):
+        for head in range(h):
+            for q_idx in range(num_q_blocks):
+                kv_indices = q2k_block_sparse_index[b, head, q_idx]
+                block_sparse_mask[b, head, q_idx, kv_indices] = True
+    
+    # Create the reverse mapping (k2q)
+    # First, initialize lists to collect q indices for each kv block
+    k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)]
+    
+    # Populate the lists based on q2k mapping
+    for b in range(bs):
+        for head in range(h):
+            flat_idx = b * h + head
+            for q_idx in range(num_q_blocks):
+                kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist()
+                for kv_idx in kv_indices:
+                    k2q_indices_list[flat_idx][kv_idx].append(q_idx)
+    
+    # Find the maximum number of q blocks that attend to any kv block
+    max_q_per_kv = 0
+    for flat_idx in range(bs * h):
+        for kv_idx in range(num_kv_blocks):
+            max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx]))
+    
+    # Create tensors for k2q mapping
+    k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, 
+                                        dtype=torch.int32, device=device)
+    k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), 
+                                       dtype=torch.int32, device=device)
+    
+    # Fill the tensors
+    for b in range(bs):
+        for head in range(h):
+            flat_idx = b * h + head
+            for kv_idx in range(num_kv_blocks):
+                q_indices = k2q_indices_list[flat_idx][kv_idx]
+                num_q = len(q_indices)
+                k2q_block_sparse_num[b, head, kv_idx] = num_q
+                if num_q > 0:
+                    k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor(
+                        q_indices, dtype=torch.int32, device=device)
+                
+    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
+
+def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops):
+    """Benchmark block sparse attention forward and backward passes."""
+    print("\n=== BLOCK SPARSE ATTENTION BENCHMARK ===")
+    
+    # Forward pass
+    # Warm-up run
+    o, l_vec = block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
+    torch.cuda.synchronize()
+    
+    # Benchmark forward
+    _, fwd_time = benchmark_forward(
+        block_sparse_attention_fwd,
+        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num,
+        repeats=20,
+        verbose=False,
+        desc='Block Sparse Forward'
+    )
+    
+    sparse_tflops = flops / fwd_time.mean * 1e-12
+    print(f"Block Sparse Forward - TFLOPS: {sparse_tflops:.2f}")
+    
+    # Backward pass
+    grad_output = torch.randn_like(o)
+    
+    # Warm-up runs
+    for _ in range(5):
+        block_sparse_attention_backward(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
+    torch.cuda.synchronize()
+    
+    # Benchmark backward
+    _, bwd_time = benchmark_forward(
+        block_sparse_attention_backward,
+        q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num,
+        repeats=20,
+        verbose=False,
+        desc='Block Sparse Backward'
+    )
+    bwd_flops = 2.5 * flops  # Approximation
+
+    sparse_bwd_tflops = bwd_flops / bwd_time.mean * 1e-12
+    print(f"Block Sparse Backward - TFLOPS: {sparse_bwd_tflops:.2f}")
+    
+    return sparse_tflops, sparse_bwd_tflops
+
+def main():
+    args = parse_arguments()
+
+    set_seed(42)
+    
+    # Extract parameters
+    batch = args.batch_size
+    head = args.num_heads
+    headdim = args.head_dim
+    
+    print(f"Block Sparse Attention Benchmark")
+    print(f"batch: {batch}, head: {head}, headdim: {headdim}")
+    
+    # Test with different sequence lengths
+    for seq_len in args.seq_lengths:
+        # Skip very long sequences if they might cause OOM
+        if seq_len > 16384 and batch > 1:
+            continue
+            
+        print("="*100)
+        print(f"\nSequence length: {seq_len}")
+        
+        # Calculate theoretical FLOPs for attention
+        flops = 4 * batch * head * headdim * seq_len * seq_len
+        
+        # Create input tensors
+        q, k, v = create_input_tensors(batch, head, seq_len, headdim)
+        
+        # Setup block sparse parameters
+        num_q_blocks = seq_len // BLOCK_M
+        num_kv_blocks = seq_len // BLOCK_N
+        
+        # Determine k value (number of kv blocks per q block)
+        topk = args.topk
+        if topk is None:
+            topk = num_kv_blocks // 10  # Default to ~90% sparsity if k is not specified
+        topk = max(1, topk)       
+        print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)")
+        
+        # Generate block sparse pattern
+        q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, _ = generate_block_sparse_pattern(
+            batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda")
+        
+        # Benchmark block sparse attention
+        sparse_fwd, sparse_bwd = benchmark_block_sparse_attention(
+            q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops
+        )
+        
+        # Print results
+        print("\n=== PERFORMANCE RESULTS ===")
+        print(f"Block Sparse Forward - TFLOPS: {sparse_fwd:.2f}")
+        print(f"Block Sparse Backward - TFLOPS: {sparse_bwd:.2f}")
+
+if __name__ == "__main__":
+    main()

csrc/attn/config_sta.py

@@ -1,6 +1,6 @@
 ### ADD TO THIS TO REGISTER NEW KERNELS
 sources = {
-    'attn': {
+    'st_attn': {
         'source_files': {
             'h100': 'st_attn/st_attn_h100.cu'  # define these source files for each GPU target desired.
         }
@@ -9,7 +9,7 @@ sources = {
 
 ### WHICH KERNELS DO WE WANT TO BUILD?
 # (oftentimes during development work you don't need to redefine them all.)
-kernels = ['attn']
+kernels = ['st_attn']
 
 ### WHICH GPU TARGET DO WE WANT TO BUILD FOR?
 target = 'h100'

csrc/attn/config_vsa.py

@@ -0,0 +1,15 @@
+### ADD TO THIS TO REGISTER NEW KERNELS
+sources = {
+    'block_sparse': {
+        'source_files': {
+            'h100': 'vsa/block_sparse_h100.cu'
+        }
+    }
+}
+
+### WHICH KERNELS DO WE WANT TO BUILD?
+# (oftentimes during development work you don't need to redefine them all.)
+kernels = ['block_sparse']
+
+### WHICH GPU TARGET DO WE WANT TO BUILD FOR?
+target = 'h100'

csrc/attn/setup_sta.py

@@ -1,7 +1,7 @@
 import os
 import subprocess
 
-from config import kernels, sources, target
+from csrc.attn.config_sta import kernels, sources, target
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
@@ -9,7 +9,7 @@ target = target.lower()
 
 # Package metadata
 PACKAGE_NAME = "st_attn"
-VERSION = "0.0.2"
+VERSION = "0.0.4"
 AUTHOR = "Hao AI Lab"
 DESCRIPTION = "Sliding Tile Atteniton Kernel Used in FastVideo"
 URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/sliding_tile_attention"

csrc/attn/setup_vsa.py

@@ -0,0 +1,76 @@
+import os
+import subprocess
+
+from csrc.attn.config_vsa import kernels, sources, target
+from setuptools import find_packages, setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+target = target.lower()
+
+# Package metadata
+PACKAGE_NAME = "vsa"
+VERSION = "0.0.1"
+AUTHOR = "Hao AI Lab"
+DESCRIPTION = "Video Sparse Attention Kernel Used in FastVideo"
+URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn"
+
+# Set environment variables
+tk_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), 'tk/')))
+python_include = subprocess.check_output(['python', '-c',
+                                          "import sysconfig; print(sysconfig.get_path('include'))"]).decode().strip()
+torch_include = subprocess.check_output([
+    'python', '-c',
+    "import torch; from torch.utils.cpp_extension import include_paths; print(' '.join(['-I' + p for p in include_paths()]))"
+]).decode().strip()
+print('vsa root:', tk_root)
+print('Python include:', python_include)
+print('Torch include directories:', torch_include)
+
+# CUDA flags
+cuda_flags = [
+    '-DNDEBUG', '-Xcompiler=-Wno-psabi', '-Xcompiler=-fno-strict-aliasing', '--expt-extended-lambda',
+    '--expt-relaxed-constexpr', '-forward-unknown-to-host-compiler', '--use_fast_math', '-std=c++20', '-O3',
+    '-Xnvlink=--verbose', '-Xptxas=--verbose', '-Xptxas=--warn-on-spills', f'-I{tk_root}/include',
+    f'-I{tk_root}/prototype', f'-I{python_include}', '-DTORCH_COMPILE'
+] + torch_include.split()
+cpp_flags = ['-std=c++20', '-O3']
+
+if target == 'h100':
+    cuda_flags.append('-DKITTENS_HOPPER')
+    cuda_flags.append('-arch=sm_90a')
+else:
+    raise ValueError(f'Target {target} not supported')
+
+source_files = ['vsa.cpp']
+for k in kernels:
+    if target not in sources[k]['source_files']:
+        raise KeyError(f'Target {target} not found in source files for kernel {k}')
+    if isinstance(sources[k]['source_files'][target], list):
+        source_files.extend(sources[k]['source_files'][target])
+    else:
+        source_files.append(sources[k]['source_files'][target])
+    cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}')
+
+setup(name=PACKAGE_NAME,
+      version=VERSION,
+      author=AUTHOR,
+      description=DESCRIPTION,
+      url=URL,
+      packages=find_packages(),
+      ext_modules=[
+          CUDAExtension('vsa_cuda',
+                        sources=source_files,
+                        extra_compile_args={
+                            'cxx': cpp_flags,
+                            'nvcc': cuda_flags
+                        },
+                        libraries=['cuda'])
+      ],
+      cmdclass={'build_ext': BuildExtension},
+      classifiers=[
+          "Programming Language :: Python :: 3",
+          "Environment :: GPU :: NVIDIA CUDA :: 12",
+          "License :: OSI Approved :: Apache Software License",
+      ],
+      python_requires='>=3.10',
+      install_requires=["torch>=2.5.0"])

csrc/attn/st_attn.cpp

@@ -7,18 +7,17 @@
 
 #include <cuda_runtime.h>
 
-
-#ifdef TK_COMPILE_ATTN
+#ifdef TK_COMPILE_ST_ATTN
 extern torch::Tensor sta_forward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text, bool has_text
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag
 ); 
 #endif
 
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.doc() = "Sliding Block Attention Kernels"; // optional module docstring
-
-#ifdef TK_COMPILE_ATTN
+    
+#ifdef TK_COMPILE_ST_ATTN
     m.def("sta_fwd",  torch::wrap_pybind_function(sta_forward), "sliding tile attention, assuming tile size is (6,8,8)");
 #endif
-}
+}

csrc/attn/st_attn/__init__.py

@@ -1,14 +1,22 @@
 import math
 
 import torch
-from st_attn_cuda import sta_fwd
+from torch.utils.checkpoint import detach_variable
+try:
+    from st_attn_cuda import sta_fwd
+except ImportError:
+    sta_fwd = None
 
-
-def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True):
+def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, dit_seq_shape='30x48x80'):
     seq_length = q_all.shape[2]
+    dit_seq_shape_mapping = {
+        '30x48x80':1,
+        '36x48x48':2,
+        '18x48x80':3, 
+    }
     if has_text:
         assert q_all.shape[
-            2] == 115456, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels"
+            2] >= 115200 and q_all.shape[2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '30x48x80' for HunyuanVideo"
         assert q_all.shape[1] == len(window_size), "Number of heads must match the number of window sizes"
         target_size = math.ceil(seq_length / 384) * 384
         pad_size = target_size - seq_length
@@ -17,8 +25,14 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te
             k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2)
             v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2)
     else:
-        assert q_all.shape[2] == 82944
+        if dit_seq_shape == '36x48x48': # Stepvideo 204x768x68
+            assert q_all.shape[2] == 82944
+        elif dit_seq_shape == '18x48x80': # Wan 69x768x1280
+            assert q_all.shape[2] == 69120
+        else:
+            raise ValueError(f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
 
+    kernel_aspect_ratio_flag = dit_seq_shape_mapping[dit_seq_shape]
     hidden_states = torch.empty_like(q_all)
     # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
     for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
@@ -29,7 +43,7 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te
                                                                                       head_index:head_index + 1],
                                               hidden_states[batch:batch + 1, head_index:head_index + 1])
 
-            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text)
+            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text, kernel_aspect_ratio_flag)
     if has_text:
-        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True)
-    return hidden_states[:, :, :seq_length]
+        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True, kernel_aspect_ratio_flag)
+    return hidden_states[:, :, :seq_length]

csrc/attn/st_attn/st_attn_h100.cu

@@ -4,9 +4,17 @@
 #include <cooperative_groups.h>
 #include <iostream>
 #include <stdio.h>
+#include <c10/cuda/CUDAGuard.h>
+
+// #define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value)))
+__device__ __forceinline__ int clamp_int(int value, int min, int max) {
+    return (value < min) ? min : ((value > max) ? max : value);
+}
+// #define ABS(x) ((x) < 0 ? -(x) : (x))
+__device__ __forceinline__ int abs_int(int value) {
+    return (value < 0) ? -value : value;
+}
 
-#define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value)))
-#define ABS(x) ((x) < 0 ? -(x) : (x))
 
 constexpr int CONSUMER_WARPGROUPS = (3); 
 constexpr int PRODUCER_WARPGROUPS = (1); 
@@ -117,16 +125,16 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
             int qt = seq_idx / 6 / (CH * CW);
             int qh = (seq_idx / 6) % (CH * CW) / CW;
             int qw = (seq_idx / 6) % CW;
-            qt = CLAMP(qt, DT, CT-DT-1);
-            qh = CLAMP(qh, DH, CH-DH-1);
-            qw = CLAMP(qw, DW, CW-DW-1);
+            qt = clamp_int(qt, DT, CT-DT-1);
+            qh = clamp_int(qh, DH, CH-DH-1);
+            qw = clamp_int(qw, DW, CW-DW-1);
             int count = 0;
             int j = 0;
             while (count < K::stages - 1) {
                 int kt = j / 3 / (CH * CW);
                 int kh = (j / 3) % (CH * CW) / CW;
                 int kw = (j / 3) % CW;
-                bool mask = (ABS(qt - kt) <= DT) && (ABS(qh - kh) <= DH) && (ABS(qw - kw) <= DW);
+                bool mask = (abs_int(qt - kt) <= DT) && (abs_int(qh - kh) <= DH) && (abs_int(qw - kw) <= DW);
                 if (mask){
                     coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0};
                     tma::expect_bytes(k_smem_arrived[count], sizeof(k_tile));
@@ -167,15 +175,15 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
                 int qt = seq_idx / 6 / (CH * CW);
                 int qh = (seq_idx / 6) % (CH * CW) / CW;
                 int qw = (seq_idx / 6) % CW;
-                qt = CLAMP(qt, DT, CT-DT-1);
-                qh = CLAMP(qh, DH, CH-DH-1);
-                qw = CLAMP(qw, DW, CW-DW-1);
-                int k_t_min = CLAMP(qt-DT, 0, CT-1);
-                int k_t_max = CLAMP(qt+DT, 0, CT-1);
-                int k_h_min = CLAMP(qh-DH, 0, CH-1);
-                int k_h_max = CLAMP(qh+DH, 0, CH-1);
-                int k_w_min = CLAMP(qw-DW, 0, CW-1);
-                int k_w_max = CLAMP(qw+DW, 0, CW-1);
+                qt = clamp_int(qt, DT, CT-DT-1);
+                qh = clamp_int(qh, DH, CH-DH-1);
+                qw = clamp_int(qw, DW, CW-DW-1);
+                int k_t_min = clamp_int(qt-DT, 0, CT-1);
+                int k_t_max = clamp_int(qt+DT, 0, CT-1);
+                int k_h_min = clamp_int(qh-DH, 0, CH-1);
+                int k_h_max = clamp_int(qh+DH, 0, CH-1);
+                int k_w_min = clamp_int(qw-DW, 0, CW-1);
+                int k_w_max = clamp_int(qw+DW, 0, CW-1);
                 int count = 0;
                 for (int kt = k_t_min; kt <= k_t_max; kt++) {
                     for (int kh = k_h_min; kh <= k_h_max; kh++) {
@@ -234,7 +242,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
             // the last three kv blocks are for text, we process them separately
             kv_iters = img_kv_blocks - 1;
         } else {
-            kv_iters = CLAMP(DT*2+1, 1, CT) * CLAMP(DH*2+1, 1, CH) * CLAMP(DW*2+1, 1, CW) * 3 - 1 ; 
+            kv_iters = clamp_int(DT*2+1, 1, CT) * clamp_int(DH*2+1, 1, CH) * clamp_int(DW*2+1, 1, CW) * 3 - 1 ; 
         }
 
         kittens::wait(qsmem_semaphore, 0);
@@ -359,7 +367,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
 #include <iostream>
 
 torch::Tensor 
-sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text, bool has_text)
+sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag)
 {
     CHECK_INPUT(q);
     CHECK_INPUT(k);
@@ -415,8 +423,9 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
     float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
     float* d_l   = reinterpret_cast<float*>(l_ptr);
 
-    cudaDeviceSynchronize();
-    auto stream = at::cuda::getCurrentCUDAStream().stream(); 
+    //cudadevicesynchronize();
+    const c10::cuda::OptionalCUDAGuard device_guard(q.device());
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
 
 
     if (head_dim == 128) {
@@ -442,11 +451,11 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
 
         globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len),  static_cast<int>(text_length), static_cast<int>(hr)};
 
-        auto mem_size = kittens::MAX_SHARED_MEMORY;
-        auto threads  = NUM_WORKERS * kittens::WARP_THREADS;
+        constexpr int mem_size = kittens::MAX_SHARED_MEMORY;
+        int threads = NUM_WORKERS * kittens::WARP_THREADS;
         if (has_text) {
             // TORCH_CHECK(seq_len % (CONSUMER_WARPGROUPS*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 192");
-            dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4-2), qo_heads, batch);
+            dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4)-2, qo_heads, batch);
             dim3 grid_text(2, qo_heads, batch);
             if (!process_text) {
                 if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
@@ -558,130 +567,275 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
 
         } else {
             dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
+            if (kernel_aspect_ratio_flag == 2){
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) {
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) {
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }  else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }  else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else {
+                    // print error
+                    std::cout << "Invalid kernel size" << std::endl;
+                    //print kernel size
+                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                }
+            }
+            else if (kernel_aspect_ratio_flag == 3) {
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
 
-            } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            }  else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            }  else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else {
-                // print error
-                std::cout << "Invalid kernel size" << std::endl;
-                //print kernel size
-                std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 2, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 1, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 9){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 9){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 0, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 0, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 3 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 1, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,0, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 2, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,0, 3, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else {
+                    // print error
+                    std::cout << "Invalid kernel size" << std::endl;
+                    //print kernel size
+                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                }
+            }
+
+            else {
+                std::cout << "Unsupported kernel_aspect_ratio_flag: " << kernel_aspect_ratio_flag << std::endl;
             }
 
         }
         CHECK_CUDA_ERROR(cudaGetLastError());
-        cudaStreamSynchronize(stream);
+        // cudaStreamSynchronize(stream);
     }
 
     return o;
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
 }
+

csrc/attn/tests/test_block_sparse.py

@@ -0,0 +1,289 @@
+import torch
+import argparse
+from flash_attn.utils.benchmark import benchmark_forward
+from flash_attn import flash_attn_func
+from vsa import block_sparse_attention_fwd, block_sparse_attention_backward, BlockSparseAttentionFunction
+from vsa import BLOCK_M, BLOCK_N
+
+import numpy as np
+import random
+import gc
+
+def set_seed(seed: int = 42):
+    # Python random module
+    random.seed(seed)
+
+    # NumPy
+    np.random.seed(seed)
+
+    # PyTorch
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # if using multi-GPU
+
+
+@torch.no_grad
+def precision_metric(quant_o, fa2_o): 
+    x, xx = quant_o.float(), fa2_o.float() 
+    sim = torch.nn.functional.cosine_similarity(x.reshape(1, -1), xx.reshape(1, -1)).item()
+    l1 =   ((x - xx).abs().sum() / xx.abs().sum() ).item()
+    rmse = torch.sqrt(torch.mean((x -xx) ** 2)).item()
+
+    return sim, l1, rmse
+
+def create_input_tensors(batch, head, seq_len, headdim):
+    """Create random input tensors for attention."""
+    q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
+    return q, k, v
+
+def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"):
+    """
+    Generate a block sparse pattern where each q block attends to exactly k kv blocks.
+    
+    Args:
+        bs: batch size
+        h: number of heads
+        num_q_blocks: number of query blocks
+        num_kv_blocks: number of key-value blocks
+        k: number of kv blocks each q block attends to
+        device: device to create tensors on
+        
+    Returns:
+        q2k_block_sparse_index: [bs, h, num_q_blocks, k]
+            Contains the indices of kv blocks that each q block attends to.
+        q2k_block_sparse_num: [bs, h, num_q_blocks]
+            Contains the number of kv blocks that each q block attends to (all equal to k).
+        k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks]
+            Contains the indices of q blocks that attend to each kv block.
+        k2q_block_sparse_num: [bs, h, num_kv_blocks]
+            Contains the number of q blocks that attend to each kv block.
+        block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]
+            Binary mask where 1 indicates attention connection.
+    """
+    # Ensure k is not larger than num_kv_blocks
+    k = min(k, num_kv_blocks)
+    
+    # Create random scores for sampling
+    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device)
+    
+    # Get top-k indices for each q block
+    _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1)
+    q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32)
+
+    # sort q2k_block_sparse_index
+    q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1)
+
+    # All q blocks attend to exactly k kv blocks
+    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device)
+    
+    # Create the corresponding mask
+    block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
+    
+    # Fill in the mask based on the indices
+    for b in range(bs):
+        for head in range(h):
+            for q_idx in range(num_q_blocks):
+                kv_indices = q2k_block_sparse_index[b, head, q_idx]
+                block_sparse_mask[b, head, q_idx, kv_indices] = True
+    
+    # Create the reverse mapping (k2q)
+    # First, initialize lists to collect q indices for each kv block
+    k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)]
+    
+    # Populate the lists based on q2k mapping
+    for b in range(bs):
+        for head in range(h):
+            flat_idx = b * h + head
+            for q_idx in range(num_q_blocks):
+                kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist()
+                for kv_idx in kv_indices:
+                    k2q_indices_list[flat_idx][kv_idx].append(q_idx)
+    
+    # Find the maximum number of q blocks that attend to any kv block
+    max_q_per_kv = 0
+    for flat_idx in range(bs * h):
+        for kv_idx in range(num_kv_blocks):
+            max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx]))
+    
+    # Create tensors for k2q mapping
+    k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, 
+                                        dtype=torch.int32, device=device)
+    k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), 
+                                       dtype=torch.int32, device=device)
+    
+    # Fill the tensors
+    for b in range(bs):
+        for head in range(h):
+            flat_idx = b * h + head
+            for kv_idx in range(num_kv_blocks):
+                q_indices = k2q_indices_list[flat_idx][kv_idx]
+                num_q = len(q_indices)
+                k2q_block_sparse_num[b, head, kv_idx] = num_q
+                if num_q > 0:
+                    k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor(
+                        q_indices, dtype=torch.int32, device=device)
+                
+    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
+
+def main(args):
+    set_seed(42)
+    
+    # Extract parameters
+    batch = args.batch_size
+    head = args.num_heads
+    headdim = args.head_dim
+    num_iterations = args.num_iterations
+    
+    print(f"Block Sparse Attention Benchmark")
+    print(f"batch: {batch}, head: {head}, headdim: {headdim}, iterations: {num_iterations}")
+    
+    # Test with different sequence lengths
+    for seq_len in args.seq_lengths:
+        # Skip very long sequences if they might cause OOM
+        # if seq_len > 16384 and batch > 1:
+        #     continue
+            
+        print("="*100)
+        print(f"\nSequence length: {seq_len}")
+        
+        # Collect metrics across iterations
+        forward_metrics = {'sim': [], 'l1': [], 'rmse': []}
+        grad_q_metrics = {'sim': [], 'l1': [], 'rmse': []}
+        grad_k_metrics = {'sim': [], 'l1': [], 'rmse': []}
+        grad_v_metrics = {'sim': [], 'l1': [], 'rmse': []}
+        
+        for iter_idx in range(num_iterations):
+            if num_iterations > 1:
+                print(f"\nIteration {iter_idx+1}/{num_iterations}")
+            
+            # Create input tensors
+            q, k, v = create_input_tensors(batch, head, seq_len, headdim)
+            
+            # Setup block sparse parameters
+            num_q_blocks = seq_len // BLOCK_M
+            num_kv_blocks = seq_len // BLOCK_N
+            
+            # Determine k value (number of kv blocks per q block)
+            topk = args.topk
+            if topk is None:
+                topk = num_kv_blocks // 10  # Default to ~90% sparsity if k is not specified
+            topk = max(1, topk)
+            if iter_idx == 0:  # Only print this once
+                print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)")
+
+            # Generate block sparse pattern
+            q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask = generate_block_sparse_pattern(
+                batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda")
+
+            # expand block_sparse_mask to full mask
+            block_mask_expanded = block_sparse_mask.unsqueeze(-1).unsqueeze(-2)  # [b, h, num_q_blocks, num_kv_blocks, 1, 1]
+            block_mask_expanded = block_mask_expanded.expand(-1, -1, -1, -1, BLOCK_M, BLOCK_N)  # [b, h, num_q_blocks, num_kv_blocks, BLOCK_M, BLOCK_N]
+            full_mask = block_mask_expanded.permute(0, 1, 2, 4, 3, 5).reshape(batch, head, seq_len, seq_len)
+
+            q.requires_grad = True
+            k.requires_grad = True
+            v.requires_grad = True
+
+
+            # testing forward
+            o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
+            del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded
+            grad_o = torch.randn_like(o)
+            o.backward(grad_o)
+            # clear memory
+            q_sdpa = q.detach().clone()
+            k_sdpa = k.detach().clone()
+            v_sdpa = v.detach().clone()
+            q_sdpa.requires_grad = True
+            k_sdpa.requires_grad = True
+            v_sdpa.requires_grad = True
+            q.data = torch.empty(0, device=q.device)
+            k.data = torch.empty(0, device=k.device)
+            v.data = torch.empty(0, device=v.device)
+            torch.cuda.empty_cache()
+
+            o_sdpa = torch.nn.functional.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=full_mask)
+
+
+            sim, l1, rmse = precision_metric(o, o_sdpa)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 8e-5, f"l1 too large: {l1}"
+            assert rmse < 2e-5, f"RMSE too large: {rmse}"
+            forward_metrics['sim'].append(sim)
+            forward_metrics['l1'].append(l1)
+            forward_metrics['rmse'].append(rmse)
+
+            print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+
+            # test backward
+            o_sdpa.backward(grad_o)
+
+            sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad)
+            # Error bounds collected on H100
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 4e-3, f"l1 too large: {l1}"
+            assert rmse < 3e-4, f"RMSE too large: {rmse}"
+            grad_q_metrics['sim'].append(sim)
+            grad_q_metrics['l1'].append(l1)
+            grad_q_metrics['rmse'].append(rmse)
+            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}")        
+            
+            sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 4e-3, f"l1 too large: {l1}"
+            assert rmse < 2e-4, f"RMSE too large: {rmse}"
+            grad_k_metrics['sim'].append(sim)
+            grad_k_metrics['l1'].append(l1)
+            grad_k_metrics['rmse'].append(rmse)
+            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            
+            sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 1e-4, f"l1 too large: {l1}"
+            assert rmse < 2e-5, f"RMSE too large: {rmse}"
+            grad_v_metrics['sim'].append(sim)
+            grad_v_metrics['l1'].append(l1)
+            grad_v_metrics['rmse'].append(rmse)
+            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            
+            del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        # Print summary statistics if multiple iterations were run
+        if num_iterations > 1:
+            print("\n" + "="*50)
+            print(f"Summary Statistics (over {num_iterations} iterations):")
+            
+            print("\nForward metrics:")
+            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}")
+            
+            print("\nGradient Q metrics:")
+            print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}, min={np.min(grad_q_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}, max={np.max(grad_q_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}, max={np.max(grad_q_metrics['rmse']):.6f}")
+            
+            print("\nGradient K metrics:")
+            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}, min={np.min(grad_k_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}, max={np.max(grad_k_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}, max={np.max(grad_k_metrics['rmse']):.6f}")
+            
+            print("\nGradient V metrics:")
+            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}, min={np.min(grad_v_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}, max={np.max(grad_v_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}, max={np.max(grad_v_metrics['rmse']):.6f}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
+    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
+    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
+    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
+    parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to')
+    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark')
+    parser.add_argument('--num_iterations', type=int, default=50, help='Number of test iterations to run')
+    args = parser.parse_args()
+    main(args)

csrc/attn/tests/test_fa2.py

@@ -0,0 +1,136 @@
+import torch
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import numpy as np
+
+def pytorch_test(Q, K, V, dO):
+    q_ = Q.to(torch.float64).requires_grad_()
+    k_ = K.to(torch.float64).requires_grad_()
+    v_ = V.to(torch.float64).requires_grad_()
+    dO_ = dO.to(torch.float64)
+    
+    # manual pytorch implementation of scaled dot product attention
+    QK = torch.matmul(q_, k_.transpose(-2, -1))
+    QK /= (q_.size(-1) ** 0.5)
+    
+    # Causal mask removed since causal is always false
+    
+    QK = torch.nn.functional.softmax(QK, dim=-1)
+    output = torch.matmul(QK, v_)
+    
+    output.backward(dO_)
+    
+    q_grad = q_.grad
+    k_grad = k_.grad
+    v_grad = v_.grad
+    
+    return output, q_grad, k_grad, v_grad
+
+def fa2_test(Q, K, V, dO):
+    Q.requires_grad = True
+    K.requires_grad = True
+    V.requires_grad = True
+    output = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=False)
+    output.backward(dO)
+    
+    return output, Q.grad, K.grad, V.grad
+
+def generate_tensor(shape, mean, std, dtype, device):
+    tensor = torch.randn(shape, dtype=dtype, device=device)
+
+    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
+    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
+    
+    return scaled_tensor.contiguous()
+
+def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all', test_mode='forward_backward'):
+    results = {
+        'FA2 vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0},
+    }
+
+    for _ in range(num_iterations):
+        torch.manual_seed(0)
+        
+        Q  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        K  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        V  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        dO = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        
+        pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, dO)
+        fa2_o, fa2_qg, fa2_kg, fa2_vg = fa2_test(Q, K, V, dO)
+        
+        if test_mode == 'forward_only':
+            tensors_fa2_pt = [(pt_o, fa2_o)]
+        else:  # 'forward_backward'
+            if error_mode == 'output':
+                tensors_fa2_pt = [(pt_o, fa2_o)]
+            elif error_mode == 'backward':
+                tensors_fa2_pt = [(pt_qg, fa2_qg),
+                               (pt_kg, fa2_kg),
+                               (pt_vg, fa2_vg)]
+            else:  # 'all'
+                tensors_fa2_pt = [(pt_o, fa2_o),
+                               (pt_qg, fa2_qg),
+                               (pt_kg, fa2_kg),
+                               (pt_vg, fa2_vg)]
+                
+        for pt, fa2 in tensors_fa2_pt:
+            diff = pt - fa2
+            abs_diff = torch.abs(diff)
+            results['FA2 vs PT']['sum_diff'] += torch.sum(abs_diff).item()
+            results['FA2 vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
+            results['FA2 vs PT']['max_diff'] = max(results['FA2 vs PT']['max_diff'], torch.max(abs_diff).item())
+                
+        torch.cuda.empty_cache()
+
+    # Calculate total elements based on test mode and error mode
+    if test_mode == 'forward_only':
+        total_elements = b * h * n * d * num_iterations
+    else:  # 'forward_backward'
+        total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4)
+    
+    for name, data in results.items():
+        avg_diff = data['sum_diff'] / total_elements
+        max_diff = data['max_diff']
+        results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff}
+
+    return results
+
+def generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward'):
+    seq_lengths = [768 * (2**i) for i in range(1)]
+
+    print(f"\n{'='*80}")
+    print(f"ATTENTION ERROR COMPARISON TABLE (b={b}, h={h}, d={d}, mean={mean}, std={std})")
+    print(f"Mode: {error_mode}, Test: {test_mode}")
+    print(f"{'='*80}")
+    
+    # Print header
+    print(f"{'Seq Length':<12} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15}")
+    print(f"{'-'*12} | {'-'*15} | {'-'*15}")
+    
+    for n in seq_lengths:
+        results = check_correctness(b, h, n, d, mean, std, error_mode=error_mode, test_mode=test_mode)
+        
+        fa2_pt_avg = results['FA2 vs PT']['avg_diff']
+        fa2_pt_max = results['FA2 vs PT']['max_diff']
+        
+        # Print row
+        print(f"{n:<12} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e}")
+    
+    print(f"{'='*80}\n")
+
+# fix random seed
+torch.manual_seed(0)
+
+# Example usage
+b, h, d = 2, 2, 64
+mean = 1e-1
+std = 10
+
+# Test forward only
+generate_error_tables(b, h, d, mean, std, error_mode='output', test_mode='forward_only')
+
+# Test forward and backward
+generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward')
+
+print("Attention error comparison completed.")

csrc/attn/tests/test_full_attn.py

@@ -0,0 +1,175 @@
+import torch
+from flash_attn_interface import flash_attn_func
+from st_attn import mha_forward, mha_backward
+import random
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import numpy as np
+
+def pytorch_test(Q, K, V, dO):
+    q_ = Q.to(torch.float64).requires_grad_()
+    k_ = K.to(torch.float64).requires_grad_()
+    v_ = V.to(torch.float64).requires_grad_()
+    dO_ = dO.to(torch.float64)
+    
+    # manual pytorch implementation of scaled dot product attention
+    QK = torch.matmul(q_, k_.transpose(-2, -1))
+    QK /= (q_.size(-1) ** 0.5)
+    
+    # Causal mask removed since causal is always false
+    
+    QK = torch.nn.functional.softmax(QK, dim=-1)
+    output = torch.matmul(QK, v_)
+    
+    output.backward(dO_)
+    
+    q_grad = q_.grad
+    k_grad = k_.grad
+    v_grad = v_.grad
+    
+    return output, q_grad, k_grad, v_grad
+
+def fa2_test(Q, K, V, dO):
+    Q.requires_grad = True
+    K.requires_grad = True
+    V.requires_grad = True
+    output = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=False)
+    output.backward(dO)
+    
+    return output, Q.grad, K.grad, V.grad
+
+
+def mha_kernel_test(Q, K, V, dO, mode): 
+    Q.requires_grad = True
+    K.requires_grad = True
+    V.requires_grad = True
+    
+    o, l_vec = mha_forward(Q, K, V)
+    
+    if mode == 'forward_only':
+        return o, None, None, None
+    else:  # 'forward_backward'
+        qg, kg, vg = mha_backward(Q, K, V, o, l_vec, dO)
+        return o, qg, kg, vg
+
+def generate_tensor(shape, mean, std, dtype, device):
+    tensor = torch.randn(shape, dtype=dtype, device=device)
+
+    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
+    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
+    
+    return scaled_tensor.contiguous()
+
+def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all', test_mode='forward_backward'):
+    results = {
+        'MHA vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0},
+        'FA2 vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0},
+    }
+
+    for _ in range(num_iterations):
+        torch.manual_seed(0)
+        
+        Q  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        K  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        V  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        dO = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        
+        pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, dO)
+        fa2_o, fa2_qg, fa2_kg, fa2_vg = fa2_test(Q, K, V, dO)
+        
+        if test_mode == 'forward_only':
+            mha_o, _, _, _ = mha_kernel_test(Q, K, V, dO, 'forward_only')
+            tensors_mha_pt = [(pt_o, mha_o)]
+            tensors_fa2_pt = [(pt_o, fa2_o)]
+        else:  # 'forward_backward'
+            mha_o, mha_qg, mha_kg, mha_vg = mha_kernel_test(Q, K, V, dO, 'forward_backward')
+            
+            if error_mode == 'output':
+                tensors_mha_pt = [(pt_o, mha_o)]
+                tensors_fa2_pt = [(pt_o, fa2_o)]
+            elif error_mode == 'backward':
+                tensors_mha_pt = [(pt_qg, mha_qg),
+                               (pt_kg, mha_kg),
+                               (pt_vg, mha_vg)]
+                tensors_fa2_pt = [(pt_qg, fa2_qg),
+                               (pt_kg, fa2_kg),
+                               (pt_vg, fa2_vg)]
+            else:  # 'all'
+                tensors_mha_pt = [(pt_o, mha_o),
+                               (pt_qg, mha_qg),
+                               (pt_kg, mha_kg),
+                               (pt_vg, mha_vg)]
+                tensors_fa2_pt = [(pt_o, fa2_o),
+                               (pt_qg, fa2_qg),
+                               (pt_kg, fa2_kg),
+                               (pt_vg, fa2_vg)]
+                
+        for pt, mha in tensors_mha_pt:
+            diff = pt - mha
+            abs_diff = torch.abs(diff)
+            results['MHA vs PT']['sum_diff'] += torch.sum(abs_diff).item()
+            results['MHA vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
+            results['MHA vs PT']['max_diff'] = max(results['MHA vs PT']['max_diff'], torch.max(abs_diff).item())
+                
+        for pt, fa2 in tensors_fa2_pt:
+            diff = pt - fa2
+            abs_diff = torch.abs(diff)
+            results['FA2 vs PT']['sum_diff'] += torch.sum(abs_diff).item()
+            results['FA2 vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
+            results['FA2 vs PT']['max_diff'] = max(results['FA2 vs PT']['max_diff'], torch.max(abs_diff).item())
+                
+        torch.cuda.empty_cache()
+
+    # Calculate total elements based on test mode and error mode
+    if test_mode == 'forward_only':
+        total_elements = b * h * n * d * num_iterations
+    else:  # 'forward_backward'
+        total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4)
+    
+    for name, data in results.items():
+        avg_diff = data['sum_diff'] / total_elements
+        max_diff = data['max_diff']
+        results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff}
+
+    return results
+
+def generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward'):
+    seq_lengths = [768 * (2**i) for i in range(1)]
+
+    print(f"\n{'='*80}")
+    print(f"MHA ERROR COMPARISON TABLE (b={b}, h={h}, d={d}, mean={mean}, std={std})")
+    print(f"Mode: {error_mode}, Test: {test_mode}")
+    print(f"{'='*80}")
+    
+    # Print header
+    print(f"{'Seq Length':<12} | {'MHA vs PT Avg':<15} | {'MHA vs PT Max':<15} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15}")
+    print(f"{'-'*12} | {'-'*15} | {'-'*15} | {'-'*15} | {'-'*15}")
+    
+    for n in seq_lengths:
+        results = check_correctness(b, h, n, d, mean, std, error_mode=error_mode, test_mode=test_mode)
+        
+        mha_pt_avg = results['MHA vs PT']['avg_diff']
+        mha_pt_max = results['MHA vs PT']['max_diff']
+        fa2_pt_avg = results['FA2 vs PT']['avg_diff']
+        fa2_pt_max = results['FA2 vs PT']['max_diff']
+        
+        # Print row
+        print(f"{n:<12} | {mha_pt_avg:<15.6e} | {mha_pt_max:<15.6e} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e}")
+    
+    print(f"{'='*80}\n")
+
+# fix random seed
+torch.manual_seed(0)
+
+# Example usage
+b, h, d = 2, 2, 64
+mean = 1e-1
+std = 10
+
+# Test forward only
+generate_error_tables(b, h, d, mean, std, error_mode='output', test_mode='forward_only')
+
+# Test forward and backward
+generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward')
+
+print("MHA attention error comparison completed.")

csrc/attn/tests/test_sta.py

@@ -9,14 +9,14 @@ flex_attention = torch.compile(flex_attention, dynamic=False)
 
 
 def flex_test(Q, K, V, kernel_size):
-    mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (36, 48, 48), 39, 'cuda', 0)
+    mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (18, 48, 80), 0, 'cuda', 0)
     output = flex_attention(Q, K, V, block_mask=mask)
 
     return output
 
 
 def h100_fwd_kernel_test(Q, K, V, kernel_size):
-    o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 39, False)
+    o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 0, False, '18x48x80')
     return o
 
 
@@ -37,7 +37,7 @@ def check_correctness(b, h, n, d, causal, mean, std, num_iterations=50, error_mo
             'max_diff': 0
         },
     }
-    kernel_size_ls = [(6, 1, 6), (6, 6, 1)]
+    kernel_size_ls = [(3, 3, 5), (3, 1, 10)]
     from tqdm import tqdm
     for kernel_size in tqdm(kernel_size_ls):
         for _ in range(num_iterations):
@@ -72,25 +72,16 @@ def check_correctness(b, h, n, d, causal, mean, std, num_iterations=50, error_mo
     return results
 
 
-def generate_error_graphs(b, h, d, causal, mean, std, error_mode='all'):
-    seq_lengths = [82944]
-
-    tk_avg_errors, tk_max_errors = [], []
-
-    for n in tqdm(seq_lengths, desc="Generating error data"):
-        results = check_correctness(b, h, n, d, causal, mean, std, error_mode=error_mode)
-
-        tk_avg_errors.append(results['TK vs FLEX']['avg_diff'])
-        tk_max_errors.append(results['TK vs FLEX']['max_diff'])
-
-
 # Example usage
 b, h, d = 2, 24, 128
+n = 69120  # Sequence length
 causal = False
 mean = 1e-1
 std = 10
 
-for mode in ['output']:
-    generate_error_graphs(b, h, d, causal, mean, std, error_mode=mode)
-
-print("Error graphs generated and saved for all modes.")
+# Run correctness check directly
+results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output')
+assert results['TK vs FLEX']['avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large"
+assert results['TK vs FLEX']['max_diff'] < 4e-2, f"Maximum difference: {results['TK vs FLEX']['max_diff']} is too large"
+print(f"Average difference: {results['TK vs FLEX']['avg_diff']}")
+print(f"Maximum difference: {results['TK vs FLEX']['max_diff']}")

csrc/attn/vsa.cpp

@@ -0,0 +1,27 @@
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+
+#include <vector>
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+
+#include <cuda_runtime.h>
+
+#ifdef TK_COMPILE_BLOCK_SPARSE
+extern std::vector<torch::Tensor> block_sparse_attention_forward(
+    torch::Tensor q, torch::Tensor k, torch::Tensor v,  torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num
+); 
+extern std::vector<torch::Tensor> block_sparse_attention_backward(
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og, torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num
+);
+#endif
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "Video Sparse Attention Kernels"; // optional module docstring
+
+#ifdef TK_COMPILE_BLOCK_SPARSE
+    m.def("block_sparse_fwd",  torch::wrap_pybind_function(block_sparse_attention_forward), "block sparse attention");
+    m.def("block_sparse_bwd",  torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward");
+#endif
+}

csrc/attn/vsa/__init__.py

@@ -0,0 +1,470 @@
+import math
+
+import torch
+from torch.utils.checkpoint import detach_variable
+from typing import Tuple
+try:
+    from vsa_cuda import block_sparse_fwd, block_sparse_bwd
+except ImportError:
+    block_sparse_fwd = None
+    block_sparse_bwd = None
+
+
+BLOCK_M = 64
+BLOCK_N = 64
+
+def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None):
+    """
+    q: [batch_size, num_heads, seq_len, head_dim]
+    k: [batch_size, num_heads, seq_len, head_dim]
+    v: [batch_size, num_heads, seq_len, head_dim]
+    topk: int
+    block_size: int or tuple of 3 ints
+    video_shape: tuple of (T, H, W)
+    compress_attn_weight: [batch_size, num_heads, seq_len, head_dim]
+    select_attn_weight: [batch_size, num_heads, seq_len, head_dim]
+    
+    V1 of sparse attention. Include compress attn and sparse attn branch, use average pooling to compress. 
+    Assume q, k, v is flattened in this way: [batch_size, num_heads, T//block_size[0], H//block_size[1], W//block_size[2], block_size[0], block_size[1], block_size[2]]
+    """
+
+    if isinstance(block_size, int):
+        block_size = (block_size, block_size, block_size)
+
+    block_elements = block_size[0] * block_size[1] * block_size[2]
+    assert block_elements % 64 == 0 and block_elements >= 64
+    assert q.shape[2] % block_elements == 0
+    batch_size, num_heads, seq_len, head_dim = q.shape
+    # compress attn
+    q_compress = q.view(batch_size, num_heads, seq_len // block_elements,
+                        block_elements, head_dim).mean(dim=3)
+    k_compress = k.view(batch_size, num_heads, seq_len // block_elements,
+                        block_elements, head_dim).mean(dim=3)
+    v_compress = v.view(batch_size, num_heads, seq_len // block_elements,
+                        block_elements, head_dim).mean(dim=3)
+
+    output_compress, block_attn_score = torch_attention(q_compress, k_compress,
+                                                        v_compress)
+
+    output_compress = output_compress.view(batch_size, num_heads,
+                                           seq_len // block_elements, 1,
+                                           head_dim)
+    output_compress = output_compress.repeat(1, 1, 1, block_elements,
+                                             1).view(batch_size, num_heads,
+                                                     seq_len, head_dim)
+
+    q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num = generate_topk_block_sparse_pattern(
+        block_attn_score, topk)
+
+    output_select = block_sparse_attn(q, k, v, q2k_block_sparse_index,
+                                      q2k_block_sparse_num,
+                                      k2q_block_sparse_index,
+                                      k2q_block_sparse_num)
+
+    if compress_attn_weight is not None:
+        final_output = output_compress * compress_attn_weight + output_select
+    else:
+        final_output = output_compress + output_select
+    return final_output
+
+def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
+    QK = torch.matmul(q, k.transpose(-2, -1))
+    QK /= (q.size(-1)**0.5)
+
+    # Causal mask removed since causal is always false
+
+    QK = torch.nn.functional.softmax(QK, dim=-1)
+    output = torch.matmul(QK, v)
+    return output, QK
+
+def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor,
+                                       topk: int):
+    """
+    Generate a block sparse pattern where each q block attends to exactly topk kv blocks,
+    based on the provided attention scores.
+    
+    Args:
+        block_attn_score: [bs, h, num_q_blocks, num_kv_blocks]
+            Attention scores between query and key blocks
+        topk: int
+            Number of kv blocks each q block attends to
+        
+    Returns:
+        q2k_block_sparse_index: [bs, h, num_q_blocks, topk]
+            Contains the indices of kv blocks that each q block attends to.
+        q2k_block_sparse_num: [bs, h, num_q_blocks]
+            Contains the number of kv blocks that each q block attends to (all equal to topk).
+        k2q_block_sparse_index: [bs, h, num_kv_blocks, max_q_per_kv]
+            Contains the indices of q blocks that attend to each kv block.
+        k2q_block_sparse_num: [bs, h, num_kv_blocks]
+            Contains the number of q blocks that attend to each kv block.
+    """
+    device = block_attn_score.device
+    # Extract dimensions from block_attn_score
+    bs, h, num_q_blocks, num_kv_blocks = block_attn_score.shape
+
+    sorted_result = torch.sort(block_attn_score, dim=-1, descending=True)
+
+    sorted_indice = sorted_result.indices
+
+    q2k_block_sparse_index, _ = torch.sort(sorted_indice[:, :, :, :topk],
+                                           dim=-1)
+    q2k_block_sparse_index = q2k_block_sparse_index.to(dtype=torch.int32)
+    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks),
+                                      topk,
+                                      device=device,
+                                      dtype=torch.int32)
+
+    block_map = topk_index_to_map(q2k_block_sparse_index,
+                                  num_kv_blocks,
+                                  transpose_map=True)
+    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(
+        block_map.transpose(2, 3))
+
+    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
+
+@torch._dynamo.disable
+def block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
+    """
+    Differentiable block sparse attention function.
+    
+    Args:
+        q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
+        k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
+        v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
+        q2k_block_sparse_index: Indices for query-to-key sparse blocks
+        q2k_block_sparse_num: Number of sparse blocks for each query block
+        k2q_block_sparse_index: Indices for key-to-query sparse blocks (for backward pass)
+        k2q_block_sparse_num: Number of sparse blocks for each key block (for backward pass)
+    
+    Returns:
+        output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
+    """
+    return BlockSparseAttentionFunction.apply(
+        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
+    )
+
+def block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num):
+    """
+    block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]. 
+        [*, *, i, j] = 1 means the i-th q block should attend to the j-th kv block.
+    """
+    # assert all elements in q2k_block_sparse_num can be devisible by 2
+    o, lse = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
+    return o, lse
+
+def block_sparse_attention_backward(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num):
+    grad_output = grad_output.contiguous()
+    grad_q, grad_k, grad_v = block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
+    return grad_q, grad_k, grad_v
+
+## pytorch sdpa version of block sparse ##
+import triton
+import triton.language as tl
+
+@triton.jit
+def index_to_mask_kernel(
+    q2k_block_sparse_index_ptr,
+    q2k_block_sparse_num_ptr,
+    mask_ptr,
+    batch_size: tl.constexpr,
+    num_heads: tl.constexpr,
+    num_q_blocks: tl.constexpr,
+    num_k_blocks: tl.constexpr,
+    max_kv_blocks: tl.constexpr,
+    BLOCK_Q: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    bh, q, id = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64)
+    b = bh // num_heads
+    h = bh % num_heads
+
+    num_valid_blocks = tl.load(q2k_block_sparse_num_ptr + b * num_heads * num_q_blocks + h * num_q_blocks + q)
+
+    if num_valid_blocks <= id:
+        return
+    k = tl.load(q2k_block_sparse_index_ptr + b * num_heads * num_q_blocks * max_kv_blocks + h * num_q_blocks * max_kv_blocks + q * max_kv_blocks + id)
+
+    full_mask = (tl.arange(0, BLOCK_Q)[:, None] < BLOCK_Q) & (tl.arange(0, BLOCK_K)[None, :] < BLOCK_K)
+
+    q_lengths = num_q_blocks * BLOCK_Q
+    k_lengths = num_k_blocks * BLOCK_K
+    mask_ptr_base = mask_ptr + b * num_heads * q_lengths * k_lengths + h * q_lengths * k_lengths + q * BLOCK_Q * k_lengths + k * BLOCK_K
+
+    tl.store(mask_ptr_base + tl.arange(0, BLOCK_Q)[:, None] * k_lengths + tl.arange(0, BLOCK_K)[None, :], full_mask)
+
+def index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, BLOCK_Q, BLOCK_K, num_k_blocks):
+    """
+    Convert block sparse indices to a mask.
+    
+    Args:
+        q2k_block_sparse_index: Indices for query-to-key sparse blocks
+        q2k_block_sparse_num: Number of sparse blocks for each query block
+    
+    Returns:
+        mask: Block sparse mask tensor
+    """
+    batch_size, num_heads, num_q_blocks, max_kv_blocks = q2k_block_sparse_index.shape
+    assert q2k_block_sparse_num.shape == (batch_size, num_heads, num_q_blocks)
+
+    mask = torch.zeros((batch_size, num_heads, num_q_blocks * BLOCK_Q, num_k_blocks * BLOCK_K), dtype=torch.bool, device=q2k_block_sparse_index.device)
+
+    grid = (batch_size * num_heads, num_q_blocks, max_kv_blocks)
+    index_to_mask_kernel[grid](
+        q2k_block_sparse_index,
+        q2k_block_sparse_num,
+        mask,
+        batch_size,
+        num_heads,
+        num_q_blocks,
+        num_k_blocks,
+        max_kv_blocks,
+        BLOCK_Q=BLOCK_Q,
+        BLOCK_K=BLOCK_K,
+    )
+
+    return mask
+
+@triton.jit
+def topk_index_to_map_kernel(
+    map_ptr,
+    index_ptr,
+    map_bs_stride,
+    map_h_stride,
+    map_q_stride,
+    map_kv_stride,
+    index_bs_stride,
+    index_h_stride,
+    index_q_stride,
+    index_kv_stride,
+    topk: tl.constexpr,
+):
+    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
+    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
+
+    for i in tl.static_range(topk):
+        index = tl.load(index_ptr_base + i * index_kv_stride)
+        tl.store(map_ptr_base + index * map_kv_stride, 1.0)
+
+@triton.jit
+def map_to_index_kernel(
+    map_ptr,
+    index_ptr,
+    index_num_ptr,
+    map_bs_stride,
+    map_h_stride,
+    map_q_stride,
+    map_kv_stride,
+    index_bs_stride,
+    index_h_stride,
+    index_q_stride,
+    index_kv_stride,
+    index_num_bs_stride,
+    index_num_h_stride,
+    index_num_q_stride,
+    num_kv_blocks: tl.constexpr,
+):
+    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
+    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
+
+    num = 0
+    for i in tl.static_range(num_kv_blocks):
+        map_entry = tl.load(map_ptr_base + i * map_kv_stride)
+        if map_entry:
+            tl.store(index_ptr_base + num * index_kv_stride, i)
+            num += 1
+
+    tl.store(
+        index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride +
+        q * index_num_q_stride, num)
+
+def topk_index_to_map(index: torch.Tensor,
+                      num_kv_blocks: int,
+                      transpose_map: bool = False):
+    """
+    Convert topk indices to a map.
+    
+    Args:
+        index: [bs, h, num_q_blocks, topk]
+            The topk indices tensor.
+        num_kv_blocks: int
+            The number of key-value blocks in the block_map returned
+        transpose_map: bool
+            If True, the block_map will be transposed on the final two dimensions.
+    
+    Returns:
+        block_map: [bs, h, num_q_blocks, num_kv_blocks]
+            A binary map where 1 indicates that the q block attends to the kv block.
+    """
+    bs, h, num_q_blocks, topk = index.shape
+
+    if transpose_map is False:
+        block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks),
+                                dtype=torch.bool,
+                                device=index.device)
+    else:
+        block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks),
+                                dtype=torch.bool,
+                                device=index.device)
+        block_map = block_map.transpose(2, 3)
+
+    grid = (bs, h, num_q_blocks)
+    topk_index_to_map_kernel[grid](
+        block_map,
+        index,
+        block_map.stride(0),
+        block_map.stride(1),
+        block_map.stride(2),
+        block_map.stride(3),
+        index.stride(0),
+        index.stride(1),
+        index.stride(2),
+        index.stride(3),
+        topk=topk,
+    )
+
+    return block_map
+
+def map_to_index(block_map: torch.Tensor):
+    """
+    Convert a block map to indices and counts.
+    
+    Args:
+        block_map: [bs, h, num_q_blocks, num_kv_blocks]
+            The block map tensor.
+    
+    Returns:
+        index: [bs, h, num_q_blocks, num_kv_blocks]
+            The indices of the blocks.
+        index_num: [bs, h, num_q_blocks]
+            The number of blocks for each q block.
+    """
+    bs, h, num_q_blocks, num_kv_blocks = block_map.shape
+
+    index = torch.full((block_map.shape),
+                       -1,
+                       dtype=torch.int32,
+                       device=block_map.device)
+    index_num = torch.empty((bs, h, num_q_blocks),
+                            dtype=torch.int32,
+                            device=block_map.device)
+
+    grid = (bs, h, num_q_blocks)
+    map_to_index_kernel[grid](
+        block_map,
+        index,
+        index_num,
+        block_map.stride(0),
+        block_map.stride(1),
+        block_map.stride(2),
+        block_map.stride(3),
+        index.stride(0),
+        index.stride(1),
+        index.stride(2),
+        index.stride(3),
+        index_num.stride(0),
+        index_num.stride(1),
+        index_num.stride(2),
+        num_kv_blocks=num_kv_blocks,
+    )
+
+    return index, index_num
+
+class BlockSparseAttentionFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
+        o, lse = block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
+        ctx.save_for_backward(q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num)
+        return o
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num = ctx.saved_tensors
+        grad_q, grad_k, grad_v = block_sparse_attention_backward(
+            q, k, v, o, lse, grad_output, k2q_block_sparse_index, k2q_block_sparse_num
+        )
+        return grad_q, grad_k, grad_v, None, None, None, None
+
+
+class DummyOperator(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        return x
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+class CheckpointSDPA(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, obj, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k):
+        """Forward pass."""
+        with torch.no_grad():
+            mask = index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k, k.shape[2] // block_k)
+            outputs = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        ctx.save_for_backward(*detach_variable((q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)))
+        ctx.block_q = block_q
+        ctx.block_k = block_k
+        # the obj is passed in, then it can access the saved input
+        # tensors later for recomputation
+        obj.ctx = ctx
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """Backward pass."""
+        inputs = ctx.saved_tensors
+        output = ctx.output
+        torch.autograd.backward(output, grad_output)
+        ctx.output = None
+        grads = tuple(inp.grad for inp in inputs)
+        return (None, ) + grads + (None, None)
+
+
+class BlockSparseAttnTorch:
+    def __init__(self):
+        self.ctx = None
+    
+    def recompute_mask(self, _):
+        recomputed_mask = index_to_mask(self.q2k_block_sparse_index, self.q2k_block_sparse_num, self.block_q, self.block_k, self.num_kv_blocks)
+        mask_size = recomputed_mask.untyped_storage().size()
+        self.mask.untyped_storage().resize_(mask_size)
+        self.mask.untyped_storage().copy_(recomputed_mask.untyped_storage()) 
+    
+    def recompute(self, _):
+        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num = self.ctx.saved_tensors
+        block_q = self.ctx.block_q
+        block_k = self.ctx.block_k
+        mask = index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k, k.shape[2] // block_k)
+        with torch.enable_grad():
+            output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        self.ctx.output = output
+        self.ctx = None
+
+    @torch._dynamo.disable
+    def forward(self, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k):
+        """
+        Differentiable block sparse attention function using PyTorch.
+        
+        Args:
+            q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
+            k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
+            v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
+            q2k_block_sparse_index: Indices for query-to-key sparse blocks
+            q2k_block_sparse_num: Number of sparse blocks for each query block
+            block_q: Block size for query
+            block_k: Block size for key-value
+
+        Returns:
+            output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
+        """
+
+        output = CheckpointSDPA.apply(
+            self, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k
+        )
+
+        o = DummyOperator.apply(output)
+        o.register_hook(self.recompute)
+        return o

docker/Dockerfile.python3.10

@@ -0,0 +1,48 @@
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+ENV PATH=/opt/conda/bin:$PATH
+
+RUN conda create --name fastvideo-dev python=3.10.0 -y
+
+SHELL ["/bin/bash", "-c"]
+
+# Copy just the pyproject.toml first to leverage Docker cache
+COPY pyproject.toml ./
+
+# Create a dummy README to satisfy the installation
+RUN echo "# Placeholder" > README.md
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
+    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
+    conda clean -afy
+
+COPY . .
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
+
+# Remove authentication headers
+RUN git config --unset-all http.https://github.com/.extraheader || true
+
+# Set up automatic conda environment activation for all shells
+RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
+    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
+    # Ensure .bashrc is sourced for SSH login shells
+    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+
+EXPOSE 22

docker/Dockerfile.python3.11

@@ -0,0 +1,48 @@
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+ENV PATH=/opt/conda/bin:$PATH
+
+RUN conda create --name fastvideo-dev python=3.11.11 -y
+
+SHELL ["/bin/bash", "-c"]
+
+# Copy just the pyproject.toml first to leverage Docker cache
+COPY pyproject.toml ./
+
+# Create a dummy README to satisfy the installation
+RUN echo "# Placeholder" > README.md
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
+    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
+    conda clean -afy
+
+COPY . .
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
+
+# Remove authentication headers
+RUN git config --unset-all http.https://github.com/.extraheader || true
+
+# Set up automatic conda environment activation for all shells
+RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
+    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
+    # Ensure .bashrc is sourced for SSH login shells
+    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+
+EXPOSE 22

docker/Dockerfile.python3.12

@@ -0,0 +1,72 @@
+FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+SHELL ["/bin/bash", "-c"]
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    zsh \
+    vim \
+    curl \
+    gcc-11 \
+    g++-11 \
+    clang-11 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set up C++20 compilers for ThunderKittens
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+# Set CUDA environment variables
+ENV CUDA_HOME=/usr/local/cuda-12.8
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+
+# Install uv and source its environment
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    echo 'source $HOME/.local/bin/env' >> /root/.bashrc
+
+# Copy just the pyproject.toml first to leverage Docker cache
+COPY pyproject.toml ./
+
+# Create a dummy README to satisfy the installation
+RUN echo "# Placeholder" > README.md
+
+# Create and activate virtual environment with specific Python version and seed
+RUN source $HOME/.local/bin/env && \
+    uv venv --python 3.12 --seed /opt/venv && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir --upgrade pip && \
+    uv pip install --no-cache-dir .[dev] && \
+    uv pip install --no-cache-dir flash-attn==2.8.0.post2 --no-build-isolation
+
+COPY . .
+
+# Install dependencies using uv and set up shell configuration
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    uv pip install --no-cache-dir -e .[dev] && \
+    git config --unset-all http.https://github.com/.extraheader || true && \
+    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
+    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+
+# Install STA (Sliding Tile Attention)
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_sta.py install
+
+# Install VSA 
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn && \
+    git submodule update --init --recursive && \
+    python setup_vsa.py install
+
+EXPOSE 22

docs/Makefile

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

docs/README.md

@@ -0,0 +1,20 @@
+# FastVideo documents
+
+## Build the docs
+
+```bash
+# Install dependencies.
+pip install -r requirements-docs.txt
+
+# Build the docs.
+make clean
+make html
+```
+
+## Open the docs with your browser
+
+```bash
+python -m http.server -d build/html/
+```
+
+Launch your browser and open localhost:8000.

docs/data_preprocess.md

@@ -1,68 +0,0 @@
-
-
-
-## 🧱 Data Preprocess
-
-To save GPU memory, we precompute text embeddings and VAE latents to eliminate the need to load the text encoder and VAE during training.
-
-
-We provide a sample dataset to help you get started. Download the source media using the following command:
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/Image-Vid-Finetune-Src --local_dir=data/Image-Vid-Finetune-Src --repo_type=dataset
-```
-To preprocess the dataset for fine-tuning or distillation, run:
-```
-bash scripts/preprocess/preprocess_mochi_data.sh # for mochi
-bash scripts/preprocess/preprocess_hunyuan_data.sh # for hunyuan
-```
-
-The preprocessed dataset will be stored in `Image-Vid-Finetune-Mochi` or `Image-Vid-Finetune-HunYuan` correspondingly.
-
-### Process your own dataset
-
-If you wish to create your own dataset for finetuning or distillation, please structure you video dataset in the following format:
-
-path_to_dataset_folder/
-├── media/
-│   ├── 0.jpg
-│   ├── 1.mp4
-│   ├── 2.jpg
-├── video2caption.json
-└── merge.txt
-
-Format the JSON file as a list, where each item represents a media source:
-
-For image media,
-```
-{
-    "path": "0.jpg",
-    "cap": ["captions"]
-}
-```
-For video media, 
-```
-{
-    "path": "1.mp4",
-    "resolution": {
-      "width": 848,
-      "height": 480
-    },
-    "fps": 30.0,
-    "duration": 6.033333333333333,
-    "cap": [
-      "caption"
-    ]
-  }
-```
-
-Use a txt file (merge.txt) to contain the source folder for media and the JSON file for meta information:
-
-```
-path_to_media_source_foder,path_to_json_file
-```
-
-Adjust the `DATA_MERGE_PATH` and `OUTPUT_DIR` in `scripts/preprocess/preprocess_****_data.sh` accordingly and run:
-```
-bash scripts/preprocess/preprocess_****_data.sh
-```
-The preprocessed data will be put into the `OUTPUT_DIR` and the `videos2caption.json` can be used in finetune and distill scripts.

docs/make.bat

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

docs/requirements-docs.txt

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

docs/seed_parameter_behavior.md

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

docs/source/_static/custom.css

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

docs/source/_static/custom.js

@@ -0,0 +1,18 @@
+// Update URL search params when tab is clicked
+  document.addEventListener("DOMContentLoaded", function () {
+    const tabs = document.querySelectorAll(".sd-tab-label");
+
+    function updateURL(tab) {
+      const syncGroup = tab.getAttribute("data-sync-group");
+      const syncId = tab.getAttribute("data-sync-id");
+      if (syncGroup && syncId) {
+          const url = new URL(window.location);
+          url.searchParams.set(syncGroup, syncId);
+          window.history.replaceState(null, "", url);
+      }
+    }
+
+    tabs.forEach(tab => {
+        tab.addEventListener("click", () => updateURL(tab));
+    });
+});

docs/source/_templates/sections/header.html

@@ -0,0 +1,39 @@
+<style>
+  .notification-bar {
+    width: 100vw;
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    font-size: 16px;
+    padding: 0 6px 0 6px;
+  }
+  .notification-bar p {
+    margin: 0;
+  }
+  .notification-bar a {
+    font-weight: bold;
+    text-decoration: none;
+  }
+
+  /* Light mode styles (default) */
+  .notification-bar {
+    background-color: #fff3cd;
+    color: #856404;
+  }
+  .notification-bar a {
+    color: #d97706;
+  }
+
+  /* Dark mode styles */
+  html[data-theme=dark] .notification-bar {
+    background-color: #333;
+    color: #ddd;
+  }
+  html[data-theme=dark] .notification-bar a {
+    color: #ffa500; /* Brighter color for visibility */
+  }
+</style>
+
+<!-- <div class="notification-bar">
+  <p>You are viewing the latest developer preview docs. <a href="https://docs.vllm.ai/en/stable/">Click here</a> to view docs for the latest stable release.</p>
+</div> -->

docs/source/api/summary.md

@@ -0,0 +1,19 @@
+# Summary
+
+## Video Generator
+
+```{autodoc2-summary}
+    fastvideo.VideoGenerator
+```
+
+## Initialization Configuration
+
+```{autodoc2-summary}
+    fastvideo.v1.configs.pipelines.PipelineConfig
+```
+
+## Sampling Configuration
+
+```{autodoc2-summary}
+    fastvideo.v1.configs.sample.SamplingParam
+```

docs/source/autodoc2_docstring_parser.py

@@ -0,0 +1,22 @@
+# type: ignore
+# SPDX-License-Identifier: Apache-2.0
+from docutils import nodes
+from myst_parser.parsers.sphinx_ import MystParser
+from sphinx.ext.napoleon import docstring
+
+
+class NapoleonParser(MystParser):
+
+    def parse(self, input_string: str, document: nodes.document) -> None:
+        # Get the Sphinx configuration
+        config = document.settings.env.config
+
+        parsed_content = str(
+            docstring.GoogleDocstring(
+                str(docstring.NumpyDocstring(input_string, config)),
+                config,
+            ))
+        return super().parse(parsed_content, document)
+
+
+Parser = NapoleonParser

docs/source/conf.py

@@ -0,0 +1,276 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Configuration file for the Sphinx documentation builder.
+#
+# This file only contains a selection of the most common options. For a full
+# list see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+
+import datetime
+import logging
+import os
+import re
+import sys
+from pathlib import Path
+
+import requests
+
+logger = logging.getLogger(__name__)
+REPO_ROOT = Path(__file__).resolve().parent.parent.parent
+print(os.path.abspath(REPO_ROOT))
+sys.path.append(os.path.abspath(REPO_ROOT))
+
+# -- Project information -----------------------------------------------------
+
+project = 'FastVideo'
+copyright = f'{datetime.datetime.now().year}, FastVideo Team'
+author = 'the FastVideo Team'
+
+# -- General configuration ---------------------------------------------------
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+    "sphinx.ext.napoleon",
+    "sphinx.ext.linkcode",
+    "sphinx.ext.intersphinx",
+    "sphinx_copybutton",
+    "autodoc2",
+    "myst_parser",
+    "sphinxarg.ext",
+    "sphinx_design",
+    "sphinx_togglebutton",
+]
+myst_enable_extensions = [
+    "colon_fence",
+    "fieldlist",
+]
+autodoc2_packages = [
+    {
+        "path": "../../fastvideo",
+        "exclude_dirs": ["__pycache__", "third_party"],
+    },
+]
+autodoc2_output_dir = "api"
+autodoc2_render_plugin = "myst"
+autodoc2_hidden_objects = ["dunder", "private", "inherited"]
+autodoc2_docstring_parser_regexes = [
+    (".*", "docs.source.autodoc2_docstring_parser"),
+]
+autodoc2_sort_names = True
+autodoc2_index_template = None
+autodoc2_skip_module_regexes = [
+    "fastvideo.dataset",
+    "fastvideo.distill",
+    "fastvideo.data_preprocess",
+    "fastvideo.models",
+    "fastvideo.sample",
+    "fastvideo.utils",
+    "fastvideo.distill_adv",
+    "fastvideo.train",
+]
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ['_templates']
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns: list[str] = ["**/*.template.md", "**/*.inc.md"]
+
+# Exclude the prompt "$" when copying code
+copybutton_prompt_text = r"\$ "
+copybutton_prompt_is_regexp = True
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+html_title = project
+html_theme = 'sphinx_book_theme'
+html_logo = '../../assets/logo.jpg'
+#html_favicon = 'assets/logos/vllm-logo-only-light.ico'
+html_theme_options = {
+    'path_to_docs': 'docs/source',
+    'repository_url': 'https://github.com/hao-ai-lab/FastVideo/',
+    'use_repository_button': True,
+    'use_edit_page_button': True,
+    # Prevents the full API being added to the left sidebar of every page.
+    # Reduces build time by 2.5x and reduces build size from ~225MB to ~95MB.
+    'collapse_navbar': True,
+    # Makes API visible in the right sidebar on API reference pages.
+    'show_toc_level': 3,
+}
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ["_static"]
+html_js_files = ["custom.js"]
+html_css_files = ["custom.css"]
+
+myst_url_schemes = {
+    'http': None,
+    'https': None,
+    'mailto': None,
+    'ftp': None,
+    "gh-issue": {
+        "url":
+        "https://github.com/hao-ai-lab/FastVideo/issues/{{path}}#{{fragment}}",
+        "title": "Issue #{{path}}",
+        "classes": ["github"],
+    },
+    "gh-pr": {
+        "url":
+        "https://github.com/hao-ai-lab/FastVideo/pull/{{path}}#{{fragment}}",
+        "title": "Pull Request #{{path}}",
+        "classes": ["github"],
+    },
+    "gh-dir": {
+        "url": "https://github.com/hao-ai-lab/FastVideo/tree/main/{{path}}",
+        "title": "{{path}}",
+        "classes": ["github"],
+    },
+    "gh-file": {
+        "url": "https://github.com/hao-ai-lab/FastVideo/blob/main/{{path}}",
+        "title": "{{path}}",
+        "classes": ["github"],
+    },
+}
+
+# see https://docs.readthedocs.io/en/stable/reference/environment-variables.html # noqa
+READTHEDOCS_VERSION_TYPE = os.environ.get('READTHEDOCS_VERSION_TYPE')
+if READTHEDOCS_VERSION_TYPE == "tag":
+    # remove the warning banner if the version is a tagged release
+    header_file = os.path.join(os.path.dirname(__file__),
+                               "_templates/sections/header.html")
+    # The file might be removed already if the build is triggered multiple times
+    # (readthedocs build both HTML and PDF versions separately)
+    if os.path.exists(header_file):
+        os.remove(header_file)
+
+
+# Generate additional rst documentation here.
+def setup(app):
+    from docs.source.generate_examples import generate_examples
+    generate_examples()
+
+
+_cached_base: str = ""
+_cached_branch: str = ""
+
+
+def get_repo_base_and_branch(pr_number: str) -> tuple[str | None, str | None]:
+    global _cached_base, _cached_branch
+    if _cached_base and _cached_branch:
+        return _cached_base, _cached_branch
+
+    url = f"https://api.github.com/repos/hao-ai-lab/FastVideo/pulls/{pr_number}"
+    response = requests.get(url)
+    if response.status_code == 200:
+        data = response.json()
+        _cached_base = data['head']['repo']['full_name']
+        _cached_branch = data['head']['ref']
+        return _cached_base, _cached_branch
+    else:
+        logger.error("Failed to fetch PR details: %s", response)
+        return None, None
+
+
+def linkcode_resolve(domain, info):
+    if domain != 'py':
+        return None
+    if not info['module']:
+        return None
+
+    # Get path from module name
+    file = Path(f"{info['module'].replace('.', '/')}.py")
+    path = REPO_ROOT / file
+    if not path.exists():
+        path = REPO_ROOT / file.with_suffix("") / "__init__.py"
+    if not path.exists():
+        return None
+
+    # Get the line number of the object
+    with open(path) as f:
+        lines = f.readlines()
+    name = info['fullname'].split(".")[-1]
+    pattern = fr"^( {{4}})*((def|class) )?{name}\b.*"
+    for lineno, line in enumerate(lines, 1):
+        if not line or line.startswith("#"):
+            continue
+        if re.match(pattern, line):
+            break
+
+    # If the line number is not found, return None
+    if lineno == len(lines):
+        return None
+
+    # If the line number is found, create the URL
+    filename = path.relative_to(REPO_ROOT)
+    if "checkouts" in path.parts:
+        # a PR build on readthedocs
+        pr_number = REPO_ROOT.name
+        base, branch = get_repo_base_and_branch(pr_number)
+        if base and branch:
+            return f"https://github.com/{base}/blob/{branch}/{filename}#L{lineno}"
+    # Otherwise, link to the source file on the main branch
+    return f"https://github.com/hao-ai-lab/FastVideo/blob/main/{filename}#L{lineno}"
+
+
+# Mock out external dependencies here, otherwise the autodoc pages may be blank.
+autodoc_mock_imports = [
+    "blake3",
+    "compressed_tensors",
+    "cpuinfo",
+    "cv2",
+    "torch",
+    "huggingface_hub",
+    "torchvision",
+    "transformers",
+    "psutil",
+    "prometheus_client",
+    "sentencepiece",
+    "vllm._C",
+    "PIL",
+    "numpy",
+    'triton',
+    "tqdm",
+    "tensorizer",
+    "pynvml",
+    "outlines",
+    "xgrammar",
+    "librosa",
+    "soundfile",
+    "gguf",
+    "lark",
+    "decord",
+]
+
+for mock_target in autodoc_mock_imports:
+    if mock_target in sys.modules:
+        logger.info(
+            "Potentially problematic mock target (%s) found; "
+            "autodoc_mock_imports cannot mock modules that have already "
+            "been loaded into sys.modules when the sphinx build starts.",
+            mock_target)
+
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3", None),
+    "typing_extensions":
+    ("https://typing-extensions.readthedocs.io/en/latest", None),
+    "aiohttp": ("https://docs.aiohttp.org/en/stable", None),
+    "pillow": ("https://pillow.readthedocs.io/en/stable", None),
+    "numpy": ("https://numpy.org/doc/stable", None),
+    "torch": ("https://pytorch.org/docs/stable", None),
+    "psutil": ("https://psutil.readthedocs.io/en/stable", None),
+}
+
+navigation_with_keys = False

docs/source/contributing/developer_env/docker.md

@@ -0,0 +1,32 @@
+(docker)=
+# 🐳 Using the FastVideo Docker Image
+
+If you prefer a containerized development environment or want to avoid managing dependencies manually, you can use our prebuilt Docker image:
+
+**Image:** [`ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest`](https://ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev)
+
+## Starting the container
+
+```bash
+docker run --gpus all -it ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
+```
+
+This will:
+
+- Start the container with GPU access  
+- Drop you into a shell with the `fastvideo-dev` Conda environment preconfigured
+
+## Using the container
+
+```bash
+# Conda environment should already be active
+# FastVideo package installed in editable mode
+
+# Pull the latest changes from remote
+cd /FastVideo
+git pull
+
+# Run linters and tests
+pre-commit run --all-files
+pytest tests/
+```

docs/source/contributing/developer_env/index.md

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

docs/source/contributing/developer_env/runpod.md

@@ -0,0 +1,54 @@
+(runpod)=
+
+# 📦 Developing FastVideo on RunPod
+
+You can easily use the FastVideo Docker image as a custom container on [RunPod](https://www.runpod.io) for development or experimentation.
+
+## Creating a new pod
+
+Choose a GPU that supports CUDA 12.4
+
+Pick 1 or 2 L40S GPU(s)
+
+![RunPod CUDA selection](../../_static/images/runpod_cuda.png)
+
+When creating your pod 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)):
+
+```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](../../_static/images/runpod_template.png)
+
+After deploying, the pod will take a few minutes to pull the image and start the SSH service.
+
+![RunPod ssh](../../_static/images/runpod_ssh.png)
+
+## Working with the pod
+
+After SSH'ing into your pod, you'll find the `fastvideo-dev` Conda environment already activated.
+
+To pull in the latest changes from the GitHub repo:
+
+```bash
+cd /FastVideo
+git pull
+```
+
+`If you have a persistent volume and want to keep your code changes, you can move /FastVideo to /workspace/FastVideo, or simply clone the repository there.`
+
+Run your development workflows as usual:
+
+```bash
+# Run linters
+pre-commit run --all-files
+
+# Run tests
+pytest tests/
+```

docs/source/contributing/overview.md

@@ -0,0 +1,52 @@
+(developer-overview)=
+
+# 🛠️ Contributing to FastVideo
+
+Thank you for your interest in contributing to FastVideo. We want to make the process as smooth for you as possible and this is a guide to help get you started!
+
+Our community is open to everyone and welcomes any contributions no matter how large or small.
+
+# Developer Environment:
+Do make sure you have CUDA 12.4 installed and supported. FastVideo currently only support Linux and CUDA GPUs, but we hope to support other platforms in the future.
+
+We recommend using a fresh Python 3.10 Conda environment to develop FastVideo:
+
+Install Miniconda:
+
+```
+wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+bash Miniconda3-latest-Linux-x86_64.sh
+source ~/.bashrc
+```
+
+Create and activate a Conda environment for FastVideo:
+
+```
+conda create -n fastvideo python=3.10 -y
+conda activate fastvideo
+```
+
+Clone the FastVideo repository and go to the FastVideo directory:
+
+```
+git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
+
+```
+
+Now you can install FastVideo and setup git hooks for running linting. By using `pre-commit`, the linters will run and have to pass before you'll be able to make a commit.
+
+```bash
+pip install -e .[dev]
+
+# Can also install flash-attn (optional)
+pip install flash-attn==2.7.4.post1 --no-build-isolation 
+
+# Linting, formatting and static type checking
+pre-commit install --hook-type pre-commit --hook-type commit-msg
+
+# You can manually run pre-commit with
+pre-commit run --all-files
+
+# Unit tests
+pytest tests/
+```

docs/source/design/overview.md

@@ -0,0 +1,410 @@
+# 🔍 FastVideo Overview
+
+This document outlines FastVideo's architecture for developers interested in framework internals or contributions. It serves as an onboarding guide for new contributors by providing an overview of the most important directories and files within the `fastvideo/v1/` codebase.
+
+## Table of Contents - V1 Directory Structure and Files
+
+- [`fastvideo/v1/pipelines/`](#design-pipeline-system) - Core diffusion pipeline components
+- [`fastvideo/v1/models/`](#design-model-components) - Model implementations
+  - [`dits/`](#design-transformer-models) - Transformer-based diffusion models
+  - [`vaes/`](#design-vae-variational-auto-encoder) - Variational autoencoders
+  - [`encoders/`](#design-text-and-image-encoders) - Text and image encoders
+  - [`schedulers/`](#design-schedulers) - Diffusion schedulers
+- [`fastvideo/v1/attention/`](#design-optimized-attention) - Optimized attention implementations
+- [`fastvideo/v1/distributed/`](#design-distributed-processing) - Distributed computing utilities
+- [`fastvideo/v1/layers/`](#design-tensor-parallelism) - Custom neural network layers
+- [`fastvideo/v1/platforms/`](#design-platforms) - Hardware platform abstractions
+- [`fastvideo/v1/worker/`](#design-executor-and-worker-abstractions) - Multi-GPU process management
+- [`fastvideo/v1/fastvideo_args.py`](#design-fastvideo-args) - Argument handling
+- [`fastvideo/v1/forward_context.py`](#design-forwardcontext) - Forward pass context management
+- `fastvideo/v1/utils.py` - Utility functions
+- [`fastvideo/v1/logger.py`](#design-logger) - Logging infrastructure
+
+## Core Architecture
+
+FastVideo separates model components from execution logic with these principles:
+- **Component Isolation**: Models (encoders, VAEs, transformers) are isolated from execution (pipelines, stages, distributed processing)
+- **Modular Design**: Components can be independently replaced
+- **Distributed Execution**: Supports various parallelism strategies (Tensor, Sequence)
+- **Custom Attention Backends**: Components can support and use different Attention implementations
+- **Pipeline Abstraction**: Consistent interface across diffusion models
+
+(design-fastvideo-args)=
+## FastVideoArgs
+
+The `FastVideoArgs` class in `fastvideo/v1/fastvideo_args.py` serves as the central configuration system for FastVideo. It contains all parameters needed to control model loading, inference configuration, performance optimization settings, and more.
+
+Key features include:
+- **Command-line Interface**: Automatic conversion between CLI arguments and dataclass fields
+- **Configuration Groups**: Organized by functional areas (model loading, video params, optimization settings)
+- **Context Management**: Global access to current settings via `get_current_fastvideo_args()`
+- **Parameter Validation**: Ensures valid combinations of settings
+
+Common configuration areas:
+- **Model paths and loading options**: `model_path`, `trust_remote_code`, `revision`
+- **Distributed execution settings**: `num_gpus`, `tp_size`, `sp_size`
+- **Video generation parameters**: `height`, `width`, `num_frames`, `num_inference_steps`
+- **Precision settings**: Control computation precision for different components
+
+Example usage:
+
+```python
+# Load arguments from command line
+fastvideo_args = prepare_fastvideo_args(sys.argv[1:])
+
+# Access parameters
+model = load_model(fastvideo_args.model_path)
+
+# Set as global context
+with set_current_fastvideo_args(fastvideo_args):
+    # Code that requires access to these arguments
+    result = generate_video()
+```
+
+(design-pipeline-system)=
+## Pipeline System
+
+### `ComposedPipelineBase`
+
+This foundational class provides:
+
+- **Model Loading**: Automatically loads components from HuggingFace-Diffusers-compatible model directories
+- **Stage Management**: Creates and orchestrates processing stages
+- **Data Flow Coordination**: Ensures proper state flow between stages
+
+```python
+class MyCustomPipeline(ComposedPipelineBase):
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+    
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        # Pipeline-specific initialization
+        pass
+        
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        self.add_stage("input_validation_stage", InputValidationStage())
+        self.add_stage("text_encoding_stage", CLIPTextEncodingStage(
+            text_encoder=self.get_module("text_encoder"),
+            tokenizer=self.get_module("tokenizer")
+        ))
+        # Additional stages...
+```
+
+### Pipeline Stages
+Each stage handles a specific diffusion process component:
+- **Input Validation**: Parameter verification
+- **Text Encoding**: CLIP, LLaMA, or T5-based encoding
+- **Image Encoding**: Image input processing
+- **Timestep & Latent Preparation**: Setup for diffusion
+- **Denoising**: Core diffusion loop
+- **Decoding**: Latent-to-pixel conversion
+
+Each stage implements a standard interface:
+
+```python
+def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+    # Process batch and update state
+    return batch
+```
+
+(design-forwardbatch)=
+### ForwardBatch
+
+Defined in `fastvideo/v1/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:
+
+- **Input Data**: Prompts, images, generation parameters
+- **Intermediate State**: Embeddings, latents, timesteps, accumulated during stage execution
+- **Output Storage**: Generated results and metadata
+- **Configuration**: Sampling parameters, precision settings
+
+This structure facilitates clear state transitions between stages.
+
+(design-model-components)=
+## Model Components
+
+The `fastvideo/v1/models/` directory contains implementations of the core neural network models used in video diffusion:
+
+(design-transformer-models)=
+### Transformer Models
+
+Transformer networks perform the actual denoising during diffusion:
+
+- **Location**: `fastvideo/v1/models/dits/`
+- **Examples**:
+  - `WanTransformer3DModel`
+  - `HunyuanVideoTransformer3DModel`
+
+Features include:
+- Text/image conditioning
+- Standardized interface for model-specific optimizations
+
+```python
+def forward(
+    self, 
+    latents,                    # [B, T, C, H, W]
+    encoder_hidden_states,      # Text embeddings
+    timestep,                   # Current diffusion timestep
+    encoder_hidden_states_image=None,  # Optional image embeddings
+    **kwargs
+):
+    # Perform denoising computation
+    return noise_pred  # Predicted noise residual
+```
+
+(design-vae-variational-auto-encoder)=
+### VAE (Variational Auto-Encoder)
+
+VAEs handle conversion between pixel space and latent space:
+
+- **Location**: `fastvideo/v1/models/vaes/`
+- **Examples**:
+  - `AutoencoderKLWan`
+  - `AutoencoderKLHunyuanVideo`
+
+These models compress image/video data to a more efficient latent representation (typically 4x-8x smaller in each dimension).
+
+FastVideo's VAE implementations include:
+- Efficient video batch processing
+- Memory optimization
+- Optional tiling for large frames
+- Distributed weight support
+
+(design-text-and-image-encoders)=
+### Text and Image Encoders
+
+Encoders process conditioning inputs into embeddings:
+
+- **Location**: `fastvideo/v1/models/encoders/`
+- **Text Encoders**:
+  - `CLIPTextModel`
+  - `LlamaModel`
+  - `UMT5EncoderModel`
+- **Image Encoders**:
+  - `CLIPVisionModel`
+
+FastVideo implements optimizations such as:
+- Vocab parallelism for distributed processing
+- Caching for common prompts
+- Precision-tuned computation
+
+(design-schedulers)=
+### Schedulers
+
+Schedulers manage the diffusion sampling process:
+
+- **Location**: `fastvideo/v1/models/schedulers/`
+- **Examples**:
+  - `UniPCMultistepScheduler`
+  - `FlowMatchEulerDiscreteScheduler`
+
+These components control:
+- Diffusion timestep sequences
+- Noise prediction to latent update conversions
+- Quality/speed trade-offs
+
+```python
+def step(
+    self, 
+    model_output: torch.Tensor,
+    timestep: torch.LongTensor,
+    sample: torch.Tensor,
+    **kwargs
+) -> torch.Tensor:
+    # Process model output and update latents
+    # Return updated latents
+    return prev_sample
+```
+
+(design-optimized-attention)=
+## Optimized Attention
+
+The `fastvideo/v1/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:
+
+### Attention Backends
+Multiple implementations with automatic selection:
+- **FLASH_ATTN**: Optimized for supporting hardware
+- **TORCH_SDPA**: Built-in PyTorch scaled dot-product attention
+- **SLIDING_TILE_ATTN**: For very long sequences
+
+```python
+# Configure available attention backends for this layer
+self.attn = LocalAttention(
+    num_heads=num_heads,
+    head_size=head_dim,
+    causal=False,
+    supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+)
+
+# Override via environment variable
+# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+```
+
+### Attention Patterns
+Supports various patterns with memory optimization techniques:
+- **Cross/Self/Temporal/Global-Local Attention**
+- Chunking, progressive computation, optimized masking
+
+(design-distributed-processing)=
+## Distributed Processing
+
+The `fastvideo/v1/distributed/` directory contains implementations for distributed model execution:
+
+(design-tensor-parallelism)=
+### Tensor Parallelism
+
+Tensor parallelism splits model weights across devices:
+
+- **Implementation**: Through `RowParallelLinear` and `ColumnParallelLinear` layers
+- **Use cases**: Will be used by encoder models as their sequence lengths are shorter and enables efficient sharding.
+
+```python
+# Tensor-parallel layers in a transformer block
+from fastvideo.v1.layers.linear import ColumnParallelLinear, RowParallelLinear
+
+# Split along output dimension
+self.qkv_proj = ColumnParallelLinear(
+    input_size=hidden_size,
+    output_size=3 * hidden_size,
+    bias=True,
+    gather_output=False
+)
+
+# Split along input dimension
+self.out_proj = RowParallelLinear(
+    input_size=hidden_size,
+    output_size=hidden_size,
+    bias=True,
+    input_is_parallel=True
+)
+```
+
+### Sequence Parallelism
+
+Sequence parallelism splits sequences across devices:
+
+- **Implementation**: Through `DistributedAttention` and sequence splitting
+- **Use cases**: Long video sequences or high-resolution processing. Used by DiT models.
+
+```python
+# Distributed attention for long sequences
+from fastvideo.v1.attention import DistributedAttention
+
+self.attn = DistributedAttention(
+    num_heads=num_heads,
+    head_size=head_dim,
+    causal=False,
+    supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN)
+)
+```
+
+### Communication Primitives
+Efficient distributed operations via AllGather, AllReduce, and synchronization mechanisms.
+
+Efficient communication primitives minimize distributed overhead:
+
+- **Sequence-Parallel AllGather**: Collects sequence chunks
+- **Tensor-Parallel AllReduce**: Combines partial results
+- **Distributed Synchronization**: Coordinates execution
+
+(design-forwardcontext)=
+## Forward Context Management
+
+### ForwardContext
+
+Defined in `fastvideo/v1/forward_context.py`, `ForwardContext` manages execution-specific state *within* a forward pass, particularly for low-level optimizations. It is accessed via `get_forward_context()`.
+
+- **Attention Metadata**: Configuration for optimized attention kernels (`attn_metadata`)
+- **Profiling Data**: Potential hooks for performance metrics collection
+
+This context-based approach enables:
+- Dynamic optimization based on execution state (e.g., attention backend selection)
+- Step-specific customizations within model components
+
+Usage example:
+
+```python
+with set_forward_context(current_timestep, attn_metadata, fastvideo_args):
+    # During this forward pass, components can access context
+    # through get_forward_context()
+    output = model(inputs)
+```
+
+(design-executor-and-worker-abstractions)=
+## Executor and Worker System
+
+The `fastvideo/v1/worker/` directory contains the distributed execution framework:
+
+### Executor Abstraction
+
+FastVideo implements a flexible execution model for distributed processing:
+
+- **Executor Base Class**: An abstract base class defining the interface for all executors
+- **MultiProcExecutor**: Primary implementation that spawns and manages worker processes
+- **GPU Workers**: Handle actual model execution on individual GPUs
+
+The MultiProcExecutor implementation:
+1. Spawns worker processes for each GPU
+2. Establishes communication channels via pipes
+3. Coordinates distributed operations across workers
+4. Handles graceful startup and shutdown of the process group
+
+Each GPU worker:
+1. Initializes the distributed environment
+2. Builds the pipeline for the specified model
+3. Executes requested operations on its assigned GPU
+4. Manages local resources and communicates results back to the executor
+
+This design allows FastVideo to efficiently utilize multiple GPUs while providing a simple, unified interface for model execution.
+
+(design-platforms)=
+## Platforms
+
+The `fastvideo/v1/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:
+
+### Platform Abstraction
+
+FastVideo's platform abstraction layer enables:
+- **Hardware Detection**: Automatic detection of available hardware
+- **Backend Selection**: Appropriate selection of compute kernels
+- **Memory Management**: Efficient utilization of hardware-specific memory features
+
+The primary components include:
+- **Platform Interface**: Defines the common API for all platform implementations
+- **CUDA Platform**: Optimized implementation for NVIDIA GPUs
+- **Backend Enum**: Used throughout the codebase for feature selection
+
+Usage example:
+
+```python
+from fastvideo.v1.platforms import current_platform, _Backend
+
+# Check hardware capabilities
+if current_platform.supports_backend(_Backend.FLASH_ATTN):
+    # Use FlashAttention implementation
+else:
+    # Fall back to standard implementation
+```
+
+The platform system is designed to be extensible for future hardware targets.
+
+(design-logger)=
+## Logger
+See [PR](https://github.com/hao-ai-lab/FastVideo/pull/356)
+
+*TODO*: (help wanted) Add an environment variable that disables process-aware logging.
+
+## Contributing to FastVideo
+
+If you're a new contributor, here are some common areas to explore:
+
+1. **Adding a new model**: Implement new model types in the appropriate subdirectory of `fastvideo/v1/models/`
+2. **Optimizing performance**: Look at attention implementations or memory management
+3. **Adding a new pipeline**: Create a new pipeline subclass in `fastvideo/v1/pipelines/`
+4. **Hardware support**: Extend the `platforms` module for new hardware targets
+
+When adding code, follow these practices:
+- Use type hints for better code readability
+- Add appropriate docstrings
+- Maintain the separation between model components and execution logic
+- Follow existing patterns for distributed processing

docs/source/generate_examples.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/docs/source/generate_examples.py
+
+import itertools
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+
+ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
+ROOT_DIR_RELATIVE = '../../../..'
+EXAMPLE_DIR = ROOT_DIR / "examples"
+EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
+
+
+def fix_case(text: str) -> str:
+    subs = {
+        "api": "API",
+        "cli": "CLI",
+        "cpu": "CPU",
+        "llm": "LLM",
+        "tpu": "TPU",
+        "aqlm": "AQLM",
+        "gguf": "GGUF",
+        "lora": "LoRA",
+        "rlhf": "RLHF",
+        "vllm": "vLLM",
+        "openai": "OpenAI",
+        "multilora": "MultiLoRA",
+        "mlpspeculator": "MLPSpeculator",
+        r"fp\d+": lambda x: x.group(0).upper(),  # e.g. fp16, fp32
+        r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
+    }
+    for pattern, repl in subs.items():
+        text = re.sub(rf'\b{pattern}\b', repl, text,
+                      flags=re.IGNORECASE)  # type: ignore[call-overload]
+    return text
+
+
+@dataclass
+class Index:
+    """
+    Index class to generate a structured document index.
+
+    Attributes:
+        path (Path): The path save the index file to.
+        title (str): The title of the index.
+        description (str): A brief description of the index.
+        caption (str): An optional caption for the table of contents.
+        maxdepth (int): The maximum depth of the table of contents. Defaults to 1.
+        documents (list[str]): A list of document paths to include in the index. Defaults to an empty list.
+
+    Methods:
+        generate() -> str:
+            Generates the index content as a string in the specified format.
+    """ # noqa: E501
+    path: Path
+    title: str
+    description: str
+    caption: str
+    maxdepth: int = 1
+    documents: list[str] = field(default_factory=list)
+
+    def generate(self) -> str:
+        content = f"# {self.title}\n\n{self.description}\n\n"
+        content += ":::{toctree}\n"
+        content += f":caption: {self.caption}\n:maxdepth: {self.maxdepth}\n"
+        content += "\n".join(self.documents) + "\n:::\n"
+        return content
+
+
+@dataclass
+class Example:
+    """
+    Example class for generating documentation content from a given path.
+
+    Attributes:
+        path (Path): The path to the main directory or file.
+        category (str): The category of the document.
+        main_file (Path): The main file in the directory.
+        other_files (list[Path]): list of other files in the directory.
+        title (str): The title of the document.
+
+    Methods:
+        __post_init__(): Initializes the main_file, other_files, and title attributes.
+        determine_main_file() -> Path: Determines the main file in the given path.
+        determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
+        determine_title() -> str: Determines the title of the document.
+        generate() -> str: Generates the documentation content.
+    """ # noqa: E501
+    path: Path
+    category: str | None = None
+    main_file: Path = field(init=False)
+    other_files: list[Path] = field(init=False)
+    title: str = field(init=False)
+
+    def __post_init__(self):
+        self.main_file = self.determine_main_file()
+        self.other_files = self.determine_other_files()
+        self.title = self.determine_title()
+
+    def determine_main_file(self) -> Path:
+        """
+        Determines the main file in the given path.
+        If the path is a file, it returns the path itself. Otherwise, it searches
+        for Markdown files (*.md) in the directory and returns the first one found.
+        Returns:
+            Path: The main file path, either the original path if it's a file or the first
+            Markdown file found in the directory.
+        Raises:
+            IndexError: If no Markdown files are found in the directory.
+        """ # noqa: E501
+        return self.path if self.path.is_file() else list(
+            self.path.glob("*.md")).pop()
+
+    def determine_other_files(self) -> list[Path]:
+        """
+        Determine other files in the directory excluding the main file.
+
+        This method checks if the given path is a file. If it is, it returns an empty list.
+        Otherwise, it recursively searches through the directory and returns a list of all
+        files that are not the main file.
+
+        Returns:
+            list[Path]: A list of Path objects representing the other files in the directory.
+        """ # noqa: E501
+        if self.path.is_file():
+            return []
+        is_other_file = lambda file: file.is_file() and file != self.main_file
+        return [file for file in self.path.rglob("*")
+                if is_other_file(file)]  # type: ignore[no-untyped-call]
+
+    def determine_title(self) -> str:
+        return fix_case(self.path.stem.replace("_", " ").title())
+
+    def generate(self) -> str:
+        # Convert the path to a relative path from __file__
+        make_relative = lambda path: ROOT_DIR_RELATIVE / path.relative_to(
+            ROOT_DIR)
+
+        content = f"Source <gh-file:{self.path.relative_to(ROOT_DIR)}>.\n\n"
+        include = "include" if self.main_file.suffix == ".md" else \
+            "literalinclude"
+        if include == "literalinclude":
+            content += f"# {self.title}\n\n"
+        content += f":::{{{include}}} {make_relative(self.main_file)}\n"  # type: ignore[no-untyped-call]
+        if include == "literalinclude":
+            content += f":language: {self.main_file.suffix[1:]}\n"
+        content += ":::\n\n"
+
+        if not self.other_files:
+            return content
+
+        content += "## Example materials\n\n"
+        for file in sorted(self.other_files):
+            include = "include" if file.suffix == ".md" else "literalinclude"
+            content += f":::{{admonition}} {file.relative_to(self.path)}\n"
+            content += ":class: dropdown\n\n"
+            content += f":::{{{include}}} {make_relative(file)}\n:::\n"  # type: ignore[no-untyped-call]
+            content += ":::\n\n"
+
+        return content
+
+
+def generate_examples(generate_main_index=False):
+    """
+    Generate example documentation.
+    
+    Args:
+        generate_main_index (bool): Whether to generate the main examples index.
+            If False, only category-specific indices will be generated.
+    """
+    # Create empty indices with dynamic paths
+    main_index_dir = ROOT_DIR / "docs/source/examples"
+    if not main_index_dir.exists():
+        main_index_dir.mkdir(parents=True)
+
+    # Create the main examples index only if requested
+    examples_index = None
+    if generate_main_index:
+        examples_index = Index(
+            path=main_index_dir / "examples_index.md",
+            title="💡 Examples",
+            description=
+            "A collection of examples demonstrating usage of FastVideo.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.",  # noqa: E501
+            caption="Examples",
+            maxdepth=2)
+
+    # Category indices with dynamic paths based on category names
+    category_indices = {
+        "inference":
+        Index(
+            path=ROOT_DIR /
+            "docs/source/inference/examples/examples_inference_index.md",
+            title="🚀 Examples",
+            description=
+            "Inference examples demonstrate how to use FastVideo in an offline setting, where the model is queried for predictions in batches. We recommend starting with <project:basic.md>.",  # noqa: E501
+            caption="Examples",
+        ),
+    }
+
+    # Ensure all category doc directories exist
+    for category, index in category_indices.items():
+        category_dir = index.path.parent
+        if not category_dir.exists():
+            category_dir.mkdir(parents=True)
+
+    examples = []
+    glob_patterns = ["*.py", "*.md", "*.sh"]
+    # Find categorised examples
+    for category in category_indices:
+        print(category)
+        category_dir = EXAMPLE_DIR / category
+        globs = [category_dir.glob(pattern) for pattern in glob_patterns]
+        for path in itertools.chain(*globs):
+            examples.append(Example(path, category))
+        # Find examples in subdirectories
+        for path in category_dir.glob("*/*.md"):
+            examples.append(Example(path.parent, category))
+
+    # Find uncategorised examples only if we're generating a main index
+    if generate_main_index:
+        globs = [EXAMPLE_DIR.glob(pattern) for pattern in glob_patterns]
+        for path in itertools.chain(*globs):
+            examples.append(Example(path))
+        # Find examples in subdirectories
+        for path in EXAMPLE_DIR.glob("*/*.md"):
+            # Skip categorised examples
+            if path.parent.name in category_indices:
+                continue
+            examples.append(Example(path.parent))
+
+    # Create document directories for each category based on category name and generate files
+    for example in sorted(examples, key=lambda e: e.path.stem):
+        print(example)
+
+        # Determine which index to use for this example
+        if example.category is not None and example.category in category_indices:
+            index = category_indices[example.category]
+        elif generate_main_index:
+            assert examples_index is not None
+            index = examples_index  # Default to main index if available
+        else:
+            # Skip examples without a category if no main index
+            print(f"Skipping {example.path} (no category and no main index)")
+            continue
+
+        # Place generated example markdown in the same directory as its index
+        doc_path = index.path.parent / f"{example.path.stem}.md"
+        with open(doc_path, "w+") as f:
+            f.write(example.generate())
+        # Add the example to the index
+        index.documents.append(example.path.stem)
+
+    # Generate the index files for categories
+    for category_index in category_indices.values():
+        if category_index.documents:
+            # Add to main index if it exists
+            if generate_main_index:
+                rel_path = category_index.path.relative_to(
+                    main_index_dir.parent)
+                assert examples_index is not None
+                examples_index.documents.insert(
+                    0,
+                    str(rel_path).replace(".md", ""))
+
+            # Write the category index file
+            with open(category_index.path, "w+") as f:
+                f.write(category_index.generate())
+
+    # Write the main index file if requested
+    if generate_main_index and examples_index:
+        with open(examples_index.path, "w+") as f:
+            f.write(examples_index.generate())

docs/source/getting_started/installation.md

@@ -0,0 +1,120 @@
+(fastvideo-installation)=
+
+# 🔧 Installation
+
+FastVideo currently only supports Linux and NVIDIA CUDA GPUs.
+
+## Requirements
+
+- **OS: Linux**
+- **Python: 3.10-3.12**
+- **CUDA 12.4**
+- **At least 1 NVIDIA GPU**
+
+## Set up using Python
+### Create a new Python environment
+
+#### Conda
+You can create a new python environment using [Conda](https://docs.conda.io/projects/conda/en/stable/user-guide/getting-started.html)
+##### 1. Install Miniconda (if not already installed)
+
+```bash
+wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+bash Miniconda3-latest-Linux-x86_64.sh
+source ~/.bashrc
+```
+
+##### 2. Create and activate a Conda environment for FastVideo
+
+```bash
+# (Recommended) Create a new conda environment.
+conda create -n fastvideo python=3.12 -y
+conda activate fastvideo
+```
+
+:::{note}
+[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
+:::
+
+#### uv
+
+:::{tip}
+We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
+:::
+
+Or you can create a new Python environment using [uv](https://docs.astral.sh/uv/), a very fast Python environment manager. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following command:
+
+```console
+# (Recommended) Create a new uv environment. Use `--seed` to install `pip` and `setuptools` in the environment.
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+```
+
+### Installation
+
+```bash
+pip install fastvideo
+
+# or if you are using uv
+uv pip install fastvideo
+```
+
+Also optionally install flash-attn:
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+### Installation from Source
+
+#### 1. Clone the FastVideo repository
+
+```bash
+git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
+```
+
+#### 2. Install FastVideo
+
+Basic installation:
+
+```bash
+pip install -e .
+
+# or if you are using uv
+uv pip install -e .
+```
+
+### Optional Dependencies
+
+#### Flash Attention
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+## Set up using Docker
+We also have prebuilt docker images with FastVideo dependencies pre-installed:
+[Docker Images](#docker)
+
+## Development Environment Setup
+
+If you're planning to contribute to FastVideo please see the following page:
+[Contributor Guide](#developer-overview)
+
+## Hardware Requirements
+
+### For Basic Inference
+- NVIDIA GPU with CUDA 12.4 support
+
+### For Lora Finetuning
+- 40GB GPU memory each for 2 GPUs with lora
+- 30GB GPU memory each for 2 GPUs with CPU offload and lora
+
+### For Full Finetuning/Distillation
+- Multiple high-memory GPUs recommended (e.g., H100)
+
+## Troubleshooting
+
+If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).
+
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.

docs/source/getting_started/v1_api.md

@@ -0,0 +1,83 @@
+# V1 API
+
+FastVideo's V1 API provides a streamlined interface for video generation tasks with powerful customization options. This page documents the primary components of the API.
+
+## Video Generator
+
+This class will be the primary Python API for generating videos and images.
+
+```{autodoc2-summary}
+    fastvideo.VideoGenerator
+```
+
+`````{py:class} VideoGenerator(fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs, executor_class: type[fastvideo.v1.worker.executor.Executor], log_stats: bool)
+:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator
+
+```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+`VideoGenerator.from_pretrained()` should be the primary way of creating a new video generator.
+
+````{py:method} from_pretrained(model_path: str, device: typing.Optional[str] = None, torch_dtype: typing.Optional[torch.dtype] = None, pipeline_config: typing.Optional[typing.Union[str | fastvideo.v1.configs.pipelines.PipelineConfig]] = None, **kwargs) -> fastvideo.v1.entrypoints.video_generator.VideoGenerator
+:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
+:classmethod:
+
+```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+
+## Configuring FastVideo
+
+The follow two classes `PipelineConfig` and `SamplingParam` are used to configure initialization and sampling parameters, respectively.
+
+### PipelineConfig
+```{autodoc2-summary}
+    fastvideo.PipelineConfig
+```
+
+`````{py:class} PipelineConfig
+:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig
+
+```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.pipelines.base.PipelineConfig
+:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+:classmethod:
+
+```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+
+````{py:method} dump_to_json(file_path: str)
+:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+
+```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+
+### SamplingParam
+
+```{autodoc2-summary}
+    fastvideo.SamplingParam
+```
+
+`````{py:class} SamplingParam
+:canonical: fastvideo.v1.configs.sample.base.SamplingParam
+
+```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.sample.base.SamplingParam
+:canonical: fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+:classmethod:
+
+```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+:parser: docs.source.autodoc2_docstring_parser
+```

docs/source/index.md

@@ -0,0 +1,119 @@
+# Welcome to FastVideo
+
+:::{figure} ../../assets/logo.jpg
+:align: center
+:alt: FastVideo
+:class: no-scaled-link
+:width: 60%
+:::
+
+:::{raw} html
+<p style="text-align:center">
+<strong>FastVideo is a unified framework for accelerated video generation.
+</strong>
+</p>
+
+<p style="text-align:center">
+<script async defer src="https://buttons.github.io/buttons.js"></script>
+<a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/" data-show-count="true" data-size="large" aria-label="Star">Star</a>
+<a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/subscription" data-icon="octicon-eye" data-size="large" aria-label="Watch">Watch</a>
+<a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/fork" data-icon="octicon-repo-forked" data-size="large" aria-label="Fork">Fork</a>
+</p>
+:::
+
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
+
+<div style="text-align: center;">
+  <img src=_static/images/perf.png width="100%"/>
+</div>
+
+## Key Features
+
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
+- Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
+- Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
+
+## Documentation
+
+% How to start using FastVideo?
+
+:::{toctree}
+:caption: Getting Started
+:maxdepth: 1
+
+getting_started/installation
+<!-- getting_started/v1_api -->
+:::
+
+:::{toctree}
+:caption: Inference
+:maxdepth: 1
+
+inference/inference_quick_start
+inference/configuration
+inference/optimizations
+inference/support_matrix
+inference/examples/examples_inference_index
+inference/cli
+inference/add_pipeline
+inference/v0_inference
+:::
+
+:::{toctree}
+:caption: Training
+:maxdepth: 1
+
+training/data_preprocess
+training/distillation
+training/finetune
+:::
+
+% What is STA Kernel?
+
+:::{toctree}
+:caption: Sliding Tile Attention
+:maxdepth: 1
+
+sliding_tile_attention/installation
+sliding_tile_attention/demo
+:::
+
+:::{toctree}
+:caption: Design
+:maxdepth: 1
+design/overview
+:::
+
+:::{toctree}
+:caption: Developer Guide
+:maxdepth: 2
+
+contributing/overview
+contributing/developer_env/index
+:::
+
+:::{toctree}
+:caption: API Reference
+:maxdepth: 2
+
+<!-- api/summary -->
+api/fastvideo/fastvideo
+:::
+
+## Indices and tables
+
+- {ref}`genindex`
+- {ref}`modindex`

docs/source/inference/add_pipeline.md

@@ -0,0 +1,316 @@
+(add-pipeline)=
+
+# 🏗️ Adding a New Pipeline
+
+This guide explains how to implement a custom diffusion pipeline in FastVideo, leveraging the framework's modular architecture for high-performance video generation.
+
+## Implementation Process Overview
+
+1. **Port Required Modules** - Identify and implement necessary model components
+2. **Create Directory Structure** - Set up pipeline files and folders
+3. **Implement Pipeline Class** - Build the pipeline using existing or custom stages
+4. **Register Your Pipeline** - Make it discoverable by the framework
+5. **Configure Your Pipeline** - (Coming soon)
+
+Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
+
+## Step 1: Pipeline Modules
+
+### Identifying Required Modules
+
+FastVideo uses the Hugging Face Diffusers format for model organization:
+
+1. Examine the `model_index.json` in the HF model repository:
+
+```json
+{
+    "_class_name": "WanImageToVideoPipeline",
+    "_diffusers_version": "0.33.0.dev0",
+    "image_encoder": ["transformers", "CLIPVisionModelWithProjection"],
+    "image_processor": ["transformers", "CLIPImageProcessor"],
+    "scheduler": ["diffusers", "UniPCMultistepScheduler"],
+    "text_encoder": ["transformers", "UMT5EncoderModel"],
+    "tokenizer": ["transformers", "T5TokenizerFast"],
+    "transformer": ["diffusers", "WanTransformer3DModel"],
+    "vae": ["diffusers", "AutoencoderKLWan"]
+}
+```
+
+1. For each component:
+   - Note the originating library (`transformers` or `diffusers`)
+   - Identify the class name
+   - Check if it's already available in FastVideo
+
+2. Review config files in each component's directory for architecture details
+
+### Implementing Modules
+
+Place new modules in the appropriate directories:
+- Encoders: `fastvideo/v1/models/encoders/`
+- VAEs: `fastvideo/v1/models/vaes/`
+- Transformer models: `fastvideo/v1/models/dits/`
+- Schedulers: `fastvideo/v1/models/schedulers/`
+
+### Adapting Model Layers
+
+#### Layer Replacements
+Replace standard PyTorch layers with FastVideo optimized versions:
+- nn.LayerNorm → fastvideo.v1.layers.layernorm.RMSNorm
+- Embedding layers → fastvideo.v1.layers.vocab_parallel_embedding modules
+- Activation functions → versions from fastvideo.v1.layers.activation
+
+#### Distributed Linear Layers
+Use appropriate parallel layers for distribution:
+
+```python
+# Output dimension parallelism
+from fastvideo.v1.layers.linear import ColumnParallelLinear
+self.q_proj = ColumnParallelLinear(
+    input_size=hidden_size,
+    output_size=head_size * num_heads,
+    bias=bias,
+    gather_output=False
+)
+
+# Fused QKV projection
+from fastvideo.v1.layers.linear import QKVParallelLinear
+self.qkv_proj = QKVParallelLinear(
+    hidden_size=hidden_size,
+    head_size=attention_head_dim,
+    total_num_heads=num_attention_heads,
+    bias=True
+)
+
+# Input dimension parallelism
+from fastvideo.v1.layers.linear import RowParallelLinear
+self.out_proj = RowParallelLinear(
+    input_size=head_size * num_heads,
+    output_size=hidden_size,
+    bias=bias,
+    input_is_parallel=True
+)
+```
+
+### Attention Layers
+Replace standard attention with FastVideo's optimized attention:
+
+```python
+# Local attention patterns
+from fastvideo.v1.attention import LocalAttention
+from fastvideo.v1.attention.backends.abstract import _Backend
+self.attn = LocalAttention(
+    num_heads=num_heads,
+    head_size=head_dim,
+    dropout_rate=0.0,
+    softmax_scale=None,
+    causal=False,
+    supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+)
+
+# Distributed attention for long sequences
+from fastvideo.v1.attention import DistributedAttention
+self.attn = DistributedAttention(
+    num_heads=num_heads,
+    head_size=head_dim,
+    dropout_rate=0.0,
+    softmax_scale=None,
+    causal=False,
+    supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+)
+```
+
+#### Define supported backend selection
+
+```python
+   _supported_attention_backends = (_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+```
+
+### Registering Models
+
+Register implemented modules in the model registry:
+
+```python
+# In fastvideo/v1/models/registry.py
+_TEXT_TO_VIDEO_DIT_MODELS = {
+    "YourTransformerModel": ("dits", "yourmodule", "YourTransformerClass"),
+}
+
+_VAE_MODELS = {
+    "YourVAEModel": ("vaes", "yourvae", "YourVAEClass"),
+}
+```
+
+## Step 2: Directory Structure
+
+Create a new directory for your pipeline:
+
+```
+fastvideo/v1/pipelines/
+├── your_pipeline/
+│   ├── __init__.py
+│   └── your_pipeline.py
+```
+
+## Step 3: Implement Pipeline Class
+
+Pipelines are composed of stages, each handling a specific part of the diffusion process:
+
+- **InputValidationStage**: Validates input parameters
+- **Text Encoding Stages**: Handle text encoding (CLIP/Llama/T5)
+- **CLIPImageEncodingStage**: Processes image inputs
+- **TimestepPreparationStage**: Prepares diffusion timesteps
+- **LatentPreparationStage**: Manages latent representations
+- **ConditioningStage**: Processes conditioning inputs
+- **DenoisingStage**: Performs denoising diffusion
+- **DecodingStage**: Converts latents to pixels
+
+### Creating Your Pipeline
+
+```python
+from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.v1.pipelines.stages import (
+    InputValidationStage, CLIPTextEncodingStage, TimestepPreparationStage,
+    LatentPreparationStage, DenoisingStage, DecodingStage
+)
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+import torch
+
+class MyCustomPipeline(ComposedPipelineBase):
+    """Custom diffusion pipeline implementation."""
+    
+    # Define required model components from model_index.json
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+    
+    @property
+    def required_config_modules(self) -> List[str]:
+        return self._required_config_modules
+        
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        """Initialize pipeline-specific components."""
+        pass
+        
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+        self.add_stage(
+            stage_name="input_validation_stage",
+            stage=InputValidationStage()
+        )
+        
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=CLIPTextEncodingStage(
+                text_encoder=self.get_module("text_encoder"),
+                tokenizer=self.get_module("tokenizer")
+            )
+        )
+        
+        self.add_stage(
+            stage_name="timestep_preparation_stage",
+            stage=TimestepPreparationStage(
+                scheduler=self.get_module("scheduler")
+            )
+        )
+        
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=LatentPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                vae=self.get_module("vae")
+            )
+        )
+        
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler")
+            )
+        )
+        
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=DecodingStage(
+                vae=self.get_module("vae")
+            )
+        )
+    
+# Register the pipeline class
+EntryClass = MyCustomPipeline
+```
+
+### Creating Custom Stages (Optional)
+
+If existing stages don't meet your needs, create custom ones:
+
+```python
+from fastvideo.v1.pipelines.stages.base import PipelineStage
+
+class MyCustomStage(PipelineStage):
+    """Custom processing stage for the pipeline."""
+    
+    def __init__(self, custom_module, other_param=None):
+        super().__init__()
+        self.custom_module = custom_module
+        self.other_param = other_param
+        
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        # Access input data
+        input_data = batch.some_attribute
+        
+        # Validate inputs
+        if input_data is None:
+            raise ValueError("Required input is missing")
+            
+        # Process with your module
+        result = self.custom_module(input_data)
+        
+        # Update batch with results
+        batch.some_output = result
+        
+        return batch
+```
+
+Add your custom stage to the pipeline:
+
+```python
+self.add_stage(
+    stage_name="my_custom_stage",
+    stage=MyCustomStage(
+        custom_module=self.get_module("custom_module"),
+        other_param="some_value"
+    )
+)
+```
+
+#### Stage Design Principles
+
+1. **Single Responsibility**: Focus on one specific task
+2. **Functional Pattern**: Receive and return a `ForwardBatch` object
+3. **Dependency Injection**: Pass dependencies through constructor
+4. **Input Validation**: Validate inputs for clear error messages
+
+## Step 4: Register Your Pipeline
+
+Define `EntryClass` at the end of your pipeline file:
+
+```python
+# Single pipeline class
+EntryClass = MyCustomPipeline
+
+# Or multiple pipeline classes
+EntryClass = [MyCustomPipeline, MyOtherPipeline]
+```
+
+The registry will automatically:
+1. Scan all packages under `fastvideo/v1/pipelines/`
+2. Look for `EntryClass` variables
+3. Register pipelines using their class names as identifiers
+
+## Best Practices
+
+- **Reuse Existing Components**: Leverage built-in stages and modules
+- **Follow Module Organization**: Place new modules in appropriate directories
+- **Match Model Patterns**: Follow existing code patterns and conventions

docs/source/inference/cli.md

@@ -0,0 +1,151 @@
+# FastVideo CLI Inference
+
+The FastVideo CLI provides a quick way to access the FastVideo inference pipeline for video generation. For more advanced usage,
+see the Python interface [here](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html).
+
+## Basic Usage
+
+The basic command to generate a video is:
+
+```bash
+fastvideo generate --model-path {MODEL_PATH} --prompt {PROMPT}
+```
+
+### Required Parameters
+
+- `--model-path {MODEL_PATH}`: Path to the model or model ID
+- `--prompt {PROMPT}`: Text description for the video you want to generate
+
+## Common Arguments
+
+To see all the options, you can use the `--help` flag:
+
+```bash
+fastvideo generate --help
+```
+
+### Hardware Configuration
+
+- `--num-gpus {NUM_GPUS}`: Number of GPUs to use
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (only for the encoder, should not be larger than 1 if text encoder offload is enabled, as layerwise offload + prefetch is faster)
+- `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)
+
+#### Video Configuration
+
+- `--height {HEIGHT}`: Height of the generated video
+- `--width {WIDTH}`: Width of the generated video
+- `--num-frames {NUM_FRAMES}`: Number of frames to generate
+- `--fps {FPS}`: Frames per second for the saved video
+
+#### Generation Parameters
+
+- `--num-inference-steps {STEPS}`: Number of denoising steps
+- `--negative-prompt {PROMPT}`: Negative prompt to guide generation away from certain concepts
+- `--seed {SEED}`: Random seed for reproducible generation
+
+#### Output Options
+
+- `--output-path {PATH}`: Directory to save the generated video
+- `--save-video`: Whether to save the video to disk
+- `--return-frames`: Whether to return the raw frames
+
+## Using Configuration Files
+
+Instead of specifying all parameters on the command line, you can use a configuration file:
+
+```bash
+fastvideo generate --config {CONFIG_FILE_PATH}
+```
+
+The config file should be in JSON or YAML format with the same parameter names as the CLI options. Command-line arguments will take precedence over settings in the configuration file, allowing you to override specific values while keeping the rest from the config file.
+
+Example configuration file (config.json):
+
+```json
+{
+    "model_path": "FastVideo/FastHunyuan-diffusers",
+    "prompt": "A beautiful woman in a red dress walking down a street",
+    "output_path": "outputs/",
+    "num_gpus": 2,
+    "sp_size": 2,
+    "tp_size": 1,
+    "num_frames": 45,
+    "height": 720,
+    "width": 1280,
+    "num_inference_steps": 6,
+    "seed": 1024,
+    "fps": 24,
+    "precision": "bf16",
+    "vae_precision": "fp16",
+    "vae_tiling": true,
+    "vae_sp": true,
+    "vae_config": {
+        "load_encoder": false,
+        "load_decoder": true,
+        "tile_sample_min_height": 256,
+        "tile_sample_min_width": 256
+    },
+    "text_encoder_precisions": [
+        "fp16",
+        "fp16"
+    ],
+    "mask_strategy_file_path": null,
+    "enable_torch_compile": false
+}
+```
+
+Or using YAML format (config.yaml):
+
+```yaml
+model_path: "FastVideo/FastHunyuan-diffusers"
+prompt: "A beautiful woman in a red dress walking down a street"
+output_path: "outputs/"
+num_gpus: 2
+sp_size: 2
+tp_size: 1
+num_frames: 45
+height: 720
+width: 1280
+num_inference_steps: 6
+seed: 1024
+fps: 24
+precision: "bf16"
+vae_precision: "fp16"
+vae_tiling: true
+vae_sp: true
+vae_config:
+  load_encoder: false
+  load_decoder: true
+  tile_sample_min_height: 256
+  tile_sample_min_width: 256
+text_encoder_precisions:
+  - "fp16"
+  - "fp16"
+mask_strategy_file_path: null
+enable_torch_compile: false
+```
+
+## Examples
+
+Generating a simple video:
+
+```bash
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A cat playing with a ball of yarn" --num-frames 45 --height 720 --width 1280 --num-inference-steps 6 --seed 1024 --output-path outputs/
+```
+
+Using a negative prompt to avoid certain elements:
+
+```bash
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A beautiful forest landscape" --negative-prompt "people, buildings, roads"
+```
+
+Combining command line arguments and a configuration file:
+
+```bash
+fastvideo generate --config config.json --prompt "A capybara lounging in a hammock"
+```
+
+## Troubleshooting
+
+- If you encounter CUDA out-of-memory errors, try reducing the video dimensions or number of frames, or the number of inference steps.
+- For reproducible results, set the same seed value between runs.

docs/source/inference/configuration.md

@@ -0,0 +1,77 @@
+(inference-configuration)=
+# Configuration
+
+## Multi-GPU Setup
+
+FastVideo automatically distributes the generation process when multiple GPUs are specified:
+
+```python
+# Will use 4 GPUs in parallel for faster generation
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=4,
+)
+```
+
+## Customizing Generation
+
+- `PipelineConfig`: Initialization time parameters
+- `SamplingParam`: Generation time parameters
+
+You can customize various parameters when generating videos using the `PipelineConfig` and `SamplingParam` class:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam, PipelineConfig
+
+def main():
+    model_name = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    config = PipelineConfig.from_pretrained(model_name)
+    config.vae_precision = "fp16"
+    config.use_cpu_offload = True
+
+    # Create the generator
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        num_gpus=1,
+        pipeline_config=config
+    )
+
+    # Create and customize sampling parameters
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    
+    # How many frames to generate
+    sampling_param.num_frames = 45
+    
+    # Video resolution (width, height)
+    sampling_param.width = 1024
+    sampling_param.height = 576
+    
+    # How many steps we denoise the video (higher = better quality, slower generation)
+    sampling_param.num_inference_steps = 30
+    
+    # How strongly the video conforms to the prompt (higher = more faithful to prompt)
+    sampling_param.guidance_scale = 7.5
+    
+    # Random seed for reproducibility
+    sampling_param.seed = 42  # Optional, leave unset for random results
+
+    # Generate video with custom parameters
+    prompt = "A beautiful sunset over a calm ocean, with gentle waves."
+    video = generator.generate_video(
+        prompt, 
+        sampling_param=sampling_param, 
+        output_path="my_videos/",  # Controls where videos are saved
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        save_video=True
+    )
+
+    # If return_frames=True, video contains the generated frames as a NumPy array
+    print(f"Generated {len(video)} frames")
+
+if __name__ == '__main__':
+    main()
+```
+
+## Performance Optimization
+
+For configuring optimizations, please see our [optimizations guide](#inference-optimizations)

docs/source/inference/inference_quick_start.md

@@ -0,0 +1,124 @@
+# Inference Quick Start
+
+This page contains step-by-step instructions to get you quickly started with video generation using FastVideo.
+
+## Requirements
+- **OS**: Linux (Tested on Ubuntu 22.04+)
+- **Python**: 3.10-3.12
+- **CUDA**: 12.4
+- **GPU**: At least one NVIDIA GPU
+
+## Installation
+
+We recommend using an environment manager such as `Conda` to create a clean environment:
+
+```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
+```
+
+For advanced installation options, see the [Installation Guide](installation.md).
+
+## Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+
+```python
+from fastvideo import VideoGenerator
+
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:
+
+```bash
+python example.py
+```
+
+The generated video will be saved in the current directory under `my_videos/`  
+
+More inference example scripts can be found in `scripts/inference/`
+## Available Models
+
+Please see the [support matrix](#support-matrix) for the list of supported models and their available optimizations.
+
+## Image-to-Video Generation
+
+You can generate a video starting from an initial image:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam
+
+def main():
+    # Create the generator
+    model_name = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+    generator = VideoGenerator.from_pretrained(model_name, num_gpus=1)
+
+    # Set up parameters with an initial image
+    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    sampling_param.num_frames = 107
+
+    # Generate video based on the image
+    prompt = "A photograph coming to life with gentle movement"
+    generator.generate_video(prompt, sampling_param=sampling_param,
+                             output_path="my_videos/",
+                             save_video=True)
+
+if __name__ == '__main__':
+    main()
+```
+
+## Troubleshooting
+
+Common issues and their solutions:
+
+### Out of Memory Errors
+If you encounter CUDA out of memory errors:
+- Reduce `num_frames` or video resolution
+- Enable memory optimization with `enable_model_cpu_offload`
+- Try a smaller model or use distilled versions
+- Use `num_gpus` > 1 if multiple GPUs are available
+
+### Slow Generation
+To speed up generation:
+- Reduce `num_inference_steps` (20-30 is usually sufficient)
+- Use half precision (`fp16`) for the VAE
+- Use multiple GPUs if available
+
+### Unexpected Results
+If the generated video doesn't match your prompt:
+- Try increasing `guidance_scale` (7.0-9.0 works well)
+- Make your prompt more detailed and specific
+- Experiment with different random seeds
+- Try a different model
+
+## Next Steps
+
+- Learn about [Advanced Inference Configurations](#inference-configuration)
+- Learn about using [Optimizations](#inference-optimizations)
+- See [Examples](../examples/examples_inference_index.md) for more usage scenarios
+- Join our [Community Discord](https://discord.gg/JA7cksDz86).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).

docs/source/inference/optimizations.md

@@ -0,0 +1,148 @@
+(inference-optimizations)=
+# Optimizations
+
+This page describes the various options for speeding up generation times in FastVideo.
+
+## Table of Contents
+- Optimized Attention Backends
+  - [Flash Attention](#optimizations-flash)
+  - [Sliding Tile Attention](#optimizations-sta)
+  - [Sage Attention](#optimizations-sage)
+
+- Caching Techniques
+  - [TeaCache](#optimizations-teacache)
+
+(optimizations-backends)=
+## Attention Backends
+
+### Available Backends
+- Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`
+- Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN`
+- Sliding Tile Attention: `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
+- Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
+
+### Configuring Backends
+
+There are two ways to configure the attention backend in FastVideo.
+
+#### 1. In Python
+In python, set the `FASTVIDEO_ATTENTION_BACKEND` environment variable before instantiating `VideoGenerator` like this:
+
+```python
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
+```
+
+#### 2. In CLI
+You can also set the environment variable on the command line:
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
+```
+
+(optimizations-flash)=
+### Flash Attention
+
+**`FLASH_ATTN`**
+
+We recommend always installing [Flash Attention 2](https://github.com/Dao-AILab/flash-attention):
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+And if using a Hopper+ GPU (ie H100), installing [Flash Attention 3](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#flashattention-3-beta-release) by compiling it from source (takes about 10 minutes for me):
+
+```bash
+git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention
+
+cd hopper
+pip install ninja 
+python setup.py install
+```
+
+:::{note}
+FastVideo will automatically detect and use `FA3` if it is installed when using `FLASH_ATTN` backend.
+:::
+
+(optimizations-sta)=
+### Sliding Tile Attention
+**`SLIDING_TILE_ATTN`**
+
+```bash
+pip install st_attn==0.0.4
+```
+
+Please see [this page](#sta-installation) for more installation instructions.
+
+(optimizations-sage)=
+### Sage Attention
+**`SAGE_ATTN`**
+
+To use [SageAttention](https://github.com/thu-ml/SageAttention) 2.1.1, please compile from source:
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd sageattention 
+python setup.py install  # or pip install -e .
+```
+
+(optimizations-teacache)=
+## Teacache
+TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo.
+
+### What is TeaCache?
+
+See the official [TeaCache](https://github.com/ali-vilab/TeaCache) repo and their paper for more details.
+
+### How to Enable TeaCache
+
+Enabling TeaCache is straightforward - simply add the `enable_teacache=True` parameter to your `generate_video()` call:
+
+```python
+# ... previous code
+generator.generate_video(
+    prompt="Your prompt here",
+    sampling_param=params,
+    enable_teacache=True
+)
+# more code ...
+```
+
+### Complete Example
+
+At the bottom is a complete example of using TeaCache for faster video generation. You can run it using the following command:
+
+```bash
+python examples/inference/optimizations/teacache_example.py
+```
+
+### Advanced Configuration
+
+While TeaCache works well with default settings, you can fine-tune its behavior by adjusting the threshold value:
+
+1. Lower threshold values (e.g., 0.1) will result in more skipped calculations and faster generation with slightly more potential for quality degradation
+2. Higher threshold values (e.g., 0.15-0.23) will skip fewer calculations but maintain quality closer to the original
+
+Note that the optimal threshold depends on your specific model and content.
+
+## Benchmarking different optimizations
+
+To benchmark the performance improvement, try generating the same video with and without TeaCache enabled and compare the generation times:
+
+```python
+# Without TeaCache
+start_time = time.perf_counter()
+generator.generate_video(prompt="Your prompt", enable_teacache=False)
+standard_time = time.perf_counter() - start_time
+
+# With TeaCache
+start_time = time.perf_counter()
+generator.generate_video(prompt="Your prompt", enable_teacache=True)
+teacache_time = time.perf_counter() - start_time
+
+print(f"Standard generation: {standard_time:.2f} seconds")
+print(f"TeaCache generation: {teacache_time:.2f} seconds")
+print(f"Speedup: {standard_time/teacache_time:.2f}x")
+```
+
+Note: If you want to benchmark different attention backends, you'll need to reinstantiate `VideoGenerator`.

docs/source/inference/support_matrix.md

@@ -0,0 +1,92 @@
+(support-matrix)=
+# Compatibility Matrix
+The table below shows every supported model and optimizations supported for them.
+
+The symbols used have the following meanings:
+
+- ✅ = Full compatibility
+- ❌ = No compatibility
+
+## Models x Optimization
+The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods and FastVideo will use the optimal default parameters when initializing and generating videos.
+
+:::{raw} html
+<style>
+  /* Make smaller to try to improve readability  */
+  td {
+    font-size: 0.9rem;
+    text-align: center;
+  }
+
+  th {
+    text-align: center;
+    font-size: 0.9rem;
+  }
+</style>
+:::
+
+:::{list-table}
+:header-rows: 1
+:stub-columns: 3
+:widths: auto
+:class: vertical-table-header
+
+- * Model Name
+  * HuggingFace Model ID
+  * Resolutions
+  * TeaCache
+  * Sliding Tile Attn
+  * Sage Attn
+- * HunyuanVideo
+  * `hunyuanvideo-community/HunyuanVideo`
+  * 720px1280p<br>544px960p
+  * ❌
+  * ✅
+  * ✅
+- * FastHunyuan
+  * `FastVideo/FastHunyuan-diffusers`
+  * 720px1280p<br>544px960p
+  * ❌
+  * ✅
+  * ✅
+- * Wan T2V 1.3B
+  * `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`
+  * 480P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan T2V 14B
+  * `Wan-AI/Wan2.1-T2V-14B-Diffusers`
+  * 480P, 720P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan I2V 480P
+  * `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`
+  * 480P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan I2V 720P
+  * `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers`
+  * 720P
+  * ✅
+  * ✅*
+  * ✅
+- * StepVideo T2V
+  * `FastVideo/stepvideo-t2v-diffusers`
+  * 768px768px204f<br>544px992px204f<br>544px992px136f
+  * ❌
+  * ❌
+  * ✅
+:::
+
+**Note**: there are some known quality issues with Wan2.1 + Sliding Tile Attn. We are working on fixing this issue.
+
+## Special requirements
+
+### StepVideo T2V
+- The self-attention in text-encoder (step_llm) only supports CUDA capabilities sm_80 sm_86 and sm_90
+
+### Sliding Tile Attention
+- Currently only Hopper GPUs (H100s) are supported.

docs/source/inference/v0_inference.md

@@ -0,0 +1,74 @@
+(v0-inference)=
+
+# [Deprecated] V0 Inference
+The following commands and APIs are deprecated but still supported until V1's API can completely replace all the features in this page.
+
+## Inference StepVideo with Sliding Tile Attention
+First, download the model:
+
+```
+python scripts/huggingface/download_hf.py --repo_id=stepfun-ai/stepvideo-t2v --local_dir=data/stepvideo-t2v --repo_type=model
+```
+
+Use the following scripts to run inference for StepVideo. When using STA for inference, the generated videos will have dimensions of 204×768×768 (currently, this is the only supported shape).
+
+```bash
+sh scripts/inference/inference_stepvideo_STA.sh # Inference stepvideo with STA
+sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo
+```
+
+## Inference HunyuanVideo with Sliding Tile Attention
+First, download the model:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model
+```
+
+We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
+
+```bash
+sh scripts/inference/inference_hunyuan_STA.sh
+```
+
+## Video Demos using STA + Teacache
+Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100.
+
+## Inference FastHunyuan on single RTX4090
+We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
+
+```bash
+# Download the model weight
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model
+# CLI inference
+bash scripts/inference/inference_hunyuan_hf_quantization.sh
+```
+
+For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):
+
+| Configuration                  | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time |
+|--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------|
+| BF16 + Pipeline CPU Offload    | 23.883G                   | 33.744G                                    | 81s            | 121.5s          |
+| INT8 + Pipeline CPU Offload    | 13.911G                   | 27.979G                                    | 88s            | 116.7s          |
+| NF4 + Pipeline CPU Offload     | 9.453G                    | 19.26G                                     | 78s            | 114.5s          |
+
+For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section:
+
+## FastHunyuan
+
+```bash
+# Download the model weight
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model
+# CLI inference
+bash scripts/inference/inference_hunyuan.sh
+```
+
+You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).
+
+## FastMochi
+
+```bash
+# Download the model weight
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffusers --local_dir=data/FastMochi-diffusers --repo_type=model
+# CLI inference
+bash scripts/inference/inference_mochi_sp.sh
+```

docs/source/sliding_tile_attention/demo.md

@@ -0,0 +1,17 @@
+(sta-demo)=
+
+# 🔍 Demo
+This is is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
+
+<div style="text-align: center;">
+  <video controls width="800">
+    <source src="https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747" type="video/mp4">
+    Your browser does not support the video tag.
+  </video>
+</div>
+
+You can run STA using the following command:
+
+```bash
+bash scripts/inference/v1_inference_wan_STA.sh
+```

docs/source/sliding_tile_attention/installation.md

@@ -0,0 +1,55 @@
+(sta-installation)=
+
+# 🔧 Installation
+You can install the Sliding Tile Attention package using
+
+```
+pip install st_attn==0.0.4
+``` 
+
+# Building from Source
+We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
+First, install C++20 for ThunderKittens:
+
+```bash
+cd csrc/sliding_tile_attention/
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+
+Install STA:
+
+```bash
+export CUDA_HOME=/usr/local/cuda-12.4
+export PATH=${CUDA_HOME}/bin:${PATH} 
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+git submodule update --init --recursive
+python setup.py install
+```
+
+# 🧪 Test
+
+```bash
+python test/test_sta.py
+```
+
+# 📋 Usage
+
+```python
+from st_attn import sliding_tile_attention
+# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
+# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256
+# a tile is a cube of size (6, 8, 8)
+# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video.
+# text_length: int ranging from 0 to 256
+# If your attention contains text token (Hunyuan)
+out = sliding_tile_attention(q, k, v, window_size, text_length)
+# If your attention does not contain text token (StepVideo)
+out = sliding_tile_attention(q, k, v, window_size, 0, False)
+
+```