V1 wan rebased (#335)

Co-authored-by: JerryZhou54 <zhouw.jerry2017@outlook.com>

.github/scripts/runpod_api.py

@@ -0,0 +1,240 @@
+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,
+        default='runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04',
+        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"""
+    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": args.image,
+        "allowedCudaVersions": ["12.4"]
+    }
+
+    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 = 6
+    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(10)
+        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)
+
+    setup_steps = [
+        "cd /workspace",
+        "wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh",
+        "bash Miniconda3-latest-Linux-x86_64.sh -b -p $HOME/miniconda3",
+        "source $HOME/miniconda3/bin/activate",
+        "conda create --name venv python=3.10.0 -y", "conda activate venv",
+        "mkdir -p /workspace/repo",
+        "tar -xzf /tmp/repo.tar.gz --no-same-owner -C /workspace/",
+        f"cd /workspace/{repo_name}", 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/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,77 @@
+# 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"
+  pull_request:
+    branches: 
+      - main
+    types: [opened, ready_for_review, synchronize, reopened]
+    paths:
+      - "docs/**/*.md"
+
+  # 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:
+  # Build job
+  build:
+    runs-on: ubuntu-latest
+    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

@@ -15,7 +15,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 +23,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
@@ -48,10 +48,10 @@ jobs:
 
     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,291 @@
+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"
+  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
+
+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 }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dorny/paths-filter@v3
+        id: filter
+        with:
+          filters: |
+            encoder-test:
+              - 'fastvideo/v1/models/encoders/**'
+              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/tests/encoders/**'
+            vae-test:
+              - 'fastvideo/v1/models/vaes/**'
+              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/tests/vaes/**'
+            transformer-test:
+              - 'fastvideo/v1/models/dits/**'
+              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/tests/transformers/**'
+
+  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')
+    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.10"
+
+      - 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: "encoder-test"
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: 30
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "NVIDIA A40"
+          --gpu-count 1
+          --volume-size 100
+          --test-command "pip install -e .[test] &&
+          pip install flash-attn==2.7.0.post2 --no-build-isolation &&
+          pytest ./fastvideo/v1/tests/encoders -s"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: "encoder-test"
+        run: python .github/scripts/runpod_cleanup.py
+
+  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')
+    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.10"
+
+      - 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: "vae-test"
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: 30
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "NVIDIA A40"
+          --gpu-count 1
+          --volume-size 100
+          --test-command "pip install -e .[test] &&
+          pip install flash-attn==2.7.0.post2 --no-build-isolation &&
+          pytest ./fastvideo/v1/tests/vaes -s"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: "vae-test"
+        run: python .github/scripts/runpod_cleanup.py
+
+  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')
+    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.10"
+
+      - 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: "transformer-test"
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: 30
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "NVIDIA L40S"
+          --gpu-count 1
+          --volume-size 100
+          --test-command "pip install -e .[test] &&
+          pip install flash-attn==2.7.0.post2 --no-build-isolation &&
+          pytest ./fastvideo/v1/tests/transformers -s"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: "transformer-test"
+        run: python .github/scripts/runpod_cleanup.py
+      
+  ssim-test:
+    needs: change-filter
+    if: >-
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
+    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.10"
+
+      - 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: "ssim-test"
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: 45
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "NVIDIA A40"
+          --gpu-count 2
+          --disk-size 200
+          --volume-size 200
+          --test-command "pip install -e .[test] &&
+          pip install flash-attn==2.7.0.post2 --no-build-isolation &&
+          pytest ./fastvideo/v1/tests/ssim -vs"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: "ssim-test"
+        run: python .github/scripts/runpod_cleanup.py
+
+  runpod-cleanup:
+    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
+    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"]' # JSON array of job IDs
+          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.10"
+    - 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/sta-publish.yml

@@ -15,7 +15,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
 
@@ -43,7 +43,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' }}
     runs-on: ${{ matrix.os }}
 
     strategy:
@@ -144,8 +144,8 @@ jobs:
 
   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' }}
     runs-on: ubuntu-22.04
     permissions:
       id-token: write  # Needed for OIDC Trusted Publishing

.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/sliding_tile_attention/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,5 +1,4 @@
 __pycache__
-*.mp4
 .ipynb_checkpoints
 *.pth
 UCF-101/
@@ -11,6 +10,7 @@ wandb/
 *.jpg
 *.safetensors
 *.mp4
+!fastvideo/v1/tests/ssim/reference_videos/**/*.mp4
 *.png
 *.gif
 *.pth

.pre-commit-config.yaml

@@ -0,0 +1,80 @@
+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/.*|
+    .github/workflows/fastvideo-publish.yml|
+    .github/workflows/sta-publish.yml
+  )
+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.4
+  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: 0a0b7a830386ba6a31c2ec8316849ae4d1b8240d # 6.0.0
+#   hooks:
+#   - id: isort
+- repo: https://github.com/jackdewinter/pymarkdown
+  rev: v0.9.29
+  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

@@ -4,19 +4,12 @@
 
 FastVideo is a lightweight framework for accelerating large video diffusion models.
 
-
 <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://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>
 
 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/).
@@ -28,8 +21,6 @@ FastVideo currently offers: (with more to come)
 
 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/).
@@ -37,21 +28,32 @@ Dev in progress and highly experimental.
 - ```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
+## 🔧 Installation from source
 The code is tested on Python 3.10.0, CUDA 12.4 and H100.
+
 ```
-./env_setup.sh fastvideo
+# Clone FastVideo
+git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
+
+# Install FastVideo
+pip install -e .
+
+# Install Flash Attention (optional)
+pip install flash-attn==2.7.0.post2
 ```
+
 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 
+### 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 
+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
@@ -59,44 +61,49 @@ 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 
+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):
 
+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
@@ -108,79 +115,99 @@ python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffuser
 bash scripts/inference/inference_mochi_sp.sh
 ```
 
-
 ## 🎯 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
 ```
+
 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 
+### ⚡ Lora Finetune
 
 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.  
-
+- 30 GB GPU memory each for 2 GPUs with CPU offload and lora.
 
 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!  
+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 
+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 
+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. 
+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.
 
 ## 📑 Development Plan
@@ -205,26 +232,26 @@ We learned and reused code from the following projects: [PCM](https://github.com
 
 We thank MBZUAI and Anyscale for their support throughout this project.
 
-## Citation 
+## Citation
 If you use FastVideo for your research, please cite our paper:
 
 ```bibtex
 @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},
 }
 ```

csrc/sliding_tile_attention/st_attn/__init__.py

@@ -8,7 +8,7 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te
     seq_length = q_all.shape[2]
     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, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels"
         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

docs/Makefile

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

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,16 +1,16 @@
 
-
-
 ## 🧱 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
@@ -33,13 +33,16 @@ path_to_dataset_folder/
 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, 
+
+For video media,
+
 ```
 {
     "path": "1.mp4",
@@ -62,7 +65,9 @@ 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,25 @@
+sphinx==6.2.1
+sphinx-argparse==0.4.0
+sphinx-book-theme==1.0.1
+sphinx-copybutton==0.5.2
+sphinx-design==0.6.1
+sphinx-togglebutton==0.3.2
+myst-parser==3.0.1
+msgspec
+cloudpickle
+
+# packages to install to build the documentation
+cachetools
+pydantic >= 2.8
+-f https://download.pytorch.org/whl/cpu
+torch
+py-cpuinfo
+transformers
+mistral_common >= 1.5.4
+aiohttp
+starlette
+openai # Required by docs/source/serving/openai_compatible_server.md's vllm.entrypoints.openai.cli_args
+fastapi # Required by docs/source/serving/openai_compatible_server.md's vllm.entrypoints.openai.cli_args
+partial-json-parser # Required by docs/source/serving/openai_compatible_server.md's vllm.entrypoints.openai.cli_args
+requests
+zmq

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/conf.py

@@ -0,0 +1,260 @@
+# 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 inspect
+import logging
+import os
+import sys
+from typing import Optional
+
+import requests
+from sphinx.ext import autodoc
+
+logger = logging.getLogger(__name__)
+sys.path.append(os.path.abspath("../.."))
+
+# -- 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",
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "myst_parser",
+    "sphinxarg.ext",
+    "sphinx_design",
+    "sphinx_togglebutton",
+]
+myst_enable_extensions = [
+    "colon_fence",
+]
+
+# 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,
+}
+# 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[Optional[str], Optional[str]]:
+    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
+    module = info['module']
+
+    # try to determine the correct file and line number to link to
+    obj = sys.modules[module]
+
+    # get as specific as we can
+    lineno: int = 0
+    filename: str = ""
+    try:
+        for part in info['fullname'].split('.'):
+            obj = getattr(obj, part)
+
+            if not (inspect.isclass(obj) or inspect.isfunction(obj)
+                    or inspect.ismethod(obj)):
+                obj = obj.__class__  # type: ignore[assignment]
+
+            lineno = inspect.getsourcelines(obj)[1]
+            filename = (inspect.getsourcefile(obj)
+                        or f"{filename}.py").split("FastVideo/", 1)[1]
+    except Exception:
+        # For some things, like a class member, won't work, so
+        # we'll use the line number of the parent (the class)
+        pass
+
+    if filename.startswith("checkouts/"):
+        # a PR build on readthedocs
+        pr_number = filename.split("/")[1]
+        filename = filename.split("/", 2)[2]
+        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",
+    "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)
+
+
+class MockedClassDocumenter(autodoc.ClassDocumenter):
+    """Remove note about base class when a class is derived from object."""
+
+    def add_line(self, line: str, source: str, *lineno: int) -> None:
+        if line == "   Bases: :py:class:`object`":
+            return
+        super().add_line(line, source, *lineno)
+
+
+autodoc.ClassDocumenter = MockedClassDocumenter
+
+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),
+}
+
+autodoc_preserve_defaults = True
+autodoc_warningiserror = True
+
+navigation_with_keys = False

docs/source/developer_guide/overview.md

@@ -0,0 +1,50 @@
+# 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/vllm-project/vllm.git && cd vllm
+
+```
+
+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.0.post2 --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/generate_examples.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import itertools
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+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: Optional[str] = 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():
+    # Create the EXAMPLE_DOC_DIR if it doesn't exist
+    if not EXAMPLE_DOC_DIR.exists():
+        EXAMPLE_DOC_DIR.mkdir(parents=True)
+
+    # Create empty indices
+    examples_index = Index(
+        path=EXAMPLE_DOC_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 stored in reverse order because they are inserted into
+    # examples_index.documents at index 0 in order
+    category_indices = {
+        "other":
+        Index(
+            path=EXAMPLE_DOC_DIR / "examples_other_index.md",
+            title="Other",
+            description=
+            "Other examples that don't strongly fit into the online or offline serving categories.",  # noqa: E501
+            caption="Examples",
+        ),
+        "online_serving":
+        Index(
+            path=EXAMPLE_DOC_DIR / "examples_online_serving_index.md",
+            title="Online Serving",
+            description=
+            "Online serving examples demonstrate how to use FastVideo in an online setting, where the model is queried for predictions in real-time.",  # noqa: E501
+            caption="Examples",
+        ),
+        "offline_inference":
+        Index(
+            path=EXAMPLE_DOC_DIR / "examples_offline_inference_index.md",
+            title="Offline Inference",
+            description=
+            "Offline 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",
+        ),
+    }
+
+    examples = []
+    glob_patterns = ["*.py", "*.md", "*.sh"]
+    # Find categorised examples
+    for category in category_indices:
+        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
+    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))
+
+    # Generate the example documentation
+    for example in sorted(examples, key=lambda e: e.path.stem):
+        doc_path = EXAMPLE_DOC_DIR / f"{example.path.stem}.md"
+        with open(doc_path, "w+") as f:
+            f.write(example.generate())
+        # Add the example to the appropriate index
+        assert example.category is not None
+        index = category_indices.get(example.category, examples_index)
+        index.documents.append(example.path.stem)
+
+    # Generate the index files
+    for category_index in category_indices.values():
+        if category_index.documents:
+            examples_index.documents.insert(0, category_index.path.name)
+            with open(category_index.path, "w+") as f:
+                f.write(category_index.generate())
+
+    with open(examples_index.path, "w+") as f:
+        f.write(examples_index.generate())

docs/source/getting_started/examples/examples_index.md

@@ -0,0 +1,10 @@
+# Examples
+
+A collection of examples demonstrating usage of FastVideo.
+All documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.
+
+:::{toctree}
+:caption: Examples
+:maxdepth: 2
+
+:::

docs/source/getting_started/installation.md

@@ -0,0 +1,10 @@
+(fastvideo-installation)=
+
+# 🔧 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 [here](#sta-installation) to install STA.

docs/source/index.md

@@ -0,0 +1,88 @@
+# 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 lightweight framework for accelerating large video diffusion models.
+</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>
+:::
+
+FastVideo is a lightweight framework for accelerating large video diffusion models developed by the [Hao AI Lab](https://hao-ai-lab.github.io/).
+
+<div style="text-align: center;">
+  <video controls width="800">
+    <source src="https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1" type="video/mp4">
+    Your browser does not support the video tag.
+  </video>
+</div>
+
+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.
+- 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.
+
+## Documentation
+
+% How to start using vLLM?
+
+:::{toctree}
+:caption: Getting Started
+:maxdepth: 1
+
+getting_started/installation
+getting_started/examples/examples_index
+:::
+
+% What is STA Kernel?
+
+:::{toctree}
+:caption: Sliding Tile Attention
+:maxdepth: 1
+
+sliding_tile_attention/installation
+sliding_tile_attention/usage
+sliding_tile_attention/test
+sliding_tile_attention/demo
+:::
+
+:::{toctree}
+:caption: Inference
+:maxdepth: 1
+
+inference/stepvideo
+inference/hunyuanvideo
+inference/fasthunyuan
+inference/fastmochi
+:::
+
+:::{toctree}
+:caption: Developer Guide
+:maxdepth: 1
+
+developer_guide/overview
+:::
+
+## Indices and tables
+
+- {ref}`genindex`
+- {ref}`modindex`

docs/source/inference/fasthunyuan.md

@@ -0,0 +1,33 @@
+(fasthunyuan)=
+
+# FastHunyuan
+## 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).

docs/source/inference/fastmochi.md

@@ -0,0 +1,9 @@
+(fastmochi)=
+
+# 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/inference/hunyuanvideo.md

@@ -0,0 +1,18 @@
+(hunyuanvideo)=
+
+# HunyuanVideo
+## 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.

docs/source/inference/stepvideo.md

@@ -0,0 +1,16 @@
+(stepvideo)=
+
+# StepVideo
+## 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
+```

docs/source/inference/wanvideo.md

@@ -0,0 +1,44 @@
+(wanvideo)=
+
+# WanVideo
+## Inference T2V with WanVideo
+First, download the model:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-T2V-1.3B-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model 
+```
+
+or
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-T2V-14B-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model 
+```
+
+Then run the inference using:
+
+```bash
+sh scripts/inference/v1_inference_wan.sh
+```
+
+Remember to set `MODEL_BASE` and `num_gpus` accordingly.
+
+## Inference I2V with WanVideo
+First, download the model:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-I2V-14B-480P-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model
+```
+
+or
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-I2V-14B-720P-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model
+```
+
+Then run the inference using:
+
+```bash
+sh scripts/inference/v1_inference_wan_i2v.sh
+```
+
+Remember to set `MODEL_BASE` and `num_gpus` accordingly.

docs/source/sliding_tile_attention/demo.md

@@ -0,0 +1,11 @@
+(sta-demo)=
+
+# Demo
+There 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>

docs/source/sliding_tile_attention/installation.md

@@ -0,0 +1,25 @@
+(sta-installation)=
+
+# Installation
+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
+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
+```

docs/source/sliding_tile_attention/test.md

@@ -0,0 +1,7 @@
+(sta-test)=
+
+# Test
+
+```bash
+python test/test_sta.py
+```

docs/source/sliding_tile_attention/usage.md

@@ -0,0 +1,17 @@
+(sta-usage)=
+
+# 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)
+
+```

env_setup.sh

@@ -1,12 +0,0 @@
-#!/bin/bash
-
-# install torch
-pip install torch==2.5.0 torchvision --index-url https://download.pytorch.org/whl/cu124
-
-# install FA2 and diffusers
-pip install packaging ninja && pip install flash-attn==2.7.0.post2 --no-build-isolation 
-
-pip install -r requirements-lint.txt
-
-# install fastvideo
-pip install -e .

fastvideo/v1/attention/__init__.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
+                                                      AttentionMetadata,
+                                                      AttentionMetadataBuilder)
+from fastvideo.v1.attention.layer import DistributedAttention, LocalAttention
+from fastvideo.v1.attention.selector import get_attn_backend
+
+__all__ = [
+    "DistributedAttention",
+    "LocalAttention",
+    "AttentionBackend",
+    "AttentionMetadata",
+    "AttentionMetadataBuilder",
+    # "AttentionState",
+    "get_attn_backend",
+]

fastvideo/v1/attention/backends/abstract.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, fields
+from typing import (TYPE_CHECKING, Any, Dict, Generic, Optional, Protocol, Set,
+                    Type, TypeVar)
+
+if TYPE_CHECKING:
+    from fastvideo.v1.inference_args import InferenceArgs
+    from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+
+import torch
+
+
+class AttentionBackend(ABC):
+    """Abstract class for attention backends."""
+    # For some attention backends, we allocate an output tensor before
+    # calling the custom op. When piecewise cudagraph is enabled, this
+    # makes sure the output tensor is allocated inside the cudagraph.
+    accept_output_buffer: bool = False
+
+    @staticmethod
+    @abstractmethod
+    def get_name() -> str:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_impl_cls() -> Type["AttentionImpl"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_metadata_cls() -> Type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    # @staticmethod
+    # @abstractmethod
+    # def get_state_cls() -> Type["AttentionState"]:
+    #     raise NotImplementedError
+
+    # @classmethod
+    # def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
+    #     return cls.get_metadata_cls()(*args, **kwargs)
+
+    @staticmethod
+    @abstractmethod
+    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+
+@dataclass
+class AttentionMetadata:
+    """Attention metadata for prefill and decode batched together."""
+    # Current step of diffusion process
+    current_timestep: int
+
+    # @property
+    # @abstractmethod
+    # def inference_metadata(self) -> Optional["AttentionMetadata"]:
+    #     """Return the attention metadata that's required to run prefill
+    #     attention."""
+    #     pass
+
+    # @property
+    # @abstractmethod
+    # def training_metadata(self) -> Optional["AttentionMetadata"]:
+    #     """Return the attention metadata that's required to run decode
+    #     attention."""
+    #     pass
+
+    def asdict_zerocopy(self,
+                        skip_fields: Optional[Set[str]] = None
+                        ) -> Dict[str, Any]:
+        """Similar to dataclasses.asdict, but avoids deepcopying."""
+        if skip_fields is None:
+            skip_fields = set()
+        # Note that if we add dataclasses as fields, they will need
+        # similar handling.
+        return {
+            field.name: getattr(self, field.name)
+            for field in fields(self) if field.name not in skip_fields
+        }
+
+
+T = TypeVar("T", bound=AttentionMetadata)
+
+# class AttentionState(ABC, Generic[T]):
+#     """Holds attention backend-specific objects reused during the
+#     lifetime of the model runner."""
+
+#     @abstractmethod
+#     def __init__(self, runner: "ModelRunnerBase"):
+#         ...
+
+#     @abstractmethod
+#     @contextmanager
+#     def graph_capture(self, max_batch_size: int):
+#         """Context manager used when capturing CUDA graphs."""
+#         yield
+
+#     @abstractmethod
+#     def graph_clone(self, batch_size: int) -> "AttentionState[T]":
+#         """Clone attention state to save in CUDA graph metadata."""
+#         ...
+
+#     @abstractmethod
+#     def graph_capture_get_metadata_for_batch(
+#             self,
+#             batch_size: int,
+#             is_encoder_decoder_model: bool = False) -> T:
+#         """Get attention metadata for CUDA graph capture of batch_size."""
+#         ...
+
+#     @abstractmethod
+#     def get_graph_input_buffers(
+#             self,
+#             attn_metadata: T,
+#             is_encoder_decoder_model: bool = False) -> Dict[str, Any]:
+#         """Get attention-specific input buffers for CUDA graph capture."""
+#         ...
+
+#     @abstractmethod
+#     def prepare_graph_input_buffers(
+#             self,
+#             input_buffers: Dict[str, Any],
+#             attn_metadata: T,
+#             is_encoder_decoder_model: bool = False) -> None:
+#         """In-place modify input buffers dict for CUDA graph replay."""
+#         ...
+
+#     @abstractmethod
+#     def begin_forward(self, model_input: "ModelRunnerInputBase") -> None:
+#         """Prepare state for forward pass."""
+#         ...
+
+
+class AttentionMetadataBuilder(ABC, Generic[T]):
+    """Abstract class for attention metadata builders."""
+
+    @abstractmethod
+    def __init__(self) -> None:
+        """Create the builder, remember some configuration and parameters."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def prepare(self) -> None:
+        """Prepare for one batch."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def build(
+        self,
+        current_timestep: int,
+        forward_batch: "ForwardBatch",
+        inference_args: "InferenceArgs",
+    ) -> T:
+        """Build attention metadata with on-device tensors."""
+        raise NotImplementedError
+
+
+class AttentionLayer(Protocol):
+
+    _k_scale: torch.Tensor
+    _v_scale: torch.Tensor
+    _k_scale_float: float
+    _v_scale_float: float
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        ...
+
+
+class AttentionImpl(ABC, Generic[T]):
+
+    @abstractmethod
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        dropout_rate: float = 0.0,
+        causal: bool = False,
+        num_kv_heads: Optional[int] = None,
+    ) -> None:
+        raise NotImplementedError
+
+    def preprocess_qkv(self, qkv: torch.Tensor,
+                       attn_metadata: T) -> torch.Tensor:
+        """Preprocess QKV tensor before performing attention operation.
+
+        Default implementation returns the tensor unchanged.
+        Subclasses can override this to implement custom preprocessing
+        like reshaping, tiling, scaling, or other transformations.
+
+        Called AFTER all_to_all for distributed attention
+        
+        Args:
+            qkv: The query-key-value tensor
+            attn_metadata: Metadata for the attention operation
+            
+        Returns:
+            Processed QKV tensor
+        """
+        return qkv
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        """Postprocess the output tensor after the attention operation.
+
+        Default implementation returns the tensor unchanged.
+        Subclasses can override this to implement custom postprocessing
+        like untiling, scaling, or other transformations.
+
+        Called BEFORE all_to_all for distributed attention
+
+        Args:
+            output: The output tensor from the attention operation
+            attn_metadata: Metadata for the attention operation
+
+        Returns:
+            Postprocessed output tensor
+        """
+
+        return output
+
+    @abstractmethod
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        raise NotImplementedError

fastvideo/v1/attention/backends/flash_attn.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import List, Optional, Type
+
+import torch
+from flash_attn import flash_attn_func
+
+from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
+                                                      AttentionImpl,
+                                                      AttentionMetadata,
+                                                      AttentionMetadataBuilder)
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class FlashAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> List[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> Type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> Type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+
+class FlashAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        dropout_rate: float,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: Optional[int] = None,
+    ) -> None:
+        self.dropout_rate = dropout_rate
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ):
+        output = flash_attn_func(query,
+                                 key,
+                                 value,
+                                 dropout_p=self.dropout_rate,
+                                 softmax_scale=self.softmax_scale,
+                                 causal=self.causal)
+        return output

fastvideo/v1/attention/backends/sdpa.py

@@ -0,0 +1,72 @@
+from typing import List, Optional, Type
+
+import torch
+
+from fastvideo.v1.attention.backends.abstract import (
+    AttentionBackend)  # FlashAttentionMetadata,
+from fastvideo.v1.attention.backends.abstract import (AttentionImpl,
+                                                      AttentionMetadata)
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class SDPABackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> List[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SDPA"
+
+    @staticmethod
+    def get_impl_cls() -> Type["SDPAImpl"]:
+        return SDPAImpl
+
+    # @staticmethod
+    # def get_metadata_cls() -> Type["AttentionMetadata"]:
+    #     return FlashAttentionMetadata
+
+
+class SDPAImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        dropout_rate: float,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: Optional[int] = None,
+    ) -> None:
+        self.dropout_rate = dropout_rate
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        # transpose to bs, heads, seq_len, head_dim
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        attn_kwargs = {
+            "attn_mask": None,
+            "dropout_p": self.dropout_rate,
+            "is_causal": self.causal,
+            "scale": self.softmax_scale
+        }
+        if query.shape[1] != key.shape[1]:
+            attn_kwargs["enable_gqa"] = True
+        output = torch.nn.functional.scaled_dot_product_attention(
+            query, key, value, **attn_kwargs)
+        output = output.transpose(1, 2)
+        return output

fastvideo/v1/attention/backends/sliding_tile_attn.py

@@ -0,0 +1,195 @@
+import json
+from dataclasses import dataclass
+from typing import List, Optional, Type
+
+import torch
+from einops import rearrange
+from st_attn import sliding_tile_attention
+
+import fastvideo.v1.envs as envs
+from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
+                                                      AttentionImpl,
+                                                      AttentionMetadata,
+                                                      AttentionMetadataBuilder)
+from fastvideo.v1.distributed import get_sp_group
+from fastvideo.v1.inference_args import InferenceArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+
+logger = init_logger(__name__)
+
+
+# TODO(will-refactor): move this to a utils file
+def dict_to_3d_list(mask_strategy,
+                    t_max=50,
+                    l_max=60,
+                    h_max=24) -> List[List[List[Optional[torch.Tensor]]]]:
+    result = [[[None for _ in range(h_max)] for _ in range(l_max)]
+              for _ in range(t_max)]
+    if mask_strategy is None:
+        return result
+    for key, value in mask_strategy.items():
+        t, layer, h = map(int, key.split('_'))
+        result[t][layer][h] = value
+    return result
+
+
+class SlidingTileAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> List[int]:
+        # TODO(will-refactor): check this
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SLIDING_TILE_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> Type["SlidingTileAttentionImpl"]:
+        return SlidingTileAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> Type["SlidingTileAttentionMetadata"]:
+        return SlidingTileAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> Type["SlidingTileAttentionMetadataBuilder"]:
+        return SlidingTileAttentionMetadataBuilder
+
+
+@dataclass
+class SlidingTileAttentionMetadata(AttentionMetadata):
+    text_length: int
+
+
+class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        forward_batch: ForwardBatch,
+        inference_args: InferenceArgs,
+    ) -> SlidingTileAttentionMetadata:
+
+        return SlidingTileAttentionMetadata(
+            current_timestep=current_timestep,
+            text_length=forward_batch.attention_mask.sum(),
+        )
+
+
+class SlidingTileAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        dropout_rate: float,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: Optional[int] = None,
+    ) -> None:
+        # TODO(will-refactor): for now this is the mask strategy, but maybe we should
+        # have a more general config for STA?
+        config_file = envs.FASTVIDEO_ATTENTION_CONFIG
+        if config_file is None:
+            raise ValueError("FASTVIDEO_ATTENTION_CONFIG is not set")
+
+        with open(config_file) as f:
+            mask_strategy = json.load(f)
+
+        mask_strategy = dict_to_3d_list(mask_strategy)
+
+        self.mask_strategy = mask_strategy
+        sp_group = get_sp_group()
+        self.sp_size = sp_group.world_size
+
+    def tile(self, x: torch.Tensor) -> torch.Tensor:
+        x = rearrange(x,
+                      "b (sp t h w) head d -> b (t sp h w) head d",
+                      sp=self.sp_size,
+                      t=30 // self.sp_size,
+                      h=48,
+                      w=80)
+        return rearrange(
+            x,
+            "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
+            n_t=5,
+            n_h=6,
+            n_w=10,
+            ts_t=6,
+            ts_h=8,
+            ts_w=8)
+
+    def untile(self, x: torch.Tensor) -> torch.Tensor:
+        x = rearrange(
+            x,
+            "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
+            n_t=5,
+            n_h=6,
+            n_w=10,
+            ts_t=6,
+            ts_h=8,
+            ts_w=8)
+        return rearrange(x,
+                         "b (t sp h w) head d -> b (sp t h w) head d",
+                         sp=self.sp_size,
+                         t=30 // self.sp_size,
+                         h=48,
+                         w=80)
+
+    def preprocess_qkv(
+        self,
+        qkv: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        return self.tile(qkv)
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: SlidingTileAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.untile(output)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attn_metadata: SlidingTileAttentionMetadata,
+    ) -> torch.Tensor:
+
+        assert self.mask_strategy is not None, "mask_strategy cannot be None for SlidingTileAttention"
+        assert self.mask_strategy[
+            0] is not None, "mask_strategy[0] cannot be None for SlidingTileAttention"
+
+        text_length = attn_metadata.text_length
+
+        query = q.transpose(1, 2)
+        key = k.transpose(1, 2)
+        value = v.transpose(1, 2)
+
+        head_num = query.size(1)
+        sp_group = get_sp_group()
+        current_rank = sp_group.rank_in_group
+        start_head = current_rank * head_num
+        windows = [
+            self.mask_strategy[head_idx + start_head]
+            for head_idx in range(head_num)
+        ]
+
+        hidden_states = sliding_tile_attention(query, key, value, windows,
+                                               text_length).transpose(1, 2)
+
+        hidden_states = hidden_states.transpose(1, 2)
+
+        return hidden_states

fastvideo/v1/attention/layer.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from fastvideo.v1.attention.selector import (backend_name_to_enum,
+                                             get_attn_backend)
+from fastvideo.v1.distributed.communication_op import (
+    sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D)
+from fastvideo.v1.distributed.parallel_state import (
+    get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size)
+from fastvideo.v1.forward_context import ForwardContext, get_forward_context
+
+
+class DistributedAttention(nn.Module):
+    """Distributed attention layer.
+    """
+
+    def __init__(self,
+                 num_heads: int,
+                 head_size: int,
+                 num_kv_heads: Optional[int] = None,
+                 dropout_rate: float = 0.0,
+                 softmax_scale: Optional[float] = None,
+                 causal: bool = False,
+                 **extra_impl_args) -> None:
+        super().__init__()
+        # self.dropout_rate = dropout_rate
+        # self.causal = causal
+        if softmax_scale is None:
+            self.softmax_scale = head_size**-0.5
+        else:
+            self.softmax_scale = softmax_scale
+
+        if num_kv_heads is None:
+            num_kv_heads = num_heads
+
+        dtype = torch.get_default_dtype()
+        attn_backend = get_attn_backend(head_size, dtype, distributed=True)
+        impl_cls = attn_backend.get_impl_cls()
+        self.impl = impl_cls(num_heads=num_heads,
+                             head_size=head_size,
+                             dropout_rate=dropout_rate,
+                             causal=causal,
+                             softmax_scale=self.softmax_scale,
+                             num_kv_heads=num_kv_heads,
+                             **extra_impl_args)
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.num_kv_heads = num_kv_heads
+        self.backend = backend_name_to_enum(attn_backend.get_name())
+        self.dtype = dtype
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        replicated_q: Optional[torch.Tensor] = None,
+        replicated_k: Optional[torch.Tensor] = None,
+        replicated_v: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Forward pass for distributed attention.
+        
+        Args:
+            q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
+            k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
+            v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim]
+            replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens
+            replicated_k (Optional[torch.Tensor]): Replicated key tensor
+            replicated_v (Optional[torch.Tensor]): Replicated value tensor
+            
+        Returns:
+            Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
+                - o (torch.Tensor): Output tensor after attention for the main sequence
+                - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided
+        """
+        # Check input shapes
+        assert q.dim() == 4 and k.dim() == 4 and v.dim(
+        ) == 4, "Expected 4D tensors"
+        # assert bs = 1
+        assert q.shape[
+            0] == 1, "Batch size must be 1, and there should be no padding tokens"
+        batch_size, seq_len, num_heads, head_dim = q.shape
+        local_rank = get_sequence_model_parallel_rank()
+        world_size = get_sequence_model_parallel_world_size()
+
+        forward_context: ForwardContext = get_forward_context()
+        ctx_attn_metadata = forward_context.attn_metadata
+
+        # Stack QKV
+        qkv = torch.cat([q, k, v], dim=0)  # [3, seq_len, num_heads, head_dim]
+
+        # Redistribute heads across sequence dimension
+        qkv = sequence_model_parallel_all_to_all_4D(qkv,
+                                                    scatter_dim=2,
+                                                    gather_dim=1)
+
+        # Apply backend-specific preprocess_qkv
+        qkv = self.impl.preprocess_qkv(qkv, ctx_attn_metadata)
+
+        # Concatenate with replicated QKV if provided
+        if replicated_q is not None:
+            assert replicated_k is not None and replicated_v is not None
+            replicated_qkv = torch.cat(
+                [replicated_q, replicated_k, replicated_v],
+                dim=0)  # [3, seq_len, num_heads, head_dim]
+            heads_per_rank = num_heads // world_size
+            replicated_qkv = replicated_qkv[:, :, local_rank *
+                                            heads_per_rank:(local_rank + 1) *
+                                            heads_per_rank]
+            qkv = torch.cat([qkv, replicated_qkv], dim=1)
+
+        q, k, v = qkv.chunk(3, dim=0)
+
+        output = self.impl.forward(q, k, v, ctx_attn_metadata)
+
+        # Redistribute back if using sequence parallelism
+        replicated_output = None
+        if replicated_q is not None:
+            replicated_output = output[:, seq_len * world_size:]
+            output = output[:, :seq_len * world_size]
+            # TODO: make this asynchronous
+            replicated_output = sequence_model_parallel_all_gather(
+                replicated_output, dim=2)
+
+        # Apply backend-specific postprocess_output
+        output = self.impl.postprocess_output(output, ctx_attn_metadata)
+
+        output = sequence_model_parallel_all_to_all_4D(output,
+                                                       scatter_dim=1,
+                                                       gather_dim=2)
+        return output, replicated_output
+
+
+class LocalAttention(nn.Module):
+    """Attention layer.
+    """
+
+    def __init__(self,
+                 num_heads: int,
+                 head_size: int,
+                 num_kv_heads: Optional[int] = None,
+                 dropout_rate: float = 0.0,
+                 softmax_scale: Optional[float] = None,
+                 causal: bool = False,
+                 **extra_impl_args) -> None:
+        super().__init__()
+        # self.dropout_rate = dropout_rate
+        # self.causal = causal
+        if softmax_scale is None:
+            self.softmax_scale = head_size**-0.5
+        else:
+            self.softmax_scale = softmax_scale
+        if num_kv_heads is None:
+            num_kv_heads = num_heads
+
+        dtype = torch.get_default_dtype()
+        attn_backend = get_attn_backend(head_size, dtype, distributed=False)
+        impl_cls = attn_backend.get_impl_cls()
+        self.impl = impl_cls(num_heads=num_heads,
+                             head_size=head_size,
+                             dropout_rate=dropout_rate,
+                             softmax_scale=self.softmax_scale,
+                             num_kv_heads=num_kv_heads,
+                             causal=causal,
+                             **extra_impl_args)
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.num_kv_heads = num_kv_heads
+        self.backend = backend_name_to_enum(attn_backend.get_name())
+        self.dtype = dtype
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Apply local attention between query, key and value tensors.
+        
+        Args:
+            q (torch.Tensor): Query tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            k (torch.Tensor): Key tensor of shape [batch_size, seq_len, num_heads, head_dim] 
+            v (torch.Tensor): Value tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            
+        Returns:
+            torch.Tensor: Output tensor after local attention
+        """
+        # Check input shapes
+        assert q.dim() == 4 and k.dim() == 4 and v.dim(
+        ) == 4, "Expected 4D tensors"
+
+        forward_context: ForwardContext = get_forward_context()
+        ctx_attn_metadata = forward_context.attn_metadata
+
+        output = self.impl.forward(q, k, v, ctx_attn_metadata)
+        return output

fastvideo/v1/attention/selector.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py
+
+import os
+from contextlib import contextmanager
+from functools import cache
+from typing import Generator, Optional, Type, cast
+
+import torch
+
+import fastvideo.v1.envs as envs
+from fastvideo.v1.attention.backends.abstract import AttentionBackend
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.platforms import _Backend, current_platform
+from fastvideo.v1.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
+
+logger = init_logger(__name__)
+
+
+def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:
+    """
+    Convert a string backend name to a _Backend enum value.
+
+    Returns:
+    * _Backend: enum value if backend_name is a valid in-tree type
+    * None: otherwise it's an invalid in-tree type or an out-of-tree platform is
+            loaded.
+    """
+    assert backend_name is not None
+    return _Backend[backend_name] if backend_name in _Backend.__members__ else \
+          None
+
+
+def get_env_variable_attn_backend() -> Optional[_Backend]:
+    '''
+    Get the backend override specified by the FastVideo attention
+    backend environment variable, if one is specified.
+
+    Returns:
+
+    * _Backend enum value if an override is specified
+    * None otherwise
+    '''
+    backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
+    return (None
+            if backend_name is None else backend_name_to_enum(backend_name))
+
+
+# Global state allows a particular choice of backend
+# to be forced, overriding the logic which auto-selects
+# a backend based on system & workload configuration
+# (default behavior if this variable is None)
+#
+# THIS SELECTION TAKES PRECEDENCE OVER THE
+# FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE
+forced_attn_backend: Optional[_Backend] = None
+
+
+def global_force_attn_backend(attn_backend: Optional[_Backend]) -> None:
+    '''
+    Force all attention operations to use a specified backend.
+
+    Passing `None` for the argument re-enables automatic
+    backend selection.,
+
+    Arguments:
+
+    * attn_backend: backend selection (None to revert to auto)
+    '''
+    global forced_attn_backend
+    forced_attn_backend = attn_backend
+
+
+def get_global_forced_attn_backend() -> Optional[_Backend]:
+    '''
+    Get the currently-forced choice of attention backend,
+    or None if auto-selection is currently enabled.
+    '''
+    return forced_attn_backend
+
+
+def get_attn_backend(
+    head_size: int,
+    dtype: torch.dtype,
+    distributed: bool,
+) -> Type[AttentionBackend]:
+    """Selects which attention backend to use and lazily imports it."""
+    # Accessing envs.* behind an @lru_cache decorator can cause the wrong
+    # value to be returned from the cache if the value changes between calls.
+    return _cached_get_attn_backend(
+        head_size=head_size,
+        dtype=dtype,
+        distributed=distributed,
+    )
+
+
+@cache
+def _cached_get_attn_backend(
+    head_size: int,
+    dtype: torch.dtype,
+    distributed: bool,
+) -> Type[AttentionBackend]:
+    # Check whether a particular choice of backend was
+    # previously forced.
+    #
+    # THIS SELECTION OVERRIDES THE FASTVIDEO_ATTENTION_BACKEND
+    # ENVIRONMENT VARIABLE.
+    selected_backend = None
+    backend_by_global_setting: Optional[_Backend] = (
+        get_global_forced_attn_backend())
+    if backend_by_global_setting is not None:
+        selected_backend = backend_by_global_setting
+    else:
+        # Check the environment variable and override if specified
+        backend_by_env_var: Optional[str] = envs.FASTVIDEO_ATTENTION_BACKEND
+        if backend_by_env_var is not None:
+            selected_backend = backend_name_to_enum(backend_by_env_var)
+
+    # get device-specific attn_backend
+    attention_cls = current_platform.get_attn_backend_cls(
+        selected_backend, head_size, dtype, distributed)
+    if not attention_cls:
+        raise ValueError(
+            f"Invalid attention backend for {current_platform.device_name}")
+    return cast(Type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
+
+
+@contextmanager
+def global_force_attn_backend_context_manager(
+        attn_backend: _Backend) -> Generator[None, None, None]:
+    '''
+    Globally force a FastVideo attention backend override within a
+    context manager, reverting the global attention backend
+    override to its prior state upon exiting the context
+    manager.
+
+    Arguments:
+
+    * attn_backend: attention backend to force
+
+    Returns:
+
+    * Generator
+    '''
+
+    # Save the current state of the global backend override (if any)
+    original_value = get_global_forced_attn_backend()
+
+    # Globally force the new backend override
+    global_force_attn_backend(attn_backend)
+
+    # Yield control back to the enclosed code block
+    try:
+        yield
+    finally:
+        # Revert the original global backend override, if any
+        global_force_attn_backend(original_value)

fastvideo/v1/default_configs/v1_inference_hunyuan_config.yaml

@@ -0,0 +1,16 @@
+num_gpus: 4
+model_path: FastVideo/FastHunyuan-diffusers
+master_port: 29503
+sp_size: 4 
+tp_size: 4 
+height: 720 
+width: 1280 
+num_frames: 125 
+num_inference_steps: 6 
+guidance_scale: 1 
+embedded_cfg_scale: 6 
+flow_shift: 17 
+prompt_path: ./assets/prompt.txt 
+seed: 1024 
+output_path: outputs_video/ 
+vae-sp: True

fastvideo/v1/distributed/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.v1.distributed.communication_op import *
+from fastvideo.v1.distributed.parallel_state import *
+from fastvideo.v1.distributed.utils import *

fastvideo/v1/distributed/communication_op.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/communication_op.py
+
+import torch
+import torch.distributed
+
+from fastvideo.v1.distributed.parallel_state import get_sp_group, get_tp_group
+
+
+def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    return get_tp_group().all_reduce(input_)
+
+
+def tensor_model_parallel_all_gather(input_: torch.Tensor,
+                                     dim: int = -1) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    return get_tp_group().all_gather(input_, dim)
+
+
+# TODO: remove model, make it sequence_parallel
+def sequence_model_parallel_all_to_all_4D(input_: torch.Tensor,
+                                          scatter_dim: int = 2,
+                                          gather_dim: int = 1) -> torch.Tensor:
+    """All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group."""
+    return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim)
+
+
+def sequence_model_parallel_all_gather(input_: torch.Tensor,
+                                       dim: int = -1) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    return get_sp_group().all_gather(input_, dim)

fastvideo/v1/distributed/device_communicators/base_device_communicator.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py
+
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+
+class DeviceCommunicatorBase:
+    """
+    Base class for device-specific communicator.
+    It can use the `cpu_group` to initialize the communicator.
+    If the device has PyTorch integration (PyTorch can recognize its
+    communication backend), the `device_group` will also be given.
+    """
+
+    def __init__(self,
+                 cpu_group: ProcessGroup,
+                 device: Optional[torch.device] = None,
+                 device_group: Optional[ProcessGroup] = None,
+                 unique_name: str = ""):
+        self.device = device or torch.device("cpu")
+        self.cpu_group = cpu_group
+        self.device_group = device_group
+        self.unique_name = unique_name
+        self.rank = dist.get_rank(cpu_group)
+        self.world_size = dist.get_world_size(cpu_group)
+        self.ranks = dist.get_process_group_ranks(cpu_group)
+        self.global_rank = dist.get_rank()
+        self.global_world_size = dist.get_world_size()
+        self.rank_in_group = dist.get_group_rank(self.cpu_group,
+                                                 self.global_rank)
+
+    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
+        dist.all_reduce(input_, group=self.device_group)
+        return input_
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+        input_size = input_.size()
+        # NOTE: we have to use concat-style all-gather here,
+        # stack-style all-gather has compatibility issues with
+        # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
+        output_size = (input_size[0] * self.world_size, ) + input_size[1:]
+        # Allocate output tensor.
+        output_tensor = torch.empty(output_size,
+                                    dtype=input_.dtype,
+                                    device=input_.device)
+        # All-gather.
+        dist.all_gather_into_tensor(output_tensor,
+                                    input_,
+                                    group=self.device_group)
+        # Reshape
+        output_tensor = output_tensor.reshape((self.world_size, ) + input_size)
+        output_tensor = output_tensor.movedim(0, dim)
+        output_tensor = output_tensor.reshape(input_size[:dim] +
+                                              (self.world_size *
+                                               input_size[dim], ) +
+                                              input_size[dim + 1:])
+        return output_tensor
+
+    def gather(self,
+               input_: torch.Tensor,
+               dst: int = 0,
+               dim: int = -1) -> Optional[torch.Tensor]:
+        """
+        NOTE: We assume that the input tensor is on the same device across
+        all the ranks.
+        NOTE: `dst` is the local rank of the destination rank.
+        """
+        world_size = self.world_size
+        assert -input_.dim() <= dim < input_.dim(), (
+            f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Allocate output tensor.
+        if self.rank_in_group == dst:
+            gather_list = [torch.empty_like(input_) for _ in range(world_size)]
+        else:
+            gather_list = None
+        # Gather.
+        torch.distributed.gather(input_,
+                                 gather_list,
+                                 dst=self.ranks[dst],
+                                 group=self.device_group)
+        if self.rank_in_group == dst:
+            output_tensor = torch.cat(gather_list, dim=dim)
+        else:
+            output_tensor = None
+        return output_tensor
+
+    def all_to_all_4D(self,
+                      input_: torch.Tensor,
+                      scatter_dim: int = 2,
+                      gather_dim: int = 1) -> torch.Tensor:
+        """Specialized all-to-all operation for 4D tensors (e.g., for QKV matrices).
+        
+        Args:
+            input_ (torch.Tensor): 4D input tensor to be scattered and gathered.
+            scatter_dim (int, optional): Dimension along which to scatter. Defaults to 2.
+            gather_dim (int, optional): Dimension along which to gather. Defaults to 1.
+            
+        Returns:
+            torch.Tensor: Output tensor after all-to-all operation.
+        """
+        # Bypass the function if we are using only 1 GPU.
+        if self.world_size == 1:
+            return input_
+
+        assert input_.dim(
+        ) == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
+
+        if scatter_dim == 2 and gather_dim == 1:
+            # input: (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs)
+            bs, shard_seqlen, hc, hs = input_.shape
+            seqlen = shard_seqlen * self.world_size
+            shard_hc = hc // self.world_size
+
+            # Reshape and transpose for scattering
+            input_t = (input_.reshape(bs, shard_seqlen, self.world_size,
+                                      shard_hc, hs).transpose(0,
+                                                              2).contiguous())
+
+            output = torch.empty_like(input_t)
+
+            torch.distributed.all_to_all_single(output,
+                                                input_t,
+                                                group=self.device_group)
+            torch.cuda.synchronize()
+
+            # Reshape and transpose back
+            output = output.reshape(seqlen, bs, shard_hc,
+                                    hs).transpose(0, 1).contiguous().reshape(
+                                        bs, seqlen, shard_hc, hs)
+
+            return output
+
+        elif scatter_dim == 1 and gather_dim == 2:
+            # input: (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs)
+            bs, seqlen, shard_hc, hs = input_.shape
+            hc = shard_hc * self.world_size
+            shard_seqlen = seqlen // self.world_size
+
+            # Reshape and transpose for scattering
+            input_t = (input_.reshape(bs, self.world_size, shard_seqlen,
+                                      shard_hc, hs).transpose(0, 3).transpose(
+                                          0, 1).contiguous().reshape(
+                                              self.world_size, shard_hc,
+                                              shard_seqlen, bs, hs))
+            output = torch.empty_like(input_t)
+
+            torch.distributed.all_to_all_single(output,
+                                                input_t,
+                                                group=self.device_group)
+            torch.cuda.synchronize()
+
+            # Reshape and transpose back
+            output = output.reshape(hc, shard_seqlen, bs,
+                                    hs).transpose(0, 2).contiguous().reshape(
+                                        bs, shard_seqlen, hc, hs)
+
+            return output
+        else:
+            raise RuntimeError(
+                "scatter_dim must be 1 or 2 and gather_dim must be 1 or 2")
+
+    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+        """Sends a tensor to the destination rank in a non-blocking way"""
+        """NOTE: `dst` is the local rank of the destination rank."""
+        if dst is None:
+            dst = (self.rank_in_group + 1) % self.world_size
+        torch.distributed.send(tensor, self.ranks[dst], self.device_group)
+
+    def recv(self,
+             size: torch.Size,
+             dtype: torch.dtype,
+             src: Optional[int] = None) -> torch.Tensor:
+        """Receives a tensor from the source rank."""
+        """NOTE: `src` is the local rank of the source rank."""
+        if src is None:
+            src = (self.rank_in_group - 1) % self.world_size
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        torch.distributed.recv(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def destroy(self):
+        pass

fastvideo/v1/distributed/device_communicators/cuda_communicator.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py
+
+from typing import Optional
+
+import torch
+from torch.distributed import ProcessGroup
+
+from fastvideo.v1.distributed.device_communicators.base_device_communicator import (
+    DeviceCommunicatorBase)
+
+
+class CudaCommunicator(DeviceCommunicatorBase):
+
+    def __init__(self,
+                 cpu_group: ProcessGroup,
+                 device: Optional[torch.device] = None,
+                 device_group: Optional[ProcessGroup] = None,
+                 unique_name: str = ""):
+        super().__init__(cpu_group, device, device_group, unique_name)
+
+        from fastvideo.v1.distributed.device_communicators.pynccl import (
+            PyNcclCommunicator)
+
+        self.pynccl_comm: Optional[PyNcclCommunicator] = None
+        if self.world_size > 1:
+            self.pynccl_comm = PyNcclCommunicator(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+    def all_reduce(self, input_):
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        out = pynccl_comm.all_reduce(input_)
+        if out is None:
+            # fall back to the default all-reduce using PyTorch.
+            # this usually happens during testing.
+            # when we run the model, allreduce only happens for the TP
+            # group, where we always have either custom allreduce or pynccl.
+            out = input_.clone()
+            torch.distributed.all_reduce(out, group=self.device_group)
+        return out
+
+    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+        """Sends a tensor to the destination rank in a non-blocking way"""
+        """NOTE: `dst` is the local rank of the destination rank."""
+        if dst is None:
+            dst = (self.rank_in_group + 1) % self.world_size
+
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.send(tensor, dst)
+        else:
+            torch.distributed.send(tensor, self.ranks[dst], self.device_group)
+
+    def recv(self,
+             size: torch.Size,
+             dtype: torch.dtype,
+             src: Optional[int] = None) -> torch.Tensor:
+        """Receives a tensor from the source rank."""
+        """NOTE: `src` is the local rank of the source rank."""
+        if src is None:
+            src = (self.rank_in_group - 1) % self.world_size
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.recv(tensor, src)
+        else:
+            torch.distributed.recv(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def destroy(self) -> None:
+        if self.pynccl_comm is not None:
+            self.pynccl_comm = None

fastvideo/v1/distributed/device_communicators/pynccl.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py
+
+from typing import Optional, Union
+
+# ===================== import region =====================
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup, ReduceOp
+
+from fastvideo.v1.distributed.device_communicators.pynccl_wrapper import (
+    NCCLLibrary, buffer_type, cudaStream_t, ncclComm_t, ncclDataTypeEnum,
+    ncclRedOpTypeEnum, ncclUniqueId)
+from fastvideo.v1.distributed.utils import StatelessProcessGroup
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.utils import current_stream
+
+logger = init_logger(__name__)
+
+
+class PyNcclCommunicator:
+
+    def __init__(
+        self,
+        group: Union[ProcessGroup, StatelessProcessGroup],
+        device: Union[int, str, torch.device],
+        library_path: Optional[str] = None,
+    ):
+        """
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the PyNcclCommunicator to. If None,
+                it will be bind to f"cuda:{local_rank}".
+            library_path: the path to the NCCL library. If None, it will
+                use the default library path.
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device.
+        """
+        if not isinstance(group, StatelessProcessGroup):
+            assert dist.is_initialized()
+            assert dist.get_backend(group) != dist.Backend.NCCL, (
+                "PyNcclCommunicator should be attached to a non-NCCL group.")
+            # note: this rank is the rank in the group
+            self.rank = dist.get_rank(group)
+            self.world_size = dist.get_world_size(group)
+        else:
+            self.rank = group.rank
+            self.world_size = group.world_size
+
+        self.group = group
+
+        # if world_size == 1, no need to create communicator
+        if self.world_size == 1:
+            self.available = False
+            self.disabled = True
+            return
+        try:
+            self.nccl = NCCLLibrary(library_path)
+        except Exception:
+            # disable because of missing NCCL library
+            # e.g. in a non-GPU environment
+            self.available = False
+            self.disabled = True
+            return
+
+        self.available = True
+        self.disabled = False
+
+        logger.info("FastVideo is using nccl==%s", self.nccl.ncclGetVersion())
+
+        if self.rank == 0:
+            # get the unique id from NCCL
+            self.unique_id = self.nccl.ncclGetUniqueId()
+        else:
+            # construct an empty unique id
+            self.unique_id = ncclUniqueId()
+
+        if not isinstance(group, StatelessProcessGroup):
+            tensor = torch.ByteTensor(list(self.unique_id.internal))
+            ranks = dist.get_process_group_ranks(group)
+            # arg `src` in `broadcast` is the global rank
+            dist.broadcast(tensor, src=ranks[0], group=group)
+            byte_list = tensor.tolist()
+            for i, byte in enumerate(byte_list):
+                self.unique_id.internal[i] = byte
+        else:
+            self.unique_id = group.broadcast_obj(self.unique_id, src=0)
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        # now `device` is a `torch.device` object
+        assert isinstance(device, torch.device)
+        self.device = device
+        # nccl communicator and stream will use this device
+        # `torch.cuda.device` is a context manager that changes the
+        # current cuda device to the specified one
+        with torch.cuda.device(device):
+            self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
+                self.world_size, self.unique_id, self.rank)
+
+            stream = current_stream()
+            # A small all_reduce for warmup.
+            data = torch.zeros(1, device=device)
+            self.all_reduce(data)
+            stream.synchronize()
+            del data
+
+    def all_reduce(self,
+                   in_tensor: torch.Tensor,
+                   op: ReduceOp = ReduceOp.SUM,
+                   stream=None) -> torch.Tensor:
+        if self.disabled:
+            return None
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert in_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {in_tensor.device}")
+
+        out_tensor = torch.empty_like(in_tensor)
+
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclAllReduce(buffer_type(in_tensor.data_ptr()),
+                                buffer_type(out_tensor.data_ptr()),
+                                in_tensor.numel(),
+                                ncclDataTypeEnum.from_torch(in_tensor.dtype),
+                                ncclRedOpTypeEnum.from_torch(op), self.comm,
+                                cudaStream_t(stream.cuda_stream))
+        return out_tensor
+
+    def all_gather(self,
+                   output_tensor: torch.Tensor,
+                   input_tensor: torch.Tensor,
+                   stream=None):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}")
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclAllGather(buffer_type(input_tensor.data_ptr()),
+                                buffer_type(output_tensor.data_ptr()),
+                                input_tensor.numel(),
+                                ncclDataTypeEnum.from_torch(input_tensor.dtype),
+                                self.comm, cudaStream_t(stream.cuda_stream))
+
+    def reduce_scatter(self,
+                       output_tensor: torch.Tensor,
+                       input_tensor: torch.Tensor,
+                       op: ReduceOp = ReduceOp.SUM,
+                       stream=None):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}")
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclReduceScatter(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()), output_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op), self.comm,
+            cudaStream_t(stream.cuda_stream))
+
+    def send(self, tensor: torch.Tensor, dst: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}")
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclSend(buffer_type(tensor.data_ptr()), tensor.numel(),
+                           ncclDataTypeEnum.from_torch(tensor.dtype), dst,
+                           self.comm, cudaStream_t(stream.cuda_stream))
+
+    def recv(self, tensor: torch.Tensor, src: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}")
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclRecv(buffer_type(tensor.data_ptr()), tensor.numel(),
+                           ncclDataTypeEnum.from_torch(tensor.dtype), src,
+                           self.comm, cudaStream_t(stream.cuda_stream))
+
+    def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}")
+        if stream is None:
+            stream = current_stream()
+        if src == self.rank:
+            sendbuff = buffer_type(tensor.data_ptr())
+            # NCCL requires the sender also to have a receive buffer
+            recvbuff = buffer_type(tensor.data_ptr())
+        else:
+            sendbuff = buffer_type()
+            recvbuff = buffer_type(tensor.data_ptr())
+        self.nccl.ncclBroadcast(sendbuff, recvbuff, tensor.numel(),
+                                ncclDataTypeEnum.from_torch(tensor.dtype), src,
+                                self.comm, cudaStream_t(stream.cuda_stream))

fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py

@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl_wrapper.py
+
+# This file is a pure Python wrapper for the NCCL library.
+# The main purpose is to use NCCL combined with CUDA graph.
+# Before writing this script, we tried the following approach:
+# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself
+#  often gets stuck when initializing the NCCL communicator.
+# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce`
+#  contains many other potential cuda APIs, that are not allowed during
+#  capturing the CUDA graph. For further details, please check
+# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ .
+#
+# Another rejected idea is to write a C/C++ binding for NCCL. It is usually
+# doable, but we often encounter issues related with nccl versions, and need
+# to switch between different versions of NCCL. See
+# https://github.com/NVIDIA/nccl/issues/1234 for more details.
+# A C/C++ binding is not flexible enough to handle this. It requires
+# recompilation of the code every time we want to switch between different
+# versions. This current implementation, with a **pure** Python wrapper, is
+# more flexible. We can easily switch between different versions of NCCL by
+# changing the environment variable `FASTVIDEO_NCCL_SO_PATH`, or the `so_file`
+# variable in the code.
+
+#TODO(will): support FASTVIDEO_NCCL_SO_PATH
+
+import ctypes
+import platform
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch.distributed import ReduceOp
+
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.utils import find_nccl_library
+
+logger = init_logger(__name__)
+
+# === export types and functions from nccl to Python ===
+# for the original nccl definition, please check
+# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in
+
+ncclResult_t = ctypes.c_int
+ncclComm_t = ctypes.c_void_p
+
+
+class ncclUniqueId(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * 128)]
+
+
+cudaStream_t = ctypes.c_void_p
+buffer_type = ctypes.c_void_p
+
+ncclDataType_t = ctypes.c_int
+
+
+class ncclDataTypeEnum:
+    ncclInt8 = 0
+    ncclChar = 0
+    ncclUint8 = 1
+    ncclInt32 = 2
+    ncclInt = 2
+    ncclUint32 = 3
+    ncclInt64 = 4
+    ncclUint64 = 5
+    ncclFloat16 = 6
+    ncclHalf = 6
+    ncclFloat32 = 7
+    ncclFloat = 7
+    ncclFloat64 = 8
+    ncclDouble = 8
+    ncclBfloat16 = 9
+    ncclNumTypes = 10
+
+    @classmethod
+    def from_torch(cls, dtype: torch.dtype) -> int:
+        if dtype == torch.int8:
+            return cls.ncclInt8
+        if dtype == torch.uint8:
+            return cls.ncclUint8
+        if dtype == torch.int32:
+            return cls.ncclInt32
+        if dtype == torch.int64:
+            return cls.ncclInt64
+        if dtype == torch.float16:
+            return cls.ncclFloat16
+        if dtype == torch.float32:
+            return cls.ncclFloat32
+        if dtype == torch.float64:
+            return cls.ncclFloat64
+        if dtype == torch.bfloat16:
+            return cls.ncclBfloat16
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+
+ncclRedOp_t = ctypes.c_int
+
+
+class ncclRedOpTypeEnum:
+    ncclSum = 0
+    ncclProd = 1
+    ncclMax = 2
+    ncclMin = 3
+    ncclAvg = 4
+    ncclNumOps = 5
+
+    @classmethod
+    def from_torch(cls, op: ReduceOp) -> int:
+        if op == ReduceOp.SUM:
+            return cls.ncclSum
+        if op == ReduceOp.PRODUCT:
+            return cls.ncclProd
+        if op == ReduceOp.MAX:
+            return cls.ncclMax
+        if op == ReduceOp.MIN:
+            return cls.ncclMin
+        if op == ReduceOp.AVG:
+            return cls.ncclAvg
+        raise ValueError(f"Unsupported op: {op}")
+
+
+@dataclass
+class Function:
+    name: str
+    restype: Any
+    argtypes: List[Any]
+
+
+class NCCLLibrary:
+    exported_functions = [
+        # const char* ncclGetErrorString(ncclResult_t result)
+        Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
+        # ncclResult_t  ncclGetVersion(int *version);
+        Function("ncclGetVersion", ncclResult_t,
+                 [ctypes.POINTER(ctypes.c_int)]),
+        # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId);
+        Function("ncclGetUniqueId", ncclResult_t,
+                 [ctypes.POINTER(ncclUniqueId)]),
+        # ncclResult_t  ncclCommInitRank(
+        #   ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
+        # note that ncclComm_t is a pointer type, so the first argument
+        # is a pointer to a pointer
+        Function("ncclCommInitRank", ncclResult_t, [
+            ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int
+        ]),
+        # ncclResult_t  ncclAllReduce(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function("ncclAllReduce", ncclResult_t, [
+            buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
+            ncclRedOp_t, ncclComm_t, cudaStream_t
+        ]),
+
+        # ncclResult_t  ncclAllGather(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function("ncclAllGather", ncclResult_t, [
+            buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
+            ncclComm_t, cudaStream_t
+        ]),
+
+        # ncclResult_t  ncclReduceScatter(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function("ncclReduceScatter", ncclResult_t, [
+            buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
+            ncclRedOp_t, ncclComm_t, cudaStream_t
+        ]),
+
+        # ncclResult_t  ncclSend(
+        #   const void* sendbuff, size_t count, ncclDataType_t datatype,
+        #   int dest, ncclComm_t comm, cudaStream_t stream);
+        Function("ncclSend", ncclResult_t, [
+            buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int,
+            ncclComm_t, cudaStream_t
+        ]),
+
+        # ncclResult_t  ncclRecv(
+        #   void* recvbuff, size_t count, ncclDataType_t datatype,
+        #   int src, ncclComm_t comm, cudaStream_t stream);
+        Function("ncclRecv", ncclResult_t, [
+            buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int,
+            ncclComm_t, cudaStream_t
+        ]),
+
+        # ncclResult_t ncclBroadcast(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, int root, ncclComm_t comm,
+        #   cudaStream_t stream);
+        Function("ncclBroadcast", ncclResult_t, [
+            buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
+            ctypes.c_int, ncclComm_t, cudaStream_t
+        ]),
+
+        # be cautious! this is a collective call, it will block until all
+        # processes in the communicator have called this function.
+        # because Python object destruction can happen in random order,
+        # it is better not to call it at all.
+        # ncclResult_t  ncclCommDestroy(ncclComm_t comm);
+        Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
+    ]
+
+    # class attribute to store the mapping from the path to the library
+    # to avoid loading the same library multiple times
+    path_to_library_cache: Dict[str, Any] = {}
+
+    # class attribute to store the mapping from library path
+    #  to the corresponding dictionary
+    path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
+
+    def __init__(self, so_file: Optional[str] = None):
+
+        so_file = so_file or find_nccl_library()
+
+        try:
+            if so_file not in NCCLLibrary.path_to_dict_mapping:
+                lib = ctypes.CDLL(so_file)
+                NCCLLibrary.path_to_library_cache[so_file] = lib
+            self.lib = NCCLLibrary.path_to_library_cache[so_file]
+        except Exception as e:
+            logger.error(
+                "Failed to load NCCL library from %s ."
+                "It is expected if you are not running on NVIDIA/AMD GPUs."
+                "Otherwise, the nccl library might not exist, be corrupted "
+                "or it does not support the current platform %s."
+                "If you already have the library, please set the "
+                "environment variable FASTVIDEO_NCCL_SO_PATH"
+                " to point to the correct nccl library path.", so_file,
+                platform.platform())
+            raise e
+
+        if so_file not in NCCLLibrary.path_to_dict_mapping:
+            _funcs: Dict[str, Any] = {}
+            for func in NCCLLibrary.exported_functions:
+                f = getattr(self.lib, func.name)
+                f.restype = func.restype
+                f.argtypes = func.argtypes
+                _funcs[func.name] = f
+            NCCLLibrary.path_to_dict_mapping[so_file] = _funcs
+        self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]
+
+    def ncclGetErrorString(self, result: ncclResult_t) -> str:
+        return str(self._funcs["ncclGetErrorString"](result).decode("utf-8"))
+
+    def NCCL_CHECK(self, result: ncclResult_t) -> None:
+        if result != 0:
+            error_str = self.ncclGetErrorString(result)
+            raise RuntimeError(f"NCCL error: {error_str}")
+
+    def ncclGetVersion(self) -> str:
+        version = ctypes.c_int()
+        self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
+        version_str = str(version.value)
+        # something like 21903 --> "2.19.3"
+        major = version_str[0].lstrip("0")
+        minor = version_str[1:3].lstrip("0")
+        patch = version_str[3:].lstrip("0")
+        return f"{major}.{minor}.{patch}"
+
+    def ncclGetUniqueId(self) -> ncclUniqueId:
+        unique_id = ncclUniqueId()
+        self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id)))
+        return unique_id
+
+    def ncclCommInitRank(self, world_size: int, unique_id: ncclUniqueId,
+                         rank: int) -> ncclComm_t:
+        comm = ncclComm_t()
+        self.NCCL_CHECK(self._funcs["ncclCommInitRank"](ctypes.byref(comm),
+                                                        world_size, unique_id,
+                                                        rank))
+        return comm
+
+    def ncclAllReduce(self, sendbuff: buffer_type, recvbuff: buffer_type,
+                      count: int, datatype: int, op: int, comm: ncclComm_t,
+                      stream: cudaStream_t) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(self._funcs["ncclAllReduce"](sendbuff, recvbuff, count,
+                                                     datatype, op, comm,
+                                                     stream))
+
+    def ncclReduceScatter(self, sendbuff: buffer_type, recvbuff: buffer_type,
+                          count: int, datatype: int, op: int, comm: ncclComm_t,
+                          stream: cudaStream_t) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(self._funcs["ncclReduceScatter"](sendbuff, recvbuff,
+                                                         count, datatype, op,
+                                                         comm, stream))
+
+    def ncclAllGather(self, sendbuff: buffer_type, recvbuff: buffer_type,
+                      count: int, datatype: int, comm: ncclComm_t,
+                      stream: cudaStream_t) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # which is an aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(self._funcs["ncclAllGather"](sendbuff, recvbuff, count,
+                                                     datatype, comm, stream))
+
+    def ncclSend(self, sendbuff: buffer_type, count: int, datatype: int,
+                 dest: int, comm: ncclComm_t, stream: cudaStream_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclSend"](sendbuff, count, datatype, dest,
+                                                comm, stream))
+
+    def ncclRecv(self, recvbuff: buffer_type, count: int, datatype: int,
+                 src: int, comm: ncclComm_t, stream: cudaStream_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclRecv"](recvbuff, count, datatype, src,
+                                                comm, stream))
+
+    def ncclBroadcast(self, sendbuff: buffer_type, recvbuff: buffer_type,
+                      count: int, datatype: int, root: int, comm: ncclComm_t,
+                      stream: cudaStream_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclBroadcast"](sendbuff, recvbuff, count,
+                                                     datatype, root, comm,
+                                                     stream))
+
+    def ncclCommDestroy(self, comm: ncclComm_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
+
+
+__all__ = [
+    "NCCLLibrary", "ncclDataTypeEnum", "ncclRedOpTypeEnum", "ncclUniqueId",
+    "ncclComm_t", "cudaStream_t", "buffer_type"
+]

fastvideo/v1/distributed/utils.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/utils.py
+
+# Copyright 2023 The vLLM team.
+# Adapted from
+# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+import dataclasses
+import pickle
+import time
+from collections import deque
+from typing import Any, Deque, Dict, Optional, Sequence, Tuple
+
+import torch
+from torch.distributed import TCPStore
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def ensure_divisibility(numerator, denominator) -> None:
+    """Ensure that numerator is divisible by the denominator."""
+    assert numerator % denominator == 0, "{} is not divisible by {}".format(
+        numerator, denominator)
+
+
+def divide(numerator: int, denominator: int) -> int:
+    """Ensure that numerator is divisible by the denominator and return
+    the division value."""
+    ensure_divisibility(numerator, denominator)
+    return numerator // denominator
+
+
+def split_tensor_along_last_dim(
+    tensor: torch.Tensor,
+    num_partitions: int,
+    contiguous_split_chunks: bool = False,
+) -> Sequence[torch.Tensor]:
+    """ Split a tensor along its last dimension.
+
+        Arguments:
+            tensor: input tensor.
+            num_partitions: number of partitions to split the tensor
+            contiguous_split_chunks: If True, make each chunk contiguous
+                                     in memory.
+
+        Returns:
+            A list of Tensors
+    """
+    # Get the size and dimension.
+    last_dim = tensor.dim() - 1
+    last_dim_size = divide(tensor.size()[last_dim], num_partitions)
+    # Split.
+    tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
+    # NOTE: torch.split does not create contiguous tensors by default.
+    if contiguous_split_chunks:
+        return tuple(chunk.contiguous() for chunk in tensor_list)
+
+    return tuple(tensor_list)
+
+
+@dataclasses.dataclass
+class StatelessProcessGroup:
+    """A dataclass to hold a metadata store, and the rank, world_size of the
+    group. Only use it to communicate metadata between processes.
+    For data-plane communication, create NCCL-related objects.
+    """
+    rank: int
+    world_size: int
+    store: torch._C._distributed_c10d.Store
+    data_expiration_seconds: int = 3600  # 1 hour
+
+    # dst rank -> counter
+    send_dst_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    # src rank -> counter
+    recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    broadcast_send_counter: int = 0
+    broadcast_recv_src_counter: Dict[int, int] = dataclasses.field(
+        default_factory=dict)
+
+    # A deque to store the data entries, with key and timestamp.
+    entries: Deque[Tuple[str, float]] = dataclasses.field(default_factory=deque)
+
+    def __post_init__(self):
+        assert self.rank < self.world_size
+        self.send_dst_counter = {i: 0 for i in range(self.world_size)}
+        self.recv_src_counter = {i: 0 for i in range(self.world_size)}
+        self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)}
+
+    def send_obj(self, obj: Any, dst: int):
+        """Send an object to a destination rank."""
+        self.expire_data()
+        key = f"send_to/{dst}/{self.send_dst_counter[dst]}"
+        self.store.set(key, pickle.dumps(obj))
+        self.send_dst_counter[dst] += 1
+        self.entries.append((key, time.time()))
+
+    def expire_data(self) -> None:
+        """Expire data that is older than `data_expiration_seconds` seconds."""
+        while self.entries:
+            # check the oldest entry
+            key, timestamp = self.entries[0]
+            if time.time() - timestamp > self.data_expiration_seconds:
+                self.store.delete_key(key)
+                self.entries.popleft()
+            else:
+                break
+
+    def recv_obj(self, src: int) -> Any:
+        """Receive an object from a source rank."""
+        obj = pickle.loads(
+            self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}"))
+        self.recv_src_counter[src] += 1
+        return obj
+
+    def broadcast_obj(self, obj: Optional[Any], src: int) -> Any:
+        """Broadcast an object from a source rank to all other ranks.
+        It does not clean up after all ranks have received the object.
+        Use it for limited times, e.g., for initialization.
+        """
+        if self.rank == src:
+            self.expire_data()
+            key = (f"broadcast_from/{src}/"
+                   f"{self.broadcast_send_counter}")
+            self.store.set(key, pickle.dumps(obj))
+            self.broadcast_send_counter += 1
+            self.entries.append((key, time.time()))
+            return obj
+        else:
+            key = (f"broadcast_from/{src}/"
+                   f"{self.broadcast_recv_src_counter[src]}")
+            recv_obj = pickle.loads(self.store.get(key))
+            self.broadcast_recv_src_counter[src] += 1
+            return recv_obj
+
+    def all_gather_obj(self, obj: Any) -> list[Any]:
+        """All gather an object from all ranks."""
+        gathered_objs = []
+        for i in range(self.world_size):
+            if i == self.rank:
+                gathered_objs.append(obj)
+                self.broadcast_obj(obj, src=self.rank)
+            else:
+                recv_obj = self.broadcast_obj(None, src=i)
+                gathered_objs.append(recv_obj)
+        return gathered_objs
+
+    def barrier(self):
+        """A barrier to synchronize all ranks."""
+        for i in range(self.world_size):
+            if i == self.rank:
+                self.broadcast_obj(None, src=self.rank)
+            else:
+                self.broadcast_obj(None, src=i)
+
+    @staticmethod
+    def create(
+        host: str,
+        port: int,
+        rank: int,
+        world_size: int,
+        data_expiration_seconds: int = 3600,
+    ) -> "StatelessProcessGroup":
+        """A replacement for `torch.distributed.init_process_group` that does not
+        pollute the global state.
+
+        If we have process A and process B called `torch.distributed.init_process_group`
+        to form a group, and then we want to form another group with process A, B, C,
+        D, it is not possible in PyTorch, because process A and process B have already
+        formed a group, and process C and process D cannot join that group. This
+        function is a workaround for this issue.
+
+        `torch.distributed.init_process_group` is a global call, while this function
+        is a stateless call. It will return a `StatelessProcessGroup` object that can be
+        used for exchanging metadata. With this function, process A and process B
+        can call `StatelessProcessGroup.create` to form a group, and then process A, B,
+        C, and D can call `StatelessProcessGroup.create` to form another group.
+        """ # noqa
+        store = TCPStore(
+            host_name=host,
+            port=port,
+            world_size=world_size,
+            is_master=(rank == 0),
+        )
+
+        return StatelessProcessGroup(
+            rank=rank,
+            world_size=world_size,
+            store=store,
+            data_expiration_seconds=data_expiration_seconds)

fastvideo/v1/entrypoints/cli/cli_types.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/types.py
+
+import argparse
+
+from fastvideo.v1.utils import FlexibleArgumentParser
+
+
+class CLISubcommand:
+    """Base class for CLI subcommands"""
+
+    def __init__(self):
+        self.name = ""
+
+    def cmd(self, args: argparse.Namespace) -> None:
+        """Execute the command with the given arguments"""
+        raise NotImplementedError
+
+    def validate(self, args: argparse.Namespace) -> None:
+        """Validate the arguments for this command"""
+        pass
+
+    def subparser_init(
+            self,
+            subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+        """Initialize the subparser for this command"""
+        raise NotImplementedError

fastvideo/v1/entrypoints/cli/generate.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py
+
+import argparse
+from typing import List
+
+from fastvideo.v1.entrypoints.cli import utils
+from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
+from fastvideo.v1.inference_args import InferenceArgs
+from fastvideo.v1.utils import FlexibleArgumentParser
+
+
+class GenerateSubcommand(CLISubcommand):
+    """The `generate` subcommand for the FastVideo CLI"""
+
+    def __init__(self) -> None:
+        self.name = "generate"
+        super().__init__()
+
+    def cmd(self, args: argparse.Namespace) -> None:
+        excluded_args = [
+            'subparser', 'config', 'num_gpus', 'master_port',
+            'dispatch_function'
+        ]
+
+        # Create a filtered dictionary of arguments
+        filtered_args = {
+            k: v
+            for k, v in vars(args).items()
+            if k not in excluded_args and v is not None
+        }
+
+        main_args = []
+
+        for key, value in filtered_args.items():
+            # Convert underscores to dashes in argument names
+            arg_name = f"--{key.replace('_', '-')}"
+
+            # Handle boolean flags
+            if isinstance(value, bool):
+                if value:
+                    main_args.append(arg_name)
+            else:
+                main_args.append(arg_name)
+                main_args.append(str(value))
+
+        utils.launch_distributed(args.num_gpus,
+                                 main_args,
+                                 master_port=args.master_port)
+
+    def validate(self, args: argparse.Namespace) -> None:
+        if args.num_gpus is not None and args.num_gpus <= 0:
+            raise ValueError("Number of gpus must be positive")
+
+        if args.master_port is not None and (args.master_port < 1024
+                                             or args.master_port > 65535):
+            raise ValueError("Master port must be between 1024 and 65535")
+
+    def subparser_init(
+            self,
+            subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
+        generate_parser = subparsers.add_parser(
+            "generate",
+            help="Run inference on a model",
+            usage=
+            "fastvideo generate --model-path MODEL_PATH_OR_ID --prompt PROMPT [OPTIONS]"
+        )
+
+        generate_parser.add_argument(
+            "--config",
+            type=str,
+            default='',
+            required=False,
+            help="Read CLI options from a config YAML file.")
+
+        generate_parser.add_argument("--num-gpus",
+                                     type=int,
+                                     default=1,
+                                     help="Number of GPUs to use")
+        generate_parser.add_argument("--master-port",
+                                     type=int,
+                                     default=None,
+                                     help="Port for the master process")
+
+        generate_parser = InferenceArgs.add_cli_args(generate_parser)
+
+        return generate_parser
+
+
+def cmd_init() -> List[CLISubcommand]:
+    return [GenerateSubcommand()]

fastvideo/v1/entrypoints/cli/main.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py
+
+from typing import List
+
+from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
+from fastvideo.v1.entrypoints.cli.generate import cmd_init as generate_cmd_init
+from fastvideo.v1.utils import FlexibleArgumentParser
+
+
+def cmd_init() -> List[CLISubcommand]:
+    """Initialize all commands from separate modules"""
+    commands = []
+    commands.extend(generate_cmd_init())
+    return commands
+
+
+def main() -> None:
+    parser = FlexibleArgumentParser(description="FastVideo CLI")
+    parser.add_argument('-v', '--version', action='version', version='0.1.0')
+
+    subparsers = parser.add_subparsers(required=False, dest="subparser")
+
+    cmds = {}
+    for cmd in cmd_init():
+        cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
+        cmds[cmd.name] = cmd
+    args = parser.parse_args()
+    if args.subparser in cmds:
+        cmds[args.subparser].validate(args)
+
+    if hasattr(args, "dispatch_function"):
+        args.dispatch_function(args)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

fastvideo/v1/entrypoints/cli/utils.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import subprocess
+import sys
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def launch_distributed(num_gpus=None, args=None, master_port=None):
+    """
+    Launch a distributed job with the given arguments
+    
+    Args:
+        num_gpus: Number of GPUs to use
+        args: Arguments to pass to v1_fastvideo_inference.py (defaults to sys.argv[1:])
+        master_port: Port for the master process (default: random)
+    """
+
+    current_env = os.environ.copy()
+    python_executable = sys.executable
+    project_root = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "../../../.."))
+    main_script = os.path.join(project_root,
+                               "fastvideo/v1/sample/v1_fastvideo_inference.py")
+
+    cmd = [
+        python_executable, "-m", "torch.distributed.run",
+        f"--nproc_per_node={num_gpus}"
+    ]
+
+    if master_port is not None:
+        cmd.append(f"--master_port={master_port}")
+
+    cmd.append(main_script)
+    cmd.extend(args)
+
+    logger.info("Running inference with %d GPU(s)", num_gpus)
+    logger.info("Launching command: %s", " ".join(cmd))
+
+    current_env["PYTHONIOENCODING"] = "utf-8"
+    process = subprocess.Popen(cmd,
+                               env=current_env,
+                               stdout=subprocess.PIPE,
+                               stderr=subprocess.STDOUT,
+                               universal_newlines=True,
+                               bufsize=1,
+                               encoding='utf-8',
+                               errors='replace')
+
+    if process.stdout:
+        for line in iter(process.stdout.readline, ''):
+            print(line.strip())
+
+    return process.wait()

fastvideo/v1/envs.py

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py
+
+import os
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+if TYPE_CHECKING:
+    FASTVIDEO_RINGBUFFER_WARNING_INTERVAL: int = 60
+    FASTVIDEO_NCCL_SO_PATH: Optional[str] = None
+    LD_LIBRARY_PATH: Optional[str] = None
+    FASTVIDEO_USE_TRITON_FLASH_ATTN: bool = False
+    FASTVIDEO_FLASH_ATTN_VERSION: Optional[int] = None
+    LOCAL_RANK: int = 0
+    CUDA_VISIBLE_DEVICES: Optional[str] = None
+    FASTVIDEO_CACHE_ROOT: str = os.path.expanduser("~/.cache/fastvideo")
+    FASTVIDEO_CONFIG_ROOT: str = os.path.expanduser("~/.config/fastvideo")
+    FASTVIDEO_CONFIGURE_LOGGING: int = 1
+    FASTVIDEO_LOGGING_LEVEL: str = "INFO"
+    FASTVIDEO_LOGGING_PREFIX: str = ""
+    FASTVIDEO_LOGGING_CONFIG_PATH: Optional[str] = None
+    FASTVIDEO_TRACE_FUNCTION: int = 0
+    FASTVIDEO_ATTENTION_BACKEND: Optional[str] = None
+    FASTVIDEO_ATTENTION_CONFIG: Optional[str] = None
+    FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork"
+    FASTVIDEO_TARGET_DEVICE: str = "cuda"
+    MAX_JOBS: Optional[str] = None
+    NVCC_THREADS: Optional[str] = None
+    CMAKE_BUILD_TYPE: Optional[str] = None
+    VERBOSE: bool = False
+    FASTVIDEO_SERVER_DEV_MODE: bool = False
+
+
+def get_default_cache_root() -> str:
+    return os.getenv(
+        "XDG_CACHE_HOME",
+        os.path.join(os.path.expanduser("~"), ".cache"),
+    )
+
+
+def get_default_config_root() -> str:
+    return os.getenv(
+        "XDG_CONFIG_HOME",
+        os.path.join(os.path.expanduser("~"), ".config"),
+    )
+
+
+def maybe_convert_int(value: Optional[str]) -> Optional[int]:
+    if value is None:
+        return None
+    return int(value)
+
+
+# The begin-* and end* here are used by the documentation generator
+# to extract the used env vars.
+
+# begin-env-vars-definition
+
+environment_variables: Dict[str, Callable[[], Any]] = {
+
+    # ================== Installation Time Env Vars ==================
+
+    # Target device of FastVideo, supporting [cuda (by default),
+    # rocm, neuron, cpu, openvino]
+    "FASTVIDEO_TARGET_DEVICE":
+    lambda: os.getenv("FASTVIDEO_TARGET_DEVICE", "cuda"),
+
+    # Maximum number of compilation jobs to run in parallel.
+    # By default this is the number of CPUs
+    "MAX_JOBS":
+    lambda: os.getenv("MAX_JOBS", None),
+
+    # Number of threads to use for nvcc
+    # By default this is 1.
+    # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU.
+    "NVCC_THREADS":
+    lambda: os.getenv("NVCC_THREADS", None),
+
+    # If set, fastvideo will use precompiled binaries (*.so)
+    "FASTVIDEO_USE_PRECOMPILED":
+    lambda: bool(os.environ.get("FASTVIDEO_USE_PRECOMPILED")) or bool(
+        os.environ.get("FASTVIDEO_PRECOMPILED_WHEEL_LOCATION")),
+
+    # CMake build type
+    # If not set, defaults to "Debug" or "RelWithDebInfo"
+    # Available options: "Debug", "Release", "RelWithDebInfo"
+    "CMAKE_BUILD_TYPE":
+    lambda: os.getenv("CMAKE_BUILD_TYPE"),
+
+    # If set, fastvideo will print verbose logs during installation
+    "VERBOSE":
+    lambda: bool(int(os.getenv('VERBOSE', '0'))),
+
+    # Root directory for FASTVIDEO configuration files
+    # Defaults to `~/.config/fastvideo` unless `XDG_CONFIG_HOME` is set
+    # Note that this not only affects how fastvideo finds its configuration files
+    # during runtime, but also affects how fastvideo installs its configuration
+    # files during **installation**.
+    "FASTVIDEO_CONFIG_ROOT":
+    lambda: os.path.expanduser(
+        os.getenv(
+            "FASTVIDEO_CONFIG_ROOT",
+            os.path.join(get_default_config_root(), "fastvideo"),
+        )),
+
+    # ================== Runtime Env Vars ==================
+
+    # Root directory for FASTVIDEO cache files
+    # Defaults to `~/.cache/fastvideo` unless `XDG_CACHE_HOME` is set
+    "FASTVIDEO_CACHE_ROOT":
+    lambda: os.path.expanduser(
+        os.getenv(
+            "FASTVIDEO_CACHE_ROOT",
+            os.path.join(get_default_cache_root(), "fastvideo"),
+        )),
+
+    # Interval in seconds to log a warning message when the ring buffer is full
+    "FASTVIDEO_RINGBUFFER_WARNING_INTERVAL":
+    lambda: int(os.environ.get("FASTVIDEO_RINGBUFFER_WARNING_INTERVAL", "60")),
+
+    # Path to the NCCL library file. It is needed because nccl>=2.19 brought
+    # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234
+    "FASTVIDEO_NCCL_SO_PATH":
+    lambda: os.environ.get("FASTVIDEO_NCCL_SO_PATH", None),
+
+    # when `FASTVIDEO_NCCL_SO_PATH` is not set, fastvideo will try to find the nccl
+    # library file in the locations specified by `LD_LIBRARY_PATH`
+    "LD_LIBRARY_PATH":
+    lambda: os.environ.get("LD_LIBRARY_PATH", None),
+
+    # flag to control if fastvideo should use triton flash attention
+    "FASTVIDEO_USE_TRITON_FLASH_ATTN":
+    lambda:
+    (os.environ.get("FASTVIDEO_USE_TRITON_FLASH_ATTN", "True").lower() in
+     ("true", "1")),
+
+    # Force fastvideo to use a specific flash-attention version (2 or 3), only valid
+    # when using the flash-attention backend.
+    "FASTVIDEO_FLASH_ATTN_VERSION":
+    lambda: maybe_convert_int(
+        os.environ.get("FASTVIDEO_FLASH_ATTN_VERSION", None)),
+
+    # Internal flag to enable Dynamo fullgraph capture
+    "FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE":
+    lambda: bool(
+        os.environ.get("FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0"),
+
+    # local rank of the process in the distributed setting, used to determine
+    # the GPU device id
+    "LOCAL_RANK":
+    lambda: int(os.environ.get("LOCAL_RANK", "0")),
+
+    # used to control the visible devices in the distributed setting
+    "CUDA_VISIBLE_DEVICES":
+    lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None),
+
+    # timeout for each iteration in the engine
+    "FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S":
+    lambda: int(os.environ.get("FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S", "60")),
+
+    # Logging configuration
+    # If set to 0, fastvideo will not configure logging
+    # If set to 1, fastvideo will configure logging using the default configuration
+    #    or the configuration file specified by FASTVIDEO_LOGGING_CONFIG_PATH
+    "FASTVIDEO_CONFIGURE_LOGGING":
+    lambda: int(os.getenv("FASTVIDEO_CONFIGURE_LOGGING", "1")),
+    "FASTVIDEO_LOGGING_CONFIG_PATH":
+    lambda: os.getenv("FASTVIDEO_LOGGING_CONFIG_PATH"),
+
+    # this is used for configuring the default logging level
+    "FASTVIDEO_LOGGING_LEVEL":
+    lambda: os.getenv("FASTVIDEO_LOGGING_LEVEL", "INFO"),
+
+    # if set, FASTVIDEO_LOGGING_PREFIX will be prepended to all log messages
+    "FASTVIDEO_LOGGING_PREFIX":
+    lambda: os.getenv("FASTVIDEO_LOGGING_PREFIX", ""),
+
+    # Trace function calls
+    # If set to 1, fastvideo will trace function calls
+    # Useful for debugging
+    "FASTVIDEO_TRACE_FUNCTION":
+    lambda: int(os.getenv("FASTVIDEO_TRACE_FUNCTION", "0")),
+
+    # Backend for attention computation
+    # Available options:
+    # - "TORCH_SDPA": use torch.nn.MultiheadAttention
+    # - "FLASH_ATTN": use FlashAttention
+    # - "STA" : use sliding tile attention
+    "FASTVIDEO_ATTENTION_BACKEND":
+    lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
+
+    # Path to the attention configuration file. Only used for sliding tile
+    # attention for now.
+    "FASTVIDEO_ATTENTION_CONFIG":
+    lambda: (None if os.getenv("FASTVIDEO_ATTENTION_CONFIG", None) is None else
+             os.path.expanduser(os.getenv("FASTVIDEO_ATTENTION_CONFIG", "."))),
+
+    # Use dedicated multiprocess context for workers.
+    # Both spawn and fork work
+    "FASTVIDEO_WORKER_MULTIPROC_METHOD":
+    lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "fork"),
+
+    # Enables torch profiler if set. Path to the directory where torch profiler
+    # traces are saved. Note that it must be an absolute path.
+    "FASTVIDEO_TORCH_PROFILER_DIR":
+    lambda: (None
+             if os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", None) is None else os.
+             path.expanduser(os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", "."))),
+
+    # If set, fastvideo will run in development mode, which will enable
+    # some additional endpoints for developing and debugging,
+    # e.g. `/reset_prefix_cache`
+    "FASTVIDEO_SERVER_DEV_MODE":
+    lambda: bool(int(os.getenv("FASTVIDEO_SERVER_DEV_MODE", "0"))),
+}
+
+# end-env-vars-definition
+
+
+def __getattr__(name: str):
+    # lazy evaluation of environment variables
+    if name in environment_variables:
+        return environment_variables[name]()
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    return list(environment_variables.keys())

fastvideo/v1/forward_context.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/forward_context.py
+
+import time
+from collections import defaultdict
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from fastvideo.v1.inference_args import InferenceArgs
+from fastvideo.v1.logger import init_logger
+
+if TYPE_CHECKING:
+    from fastvideo.v1.attention import AttentionMetadata
+
+logger = init_logger(__name__)
+
+# TODO(will): check if this is needed
+# track_batchsize: bool = envs.FASTVIDEO_LOG_BATCHSIZE_INTERVAL >= 0
+track_batchsize: bool = False
+last_logging_time: float = 0
+forward_start_time: float = 0
+# batchsize_logging_interval: float = envs.FASTVIDEO_LOG_BATCHSIZE_INTERVAL
+batchsize_logging_interval: float = 1000
+batchsize_forward_time: defaultdict = defaultdict(list)
+
+
+#
+@dataclass
+class ForwardContext:
+    # TODO(will): check this arg
+    # copy from vllm_config.compilation_config.static_forward_context
+    # attn_layers: Dict[str, Any]
+    # TODO: extend to support per-layer dynamic forward context
+    attn_metadata: "AttentionMetadata"  # set dynamically for each forward pass
+
+
+_forward_context: Optional[ForwardContext] = None
+
+
+def get_forward_context() -> ForwardContext:
+    """Get the current forward context."""
+    assert _forward_context is not None, (
+        "Forward context is not set. "
+        "Please use `set_forward_context` to set the forward context.")
+    return _forward_context
+
+
+# TODO(will): finalize the interface
+@contextmanager
+def set_forward_context(current_timestep,
+                        attn_metadata,
+                        inference_args: InferenceArgs = None):
+    """A context manager that stores the current forward context,
+    can be attention metadata, etc.
+    Here we can inject common logic for every model forward pass.
+    """
+    global forward_start_time
+    need_to_track_batchsize = track_batchsize and attn_metadata is not None
+    if need_to_track_batchsize:
+        forward_start_time = time.perf_counter()
+    global _forward_context
+    prev_context = _forward_context
+    _forward_context = ForwardContext(attn_metadata=attn_metadata)
+    try:
+        yield
+    finally:
+        global last_logging_time, batchsize_logging_interval
+        if need_to_track_batchsize:
+            if hasattr(attn_metadata, "num_prefill_tokens"):
+                # for v0 attention backends
+                batchsize = attn_metadata.num_prefill_tokens + \
+                    attn_metadata.num_decode_tokens
+            else:
+                # for v1 attention backends
+                batchsize = attn_metadata.num_input_tokens
+            # we use synchronous scheduling right now,
+            # adding a sync point here should not affect
+            # scheduling of the next batch
+            torch.cuda.synchronize()
+            now = time.perf_counter()
+            # time measurement is in milliseconds
+            batchsize_forward_time[batchsize].append(
+                (now - forward_start_time) * 1000)
+            if now - last_logging_time > batchsize_logging_interval:
+                last_logging_time = now
+                forward_stats = []
+                for bs, times in batchsize_forward_time.items():
+                    if len(times) <= 1:
+                        # can be cudagraph / profiling run
+                        continue
+                    medium = torch.quantile(torch.tensor(times), q=0.5).item()
+                    medium = round(medium, 2)
+                    forward_stats.append((bs, len(times), medium))
+                forward_stats.sort(key=lambda x: x[1], reverse=True)
+                if forward_stats:
+                    logger.info(("Batchsize forward time stats "
+                                 "(batchsize, count, median_time(ms)): %s"),
+                                forward_stats)
+        _forward_context = prev_context

fastvideo/v1/inference_args.py

@@ -0,0 +1,454 @@
+# SPDX-License-Identifier: Apache-2.0
+# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py
+"""The arguments of FastVideo Inference."""
+
+import argparse
+import dataclasses
+from typing import List, Optional
+
+from fastvideo.v1.utils import FlexibleArgumentParser
+
+
+@dataclasses.dataclass
+class InferenceArgs:
+    # Model and path configuration
+    model_path: str
+
+    # HuggingFace specific parameters
+    trust_remote_code: bool = False
+    revision: Optional[str] = None
+
+    # Parallelism
+    tp_size: int = 1
+    sp_size: int = 1
+    dist_timeout: Optional[int] = None  # timeout for torch.distributed
+
+    # Video generation parameters
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 117
+    num_inference_steps: int = 50
+    guidance_scale: float = 1.0
+    guidance_rescale: float = 0.0
+    embedded_cfg_scale: float = 6.0
+    flow_shift: Optional[float] = None
+
+    output_type: str = "pil"
+
+    # Model configuration
+    precision: str = "bf16"
+
+    # VAE configuration
+    vae_precision: str = "fp16"
+    vae_tiling: bool = True
+    vae_sp: bool = False
+
+    # Image encoder configuration
+    image_encoder_precision: str = "fp32"
+
+    # Text encoder configuration
+    text_encoder_precision: str = "fp16"
+    text_len: int = 256
+    hidden_state_skip_layer: int = 2
+
+    # Secondary text encoder
+    text_encoder_precision_2: str = "fp16"
+    text_len_2: int = 77
+
+    # Flow Matching parameters
+    flow_solver: str = "euler"
+    denoise_type: str = "flow"  # Deprecated. Will use scheduler_config.json
+
+    # STA (Spatial-Temporal Attention) parameters
+    mask_strategy_file_path: Optional[str] = None
+    enable_torch_compile: bool = False
+
+    # Scheduler options
+    scheduler_type: str = "euler"  # Deprecated. Will use the param in scheduler_config.json
+
+    neg_prompt: Optional[str] = None
+    num_videos: int = 1
+    fps: int = 24
+    use_cpu_offload: bool = False
+    disable_autocast: bool = False
+
+    # Logging
+    log_level: str = "info"
+
+    # Inference parameters
+    image_path: Optional[str] = None
+    prompt: Optional[str] = None
+    prompt_path: Optional[str] = None
+    output_path: str = "outputs/"
+    seed: int = 1024
+    device_str: Optional[str] = None
+    device = None
+
+    def __post_init__(self):
+        pass
+
+    @staticmethod
+    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
+        # Model and path configuration
+        parser.add_argument(
+            "--model-path",
+            type=str,
+            required=True,
+            help=
+            "The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
+        )
+        parser.add_argument(
+            "--dit-weight",
+            type=str,
+            help="Path to the DiT model weights",
+        )
+        parser.add_argument(
+            "--model-dir",
+            type=str,
+            help="Directory containing StepVideo model",
+        )
+
+        # HuggingFace specific parameters
+        parser.add_argument(
+            "--trust-remote-code",
+            action="store_true",
+            default=InferenceArgs.trust_remote_code,
+            help="Trust remote code when loading HuggingFace models",
+        )
+        parser.add_argument(
+            "--revision",
+            type=str,
+            default=InferenceArgs.revision,
+            help=
+            "The specific model version to use (can be a branch name, tag name, or commit id)",
+        )
+
+        # Parallelism
+        parser.add_argument(
+            "--tensor-parallel-size",
+            "--tp-size",
+            type=int,
+            default=InferenceArgs.tp_size,
+            help="The tensor parallelism size.",
+        )
+        parser.add_argument(
+            "--sequence-parallel-size",
+            "--sp-size",
+            type=int,
+            default=InferenceArgs.sp_size,
+            help="The sequence parallelism size.",
+        )
+        parser.add_argument(
+            "--dist-timeout",
+            type=int,
+            default=InferenceArgs.dist_timeout,
+            help="Set timeout for torch.distributed initialization.",
+        )
+
+        # Video generation parameters
+        parser.add_argument(
+            "--height",
+            type=int,
+            default=InferenceArgs.height,
+            help="Height of generated video",
+        )
+        parser.add_argument(
+            "--width",
+            type=int,
+            default=InferenceArgs.width,
+            help="Width of generated video",
+        )
+        parser.add_argument(
+            "--num-frames",
+            type=int,
+            default=InferenceArgs.num_frames,
+            help="Number of frames to generate",
+        )
+        parser.add_argument(
+            "--num-inference-steps",
+            type=int,
+            default=InferenceArgs.num_inference_steps,
+            help="Number of inference steps",
+        )
+        parser.add_argument(
+            "--guidance-scale",
+            type=float,
+            default=InferenceArgs.guidance_scale,
+            help="Guidance scale for classifier-free guidance",
+        )
+        parser.add_argument(
+            "--guidance-rescale",
+            type=float,
+            default=InferenceArgs.guidance_rescale,
+            help="Guidance rescale for classifier-free guidance",
+        )
+        parser.add_argument(
+            "--embedded-cfg-scale",
+            type=float,
+            default=InferenceArgs.embedded_cfg_scale,
+            help="Embedded CFG scale",
+        )
+        parser.add_argument(
+            "--flow-shift",
+            "--shift",
+            type=float,
+            default=InferenceArgs.flow_shift,
+            help="Flow shift parameter",
+        )
+        parser.add_argument(
+            "--output-type",
+            type=str,
+            default=InferenceArgs.output_type,
+            choices=["pil"],
+            help="Output type for the generated video",
+        )
+
+        parser.add_argument(
+            "--precision",
+            type=str,
+            default=InferenceArgs.precision,
+            choices=["fp32", "fp16", "bf16"],
+            help="Precision for the model",
+        )
+
+        # VAE configuration
+        parser.add_argument(
+            "--vae-precision",
+            type=str,
+            default=InferenceArgs.vae_precision,
+            choices=["fp32", "fp16", "bf16"],
+            help="Precision for VAE",
+        )
+        parser.add_argument(
+            "--vae-tiling",
+            action="store_true",
+            default=InferenceArgs.vae_tiling,
+            help="Enable VAE tiling",
+        )
+        parser.add_argument(
+            "--vae-sp",
+            action="store_true",
+            help="Enable VAE spatial parallelism",
+        )
+
+        parser.add_argument(
+            "--text-encoder-precision",
+            type=str,
+            default=InferenceArgs.text_encoder_precision,
+            choices=["fp32", "fp16", "bf16"],
+            help="Precision for text encoder",
+        )
+        parser.add_argument(
+            "--text-len",
+            type=int,
+            default=InferenceArgs.text_len,
+            help="Maximum text length",
+        )
+
+        # Image encoder config
+        parser.add_argument(
+            "--image-encoder-precision",
+            type=str,
+            default=InferenceArgs.image_encoder_precision,
+            choices=["fp32", "fp16", "bf16"],
+            help="Precision for image encoder",
+        )
+
+        # Secondary text encoder
+
+        parser.add_argument(
+            "--text-encoder-precision-2",
+            type=str,
+            default=InferenceArgs.text_encoder_precision_2,
+            choices=["fp32", "fp16", "bf16"],
+            help="Precision for secondary text encoder",
+        )
+        parser.add_argument(
+            "--text-len-2",
+            type=int,
+            default=InferenceArgs.text_len_2,
+            help="Maximum secondary text length",
+        )
+
+        # Flow Matching parameters
+        parser.add_argument(
+            "--flow-solver",
+            type=str,
+            default=InferenceArgs.flow_solver,
+            help="Solver for flow matching",
+        )
+        parser.add_argument(
+            "--denoise-type",
+            type=str,
+            default=InferenceArgs.denoise_type,
+            help="Denoise type for noised inputs",
+        )
+
+        # STA (Spatial-Temporal Attention) parameters
+        parser.add_argument(
+            "--mask-strategy-file-path",
+            type=str,
+            help="Path to mask strategy JSON file for STA",
+        )
+        parser.add_argument(
+            "--enable-torch-compile",
+            action="store_true",
+            help=
+            "Use torch.compile for speeding up STA inference without teacache",
+        )
+
+        # Scheduler options
+        parser.add_argument(
+            "--scheduler-type",
+            type=str,
+            default=InferenceArgs.scheduler_type,
+            help="Type of scheduler to use",
+        )
+
+        # HunYuan specific parameters
+        parser.add_argument(
+            "--neg-prompt",
+            type=str,
+            default=InferenceArgs.neg_prompt,
+            help="Negative prompt for sampling",
+        )
+        parser.add_argument(
+            "--num-videos",
+            type=int,
+            default=InferenceArgs.num_videos,
+            help="Number of videos to generate per prompt",
+        )
+        parser.add_argument(
+            "--fps",
+            type=int,
+            default=InferenceArgs.fps,
+            help="Frames per second for output video",
+        )
+        parser.add_argument(
+            "--use-cpu-offload",
+            action="store_true",
+            help="Use CPU offload for the model load",
+        )
+        parser.add_argument(
+            "--disable-autocast",
+            action="store_true",
+            help=
+            "Disable autocast for denoising loop and vae decoding in pipeline sampling",
+        )
+
+        # Logging
+        parser.add_argument(
+            "--log-level",
+            type=str,
+            default=InferenceArgs.log_level,
+            help="The logging level of all loggers.",
+        )
+
+        # Inference parameters
+        prompt_group = parser.add_mutually_exclusive_group(required=True)
+        prompt_group.add_argument(
+            "--prompt",
+            type=str,
+            help="Text prompt for video generation",
+        )
+        prompt_group.add_argument(
+            "--prompt-path",
+            type=str,
+            help="Path to a text file containing the prompt",
+        )
+
+        parser.add_argument("--image-path",
+                            type=str,
+                            help="Path to the image for I2V generation")
+
+        parser.add_argument(
+            "--output-path",
+            type=str,
+            default=InferenceArgs.output_path,
+            help="Directory to save generated videos",
+        )
+        parser.add_argument(
+            "--seed",
+            type=int,
+            default=InferenceArgs.seed,
+            help="Random seed for reproducibility",
+        )
+
+        return parser
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "InferenceArgs":
+        args.tp_size = args.tensor_parallel_size
+        args.sp_size = args.sequence_parallel_size
+        args.flow_shift = getattr(args, "shift", args.flow_shift)
+
+        # Get all fields from the dataclass
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+
+        # Create a dictionary of attribute values, with defaults for missing attributes
+        kwargs = {}
+        for attr in attrs:
+            # Handle renamed attributes or those with multiple CLI names
+            if attr == 'tp_size' and hasattr(args, 'tensor_parallel_size'):
+                kwargs[attr] = args.tensor_parallel_size
+            elif attr == 'sp_size' and hasattr(args, 'sequence_parallel_size'):
+                kwargs[attr] = args.sequence_parallel_size
+            elif attr == 'flow_shift' and hasattr(args, 'shift'):
+                kwargs[attr] = args.shift
+            # Use getattr with default value from the dataclass for potentially missing attributes
+            else:
+                default_value = getattr(cls, attr, None)
+                kwargs[attr] = getattr(args, attr, default_value)
+
+        return cls(**kwargs)
+
+    def check_inference_args(self) -> None:
+        """Validate inference arguments for consistency"""
+
+        # Validate VAE spatial parallelism with VAE tiling
+        if self.vae_sp and not self.vae_tiling:
+            raise ValueError(
+                "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True."
+            )
+        if self.prompt_path and not self.prompt_path.endswith(".txt"):
+            raise ValueError("prompt_path must be a text file")
+
+
+_inference_args = None
+
+
+def prepare_inference_args(argv: List[str]) -> InferenceArgs:
+    """
+    Prepare the inference arguments from the command line arguments.
+
+    Args:
+        argv: The command line arguments. Typically, it should be `sys.argv[1:]`
+            to ensure compatibility with `parse_args` when no arguments are passed.
+
+    Returns:
+        The inference arguments.
+    """
+    parser = FlexibleArgumentParser()
+    InferenceArgs.add_cli_args(parser)
+    raw_args = parser.parse_args(argv)
+    inference_args = InferenceArgs.from_cli_args(raw_args)
+    inference_args.check_inference_args()
+    global _inference_args
+    _inference_args = inference_args
+    return inference_args
+
+
+def get_inference_args() -> InferenceArgs:
+    global _inference_args
+    if _inference_args is None:
+        raise ValueError("Inference arguments not set")
+    return _inference_args
+
+
+class DeprecatedAction(argparse.Action):
+
+    def __init__(self, option_strings, dest, nargs=0, **kwargs):
+        super().__init__(option_strings, dest, nargs=nargs, **kwargs)
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        raise ValueError(self.help)

fastvideo/v1/inference_engine.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Inference module for diffusion models.
+
+This module provides classes and functions for running inference with diffusion models.
+"""
+
+import time
+from typing import Any, Dict
+
+import torch
+
+from fastvideo.v1.inference_args import InferenceArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch,
+                                    build_pipeline)
+# TODO(will): remove, check if this is hunyuan specific
+from fastvideo.v1.utils import align_to
+
+logger = init_logger(__name__)
+
+
+class InferenceEngine:
+    """
+    Engine for running inference with diffusion models.
+    """
+
+    def __init__(
+        self,
+        pipeline: ComposedPipelineBase,
+        inference_args: InferenceArgs,
+    ):
+        """
+        Initialize the inference engine.
+        
+        Args:
+            pipeline: The pipeline to use for inference.
+            inference_args: The inference arguments.
+            default_negative_prompt: The default negative prompt to use.
+        """
+        self.pipeline = pipeline
+        self.inference_args = inference_args
+
+    @classmethod
+    def create_engine(
+        cls,
+        inference_args: InferenceArgs,
+    ) -> "InferenceEngine":
+        """
+        Create an inference engine with the specified arguments.
+        
+        Args:
+            inference_args: The inference arguments.
+            model_loader_cls: The model loader class to use. If None, it will be
+                determined from the model type.
+            pipeline_type: The type of pipeline to create. If None, it will be
+                determined from the model type.
+                
+        Returns:
+            The created inference engine.
+            
+        Raises:
+            ValueError: If the model type is not recognized or if the pipeline type
+                is not recognized.
+        """
+
+        logger.info("Building pipeline...")
+
+        # TODO(will): I don't really like this api.
+        # it should be something closer to pipeline_cls.from_pretrained(...)
+        # this way for training we can just do pipeline_cls.from_pretrained(
+        # checkpoint_path) and have it handle everything.
+        # TODO(Peiyuan): Then maybe we should only pass in model path and device, not the entire inference args?
+        pipeline = build_pipeline(inference_args)
+        logger.info("Pipeline Ready")
+
+        # Create the inference engine
+        return cls(pipeline, inference_args)
+
+    def run(
+        self,
+        prompt: str,
+        inference_args: InferenceArgs,
+    ) -> Dict[str, Any]:
+        """
+        Run inference with the pipeline.
+        
+        Args:
+            prompt: The prompt to use for generation.
+            negative_prompt: The negative prompt to use. If None, the default will be used.
+            seed: The random seed to use. If None, a random seed will be used.
+            **kwargs: Additional arguments to pass to the pipeline.
+            
+        Returns:
+            A dictionary containing the generated videos and metadata.
+        """
+        out_dict: Dict[str, Any] = dict()
+
+        num_videos_per_prompt = inference_args.num_videos
+        seed = inference_args.seed
+        height = inference_args.height
+        width = inference_args.width
+        video_length = inference_args.num_frames
+        negative_prompt = inference_args.neg_prompt
+        infer_steps = inference_args.num_inference_steps
+        guidance_scale = inference_args.guidance_scale
+        flow_shift = inference_args.flow_shift
+        embedded_guidance_scale = inference_args.embedded_cfg_scale
+        image_path = inference_args.image_path
+
+        # ========================================================================
+        # Arguments: target_width, target_height, target_video_length
+        # ========================================================================
+        if width <= 0 or height <= 0 or video_length <= 0:
+            raise ValueError(
+                f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}"
+            )
+        if (video_length - 1) % 4 != 0:
+            raise ValueError(
+                f"`video_length-1` must be a multiple of 4, got {video_length}")
+
+        target_height = align_to(height, 16)
+        target_width = align_to(width, 16)
+        target_video_length = video_length
+
+        out_dict["size"] = (target_height, target_width, target_video_length)
+
+        # ========================================================================
+        # Arguments: prompt, new_prompt, negative_prompt
+        # ========================================================================
+        if not isinstance(prompt, str):
+            raise TypeError(
+                f"`prompt` must be a string, but got {type(prompt)}")
+        prompt = prompt.strip()
+
+        # negative prompt
+        if negative_prompt is not None:
+            negative_prompt = negative_prompt.strip()
+
+        # TODO(PY): move to hunyuan stage
+        latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+
+        # ========================================================================
+        # Print infer args
+        # ========================================================================
+        debug_str = f"""
+                        height: {target_height}
+                         width: {target_width}
+                  video_length: {target_video_length}
+                        prompt: {prompt}
+                    neg_prompt: {negative_prompt}
+                          seed: {seed}
+                   infer_steps: {infer_steps}
+         num_videos_per_prompt: {num_videos_per_prompt}
+                guidance_scale: {guidance_scale}
+                      n_tokens: {n_tokens}
+                    flow_shift: {flow_shift}
+       embedded_guidance_scale: {embedded_guidance_scale}"""
+        logger.info(debug_str)
+        # return
+        # sp_group = get_sp_group()
+        # local_rank = sp_group.rank
+        device = torch.device(inference_args.device_str)
+        batch = ForwardBatch(
+            image_path=image_path,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_videos_per_prompt=num_videos_per_prompt,
+            height=inference_args.height,
+            width=inference_args.width,
+            num_frames=inference_args.num_frames,
+            num_inference_steps=inference_args.num_inference_steps,
+            guidance_scale=inference_args.guidance_scale,
+            # generator=generator,
+            eta=0.0,
+            n_tokens=n_tokens,
+            data_type="video" if inference_args.num_frames > 1 else "image",
+            device=device,
+            extra={},  # Any additional parameters
+        )
+
+        print('===============================================')
+        print(batch)
+        print('===============================================')
+        print('===============================================')
+        print(inference_args)
+
+        # ========================================================================
+        # Pipeline inference
+        # ========================================================================
+        start_time = time.time()
+        samples = self.pipeline.forward(
+            batch=batch,
+            inference_args=inference_args,
+        ).output
+        # TODO(will): fix and move to hunyuan stage
+        # out_dict["seeds"] = batch.seeds
+        out_dict["samples"] = samples
+        out_dict["prompts"] = prompt
+
+        gen_time = time.time() - start_time
+        logger.info("Success, time: %s", gen_time)
+
+        return out_dict

fastvideo/v1/layers/activation.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/activation.py
+"""Custom activation functions."""
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# TODO (will): remove this dependency
+from fastvideo.v1.layers.custom_op import CustomOp
+from fastvideo.v1.platforms import current_platform
+
+
+@CustomOp.register("silu_and_mul")
+class SiluAndMul(CustomOp):
+    """An activation function for SwiGLU.
+
+    The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.
+
+    Shapes:
+        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
+        return: (num_tokens, d) or (batch_size, seq_len, d)
+    """
+
+    def __init__(self):
+        super().__init__()
+        if current_platform.is_cuda_alike() or current_platform.is_cpu():
+            self.op = torch.ops._C.silu_and_mul
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        d = x.shape[-1] // 2
+        return F.silu(x[..., :d]) * x[..., d:]
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        output_shape = (x.shape[:-1] + (d, ))
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        self.op(out, x)
+        return out
+
+
+@CustomOp.register("gelu_and_mul")
+class GeluAndMul(CustomOp):
+    """An activation function for GeGLU.
+
+    The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2.
+
+    Shapes:
+        x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d)
+        return: (batch_size, seq_len, d) or (num_tokens, d)
+    """
+
+    def __init__(self, approximate: str = "none"):
+        super().__init__()
+        self.approximate = approximate
+        if approximate not in ("none", "tanh"):
+            raise ValueError(f"Unknown approximate mode: {approximate}")
+        if current_platform.is_cuda_alike() or current_platform.is_cpu():
+            if approximate == "none":
+                self.op = torch.ops._C.gelu_and_mul
+            elif approximate == "tanh":
+                self.op = torch.ops._C.gelu_tanh_and_mul
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        d = x.shape[-1] // 2
+        return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:]
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        output_shape = (x.shape[:-1] + (d, ))
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        self.op(out, x)
+        return out
+
+    def extra_repr(self) -> str:
+        return f'approximate={repr(self.approximate)}'
+
+
+@CustomOp.register("gelu_new")
+class NewGELU(CustomOp):
+
+    def __init__(self):
+        super().__init__()
+        if current_platform.is_cuda_alike() or current_platform.is_cpu():
+            self.op = torch.ops._C.gelu_new
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        c = math.sqrt(2.0 / math.pi)
+        return 0.5 * x * (1.0 + torch.tanh(c *
+                                           (x + 0.044715 * torch.pow(x, 3.0))))
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        out = torch.empty_like(x)
+        self.op(out, x)
+        return out
+
+    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
+        return self.op(x)
+
+
+@CustomOp.register("quick_gelu")
+class QuickGELU(CustomOp):
+    # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90
+    def __init__(self):
+        super().__init__()
+        if current_platform.is_cuda_alike() or current_platform.is_cpu():
+            self.op = torch.ops._C.gelu_quick
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        return x * torch.sigmoid(1.702 * x)
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        out = torch.empty_like(x)
+        self.op(out, x)
+        return out
+
+
+_ACTIVATION_REGISTRY = {
+    "gelu": nn.GELU,
+    "gelu_new": NewGELU,
+    "gelu_pytorch_tanh": lambda: nn.GELU(approximate="tanh"),
+    "relu": nn.ReLU,
+    "silu": nn.SiLU,
+    "quick_gelu": QuickGELU,
+}
+
+
+def get_act_fn(act_fn_name: str) -> nn.Module:
+    """Get an activation function by name."""
+    act_fn_name = act_fn_name.lower()
+    if act_fn_name not in _ACTIVATION_REGISTRY:
+        raise ValueError(
+            f"Activation function {act_fn_name!r} is not supported.")
+
+    return _ACTIVATION_REGISTRY[act_fn_name]()
+
+
+_ACTIVATION_AND_MUL_REGISTRY = {
+    "gelu": GeluAndMul,
+    "silu": SiluAndMul,
+}
+
+
+def get_act_and_mul_fn(act_fn_name: str) -> nn.Module:
+    """Get an activation-and-mul (i.e. SiluAndMul) function by name."""
+    act_fn_name = act_fn_name.lower()
+    if act_fn_name not in _ACTIVATION_AND_MUL_REGISTRY:
+        raise ValueError(
+            f"Activation function {act_fn_name!r} is not supported.")
+
+    return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]()

fastvideo/v1/layers/custom_op.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/custom_op.py
+
+from typing import Any, Callable, Dict, Type
+
+import torch.nn as nn
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class CustomOp(nn.Module):
+    """
+    Base class for custom ops.
+    Dispatches the forward method to the appropriate backend.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._forward_method = self.dispatch_forward()
+
+    def forward(self, *args, **kwargs) -> Any:
+        return self._forward_method(*args, **kwargs)
+
+    def forward_native(self, *args, **kwargs) -> Any:
+        """PyTorch-native implementation of the forward method.
+        This method is optional. If implemented, it can be used with compilers
+        such as torch.compile or PyTorch XLA. Also, it can be used for testing
+        purposes.
+        """
+        raise NotImplementedError
+
+    def forward_cuda(self, *args, **kwargs) -> Any:
+        raise NotImplementedError
+
+    def forward_cpu(self, *args, **kwargs) -> Any:
+        # By default, we assume that CPU ops are compatible with CUDA ops.
+        return self.forward_cuda(*args, **kwargs)
+
+    def forward_tpu(self, *args, **kwargs) -> Any:
+        # By default, we assume that TPU ops are compatible with the
+        # PyTorch-native implementation.
+        # NOTE(woosuk): This is a placeholder for future extensions.
+        return self.forward_native(*args, **kwargs)
+
+    def forward_oot(self, *args, **kwargs) -> Any:
+        # By default, we assume that OOT ops are compatible with the
+        # PyTorch-native implementation.
+        return self.forward_native(*args, **kwargs)
+
+    def dispatch_forward(self) -> Callable:
+        # NOTE(woosuk): Here we assume that vLLM was built for only one
+        # specific backend. Currently, we do not support dynamic dispatching.
+        enabled = self.enabled()
+
+        if not enabled:
+            return self.forward_native
+
+        return self.forward_cuda
+
+    @classmethod
+    def enabled(cls) -> bool:
+        # since we are not using Inductor, we always return True
+        return True
+
+    @staticmethod
+    def default_on() -> bool:
+        """
+        On by default if level < CompilationLevel.PIECEWISE
+        Specifying 'all' or 'none' in custom_op takes precedence.
+        """
+        raise NotImplementedError
+
+    # Dictionary of all custom ops (classes, indexed by registered name).
+    # To check if an op with a name is enabled, call .enabled() on the class.
+    # Examples:
+    # - MyOp.enabled()
+    # - op_registry["my_op"].enabled()
+    op_registry: Dict[str, Type['CustomOp']] = {}
+
+    # Decorator to register custom ops.
+    @classmethod
+    def register(cls, name: str) -> Callable:
+
+        def decorator(op_cls):
+            assert name not in cls.op_registry, f"Duplicate op name: {name}"
+            op_cls.name = name
+            cls.op_registry[name] = op_cls
+            return op_cls
+
+        return decorator

fastvideo/v1/layers/layernorm.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/layernorm.py
+"""Custom normalization layers."""
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from fastvideo.v1.layers.custom_op import CustomOp
+
+
+@CustomOp.register("rms_norm")
+class RMSNorm(CustomOp):
+    """Root mean square normalization.
+
+    Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight.
+    Refer to https://arxiv.org/abs/1910.07467
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        eps: float = 1e-6,
+        dtype: torch.dtype = torch.float32,
+        var_hidden_size: Optional[int] = None,
+        has_weight: bool = True,
+    ) -> None:
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.variance_epsilon = eps
+        self.variance_size_override = (None if var_hidden_size == hidden_size
+                                       else var_hidden_size)
+        self.has_weight = has_weight
+
+        self.weight = torch.ones(hidden_size)
+        if self.has_weight:
+            self.weight = nn.Parameter(self.weight)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """PyTorch-native implementation equivalent to forward()."""
+        orig_dtype = x.dtype
+        x = x.to(torch.float32)
+        if residual is not None:
+            x = x + residual.to(torch.float32)
+            residual = x.to(orig_dtype)
+
+        hidden_size = x.shape[-1]
+        if hidden_size != self.hidden_size:
+            raise ValueError("Expected hidden_size to be "
+                             f"{self.hidden_size}, but found: {hidden_size}")
+
+        if self.variance_size_override is None:
+            x_var = x
+        else:
+            if hidden_size < self.variance_size_override:
+                raise ValueError(
+                    "Expected hidden_size to be at least "
+                    f"{self.variance_size_override}, but found: {hidden_size}")
+
+            x_var = x[:, :, :self.variance_size_override]
+
+        variance = x_var.pow(2).mean(dim=-1, keepdim=True)
+
+        x = x * torch.rsqrt(variance + self.variance_epsilon)
+        x = x.to(orig_dtype)
+        if self.has_weight:
+            x = x * self.weight
+        if residual is None:
+            return x
+        else:
+            return x, residual
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if self.variance_size_override is not None:
+            return self.forward_native(x, residual)
+
+        from vllm import _custom_ops as ops
+
+        if residual is not None:
+            ops.fused_add_rms_norm(
+                x,
+                residual,
+                self.weight.data,
+                self.variance_epsilon,
+            )
+            return x, residual
+        out = torch.empty_like(x)
+        ops.rms_norm(
+            out,
+            x,
+            self.weight.data,
+            self.variance_epsilon,
+        )
+        return out
+
+    def extra_repr(self) -> str:
+        s = f"hidden_size={self.weight.data.size(0)}"
+        s += f", eps={self.variance_epsilon}"
+        return s
+
+
+class ScaleResidual(nn.Module):
+    """
+    Applies gated residual connection.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, residual: torch.Tensor, x: torch.Tensor,
+                gate: torch.Tensor) -> torch.Tensor:
+        """Apply gated residual connection."""
+        return residual + x * gate
+
+
+class ScaleResidualLayerNormScaleShift(nn.Module):
+    """
+    Fused operation that combines:
+    1. Gated residual connection
+    2. LayerNorm
+    3. Scale and shift operations
+    
+    This reduces memory bandwidth by combining memory-bound operations.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        norm_type: str = "rms",
+        eps: float = 1e-6,
+        elementwise_affine: bool = False,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+        if norm_type == "rms":
+            self.norm = RMSNorm(hidden_size,
+                                has_weight=elementwise_affine,
+                                eps=eps,
+                                dtype=dtype)
+        elif norm_type == "layer":
+            self.norm = nn.LayerNorm(hidden_size,
+                                     elementwise_affine=elementwise_affine,
+                                     eps=eps,
+                                     dtype=dtype)
+        else:
+            raise NotImplementedError(f"Norm type {norm_type} not implemented")
+
+    def forward(self, residual: torch.Tensor, x: torch.Tensor,
+                gate: torch.Tensor, shift: torch.Tensor,
+                scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Apply gated residual connection, followed by layernorm and 
+        scale/shift in a single fused operation.
+        
+        Returns:
+            Tuple containing:
+            - normalized and modulated output
+            - residual value (value after residual connection 
+              but before normalization)
+        """
+        # Apply residual connection with gating
+        residual_output = residual + x * gate
+        # Apply normalization
+        normalized = self.norm(residual_output)
+        # Apply scale and shift
+        modulated = normalized * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        return modulated, residual_output
+
+
+class LayerNormScaleShift(nn.Module):
+    """
+    Fused operation that combines LayerNorm with scale and shift operations.
+    This reduces memory bandwidth by combining memory-bound operations.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        norm_type: str = "rms",
+        eps: float = 1e-6,
+        elementwise_affine: bool = False,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+        if norm_type == "rms":
+            self.norm = RMSNorm(hidden_size,
+                                has_weight=elementwise_affine,
+                                eps=eps)
+        elif norm_type == "layer":
+            self.norm = nn.LayerNorm(hidden_size,
+                                     elementwise_affine=elementwise_affine,
+                                     eps=eps,
+                                     dtype=dtype)
+        else:
+            raise NotImplementedError(f"Norm type {norm_type} not implemented")
+
+    def forward(self, x: torch.Tensor, shift: torch.Tensor,
+                scale: torch.Tensor) -> torch.Tensor:
+        """Apply ln followed by scale and shift in a single fused operation."""
+        normalized = self.norm(x)
+        return normalized * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)

fastvideo/v1/layers/linear.py

@@ -0,0 +1,955 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/linear.py
+
+from abc import abstractmethod
+from typing import Optional, Union
+
+import torch
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter
+# TODO(will): remove this import by copying the definition from vLLM then
+# manually import each quantization method we want to use. Refer to SGLang
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig, QuantizeMethodBase)
+
+from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank,
+                                      get_tensor_model_parallel_world_size,
+                                      split_tensor_along_last_dim,
+                                      tensor_model_parallel_all_gather,
+                                      tensor_model_parallel_all_reduce)
+from fastvideo.v1.logger import init_logger
+# yapf: disable
+from fastvideo.v1.models.parameter import (BasevLLMParameter,
+                                           BlockQuantScaleParameter,
+                                           PackedColumnParameter,
+                                           PackedvLLMParameter,
+                                           PerTensorScaleParameter,
+                                           RowvLLMParameter)
+# yapf: enable
+from fastvideo.v1.models.utils import set_weight_attrs
+
+logger = init_logger(__name__)
+
+WEIGHT_LOADER_V2_SUPPORTED = [
+    "CompressedTensorsLinearMethod", "AWQMarlinLinearMethod", "AWQLinearMethod",
+    "GPTQMarlinLinearMethod", "Fp8LinearMethod", "MarlinLinearMethod",
+    "QQQLinearMethod", "GPTQMarlin24LinearMethod", "TPUInt8LinearMethod",
+    "GPTQLinearMethod", "FBGEMMFp8LinearMethod", "ModelOptFp8LinearMethod",
+    "IPEXAWQLinearMethod", "IPEXGPTQLinearMethod", "HQQMarlinMethod",
+    "QuarkLinearMethod"
+]
+
+
+def adjust_scalar_to_fused_array(
+        param: torch.Tensor, loaded_weight: torch.Tensor,
+        shard_id: Union[str, int]) -> tuple[torch.Tensor, torch.Tensor]:
+    """For fused modules (QKV and MLP) we have an array of length
+    N that holds 1 scale for each "logical" matrix. So the param
+    is an array of length N. The loaded_weight corresponds to 
+    one of the shards on disk. Here, we slice the param based on 
+    the shard_id for loading.
+    """
+    qkv_idxs = {"q": 0, "k": 1, "v": 2}
+
+    if isinstance(shard_id, str):
+        shard_id = qkv_idxs[shard_id]
+    elif not isinstance(shard_id, int):
+        raise ValueError(f"Unknown Shard Id {shard_id}")
+
+    # AutoFP8 scales do not have a shape
+    # compressed-tensors scales do have a shape
+    if len(loaded_weight.shape) != 0:
+        assert loaded_weight.shape[0] == 1
+        loaded_weight = loaded_weight[0]
+
+    return param[shard_id], loaded_weight
+
+
+class LinearMethodBase(QuantizeMethodBase):
+    """Base class for different (maybe quantized) linear methods."""
+
+    @abstractmethod
+    def create_weights(self, layer: torch.nn.Module,
+                       input_size_per_partition: int,
+                       output_partition_sizes: list[int], input_size: int,
+                       output_size: int, params_dtype: torch.dtype,
+                       **extra_weight_attrs) -> None:
+        """Create weights for a linear layer. 
+           The weights will be set as attributes of the layer.
+
+        Args:
+            layer: The layer that is using the LinearMethodBase factory.
+            input_size_per_partition: Size of the weight input dim on rank X.
+            output_partition_sizes: Sizes of the output dim of each logical 
+                weight on rank X. E.g., output_partition_sizes for QKVLinear
+                is a list contains the width of Wq, Wk, Wv on rank X.
+            input_size: Size of the input dim of the weight across all ranks.
+            output_size: Size of the output dim of the weight across all ranks.
+            params_dtype: Datatype of the parameters.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Apply the weights in layer to the input tensor.
+        Expects create_weights to have been called before on the layer."""
+        raise NotImplementedError
+
+
+class UnquantizedLinearMethod(LinearMethodBase):
+    """Linear method without quantization."""
+
+    def create_weights(self, layer: torch.nn.Module,
+                       input_size_per_partition: int,
+                       output_partition_sizes: list[int], input_size: int,
+                       output_size: int, params_dtype: torch.dtype,
+                       **extra_weight_attrs) -> None:
+        weight = Parameter(torch.empty(sum(output_partition_sizes),
+                                       input_size_per_partition,
+                                       dtype=params_dtype),
+                           requires_grad=False)
+        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
+        layer.register_parameter("weight", weight)
+        set_weight_attrs(weight, extra_weight_attrs)
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        return F.linear(x, layer.weight, bias)
+
+
+class LinearBase(torch.nn.Module):
+    """Base linear layer.
+
+    Args:
+        input_size: input dimension of the linear layer.
+        output_size: output dimension of the linear layer.
+        bias: If true, add bias.
+        skip_bias_add: If true, skip adding bias but instead return it.
+        params_dtype: Data type for the parameters.
+        quant_config: Quantization configure.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int,
+        skip_bias_add: bool = False,
+        params_dtype: Optional[torch.dtype] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        # Keep input parameters
+        self.input_size = input_size
+        self.output_size = output_size
+        self.skip_bias_add = skip_bias_add
+        if params_dtype is None:
+            params_dtype = torch.get_default_dtype()
+        self.params_dtype = params_dtype
+        if quant_config is None:
+            self.quant_method: Optional[
+                QuantizeMethodBase] = UnquantizedLinearMethod()
+        else:
+            self.quant_method = quant_config.get_quant_method(self,
+                                                              prefix=prefix)
+
+    def forward(self,
+                x: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:
+        raise NotImplementedError
+
+
+class ReplicatedLinear(LinearBase):
+    """Replicated linear layer.
+
+    Args:
+        input_size: input dimension of the linear layer.
+        output_size: output dimension of the linear layer.
+        bias: If true, add bias.
+        skip_bias_add: If true, skip adding bias but instead return it.
+        params_dtype: Data type for the parameters.
+        quant_config: Quantization configure.
+        prefix: The name of the layer in the state dict, including all parents
+                        (e.g. model.layers.0.qkv_proj)
+    """
+
+    def __init__(self,
+                 input_size: int,
+                 output_size: int,
+                 bias: bool = True,
+                 skip_bias_add: bool = False,
+                 params_dtype: Optional[torch.dtype] = None,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        super().__init__(input_size,
+                         output_size,
+                         skip_bias_add,
+                         params_dtype,
+                         quant_config,
+                         prefix=prefix)
+
+        # All the linear layer supports quant method.
+        assert self.quant_method is not None
+        self.quant_method.create_weights(self,
+                                         self.input_size, [self.output_size],
+                                         self.input_size,
+                                         self.output_size,
+                                         self.params_dtype,
+                                         weight_loader=self.weight_loader)
+
+        if bias:
+            self.bias = Parameter(
+                torch.empty(self.output_size, dtype=self.params_dtype))
+            set_weight_attrs(self.bias, {
+                "output_dim": 0,
+                "weight_loader": self.weight_loader,
+            })
+        else:
+            self.register_parameter("bias", None)
+
+    def weight_loader(self, param: Parameter,
+                      loaded_weight: torch.Tensor) -> None:
+        # If the weight on disk does not have a shape, give it one
+        # (such scales for AutoFp8).
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+
+        assert param.size() == loaded_weight.size(), (
+            f"Tried to load weights of size {loaded_weight.size()}"
+            f"to a parameter of size {param.size()}")
+        param.data.copy_(loaded_weight)
+
+    def forward(self,
+                x: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:
+        bias = self.bias if not self.skip_bias_add else None
+        assert self.quant_method is not None
+        output = self.quant_method.apply(self, x, bias)
+        output_bias = self.bias if self.skip_bias_add else None
+        return output, output_bias
+
+    def extra_repr(self) -> str:
+        s = f"in_features={self.input_size}"
+        s += f", output_features={self.output_size}"
+        s += f", bias={self.bias is not None}"
+        return s
+
+
+class ColumnParallelLinear(LinearBase):
+    """Linear layer with column parallelism.
+
+    The linear layer is defined as Y = XA + b. A is parallelized along
+    its second dimension as A = [A_1, ..., A_p].
+
+    Args:
+        input_size: first dimension of matrix A.
+        output_size: second dimension of matrix A.
+        bias: If true, add bias.
+        gather_output: If true, call all-gather on output and make Y available
+                       to all GPUs, otherwise, every GPU will have its output
+                       which is Y_i = XA_i
+        skip_bias_add: This was added to enable performance optimizations where
+                       bias can be fused with other element-wise operations. we
+                       skip adding bias but instead return it.
+        params_dtype: Data type for the parameters.
+        quant_config: Quantization configure.
+        output_sizes: list of output sizes packed into one output, like for QKV
+                       the list would be size 3.
+        prefix: The name of the layer in the state dict, including all parents
+                        (e.g. model.layers.0.qkv_proj) 
+    """
+
+    def __init__(self,
+                 input_size: int,
+                 output_size: int,
+                 bias: bool = True,
+                 gather_output: bool = False,
+                 skip_bias_add: bool = False,
+                 params_dtype: Optional[torch.dtype] = None,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 output_sizes: Optional[list[int]] = None,
+                 prefix: str = ""):
+        # Divide the weight matrix along the last dimension.
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.input_size_per_partition = input_size
+        self.output_size_per_partition = divide(output_size, self.tp_size)
+        self.output_partition_sizes = [self.output_size_per_partition]
+        # If QKV or MergedColumn, use output size of each partition.
+        if hasattr(self, "output_sizes"):
+            self.output_partition_sizes = [
+                divide(output_size, self.tp_size)
+                for output_size in self.output_sizes
+            ]
+
+        super().__init__(input_size, output_size, skip_bias_add, params_dtype,
+                         quant_config, prefix)
+
+        self.gather_output = gather_output
+
+        if output_sizes is None:
+            output_sizes = [output_size]
+
+        assert self.quant_method is not None
+        self.quant_method.create_weights(
+            layer=self,
+            input_size_per_partition=self.input_size_per_partition,
+            output_partition_sizes=self.output_partition_sizes,
+            input_size=self.input_size,
+            output_size=self.output_size,
+            params_dtype=self.params_dtype,
+            weight_loader=(
+                self.weight_loader_v2 if self.quant_method.__class__.__name__
+                in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader))
+        if bias:
+            self.bias = Parameter(
+                torch.empty(self.output_size_per_partition, dtype=params_dtype))
+            set_weight_attrs(self.bias, {
+                "output_dim": 0,
+                "weight_loader": self.weight_loader,
+            })
+        else:
+            self.register_parameter("bias", None)
+
+    def weight_loader(self, param: Parameter,
+                      loaded_weight: torch.Tensor) -> None:
+        tp_rank = get_tensor_model_parallel_rank()
+        output_dim = getattr(param, "output_dim", None)
+
+        is_sharded_weight = getattr(param, "is_sharded_weight", False)
+        is_sharded_weight = is_sharded_weight
+
+        param_data = param.data
+        if output_dim is not None and not is_sharded_weight:
+            shard_size = param_data.shape[output_dim]
+            start_idx = tp_rank * shard_size
+            loaded_weight = loaded_weight.narrow(output_dim, start_idx,
+                                                 shard_size)
+
+        # Special case for loading scales off disk, which often do not
+        # have a shape (such as in the case of AutoFP8).
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+    def weight_loader_v2(self, param: Parameter,
+                         loaded_weight: torch.Tensor) -> None:
+        # Special case for loading scales off disk, which often do not
+        # have a shape (such as in the case of AutoFP8).
+        if len(loaded_weight.shape) == 0:
+            assert loaded_weight.numel() == 1
+            loaded_weight = loaded_weight.reshape(1)
+        param.load_column_parallel_weight(loaded_weight=loaded_weight)
+
+    def forward(
+            self,
+            input_: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:
+        bias = self.bias if not self.skip_bias_add else None
+
+        # Matrix multiply.
+        assert self.quant_method is not None
+        output_parallel = self.quant_method.apply(self, input_, bias)
+        if self.gather_output:
+            # All-gather across the partitions.
+            output = tensor_model_parallel_all_gather(output_parallel)
+        else:
+            output = output_parallel
+        output_bias = self.bias if self.skip_bias_add else None
+        return output, output_bias
+
+    def extra_repr(self) -> str:
+        s = f"in_features={self.input_size}"
+        s += f", output_features={self.output_size_per_partition}"
+        s += f", bias={self.bias is not None}"
+        s += f", tp_size={get_tensor_model_parallel_world_size()}"
+        s += f", gather_output={self.gather_output}"
+        return s
+
+
+class MergedColumnParallelLinear(ColumnParallelLinear):
+    """Packed linear layers with column parallelism.
+
+    Similar to ColumnParallelLinear, but the weight matrix is concatenated
+    along the output dimension. When the weight matrix is loaded, the
+    different partitions are sharded separately.
+
+    Args:
+        input_size: input dimension of the linear layer.
+        output_sizes: list of output dimensions of the linear layer.
+        bias: If true, add bias.
+        gather_output: If true, call all-gather on output and make the output
+                       available to all GPUs, otherwise, every GPU will have
+                       its own output.
+        skip_bias_add: This was added to enable performance optimizations where
+                       bias can be fused with other element-wise operations. we
+                       skip adding bias but instead return it.
+        params_dtype: Data type for the parameters.
+        quant_config: Quantization configure.
+        prefix: The name of the layer in the state dict, including all parents
+                        (e.g. model.layers.0.qkv_proj)
+    """
+
+    def __init__(self,
+                 input_size: int,
+                 output_sizes: list[int],
+                 bias: bool = True,
+                 gather_output: bool = False,
+                 skip_bias_add: bool = False,
+                 params_dtype: Optional[torch.dtype] = None,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        self.output_sizes = output_sizes
+        tp_size = get_tensor_model_parallel_world_size()
+        assert all(output_size % tp_size == 0 for output_size in output_sizes)
+        super().__init__(input_size=input_size,
+                         output_size=sum(output_sizes),
+                         bias=bias,
+                         gather_output=gather_output,
+                         skip_bias_add=skip_bias_add,
+                         params_dtype=params_dtype,
+                         quant_config=quant_config,
+                         prefix=prefix)
+
+    def weight_loader(self,
+                      param: Parameter,
+                      loaded_weight: torch.Tensor,
+                      loaded_shard_id: Optional[int] = None) -> None:
+
+        param_data = param.data
+        output_dim = getattr(param, "output_dim", None)
+        # Special case for AQLM codebooks.
+        is_metadata = getattr(param, "is_metadata", False)
+        # Special case for per-tensor scale to load scalar into fused array.
+        needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
+
+        if loaded_shard_id is None:
+            # Loaded weight is already fused on disk (mlp).
+            # (e.g., Phi-3's gate_up_proj).
+            if output_dim is None:
+                if needs_scalar_to_array:
+                    param_data, loaded_weight = adjust_scalar_to_fused_array(
+                        param_data, loaded_weight, 0)
+
+                assert param_data.shape == loaded_weight.shape
+                param_data.copy_(loaded_weight)
+                return
+            current_shard_offset = 0
+            shard_offsets: list[tuple[int, int, int]] = []
+            for i, output_size in enumerate(self.output_sizes):
+                shard_offsets.append((i, current_shard_offset, output_size))
+                current_shard_offset += output_size
+            for shard_id, shard_offset, shard_size in shard_offsets:
+                loaded_weight_shard = loaded_weight.narrow(
+                    output_dim, shard_offset, shard_size)
+                self.weight_loader(param, loaded_weight_shard, shard_id)
+            return
+
+        assert loaded_shard_id < len(self.output_sizes)
+        tp_rank = get_tensor_model_parallel_rank()
+        tp_size = get_tensor_model_parallel_world_size()
+        if output_dim is not None:
+            shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
+            shard_size = self.output_sizes[loaded_shard_id] // tp_size
+
+            is_sharded_weight = getattr(param, "is_sharded_weight", False)
+            # bitsandbytes loads the weights of the specific portion
+            # no need to narrow
+            is_sharded_weight = is_sharded_weight
+
+            param_data = param_data.narrow(output_dim, shard_offset, shard_size)
+            start_idx = tp_rank * shard_size
+            if not is_sharded_weight:
+                loaded_weight = loaded_weight.narrow(output_dim, start_idx,
+                                                     shard_size)
+        # Special case for AQLM codebooks.
+        elif is_metadata:
+            # metadata indicates fixed size concatenated along dim 0
+            shard_size = loaded_weight.shape[0]
+            shard_offset = loaded_shard_id * shard_size
+            param_data = param_data.narrow(0, shard_offset, shard_size)
+
+        # Special case for per-tensor scales in fused case.
+        elif needs_scalar_to_array:
+            param_data, loaded_weight = adjust_scalar_to_fused_array(
+                param_data, loaded_weight, loaded_shard_id)
+
+        else:
+            ignore_warning = getattr(param, "ignore_warning", False)
+            if not ignore_warning:
+                logger.warning(
+                    "Loading a weight without `output_dim` attribute in "
+                    "MergedColumnParallelLinear, assume the weight is "
+                    "the same for all partitions.")
+
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+    def _load_fused_module_from_checkpoint(self, param: BasevLLMParameter,
+                                           loaded_weight: torch.Tensor) -> None:
+        """
+        Handle special case for models where MLP layers are already
+        fused on disk. In this case, we have no shard id. This function
+        determmines the shard id by splitting these layers and then calls
+        the weight loader using the shard id.
+
+        An example of a model with these fused layers:
+        https://huggingface.co/microsoft/Phi-3-mini-4k-instruct
+        """
+
+        current_shard_offset = 0
+        shard_offsets: list[tuple[int, int, int]] = []
+        for i, output_size in enumerate(self.output_sizes):
+            shard_offsets.append((i, current_shard_offset, output_size))
+            current_shard_offset += output_size
+
+        for shard_id, shard_offset, shard_size in shard_offsets:
+            # Special case for Quantization.
+            # If quantized, we need to adjust the offset and size to account
+            # for the packing.
+            if isinstance(
+                    param,
+                (PackedColumnParameter,
+                 PackedvLLMParameter)) and param.packed_dim == param.output_dim:
+                shard_size, shard_offset = \
+                    param.adjust_shard_indexes_for_packing(
+                    shard_size=shard_size, shard_offset=shard_offset)
+
+            loaded_weight_shard = loaded_weight.narrow(param.output_dim,
+                                                       shard_offset, shard_size)
+            self.weight_loader_v2(param, loaded_weight_shard, shard_id)
+
+    def weight_loader_v2(self,
+                         param: BasevLLMParameter,
+                         loaded_weight: torch.Tensor,
+                         loaded_shard_id: Optional[int] = None) -> None:
+        if loaded_shard_id is None:
+            if isinstance(param, PerTensorScaleParameter):
+                param.load_merged_column_weight(loaded_weight=loaded_weight,
+                                                shard_id=0)
+                return
+            elif type(param) in (RowvLLMParameter, BasevLLMParameter):
+                param.load_merged_column_weight(loaded_weight=loaded_weight)
+                return
+            # TODO: @dsikka - move to parameter.py
+            self._load_fused_module_from_checkpoint(param, loaded_weight)
+            return
+
+        assert loaded_shard_id < len(self.output_sizes)
+
+        tp_size = get_tensor_model_parallel_world_size()
+
+        if isinstance(param, BlockQuantScaleParameter):
+            from vllm.model_executor.layers.quantization.fp8 import (
+                Fp8LinearMethod, Fp8MoEMethod)
+            assert self.quant_method is not None
+            assert isinstance(self.quant_method,
+                              (Fp8LinearMethod, Fp8MoEMethod))
+            weight_block_size = self.quant_method.quant_config.weight_block_size
+            assert weight_block_size is not None
+            block_n, _ = weight_block_size[0], weight_block_size[1]
+            shard_offset = (
+                (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //
+                block_n) // tp_size
+            shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //
+                          block_n // tp_size)
+        else:
+            shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
+            shard_size = self.output_sizes[loaded_shard_id] // tp_size
+
+        param.load_merged_column_weight(loaded_weight=loaded_weight,
+                                        shard_id=loaded_shard_id,
+                                        shard_offset=shard_offset,
+                                        shard_size=shard_size)
+
+
+class QKVParallelLinear(ColumnParallelLinear):
+    """Linear layers for the attention's QKV transformation.
+
+    Linear layers for the linear transformation of the query, key, and value
+    vectors in the attention layer. The weight matrix is concatenated along
+    the output dimension. The layer is parallelized along the head dimension.
+    When the number of key/value heads is smaller than the number of query
+    heads (e.g., multi-query/grouped-query attention), the key/value head may
+    be replicated while the query heads are partitioned.
+
+    Args:
+        hidden_size: input hidden state size of the transformer.
+        head_size: size of each attention head.
+        total_num_heads: total number of attention query heads.
+        total_num_kv_heads: total number of attention key/value heads. If
+                            None, assume total_num_kv_heads = total_num_heads.
+        bias: If true, add bias.
+        skip_bias_add: This was added to enable performance optimizations where
+                       bias can be fused with other element-wise operations. we
+                       skip adding bias but instead return it.
+        params_dtype: Data type for the parameters.
+        quant_config: Quantization configure.
+        prefix: The name of the layer in the state dict, including all parents
+                        (e.g. model.layers.0.qkv_proj)
+    """
+
+    def __init__(self,
+                 hidden_size: int,
+                 head_size: int,
+                 total_num_heads: int,
+                 total_num_kv_heads: Optional[int] = None,
+                 bias: bool = True,
+                 skip_bias_add: bool = False,
+                 params_dtype: Optional[torch.dtype] = None,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        self.hidden_size = hidden_size
+        self.head_size = head_size
+        self.total_num_heads = total_num_heads
+        if total_num_kv_heads is None:
+            total_num_kv_heads = total_num_heads
+        self.total_num_kv_heads = total_num_kv_heads
+        # Divide the weight matrix along the last dimension.
+        tp_size = get_tensor_model_parallel_world_size()
+        self.num_heads = divide(self.total_num_heads, tp_size)
+        if tp_size >= self.total_num_kv_heads:
+            self.num_kv_heads = 1
+            self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
+        else:
+            self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
+            self.num_kv_head_replicas = 1
+        input_size = self.hidden_size
+        output_size = (self.num_heads +
+                       2 * self.num_kv_heads) * tp_size * self.head_size
+        self.output_sizes = [
+            self.num_heads * self.head_size * tp_size,  # q_proj
+            self.num_kv_heads * self.head_size * tp_size,  # k_proj
+            self.num_kv_heads * self.head_size * tp_size,  # v_proj 
+        ]
+
+        super().__init__(input_size=input_size,
+                         output_size=output_size,
+                         bias=bias,
+                         gather_output=False,
+                         skip_bias_add=skip_bias_add,
+                         params_dtype=params_dtype,
+                         quant_config=quant_config,
+                         prefix=prefix)
+
+    def _get_shard_offset_mapping(self, loaded_shard_id: str) -> Optional[int]:
+        shard_offset_mapping = {
+            "q": 0,
+            "k": self.num_heads * self.head_size,
+            "v": (self.num_heads + self.num_kv_heads) * self.head_size,
+            "total": (self.num_heads + 2 * self.num_kv_heads) * self.head_size
+        }
+        return shard_offset_mapping.get(loaded_shard_id)
+
+    def _get_shard_size_mapping(self, loaded_shard_id: str) -> Optional[int]:
+        shard_size_mapping = {
+            "q": self.num_heads * self.head_size,
+            "k": self.num_kv_heads * self.head_size,
+            "v": self.num_kv_heads * self.head_size,
+        }
+        return shard_size_mapping.get(loaded_shard_id)
+
+    def _load_fused_module_from_checkpoint(self, param: BasevLLMParameter,
+                                           loaded_weight: torch.Tensor):
+        """
+        Handle special case for models where QKV layers are already 
+        fused on disk. In this case, we have no shard id. This function
+        determmines the shard id by splitting these layers and then calls
+        the weight loader using the shard id.
+
+        An example of a model with these fused layers:
+        https://huggingface.co/microsoft/Phi-3-mini-4k-instruct
+        """
+        shard_offsets = [
+            # (shard_id, shard_offset, shard_size)
+            ("q", 0, self.total_num_heads * self.head_size),
+            ("k", self.total_num_heads * self.head_size,
+             self.total_num_kv_heads * self.head_size),
+            ("v",
+             (self.total_num_heads + self.total_num_kv_heads) * self.head_size,
+             self.total_num_kv_heads * self.head_size),
+        ]
+
+        for shard_id, shard_offset, shard_size in shard_offsets:
+            # Special case for Quantization.
+            # If quantized, we need to adjust the offset and size to account
+            # for the packing.
+            if isinstance(
+                    param,
+                (PackedColumnParameter,
+                 PackedvLLMParameter)) and param.packed_dim == param.output_dim:
+                shard_size, shard_offset = \
+                    param.adjust_shard_indexes_for_packing(
+                    shard_size=shard_size, shard_offset=shard_offset)
+
+            loaded_weight_shard = loaded_weight.narrow(param.output_dim,
+                                                       shard_offset, shard_size)
+            self.weight_loader_v2(param, loaded_weight_shard, shard_id)
+
+    def weight_loader_v2(self,
+                         param: BasevLLMParameter,
+                         loaded_weight: torch.Tensor,
+                         loaded_shard_id: Optional[str] = None):
+        if loaded_shard_id is None:  # special case for certain models
+            if isinstance(param, PerTensorScaleParameter):
+                param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0)
+                return
+            elif type(param) in (RowvLLMParameter, BasevLLMParameter):
+                param.load_qkv_weight(loaded_weight=loaded_weight)
+                return
+            # TODO: @dsikka - move to parameter.py
+            self._load_fused_module_from_checkpoint(param, loaded_weight)
+            return
+
+        assert loaded_shard_id in ["q", "k", "v"]
+
+        shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
+        shard_size = self._get_shard_size_mapping(loaded_shard_id)
+
+        param.load_qkv_weight(loaded_weight=loaded_weight,
+                              num_heads=self.num_kv_head_replicas,
+                              shard_id=loaded_shard_id,
+                              shard_offset=shard_offset,
+                              shard_size=shard_size)
+
+    def weight_loader(self,
+                      param: Parameter,
+                      loaded_weight: torch.Tensor,
+                      loaded_shard_id: Optional[str] = None):
+
+        param_data = param.data
+        output_dim = getattr(param, "output_dim", None)
+        # Special case for AQLM codebooks.
+        is_metadata = getattr(param, "is_metadata", False)
+
+        # Special case for per-tensor scales in fused case.
+        needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
+
+        if loaded_shard_id is None:
+            # Loaded weight is already fused on disk (qkv).
+            # (e.g., Phi-3's qkv_proj).
+            if output_dim is None:
+                if needs_scalar_to_array:
+                    param_data, loaded_weight = adjust_scalar_to_fused_array(
+                        param_data, loaded_weight, 0)
+
+                assert param_data.shape == loaded_weight.shape
+                param_data.copy_(loaded_weight)
+                return
+            shard_offsets = [
+                # (shard_id, shard_offset, shard_size)
+                ("q", 0, self.total_num_heads * self.head_size),
+                ("k", self.total_num_heads * self.head_size,
+                 self.total_num_kv_heads * self.head_size),
+                ("v", (self.total_num_heads + self.total_num_kv_heads) *
+                 self.head_size, self.total_num_kv_heads * self.head_size),
+            ]
+
+            for shard_id, shard_offset, shard_size in shard_offsets:
+
+                loaded_weight_shard = loaded_weight.narrow(
+                    output_dim, shard_offset, shard_size)
+                self.weight_loader(param, loaded_weight_shard, shard_id)
+            return
+
+        tp_rank = get_tensor_model_parallel_rank()
+        assert loaded_shard_id in ["q", "k", "v"]
+
+        # If output dim is defined, use the default loading process.
+        if output_dim is not None:
+            if loaded_shard_id == "q":
+                shard_offset = 0
+                shard_size = self.num_heads * self.head_size
+            elif loaded_shard_id == "k":
+                shard_offset = self.num_heads * self.head_size
+                shard_size = self.num_kv_heads * self.head_size
+            elif loaded_shard_id == "v":
+                shard_offset = (self.num_heads +
+                                self.num_kv_heads) * self.head_size
+                shard_size = self.num_kv_heads * self.head_size
+
+            is_sharded_weight = getattr(param, "is_sharded_weight", False)
+            # bitsandbytes loads the weights of the specific portion
+            # no need to narrow
+            is_sharded_weight = is_sharded_weight
+
+            param_data = param_data.narrow(output_dim, shard_offset, shard_size)
+            if loaded_shard_id == "q":
+                shard_id = tp_rank
+            else:
+                shard_id = tp_rank // self.num_kv_head_replicas
+            start_idx = shard_id * shard_size
+
+            if not is_sharded_weight:
+                loaded_weight = loaded_weight.narrow(output_dim, start_idx,
+                                                     shard_size)
+
+        # Special case for for AQLM codebooks.
+        elif is_metadata:
+            # metadata indicates fixed size concatenated along dim 0
+            shard_size = loaded_weight.shape[0]
+            shard_index = ["q", "k", "v"].index(loaded_shard_id)
+            param_data = param_data.narrow(0, shard_index * shard_size,
+                                           shard_size)
+        # Special case for per-tensor scales in fused case.
+        elif needs_scalar_to_array:
+            param_data, loaded_weight = adjust_scalar_to_fused_array(
+                param_data, loaded_weight, loaded_shard_id)
+        else:
+            ignore_warning = getattr(param, "ignore_warning", False)
+            if not ignore_warning:
+                logger.warning(
+                    "Loading a weight without `output_dim` attribute in "
+                    "QKVParallelLinear, assume the weight is the same "
+                    "for all partitions.")
+
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+
+class RowParallelLinear(LinearBase):
+    """Linear layer with row parallelism.
+
+    The linear layer is defined as Y = XA + b. A is parallelized along
+    its first dimension and X along its second dimension as:
+               -   -
+              | A_1 |
+              | .   |
+          A = | .   |        X = [X_1, ..., X_p]
+              | .   |
+              | A_p |
+               -   -
+    Arguments:
+        input_size: first dimension of matrix A.
+        output_size: second dimension of matrix A.
+        bias: If true, add bias. Note that bias is not parallelized.
+        input_is_parallel: If true, we assume that the input is already
+                           split across the GPUs and we do not split
+                           again.
+        skip_bias_add: This was added to enable performance optimization where
+                       bias can be fused with other element-wise operations.
+                       We skip adding bias but instead return it.
+        params_dtype: Data type for the parameters.
+        quant_config: Quantization configure.
+    """
+
+    def __init__(self,
+                 input_size: int,
+                 output_size: int,
+                 bias: bool = True,
+                 input_is_parallel: bool = True,
+                 skip_bias_add: bool = False,
+                 params_dtype: Optional[torch.dtype] = None,
+                 reduce_results: bool = True,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        # Divide the weight matrix along the first dimension.
+        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.input_size_per_partition = divide(input_size, self.tp_size)
+        self.output_size_per_partition = output_size
+        self.output_partition_sizes = [output_size]
+
+        super().__init__(input_size, output_size, skip_bias_add, params_dtype,
+                         quant_config, prefix)
+
+        self.input_is_parallel = input_is_parallel
+        self.reduce_results = reduce_results
+
+        assert self.quant_method is not None
+        self.quant_method.create_weights(
+            layer=self,
+            input_size_per_partition=self.input_size_per_partition,
+            output_partition_sizes=self.output_partition_sizes,
+            input_size=self.input_size,
+            output_size=self.output_size,
+            params_dtype=self.params_dtype,
+            weight_loader=(
+                self.weight_loader_v2 if self.quant_method.__class__.__name__
+                in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader))
+        if not reduce_results and (bias and not skip_bias_add):
+            raise ValueError("When not reduce the results, adding bias to the "
+                             "results can lead to incorrect results")
+
+        if bias:
+            self.bias = Parameter(
+                torch.empty(self.output_size, dtype=params_dtype))
+            set_weight_attrs(self.bias, {
+                "output_dim": 0,
+                "weight_loader": self.weight_loader,
+            })
+        else:
+            self.register_parameter("bias", None)
+
+    def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
+        tp_rank = get_tensor_model_parallel_rank()
+        input_dim = getattr(param, "input_dim", None)
+        is_sharded_weight = getattr(param, "is_sharded_weight", False)
+        # bitsandbytes loads the weights of the specific portion
+        # no need to narrow
+        is_sharded_weight = is_sharded_weight
+
+        param_data = param.data
+        if input_dim is not None and not is_sharded_weight:
+            shard_size = param_data.shape[input_dim]
+            start_idx = tp_rank * shard_size
+            loaded_weight = loaded_weight.narrow(input_dim, start_idx,
+                                                 shard_size)
+
+        # Special case for loading scales off disk, which often do not
+        # have a shape (such as in the case of AutoFP8).
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+    def weight_loader_v2(self, param: BasevLLMParameter,
+                         loaded_weight: torch.Tensor):
+
+        # Special case for loading scales off disk, which often do not
+        # have a shape (such as in the case of AutoFP8).
+        if len(loaded_weight.shape) == 0:
+            assert loaded_weight.numel() == 1
+            loaded_weight = loaded_weight.reshape(1)
+
+        param.load_row_parallel_weight(loaded_weight=loaded_weight)
+
+    def forward(self, input_) -> tuple[torch.Tensor, Optional[Parameter]]:
+        if self.input_is_parallel:
+            input_parallel = input_
+        else:
+            tp_rank = get_tensor_model_parallel_rank()
+            splitted_input = split_tensor_along_last_dim(
+                input_, num_partitions=self.tp_size)
+            input_parallel = splitted_input[tp_rank].contiguous()
+
+        # Matrix multiply.
+        assert self.quant_method is not None
+        # Only fuse bias add into GEMM for rank 0 (this ensures that
+        # bias will not get added more than once in TP>1 case)
+        bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
+        output_parallel = self.quant_method.apply(self,
+                                                  input_parallel,
+                                                  bias=bias_)
+        if self.reduce_results and self.tp_size > 1:
+            output = tensor_model_parallel_all_reduce(output_parallel)
+        else:
+            output = output_parallel
+
+        output_bias = self.bias if self.skip_bias_add else None
+
+        return output, output_bias
+
+    def extra_repr(self) -> str:
+        s = f"input_features={self.input_size_per_partition}"
+        s += f", output_features={self.output_size}"
+        s += f", bias={self.bias is not None}"
+        s += f", tp_size={self.tp_size}"
+        s += f", reduce_results={self.reduce_results}"
+        return s

fastvideo/v1/layers/mlp.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from fastvideo.v1.layers.activation import get_act_fn
+from fastvideo.v1.layers.linear import ReplicatedLinear
+
+
+class MLP(nn.Module):
+    """
+    MLP for DiT blocks, NO gated linear units
+    TODO: add Tensor Parallel
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        mlp_hidden_dim: int,
+        output_dim: Optional[int] = None,
+        bias: bool = True,
+        act_type: str = "gelu_pytorch_tanh",
+        dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+        self.fc_in = ReplicatedLinear(
+            input_dim,
+            mlp_hidden_dim,  # For activation func like SiLU that need 2x width
+            bias=bias,
+            params_dtype=dtype)
+
+        self.act = get_act_fn(act_type)
+        if output_dim is None:
+            output_dim = input_dim
+        self.fc_out = ReplicatedLinear(mlp_hidden_dim,
+                                       output_dim,
+                                       bias=bias,
+                                       params_dtype=dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.fc_in(x)
+        x = self.act(x)
+        x, _ = self.fc_out(x)
+        return x

fastvideo/v1/layers/rotary_embedding.py

@@ -0,0 +1,479 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/rotary_embedding.py
+
+# Adapted from
+# https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Rotary Positional Embeddings."""
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+
+from fastvideo.v1.distributed.parallel_state import get_sp_group
+from fastvideo.v1.layers.custom_op import CustomOp
+
+
+def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., :x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    x = torch.stack((-x2, x1), dim=-1)
+    return x.flatten(-2)
+
+
+def _apply_rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    is_neox_style: bool,
+) -> torch.Tensor:
+    """
+    Args:
+        x: [num_tokens, num_heads, head_size]
+        cos: [num_tokens, head_size // 2]
+        sin: [num_tokens, head_size // 2]
+        is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
+            positional embeddings.
+    """
+    # cos = cos.unsqueeze(-2).to(x.dtype)
+    # sin = sin.unsqueeze(-2).to(x.dtype)
+    cos = cos.unsqueeze(-2)
+    sin = sin.unsqueeze(-2)
+    if is_neox_style:
+        x1, x2 = torch.chunk(x, 2, dim=-1)
+    else:
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+    o1 = (x1.float() * cos - x2.float() * sin).type_as(x)
+    o2 = (x2.float() * cos + x1.float() * sin).type_as(x)
+    if is_neox_style:
+        return torch.cat((o1, o2), dim=-1)
+    else:
+        return torch.stack((o1, o2), dim=-1).flatten(-2)
+
+
+@CustomOp.register("rotary_embedding")
+class RotaryEmbedding(CustomOp):
+    """Original rotary positional embedding."""
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+    ) -> None:
+        super().__init__()
+        self.head_size = head_size
+        self.rotary_dim = rotary_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        self.is_neox_style = is_neox_style
+        self.dtype = dtype
+
+        cache = self._compute_cos_sin_cache()
+        cache = cache.to(dtype)
+        self.cos_sin_cache: torch.Tensor
+        self.register_buffer("cos_sin_cache", cache, persistent=False)
+
+    def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
+        """Compute the inverse frequency."""
+        # NOTE(woosuk): To exactly match the HF implementation, we need to
+        # use CPU to compute the cache and then move it to GPU. However, we
+        # create the cache on GPU for faster initialization. This may cause
+        # a slight numerical difference between the HF implementation and ours.
+        inv_freq = 1.0 / (base**(torch.arange(
+            0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim))
+        return inv_freq
+
+    def _compute_cos_sin_cache(self) -> torch.Tensor:
+        """Compute the cos and sin cache."""
+        inv_freq = self._compute_inv_freq(self.base)
+        t = torch.arange(self.max_position_embeddings, dtype=torch.float)
+
+        freqs = torch.einsum("i,j -> ij", t, inv_freq)
+        cos = freqs.cos()
+        sin = freqs.sin()
+        cache = torch.cat((cos, sin), dim=-1)
+        return cache
+
+    def forward_native(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """A PyTorch-native implementation of forward()."""
+        if offsets is not None:
+            positions = positions + offsets
+        positions = positions.flatten()
+        num_tokens = positions.shape[0]
+        cos_sin = self.cos_sin_cache.index_select(0, positions)
+        cos, sin = cos_sin.chunk(2, dim=-1)
+
+        query_shape = query.shape
+        query = query.view(num_tokens, -1, self.head_size)
+        query_rot = query[..., :self.rotary_dim]
+        query_pass = query[..., self.rotary_dim:]
+        query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
+        query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
+
+        key_shape = key.shape
+        key = key.view(num_tokens, -1, self.head_size)
+        key_rot = key[..., :self.rotary_dim]
+        key_pass = key[..., self.rotary_dim:]
+        key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
+        key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
+        return query, key
+
+    def forward_cuda(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        from vllm import _custom_ops as ops
+
+        self.cos_sin_cache = self.cos_sin_cache.to(query.device,
+                                                   dtype=query.dtype)
+        # ops.rotary_embedding()/batched_rotary_embedding()
+        # are in-place operations that update the query and key tensors.
+        if offsets is not None:
+            ops.batched_rotary_embedding(positions, query, key, self.head_size,
+                                         self.cos_sin_cache, self.is_neox_style,
+                                         self.rotary_dim, offsets)
+        else:
+            ops.rotary_embedding(positions, query, key, self.head_size,
+                                 self.cos_sin_cache, self.is_neox_style)
+        return query, key
+
+    def extra_repr(self) -> str:
+        s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
+        s += f", max_position_embeddings={self.max_position_embeddings}"
+        s += f", base={self.base}, is_neox_style={self.is_neox_style}"
+        return s
+
+
+def _to_tuple(x: Union[int, Tuple[int, ...]], dim: int = 2) -> Tuple[int, ...]:
+    if isinstance(x, int):
+        return (x, ) * dim
+    elif len(x) == dim:
+        return x
+    else:
+        raise ValueError(f"Expected length {dim} or int, but got {x}")
+
+
+def get_meshgrid_nd(start: Union[int, Tuple[int, ...]],
+                    *args: Union[int, Tuple[int, ...]],
+                    dim: int = 2) -> torch.Tensor:
+    """
+    Get n-D meshgrid with start, stop and num.
+
+    Args:
+        start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop,
+            step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num
+            should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in
+            n-tuples.
+        *args: See above.
+        dim (int): Dimension of the meshgrid. Defaults to 2.
+
+    Returns:
+        grid (np.ndarray): [dim, ...]
+    """
+    if len(args) == 0:
+        # start is grid_size
+        num = _to_tuple(start, dim=dim)
+        start = (0, ) * dim
+        stop = num
+    elif len(args) == 1:
+        # start is start, args[0] is stop, step is 1
+        start = _to_tuple(start, dim=dim)
+        stop = _to_tuple(args[0], dim=dim)
+        num = tuple(stop[i] - start[i] for i in range(dim))
+    elif len(args) == 2:
+        # start is start, args[0] is stop, args[1] is num
+        start = _to_tuple(start, dim=dim)  # Left-Top       eg: 12,0
+        stop = _to_tuple(args[0], dim=dim)  # Right-Bottom   eg: 20,32
+        num = _to_tuple(args[1], dim=dim)  # Target Size    eg: 32,124
+    else:
+        raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}")
+
+    # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False)
+    axis_grid = []
+    for i in range(dim):
+        a, b, n = start[i], stop[i], num[i]
+        g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n]
+        axis_grid.append(g)
+    grid = torch.meshgrid(*axis_grid, indexing="ij")  # dim x [W, H, D]
+    grid = torch.stack(grid, dim=0)  # [dim, W, H, D]
+
+    return grid
+
+
+def get_1d_rotary_pos_embed(
+    dim: int,
+    pos: Union[torch.FloatTensor, int],
+    theta: float = 10000.0,
+    theta_rescale_factor: float = 1.0,
+    interpolation_factor: float = 1.0,
+    dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Precompute the frequency tensor for complex exponential (cis) with given dimensions.
+    (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
+
+    This function calculates a frequency tensor with complex exponential using the given dimension 'dim'
+    and the end index 'end'. The 'theta' parameter scales the frequencies.
+
+    Args:
+        dim (int): Dimension of the frequency tensor.
+        pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar
+        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
+        theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0.
+        interpolation_factor (float, optional): Factor to scale positions. Defaults to 1.0.
+
+    Returns:
+        freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D]
+    """
+    if isinstance(pos, int):
+        pos = torch.arange(pos).float()
+
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    if theta_rescale_factor != 1.0:
+        theta *= theta_rescale_factor**(dim / (dim - 2))
+
+    freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].to(dtype) / dim)
+                   )  # [D/2]
+    freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
+    freqs_cos = freqs.cos()  # [S, D/2]
+    freqs_sin = freqs.sin()  # [S, D/2]
+    return freqs_cos, freqs_sin
+
+
+def get_nd_rotary_pos_embed(
+    rope_dim_list,
+    start,
+    *args,
+    theta=10000.0,
+    theta_rescale_factor: Union[float, List[float]] = 1.0,
+    interpolation_factor: Union[float, List[float]] = 1.0,
+    shard_dim: int = 0,
+    sp_rank: int = 0,
+    sp_world_size: int = 1,
+    dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
+    Supports sequence parallelism by allowing sharding of a specific dimension.
+
+    Args:
+        rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n.
+            sum(rope_dim_list) should equal to head_dim of attention layer.
+        start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start,
+            args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num.
+        *args: See above.
+        theta (float): Scaling factor for frequency computation. Defaults to 10000.0.
+        theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0.
+        interpolation_factor (float): Factor to scale positions. Defaults to 1.0.
+        shard_dim (int): Which dimension to shard for sequence parallelism. Defaults to 0.
+        sp_rank (int): Rank in the sequence parallel group. Defaults to 0.
+        sp_world_size (int): World size of the sequence parallel group. Defaults to 1.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: (cos, sin) tensors of shape [HW, D/2]
+    """
+    # Get the full grid
+    full_grid = get_meshgrid_nd(
+        start, *args, dim=len(rope_dim_list))  # [3, W, H, D] / [2, W, H]
+
+    # Shard the grid if using sequence parallelism (sp_world_size > 1)
+    assert shard_dim < len(
+        rope_dim_list
+    ), f"shard_dim {shard_dim} must be less than number of dimensions {len(rope_dim_list)}"
+    if sp_world_size > 1:
+        # Get the shape of the full grid
+        grid_shape = list(full_grid.shape[1:])
+
+        # Ensure the dimension to shard is divisible by sp_world_size
+        assert grid_shape[shard_dim] % sp_world_size == 0, (
+            f"Dimension {shard_dim} with size {grid_shape[shard_dim]} is not divisible "
+            f"by sequence parallel world size {sp_world_size}")
+
+        # Compute the start and end indices for this rank's shard
+        shard_size = grid_shape[shard_dim] // sp_world_size
+        start_idx = sp_rank * shard_size
+        end_idx = (sp_rank + 1) * shard_size
+
+        # Create slicing indices for each dimension
+        slice_indices = [slice(None) for _ in range(len(grid_shape))]
+        slice_indices[shard_dim] = slice(start_idx, end_idx)
+
+        # Shard the grid
+        # Update grid shape for the sharded dimension
+        grid_shape[shard_dim] = grid_shape[shard_dim] // sp_world_size
+        grid = torch.empty((len(rope_dim_list), ) + tuple(grid_shape),
+                           dtype=full_grid.dtype)
+        for i in range(len(rope_dim_list)):
+            grid[i] = full_grid[i][tuple(slice_indices)]
+    else:
+        grid = full_grid
+
+    if isinstance(theta_rescale_factor, (int, float)):
+        theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
+    elif isinstance(theta_rescale_factor,
+                    list) and len(theta_rescale_factor) == 1:
+        theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list)
+    assert len(theta_rescale_factor) == len(
+        rope_dim_list
+    ), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
+
+    if isinstance(interpolation_factor, (int, float)):
+        interpolation_factor = [interpolation_factor] * len(rope_dim_list)
+    elif isinstance(interpolation_factor,
+                    list) and len(interpolation_factor) == 1:
+        interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list)
+    assert len(interpolation_factor) == len(
+        rope_dim_list
+    ), "len(interpolation_factor) should equal to len(rope_dim_list)"
+
+    # use 1/ndim of dimensions to encode grid_axis
+    embs = []
+    for i in range(len(rope_dim_list)):
+        emb = get_1d_rotary_pos_embed(
+            rope_dim_list[i],
+            grid[i].reshape(-1),
+            theta,
+            theta_rescale_factor=theta_rescale_factor[i],
+            interpolation_factor=interpolation_factor[i],
+            dtype=dtype,
+        )  # 2 x [WHD, rope_dim_list[i]]
+        embs.append(emb)
+
+    cos = torch.cat([emb[0] for emb in embs], dim=1)  # (WHD, D/2)
+    sin = torch.cat([emb[1] for emb in embs], dim=1)  # (WHD, D/2)
+    return cos, sin
+
+
+def get_rotary_pos_embed(
+    rope_sizes,
+    hidden_size,
+    heads_num,
+    rope_dim_list,
+    rope_theta,
+    theta_rescale_factor=1.0,
+    interpolation_factor=1.0,
+    shard_dim: int = 0,
+    dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generate rotary positional embeddings for the given sizes.
+    
+    Args:
+        rope_sizes: Tuple of dimensions (t, h, w)
+        hidden_size: Hidden dimension size
+        heads_num: Number of attention heads
+        rope_dim_list: List of dimensions for each axis, or None
+        rope_theta: Base for frequency calculations
+        theta_rescale_factor: Rescale factor for theta. Defaults to 1.0
+        interpolation_factor: Factor to scale positions. Defaults to 1.0
+        shard_dim: Which dimension to shard for sequence parallelism. Defaults to 0.
+        
+    Returns:
+        Tuple of (cos, sin) tensors for rotary embeddings
+    """
+
+    target_ndim = 3
+    head_dim = hidden_size // heads_num
+
+    if rope_dim_list is None:
+        rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
+
+    assert sum(
+        rope_dim_list
+    ) == head_dim, "sum(rope_dim_list) should equal to head_dim of attention layer"
+
+    # Get SP info
+    sp_group = get_sp_group()
+    sp_rank = sp_group.rank_in_group
+    sp_world_size = sp_group.world_size
+
+    freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
+        rope_dim_list,
+        rope_sizes,
+        theta=rope_theta,
+        theta_rescale_factor=theta_rescale_factor,
+        interpolation_factor=interpolation_factor,
+        shard_dim=shard_dim,
+        sp_rank=sp_rank,
+        sp_world_size=sp_world_size,
+        dtype=dtype,
+    )
+    return freqs_cos, freqs_sin
+
+
+_ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {}
+
+
+def get_rope(
+    head_size: int,
+    rotary_dim: int,
+    max_position: int,
+    base: int,
+    is_neox_style: bool = True,
+    rope_scaling: Optional[Dict[str, Any]] = None,
+    dtype: Optional[torch.dtype] = None,
+    partial_rotary_factor: float = 1.0,
+) -> RotaryEmbedding:
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+    if rope_scaling is not None:
+        # Transforms every value that is a list into a tuple for caching calls
+        rope_scaling_tuple = {
+            k: tuple(v) if isinstance(v, list) else v
+            for k, v in rope_scaling.items()
+        }
+        rope_scaling_args = tuple(rope_scaling_tuple.items())
+    else:
+        rope_scaling_args = None
+    if partial_rotary_factor < 1.0:
+        rotary_dim = int(rotary_dim * partial_rotary_factor)
+    key = (head_size, rotary_dim, max_position, base, is_neox_style,
+           rope_scaling_args, dtype)
+    if key in _ROPE_DICT:
+        return _ROPE_DICT[key]
+
+    if rope_scaling is None:
+        rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base,
+                                     is_neox_style, dtype)
+    else:
+        raise ValueError(f"Unknown RoPE scaling {rope_scaling}")
+    _ROPE_DICT[key] = rotary_emb
+    return rotary_emb

fastvideo/v1/layers/utils.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/utils.py
+"""Utility methods for model layers."""
+from typing import Tuple
+
+import torch
+
+
+def get_token_bin_counts_and_mask(
+    tokens: torch.Tensor,
+    vocab_size: int,
+    num_seqs: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # Compute the bin counts for the tokens.
+    # vocab_size + 1 for padding.
+    bin_counts = torch.zeros((num_seqs, vocab_size + 1),
+                             dtype=torch.long,
+                             device=tokens.device)
+    bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens))
+    bin_counts = bin_counts[:, :vocab_size]
+    mask = bin_counts > 0
+
+    return bin_counts, mask

fastvideo/v1/layers/visual_embedding.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from fastvideo.v1.layers.activation import get_act_fn
+from fastvideo.v1.layers.linear import ReplicatedLinear
+from fastvideo.v1.layers.mlp import MLP
+
+
+class PatchEmbed(nn.Module):
+    """2D Image to Patch Embedding
+
+    Image to Patch Embedding using Conv2d
+
+    A convolution based approach to patchifying a 2D image w/ embedding projection.
+
+    Based on the impl in https://github.com/google-research/vision_transformer
+
+    Hacked together by / Copyright 2020 Ross Wightman
+
+    Remove the _assert function in forward function to be compatible with multi-resolution images.
+    """
+
+    def __init__(self,
+                 patch_size=16,
+                 in_chans=3,
+                 embed_dim=768,
+                 norm_layer=None,
+                 flatten=True,
+                 bias=True,
+                 dtype=None):
+        super().__init__()
+        # Convert patch_size to 2-tuple
+        if isinstance(patch_size, (list, tuple)):
+            if len(patch_size) == 1:
+                patch_size = (patch_size[0], patch_size[0])
+        else:
+            patch_size = (patch_size, patch_size)
+
+        self.patch_size = patch_size
+        self.flatten = flatten
+
+        self.proj = nn.Conv3d(in_chans,
+                              embed_dim,
+                              kernel_size=patch_size,
+                              stride=patch_size,
+                              bias=bias,
+                              dtype=dtype)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        act_layer="silu",
+        frequency_embedding_size=256,
+        max_period=10000,
+        dtype=None,
+        freq_dtype=torch.float32,
+    ):
+        super().__init__()
+        self.frequency_embedding_size = frequency_embedding_size
+        self.max_period = max_period
+
+        self.mlp = MLP(frequency_embedding_size,
+                       hidden_size,
+                       hidden_size,
+                       act_type=act_layer,
+                       dtype=dtype)
+        self.freq_dtype = freq_dtype
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        t_freq = timestep_embedding(t,
+                                    self.frequency_embedding_size,
+                                    self.max_period,
+                                    dtype=self.freq_dtype).to(
+                                        self.mlp.fc_in.weight.dtype)
+        # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+def timestep_embedding(t: torch.Tensor,
+                       dim: int,
+                       max_period: int = 10000,
+                       dtype: torch.dtype = torch.float32) -> torch.Tensor:
+    """
+    Create sinusoidal timestep embeddings.
+    
+    Args:
+        t: Tensor of shape [B] with timesteps
+        dim: Embedding dimension
+        max_period: Controls the minimum frequency of the embeddings
+        
+    Returns:
+        Tensor of shape [B, dim] with embeddings
+    """
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) *
+                      torch.arange(start=0, end=half, dtype=dtype) /
+                      half).to(device=t.device)
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat(
+            [embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+class ModulateProjection(nn.Module):
+    """Modulation layer for DiT blocks."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        factor: int = 2,
+        act_layer: str = "silu",
+        dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+        self.factor = factor
+        self.hidden_size = hidden_size
+        self.linear = ReplicatedLinear(hidden_size,
+                                       hidden_size * factor,
+                                       bias=True,
+                                       params_dtype=dtype)
+        self.act = get_act_fn(act_layer)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.act(x)
+        x, _ = self.linear(x)
+        return x
+
+
+def unpatchify(x, t, h, w, patch_size, channels):
+    """
+    Convert patched representation back to image space.
+    
+    Args:
+        x: Tensor of shape [B, T*H*W, C*P_t*P_h*P_w]
+        t, h, w: Temporal and spatial dimensions
+        
+    Returns:
+        Unpatchified tensor of shape [B, C, T*P_t, H*P_h, W*P_w]
+    """
+    assert x.ndim == 3, f"x.ndim: {x.ndim}"
+    assert len(patch_size) == 3, f"patch_size: {patch_size}"
+    assert t * h * w == x.shape[
+        1], f"t * h * w: {t * h * w}, x.shape[1]: {x.shape[1]}"
+    c = channels
+    pt, ph, pw = patch_size
+
+    x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw))
+    x = torch.einsum("nthwcopq->nctohpwq", x)
+    imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
+
+    return imgs

fastvideo/v1/layers/vocab_parallel_embedding.py

@@ -0,0 +1,417 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass
+from typing import List, Optional, Sequence, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter, UninitializedParameter
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
+
+from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank,
+                                      get_tensor_model_parallel_world_size,
+                                      tensor_model_parallel_all_reduce)
+from fastvideo.v1.models.parameter import BasevLLMParameter
+from fastvideo.v1.models.utils import set_weight_attrs
+from fastvideo.v1.platforms import current_platform
+
+DEFAULT_VOCAB_PADDING_SIZE = 64
+
+
+class UnquantizedEmbeddingMethod(QuantizeMethodBase):
+    """Unquantized method for embeddings."""
+
+    def create_weights(self, layer: torch.nn.Module,
+                       input_size_per_partition: int,
+                       output_partition_sizes: List[int], input_size: int,
+                       output_size: int, params_dtype: torch.dtype,
+                       **extra_weight_attrs):
+        """Create weights for embedding layer."""
+        weight = Parameter(torch.empty(sum(output_partition_sizes),
+                                       input_size_per_partition,
+                                       dtype=params_dtype),
+                           requires_grad=False)
+        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
+        layer.register_parameter("weight", weight)
+        set_weight_attrs(weight, extra_weight_attrs)
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return F.linear(x, layer.weight, bias)
+
+    def embedding(self, layer: torch.nn.Module,
+                  input_: torch.Tensor) -> torch.Tensor:
+        return F.embedding(input_, layer.weight)
+
+
+def pad_vocab_size(vocab_size: int,
+                   pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int:
+    """Pad the vocab size to the given value."""
+    return ((vocab_size + pad_to - 1) // pad_to) * pad_to
+
+
+def vocab_range_from_per_partition_vocab_size(per_partition_vocab_size: int,
+                                              rank: int,
+                                              offset: int = 0) -> Sequence[int]:
+    index_f = rank * per_partition_vocab_size
+    index_l = index_f + per_partition_vocab_size
+    return index_f + offset, index_l + offset
+
+
+def vocab_range_from_global_vocab_size(global_vocab_size: int,
+                                       rank: int,
+                                       world_size: int,
+                                       offset: int = 0) -> Sequence[int]:
+    per_partition_vocab_size = divide(global_vocab_size, world_size)
+    return vocab_range_from_per_partition_vocab_size(per_partition_vocab_size,
+                                                     rank,
+                                                     offset=offset)
+
+
+@dataclass
+class VocabParallelEmbeddingShardIndices:
+    """Indices for a shard of a vocab parallel embedding."""
+    padded_org_vocab_start_index: int
+    padded_org_vocab_end_index: int
+    padded_added_vocab_start_index: int
+    padded_added_vocab_end_index: int
+
+    org_vocab_start_index: int
+    org_vocab_end_index: int
+    added_vocab_start_index: int
+    added_vocab_end_index: int
+
+    @property
+    def num_org_elements(self) -> int:
+        return self.org_vocab_end_index - self.org_vocab_start_index
+
+    @property
+    def num_added_elements(self) -> int:
+        return self.added_vocab_end_index - self.added_vocab_start_index
+
+    @property
+    def num_org_elements_padded(self) -> int:
+        return (self.padded_org_vocab_end_index -
+                self.padded_org_vocab_start_index)
+
+    @property
+    def num_added_elements_padded(self) -> int:
+        return (self.padded_added_vocab_end_index -
+                self.padded_added_vocab_start_index)
+
+    @property
+    def num_org_vocab_padding(self) -> int:
+        return self.num_org_elements_padded - self.num_org_elements
+
+    @property
+    def num_added_vocab_padding(self) -> int:
+        return self.num_added_elements_padded - self.num_added_elements
+
+    @property
+    def num_elements_padded(self) -> int:
+        return self.num_org_elements_padded + self.num_added_elements_padded
+
+    def __post_init__(self):
+        # sanity checks
+        assert (self.padded_org_vocab_start_index
+                <= self.padded_org_vocab_end_index)
+        assert (self.padded_added_vocab_start_index
+                <= self.padded_added_vocab_end_index)
+
+        assert self.org_vocab_start_index <= self.org_vocab_end_index
+        assert self.added_vocab_start_index <= self.added_vocab_end_index
+
+        assert self.org_vocab_start_index <= self.padded_org_vocab_start_index
+        assert (self.added_vocab_start_index
+                <= self.padded_added_vocab_start_index)
+        assert self.org_vocab_end_index <= self.padded_org_vocab_end_index
+        assert self.added_vocab_end_index <= self.padded_added_vocab_end_index
+
+        assert self.num_org_elements <= self.num_org_elements_padded
+        assert self.num_added_elements <= self.num_added_elements_padded
+
+
+@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+def get_masked_input_and_mask(
+        input_: torch.Tensor, org_vocab_start_index: int,
+        org_vocab_end_index: int, num_org_vocab_padding: int,
+        added_vocab_start_index: int,
+        added_vocab_end_index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    # torch.compile will fuse all of the pointwise ops below
+    # into a single kernel, making it very fast
+    org_vocab_mask = (input_ >= org_vocab_start_index) & (input_
+                                                          < org_vocab_end_index)
+    added_vocab_mask = (input_ >= added_vocab_start_index) & (
+        input_ < added_vocab_end_index)
+    added_offset = added_vocab_start_index - (
+        org_vocab_end_index - org_vocab_start_index) - num_org_vocab_padding
+    valid_offset = (org_vocab_start_index * org_vocab_mask) + (added_offset *
+                                                               added_vocab_mask)
+    vocab_mask = org_vocab_mask | added_vocab_mask
+    input_ = vocab_mask * (input_ - valid_offset)
+    return input_, ~vocab_mask
+
+
+class VocabParallelEmbedding(torch.nn.Module):
+    """Embedding parallelized in the vocabulary dimension.
+
+    Adapted from torch.nn.Embedding, note that we pad the vocabulary size to
+    make sure it is divisible by the number of model parallel GPUs.
+
+    In order to support various loading methods, we ensure that LoRA-added
+    embeddings are always at the end of TP-sharded tensors. In other words,
+    we shard base embeddings and LoRA embeddings separately (both padded),
+    and place them in the same tensor.
+    In this example, we will have the original vocab size = 1010,
+    added vocab size = 16 and padding to 64. Therefore, the total
+    vocab size with padding will be 1088 (because we first pad 1010 to
+    1024, add 16, and then pad to 1088).
+    Therefore, the tensor format looks like the following:
+    TP1, rank 0 (no sharding):
+                            |< --------BASE-------- >|< -BASE PADDING-- >|< -----LORA------ >|< -LORA PADDING-- >|
+    corresponding token_id: |  0  |  1  | ... | 1009 |  -1  | ... |  -1  | 1010 | ... | 1015 |  -1  | ... |  -1  |
+                     index: |  0  |  1  | ... | 1009 | 1010 | ... | 1023 | 1024 | ... | 1039 | 1040 | ... | 1087 |
+
+    TP2, rank 0:
+                            |< --------------------BASE--------------------- >|< -----LORA------ >|< -LORA PADDING- >|
+    corresponding token_id: |  0  |  1  |  2  | ... | 497  | 498 | ...  | 511 | 1000 | ... | 1015 |  -1  | ... |  -1 |
+                     index: |  0  |  1  |  2  | ... | 497  | 498 | ...  | 511 | 512  | ... | 527  |  520 | ... | 543 |
+    TP2, rank 1:
+                            |< -----------BASE----------- >|< -BASE PADDING- >|< -----------LORA PADDING----------- >|
+    corresponding token_id: | 512 | 513 | 514 | ... | 1009 | -1  | ...  | -1  |  -1  | ... |  -1  | -1  | ... |   -1 |
+                     index: |  0  |  1  |  2  | ... | 497  | 498 | ...  | 511 | 512  | ... | 519  | 520 | ... |  543 |
+
+    Args:
+        num_embeddings: vocabulary size.
+        embedding_dim: size of hidden state.
+        params_dtype: type of the parameters.
+        org_num_embeddings: original vocabulary size (without LoRA).
+        padding_size: padding size for the vocabulary.
+        quant_config: quant config for the layer
+        prefix: full name of the layer in the state dict
+    """  # noqa: E501
+
+    def __init__(self,
+                 num_embeddings: int,
+                 embedding_dim: int,
+                 params_dtype: Optional[torch.dtype] = None,
+                 org_num_embeddings: Optional[int] = None,
+                 padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        super().__init__()
+
+        # Keep the input dimensions.
+        tp_rank = get_tensor_model_parallel_rank()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.num_embeddings = num_embeddings
+        self.padding_size = padding_size
+        self.org_vocab_size = org_num_embeddings or num_embeddings
+        num_added_embeddings = num_embeddings - self.org_vocab_size
+        self.org_vocab_size_padded = pad_vocab_size(self.org_vocab_size,
+                                                    self.padding_size)
+        self.num_embeddings_padded = pad_vocab_size(
+            self.org_vocab_size_padded + num_added_embeddings,
+            self.padding_size)
+        assert self.org_vocab_size_padded <= self.num_embeddings_padded
+
+        self.shard_indices = self._get_indices(self.num_embeddings_padded,
+                                               self.org_vocab_size_padded,
+                                               self.num_embeddings,
+                                               self.org_vocab_size, tp_rank,
+                                               self.tp_size)
+        self.embedding_dim = embedding_dim
+
+        quant_method = None
+        if quant_config is not None:
+            quant_method = quant_config.get_quant_method(self, prefix=prefix)
+        if quant_method is None:
+            quant_method = UnquantizedEmbeddingMethod()
+
+        # If we are making an embedding layer, then our quantization linear
+        # method must implement the embedding operation. If we are another
+        # layer type like ParallelLMHead, this is not important.
+        is_embedding_layer = type(self.__class__) is VocabParallelEmbedding
+        quant_method_implements_embedding = method_has_implemented_embedding(
+            type(quant_method))
+        if is_embedding_layer and not quant_method_implements_embedding:
+            raise NotImplementedError(
+                f"The class {type(quant_method).__name__} must implement "
+                "the 'embedding' method, see UnquantizedEmbeddingMethod.")
+
+        self.quant_method: QuantizeMethodBase = quant_method
+
+        if params_dtype is None:
+            params_dtype = torch.get_default_dtype()
+        # Divide the weight matrix along the vocaburaly dimension.
+        self.num_added_embeddings = self.num_embeddings - self.org_vocab_size
+        self.num_embeddings_per_partition = divide(self.num_embeddings_padded,
+                                                   self.tp_size)
+        assert (self.shard_indices.num_elements_padded ==
+                self.num_embeddings_per_partition)
+        self.num_org_embeddings_per_partition = (
+            self.shard_indices.org_vocab_end_index -
+            self.shard_indices.org_vocab_start_index)
+        self.num_added_embeddings_per_partition = (
+            self.shard_indices.added_vocab_end_index -
+            self.shard_indices.added_vocab_start_index)
+
+        self.quant_method.create_weights(self,
+                                         self.embedding_dim,
+                                         [self.num_embeddings_per_partition],
+                                         self.embedding_dim,
+                                         self.num_embeddings_padded,
+                                         params_dtype=params_dtype,
+                                         weight_loader=self.weight_loader)
+
+    @classmethod
+    def _get_indices(cls, vocab_size_padded: int, org_vocab_size_padded: int,
+                     vocab_size: int, org_vocab_size: int, tp_rank: int,
+                     tp_size: int) -> VocabParallelEmbeddingShardIndices:
+        """Get start and end indices for vocab parallel embedding, following the
+        layout outlined in the class docstring, based on the given tp_rank and
+        tp_size."""
+        num_added_embeddings_padded = vocab_size_padded - org_vocab_size_padded
+        padded_org_vocab_start_index, padded_org_vocab_end_index = (
+            vocab_range_from_global_vocab_size(org_vocab_size_padded, tp_rank,
+                                               tp_size))
+        padded_added_vocab_start_index, padded_added_vocab_end_index = (
+            vocab_range_from_global_vocab_size(num_added_embeddings_padded,
+                                               tp_rank,
+                                               tp_size,
+                                               offset=org_vocab_size))
+        # remove padding
+        org_vocab_start_index = min(padded_org_vocab_start_index,
+                                    org_vocab_size)
+        org_vocab_end_index = min(padded_org_vocab_end_index, org_vocab_size)
+        added_vocab_start_index = min(padded_added_vocab_start_index,
+                                      vocab_size)
+        added_vocab_end_index = min(padded_added_vocab_end_index, vocab_size)
+        return VocabParallelEmbeddingShardIndices(
+            padded_org_vocab_start_index, padded_org_vocab_end_index,
+            padded_added_vocab_start_index, padded_added_vocab_end_index,
+            org_vocab_start_index, org_vocab_end_index, added_vocab_start_index,
+            added_vocab_end_index)
+
+    def get_sharded_to_full_mapping(self) -> Optional[List[int]]:
+        """Get a mapping that can be used to reindex the gathered
+        logits for sampling.
+        
+        During sampling, we gather logits from all ranks. The relationship
+        of index->token_id will follow the same format as outlined in the class
+        docstring. However, after the gather, we want to reindex the final
+        logits tensor to map index->token_id one-to-one (the index is always
+        equal the token_id it corresponds to). The indices returned by this
+        method allow us to do that.
+        """
+        if self.tp_size < 2:
+            return None
+
+        base_embeddings: List[int] = []
+        added_embeddings: List[int] = []
+        padding: List[int] = []
+        for tp_rank in range(self.tp_size):
+            shard_indices = self._get_indices(self.num_embeddings_padded,
+                                              self.org_vocab_size_padded,
+                                              self.num_embeddings,
+                                              self.org_vocab_size, tp_rank,
+                                              self.tp_size)
+            range_start = self.num_embeddings_per_partition * tp_rank
+            range_end = self.num_embeddings_per_partition * (tp_rank + 1)
+            base_embeddings.extend(
+                range(range_start,
+                      range_start + shard_indices.num_org_elements))
+            padding.extend(
+                range(range_start + shard_indices.num_org_elements,
+                      range_start + shard_indices.num_org_elements_padded))
+            added_embeddings.extend(
+                range(
+                    range_start + shard_indices.num_org_elements_padded,
+                    range_start + shard_indices.num_org_elements_padded +
+                    shard_indices.num_added_elements))
+            padding.extend(
+                range(
+                    range_start + shard_indices.num_org_elements_padded +
+                    shard_indices.num_added_elements,
+                    range_start + shard_indices.num_org_elements_padded +
+                    shard_indices.num_added_elements_padded))
+            assert (range_start + shard_indices.num_org_elements_padded +
+                    shard_indices.num_added_elements_padded == range_end)
+        ret = base_embeddings + added_embeddings + padding
+        assert len(ret) == self.num_embeddings_padded
+        return ret
+
+    def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
+        output_dim = getattr(param, "output_dim", None)
+        packed_dim = getattr(param, "packed_dim", None)
+
+        # If the parameter is a gguf weight, then load it directly.
+        if getattr(param, "is_gguf_weight_type", None):
+            param.data.copy_(loaded_weight)
+            param.weight_type = loaded_weight.item()
+            return
+        elif isinstance(param, UninitializedParameter):
+            shape = list(loaded_weight.shape)
+            if output_dim is not None:
+                shape[output_dim] = self.num_embeddings_per_partition
+            param.materialize(tuple(shape), dtype=loaded_weight.dtype)
+
+        # If parameter does not have output dim, then it should
+        # be copied onto all gpus (e.g. g_idx for act_order gptq).
+        if output_dim is None:
+            assert param.data.shape == loaded_weight.shape
+            param.data.copy_(loaded_weight)
+            return
+
+        # Shard indexes for loading the weight
+        start_idx = self.shard_indices.org_vocab_start_index
+        shard_size = self.shard_indices.org_vocab_end_index - start_idx
+
+        # If param packed on the same dim we are sharding on, then
+        # need to adjust offsets of loaded weight by pack_factor.
+        if packed_dim is not None and packed_dim == output_dim:
+            packed_factor = param.packed_factor if isinstance(
+                param, BasevLLMParameter) else param.pack_factor
+            assert loaded_weight.shape[output_dim] == (self.org_vocab_size //
+                                                       param.packed_factor)
+            start_idx = start_idx // packed_factor
+            shard_size = shard_size // packed_factor
+        else:
+            assert loaded_weight.shape[output_dim] == self.org_vocab_size
+
+        # Copy the data. Select chunk corresponding to current shard.
+        loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
+
+        param[:loaded_weight.shape[0]].data.copy_(loaded_weight)
+        param[loaded_weight.shape[0]:].data.fill_(0)
+
+    def forward(self, input_):
+        if self.tp_size > 1:
+            # Build the mask.
+            masked_input, input_mask = get_masked_input_and_mask(
+                input_, self.shard_indices.org_vocab_start_index,
+                self.shard_indices.org_vocab_end_index,
+                self.shard_indices.num_org_vocab_padding,
+                self.shard_indices.added_vocab_start_index,
+                self.shard_indices.added_vocab_end_index)
+        else:
+            masked_input = input_
+        # Get the embeddings.
+        output_parallel = self.quant_method.embedding(self, masked_input.long())
+        # Mask the output embedding.
+        if self.tp_size > 1:
+            output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
+        # Reduce across all the model parallel GPUs.
+        output = tensor_model_parallel_all_reduce(output_parallel)
+        return output
+
+    def extra_repr(self) -> str:
+        s = f"num_embeddings={self.num_embeddings_per_partition}"
+        s += f", embedding_dim={self.embedding_dim}"
+        s += f", org_vocab_size={self.org_vocab_size}"
+        s += f', num_embeddings_padded={self.num_embeddings_padded}'
+        s += f', tp_size={self.tp_size}'
+        return s

fastvideo/v1/logger.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logger.py
+"""Logging configuration for fastvideo.v1."""
+import datetime
+import json
+import logging
+import os
+import sys
+from functools import lru_cache, partial
+from logging import Logger
+from logging.config import dictConfig
+from os import path
+from types import MethodType
+from typing import Any, Optional, cast
+
+import fastvideo.v1.envs as envs
+
+FASTVIDEO_CONFIGURE_LOGGING = envs.FASTVIDEO_CONFIGURE_LOGGING
+FASTVIDEO_LOGGING_CONFIG_PATH = envs.FASTVIDEO_LOGGING_CONFIG_PATH
+FASTVIDEO_LOGGING_LEVEL = envs.FASTVIDEO_LOGGING_LEVEL
+FASTVIDEO_LOGGING_PREFIX = envs.FASTVIDEO_LOGGING_PREFIX
+
+_FORMAT = (f"{FASTVIDEO_LOGGING_PREFIX}%(levelname)s %(asctime)s "
+           "[%(filename)s:%(lineno)d] %(message)s")
+_DATE_FORMAT = "%m-%d %H:%M:%S"
+
+DEFAULT_LOGGING_CONFIG = {
+    "formatters": {
+        "fastvideo": {
+            "class": "fastvideo.v1.logging_utils.NewLineFormatter",
+            "datefmt": _DATE_FORMAT,
+            "format": _FORMAT,
+        },
+    },
+    "handlers": {
+        "fastvideo": {
+            "class": "logging.StreamHandler",
+            "formatter": "fastvideo",
+            "level": FASTVIDEO_LOGGING_LEVEL,
+            "stream": "ext://sys.stdout",
+        },
+    },
+    "loggers": {
+        "fastvideo": {
+            "handlers": ["fastvideo"],
+            "level": "DEBUG",
+            "propagate": False,
+        },
+    },
+    "root": {
+        "handlers": ["fastvideo"],
+        "level": "DEBUG",
+    },
+    "version": 1,
+    "disable_existing_loggers": False
+}
+
+
+@lru_cache
+def _print_info_once(logger: Logger, msg: str) -> None:
+    # Set the stacklevel to 2 to print the original caller's line info
+    logger.info(msg, stacklevel=2)
+
+
+@lru_cache
+def _print_warning_once(logger: Logger, msg: str) -> None:
+    # Set the stacklevel to 2 to print the original caller's line info
+    logger.warning(msg, stacklevel=2)
+
+
+class _FastvideoLogger(Logger):
+    """
+    Note:
+        This class is just to provide type information.
+        We actually patch the methods directly on the :class:`logging.Logger`
+        instance to avoid conflicting with other libraries such as
+        `intel_extension_for_pytorch.utils._logger`.
+    """
+
+    def info_once(self, msg: str) -> None:
+        """
+        As :meth:`info`, but subsequent calls with the same message
+        are silently dropped.
+        """
+        _print_info_once(self, msg)
+
+    def warning_once(self, msg: str) -> None:
+        """
+        As :meth:`warning`, but subsequent calls with the same message
+        are silently dropped.
+        """
+        _print_warning_once(self, msg)
+
+
+def _configure_fastvideo_root_logger() -> None:
+    logging_config = dict[str, Any]()
+
+    if not FASTVIDEO_CONFIGURE_LOGGING and FASTVIDEO_LOGGING_CONFIG_PATH:
+        raise RuntimeError(
+            "FASTVIDEO_CONFIGURE_LOGGING evaluated to false, but "
+            "FASTVIDEO_LOGGING_CONFIG_PATH was given. FASTVIDEO_LOGGING_CONFIG_PATH "
+            "implies FASTVIDEO_CONFIGURE_LOGGING. Please enable "
+            "FASTVIDEO_CONFIGURE_LOGGING or unset FASTVIDEO_LOGGING_CONFIG_PATH."
+        )
+
+    if FASTVIDEO_CONFIGURE_LOGGING:
+        logging_config = DEFAULT_LOGGING_CONFIG
+
+    if FASTVIDEO_LOGGING_CONFIG_PATH:
+        if not path.exists(FASTVIDEO_LOGGING_CONFIG_PATH):
+            raise RuntimeError(
+                "Could not load logging config. File does not exist: %s",
+                FASTVIDEO_LOGGING_CONFIG_PATH)
+        with open(FASTVIDEO_LOGGING_CONFIG_PATH, encoding="utf-8") as file:
+            custom_config = json.loads(file.read())
+
+        if not isinstance(custom_config, dict):
+            raise ValueError("Invalid logging config. Expected Dict, got %s.",
+                             type(custom_config).__name__)
+        logging_config = custom_config
+
+    for formatter in logging_config.get("formatters", {}).values():
+        # This provides backwards compatibility after #10134.
+        if formatter.get("class") == "fastvideo.v1.logging.NewLineFormatter":
+            formatter["class"] = "fastvideo.v1.logging_utils.NewLineFormatter"
+
+    if logging_config:
+        dictConfig(logging_config)
+
+
+# TODO: add rank_zero_only log
+def init_logger(name: str) -> _FastvideoLogger:
+    """The main purpose of this function is to ensure that loggers are
+    retrieved in such a way that we can be sure the root fastvideo logger has
+    already been configured."""
+
+    logger = logging.getLogger(name)
+
+    methods_to_patch = {
+        "info_once": _print_info_once,
+        "warning_once": _print_warning_once,
+    }
+
+    for method_name, method in methods_to_patch.items():
+        setattr(logger, method_name, MethodType(method, logger))
+
+    return cast(_FastvideoLogger, logger)
+
+
+# The root logger is initialized when the module is imported.
+# This is thread-safe as the module is only imported once,
+# guaranteed by the Python GIL.
+_configure_fastvideo_root_logger()
+
+logger = init_logger(__name__)
+
+
+def _trace_calls(log_path, root_dir, frame, event, arg=None):
+    if event in ['call', 'return']:
+        # Extract the filename, line number, function name, and the code object
+        filename = frame.f_code.co_filename
+        lineno = frame.f_lineno
+        func_name = frame.f_code.co_name
+        if not filename.startswith(root_dir):
+            # only log the functions in the fastvideo root_dir
+            return
+        # Log every function call or return
+        try:
+            last_frame = frame.f_back
+            if last_frame is not None:
+                last_filename = last_frame.f_code.co_filename
+                last_lineno = last_frame.f_lineno
+                last_func_name = last_frame.f_code.co_name
+            else:
+                # initial frame
+                last_filename = ""
+                last_lineno = 0
+                last_func_name = ""
+            with open(log_path, 'a') as f:
+                ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
+                if event == 'call':
+                    f.write(f"{ts} Call to"
+                            f" {func_name} in {filename}:{lineno}"
+                            f" from {last_func_name} in {last_filename}:"
+                            f"{last_lineno}\n")
+                else:
+                    f.write(f"{ts} Return from"
+                            f" {func_name} in {filename}:{lineno}"
+                            f" to {last_func_name} in {last_filename}:"
+                            f"{last_lineno}\n")
+        except NameError:
+            # modules are deleted during shutdown
+            pass
+    return partial(_trace_calls, log_path, root_dir)
+
+
+def enable_trace_function_call(log_file_path: str,
+                               root_dir: Optional[str] = None):
+    """
+    Enable tracing of every function call in code under `root_dir`.
+    This is useful for debugging hangs or crashes.
+    `log_file_path` is the path to the log file.
+    `root_dir` is the root directory of the code to trace. If None, it is the
+    fastvideo root directory.
+
+    Note that this call is thread-level, any threads calling this function
+    will have the trace enabled. Other threads will not be affected.
+    """
+    logger.warning(
+        "FASTVIDEO_TRACE_FUNCTION is enabled. It will record every"
+        " function executed by Python. This will slow down the code. It "
+        "is suggested to be used for debugging hang or crashes only.")
+    logger.info("Trace frame log is saved to %s", log_file_path)
+    if root_dir is None:
+        # by default, this is the fastvideo root directory
+        root_dir = os.path.dirname(os.path.dirname(__file__))
+    sys.settrace(partial(_trace_calls, log_file_path, root_dir))

fastvideo/v1/logging_utils/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.v1.logging_utils.formatter import NewLineFormatter
+
+__all__ = [
+    "NewLineFormatter",
+]

fastvideo/v1/logging_utils/formatter.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/logging_utils/formatter.py
+
+import logging
+
+
+class NewLineFormatter(logging.Formatter):
+    """Adds logging prefix to newlines to align multi-line messages."""
+
+    def __init__(self, fmt, datefmt=None, style="%"):
+        logging.Formatter.__init__(self, fmt, datefmt, style)
+
+    def format(self, record):
+        msg = logging.Formatter.format(self, record)
+        if record.message != "":
+            parts = msg.split(record.message)
+            msg = msg.replace("\n", "\r\n" + parts[0])
+        return msg

fastvideo/v1/models/dits/base.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+from abc import ABC, abstractmethod
+from typing import List, Union
+
+import torch
+from torch import nn
+
+
+# TODO
+class BaseDiT(nn.Module, ABC):
+    _fsdp_shard_conditions: list = []
+    attention_head_dim: int | None = None
+    _param_names_mapping: dict
+    hidden_size: int
+    num_attention_heads: int
+
+    def __init_subclass__(cls):
+        required_class_attrs = [
+            "_fsdp_shard_conditions", "_param_names_mapping"
+        ]
+        super().__init_subclass__()
+        for attr in required_class_attrs:
+            if not hasattr(cls, attr):
+                raise AttributeError(
+                    f"Subclasses of BaseDiT must define '{attr}' class variable"
+                )
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__()
+
+    @abstractmethod
+    def forward(self,
+                hidden_states: torch.Tensor,
+                encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+                timestep: torch.LongTensor,
+                guidance=None,
+                **kwargs) -> torch.Tensor:
+        pass
+
+    def __post_init__(self):
+        required_attrs = ["hidden_size", "num_attention_heads"]
+        for attr in required_attrs:
+            if not hasattr(self, attr):
+                raise AttributeError(
+                    f"Subclasses of BaseDiT must define '{attr}' instance variable"
+                )

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -0,0 +1,932 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from fastvideo.v1.attention import DistributedAttention, LocalAttention
+from fastvideo.v1.distributed.parallel_state import (
+    get_sequence_model_parallel_world_size)
+from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, ScaleResidual,
+                                           ScaleResidualLayerNormScaleShift)
+from fastvideo.v1.layers.linear import ReplicatedLinear
+# TODO(will-PY-refactor): RMSNorm ....
+from fastvideo.v1.layers.mlp import MLP
+from fastvideo.v1.layers.rotary_embedding import (_apply_rotary_emb,
+                                                  get_rotary_pos_embed)
+from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
+                                                  PatchEmbed, TimestepEmbedder,
+                                                  unpatchify)
+from fastvideo.v1.models.dits.base import BaseDiT
+
+
+class HunyuanRMSNorm(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        elementwise_affine=True,
+        eps: float = 1e-6,
+        device=None,
+        dtype=None,
+    ):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.eps = eps
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
+
+    def _norm(self, x) -> torch.Tensor:
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        if hasattr(self, "weight"):
+            output = output * self.weight
+        return output
+
+
+class MMDoubleStreamBlock(nn.Module):
+    """
+    A multimodal DiT block with separate modulation for text and image/video,
+    using distributed attention and linear layers.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        mlp_ratio: float,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+
+        self.deterministic = False
+        self.num_attention_heads = num_attention_heads
+        head_dim = hidden_size // num_attention_heads
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        # Image modulation components
+        self.img_mod = ModulateProjection(
+            hidden_size,
+            factor=6,
+            act_layer="silu",
+            dtype=dtype,
+        )
+
+        # Fused operations for image stream
+        self.img_attn_norm = LayerNormScaleShift(hidden_size,
+                                                 norm_type="layer",
+                                                 elementwise_affine=False,
+                                                 dtype=dtype)
+        self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            elementwise_affine=False,
+            dtype=dtype)
+        self.img_mlp_residual = ScaleResidual()
+
+        # Image attention components
+        self.img_attn_qkv = ReplicatedLinear(hidden_size,
+                                             hidden_size * 3,
+                                             bias=True,
+                                             params_dtype=dtype)
+
+        self.img_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+        self.img_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+
+        self.img_attn_proj = ReplicatedLinear(hidden_size,
+                                              hidden_size,
+                                              bias=True,
+                                              params_dtype=dtype)
+
+        self.img_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)
+
+        # Text modulation components
+        self.txt_mod = ModulateProjection(
+            hidden_size,
+            factor=6,
+            act_layer="silu",
+            dtype=dtype,
+        )
+
+        # Fused operations for text stream
+        self.txt_attn_norm = LayerNormScaleShift(hidden_size,
+                                                 norm_type="layer",
+                                                 elementwise_affine=False,
+                                                 dtype=dtype)
+        self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            elementwise_affine=False,
+            dtype=dtype)
+        self.txt_mlp_residual = ScaleResidual()
+
+        # Text attention components
+        self.txt_attn_qkv = ReplicatedLinear(hidden_size,
+                                             hidden_size * 3,
+                                             bias=True,
+                                             params_dtype=dtype)
+
+        # QK norm layers for text
+        self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+        self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+
+        self.txt_attn_proj = ReplicatedLinear(hidden_size,
+                                              hidden_size,
+                                              bias=True,
+                                              params_dtype=dtype)
+
+        self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)
+
+        # Distributed attention
+        self.attn = DistributedAttention(num_heads=num_attention_heads,
+                                         head_size=head_dim,
+                                         dropout_rate=0.0,
+                                         causal=False)
+
+        # QK norm layers for text
+        self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+        self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+
+        self.txt_attn_proj = ReplicatedLinear(hidden_size,
+                                              hidden_size,
+                                              bias=True,
+                                              params_dtype=dtype)
+
+        self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)
+
+        # Distributed attention
+        self.attn = DistributedAttention(num_heads=num_attention_heads,
+                                         head_size=head_dim,
+                                         dropout_rate=0.0,
+                                         causal=False)
+
+    def forward(
+        self,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        vec: torch.Tensor,
+        freqs_cis: tuple,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Process modulation vectors
+        img_mod_outputs = self.img_mod(vec)
+        (
+            img_attn_shift,
+            img_attn_scale,
+            img_attn_gate,
+            img_mlp_shift,
+            img_mlp_scale,
+            img_mlp_gate,
+        ) = torch.chunk(img_mod_outputs, 6, dim=-1)
+
+        txt_mod_outputs = self.txt_mod(vec)
+        (
+            txt_attn_shift,
+            txt_attn_scale,
+            txt_attn_gate,
+            txt_mlp_shift,
+            txt_mlp_scale,
+            txt_mlp_gate,
+        ) = torch.chunk(txt_mod_outputs, 6, dim=-1)
+
+        # Prepare image for attention using fused operation
+        img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale)
+        # Get QKV for image
+        img_qkv, _ = self.img_attn_qkv(img_attn_input)
+        batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1]
+
+        # Split QKV
+        img_qkv = img_qkv.view(batch_size, image_seq_len, 3,
+                               self.num_attention_heads, -1)
+        img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :,
+                                                                          2]
+
+        # Apply QK-Norm if needed
+
+        img_q = self.img_attn_q_norm(img_q).to(img_v)
+        img_k = self.img_attn_k_norm(img_k).to(img_v)
+        # Apply rotary embeddings
+        cos, sin = freqs_cis
+        img_q, img_k = _apply_rotary_emb(
+            img_q, cos, sin,
+            is_neox_style=False), _apply_rotary_emb(img_k,
+                                                    cos,
+                                                    sin,
+                                                    is_neox_style=False)
+        # Prepare text for attention using fused operation
+        txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale)
+
+        # Get QKV for text
+        txt_qkv, _ = self.txt_attn_qkv(txt_attn_input)
+        batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1]
+
+        # Split QKV
+        txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3,
+                               self.num_attention_heads, -1)
+        txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :,
+                                                                          2]
+
+        # Apply QK-Norm if needed
+        txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
+        txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
+
+        # Run distributed attention
+        img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v)
+        img_attn_out, _ = self.img_attn_proj(
+            img_attn.view(batch_size, image_seq_len, -1))
+        # Use fused operation for residual connection, normalization, and modulation
+        img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(
+            img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale)
+
+        # Process image MLP
+        img_mlp_out = self.img_mlp(img_mlp_input)
+        img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate)
+
+        # Process text attention output
+        txt_attn_out, _ = self.txt_attn_proj(
+            txt_attn.reshape(batch_size, text_seq_len, -1))
+
+        # Use fused operation for residual connection, normalization, and modulation
+        txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(
+            txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale)
+
+        # Process text MLP
+        txt_mlp_out = self.txt_mlp(txt_mlp_input)
+        txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate)
+
+        return img, txt
+
+
+class MMSingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers using distributed attention
+    and tensor parallelism.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        mlp_ratio: float = 4.0,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+
+        self.deterministic = False
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        head_dim = hidden_size // num_attention_heads
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp_hidden_dim = mlp_hidden_dim
+
+        # Combined QKV and MLP input projection
+        self.linear1 = ReplicatedLinear(hidden_size,
+                                        hidden_size * 3 + mlp_hidden_dim,
+                                        bias=True,
+                                        params_dtype=dtype)
+
+        # Combined projection and MLP output
+        self.linear2 = ReplicatedLinear(hidden_size + mlp_hidden_dim,
+                                        hidden_size,
+                                        bias=True,
+                                        params_dtype=dtype)
+
+        # QK norm layers
+        self.q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+        self.k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype)
+
+        # Fused operations with better naming
+        self.input_norm_scale_shift = LayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            eps=1e-6,
+            elementwise_affine=False,
+            dtype=dtype)
+        self.output_residual = ScaleResidual()
+
+        # Activation function
+        self.mlp_act = nn.GELU(approximate="tanh")
+
+        # Modulation
+        self.modulation = ModulateProjection(hidden_size,
+                                             factor=3,
+                                             act_layer="silu",
+                                             dtype=dtype)
+
+        # Distributed attention
+        self.attn = DistributedAttention(num_heads=num_attention_heads,
+                                         head_size=head_dim,
+                                         dropout_rate=0.0,
+                                         causal=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        vec: torch.Tensor,
+        txt_len: int,
+        freqs_cis: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        # Process modulation
+        mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
+
+        # Apply pre-norm and modulation using fused operation
+        x_mod = self.input_norm_scale_shift(x, mod_shift, mod_scale)
+
+        # Get combined projections
+        linear1_out, _ = self.linear1(x_mod)
+
+        # Split into QKV and MLP parts
+        qkv, mlp = torch.split(linear1_out,
+                               [3 * self.hidden_size, self.mlp_hidden_dim],
+                               dim=-1)
+
+        # Process QKV
+        batch_size, seq_len = qkv.shape[0], qkv.shape[1]
+        qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+
+        # Apply QK-Norm
+        q = self.q_norm(q).to(v.dtype)
+        k = self.k_norm(k).to(v.dtype)
+
+        # Split into image and text parts
+        img_q, txt_q = q[:, :-txt_len], q[:, -txt_len:]
+        img_k, txt_k = k[:, :-txt_len], k[:, -txt_len:]
+        img_v, txt_v = v[:, :-txt_len], v[:, -txt_len:]
+        # Apply rotary embeddings to image parts
+        cos, sin = freqs_cis
+        img_q, img_k = _apply_rotary_emb(
+            img_q, cos, sin,
+            is_neox_style=False), _apply_rotary_emb(img_k,
+                                                    cos,
+                                                    sin,
+                                                    is_neox_style=False)
+
+        # Run distributed attention
+        img_attn_output, txt_attn_output = self.attn(img_q, img_k, img_v, txt_q,
+                                                     txt_k, txt_v)
+        attn_output = torch.cat((img_attn_output, txt_attn_output),
+                                dim=1).view(batch_size, seq_len, -1)
+        # Process MLP activation
+        mlp_output = self.mlp_act(mlp)
+
+        # Combine attention and MLP outputs
+        combined = torch.cat((attn_output, mlp_output), dim=-1)
+
+        # Final projection
+        output, _ = self.linear2(combined)
+
+        # Apply residual connection with gating using fused operation
+        return self.output_residual(x, output, mod_gate)
+
+
+class HunyuanVideoTransformer3DModel(BaseDiT):
+    """
+    HunyuanVideo Transformer backbone adapted for distributed training.
+    
+    This implementation uses distributed attention and linear layers for efficient
+    parallel processing across multiple GPUs.
+    
+    Based on the architecture from:
+    - Flux.1: https://github.com/black-forest-labs/flux
+    - MMDiT: http://arxiv.org/abs/2403.03206
+    """
+    # PY: we make the input args the same as HF config
+
+    # shard single stream, double stream blocks, and refiner_blocks
+    _fsdp_shard_conditions = [
+        lambda n, m: "double" in n and str.isdigit(n.split(".")[-1]),
+        lambda n, m: "single" in n and str.isdigit(n.split(".")[-1]),
+        lambda n, m: "refiner" in n and str.isdigit(n.split(".")[-1]),
+    ]
+    _param_names_mapping = {
+        # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
+        r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
+        r"txt_in.t_embedder.mlp.fc_in.\1",
+        r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
+        r"txt_in.t_embedder.mlp.fc_out.\1",
+        r"^context_embedder\.proj_in\.(.*)$":
+        r"txt_in.input_embedder.\1",
+        r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$":
+        r"txt_in.c_embedder.fc_in.\1",
+        r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$":
+        r"txt_in.c_embedder.fc_out.\1",
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$":
+        r"txt_in.refiner_blocks.\1.norm1.\2",
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$":
+        r"txt_in.refiner_blocks.\1.norm2.\2",
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+        (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 0, 3),
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+        (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 1, 3),
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+        (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 2, 3),
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
+        r"txt_in.refiner_blocks.\1.self_attn_proj.\2",
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
+        r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
+        r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
+        r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$":
+        r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
+
+        # 3. x_embedder mapping:
+        r"^x_embedder\.proj\.(.*)$":
+        r"img_in.proj.\1",
+
+        # 4. Top-level time_text_embed mappings:
+        r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
+        r"time_in.mlp.fc_in.\1",
+        r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
+        r"time_in.mlp.fc_out.\1",
+        r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$":
+        r"guidance_in.mlp.fc_in.\1",
+        r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$":
+        r"guidance_in.mlp.fc_out.\1",
+        r"^time_text_embed\.text_embedder\.linear_1\.(.*)$":
+        r"vector_in.fc_in.\1",
+        r"^time_text_embed\.text_embedder\.linear_2\.(.*)$":
+        r"vector_in.fc_out.\1",
+
+        # 5. transformer_blocks mapping:
+        r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$":
+        r"double_blocks.\1.img_mod.linear.\2",
+        r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$":
+        r"double_blocks.\1.txt_mod.linear.\2",
+        r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
+        r"double_blocks.\1.img_attn_q_norm.\2",
+        r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
+        r"double_blocks.\1.img_attn_k_norm.\2",
+        r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+        (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
+        r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+        (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
+        r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+        (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
+        r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
+        (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
+        r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
+        (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
+        r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
+        (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
+        r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
+        r"double_blocks.\1.img_attn_proj.\2",
+        # Corrected: merge attn.to_add_out into the main projection.
+        r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$":
+        r"double_blocks.\1.txt_attn_proj.\2",
+        r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$":
+        r"double_blocks.\1.txt_attn_q_norm.\2",
+        r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$":
+        r"double_blocks.\1.txt_attn_k_norm.\2",
+        r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
+        r"double_blocks.\1.img_mlp.fc_in.\2",
+        r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
+        r"double_blocks.\1.img_mlp.fc_out.\2",
+        r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$":
+        r"double_blocks.\1.txt_mlp.fc_in.\2",
+        r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$":
+        r"double_blocks.\1.txt_mlp.fc_out.\2",
+
+        # 6. single_transformer_blocks mapping:
+        r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
+        r"single_blocks.\1.q_norm.\2",
+        r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
+        r"single_blocks.\1.k_norm.\2",
+        r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+        (r"single_blocks.\1.linear1.\2", 0, 4),
+        r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+        (r"single_blocks.\1.linear1.\2", 1, 4),
+        r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+        (r"single_blocks.\1.linear1.\2", 2, 4),
+        r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$":
+        (r"single_blocks.\1.linear1.\2", 3, 4),
+        # Corrected: map proj_out to modulation.linear rather than a separate proj_out branch.
+        r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$":
+        r"single_blocks.\1.linear2.\2",
+        r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$":
+        r"single_blocks.\1.modulation.linear.\2",
+
+        # 7. Final layers mapping:
+        r"^norm_out\.linear\.(.*)$":
+        r"final_layer.adaLN_modulation.linear.\1",
+        r"^proj_out\.(.*)$":
+        r"final_layer.linear.\1",
+    }
+
+    def __init__(
+            self,
+            patch_size: int = 2,
+            patch_size_t: int = 1,
+            in_channels: int = 16,
+            out_channels: int = 16,
+            num_attention_heads: int = 24,
+            attention_head_dim: int = 128,
+            mlp_ratio: float = 4.0,
+            num_layers: int = 20,
+            num_single_layers: int = 40,
+            num_refiner_layers: int = 2,
+            rope_axes_dim: Tuple[int, int, int] = (16, 56, 56),
+            guidance_embeds: bool = False,
+            dtype: Optional[torch.dtype] = None,
+            text_embed_dim: int = 4096,
+            pooled_projection_dim: int = 768,
+            rope_theta: int = 256,
+            qk_norm: str = "rms_norm",  #TODO(PY)
+    ):
+        super().__init__()
+        hidden_size = attention_head_dim * num_attention_heads
+        self.patch_size = [patch_size_t, patch_size, patch_size]
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.unpatchify_channels = self.out_channels
+        self.guidance_embeds = guidance_embeds
+        self.rope_dim_list = list(rope_axes_dim)
+        self.rope_theta = rope_theta
+        self.text_states_dim = text_embed_dim
+        self.text_states_dim_2 = pooled_projection_dim
+        # TODO(will): hack?
+        self.dtype = dtype
+
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError(
+                f"Hidden size {hidden_size} must be divisible by num_attention_heads {num_attention_heads}"
+            )
+
+        pe_dim = hidden_size // num_attention_heads
+        if sum(rope_axes_dim) != pe_dim:
+            raise ValueError(
+                f"Got {rope_axes_dim} but expected positional dim {pe_dim}")
+
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+
+        # Image projection
+        self.img_in = PatchEmbed(self.patch_size,
+                                 self.in_channels,
+                                 self.hidden_size,
+                                 dtype=dtype)
+
+        self.txt_in = SingleTokenRefiner(self.text_states_dim,
+                                         hidden_size,
+                                         num_attention_heads,
+                                         depth=num_refiner_layers,
+                                         dtype=dtype)
+
+        # Time modulation
+        self.time_in = TimestepEmbedder(self.hidden_size,
+                                        act_layer="silu",
+                                        dtype=dtype)
+
+        # Text modulation
+        self.vector_in = MLP(self.text_states_dim_2,
+                             self.hidden_size,
+                             self.hidden_size,
+                             act_type="silu",
+                             dtype=dtype)
+
+        # Guidance modulation
+        self.guidance_in = (TimestepEmbedder(
+            self.hidden_size, act_layer="silu", dtype=dtype)
+                            if self.guidance_embeds else None)
+
+        # Double blocks
+        self.double_blocks = nn.ModuleList([
+            MMDoubleStreamBlock(
+                hidden_size,
+                num_attention_heads,
+                mlp_ratio=mlp_ratio,
+                dtype=dtype,
+            ) for _ in range(num_layers)
+        ])
+
+        # Single blocks
+        self.single_blocks = nn.ModuleList([
+            MMSingleStreamBlock(
+                hidden_size,
+                num_attention_heads,
+                mlp_ratio=mlp_ratio,
+                dtype=dtype,
+            ) for _ in range(num_single_layers)
+        ])
+
+        self.final_layer = FinalLayer(hidden_size,
+                                      self.patch_size,
+                                      self.out_channels,
+                                      dtype=dtype)
+
+        self.__post_init__()
+
+    # TODO: change the input the FORWAD_BACTCH Dict
+    # TODO: change output to a dict
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+        timestep: torch.LongTensor,
+        guidance=None,
+    ):
+        """
+        Forward pass of the HunyuanDiT model.
+        
+        Args:
+            hidden_states: Input image/video latents [B, C, T, H, W]
+            encoder_hidden_states: Text embeddings [B, L, D]
+            timestep: Diffusion timestep
+            guidance: Guidance scale for CFG
+            
+        Returns:
+            Tuple of (output)
+        """
+        if guidance is None:
+            guidance = torch.tensor([6016.0],
+                                    device=hidden_states.device,
+                                    dtype=hidden_states.dtype)
+
+        img = x = hidden_states
+        t = timestep
+
+        # Split text embeddings - first token is global, rest are per-token
+        if isinstance(encoder_hidden_states, torch.Tensor):
+            txt = encoder_hidden_states[:, 1:]
+            text_states_2 = encoder_hidden_states[:, 0, :self.text_states_dim_2]
+        else:
+            txt = encoder_hidden_states[0]
+            text_states_2 = encoder_hidden_states[1]
+
+        # Get spatial dimensions
+        _, _, ot, oh, ow = x.shape  # codespell:ignore
+        tt, th, tw = (
+            ot // self.patch_size[0],  # codespell:ignore
+            oh // self.patch_size[1],
+            ow // self.patch_size[2],
+        )
+
+        # Get rotary embeddings
+        freqs_cos, freqs_sin = get_rotary_pos_embed(
+            (tt * get_sequence_model_parallel_world_size(), th, tw),
+            self.hidden_size, self.num_attention_heads, self.rope_dim_list,
+            self.rope_theta)
+        freqs_cos = freqs_cos.to(x.device)
+        freqs_sin = freqs_sin.to(x.device)
+        # Prepare modulation vectors
+        vec = self.time_in(t)
+
+        # Add text modulation
+        vec = vec + self.vector_in(text_states_2)
+
+        # Add guidance modulation if needed
+        if self.guidance_in and guidance is not None:
+            vec = vec + self.guidance_in(guidance)
+        # Embed image and text
+        img = self.img_in(img)
+        txt = self.txt_in(txt, t)
+        txt_seq_len = txt.shape[1]
+        img_seq_len = img.shape[1]
+
+        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
+        # Process through double stream blocks
+        for index, block in enumerate(self.double_blocks):
+            double_block_args = [img, txt, vec, freqs_cis]
+            img, txt = block(*double_block_args)
+        # Merge txt and img to pass through single stream blocks
+        x = torch.cat((img, txt), 1)
+
+        # Process through single stream blocks
+        if len(self.single_blocks) > 0:
+            for index, block in enumerate(self.single_blocks):
+                single_block_args = [
+                    x,
+                    vec,
+                    txt_seq_len,
+                    freqs_cis,
+                ]
+                x = block(*single_block_args)
+
+        # Extract image features
+        img = x[:, :img_seq_len, ...]
+        # Final layer processing
+        img = self.final_layer(img, vec)
+        # Unpatchify to get original shape
+        img = unpatchify(img, tt, th, tw, self.patch_size, self.out_channels)
+
+        return img
+
+
+class SingleTokenRefiner(nn.Module):
+    """
+    A token refiner that processes text embeddings with attention to improve
+    their representation for cross-attention with image features.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_size,
+        num_attention_heads,
+        depth=2,
+        qkv_bias=True,
+        dtype=None,
+    ) -> None:
+        super().__init__()
+
+        # Input projection
+        self.input_embedder = ReplicatedLinear(in_channels,
+                                               hidden_size,
+                                               bias=True,
+                                               params_dtype=dtype)
+
+        # Timestep embedding
+        self.t_embedder = TimestepEmbedder(hidden_size,
+                                           act_layer="silu",
+                                           dtype=dtype)
+
+        # Context embedding
+        self.c_embedder = MLP(in_channels,
+                              hidden_size,
+                              hidden_size,
+                              act_type="silu",
+                              dtype=dtype)
+
+        # Refiner blocks
+        self.refiner_blocks = nn.ModuleList([
+            IndividualTokenRefinerBlock(
+                hidden_size,
+                num_attention_heads,
+                qkv_bias=qkv_bias,
+                dtype=dtype,
+            ) for _ in range(depth)
+        ])
+
+    def forward(self, x, t):
+        # Get timestep embeddings
+        timestep_aware_representations = self.t_embedder(t)
+
+        # Get context-aware representations
+
+        context_aware_representations = torch.mean(x, dim=1)
+
+        context_aware_representations = self.c_embedder(
+            context_aware_representations)
+        c = timestep_aware_representations + context_aware_representations
+        # Project input
+        x, _ = self.input_embedder(x)
+        # Process through refiner blocks
+        for block in self.refiner_blocks:
+            x = block(x, c)
+        return x
+
+
+class IndividualTokenRefinerBlock(nn.Module):
+    """
+    A transformer block for refining individual tokens with self-attention.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        num_attention_heads,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        dtype=None,
+    ) -> None:
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        # Normalization and attention
+        self.norm1 = nn.LayerNorm(hidden_size,
+                                  eps=1e-6,
+                                  elementwise_affine=True,
+                                  dtype=dtype)
+
+        self.self_attn_qkv = ReplicatedLinear(hidden_size,
+                                              hidden_size * 3,
+                                              bias=qkv_bias,
+                                              params_dtype=dtype)
+
+        self.self_attn_proj = ReplicatedLinear(hidden_size,
+                                               hidden_size,
+                                               bias=qkv_bias,
+                                               params_dtype=dtype)
+
+        # MLP
+        self.norm2 = nn.LayerNorm(hidden_size,
+                                  eps=1e-6,
+                                  elementwise_affine=True,
+                                  dtype=dtype)
+        self.mlp = MLP(hidden_size,
+                       mlp_hidden_dim,
+                       bias=True,
+                       act_type="silu",
+                       dtype=dtype)
+
+        # Modulation
+        self.adaLN_modulation = ModulateProjection(hidden_size,
+                                                   factor=2,
+                                                   act_layer="silu",
+                                                   dtype=dtype)
+
+        # Scaled dot product attention
+        self.attn = LocalAttention(
+            num_heads=num_attention_heads,
+            head_size=hidden_size // num_attention_heads,
+        )
+
+    def forward(self, x, c):
+        # Get modulation parameters
+        gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1)
+        # Self-attention
+        norm_x = self.norm1(x)
+        qkv, _ = self.self_attn_qkv(norm_x)
+
+        batch_size, seq_len = qkv.shape[0], qkv.shape[1]
+        qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+
+        # Run scaled dot product attention
+        attn_output = self.attn(q, k, v)  # [B, L, H, D]
+        attn_output = attn_output.reshape(batch_size, seq_len,
+                                          -1)  # [B, L, H*D]
+
+        # Project and apply residual connection with gating
+        attn_out, _ = self.self_attn_proj(attn_output)
+        x = x + attn_out * gate_msa.unsqueeze(1)
+
+        # MLP
+        mlp_out = self.mlp(self.norm2(x))
+        x = x + mlp_out * gate_mlp.unsqueeze(1)
+
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT that projects features to pixel space.
+    """
+
+    def __init__(self,
+                 hidden_size,
+                 patch_size,
+                 out_channels,
+                 dtype=None) -> None:
+        super().__init__()
+
+        # Normalization
+        self.norm_final = nn.LayerNorm(hidden_size,
+                                       eps=1e-6,
+                                       elementwise_affine=False,
+                                       dtype=dtype)
+
+        output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels
+
+        self.linear = ReplicatedLinear(hidden_size,
+                                       output_dim,
+                                       bias=True,
+                                       params_dtype=dtype)
+
+        # Modulation
+        self.adaLN_modulation = ModulateProjection(hidden_size,
+                                                   factor=2,
+                                                   act_layer="silu",
+                                                   dtype=dtype)
+
+    def forward(self, x, c):
+        # What the heck HF? Why you change the scale and shift order here???
+        scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        x, _ = self.linear(x)
+        return x

fastvideo/v1/models/dits/wanvideo.py

@@ -0,0 +1,528 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from fastvideo.v1.attention import DistributedAttention, LocalAttention
+from fastvideo.v1.distributed.parallel_state import (
+    get_sequence_model_parallel_world_size)
+from fastvideo.v1.layers.layernorm import (LayerNormScaleShift, RMSNorm,
+                                           ScaleResidual,
+                                           ScaleResidualLayerNormScaleShift)
+from fastvideo.v1.layers.linear import ReplicatedLinear
+# from torch.nn import RMSNorm
+# TODO: RMSNorm ....
+from fastvideo.v1.layers.mlp import MLP
+from fastvideo.v1.layers.rotary_embedding import (_apply_rotary_emb,
+                                                  get_rotary_pos_embed)
+from fastvideo.v1.layers.visual_embedding import (ModulateProjection,
+                                                  PatchEmbed, TimestepEmbedder)
+from fastvideo.v1.models.dits.base import BaseDiT
+
+
+class WanImageEmbedding(torch.nn.Module):
+
+    def __init__(self, in_features: int, out_features: int):
+        super().__init__()
+
+        self.norm1 = nn.LayerNorm(in_features)
+        self.ff = MLP(in_features, in_features, out_features, act_type="gelu")
+        self.norm2 = nn.LayerNorm(out_features)
+
+    def forward(self,
+                encoder_hidden_states_image: torch.Tensor) -> torch.Tensor:
+        dtype = encoder_hidden_states_image.dtype
+        hidden_states = self.norm1(encoder_hidden_states_image)
+        hidden_states = self.ff(hidden_states)
+        hidden_states = self.norm2(hidden_states).to(dtype)
+        return hidden_states
+
+
+class WanTimeTextImageEmbedding(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        time_freq_dim: int,
+        text_embed_dim: int,
+        image_embed_dim: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.time_embedder = TimestepEmbedder(
+            dim, frequency_embedding_size=time_freq_dim, act_layer="silu")
+        self.time_modulation = ModulateProjection(dim,
+                                                  factor=6,
+                                                  act_layer="silu")
+        self.text_embedder = MLP(text_embed_dim,
+                                 dim,
+                                 dim,
+                                 bias=True,
+                                 act_type="gelu_pytorch_tanh")
+
+        self.image_embedder = None
+        if image_embed_dim is not None:
+            self.image_embedder = WanImageEmbedding(image_embed_dim, dim)
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_states_image: Optional[torch.Tensor] = None,
+    ):
+        temb = self.time_embedder(timestep)
+        timestep_proj = self.time_modulation(temb)
+
+        encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+        if encoder_hidden_states_image is not None:
+            assert self.image_embedder is not None
+            encoder_hidden_states_image = self.image_embedder(
+                encoder_hidden_states_image)
+
+        return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image
+
+
+class WanSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6,
+                 parallel_attention=False) -> None:
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+        self.parallel_attention = parallel_attention
+
+        # layers
+        self.to_q = ReplicatedLinear(dim, dim)
+        self.to_k = ReplicatedLinear(dim, dim)
+        self.to_v = ReplicatedLinear(dim, dim)
+        self.to_out = ReplicatedLinear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+        # Scaled dot product attention
+        self.attn = LocalAttention(num_heads=num_heads,
+                                   head_size=self.head_dim,
+                                   dropout_rate=0,
+                                   softmax_scale=None,
+                                   causal=False)
+
+    def forward(self, x: torch.Tensor, context: torch.Tensor,
+                context_lens: int):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        pass
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q.forward_native(self.to_q(x)[0]).view(b, -1, n, d)
+        k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d)
+        v = self.to_v(context)[0].view(b, -1, n, d)
+
+        # compute attention
+        x = self.attn(q, k, v)
+
+        # output
+        x = x.flatten(2)
+        x, _ = self.to_out(x)
+        return x
+
+
+class WanI2VCrossAttention(WanSelfAttention):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6) -> None:
+        super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+        self.add_k_proj = ReplicatedLinear(dim, dim)
+        self.add_v_proj = ReplicatedLinear(dim, dim)
+        self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        context_img = context[:, :257]
+        context = context[:, 257:]
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q.forward_native(self.to_q(x)[0]).view(b, -1, n, d)
+        k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d)
+        v = self.to_v(context)[0].view(b, -1, n, d)
+        k_img = self.norm_added_k.forward_native(
+            self.add_k_proj(context_img)[0]).view(b, -1, n, d)
+        v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d)
+        img_x = self.attn(q, k_img, v_img)
+        # compute attention
+        x = self.attn(q, k, v)
+
+        # output
+        x = x.flatten(2)
+        img_x = img_x.flatten(2)
+        x = x + img_x
+        x, _ = self.to_out(x)
+        return x
+
+
+class WanTransformerBlock(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        qk_norm: str = "rms_norm_across_heads",
+        cross_attn_norm: bool = False,
+        eps: float = 1e-6,
+        added_kv_proj_dim: Optional[int] = None,
+    ):
+        super().__init__()
+
+        # 1. Self-attention
+        self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False)
+        self.to_q = ReplicatedLinear(dim, dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, dim, bias=True)
+        self.to_out = ReplicatedLinear(dim, dim, bias=True)
+        self.attn1 = DistributedAttention(num_heads=num_heads,
+                                          head_size=dim // num_heads,
+                                          dropout_rate=0.0,
+                                          causal=False)
+        self.hidden_dim = dim
+        self.num_attention_heads = num_heads
+        dim_head = dim // num_heads
+        if qk_norm == "rms_norm":
+            self.norm_q = RMSNorm(dim_head, eps=eps)
+            self.norm_k = RMSNorm(dim_head, eps=eps)
+        elif qk_norm == "rms_norm_across_heads":
+            # LTX applies qk norm across all heads
+            self.norm_q = RMSNorm(dim, eps=eps)
+            self.norm_k = RMSNorm(dim, eps=eps)
+        else:
+            print("QK Norm type not supported")
+            raise Exception
+        assert cross_attn_norm is True
+        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
+            dim,
+            norm_type="layer",
+            eps=eps,
+            elementwise_affine=True,
+            dtype=torch.float32)
+
+        # 2. Cross-attention
+        if added_kv_proj_dim is not None:
+            # I2V
+            self.attn2 = WanI2VCrossAttention(dim,
+                                              num_heads,
+                                              qk_norm=qk_norm,
+                                              eps=eps)
+        else:
+            # T2V
+            self.attn2 = WanT2VCrossAttention(dim,
+                                              num_heads,
+                                              qk_norm=qk_norm,
+                                              eps=eps)
+        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
+            dim,
+            norm_type="layer",
+            eps=eps,
+            elementwise_affine=False,
+            dtype=torch.float32)
+
+        # 3. Feed-forward
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.mlp_residual = ScaleResidual()
+
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        freqs_cis: Tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        if hidden_states.dim() == 4:
+            hidden_states = hidden_states.squeeze(1)
+        bs, seq_length, _ = hidden_states.shape
+        orig_dtype = hidden_states.dtype
+        assert orig_dtype != torch.float32
+        e = self.scale_shift_table + temb.float()
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
+            6, dim=1)
+        assert shift_msa.dtype == torch.float32
+
+        # 1. Self-attention
+        norm_hidden_states = (self.norm1(hidden_states.float()) *
+                              (1 + scale_msa) + shift_msa).to(orig_dtype)
+        query, _ = self.to_q(norm_hidden_states)
+        key, _ = self.to_k(norm_hidden_states)
+        value, _ = self.to_v(norm_hidden_states)
+
+        if self.norm_q is not None:
+            query = self.norm_q.forward_native(query)
+        if self.norm_k is not None:
+            key = self.norm_k.forward_native(key)
+
+        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        # Apply rotary embeddings
+        cos, sin = freqs_cis
+        query, key = _apply_rotary_emb(query, cos, sin,
+                                       is_neox_style=False), _apply_rotary_emb(
+                                           key, cos, sin, is_neox_style=False)
+
+        attn_output, _ = self.attn1(query, key, value)
+        attn_output = attn_output.flatten(2)
+        attn_output, _ = self.to_out(attn_output)
+        attn_output = attn_output.squeeze(1)
+
+        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
+        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
+            hidden_states, attn_output, gate_msa, null_shift, null_scale)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(
+            orig_dtype), hidden_states.to(orig_dtype)
+
+        # 2. Cross-attention
+        attn_output = self.attn2(norm_hidden_states,
+                                 context=encoder_hidden_states,
+                                 context_lens=None)
+        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
+            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(
+            orig_dtype), hidden_states.to(orig_dtype)
+
+        # 3. Feed-forward
+        ff_output = self.ffn(norm_hidden_states)
+        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
+        hidden_states = hidden_states.to(orig_dtype)
+
+        return hidden_states
+
+
+class WanTransformer3DModel(BaseDiT):
+    _fsdp_shard_conditions = [
+        lambda n, m: "blocks" in n and str.isdigit(n.split(".")[-1]),
+    ]
+    _param_names_mapping = {
+        r"^patch_embedding\.(.*)$":
+        r"patch_embedding.proj.\1",
+        r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
+        r"condition_embedder.text_embedder.fc_in.\1",
+        r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":
+        r"condition_embedder.text_embedder.fc_out.\1",
+        r"^condition_embedder\.time_embedder\.linear_1\.(.*)$":
+        r"condition_embedder.time_embedder.mlp.fc_in.\1",
+        r"^condition_embedder\.time_embedder\.linear_2\.(.*)$":
+        r"condition_embedder.time_embedder.mlp.fc_out.\1",
+        r"^condition_embedder\.time_proj\.(.*)$":
+        r"condition_embedder.time_modulation.linear.\1",
+        r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
+        r"condition_embedder.image_embedder.ff.fc_in.\1",
+        r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$":
+        r"condition_embedder.image_embedder.ff.fc_out.\1",
+        r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$":
+        r"blocks.\1.to_q.\2",
+        r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$":
+        r"blocks.\1.to_k.\2",
+        r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$":
+        r"blocks.\1.to_v.\2",
+        r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
+        r"blocks.\1.to_out.\2",
+        r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
+        r"blocks.\1.norm_q.\2",
+        r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
+        r"blocks.\1.norm_k.\2",
+        r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
+        r"blocks.\1.attn2.to_out.\2",
+        r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$":
+        r"blocks.\1.ffn.fc_in.\2",
+        r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
+        r"blocks.\1.ffn.fc_out.\2",
+        r"blocks\.(\d+)\.norm2\.(.*)$":
+        r"blocks.\1.self_attn_residual_norm.norm.\2",
+    }
+
+    def __init__(
+        self,
+        patch_size: Tuple[int, int, int] = (1, 2, 2),
+        text_len=512,
+        num_attention_heads: int = 40,
+        attention_head_dim: int = 128,
+        in_channels: int = 16,
+        out_channels: int = 16,
+        text_dim: int = 4096,
+        freq_dim: int = 256,
+        ffn_dim: int = 13824,
+        num_layers: int = 40,
+        cross_attn_norm: bool = True,
+        qk_norm: str = "rms_norm_across_heads",
+        eps: float = 1e-6,
+        image_dim: Optional[int] = None,
+        added_kv_proj_dim: Optional[int] = None,
+        rope_max_seq_len: int = 1024,
+    ) -> None:
+        super().__init__()
+
+        inner_dim = num_attention_heads * attention_head_dim
+        self.hidden_size = inner_dim
+        self.num_attention_heads = num_attention_heads
+        self.in_channels = in_channels
+        self.out_channels = out_channels or in_channels
+        self.patch_size = patch_size
+        self.text_len = text_len
+
+        # 1. Patch & position embedding
+        self.patch_embedding = PatchEmbed(in_chans=in_channels,
+                                          embed_dim=inner_dim,
+                                          patch_size=patch_size,
+                                          flatten=False)
+
+        # 2. Condition embeddings
+        self.condition_embedder = WanTimeTextImageEmbedding(
+            dim=inner_dim,
+            time_freq_dim=freq_dim,
+            text_embed_dim=text_dim,
+            image_embed_dim=image_dim,
+        )
+
+        # 3. Transformer blocks
+        self.blocks = nn.ModuleList([
+            WanTransformerBlock(inner_dim, ffn_dim, num_attention_heads,
+                                qk_norm, cross_attn_norm, eps,
+                                added_kv_proj_dim) for _ in range(num_layers)
+        ])
+
+        # 4. Output norm & projection
+        self.norm_out = LayerNormScaleShift(inner_dim,
+                                            norm_type="layer",
+                                            eps=eps,
+                                            elementwise_affine=False,
+                                            dtype=torch.float32)
+        self.proj_out = nn.Linear(inner_dim,
+                                  out_channels * math.prod(patch_size))
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(1, 2, inner_dim) / inner_dim**0.5)
+
+        self.gradient_checkpointing = False
+
+        self.__post_init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+        timestep: torch.LongTensor,
+        seq_len: Optional[int] = None,
+        encoder_hidden_states_image: Optional[Union[torch.Tensor,
+                                                    List[torch.Tensor]]] = None,
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance=None,
+    ) -> torch.Tensor:
+        orig_dtype = hidden_states.dtype
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+        if isinstance(encoder_hidden_states_image,
+                      list) and len(encoder_hidden_states_image) > 0:
+            encoder_hidden_states_image = encoder_hidden_states_image[0]
+        else:
+            encoder_hidden_states_image = None
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # Get rotary embeddings
+        d = self.hidden_size // self.num_attention_heads
+        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
+        freqs_cos, freqs_sin = get_rotary_pos_embed(
+            (post_patch_num_frames * get_sequence_model_parallel_world_size(),
+             post_patch_height, post_patch_width),
+            self.hidden_size,
+            self.num_attention_heads,
+            rope_dim_list,
+            dtype=torch.float64,
+            rope_theta=10000)
+        freqs_cos = freqs_cos.to(hidden_states.device)
+        freqs_sin = freqs_sin.to(hidden_states.device)
+        freqs_cis = (freqs_cos.float(),
+                     freqs_sin.float()) if freqs_cos is not None else None
+
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
+            timestep, encoder_hidden_states, encoder_hidden_states_image)
+        timestep_proj = timestep_proj.unflatten(1, (6, -1))
+
+        if encoder_hidden_states_image is not None:
+            encoder_hidden_states = torch.concat(
+                [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+
+        assert encoder_hidden_states.dtype == orig_dtype
+        # 4. Transformer blocks
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for block in self.blocks:
+                hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, encoder_hidden_states, timestep_proj,
+                    freqs_cis)
+        else:
+            for block in self.blocks:
+                hidden_states = block(hidden_states, encoder_hidden_states,
+                                      timestep_proj, freqs_cis)
+
+        # 5. Output norm, projection & unpatchify
+        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
+                                                                          dim=1)
+        hidden_states = self.norm_out(hidden_states.float(), shift, scale)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
+                                              post_patch_height,
+                                              post_patch_width, p_t, p_h, p_w,
+                                              -1)
+        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return output

fastvideo/v1/models/encoders/clip.py

@@ -0,0 +1,688 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/clip.py
+# Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py
+"""Minimal implementation of CLIPVisionModel intended to be only used
+within a vision language model."""
+from typing import Iterable, Optional, Set, Tuple, Union, cast
+
+import torch
+import torch.nn as nn
+from transformers import CLIPTextConfig, CLIPVisionConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from vllm.model_executor.models.interfaces import SupportsQuant
+
+# from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask
+from fastvideo.v1.attention import LocalAttention
+from fastvideo.v1.distributed import (divide,
+                                      get_tensor_model_parallel_world_size)
+from fastvideo.v1.layers.activation import get_act_fn
+from fastvideo.v1.layers.linear import (ColumnParallelLinear, QKVParallelLinear,
+                                        RowParallelLinear)
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.encoders.vision import (VisionEncoderInfo,
+                                                 resolve_visual_encoder_outputs)
+# TODO: support quantization
+# from vllm.model_executor.layers.quantization import QuantizationConfig
+from fastvideo.v1.models.loader.weight_utils import default_weight_loader
+
+logger = init_logger(__name__)
+
+
+class QuantizationConfig:
+    pass
+
+
+class CLIPEncoderInfo(VisionEncoderInfo[CLIPVisionConfig]):
+
+    def get_num_image_tokens(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+    ) -> int:
+        return self.get_patch_grid_length()**2 + 1
+
+    def get_max_image_tokens(self) -> int:
+        return self.get_patch_grid_length()**2 + 1
+
+    def get_image_size(self) -> int:
+        return cast(int, self.vision_config.image_size)
+
+    def get_patch_size(self) -> int:
+        return cast(int, self.vision_config.patch_size)
+
+    def get_patch_grid_length(self) -> int:
+        image_size, patch_size = self.get_image_size(), self.get_patch_size()
+        assert image_size % patch_size == 0
+        return image_size // patch_size
+
+
+# Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py#L164 # noqa
+class CLIPVisionEmbeddings(nn.Module):
+
+    def __init__(self, config: CLIPVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+        assert self.image_size % self.patch_size == 0
+
+        self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=False,
+        )
+
+        self.num_patches = (self.image_size // self.patch_size)**2
+        self.num_positions = self.num_patches + 1
+        self.position_embedding = nn.Embedding(self.num_positions,
+                                               self.embed_dim)
+        self.register_buffer("position_ids",
+                             torch.arange(self.num_positions).expand((1, -1)),
+                             persistent=False)
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        batch_size = pixel_values.shape[0]
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(pixel_values.to(
+            dtype=target_dtype))  # shape = [*, width, grid, grid]
+        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+
+        class_embeds = self.class_embedding.expand(batch_size, 1, -1)
+        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+        embeddings = embeddings + self.position_embedding(self.position_ids)
+
+        return embeddings
+
+
+class CLIPTextEmbeddings(nn.Module):
+
+    def __init__(self, config: CLIPTextConfig):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
+        self.position_embedding = nn.Embedding(config.max_position_embeddings,
+                                               embed_dim)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer(
+            "position_ids",
+            torch.arange(config.max_position_embeddings).expand((1, -1)),
+            persistent=False)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+    ) -> torch.Tensor:
+        if input_ids is not None:
+            seq_length = input_ids.shape[-1]
+        elif inputs_embeds is not None:
+            seq_length = inputs_embeds.shape[-2]
+        else:
+            raise ValueError(
+                "Either input_ids or inputs_embeds must be provided.")
+
+        max_position_embedding = self.position_embedding.weight.shape[0]
+
+        if seq_length > max_position_embedding:
+            raise ValueError(
+                f"Sequence length must be less than max_position_embeddings (got `sequence length`: "
+                f"{seq_length} and max_position_embeddings: {max_position_embedding}"
+            )
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, :seq_length]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.token_embedding(input_ids)
+
+        position_embeddings = self.position_embedding(position_ids)
+        embeddings = inputs_embeds + position_embeddings
+
+        return embeddings
+
+
+class CLIPAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self,
+        config: CLIPVisionConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                "embed_dim must be divisible by num_heads "
+                f"(got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads}).")
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.embed_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.num_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.out_proj = RowParallelLinear(
+            input_size=self.embed_dim,
+            output_size=self.embed_dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+        )
+
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
+
+        self.attn = LocalAttention(self.num_heads_per_partition,
+                                   self.head_dim,
+                                   self.num_heads_per_partition,
+                                   softmax_scale=self.scale,
+                                   causal=True)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads,
+                           self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ):
+        """Input shape: Batch x Time x Channel"""
+
+        qkv_states, _ = self.qkv_proj(hidden_states)
+        query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
+        # use flash_attn_func
+        query_states = query_states.reshape(query_states.shape[0],
+                                            query_states.shape[1],
+                                            self.num_heads_per_partition,
+                                            self.head_dim)
+        key_states = key_states.reshape(key_states.shape[0],
+                                        key_states.shape[1],
+                                        self.num_heads_per_partition,
+                                        self.head_dim)
+        value_states = value_states.reshape(value_states.shape[0],
+                                            value_states.shape[1],
+                                            self.num_heads_per_partition,
+                                            self.head_dim)
+        attn_output = self.attn(query_states, key_states, value_states)
+        attn_output = attn_output.reshape(
+            attn_output.shape[0], attn_output.shape[1],
+            self.num_heads_per_partition * self.head_dim)
+        attn_output, _ = self.out_proj(attn_output)
+
+        return attn_output, None
+
+
+class CLIPMLP(nn.Module):
+
+    def __init__(
+        self,
+        config: CLIPVisionConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.activation_fn = get_act_fn(config.hidden_act)
+        self.fc1 = ColumnParallelLinear(config.hidden_size,
+                                        config.intermediate_size,
+                                        bias=True,
+                                        quant_config=quant_config,
+                                        prefix=f"{prefix}.fc1")
+        self.fc2 = RowParallelLinear(config.intermediate_size,
+                                     config.hidden_size,
+                                     bias=True,
+                                     quant_config=quant_config,
+                                     prefix=f"{prefix}.fc2")
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states, _ = self.fc2(hidden_states)
+
+        return hidden_states
+
+
+class CLIPEncoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: CLIPTextConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.self_attn = CLIPAttention(
+            config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.layer_norm1 = nn.LayerNorm(config.hidden_size,
+                                        eps=config.layer_norm_eps)
+        self.mlp = CLIPMLP(config,
+                           quant_config=quant_config,
+                           prefix=f"{prefix}.mlp")
+        self.layer_norm2 = nn.LayerNorm(config.hidden_size,
+                                        eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, _ = self.self_attn(hidden_states=hidden_states)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class CLIPEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self
+    attention layers. Each layer is a [`CLIPEncoderLayer`].
+
+    Args:
+        config: CLIPConfig
+    """
+
+    def __init__(
+        self,
+        config: CLIPVisionConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        num_hidden_layers_override: Optional[int] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+
+        if num_hidden_layers_override is None:
+            num_hidden_layers = config.num_hidden_layers
+        else:
+            num_hidden_layers = num_hidden_layers_override
+        self.layers = nn.ModuleList([
+            CLIPEncoderLayer(config=config,
+                             quant_config=quant_config,
+                             prefix=f"{prefix}.layers.{layer_idx}")
+            for layer_idx in range(num_hidden_layers)
+        ])
+
+    def forward(
+        self, inputs_embeds: torch.Tensor, return_all_hidden_states: bool
+    ) -> Union[torch.Tensor, list[torch.Tensor]]:
+        hidden_states_pool = [inputs_embeds]
+        hidden_states = inputs_embeds
+
+        for encoder_layer in self.layers:
+            hidden_states = encoder_layer(hidden_states)
+            if return_all_hidden_states:
+                hidden_states_pool.append(hidden_states)
+        # If we have multiple feature sample layers, we return all hidden
+        # states in order and grab the ones we need by index.
+        if return_all_hidden_states:
+            return hidden_states_pool
+        return [hidden_states]
+
+
+class CLIPTextTransformer(nn.Module):
+
+    def __init__(self,
+                 config: CLIPTextConfig,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 *,
+                 num_hidden_layers_override: Optional[int] = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = CLIPTextEmbeddings(config)
+
+        self.encoder = CLIPEncoder(
+            config,
+            quant_config=quant_config,
+            num_hidden_layers_override=num_hidden_layers_override,
+            prefix=prefix)
+
+        self.final_layer_norm = nn.LayerNorm(embed_dim,
+                                             eps=config.layer_norm_eps)
+
+        # For `pooled_output` computation
+        self.eos_token_id = config.eos_token_id
+
+        # For attention mask, it differs between `flash_attention_2` and other attention implementations
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is None:
+            raise ValueError("You have to specify input_ids")
+
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+
+        hidden_states = self.embeddings(input_ids=input_ids,
+                                        position_ids=position_ids)
+
+        # CLIP's text model uses causal mask, prepare it here.
+        # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
+        # causal_attention_mask = _create_4d_causal_attention_mask(
+        #     input_shape, hidden_states.dtype, device=hidden_states.device
+        # )
+
+        # # expand attention_mask
+        # if attention_mask is not None and not self._use_flash_attention_2:
+        #     raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer")
+        #     # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        #     attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            # attention_mask=attention_mask,
+            # causal_attention_mask=causal_attention_mask,
+            # output_attentions=output_attentions,
+            return_all_hidden_states=output_hidden_states,
+            # return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[-1]
+        last_hidden_state = self.final_layer_norm(last_hidden_state)
+
+        if self.eos_token_id == 2:
+            # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here.
+            # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added
+            # ------------------------------------------------------------
+            # text_embeds.shape = [batch_size, sequence_length, transformer.width]
+            # take features from the eot embedding (eot_token is the highest number in each sequence)
+            # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
+            pooled_output = last_hidden_state[
+                torch.arange(last_hidden_state.shape[0],
+                             device=last_hidden_state.device),
+                input_ids.to(dtype=torch.int, device=last_hidden_state.device).
+                argmax(dim=-1),
+            ]
+        else:
+            # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible)
+            pooled_output = last_hidden_state[
+                torch.arange(last_hidden_state.shape[0],
+                             device=last_hidden_state.device),
+                # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
+                # Note: we assume each sequence (along batch dim.) contains an  `eos_token_id` (e.g. prepared by the tokenizer)
+                (input_ids.to(dtype=torch.int, device=last_hidden_state.device
+                              ) == self.eos_token_id).int().argmax(dim=-1),
+            ]
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        # return last_hidden_state
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs,
+            # attentions=encoder_outputs.attentions,
+        )
+
+
+class CLIPTextModel(nn.Module):
+
+    def __init__(
+        self,
+        config: CLIPTextConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.text_model = CLIPTextTransformer(config=config,
+                                              quant_config=quant_config,
+                                              prefix=prefix)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        return self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=None,
+        )
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+
+        # Define mapping for stacked parameters
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        for name, loaded_weight in weights:
+            # Handle q_proj, k_proj, v_proj -> qkv_proj mapping
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name in name:
+                    # Replace the weight name with the parameter name
+                    model_param_name = name.replace(weight_name, param_name)
+
+                    if model_param_name in params_dict:
+                        param = params_dict[model_param_name]
+                        weight_loader = param.weight_loader
+                        weight_loader(param, loaded_weight, shard_id)
+                        loaded_params.add(model_param_name)
+                    break
+            else:
+                # Use default weight loader for all other parameters
+                if name in params_dict:
+                    param = params_dict[name]
+                    weight_loader = getattr(param, "weight_loader",
+                                            default_weight_loader)
+                    weight_loader(param, loaded_weight)
+                    loaded_params.add(name)
+
+        return loaded_params
+
+
+class CLIPVisionTransformer(nn.Module):
+
+    def __init__(
+        self,
+        config: CLIPVisionConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        *,
+        num_hidden_layers_override: Optional[int] = None,
+        require_post_norm: Optional[bool] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = CLIPVisionEmbeddings(config)
+
+        # NOTE: This typo of "layrnorm" is not fixed on purpose to match
+        # the original transformers code and name of the model weights.
+        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+        self.encoder = CLIPEncoder(
+            config=config,
+            quant_config=quant_config,
+            num_hidden_layers_override=num_hidden_layers_override,
+            prefix=f"{prefix}.encoder",
+        )
+
+        num_hidden_layers = config.num_hidden_layers
+        if len(self.encoder.layers) > config.num_hidden_layers:
+            raise ValueError(
+                f"The original encoder only has {num_hidden_layers} "
+                f"layers, but you requested {len(self.encoder.layers)} layers.")
+
+        # If possible, skip post_layernorm to conserve memory
+        if require_post_norm is None:
+            require_post_norm = len(self.encoder.layers) == num_hidden_layers
+
+        if require_post_norm:
+            self.post_layernorm = nn.LayerNorm(embed_dim,
+                                               eps=config.layer_norm_eps)
+        else:
+            self.post_layernorm = None
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        feature_sample_layers: Optional[list[int]] = None,
+    ) -> torch.Tensor:
+
+        hidden_states = self.embeddings(pixel_values)
+        hidden_states = self.pre_layrnorm(hidden_states)
+
+        return_all_hidden_states = feature_sample_layers is not None
+
+        # Produces either the last layer output or all of the hidden states,
+        # depending on if we have feature_sample_layers or not
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            return_all_hidden_states=return_all_hidden_states)
+
+        if not return_all_hidden_states:
+            encoder_outputs = encoder_outputs[0]
+
+        # Handle post-norm (if applicable) and stacks feature layers if needed
+        encoder_outputs = resolve_visual_encoder_outputs(
+            encoder_outputs, feature_sample_layers, self.post_layernorm,
+            self.config.num_hidden_layers)
+
+        return encoder_outputs
+
+
+class CLIPVisionModel(nn.Module, SupportsQuant):
+    config_class = CLIPVisionConfig
+    main_input_name = "pixel_values"
+    packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
+
+    def __init__(
+        self,
+        config: CLIPVisionConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        *,
+        num_hidden_layers_override: Optional[int] = None,
+        require_post_norm: Optional[bool] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.vision_model = CLIPVisionTransformer(
+            config=config,
+            quant_config=quant_config,
+            num_hidden_layers_override=num_hidden_layers_override,
+            require_post_norm=require_post_norm,
+            prefix=f"{prefix}.vision_model")
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        feature_sample_layers: Optional[list[int]] = None,
+    ) -> torch.Tensor:
+        return self.vision_model(pixel_values, feature_sample_layers)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    # (TODO) Add prefix argument for filtering out weights to be loaded
+    #        ref: https://github.com/vllm-project/vllm/pull/7186#discussion_r1734163986
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        layer_count = len(self.vision_model.encoder.layers)
+
+        for name, loaded_weight in weights:
+            if name.startswith("visual_projection"):
+                continue
+            # post_layernorm is not needed in CLIPVisionModel
+            if (name.startswith("vision_model.post_layernorm")
+                    and self.vision_model.post_layernorm is None):
+                continue
+
+            # omit layers when num_hidden_layers_override is set
+            if name.startswith("vision_model.encoder.layers"):
+                layer_idx = int(name.split(".")[3])
+                if layer_idx >= layer_count:
+                    continue
+
+            for (param_name, weight_name, shard_id) in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

fastvideo/v1/models/encoders/llama.py

@@ -0,0 +1,440 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/llama.py
+
+# Adapted from
+# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inference-only LLaMA model compatible with HuggingFace weights."""
+from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type
+
+import torch
+from torch import nn
+from transformers import LlamaConfig
+from transformers.modeling_outputs import BaseModelOutputWithPast
+
+# from vllm.model_executor.layers.quantization import QuantizationConfig
+from fastvideo.v1.attention import LocalAttention
+from fastvideo.v1.distributed import get_tensor_model_parallel_world_size
+from fastvideo.v1.layers.activation import SiluAndMul
+from fastvideo.v1.layers.layernorm import RMSNorm
+from fastvideo.v1.layers.linear import (MergedColumnParallelLinear,
+                                        QKVParallelLinear, RowParallelLinear)
+from fastvideo.v1.layers.rotary_embedding import get_rope
+from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from fastvideo.v1.models.loader.weight_utils import (default_weight_loader,
+                                                     maybe_remap_kv_scale_name)
+
+# from ..utils import (extract_layer_index)
+
+
+class QuantizationConfig:
+    pass
+
+
+class LlamaMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: Optional[QuantizationConfig] = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            # output_size=intermediate_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {hidden_act}. "
+                             "Only silu is supported for now.")
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class LlamaAttention(nn.Module):
+
+    def __init__(self,
+                 config: LlamaConfig,
+                 hidden_size: int,
+                 num_heads: int,
+                 num_kv_heads: int,
+                 rope_theta: float = 10000,
+                 rope_scaling: Optional[Dict[str, Any]] = None,
+                 max_position_embeddings: int = 8192,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 bias: bool = False,
+                 bias_o_proj: bool = False,
+                 prefix: str = "") -> None:
+        super().__init__()
+        # layer_idx = extract_layer_index(prefix)
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
+        self.head_dim = getattr(config, "head_dim",
+                                self.hidden_size // self.total_num_heads)
+        # Phi models introduced a partial_rotary_factor parameter in the config
+        partial_rotary_factor = getattr(config, "partial_rotary_factor", 1)
+        self.rotary_dim = int(partial_rotary_factor * self.head_dim)
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias_o_proj,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        is_neox_style = True
+        is_gguf = quant_config and hasattr(
+            quant_config, "get_name") and quant_config.get_name() == "gguf"
+        if is_gguf and config.model_type == "llama":
+            is_neox_style = False
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.rotary_dim,
+            max_position=max_position_embeddings,
+            base=rope_theta,
+            rope_scaling=rope_scaling,
+            is_neox_style=is_neox_style,
+        )
+
+        self.attn = LocalAttention(self.num_heads,
+                                   self.head_dim,
+                                   self.num_kv_heads,
+                                   softmax_scale=self.scaling,
+                                   causal=True)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        # attn_output = self.attn(q, k, v)
+        # use flash_attn_func
+        # TODO (Attn abstraction and backend)
+        # from flash_attn import flash_attn_func
+        # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim)
+        batch_size = q.shape[0]
+        seq_len = q.shape[1]
+        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        # import pdb; pdb.set_trace()
+        # attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True)
+        attn_output = self.attn(q, k, v)
+        attn_output = attn_output.reshape(batch_size, seq_len,
+                                          self.num_heads * self.head_dim)
+
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class LlamaDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: LlamaConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 10000)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        if rope_scaling is not None and getattr(
+                config, "original_max_position_embeddings", None):
+            rope_scaling["original_max_position_embeddings"] = (
+                config.original_max_position_embeddings)
+        max_position_embeddings = getattr(config, "max_position_embeddings",
+                                          8192)
+        # Support abacusai/Smaug-72B-v0.1 with attention_bias
+        # Support internlm/internlm-7b with bias
+        attention_bias = getattr(config, "attention_bias", False) or getattr(
+            config, "bias", False)
+        bias_o_proj = attention_bias
+        # support internlm/internlm3-8b with qkv_bias
+        if hasattr(config, 'qkv_bias'):
+            attention_bias = config.qkv_bias
+
+        self.self_attn = LlamaAttention(
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=getattr(config, "num_key_value_heads",
+                                 config.num_attention_heads),
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            max_position_embeddings=max_position_embeddings,
+            quant_config=quant_config,
+            bias=attention_bias,
+            bias_o_proj=bias_o_proj,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = LlamaMLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            bias=getattr(config, "mlp_bias", False),
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size,
+                                       eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size,
+                                                eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(
+                hidden_states, residual)
+
+        hidden_states = self.self_attn(positions=positions,
+                                       hidden_states=hidden_states)
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(
+            hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class LlamaModel(nn.Module):
+
+    def __init__(self,
+                 config: LlamaConfig,
+                 prefix: str = "",
+                 layer_type: Type[LlamaDecoderLayer] = LlamaDecoderLayer):
+        super().__init__()
+
+        quant_config = None
+        lora_config = None
+
+        self.config = config
+        self.quant_config = quant_config
+        if lora_config is not None:
+            max_loras = 1
+            lora_vocab_size = 1
+            if hasattr(lora_config, "max_loras"):
+                max_loras = lora_config.max_loras
+            if hasattr(lora_config, "lora_extra_vocab_size"):
+                lora_vocab_size = lora_config.lora_extra_vocab_size
+            lora_vocab = lora_vocab_size * max_loras
+        else:
+            lora_vocab = 0
+        self.vocab_size = config.vocab_size + lora_vocab
+        self.org_vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+            org_num_embeddings=config.vocab_size,
+            quant_config=quant_config,
+        )
+
+        self.layers = nn.ModuleList([
+            layer_type(config=config,
+                       quant_config=quant_config,
+                       prefix=f"{prefix}.layers.{i}")
+            for i in range(config.num_hidden_layers)
+        ])
+
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor],
+        positions: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+    ) -> torch.Tensor:
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = self.get_input_embeddings(input_ids)
+        residual = None
+
+        if positions is None:
+            positions = torch.arange(0,
+                                     hidden_states.shape[1],
+                                     device=hidden_states.device).unsqueeze(0)
+
+        all_hidden_states: Optional[Tuple[Any, ...]] = (
+        ) if output_hidden_states else None
+        for layer in self.layers:
+            if all_hidden_states is not None:
+                # TODO
+                all_hidden_states += (
+                    hidden_states, ) if residual is None else (hidden_states +
+                                                               residual, )
+            hidden_states, residual = layer(positions, hidden_states, residual)
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+
+        # add hidden states from the last decoder layer
+        if all_hidden_states is not None:
+            all_hidden_states += (hidden_states, )
+
+        # TODO(will): maybe unify the output format with other models and use
+        # our own class
+        output = BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            # past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            # attentions=all_self_attns,
+        )
+
+        return output
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if ("rotary_emb.cos_cached" in name
+                    or "rotary_emb.sin_cached" in name):
+                # Models trained using ColossalAI may include these tensors in
+                # the checkpoint. Skip them.
+                continue
+            if (self.quant_config is not None and
+                (scale_name := self.quant_config.get_cache_scale(name))):
+                # Loading kv cache quantization scales
+                param = params_dict[scale_name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
+                                 loaded_weight[0])
+                weight_loader(param, loaded_weight)
+                loaded_params.add(scale_name)
+                continue
+            if "scale" in name:
+                # Remapping the name of FP8 kv-scale.
+                kv_scale_name: Optional[str] = maybe_remap_kv_scale_name(
+                    name, params_dict)
+                if kv_scale_name is None:
+                    continue
+                else:
+                    name = kv_scale_name
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

fastvideo/v1/models/encoders/t5.py

@@ -0,0 +1,677 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from transformers: https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/t5/modeling_t5.py
+
+# Derived from T5 implementation posted on HuggingFace; license below:
+#
+# coding=utf-8
+# Copyright 2018 Mesh TensorFlow authors, T5 Authors and HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch T5 & UMT5 model."""
+
+import math
+from dataclasses import dataclass
+from typing import Iterable, Optional, Set, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import T5Config
+
+from fastvideo.v1.distributed import (get_tensor_model_parallel_rank,
+                                      get_tensor_model_parallel_world_size)
+from fastvideo.v1.layers.activation import get_act_fn
+from fastvideo.v1.layers.layernorm import RMSNorm
+from fastvideo.v1.layers.linear import (MergedColumnParallelLinear,
+                                        QKVParallelLinear, RowParallelLinear)
+from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from fastvideo.v1.models.loader.weight_utils import default_weight_loader
+
+
+class QuantizationConfig:
+    pass
+
+
+class AttentionType:
+    """
+    Attention type.
+    Use string to be compatible with `torch.compile`.
+    """
+    # Decoder attention between previous layer Q/K/V
+    DECODER = "decoder"
+    # Encoder attention between previous layer Q/K/V for encoder-decoder
+    ENCODER = "encoder"
+    # Encoder attention between previous layer Q/K/V
+    ENCODER_ONLY = "encoder_only"
+    # Attention between dec. Q and enc. K/V for encoder-decoder
+    ENCODER_DECODER = "encoder_decoder"
+
+
+@dataclass
+class AttentionMetadata:
+    attn_bias: torch.Tensor
+
+
+class T5DenseActDense(nn.Module):
+
+    def __init__(self,
+                 config: T5Config,
+                 quant_config: Optional[QuantizationConfig] = None):
+        super().__init__()
+        self.wi = MergedColumnParallelLinear(config.d_model, [config.d_ff],
+                                             bias=False)
+        self.wo = RowParallelLinear(config.d_ff,
+                                    config.d_model,
+                                    bias=False,
+                                    quant_config=quant_config)
+        self.act = get_act_fn(config.dense_act_fn)
+
+    def forward(self, hidden_states) -> torch.Tensor:
+        hidden_states, _ = self.wi(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states, _ = self.wo(hidden_states)
+        return hidden_states
+
+
+class T5DenseGatedActDense(nn.Module):
+
+    def __init__(self,
+                 config: T5Config,
+                 quant_config: Optional[QuantizationConfig] = None):
+        super().__init__()
+        self.wi_0 = MergedColumnParallelLinear(config.d_model, [config.d_ff],
+                                               bias=False,
+                                               quant_config=quant_config)
+        self.wi_1 = MergedColumnParallelLinear(config.d_model, [config.d_ff],
+                                               bias=False,
+                                               quant_config=quant_config)
+        # Should not run in fp16 unless mixed-precision is used,
+        # see https://github.com/huggingface/transformers/issues/20287.
+        self.wo = RowParallelLinear(config.d_ff,
+                                    config.d_model,
+                                    bias=False,
+                                    quant_config=quant_config)
+        self.act = get_act_fn(config.dense_act_fn)
+
+    def forward(self, hidden_states) -> torch.Tensor:
+        hidden_gelu = self.act(self.wi_0(hidden_states)[0])
+        hidden_linear, _ = self.wi_1(hidden_states)
+        hidden_states = hidden_gelu * hidden_linear
+        hidden_states, _ = self.wo(hidden_states)
+        return hidden_states
+
+
+class T5LayerFF(nn.Module):
+
+    def __init__(self,
+                 config: T5Config,
+                 quant_config: Optional[QuantizationConfig] = None):
+        super().__init__()
+        if config.is_gated_act:
+            self.DenseReluDense = T5DenseGatedActDense(
+                config, quant_config=quant_config)
+        else:
+            self.DenseReluDense = T5DenseActDense(config,
+                                                  quant_config=quant_config)
+
+        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(self, hidden_states) -> torch.Tensor:
+        forwarded_states = self.layer_norm.forward_native(hidden_states)
+        forwarded_states = self.DenseReluDense(forwarded_states)
+        hidden_states = hidden_states + forwarded_states
+        return hidden_states
+
+
+# T5 has attn_bias and does not use softmax scaling
+class T5MultiHeadAttention(nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, q, k, v, attn_bias=None):
+        b, _, n, c = q.shape
+        attn = torch.einsum('binc,bjnc->bnij', q, k)
+        if attn_bias is not None:
+            attn += attn_bias
+
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum('bnij,bjnc->binc', attn, v)
+        x = x.reshape(b, -1, n * c)
+        return x
+
+
+class T5Attention(nn.Module):
+
+    def __init__(self,
+                 config: T5Config,
+                 attn_type: str,
+                 has_relative_attention_bias=False,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.attn_type = attn_type
+        # Cross-attention has no relative pos encoding anyway
+        self.is_decoder = attn_type == AttentionType.DECODER
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.relative_attention_num_buckets = \
+            config.relative_attention_num_buckets
+        self.relative_attention_max_distance = \
+            config.relative_attention_max_distance
+        self.d_model = config.d_model
+        self.key_value_proj_dim = config.d_kv
+        self.total_num_heads = self.total_num_kv_heads = config.num_heads
+
+        # Partition heads across multiple tensor parallel GPUs.
+        tp_world_size = get_tensor_model_parallel_world_size()
+        assert config.num_heads % tp_world_size == 0
+        self.n_heads = config.num_heads // tp_world_size
+
+        self.inner_dim = self.n_heads * self.key_value_proj_dim
+        # No GQA in t5.
+        # self.n_kv_heads = self.n_heads
+
+        self.qkv_proj = QKVParallelLinear(
+            self.d_model,
+            self.d_model // self.total_num_heads,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.attn = T5MultiHeadAttention()
+
+        if self.has_relative_attention_bias:
+            self.relative_attention_bias = \
+                VocabParallelEmbedding(self.relative_attention_num_buckets,
+                                       self.total_num_heads,
+                                       org_num_embeddings=self.relative_attention_num_buckets,
+                                       padding_size=self.relative_attention_num_buckets,
+                                       quant_config=quant_config)
+        self.o = RowParallelLinear(
+            self.d_model,
+            self.d_model,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+    @staticmethod
+    def _relative_position_bucket(relative_position,
+                                  bidirectional=True,
+                                  num_buckets=32,
+                                  max_distance=128) -> torch.Tensor:
+        """
+        Adapted from Mesh Tensorflow:
+        https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593
+        Translate relative position to a bucket number for relative attention. 
+        The relative position is defined as memory_position - query_position, 
+        i.e. the distance in tokens from the attending position to the 
+        attended-to position. If bidirectional=False, then positive relative 
+        positions are invalid. We use smaller buckets for small absolute 
+        relative_position and larger buckets for larger absolute 
+        relative_positions. All relative positions >=max_distance map to the
+        same bucket. All relative positions <=-max_distance map to the same 
+        bucket. This should allow for more graceful generalization to longer
+        sequences than the model has been trained on
+        Args:
+            relative_position: an int32 Tensor
+            bidirectional: a boolean - whether the attention is bidirectional
+            num_buckets: an integer
+            max_distance: an integer
+        Returns:
+            a Tensor with the same shape as relative_position, containing int32
+            values in the range [0, num_buckets)
+        """# noqa: E501
+        relative_buckets = 0
+        if bidirectional:
+            num_buckets //= 2
+            relative_buckets += (relative_position > 0).to(
+                torch.long) * num_buckets
+            relative_position = torch.abs(relative_position)
+        else:
+            relative_position = -torch.min(relative_position,
+                                           torch.zeros_like(relative_position))
+        # now relative_position is in the range [0, inf)
+
+        # half of the buckets are for exact increments in positions
+        max_exact = num_buckets // 2
+        is_small = relative_position < max_exact
+
+        # The other half of the buckets are for logarithmically bigger bins
+        # in positions up to max_distance
+        relative_position_if_large = max_exact + (
+            torch.log(relative_position.float() / max_exact) /
+            math.log(max_distance / max_exact) *
+            (num_buckets - max_exact)).to(torch.long)
+        relative_position_if_large = torch.min(
+            relative_position_if_large,
+            torch.full_like(relative_position_if_large, num_buckets - 1))
+
+        relative_buckets += torch.where(is_small, relative_position,
+                                        relative_position_if_large)
+        return relative_buckets
+
+    def compute_bias(self,
+                     query_length,
+                     key_length,
+                     device=None) -> torch.Tensor:
+        """Compute binned relative position bias"""
+        if device is None:
+            device = self.relative_attention_bias.weight.device
+        context_position = torch.arange(query_length,
+                                        dtype=torch.long,
+                                        device=device)[:, None]
+        memory_position = torch.arange(key_length,
+                                       dtype=torch.long,
+                                       device=device)[None, :]
+        # max_seq_len, nh
+        relative_position = memory_position - context_position
+        relative_position_bucket = self._relative_position_bucket(
+            relative_position,  # shape (query_length, key_length)
+            bidirectional=(not self.is_decoder),
+            num_buckets=self.relative_attention_num_buckets,
+            max_distance=self.relative_attention_max_distance,
+        )
+        values = self.relative_attention_bias(
+            relative_position_bucket
+        )  # shape (query_length, key_length, num_heads)
+        x = values.permute([2, 0, 1]).unsqueeze(
+            0)  # shape (1, num_heads, query_length, key_length)
+        return x
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,  # (num_tokens, d_model)
+        attention_mask: torch.Tensor,
+        attn_metadata: Optional[AttentionMetadata] = None,
+    ) -> torch.Tensor:
+        bs, seq_len, _ = hidden_states.shape
+        num_seqs = bs
+        n, c = self.n_heads, self.d_model // self.total_num_heads
+        qkv, _ = self.qkv_proj(hidden_states)
+        # Projection of 'own' hidden state (self-attention). No GQA here.
+        q, k, v = qkv.split(self.inner_dim, dim=-1)
+        q = q.reshape(bs, seq_len, n, c)
+        k = k.reshape(bs, seq_len, n, c)
+        v = v.reshape(bs, seq_len, n, c)
+
+        assert attn_metadata is not None
+        attn_bias = attn_metadata.attn_bias
+        # Not compatible with CP here (as all encoder-decoder models),
+        # as it assumes homogeneous batch (prefills or decodes).
+        if self.has_relative_attention_bias:
+            # Self-attention. Compute T5 relative positional encoding.
+            # The bias term is computed on longest sequence in batch. Biases
+            # for shorter sequences are slices of the longest.
+            assert self.attn_type == AttentionType.ENCODER
+            attn_bias = self.compute_bias(seq_len,
+                                          seq_len).repeat(num_seqs, 1, 1, 1)
+            attn_metadata.attn_bias = attn_bias
+        else:
+            # Encoder/Decoder Self-Attention Layer, attn bias already cached.
+            assert attn_bias is not None
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.view(
+                bs, 1, 1,
+                -1) if attention_mask.ndim == 2 else attention_mask.unsqueeze(1)
+            attn_bias.masked_fill_(attention_mask == 0,
+                                   torch.finfo(q.dtype).min)
+
+        if get_tensor_model_parallel_world_size() > 1:
+            rank = get_tensor_model_parallel_rank()
+            attn_bias = attn_bias[:, rank * self.n_heads:(rank + 1) *
+                                  self.n_heads, :, :]
+        attn_output = self.attn(q, k, v, attn_bias)
+        output, _ = self.o(attn_output)
+        return output
+
+
+class T5LayerSelfAttention(nn.Module):
+
+    def __init__(
+        self,
+        config,
+        has_relative_attention_bias=False,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.SelfAttention = T5Attention(
+            config,
+            AttentionType.DECODER
+            if "decoder" in prefix else AttentionType.ENCODER,
+            has_relative_attention_bias=has_relative_attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.SelfAttention")
+        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        attn_metadata: Optional[AttentionMetadata] = None,
+    ) -> torch.Tensor:
+        normed_hidden_states = self.layer_norm.forward_native(hidden_states)
+        attention_output = self.SelfAttention(
+            hidden_states=normed_hidden_states,
+            attention_mask=attention_mask,
+            attn_metadata=attn_metadata,
+        )
+        hidden_states = hidden_states + attention_output
+        return hidden_states
+
+
+class T5LayerCrossAttention(nn.Module):
+
+    def __init__(self,
+                 config,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.EncDecAttention = T5Attention(config,
+                                           AttentionType.ENCODER_DECODER,
+                                           has_relative_attention_bias=False,
+                                           quant_config=quant_config,
+                                           prefix=f"{prefix}.EncDecAttention")
+        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attn_metadata: Optional[AttentionMetadata] = None,
+    ) -> torch.Tensor:
+        normed_hidden_states = self.layer_norm.forward_native(hidden_states)
+        attention_output = self.EncDecAttention(
+            hidden_states=normed_hidden_states,
+            attn_metadata=attn_metadata,
+        )
+        hidden_states = hidden_states + attention_output
+        return hidden_states
+
+
+class T5Block(nn.Module):
+
+    def __init__(self,
+                 config: T5Config,
+                 is_decoder: bool,
+                 has_relative_attention_bias=False,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.is_decoder = is_decoder
+        self.layer = nn.ModuleList()
+        self.layer.append(
+            T5LayerSelfAttention(
+                config,
+                has_relative_attention_bias=has_relative_attention_bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn"))
+
+        if self.is_decoder:
+            self.layer.append(
+                T5LayerCrossAttention(config,
+                                      quant_config=quant_config,
+                                      prefix=f"{prefix}.cross_attn"))
+
+        self.layer.append(T5LayerFF(config, quant_config=quant_config))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        attn_metadata: Optional[AttentionMetadata] = None,
+    ) -> torch.Tensor:
+
+        hidden_states = self.layer[0](hidden_states=hidden_states,
+                                      attention_mask=attention_mask,
+                                      attn_metadata=attn_metadata)
+        if self.is_decoder:
+            hidden_states = self.layer[1](hidden_states=hidden_states,
+                                          attn_metadata=attn_metadata)
+
+            # Apply Feed Forward layer
+            hidden_states = self.layer[2](hidden_states)
+        else:
+            hidden_states = self.layer[1](hidden_states)
+        return hidden_states
+
+
+class T5Stack(nn.Module):
+
+    def __init__(self,
+                 config: T5Config,
+                 is_decoder: bool,
+                 n_layers: int,
+                 embed_tokens=None,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = "",
+                 is_umt5: bool = False):
+        super().__init__()
+        self.embed_tokens = embed_tokens
+        self.is_umt5 = is_umt5
+        if is_umt5:
+            self.block = nn.ModuleList([
+                T5Block(config,
+                        is_decoder=is_decoder,
+                        has_relative_attention_bias=True,
+                        quant_config=quant_config,
+                        prefix=f"{prefix}.blocks.{i}") for i in range(n_layers)
+            ])
+        else:
+            # Only the first block has relative positional encoding.
+            self.block = nn.ModuleList([
+                T5Block(config,
+                        is_decoder=is_decoder,
+                        has_relative_attention_bias=i == 0,
+                        quant_config=quant_config,
+                        prefix=f"{prefix}.blocks.{i}") for i in range(n_layers)
+            ])
+        self.final_layer_norm = RMSNorm(config.d_model,
+                                        eps=config.layer_norm_epsilon)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        hidden_states = self.embed_tokens(input_ids)
+
+        for idx, block in enumerate(self.block):
+            hidden_states = block(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                attn_metadata=attn_metadata,
+            )
+        hidden_states = self.final_layer_norm.forward_native(hidden_states)
+        return hidden_states
+
+
+class T5EncoderModel(nn.Module):
+
+    def __init__(self, config: T5Config, prefix: str = ""):
+        super().__init__()
+
+        quant_config = None
+
+        self.shared = VocabParallelEmbedding(
+            config.vocab_size,
+            config.d_model,
+            org_num_embeddings=config.vocab_size)
+
+        self.encoder = T5Stack(config,
+                               False,
+                               config.num_layers,
+                               self.shared,
+                               quant_config=quant_config,
+                               prefix=f"{prefix}.encoder",
+                               is_umt5=False)
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> torch.Tensor:
+        attn_metadata = AttentionMetadata(None)
+        encoder_outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            attn_metadata=attn_metadata,
+        )
+
+        return encoder_outputs
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q", "q"),
+            (".qkv_proj", ".k", "k"),
+            (".qkv_proj", ".v", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        for name, loaded_weight in weights:
+            loaded = False
+            if "decoder" in name or "lm_head" in name:
+                continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                loaded = True
+                break
+            if not loaded:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class UMT5EncoderModel(nn.Module):
+
+    def __init__(self, config: T5Config, prefix: str = ""):
+        super().__init__()
+
+        quant_config = None
+
+        self.shared = VocabParallelEmbedding(
+            config.vocab_size,
+            config.d_model,
+            org_num_embeddings=config.vocab_size)
+
+        self.encoder = T5Stack(config,
+                               False,
+                               config.num_layers,
+                               self.shared,
+                               quant_config=quant_config,
+                               prefix=f"{prefix}.encoder",
+                               is_umt5=True)
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> torch.Tensor:
+        attn_metadata = AttentionMetadata(None)
+        encoder_outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            attn_metadata=attn_metadata,
+        )
+
+        return encoder_outputs
+
+    def load_weights(self, weights: Iterable[Tuple[str,
+                                                   torch.Tensor]]) -> Set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q", "q"),
+            (".qkv_proj", ".k", "k"),
+            (".qkv_proj", ".v", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        for name, loaded_weight in weights:
+            loaded = False
+            if "decoder" in name or "lm_head" in name:
+                continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                loaded = True
+                break
+            if not loaded:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader",
+                                        default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

fastvideo/v1/models/encoders/vision.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/vision.py
+
+from abc import ABC, abstractmethod
+from typing import Generic, Optional, TypeVar, Union
+
+import torch
+from transformers import PretrainedConfig
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+_C = TypeVar("_C", bound=PretrainedConfig)
+
+
+class VisionEncoderInfo(ABC, Generic[_C]):
+
+    def __init__(self, vision_config: _C) -> None:
+        super().__init__()
+
+        self.vision_config = vision_config
+
+    @abstractmethod
+    def get_num_image_tokens(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+    ) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_max_image_tokens(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_image_size(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_patch_size(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_patch_grid_length(self) -> int:
+        raise NotImplementedError
+
+
+def resolve_visual_encoder_outputs(
+    encoder_outputs: Union[torch.Tensor, list[torch.Tensor]],
+    feature_sample_layers: Optional[list[int]],
+    post_layer_norm: Optional[torch.nn.LayerNorm],
+    max_possible_layers: int,
+) -> torch.Tensor:
+    """Given the outputs a visual encoder module that may correspond to the
+    output of the last layer, or a list of hidden states to be stacked,
+    handle post normalization and resolve it into a single output tensor.
+
+    Args:
+        encoder_outputs: Output of encoder's last layer or all hidden states.
+        feature_sample_layers: Optional layer indices to grab from the encoder
+            outputs; if provided, encoder outputs must be a list.
+        post_layer_norm: Post norm to apply to the output of the encoder.
+        max_possible_layers: Total layers in the fully loaded visual encoder.
+
+    """
+    if feature_sample_layers is None:
+        if post_layer_norm is not None:
+            return post_layer_norm(encoder_outputs)
+        return encoder_outputs
+
+    # Get the hidden states corresponding to the layer indices.
+    # Negative values are relative to the full visual encoder,
+    # so offset them depending on how many layers were loaded.
+    # NOTE: this assumes that encoder_outputs is a list containing
+    # the inputs to the visual encoder, followed by the hidden states
+    # of each layer.
+    num_loaded_layers = len(encoder_outputs) - 1
+    offset = max_possible_layers - num_loaded_layers
+    hs_pool = [
+        encoder_outputs[layer_idx]
+        if layer_idx >= 0 else encoder_outputs[layer_idx + offset]
+        for layer_idx in feature_sample_layers
+    ]
+
+    # Apply post-norm on the final hidden state if we are using it
+    uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1)
+    if post_layer_norm is not None and uses_last_layer:
+        hs_pool[-1] = post_layer_norm(encoder_outputs)
+    return torch.cat(hs_pool, dim=-1)

fastvideo/v1/models/hf_transformer_utils.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/hf_transformers_utils.py
+
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for Huggingface Transformers."""
+
+import contextlib
+import json
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional, Type, Union
+
+from huggingface_hub import snapshot_download
+from transformers import AutoConfig, PretrainedConfig
+from transformers.models.auto.modeling_auto import (
+    MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
+
+_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
+    # ChatGLMConfig.model_type: ChatGLMConfig,
+    # DbrxConfig.model_type: DbrxConfig,
+    # ExaoneConfig.model_type: ExaoneConfig,
+    # Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig,
+}
+
+for name, cls in _CONFIG_REGISTRY.items():
+    with contextlib.suppress(ValueError):
+        AutoConfig.register(name, cls)
+
+
+def download_from_hf(model_path: str):
+    if os.path.exists(model_path):
+        return model_path
+
+    return snapshot_download(model_path,
+                             allow_patterns=["*.json", "*.bin", "*.model"])
+
+
+def get_hf_config(
+    model: str,
+    trust_remote_code: bool,
+    revision: Optional[str] = None,
+    model_override_args: Optional[dict] = None,
+    inference_args: Optional[dict] = None,
+    **kwargs,
+):
+    is_gguf = check_gguf_file(model)
+    if is_gguf:
+        raise NotImplementedError("GGUF models are not supported.")
+
+    config = AutoConfig.from_pretrained(model,
+                                        trust_remote_code=trust_remote_code,
+                                        revision=revision,
+                                        **kwargs)
+    if config.model_type in _CONFIG_REGISTRY:
+        config_class = _CONFIG_REGISTRY[config.model_type]
+        config = config_class.from_pretrained(model, revision=revision)
+        # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`.
+        config._name_or_path = model
+    if model_override_args:
+        config.update(model_override_args)
+
+    # Special architecture mapping check for GGUF models
+    if is_gguf:
+        if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
+            raise RuntimeError(
+                f"Can't get gguf config for {config.model_type}.")
+        model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type]
+        config.update({"architectures": [model_type]})
+
+    return config
+
+
+def get_diffusers_config(
+    model: str,
+    inference_args: Optional[dict] = None,
+) -> Dict[str, Any]:
+    """Gets a configuration for the given diffusers model.
+    
+    Args:
+        model: The model name or path.
+        inference_args: Optional inference arguments to override in the config.
+        
+    Returns:
+        The loaded configuration.
+    """
+    config_name = "config.json"
+    if "scheduler" in model:
+        config_name = "scheduler_config.json"
+    # Check if the model path exists
+    if os.path.exists(model):
+        config_file = os.path.join(model, config_name)
+        if os.path.exists(config_file):
+            try:
+                # Load the config directly from the file
+                with open(config_file) as f:
+                    config_dict: Dict[str, Any] = json.load(f)
+
+                # TODO(will): apply any overrides from inference args
+                return config_dict
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to load diffusers config from {config_file}: {e}"
+                ) from e
+        raise RuntimeError(f"Config file not found at {config_file}")
+    else:
+        raise RuntimeError(f"Diffusers config file not found at {model}")
+
+
+# Models don't use the same configuration key for determining the maximum
+# context length.  Store them here so we can sanely check them.
+# NOTE: The ordering here is important. Some models have two of these and we
+# have a preference for which value gets used.
+CONTEXT_LENGTH_KEYS = [
+    "max_sequence_length",
+    "seq_length",
+    "max_seq_len",
+    "model_max_length",
+    "max_position_embeddings",
+]
+
+
+def attach_additional_stop_token_ids(tokenizer):
+    # Special handling for stop token <|eom_id|> generated by llama 3 tool use.
+    if "<|eom_id|>" in tokenizer.get_added_vocab():
+        tokenizer.additional_stop_token_ids = set(
+            [tokenizer.get_added_vocab()["<|eom_id|>"]])
+    else:
+        tokenizer.additional_stop_token_ids = None
+
+
+def check_gguf_file(model: Union[str, os.PathLike]) -> bool:
+    """Check if the file is a GGUF model."""
+    model = Path(model)
+    if not model.is_file():
+        return False
+    elif model.suffix == ".gguf":
+        return True
+
+    with open(model, "rb") as f:
+        header = f.read(4)
+    return header == b"GGUF"

fastvideo/v1/models/loader/component_loader.py

@@ -0,0 +1,474 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+import glob
+import os
+import time
+from abc import ABC, abstractmethod
+from typing import Any, Generator, Iterable, List, Optional, Tuple, cast
+
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file as safetensors_load_file
+from transformers import AutoImageProcessor, AutoTokenizer, PretrainedConfig
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
+
+from fastvideo.v1.inference_args import InferenceArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.hf_transformer_utils import (get_diffusers_config,
+                                                      get_hf_config)
+from fastvideo.v1.models.loader.fsdp_load import load_fsdp_model
+from fastvideo.v1.models.loader.utils import set_default_torch_dtype
+from fastvideo.v1.models.loader.weight_utils import (
+    filter_duplicate_safetensors_files, filter_files_not_needed_for_inference,
+    pt_weights_iterator, safetensors_weights_iterator)
+from fastvideo.v1.models.registry import ModelRegistry
+from fastvideo.v1.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class ComponentLoader(ABC):
+    """Base class for loading a specific type of model component."""
+
+    def __init__(self, device=None) -> None:
+        self.device = device
+
+    @abstractmethod
+    def load(self, model_path: str, architecture: str,
+             inference_args: InferenceArgs):
+        """
+        Load the component based on the model path, architecture, and inference args.
+        
+        Args:
+            model_path: Path to the component model
+            architecture: Architecture of the component model
+            inference_args: Inference arguments
+            
+        Returns:
+            The loaded component
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def for_module_type(cls, module_type: str,
+                        transformers_or_diffusers: str) -> 'ComponentLoader':
+        """
+        Factory method to create a component loader for a specific module type.
+        
+        Args:
+            module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler")
+            transformers_or_diffusers: Whether the module is from transformers or diffusers
+            
+        Returns:
+            A component loader for the specified module type
+        """
+        # Map of module types to their loader classes and expected library
+        module_loaders = {
+            "scheduler": (SchedulerLoader, "diffusers"),
+            "transformer": (TransformerLoader, "diffusers"),
+            "vae": (VAELoader, "diffusers"),
+            "text_encoder": (TextEncoderLoader, "transformers"),
+            "text_encoder_2": (TextEncoderLoader, "transformers"),
+            "tokenizer": (TokenizerLoader, "transformers"),
+            "tokenizer_2": (TokenizerLoader, "transformers"),
+            "image_processor": (ImageProcessorLoader, "transformers"),
+            "image_encoder": (ImageEncoderLoader, "transformers"),
+        }
+
+        if module_type in module_loaders:
+            loader_cls, expected_library = module_loaders[module_type]
+            # Assert that the library matches what's expected for this module type
+            assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
+            return loader_cls()
+
+        # For unknown module types, use a generic loader
+        logger.warning(
+            "No specific loader found for module type: %s. Using generic loader.",
+            module_type)
+        return GenericComponentLoader(transformers_or_diffusers)
+
+
+class TextEncoderLoader(ComponentLoader):
+    """Loader for text encoders."""
+
+    @dataclasses.dataclass
+    class Source:
+        """A source for weights."""
+
+        model_or_path: str
+        """The model ID or path."""
+
+        prefix: str = ""
+        """A prefix to prepend to all weights."""
+
+        fall_back_to_pt: bool = True
+        """Whether .pt weights can be used."""
+
+        allow_patterns_overrides: Optional[list[str]] = None
+        """If defined, weights will load exclusively using these patterns."""
+
+    counter_before_loading_weights: float = 0.0
+    counter_after_loading_weights: float = 0.0
+
+    def _prepare_weights(
+        self,
+        model_name_or_path: str,
+        fall_back_to_pt: bool,
+        allow_patterns_overrides: Optional[list[str]],
+    ) -> Tuple[str, List[str], bool]:
+        """Prepare weights for the model.
+
+        If the model is not local, it will be downloaded."""
+        # model_name_or_path = (self._maybe_download_from_modelscope(
+        #     model_name_or_path, revision) or model_name_or_path)
+
+        is_local = os.path.isdir(model_name_or_path)
+        assert is_local, "Model path must be a local directory"
+
+        use_safetensors = False
+        index_file = SAFE_WEIGHTS_INDEX_NAME
+        allow_patterns = ["*.safetensors", "*.bin"]
+
+        if fall_back_to_pt:
+            allow_patterns += ["*.pt"]
+
+        if allow_patterns_overrides is not None:
+            allow_patterns = allow_patterns_overrides
+
+        hf_folder = model_name_or_path
+
+        hf_weights_files: List[str] = []
+        for pattern in allow_patterns:
+            hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
+            if len(hf_weights_files) > 0:
+                if pattern == "*.safetensors":
+                    use_safetensors = True
+                break
+
+        if use_safetensors:
+            hf_weights_files = filter_duplicate_safetensors_files(
+                hf_weights_files, hf_folder, index_file)
+        else:
+            hf_weights_files = filter_files_not_needed_for_inference(
+                hf_weights_files)
+
+        if len(hf_weights_files) == 0:
+            raise RuntimeError(
+                f"Cannot find any model weights with `{model_name_or_path}`")
+
+        return hf_folder, hf_weights_files, use_safetensors
+
+    def _get_weights_iterator(
+            self, source: "Source"
+    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+        """Get an iterator for the model weights based on the load format."""
+        hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
+            source.model_or_path, source.fall_back_to_pt,
+            source.allow_patterns_overrides)
+        if use_safetensors:
+            weights_iterator = safetensors_weights_iterator(hf_weights_files)
+        else:
+            weights_iterator = pt_weights_iterator(hf_weights_files)
+
+        if self.counter_before_loading_weights == 0.0:
+            self.counter_before_loading_weights = time.perf_counter()
+        # Apply the prefix.
+        return ((source.prefix + name, tensor)
+                for (name, tensor) in weights_iterator)
+
+    def _get_all_weights(
+        self,
+        model_config: Any,
+        model: nn.Module,
+    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+        primary_weights = TextEncoderLoader.Source(
+            model_config.model,
+            prefix="",
+            fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
+            allow_patterns_overrides=getattr(model, "allow_patterns_overrides",
+                                             None),
+        )
+        yield from self._get_weights_iterator(primary_weights)
+
+        secondary_weights = cast(
+            Iterable[TextEncoderLoader.Source],
+            getattr(model, "secondary_weights", ()),
+        )
+        for source in secondary_weights:
+            yield from self._get_weights_iterator(source)
+
+    def load(self, model_path: str, architecture: str,
+             inference_args: InferenceArgs):
+        """Load the text encoders based on the model path, architecture, and inference args."""
+        model_config: PretrainedConfig = get_hf_config(
+            model=model_path,
+            trust_remote_code=inference_args.trust_remote_code,
+            revision=inference_args.revision,
+            model_override_args=None,
+            inference_args=inference_args,
+        )
+        logger.info("HF Model config: %s", model_config)
+
+        target_device = torch.device(inference_args.device_str)
+        # TODO(will): add support for other dtypes
+        return self.load_model(model_path, model_config, target_device,
+                               inference_args.text_encoder_precision)
+
+    def load_model(self,
+                   model_path: str,
+                   model_config,
+                   target_device: torch.device,
+                   dtype: str = "fp16"):
+        with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]):
+            with target_device:
+                architectures = getattr(model_config, "architectures", [])
+                model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
+                model = model_cls(model_config)
+
+            weights_to_load = {name for name, _ in model.named_parameters()}
+            model_config.model = model_path
+            loaded_weights = model.load_weights(
+                self._get_all_weights(model_config, model))
+            self.counter_after_loading_weights = time.perf_counter()
+            logger.info(
+                "Loading weights took %.2f seconds",
+                self.counter_after_loading_weights -
+                self.counter_before_loading_weights)
+            # We only enable strict check for non-quantized models
+            # that have loaded weights tracking currently.
+            # if loaded_weights is not None:
+            weights_not_loaded = weights_to_load - loaded_weights
+            if weights_not_loaded:
+                raise ValueError("Following weights were not initialized from "
+                                 f"checkpoint: {weights_not_loaded}")
+
+        # TODO(will): add support for training/finetune
+        return model.eval()
+
+
+class ImageEncoderLoader(TextEncoderLoader):
+
+    def load(self, model_path: str, architecture: str,
+             inference_args: InferenceArgs):
+        """Load the text encoders based on the model path, architecture, and inference args."""
+        model_config: PretrainedConfig = get_hf_config(
+            model=model_path,
+            trust_remote_code=inference_args.trust_remote_code,
+            revision=inference_args.revision,
+            model_override_args=None,
+            inference_args=inference_args,
+        )
+        logger.info("HF Model config: %s", model_config)
+
+        target_device = torch.device(inference_args.device_str)
+        # TODO(will): add support for other dtypes
+        return self.load_model(model_path, model_config, target_device,
+                               inference_args.image_encoder_precision)
+
+
+class ImageProcessorLoader(ComponentLoader):
+    """Loader for image processor."""
+
+    def load(self, model_path: str, architecture: str,
+             inference_args: InferenceArgs):
+        """Load the image processor based on the model path, architecture, and inference args."""
+        logger.info("Loading image processor from %s", model_path)
+
+        image_processor = AutoImageProcessor.from_pretrained(model_path, )
+        logger.info("Loaded image processor: %s",
+                    image_processor.__class__.__name__)
+        return image_processor
+
+
+class TokenizerLoader(ComponentLoader):
+    """Loader for tokenizers."""
+
+    def load(self, model_path: str, architecture: str,
+             inference_args: InferenceArgs):
+        """Load the tokenizer based on the model path, architecture, and inference args."""
+        logger.info("Loading tokenizer from %s", model_path)
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            # TODO(will): pass these tokenizer kwargs from inference args? Maybe
+            # other method of config?
+            padding_size='right',
+        )
+        logger.info("Loaded tokenizer: %s", tokenizer.__class__.__name__)
+        return tokenizer
+
+
+class VAELoader(ComponentLoader):
+    """Loader for VAE."""
+
+    def load(self, model_path: str, architecture: str,
+             inference_args: InferenceArgs):
+        """Load the VAE based on the model path, architecture, and inference args."""
+        # TODO(will): move this to a constants file
+        config = get_diffusers_config(model=model_path)
+
+        class_name = config.pop("_class_name")
+        assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported."
+        config.pop("_diffusers_version")
+
+        vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+
+        vae = vae_cls(**config).to(inference_args.device)
+
+        # Find all safetensors files
+        safetensors_list = glob.glob(
+            os.path.join(str(model_path), "*.safetensors"))
+        # TODO(PY)
+        assert len(
+            safetensors_list
+        ) == 1, f"Found {len(safetensors_list)} safetensors files in {model_path}"
+        loaded = safetensors_load_file(safetensors_list[0])
+        vae.load_state_dict(loaded)
+        dtype = PRECISION_TO_TYPE[inference_args.vae_precision]
+        vae = vae.eval().to(dtype)
+
+        return vae
+
+
+class TransformerLoader(ComponentLoader):
+    """Loader for transformer."""
+
+    def load(self, model_path: str, architecture: str,
+             inference_args: InferenceArgs):
+        """Load the transformer based on the model path, architecture, and inference args."""
+        model_config = get_diffusers_config(model=model_path)
+        cls_name = model_config.pop("_class_name")
+        if cls_name is None:
+            raise ValueError(
+                "Model config does not contain a _class_name attribute. "
+                "Only diffusers format is supported.")
+        model_config.pop("_diffusers_version")
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
+
+        # Find all safetensors files
+        safetensors_list = glob.glob(
+            os.path.join(str(model_path), "*.safetensors"))
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {model_path}")
+
+        logger.info("Loading model from %s safetensors files in %s",
+                    len(safetensors_list), model_path)
+
+        # initialize_sequence_parallel_group(inference_args.sp_size)
+        default_dtype = PRECISION_TO_TYPE[inference_args.precision]
+
+        # Load the model using FSDP loader
+        logger.info("Loading model from %s", cls_name)
+        model = load_fsdp_model(model_cls=model_cls,
+                                init_params=model_config,
+                                weight_dir_list=safetensors_list,
+                                device=inference_args.device,
+                                cpu_offload=inference_args.use_cpu_offload,
+                                default_dtype=default_dtype)
+
+        total_params = sum(p.numel() for p in model.parameters())
+        logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
+
+        dtypes = set(param.dtype for param in model.parameters())
+        if len(dtypes) > 1:
+            model = model.to(default_dtype)
+        model = model.eval()
+        return model
+
+
+class SchedulerLoader(ComponentLoader):
+    """Loader for scheduler."""
+
+    def load(self, model_path: str, architecture: str,
+             inference_args: InferenceArgs):
+        """Load the scheduler based on the model path, architecture, and inference args."""
+        config = get_diffusers_config(model=model_path)
+
+        class_name = config.pop("_class_name")
+        assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported."
+        config.pop("_diffusers_version")
+
+        scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+
+        scheduler = scheduler_cls(**config)
+        if inference_args.flow_shift is not None:
+            scheduler.set_shift(inference_args.flow_shift)
+
+        return scheduler
+
+
+class GenericComponentLoader(ComponentLoader):
+    """Generic loader for components that don't have a specific loader."""
+
+    def __init__(self, library="transformers") -> None:
+        super().__init__()
+        self.library = library
+
+    def load(self, model_path: str, architecture: str,
+             inference_args: InferenceArgs):
+        """Load a generic component based on the model path, architecture, and inference args."""
+        logger.warning("Using generic loader for %s with library %s",
+                       model_path, self.library)
+
+        if self.library == "transformers":
+            from transformers import AutoModel
+
+            model = AutoModel.from_pretrained(
+                model_path,
+                trust_remote_code=inference_args.trust_remote_code,
+                revision=inference_args.revision,
+            )
+            logger.info("Loaded generic transformers model: %s",
+                        model.__class__.__name__)
+            return model
+        elif self.library == "diffusers":
+            logger.warning(
+                "Generic loading for diffusers components is not fully implemented"
+            )
+
+            model_config = get_diffusers_config(model=model_path)
+            logger.info("Diffusers Model config: %s", model_config)
+            # This is a placeholder - in a real implementation, you'd need to handle this properly
+            return None
+        else:
+            raise ValueError(f"Unsupported library: {self.library}")
+
+
+class PipelineComponentLoader:
+    """
+    Utility class for loading pipeline components.
+    This replaces the chain of if-else statements in load_pipeline_module.
+    """
+
+    @staticmethod
+    def load_module(module_name: str, component_model_path: str,
+                    transformers_or_diffusers: str, architecture: str,
+                    inference_args: InferenceArgs):
+        """
+        Load a pipeline module.
+        
+        Args:
+            module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler")
+            component_model_path: Path to the component model
+            transformers_or_diffusers: Whether the module is from transformers or diffusers
+            architecture: Architecture of the component model
+            inference_args: Inference arguments
+            
+        Returns:
+            The loaded module
+        """
+        logger.info(
+            "Loading %s using %s from %s",
+            module_name,
+            transformers_or_diffusers,
+            component_model_path,
+        )
+
+        # Get the appropriate loader for this module type
+        loader = ComponentLoader.for_module_type(module_name,
+                                                 transformers_or_diffusers)
+
+        # Load the module
+        return loader.load(component_model_path, architecture, inference_args)

fastvideo/v1/models/loader/fsdp_load.py

@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Adapted from torchtune
+# Copyright 2024 The TorchTune Authors.
+# Copyright 2025 The FastVideo Authors.
+
+import contextlib
+import re
+from collections import defaultdict
+from itertools import chain
+from typing import (Any, Callable, DefaultDict, Dict, Generator, Hashable, List,
+                    Optional, Tuple, Type)
+
+import torch
+from torch import nn
+from torch.distributed import DeviceMesh, init_device_mesh
+from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard
+from torch.distributed._tensor import distribute_tensor
+from torch.nn.modules.module import _IncompatibleKeys
+
+from fastvideo.v1.distributed.parallel_state import (
+    get_sequence_model_parallel_world_size)
+from fastvideo.v1.models.loader.weight_utils import safetensors_weights_iterator
+
+
+# TODO(PY): move this to utils elsewhere
+@contextlib.contextmanager
+def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
+    """
+    Context manager to set torch's default dtype.
+
+    Args:
+        dtype (torch.dtype): The desired default dtype inside the context manager.
+
+    Returns:
+        ContextManager: context manager for setting default dtype.
+
+    Example:
+        >>> with set_default_dtype(torch.bfloat16):
+        >>>     x = torch.tensor([1, 2, 3])
+        >>>     x.dtype
+        torch.bfloat16
+
+
+    """
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(old_dtype)
+
+
+def get_param_names_mapping(
+        mapping_dict: Dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]:
+    """
+    Creates a mapping function that transforms parameter names using regex patterns.
+    
+    Args:
+        mapping_dict (Dict[str, str]): Dictionary mapping regex patterns to replacement patterns
+        param_name (str): The parameter name to be transformed
+        
+    Returns:
+        Callable[[str], str]: A function that maps parameter names from source to target format
+    """
+
+    def mapping_fn(name: str) -> tuple[str, Any, Any]:
+
+        # Try to match and transform the name using the regex patterns in mapping_dict
+        for pattern, replacement in mapping_dict.items():
+            match = re.match(pattern, name)
+            if match:
+                merge_index = None
+                total_splitted_params = None
+                if isinstance(replacement, tuple):
+                    merge_index = replacement[1]
+                    total_splitted_params = replacement[2]
+                    replacement = replacement[0]
+                name = re.sub(pattern, replacement, name)
+                return name, merge_index, total_splitted_params
+
+        # If no pattern matches, return the original name
+        return name, None, None
+
+    return mapping_fn
+
+
+# TODO(PY): add compile option
+def load_fsdp_model(
+    model_cls: Type[nn.Module],
+    init_params: Dict[str, Any],
+    weight_dir_list: List[str],
+    device: torch.device,
+    cpu_offload: bool = False,
+    default_dtype: Optional[torch.dtype] = torch.bfloat16,
+) -> torch.nn.Module:
+    with set_default_dtype(default_dtype), torch.device("meta"):
+        model = model_cls(**init_params)
+    device_mesh = init_device_mesh(
+        "cuda",
+        mesh_shape=(get_sequence_model_parallel_world_size(), ),
+        mesh_dim_names=("dp", ),
+    )
+    shard_model(model,
+                cpu_offload=cpu_offload,
+                reshard_after_forward=True,
+                dp_mesh=device_mesh["dp"])
+    weight_iterator = safetensors_weights_iterator(weight_dir_list)
+    param_names_mapping_fn = get_param_names_mapping(model._param_names_mapping)
+    load_fsdp_model_from_full_model_state_dict(
+        model,
+        weight_iterator,
+        device,
+        strict=True,
+        cpu_offload=cpu_offload,
+        param_names_mapping=param_names_mapping_fn,
+    )
+    for n, p in chain(model.named_parameters(), model.named_buffers()):
+        if p.is_meta:
+            raise RuntimeError(
+                f"Unexpected param or buffer {n} on meta device.")
+    for p in model.parameters():
+        p.requires_grad = False
+    return model
+
+
+def shard_model(
+    model,
+    *,
+    cpu_offload: bool,
+    reshard_after_forward: bool = True,
+    dp_mesh: Optional[DeviceMesh] = None,
+) -> None:
+    """
+    Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API.
+
+    This method will over the model's named modules from the bottom-up and apply shard modules
+    based on whether they meet any of the criteria from shard_conditions.
+
+    Args:
+        model (TransformerDecoder): Model to shard with FSDP.
+        shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine
+            which modules to shard with FSDP. Each function should take module name (relative to root)
+            and the module itself, returning True if FSDP should shard the module and False otherwise.
+            If any of shard_conditions return True for a given module, it will be sharded by FSDP.
+        cpu_offload (bool): If set to True, FSDP will offload parameters, gradients, and optimizer
+            states to CPU.
+        reshard_after_forward (bool): Whether to reshard parameters and buffers after
+            the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy
+            from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy.
+        dp_mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
+            Default to None.
+
+    Raises:
+        ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered.
+    """
+    fsdp_kwargs = {
+        "reshard_after_forward": reshard_after_forward,
+        "mesh": dp_mesh
+    }
+    if cpu_offload:
+        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy()
+
+    # Shard the model with FSDP, iterating in reverse to start with
+    # lowest-level modules first
+    num_layers_sharded = 0
+    for n, m in reversed(list(model.named_modules())):
+        if any([
+                shard_condition(n, m)
+                for shard_condition in model._fsdp_shard_conditions
+        ]):
+            fully_shard(m, **fsdp_kwargs)
+            num_layers_sharded += 1
+
+    if num_layers_sharded == 0:
+        raise ValueError(
+            "No layer modules were sharded. Please check if shard conditions are working as expected."
+        )
+
+    # Finally shard the entire model to account for any stragglers
+    fully_shard(model, **fsdp_kwargs)
+
+
+# TODO(PY): device mesh for cfg parallel
+def load_fsdp_model_from_full_model_state_dict(
+    model: torch.nn.Module,
+    full_sd_iterator: Generator[Tuple[str, torch.Tensor], None, None],
+    device: torch.device,
+    strict: bool = False,
+    cpu_offload: bool = False,
+    param_names_mapping: Optional[Callable[[str], tuple[str, Any, Any]]] = None,
+) -> _IncompatibleKeys:
+    """
+    Converting full state dict into a sharded state dict
+    and loading it into FSDP model
+    Args:
+        model (FSDPModule): Model to generate fully qualified names for cpu_state_dict
+        full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs
+        device (torch.device): device used to move full state dict tensors
+        strict (bool): flag to check if to load the model in strict mode
+        cpu_offload (bool): flag to check if offload to CPU is enabled
+        param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
+
+    Returns:
+        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
+            * **missing_keys** is a list of str containing the missing keys
+            * **unexpected_keys** is a list of str containing the unexpected keys
+
+    Raises:
+        NotImplementedError: If got FSDP with more than 1D.
+    """
+    meta_sharded_sd = model.state_dict()
+
+    sharded_sd = {}
+    to_merge_params: DefaultDict[Hashable, Dict[Any, Any]] = defaultdict(dict)
+    for source_param_name, full_tensor in full_sd_iterator:
+        assert param_names_mapping is not None
+        target_param_name, merge_index, num_params_to_merge = param_names_mapping(
+            source_param_name)
+
+        if merge_index is not None:
+            to_merge_params[target_param_name][merge_index] = full_tensor
+            if len(to_merge_params[target_param_name]) == num_params_to_merge:
+                # cat at dim=1 according to the merge_index order
+                sorted_tensors = [
+                    to_merge_params[target_param_name][i]
+                    for i in range(num_params_to_merge)
+                ]
+                full_tensor = torch.cat(sorted_tensors, dim=0)
+                del to_merge_params[target_param_name]
+            else:
+                continue
+
+        sharded_meta_param = meta_sharded_sd.get(target_param_name)
+        if sharded_meta_param is None:
+            raise ValueError(
+                f"Parameter {source_param_name}-->{target_param_name} not found in meta sharded state dict"
+            )
+        full_tensor = full_tensor.to(sharded_meta_param.dtype).to(device)
+
+        if not hasattr(sharded_meta_param, "device_mesh"):
+            # In cases where parts of the model aren't sharded, some parameters will be plain tensors
+            sharded_tensor = full_tensor
+        else:
+            sharded_tensor = distribute_tensor(
+                full_tensor,
+                sharded_meta_param.device_mesh,
+                sharded_meta_param.placements,
+            )
+        if cpu_offload:
+            sharded_tensor = sharded_tensor.cpu()
+        sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)
+    # choose `assign=True` since we cannot call `copy_` on meta tensor
+    return model.load_state_dict(sharded_sd, strict=strict, assign=True)

fastvideo/v1/models/loader/utils.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for selecting and loading models."""
+import contextlib
+
+import torch
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@contextlib.contextmanager
+def set_default_torch_dtype(dtype: torch.dtype):
+    """Sets the default torch dtype to the given dtype."""
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    yield
+    torch.set_default_dtype(old_dtype)

fastvideo/v1/models/loader/weight_utils.py

@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/model_loader/weight_utils.py
+"""Utilities for downloading and initializing model weights."""
+import fnmatch
+import hashlib
+import json
+import os
+import tempfile
+import time
+from collections import defaultdict
+from pathlib import Path
+from typing import Generator, List, Optional, Tuple, Union
+
+import filelock
+import huggingface_hub.constants
+import torch
+from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
+from safetensors.torch import safe_open
+from tqdm.auto import tqdm
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+# use system-level temp directory for file locks, so that multiple users
+# can share the same lock without error.
+# lock files in the temp directory will be automatically deleted when the
+# system reboots, so users will not complain about annoying lock files
+temp_dir = tempfile.gettempdir()
+
+
+def enable_hf_transfer() -> None:
+    """automatically activates hf_transfer
+    """
+    if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ:
+        try:
+            # enable hf hub transfer if available
+            import hf_transfer  # type: ignore # noqa
+            huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True
+        except ImportError:
+            pass
+
+
+enable_hf_transfer()
+
+
+class DisabledTqdm(tqdm):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, disable=True)
+
+
+def get_lock(model_name_or_path: Union[str, Path],
+             cache_dir: Optional[str] = None):
+    lock_dir = cache_dir or temp_dir
+    model_name_or_path = str(model_name_or_path)
+    os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
+    model_name = model_name_or_path.replace("/", "-")
+    hash_name = hashlib.sha256(model_name.encode()).hexdigest()
+    # add hash to avoid conflict with old users' lock files
+    lock_file_name = hash_name + model_name + ".lock"
+    # mode 0o666 is required for the filelock to be shared across users
+    lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
+    return lock
+
+
+def _shared_pointers(tensors):
+    ptrs = defaultdict(list)
+    for k, v in tensors.items():
+        ptrs[v.data_ptr()].append(k)
+    failing = []
+    for _, names in ptrs.items():
+        if len(names) > 1:
+            failing.append(names)
+    return failing
+
+
+def download_weights_from_hf(
+    model_name_or_path: str,
+    cache_dir: Optional[str],
+    allow_patterns: List[str],
+    revision: Optional[str] = None,
+    ignore_patterns: Optional[Union[str, List[str]]] = None,
+) -> str:
+    """Download model weights from Hugging Face Hub.
+
+    Args:
+        model_name_or_path (str): The model name or path.
+        cache_dir (Optional[str]): The cache directory to store the model
+            weights. If None, will use HF defaults.
+        allow_patterns (List[str]): The allowed patterns for the
+            weight files. Files matched by any of the patterns will be
+            downloaded.
+        revision (Optional[str]): The revision of the model.
+        ignore_patterns (Optional[Union[str, List[str]]]): The patterns to
+            filter out the weight files. Files matched by any of the patterns
+            will be ignored.
+
+    Returns:
+        str: The path to the downloaded model weights.
+    """
+    local_only = huggingface_hub.constants.HF_HUB_OFFLINE
+    if not local_only:
+        # Before we download we look at that is available:
+        fs = HfFileSystem()
+        file_list = fs.ls(model_name_or_path, detail=False, revision=revision)
+
+        # depending on what is available we download different things
+        for pattern in allow_patterns:
+            matching = fnmatch.filter(file_list, pattern)
+            if len(matching) > 0:
+                allow_patterns = [pattern]
+                break
+
+    logger.info("Using model weights format %s", allow_patterns)
+    # Use file lock to prevent multiple processes from
+    # downloading the same model weights at the same time.
+    with get_lock(model_name_or_path, cache_dir):
+        start_time = time.perf_counter()
+        hf_folder: str = snapshot_download(
+            model_name_or_path,
+            allow_patterns=allow_patterns,
+            ignore_patterns=ignore_patterns,
+            cache_dir=cache_dir,
+            tqdm_class=DisabledTqdm,
+            revision=revision,
+            local_files_only=local_only,
+        )
+        time_taken = time.perf_counter() - start_time
+        if time_taken > 0.5:
+            logger.info("Time spent downloading weights for %s: %.6f seconds",
+                        model_name_or_path, time_taken)
+    return hf_folder
+
+
+def download_safetensors_index_file_from_hf(
+    model_name_or_path: str,
+    index_file: str,
+    cache_dir: Optional[str],
+    revision: Optional[str] = None,
+) -> None:
+    """Download hf safetensors index file from Hugging Face Hub.
+
+    Args:
+        model_name_or_path (str): The model name or path.
+        cache_dir (Optional[str]): The cache directory to store the model
+            weights. If None, will use HF defaults.
+        revision (Optional[str]): The revision of the model.
+    """
+    # Use file lock to prevent multiple processes from
+    # downloading the same model weights at the same time.
+    with get_lock(model_name_or_path, cache_dir):
+        try:
+            # Download the safetensors index file.
+            hf_hub_download(
+                repo_id=model_name_or_path,
+                filename=index_file,
+                cache_dir=cache_dir,
+                revision=revision,
+                local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+            )
+        # If file not found on remote or locally, we should not fail since
+        # only some models will have index_file.
+        except huggingface_hub.utils.EntryNotFoundError:
+            logger.info("No %s found in remote.", index_file)
+        except huggingface_hub.utils.LocalEntryNotFoundError:
+            logger.info("No %s found in local cache.", index_file)
+
+
+# For models like Mistral-7B-v0.3, there are both sharded
+# safetensors files and a consolidated safetensors file.
+# Passing both of these to the weight loader functionality breaks.
+# So, we use the index_file to
+# look up which safetensors files should be used.
+def filter_duplicate_safetensors_files(hf_weights_files: List[str],
+                                       hf_folder: str,
+                                       index_file: str) -> List[str]:
+    # model.safetensors.index.json is a mapping from keys in the
+    # torch state_dict to safetensors file holding that weight.
+    index_file_name = os.path.join(hf_folder, index_file)
+    if not os.path.isfile(index_file_name):
+        return hf_weights_files
+
+    # Iterate through the weight_map (weight_name: safetensors files)
+    # to identify weights that we should use.
+    with open(index_file_name) as f:
+        weight_map = json.load(f)["weight_map"]
+    weight_files_in_index = set()
+    for weight_name in weight_map:
+        weight_files_in_index.add(
+            os.path.join(hf_folder, weight_map[weight_name]))
+    # Filter out any fields that are not found in the index file.
+    hf_weights_files = [
+        f for f in hf_weights_files if f in weight_files_in_index
+    ]
+    return hf_weights_files
+
+
+def filter_files_not_needed_for_inference(
+        hf_weights_files: List[str]) -> List[str]:
+    """
+    Exclude files that are not needed for inference.
+
+    See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233
+    """
+    blacklist = [
+        "training_args.bin",
+        "optimizer.bin",
+        "optimizer.pt",
+        "scheduler.pt",
+        "scaler.pt",
+    ]
+    hf_weights_files = [
+        f for f in hf_weights_files
+        if not any(f.endswith(x) for x in blacklist)
+    ]
+    return hf_weights_files
+
+
+# explicitly use pure text format, with a newline at the end
+# this makes it impossible to see the animation in the progress bar
+# but will avoid messing up with ray or multiprocessing, which wraps
+# each line of output with some prefix.
+_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n"  # noqa: E501
+
+
+def safetensors_weights_iterator(
+    hf_weights_files: List[str]
+) -> Generator[Tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model safetensor files."""
+    enable_tqdm = not torch.distributed.is_initialized(
+    ) or torch.distributed.get_rank() == 0
+    for st_file in tqdm(
+            hf_weights_files,
+            desc="Loading safetensors checkpoint shards",
+            disable=not enable_tqdm,
+            bar_format=_BAR_FORMAT,
+    ):
+        with safe_open(st_file, framework="pt") as f:
+            for name in f.keys():  # noqa: SIM118
+                param = f.get_tensor(name)
+                yield name, param
+
+
+def pt_weights_iterator(
+    hf_weights_files: List[str]
+) -> Generator[Tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model bin/pt files."""
+    enable_tqdm = not torch.distributed.is_initialized(
+    ) or torch.distributed.get_rank() == 0
+    for bin_file in tqdm(
+            hf_weights_files,
+            desc="Loading pt checkpoint shards",
+            disable=not enable_tqdm,
+            bar_format=_BAR_FORMAT,
+    ):
+        state = torch.load(bin_file, map_location="cpu", weights_only=True)
+        yield from state.items()
+        del state
+
+
+def default_weight_loader(param: torch.Tensor,
+                          loaded_weight: torch.Tensor) -> None:
+    """Default weight loader."""
+    try:
+        if param.numel() == 1 and loaded_weight.numel() == 1:
+            # Sometimes scalar values aren't considered tensors with shapes
+            # so if both param and loaded_weight are a scalar,
+            # "broadcast" instead of copy
+            param.data.fill_(loaded_weight.item())
+        else:
+            assert param.size() == loaded_weight.size(), (
+                f"Attempted to load weight ({loaded_weight.size()}) "
+                f"into parameter ({param.size()})")
+
+            param.data.copy_(loaded_weight)
+    except Exception:
+        # NOTE: This exception is added for the purpose of setting breakpoint to
+        # debug weight loading issues.
+        raise
+
+
+def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
+    """Remap the name of FP8 k/v_scale parameters.
+
+    This function handles the remapping of FP8 k/v_scale parameter names.
+    It detects if the given name ends with a suffix and attempts to remap
+    it to the expected name format in the model. If the remapped name is not
+    found in the params_dict, a warning is printed and None is returned.
+
+    Args:
+        name (str): The original loaded checkpoint parameter name.
+        params_dict (dict): Dictionary containing the model's named parameters.
+
+    Returns:
+        str: The remapped parameter name if successful, or the original name
+             if no remapping is needed.
+        None: If the remapped name is not found in params_dict.
+    """
+    if name.endswith(".kv_scale"):
+        logger.warning_once(
+            "DEPRECATED. Found kv_scale in the checkpoint. "
+            "This format is deprecated in favor of separate k_scale and "
+            "v_scale tensors and will be removed in a future release. "
+            "Functionally, we will remap kv_scale to k_scale and duplicate "
+            "k_scale to v_scale")
+        # NOTE: we remap the deprecated kv_scale to k_scale
+        remapped_name = name.replace(".kv_scale", ".attn.k_scale")
+        if remapped_name not in params_dict:
+            logger.warning_once(
+                f"Found kv_scale in the checkpoint (e.g. {name}), "
+                "but not found the expected name in the model "
+                f"(e.g. {remapped_name}). kv_scale is "
+                "not loaded.")
+            return None
+        return remapped_name
+
+    possible_scale_names = [".k_scale", ".v_scale"]
+    modelopt_scale_names = [
+        ".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"
+    ]
+    for scale_name in possible_scale_names:
+        if name.endswith(scale_name):
+            if any(mo_scale_name in name
+                   for mo_scale_name in modelopt_scale_names):
+                remapped_name = name.replace(
+                    f".self_attn.{scale_name[1]}_proj{scale_name}",
+                    f".self_attn.attn{scale_name}")
+            else:
+                remapped_name = name.replace(scale_name, f".attn{scale_name}")
+            if remapped_name not in params_dict:
+                logger.warning_once(
+                    f"Found {scale_name} in the checkpoint (e.g. {name}), "
+                    "but not found the expected name in the model "
+                    f"(e.g. {remapped_name}). {scale_name} is "
+                    "not loaded.")
+                return None
+            return remapped_name
+
+    # If there were no matches, return the untouched param name
+    return name

fastvideo/v1/models/parameter.py

@@ -0,0 +1,429 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/parameter.py
+
+from fractions import Fraction
+from typing import Any, Callable, Optional, Tuple, Union
+
+import torch
+from torch.nn import Parameter
+
+from fastvideo.v1.distributed import get_tensor_model_parallel_rank
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.utils import _make_synced_weight_loader
+
+logger = init_logger(__name__)
+
+
+class BasevLLMParameter(Parameter):
+    """
+    Base parameter for vLLM linear layers. Extends the torch.nn.parameter
+    by taking in a linear weight loader. Will copy the loaded weight
+    into the parameter when the provided weight loader is called.
+    """
+
+    def __new__(cls, data: torch.Tensor, **kwargs):
+
+        return super().__new__(cls, data=data, requires_grad=False)
+
+    def __init__(self, data: torch.Tensor, weight_loader: Callable):
+        """
+        Initialize the BasevLLMParameter
+
+        :param data: torch tensor with the parameter data
+        :param weight_loader: weight loader callable
+
+        :returns: a torch.nn.parameter
+        """
+
+        # During weight loading, we often do something like:
+        # narrowed_tensor = param.data.narrow(0, offset, len)
+        # narrowed_tensor.copy_(real_weight)
+        # expecting narrowed_tensor and param.data to share the same storage.
+        # However, on TPUs, narrowed_tensor will lazily propagate to the base
+        # tensor, which is param.data, leading to the redundant memory usage.
+        # This sometimes causes OOM errors during model loading. To avoid this,
+        # we sync the param tensor after its weight loader is called.
+        from fastvideo.v1.platforms import current_platform
+        if current_platform.is_tpu():
+            weight_loader = _make_synced_weight_loader(weight_loader)
+
+        self._weight_loader = weight_loader
+
+    @property
+    def weight_loader(self):
+        return self._weight_loader
+
+    def _is_1d_and_scalar(self, loaded_weight: torch.Tensor):
+        cond1 = self.data.ndim == 1 and self.data.numel() == 1
+        cond2 = loaded_weight.ndim == 0 and loaded_weight.numel() == 1
+        return (cond1 and cond2)
+
+    def _assert_and_load(self, loaded_weight: torch.Tensor):
+        assert (self.data.shape == loaded_weight.shape
+                or self._is_1d_and_scalar(loaded_weight))
+        self.data.copy_(loaded_weight)
+
+    def load_column_parallel_weight(self, loaded_weight: torch.Tensor):
+        self._assert_and_load(loaded_weight)
+
+    def load_row_parallel_weight(self, loaded_weight: torch.Tensor):
+        self._assert_and_load(loaded_weight)
+
+    def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):
+        self._assert_and_load(loaded_weight)
+
+    def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):
+        self._assert_and_load(loaded_weight)
+
+
+class _ColumnvLLMParameter(BasevLLMParameter):
+    """
+    Private class defining weight loading functionality 
+    (load_merged_column_weight, load_qkv_weight)
+    for parameters being loaded into linear layers with column
+    parallelism. This includes QKV and MLP layers which are
+    not already fused on disk. Requires an output dimension 
+    to be defined. Called within the weight loader of
+    each of the column parallel linear layers.
+    """
+
+    def __init__(self, output_dim: int, **kwargs):
+        self._output_dim = output_dim
+        super().__init__(**kwargs)
+
+    @property
+    def output_dim(self):
+        return self._output_dim
+
+    def load_column_parallel_weight(self, loaded_weight: torch.Tensor):
+        tp_rank = get_tensor_model_parallel_rank()
+        shard_size = self.data.shape[self.output_dim]
+        loaded_weight = loaded_weight.narrow(self.output_dim,
+                                             tp_rank * shard_size, shard_size)
+        assert self.data.shape == loaded_weight.shape
+        self.data.copy_(loaded_weight)
+
+    def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):
+
+        shard_offset = kwargs.get("shard_offset")
+        shard_size = kwargs.get("shard_size")
+        if isinstance(
+                self,
+            (PackedColumnParameter,
+             PackedvLLMParameter)) and self.packed_dim == self.output_dim:
+            shard_size, shard_offset = self.adjust_shard_indexes_for_packing(
+                shard_offset=shard_offset, shard_size=shard_size)
+
+        param_data = self.data
+
+        tp_rank = get_tensor_model_parallel_rank()
+        param_data = param_data.narrow(self.output_dim, shard_offset,
+                                       shard_size)
+        loaded_weight = loaded_weight.narrow(self.output_dim,
+                                             tp_rank * shard_size, shard_size)
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+    def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):
+
+        shard_offset = kwargs.get("shard_offset")
+        shard_size = kwargs.get("shard_size")
+        shard_id = kwargs.get("shard_id")
+        num_heads = kwargs.get("num_heads")
+
+        if isinstance(
+                self,
+            (PackedColumnParameter,
+             PackedvLLMParameter)) and self.output_dim == self.packed_dim:
+            shard_size, shard_offset = self.adjust_shard_indexes_for_packing(
+                shard_offset=shard_offset, shard_size=shard_size)
+
+        param_data = self.data
+        tp_rank = get_tensor_model_parallel_rank()
+        shard_id = tp_rank if shard_id == "q" else tp_rank // num_heads
+        param_data = param_data.narrow(self.output_dim, shard_offset,
+                                       shard_size)
+        loaded_weight = loaded_weight.narrow(self.output_dim,
+                                             shard_id * shard_size, shard_size)
+
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+
+class RowvLLMParameter(BasevLLMParameter):
+    """
+    Parameter class defining weight_loading functionality
+    (load_row_parallel_weight) for parameters being loaded
+    into linear layers with row parallel functionality.
+    Requires an input_dim to be defined.
+    """
+
+    def __init__(self, input_dim: int, **kwargs):
+        self._input_dim = input_dim
+        super().__init__(**kwargs)
+
+    @property
+    def input_dim(self):
+        return self._input_dim
+
+    def load_row_parallel_weight(self, loaded_weight: torch.Tensor):
+        tp_rank = get_tensor_model_parallel_rank()
+        shard_size = self.data.shape[self.input_dim]
+        loaded_weight = loaded_weight.narrow(self.input_dim,
+                                             tp_rank * shard_size, shard_size)
+
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+
+        assert self.data.shape == loaded_weight.shape
+        self.data.copy_(loaded_weight)
+
+
+class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter):
+    """
+    Parameter class for linear layer weights. Uses both column and
+    row parallelism.
+    """
+    pass
+
+
+class GroupQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):
+    """
+    Parameter class for weight scales loaded for weights with
+    grouped quantization. Uses both column and row parallelism.
+    """
+    pass
+
+
+class ChannelQuantScaleParameter(_ColumnvLLMParameter):
+    """
+    Parameter class for weight scales loaded for weights with
+    channel-wise quantization. Equivalent to _ColumnvLLMParameter.
+    """
+    pass
+
+
+class PerTensorScaleParameter(BasevLLMParameter):
+    """
+    Parameter class for scales where the number of scales is
+    equivalent to the number of logical matrices in fused linear
+    layers (e.g. for QKV, there are 3 scales loaded from disk).
+    This is relevant to weights with per-tensor quantization.
+    Adds functionality to map the scalers to a shard during
+    weight loading. 
+
+    Note: additional parameter manipulation may be handled 
+    for each quantization config specifically, within 
+    process_weights_after_loading 
+    """
+
+    def __init__(self, **kwargs):
+        self.qkv_idxs = {"q": 0, "k": 1, "v": 2}
+        super().__init__(**kwargs)
+
+    def _shard_id_as_int(self, shard_id: Union[str, int]) -> int:
+        if isinstance(shard_id, int):
+            return shard_id
+
+        # if not int, assume shard_id for qkv
+        # map to int and return
+        assert isinstance(shard_id, str)
+        assert shard_id in self.qkv_idxs
+        return self.qkv_idxs[shard_id]
+
+    # For row parallel layers, no sharding needed
+    # load weight into parameter as is
+    def load_row_parallel_weight(self, *args, **kwargs):
+        super().load_row_parallel_weight(*args, **kwargs)
+
+    def load_merged_column_weight(self, *args, **kwargs):
+        self._load_into_shard_id(*args, **kwargs)
+
+    def load_qkv_weight(self, *args, **kwargs):
+        self._load_into_shard_id(*args, **kwargs)
+
+    def load_column_parallel_weight(self, *args, **kwargs):
+        super().load_row_parallel_weight(*args, **kwargs)
+
+    def _load_into_shard_id(self, loaded_weight: torch.Tensor,
+                            shard_id: Union[str, int], **kwargs):
+        """
+        Slice the parameter data based on the shard id for 
+        loading.
+        """
+
+        param_data = self.data
+        shard_id = self._shard_id_as_int(shard_id)
+
+        # AutoFP8 scales do not have a shape
+        # compressed-tensors scales do have a shape
+        if len(loaded_weight.shape) != 0:
+            assert loaded_weight.shape[0] == 1
+            loaded_weight = loaded_weight[0]
+
+        param_data = param_data[shard_id]
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+
+class PackedColumnParameter(_ColumnvLLMParameter):
+    """
+    Parameter for model parameters which are packed on disk
+    and support column parallelism only. See PackedvLLMParameter
+    for more details on the packed properties.
+    """
+
+    def __init__(self,
+                 packed_factor: Union[int, Fraction],
+                 packed_dim: int,
+                 marlin_tile_size: Optional[int] = None,
+                 **kwargs):
+        self._packed_factor = packed_factor
+        self._packed_dim = packed_dim
+        self._marlin_tile_size = marlin_tile_size
+        super().__init__(**kwargs)
+
+    @property
+    def packed_dim(self):
+        return self._packed_dim
+
+    @property
+    def packed_factor(self):
+        return self._packed_factor
+
+    @property
+    def marlin_tile_size(self):
+        return self._marlin_tile_size
+
+    def adjust_shard_indexes_for_packing(self, shard_size,
+                                         shard_offset) -> Tuple[Any, Any]:
+        return _adjust_shard_indexes_for_packing(
+            shard_size=shard_size,
+            shard_offset=shard_offset,
+            packed_factor=self.packed_factor,
+            marlin_tile_size=self.marlin_tile_size)
+
+
+class PackedvLLMParameter(ModelWeightParameter):
+    """
+    Parameter for model weights which are packed on disk.
+    Example: GPTQ Marlin weights are int4 or int8, packed into int32.
+    Extends the ModelWeightParameter to take in the
+    packed factor, the packed dimension, and optionally, marlin
+    tile size for marlin kernels. Adjusts the shard_size and 
+    shard_offset for fused linear layers model weight loading
+    by accounting for packing and optionally, marlin tile size.
+    """
+
+    def __init__(self,
+                 packed_factor: Union[int, Fraction],
+                 packed_dim: int,
+                 marlin_tile_size: Optional[int] = None,
+                 **kwargs):
+        self._packed_factor = packed_factor
+        self._packed_dim = packed_dim
+        self._marlin_tile_size = marlin_tile_size
+        super().__init__(**kwargs)
+
+    @property
+    def packed_dim(self):
+        return self._packed_dim
+
+    @property
+    def packed_factor(self):
+        return self._packed_factor
+
+    @property
+    def marlin_tile_size(self):
+        return self._marlin_tile_size
+
+    def adjust_shard_indexes_for_packing(self, shard_size, shard_offset):
+        return _adjust_shard_indexes_for_packing(
+            shard_size=shard_size,
+            shard_offset=shard_offset,
+            packed_factor=self.packed_factor,
+            marlin_tile_size=self.marlin_tile_size)
+
+
+class BlockQuantScaleParameter(_ColumnvLLMParameter, RowvLLMParameter):
+    """
+    Parameter class for weight scales loaded for weights with
+    block-wise quantization. Uses both column and row parallelism.
+    """
+
+    pass
+
+
+def permute_param_layout_(param: BasevLLMParameter, input_dim: int,
+                          output_dim: int, **kwargs) -> BasevLLMParameter:
+    """
+    Permute a parameter's layout to the specified input and output dimensions, 
+    useful for forcing the parameter into a known layout, for example, if I need
+    a packed (quantized) weight matrix to be in the layout 
+        {input_dim = 0, output_dim = 1, packed_dim = 0}
+    then I can call:
+        permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
+    to ensure x is in the correct layout (permuting it to the correct layout if 
+    required, asserting if it cannot get it to the correct layout)
+    """
+
+    curr_input_dim = getattr(param, "input_dim", None)
+    curr_output_dim = getattr(param, "output_dim", None)
+
+    if curr_input_dim is None or curr_output_dim is None:
+        assert param.data.dim() == 2,\
+            "permute_param_layout_ only supports 2D parameters when either "\
+            "input_dim or output_dim is not set"
+
+    # if one of the dimensions is not set, set it to the opposite of the other
+    #  we can only do this since we asserted the parameter is 2D above
+    if curr_input_dim is None:
+        assert curr_output_dim is not None,\
+            "either input or output dim must be set"
+        curr_input_dim = (curr_output_dim + 1) % 2
+    if curr_output_dim is None:
+        assert curr_input_dim is not None,\
+            "either input or output dim must be set"
+        curr_output_dim = (curr_input_dim + 1) % 2
+
+    # create permutation from the current layout to the layout with
+    # self.input_dim at input_dim and self.output_dim at output_dim preserving
+    # other dimensions
+    perm = [
+        i for i in range(param.data.dim())
+        if i not in [curr_input_dim, curr_output_dim]
+    ]
+    perm.insert(input_dim, curr_input_dim)
+    perm.insert(output_dim, curr_output_dim)
+
+    if "packed_dim" in kwargs:
+        assert hasattr(param, "packed_dim") and\
+            param.packed_dim == perm[kwargs["packed_dim"]],\
+            "permute_param_layout_ currently doesn't support repacking"
+
+    param.data = param.data.permute(*perm)
+    if hasattr(param, "_input_dim"):
+        param._input_dim = input_dim
+    if hasattr(param, "_output_dim"):
+        param._output_dim = output_dim
+    if "packed_dim" in kwargs and hasattr(param, "_packed_dim"):
+        param._packed_dim = kwargs["packed_dim"]
+
+    return param
+
+
+def _adjust_shard_indexes_for_marlin(shard_size, shard_offset,
+                                     marlin_tile_size) -> Tuple[Any, Any]:
+    return shard_size * marlin_tile_size, shard_offset * marlin_tile_size
+
+
+def _adjust_shard_indexes_for_packing(shard_size, shard_offset, packed_factor,
+                                      marlin_tile_size) -> Tuple[Any, Any]:
+    shard_size = shard_size // packed_factor
+    shard_offset = shard_offset // packed_factor
+    if marlin_tile_size is not None:
+        return _adjust_shard_indexes_for_marlin(
+            shard_size=shard_size,
+            shard_offset=shard_offset,
+            marlin_tile_size=marlin_tile_size)
+    return shard_size, shard_offset

fastvideo/v1/models/registry.py

@@ -0,0 +1,307 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/registry.py
+
+import importlib
+import os
+import pickle
+import subprocess
+import sys
+import tempfile
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import (AbstractSet, Callable, Dict, List, NoReturn, Optional,
+                    Tuple, Type, TypeVar, Union, cast)
+
+import cloudpickle
+from torch import nn
+
+from fastvideo.v1.logger import logger
+
+# huggingface class name: (component_name, fastvideo module name, fastvideo class name)
+_TEXT_TO_VIDEO_DIT_MODELS = {
+    "HunyuanVideoTransformer3DModel":
+    ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
+    "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+}
+
+_IMAGE_TO_VIDEO_DIT_MODELS = {
+    # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
+    "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+}
+
+_TEXT_ENCODER_MODELS = {
+    "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"),
+    "LlamaModel": ("encoders", "llama", "LlamaModel"),
+    "UMT5EncoderModel": ("encoders", "t5", "UMT5EncoderModel"),
+}
+
+_IMAGE_ENCODER_MODELS: dict[str, tuple] = {
+    # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"),
+    "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"),
+}
+
+_VAE_MODELS = {
+    "AutoencoderKLHunyuanVideo":
+    ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"),
+    "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"),
+}
+
+_SCHEDULERS = {
+    "FlowMatchEulerDiscreteScheduler":
+    ("schedulers", "scheduling_flow_match_euler_discrete",
+     "FlowMatchDiscreteScheduler"),
+    "UniPCMultistepScheduler":
+    ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"),
+}
+
+_FAST_VIDEO_MODELS = {
+    **_TEXT_TO_VIDEO_DIT_MODELS,
+    **_IMAGE_TO_VIDEO_DIT_MODELS,
+    **_TEXT_ENCODER_MODELS,
+    **_IMAGE_ENCODER_MODELS,
+    **_VAE_MODELS,
+    **_SCHEDULERS,
+}
+
+_SUBPROCESS_COMMAND = [
+    sys.executable, "-m", "fastvideo.v1.models.dits.registry"
+]
+
+_T = TypeVar("_T")
+
+
+@dataclass(frozen=True)
+class _ModelInfo:
+    architecture: str
+
+    @staticmethod
+    def from_model_cls(model: Type[nn.Module]) -> "_ModelInfo":
+        return _ModelInfo(architecture=model.__name__, )
+
+
+class _BaseRegisteredModel(ABC):
+
+    @abstractmethod
+    def inspect_model_cls(self) -> _ModelInfo:
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_model_cls(self) -> Type[nn.Module]:
+        raise NotImplementedError
+
+
+@dataclass(frozen=True)
+class _RegisteredModel(_BaseRegisteredModel):
+    """
+    Represents a model that has already been imported in the main process.
+    """
+
+    interfaces: _ModelInfo
+    model_cls: Type[nn.Module]
+
+    @staticmethod
+    def from_model_cls(model_cls: Type[nn.Module]):
+        return _RegisteredModel(
+            interfaces=_ModelInfo.from_model_cls(model_cls),
+            model_cls=model_cls,
+        )
+
+    def inspect_model_cls(self) -> _ModelInfo:
+        return self.interfaces
+
+    def load_model_cls(self) -> Type[nn.Module]:
+        return self.model_cls
+
+
+def _run_in_subprocess(fn: Callable[[], _T]) -> _T:
+    # NOTE: We use a temporary directory instead of a temporary file to avoid
+    # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file
+    with tempfile.TemporaryDirectory() as tempdir:
+        output_filepath = os.path.join(tempdir, "registry_output.tmp")
+
+        # `cloudpickle` allows pickling lambda functions directly
+        input_bytes = cloudpickle.dumps((fn, output_filepath))
+
+        # cannot use `sys.executable __file__` here because the script
+        # contains relative imports
+        returned = subprocess.run(_SUBPROCESS_COMMAND,
+                                  input=input_bytes,
+                                  capture_output=True)
+
+        # check if the subprocess is successful
+        try:
+            returned.check_returncode()
+        except Exception as e:
+            # wrap raised exception to provide more information
+            raise RuntimeError(f"Error raised in subprocess:\n"
+                               f"{returned.stderr.decode()}") from e
+
+        with open(output_filepath, "rb") as f:
+            return cast(_T, pickle.load(f))
+
+
+@dataclass(frozen=True)
+class _LazyRegisteredModel(_BaseRegisteredModel):
+    """
+    Represents a model that has not been imported in the main process.
+    """
+    module_name: str
+    component_name: str
+    class_name: str
+
+    # Performed in another process to avoid initializing CUDA
+    def inspect_model_cls(self) -> _ModelInfo:
+        return _run_in_subprocess(
+            lambda: _ModelInfo.from_model_cls(self.load_model_cls()))
+
+    def load_model_cls(self) -> Type[nn.Module]:
+        mod = importlib.import_module(self.module_name)
+        return cast(Type[nn.Module], getattr(mod, self.class_name))
+
+
+@lru_cache(maxsize=128)
+def _try_load_model_cls(
+    model_arch: str,
+    model: _BaseRegisteredModel,
+) -> Optional[Type[nn.Module]]:
+    from fastvideo.v1.platforms import current_platform
+    current_platform.verify_model_arch(model_arch)
+    try:
+        return model.load_model_cls()
+    except Exception:
+        logger.exception("Error in loading model architecture '%s'", model_arch)
+        return None
+
+
+@lru_cache(maxsize=128)
+def _try_inspect_model_cls(
+    model_arch: str,
+    model: _BaseRegisteredModel,
+) -> Optional[_ModelInfo]:
+    try:
+        return model.inspect_model_cls()
+    except Exception:
+        logger.exception("Error in inspecting model architecture '%s'",
+                         model_arch)
+        return None
+
+
+@dataclass
+class _ModelRegistry:
+    # Keyed by model_arch
+    models: Dict[str, _BaseRegisteredModel] = field(default_factory=dict)
+
+    def get_supported_archs(self) -> AbstractSet[str]:
+        return self.models.keys()
+
+    def register_model(
+        self,
+        model_arch: str,
+        model_cls: Union[Type[nn.Module], str],
+    ) -> None:
+        """
+        Register an external model to be used in vLLM.
+
+        :code:`model_cls` can be either:
+
+        - A :class:`torch.nn.Module` class directly referencing the model.
+        - A string in the format :code:`<module>:<class>` which can be used to
+          lazily import the model. This is useful to avoid initializing CUDA
+          when importing the model and thus the related error
+          :code:`RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
+        """
+        if model_arch in self.models:
+            logger.warning(
+                "Model architecture %s is already registered, and will be "
+                "overwritten by the new model class %s.", model_arch, model_cls)
+
+        if isinstance(model_cls, str):
+            split_str = model_cls.split(":")
+            if len(split_str) != 2:
+                msg = "Expected a string in the format `<module>:<class>`"
+                raise ValueError(msg)
+
+            model = _LazyRegisteredModel(*split_str)
+        else:
+            model = _RegisteredModel.from_model_cls(model_cls)
+
+        self.models[model_arch] = model
+
+    def _raise_for_unsupported(self, architectures: List[str]) -> NoReturn:
+        all_supported_archs = self.get_supported_archs()
+
+        if any(arch in all_supported_archs for arch in architectures):
+            raise ValueError(
+                f"Model architectures {architectures} failed "
+                "to be inspected. Please check the logs for more details.")
+
+        raise ValueError(
+            f"Model architectures {architectures} are not supported for now. "
+            f"Supported architectures: {all_supported_archs}")
+
+    def _try_load_model_cls(self, model_arch: str) -> Optional[Type[nn.Module]]:
+        if model_arch not in self.models:
+            return None
+
+        return _try_load_model_cls(model_arch, self.models[model_arch])
+
+    def _try_inspect_model_cls(self, model_arch: str) -> Optional[_ModelInfo]:
+        if model_arch not in self.models:
+            return None
+
+        return _try_inspect_model_cls(model_arch, self.models[model_arch])
+
+    def _normalize_archs(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> List[str]:
+        if isinstance(architectures, str):
+            architectures = [architectures]
+        if not architectures:
+            logger.warning("No model architectures are specified")
+
+        normalized_arch = []
+        for model in architectures:
+            if model not in self.models:
+                model = "TransformersModel"
+            normalized_arch.append(model)
+        return normalized_arch
+
+    def inspect_model_cls(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> Tuple[_ModelInfo, str]:
+        architectures = self._normalize_archs(architectures)
+
+        for arch in architectures:
+            model_info = self._try_inspect_model_cls(arch)
+            if model_info is not None:
+                return (model_info, arch)
+
+        return self._raise_for_unsupported(architectures)
+
+    def resolve_model_cls(
+        self,
+        architectures: Union[str, List[str]],
+    ) -> Tuple[Type[nn.Module], str]:
+        architectures = self._normalize_archs(architectures)
+
+        for arch in architectures:
+            model_cls = self._try_load_model_cls(arch)
+            if model_cls is not None:
+                return (model_cls, arch)
+
+        return self._raise_for_unsupported(architectures)
+
+
+ModelRegistry = _ModelRegistry({
+    model_arch:
+    _LazyRegisteredModel(
+        module_name=f"fastvideo.v1.models.{component_name}.{mod_relname}",
+        component_name=component_name,
+        class_name=cls_name,
+    )
+    for model_arch, (component_name, mod_relname,
+                     cls_name) in _FAST_VIDEO_MODELS.items()
+})

fastvideo/v1/models/schedulers/base.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple, Union
+
+import torch
+from diffusers.utils import BaseOutput
+
+
+class BaseScheduler(ABC):
+    timesteps: torch.tensor
+    order: int
+
+    def __init__(self, *args, **kwargs) -> None:
+        # Check if subclass has defined all required properties
+        required_attributes = ['timesteps', 'order']
+
+        for attr in required_attributes:
+            if not hasattr(self, attr):
+                raise AttributeError(
+                    f"Subclasses of BaseScheduler must define '{attr}' property"
+                )
+
+    @abstractmethod
+    def set_shift(self, shift: float) -> None:
+        pass
+
+    @abstractmethod
+    def set_timesteps(self, *args, **kwargs) -> None:
+        pass
+
+    @abstractmethod
+    def scale_model_input(self,
+                          sample: torch.Tensor,
+                          timestep: Optional[int] = None) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> Union[BaseOutput, Tuple]:
+        pass