update template directions

remove v1 folder
move collect env
2025-07-24 21:51:51 -07:00 · 2025-07-24 21:47:57 -07:00 · 2025-07-24 21:33:21 -07:00 · 2025-07-24 21:31:42 -07:00 · 2025-07-24 21:25:20 -07:00 · 2025-07-24 21:22:48 -07:00
423 changed files with 5833 additions and 26167 deletions
@@ -13,12 +13,12 @@ steps:
  - label: "Trigger Tests"
    plugins:
      - monorepo-diff#v1.4.0:
-          diff: "git diff --name-only $BUILDKITE_PULL_REQUEST_BASE_BRANCH...HEAD"
+          diff: 'git fetch origin "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" && git diff --name-only origin/"$BUILDKITE_PULL_REQUEST_BASE_BRANCH"...HEAD'
          watch:
            - path: 
-                - "fastvideo/v1/models/encoders/**"
-                - "fastvideo/v1/models/loader/**"
-                - "fastvideo/v1/tests/encoders/**"
+                - "fastvideo/models/encoders/**"
+                - "fastvideo/models/loader/**"
+                - "fastvideo/tests/encoders/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -29,9 +29,9 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/models/vaes/**"
-                - "fastvideo/v1/models/loader/**"
-                - "fastvideo/v1/tests/vaes/**"
+                - "fastvideo/models/vaes/**"
+                - "fastvideo/models/loader/**"
+                - "fastvideo/tests/vaes/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -42,11 +42,11 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/models/dits/**"
-                - "fastvideo/v1/models/loader/**"
-                - "fastvideo/v1/tests/transformers/**"
-                - "fastvideo/v1/layers/**"
-                - "fastvideo/v1/attention/**"
+                - "fastvideo/models/dits/**"
+                - "fastvideo/models/loader/**"
+                - "fastvideo/tests/transformers/**"
+                - "fastvideo/layers/**"
+                - "fastvideo/attention/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -57,22 +57,22 @@ steps:
                agents:
                  queue: "default"
            - path: 
-                - "fastvideo/v1/**/*.py"
+                - "fastvideo/**/*.py"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
+                command: "timeout 45m .buildkite/scripts/pr_test.sh"
                label: "SSIM Tests"
                env:
                  - TEST_TYPE=ssim
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/tests/lora/**"
-                - "fastvideo/v1/models/loader/**"
-                - "fastvideo/v1/tests/transformers/**"
-                - "fastvideo/v1/pipelines/**"
-                - "fastvideo/v1/layers/lora/**"
+                - "fastvideo/tests/lora/**"
+                - "fastvideo/models/loader/**"
+                - "fastvideo/tests/transformers/**"
+                - "fastvideo/pipelines/**"
+                - "fastvideo/layers/lora/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -83,7 +83,7 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/**"
+                - "fastvideo/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -94,7 +94,18 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/**"
+                - "fastvideo/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "LoRA Training Tests"
+                env:
+                  - TEST_TYPE=training_lora
+                agents:
+                  queue: "default"
+            - path:
+                - "fastvideo/**"
                - "csrc/attn/vsa/**"
                - "csrc/attn/tk/**"
                - "csrc/attn/setup_vsa.py"
@@ -110,7 +121,7 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "fastvideo/v1/**"
+                - "fastvideo/**"
                - "csrc/attn/st_attn/**"
                - "csrc/attn/setup_sta.py"
                - "csrc/attn/config_sta.py"
@@ -50,7 +50,7 @@ else
    exit 1
 fi

-MODAL_TEST_FILE="fastvideo/v1/tests/modal/pr_test.py"
+MODAL_TEST_FILE="fastvideo/tests/modal/pr_test.py"

 if [ -z "${TEST_TYPE:-}" ]; then
    log "Error: TEST_TYPE environment variable is not set"
@@ -81,6 +81,10 @@ case "$TEST_TYPE" in
        log "Running training tests..."
        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests"
        ;;
+    "training_lora")
+        log "Running LoRA training tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_lora_tests"
+        ;;
    "training_vsa")
        log "Running training VSA tests..."
        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests_VSA"
@@ -23,7 +23,7 @@ body:
  attributes:
    label: Environment
    description: |
-      Please share your environment with us. You can run the command **python fastvideo/utils/collect_env.py** and copy-paste its output below.
+      Please share your environment with us. You can run the command **python collect_env.py** and copy-paste its output below.
    placeholder: FastVideo version, platform, python version, cuda version...
  validations:
    required: true
@@ -0,0 +1,56 @@
+name: 💬 Request for comments (RFC).
+description: Ask for feedback on major architectural changes or design choices.
+title: "[RFC]: "
+labels: ["RFC"]
+
+body:
+- type: markdown
+  attributes:
+    value: >
+      #### Please take a look at previous [RFCs](https://github.com/hao-ai-lab/FastVideo/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
+- type: textarea
+  attributes:
+    label: Motivation.
+    description: >
+      The motivation of the RFC.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Proposed Change.
+    description: >
+      The proposed change of the RFC.
+  validations:
+    required: true
+- type: textarea
+  attributes:
+    label: Feedback Period.
+    description: >
+      The feedback period of the RFC. Usually at least one week.
+  validations:
+    required: false
+- type: textarea
+  attributes:
+    label: CC List.
+    description: >
+      The list of people you want to CC.
+  validations:
+    required: false
+- type: textarea
+  attributes:
+    label: Any Other Things.
+    description: >
+      Any other things you would like to mention.
+  validations:
+    required: false
+- type: markdown
+  attributes:
+    value: >
+      Thanks for contributing 🎉!
+- type: checkboxes
+  id: askllm
+  attributes:
+    label: Before submitting a new issue...
+    options:
+      - label: Make sure you already searched for relevant issues.
+        required: true
@@ -8,14 +8,14 @@ on:
      - main
    paths:
      - "docs/**/*.md"
-      - "fastvideo/v1/examples/**/*.py"
+      - "fastvideo/examples/**/*.py"
  pull_request:
    branches: 
      - main
    types: [opened, ready_for_review, synchronize, reopened]
    paths:
      - "docs/**/*.md"
-      - "fastvideo/v1/examples/**/*.py"
+      - "fastvideo/examples/**/*.py"

  # Allows you to run this workflow manually from the Actions tab
  workflow_dispatch:
@@ -115,37 +115,37 @@ jobs:
              - 'csrc/attn/config_vsa.py'
              - 'csrc/attn/vsa.cpp'
            vsa-paths: &vsa-paths
-              - 'fastvideo/v1/**'
+              - 'fastvideo/**'
              - *common-paths
              - *vsa-kernel-paths
            
            # Actual tests
            encoder-test:
-              - 'fastvideo/v1/models/encoders/**'
-              - 'fastvideo/v1/models/loader/**'
-              - 'fastvideo/v1/tests/encoders/**'
+              - 'fastvideo/models/encoders/**'
+              - 'fastvideo/models/loader/**'
+              - 'fastvideo/tests/encoders/**'
              - *common-paths
            vae-test:
-              - 'fastvideo/v1/models/vaes/**'
-              - 'fastvideo/v1/models/loader/**'
-              - 'fastvideo/v1/tests/vaes/**'
+              - 'fastvideo/models/vaes/**'
+              - 'fastvideo/models/loader/**'
+              - 'fastvideo/tests/vaes/**'
              - *common-paths
            transformer-test:
-              - 'fastvideo/v1/models/dits/**'
-              - 'fastvideo/v1/models/loader/**'
-              - 'fastvideo/v1/tests/transformers/**'
-              - 'fastvideo/v1/layers/**'
-              - 'fastvideo/v1/attention/**'
+              - 'fastvideo/models/dits/**'
+              - 'fastvideo/models/loader/**'
+              - 'fastvideo/tests/transformers/**'
+              - 'fastvideo/layers/**'
+              - 'fastvideo/attention/**'
              - *common-paths
            training-test:
-              - 'fastvideo/v1/**'
+              - 'fastvideo/**'
              - *common-paths
            training-test-VSA:
-              - 'fastvideo/v1/**'
+              - 'fastvideo/**'
              - *common-paths
              - *vsa-kernel-paths
            inference-test-STA:
-              - 'fastvideo/v1/**'
+              - 'fastvideo/**'
              - *common-paths
              - *sta-kernel-paths
            precision-test-STA:
@@ -167,7 +167,7 @@ jobs:
      gpu_count: 1
      volume_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/encoders -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -185,7 +185,7 @@ jobs:
      gpu_count: 1
      volume_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/vaes -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -203,7 +203,7 @@ jobs:
      gpu_count: 1
      volume_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/transformers -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -229,7 +229,7 @@ jobs:
      volume_size: 200
      disk_size: 200
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/ssim -vs"
      timeout_minutes: 60
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -248,7 +248,7 @@ jobs:
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/Vanilla -srP"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/Vanilla -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -268,7 +268,7 @@ jobs:
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/VSA -srP"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/VSA -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -288,7 +288,7 @@ jobs:
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/inference/STA -srP"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/inference/STA -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -343,7 +343,7 @@ jobs:
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/nightly/test_e2e_overfit_single_sample.py -vs"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -20,6 +20,7 @@ samples/
 data/
 outputs/
 outputs_video
+checkpoints/
 sbatch.sh
 *.out
 env
@@ -40,6 +41,7 @@ eggs/
 docs/_build/
 docs/source/getting_started/examples/
 docs/source/inference/examples/
+docs/source/training/examples/

 # VSCode
 .vscode/
@@ -55,7 +57,7 @@ docs/source/inference/examples/
 *.pkl

 # Reference videos
-!fastvideo/v1/tests/ssim/reference_videos/**/*.mp4
+!fastvideo/tests/ssim/reference_videos/**/*.mp4

 # Static images
 !docs/source/_static/images/**/*.png
@@ -3,7 +3,7 @@ default_stages:
  - manual # Run in CI
 exclude: |
  (?x)(
-    fastvideo/v1/third_party/.*|
+    fastvideo/third_party/.*|
    csrc/.*|
    assets/.*|
    tests/.*|
@@ -69,7 +69,7 @@ repos:
    entry: bash
    args:
      - -c
-      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep -v "^fastvideo/v1/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
+      - 'git ls-files | grep -v "^fastvideo/tests/ssim/" | grep -v "^fastvideo/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
    language: system
    always_run: true
    pass_filenames: false
@@ -1,24 +0,0 @@
-# Configuration for Cog ⚙️
-# Reference: https://cog.run/yaml
-
-build:
-  gpu: true
-  cuda: "12.1"
-  python_version: "3.10"
-  python_packages:
-    - "torch==2.4.0"
-    - "torchvision"
-    - "ninja==1.11.1.3"
-    - "transformers==4.46.1"
-    - "git+https://github.com/huggingface/diffusers.git@bf64b32652a63a1865a0528a73a13652b201698b"
-    - "accelerate==1.0.1"
-    - "safetensors==0.4.5"
-    - "peft==0.13.2"
-    - "packaging==24.2"
-    - "git+https://github.com/hao-ai-lab/FastVideo"
-
-  run:
-    - FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation
-    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/latest/download/pget_$(uname -s)_$(uname -m)" && chmod +x /usr/local/bin/pget
-
-predict: "predict.py:Predictor"
@@ -16,7 +16,7 @@ import sys
 # Run it with `python collect_env.py` or `python -m torch.utils.collect_env`
 from collections import namedtuple

-from fastvideo.v1.envs import environment_variables
+from fastvideo.envs import environment_variables

 try:
    import torch
@@ -81,6 +81,7 @@ DEFAULT_CONDA_PATTERNS = {
 DEFAULT_PIP_PATTERNS = {
    "torch",
    "numpy",
+    "mypy",
    "flake8",
    "triton",
    "optree",
@@ -74,8 +74,8 @@ After installation, the following nodes will be available in the ComfyUI interfa
 You may have noticed many arguments on the nodes have 'auto' as the default value. This is because FastVideo will automatically detect the best values for these parameters based on the model and the hardware. However, you can also manually configure these parameters to get the best performance for your specific use case. We plan on releasing more optimized workflow files for different models and hardware configurations in the future.

 You can see what some of the default configurations are by looking at the FastVideo repo:
- [Wan2.1-I2V-14B-480P-Diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json)
- [FastHunyuan-diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/v1/configs/fasthunyuan_t2v.json)
+- [Wan2.1-I2V-14B-480P-Diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/configs/wan_14B_i2v_480p_pipeline.json)
+- [FastHunyuan-diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/configs/fasthunyuan_t2v.json)

 ### Node Configuration

@@ -1,11 +1,23 @@


-# Sliding Tile Atteniton Kernel
+# Attention Kernel Used in FastVideo

+## Sliding Tile Attention (STA)
+We only support H100 for STA.
+```bash
+git submodule update --init --recursive
+python setup_sta.py install
+```

-## Installation
-We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only support H100/H200, because ThunderKittens uses TMA but doesn't support Blackwell yet.
-First, install C++20 for ThunderKittens:
+## Video Sparse Attention (VSA)
+We support H100 (via TK) and RTX 4090 (via triton) for VSA.
+```bash
+git submodule update --init --recursive
+python setup_vsa.py install
+```
+
+If you encounter error during installation, try below:
+Install C++20 for ThunderKittens:
 ```bash
 sudo apt update
 sudo apt install gcc-11 g++-11
@@ -15,27 +27,23 @@ sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave
 sudo apt update
 sudo apt install clang-11
 ```
-
-## Environment Setup
-First, set up your CUDA environment:
+(If you use CUDA12.4)
 ```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
 ```

-## Install Sliding Tile Attention (STA)
-```bash
-python setup_sta.py install
-```

-## Install Video Sparse Attention (VSA)
-```bash
-python setup_vsa.py install
-```

 ## Usage
+### STA
+End-2-end inference with FastVideo:
+```bash
+bash scripts/inference/v1_inference_wan_STA.sh
+```
+
+If you want to use sliding tile attention in your custom model:
 ```python
 from st_attn import sliding_tile_attention
 # assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
@@ -47,9 +55,11 @@ from st_attn import sliding_tile_attention
 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)
-
 ```

+### VSA
+We do not officially supoort end-2-end inference with VSA in FastVideo yet. Stay tuned.
+

 ## Test
 ```bash
@@ -1,7 +1,7 @@
 import torch
 import argparse
 from flash_attn.utils.benchmark import benchmark_forward
-from vsa import block_sparse_attention_fwd, block_sparse_attention_backward
+from vsa import block_sparse_fwd, block_sparse_bwd
 from vsa import BLOCK_M, BLOCK_N

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

-from csrc.attn.config_vsa import kernels, sources, target
+from config_vsa import kernels, sources, target
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension

@@ -51,21 +51,29 @@ for k in kernels:
        source_files.append(sources[k]['source_files'][target])
    cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}')

+ext_modules = []
+import torch
+major, minor = torch.cuda.get_device_capability(0)
+if major == 9 and minor == 0:# check if H100
+    ext_modules = [
+        CUDAExtension('vsa_cuda',
+            sources=source_files,
+            extra_compile_args={
+                'cxx': cpp_flags,
+                'nvcc': cuda_flags
+            },
+            libraries=['cuda'])
+    ]
+
+
+
 setup(name=PACKAGE_NAME,
      version=VERSION,
      author=AUTHOR,
      description=DESCRIPTION,
      url=URL,
      packages=find_packages(),
-      ext_modules=[
-          CUDAExtension('vsa_cuda',
-                        sources=source_files,
-                        extra_compile_args={
-                            'cxx': cpp_flags,
-                            'nvcc': cuda_flags
-                        },
-                        libraries=['cuda'])
-      ],
+      ext_modules=ext_modules,
      cmdclass={'build_ext': BuildExtension},
      classifiers=[
          "Programming Language :: Python :: 3",
@@ -2,7 +2,7 @@ import torch
 import argparse
 from flash_attn.utils.benchmark import benchmark_forward
 from flash_attn import flash_attn_func
-from vsa import block_sparse_attention_fwd, block_sparse_attention_backward, BlockSparseAttentionFunction
+from vsa import block_sparse_attn
 from vsa import BLOCK_M, BLOCK_N

 import numpy as np
@@ -188,7 +188,7 @@ def main(args):


            # testing forward
-            o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
+            o = block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
            del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded
            grad_o = torch.randn_like(o)
            o.backward(grad_o)
@@ -215,7 +215,7 @@ def main(args):
            forward_metrics['l1'].append(l1)
            forward_metrics['rmse'].append(rmse)

-            print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            print(f"block_sparse_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")

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

+from vsa import triton_attention
+
 def pytorch_test(Q, K, V, dO):
    q_ = Q.to(torch.float64).requires_grad_()
    k_ = K.to(torch.float64).requires_grad_()
    v_ = V.to(torch.float64).requires_grad_()
+    q_.grad = None
+    k_.grad = None
+    v_.grad = None
    dO_ = dO.to(torch.float64)
    
    # manual pytorch implementation of scaled dot product attention
@@ -30,11 +35,33 @@ def fa2_test(Q, K, V, dO):
    Q.requires_grad = True
    K.requires_grad = True
    V.requires_grad = True
+    Q.grad = None
+    K.grad = None
+    V.grad = None
+    
    output = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=False)
    output.backward(dO)
    
    return output, Q.grad, K.grad, V.grad

+def triton_test(Q, K, V, dO):
+    Q.requires_grad = True
+    K.requires_grad = True
+    V.requires_grad = True
+    Q.grad = None
+    K.grad = None
+    V.grad = None
+
+    output = triton_attention(Q, K, V)
+
+    output.backward(dO)
+    
+    q_grad = Q.grad
+    k_grad = K.grad
+    v_grad = V.grad
+    
+    return output.to(Q.dtype) if output is not None else None, q_grad, k_grad, v_grad
+
 def generate_tensor(shape, mean, std, dtype, device):
    tensor = torch.randn(shape, dtype=dtype, device=device)

@@ -45,7 +72,8 @@ def generate_tensor(shape, mean, std, dtype, device):

 def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all', test_mode='forward_backward'):
    results = {
-        'FA2 vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0},
+        'FA2 vs PT': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
+        'Triton vs PT': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
    }

    for _ in range(num_iterations):
@@ -58,21 +86,31 @@ def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all
        
        pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, dO)
        fa2_o, fa2_qg, fa2_kg, fa2_vg = fa2_test(Q, K, V, dO)
+        triton_o, triton_qg, triton_kg, triton_vg = triton_test(Q, K, V, dO)
        
        if test_mode == 'forward_only':
            tensors_fa2_pt = [(pt_o, fa2_o)]
+            tensors_triton_pt = [(pt_o, triton_o)]
        else:  # 'forward_backward'
            if error_mode == 'output':
                tensors_fa2_pt = [(pt_o, fa2_o)]
+                tensors_triton_pt = [(pt_o, triton_o)]
            elif error_mode == 'backward':
                tensors_fa2_pt = [(pt_qg, fa2_qg),
                               (pt_kg, fa2_kg),
                               (pt_vg, fa2_vg)]
+                tensors_triton_pt = [(pt_qg, triton_qg),
+                                 (pt_kg, triton_kg),
+                                 (pt_vg, triton_vg)]
            else:  # 'all'
                tensors_fa2_pt = [(pt_o, fa2_o),
                               (pt_qg, fa2_qg),
                               (pt_kg, fa2_kg),
                               (pt_vg, fa2_vg)]
+                tensors_triton_pt = [(pt_o, triton_o),
+                                 (pt_qg, triton_qg),
+                                 (pt_kg, triton_kg),
+                                 (pt_vg, triton_vg)]
                
        for pt, fa2 in tensors_fa2_pt:
            diff = pt - fa2
@@ -80,6 +118,13 @@ def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all
            results['FA2 vs PT']['sum_diff'] += torch.sum(abs_diff).item()
            results['FA2 vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
            results['FA2 vs PT']['max_diff'] = max(results['FA2 vs PT']['max_diff'], torch.max(abs_diff).item())
+        
+        for pt, triton in tensors_triton_pt:
+            diff = pt - triton
+            abs_diff = torch.abs(diff)
+            results['Triton vs PT']['sum_diff'] += torch.sum(abs_diff).item()
+            results['Triton vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
+            results['Triton vs PT']['max_diff'] = max(results['Triton vs PT']['max_diff'], torch.max(abs_diff).item())
                
        torch.cuda.empty_cache()

@@ -99,38 +144,40 @@ def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all
 def generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward'):
    seq_lengths = [768 * (2**i) for i in range(1)]

-    print(f"\n{'='*80}")
+    print(f"\n{'='*100}")
    print(f"ATTENTION ERROR COMPARISON TABLE (b={b}, h={h}, d={d}, mean={mean}, std={std})")
    print(f"Mode: {error_mode}, Test: {test_mode}")
-    print(f"{'='*80}")
+    print(f"{'='*100}")
    
    # Print header
-    print(f"{'Seq Length':<12} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15}")
-    print(f"{'-'*12} | {'-'*15} | {'-'*15}")
+    print(f"{'Seq Length':<12} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15} | {'Triton vs PT Avg':<15} | {'Triton vs PT Max':<15}")
+    print(f"{'-'*12} | {'-'*15} | {'-'*15} | {'-'*15} | {'-'*15}")
    
    for n in seq_lengths:
        results = check_correctness(b, h, n, d, mean, std, error_mode=error_mode, test_mode=test_mode)
        
        fa2_pt_avg = results['FA2 vs PT']['avg_diff']
        fa2_pt_max = results['FA2 vs PT']['max_diff']
+        triton_pt_avg = results['Triton vs PT']['avg_diff']
+        triton_pt_max = results['Triton vs PT']['max_diff']
        
-        # Print row
-        print(f"{n:<12} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e}")
+        # Print row with both comparisons
+        print(f"{n:<12} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e} | {triton_pt_avg:<15.6e} | {triton_pt_max:<15.6e}")
    
-    print(f"{'='*80}\n")
+    print(f"{'='*100}\n")

 # fix random seed
 torch.manual_seed(0)

-# Example usage
-b, h, d = 2, 2, 64
 mean = 1e-1
 std = 10
+configs = [
+    (4, 1, 128),  # Larger batch, single head, larger dim
+    (2, 8, 64),   # Medium batch, many heads, medium dim
+]

-# Test forward only
-generate_error_tables(b, h, d, mean, std, error_mode='output', test_mode='forward_only')
-
-# Test forward and backward
-generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward')
+for b, h, d in configs:
+    print(f"\nConfiguration: batch={b}, heads={h}, dim={d}")
+    generate_error_tables(b, h, d, mean, std, error_mode='backward', test_mode='forward_backward')

 print("Attention error comparison completed.")
@@ -1,18 +1,27 @@
-import math
-
 import torch
-from torch.utils.checkpoint import detach_variable
 from typing import Tuple
+from vsa.vsa import block_sparse_attn
 try:
    from vsa_cuda import block_sparse_fwd, block_sparse_bwd
 except ImportError:
    block_sparse_fwd = None
    block_sparse_bwd = None
-
-
+from vsa.block_sparse_attn_triton import attention as triton_attention, attention_sparse as triton_attention_sparse
 BLOCK_M = 64
 BLOCK_N = 64

+
+def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
+    QK = torch.matmul(q, k.transpose(-2, -1))
+    QK /= (q.size(-1)**0.5)
+
+    # Causal mask removed since causal is always false
+
+    QK = torch.nn.functional.softmax(QK, dim=-1)
+    output = torch.matmul(QK, v)
+    return output, QK
+
+
 def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None):
    """
    q: [batch_size, num_heads, seq_len, head_dim]
@@ -67,15 +76,6 @@ def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None):
        final_output = output_compress + output_select
    return final_output

-def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
-    QK = torch.matmul(q, k.transpose(-2, -1))
-    QK /= (q.size(-1)**0.5)
-
-    # Causal mask removed since causal is always false
-
-    QK = torch.nn.functional.softmax(QK, dim=-1)
-    output = torch.matmul(QK, v)
-    return output, QK

 def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor,
                                       topk: int):
@@ -123,107 +123,11 @@ def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor,

    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num

-@torch._dynamo.disable
-def block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
-    """
-    Differentiable block sparse attention function.
-    
-    Args:
-        q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
-        k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
-        v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
-        q2k_block_sparse_index: Indices for query-to-key sparse blocks
-        q2k_block_sparse_num: Number of sparse blocks for each query block
-        k2q_block_sparse_index: Indices for key-to-query sparse blocks (for backward pass)
-        k2q_block_sparse_num: Number of sparse blocks for each key block (for backward pass)
-    
-    Returns:
-        output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
-    """
-    return BlockSparseAttentionFunction.apply(
-        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
-    )
-
-def block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num):
-    """
-    block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]. 
-        [*, *, i, j] = 1 means the i-th q block should attend to the j-th kv block.
-    """
-    # assert all elements in q2k_block_sparse_num can be devisible by 2
-    o, lse = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
-    return o, lse
-
-def block_sparse_attention_backward(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num):
-    grad_output = grad_output.contiguous()
-    grad_q, grad_k, grad_v = block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
-    return grad_q, grad_k, grad_v

 ## pytorch sdpa version of block sparse ##
 import triton
 import triton.language as tl

-@triton.jit
-def index_to_mask_kernel(
-    q2k_block_sparse_index_ptr,
-    q2k_block_sparse_num_ptr,
-    mask_ptr,
-    batch_size: tl.constexpr,
-    num_heads: tl.constexpr,
-    num_q_blocks: tl.constexpr,
-    num_k_blocks: tl.constexpr,
-    max_kv_blocks: tl.constexpr,
-    BLOCK_Q: tl.constexpr,
-    BLOCK_K: tl.constexpr,
-):
-    bh, q, id = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64)
-    b = bh // num_heads
-    h = bh % num_heads
-
-    num_valid_blocks = tl.load(q2k_block_sparse_num_ptr + b * num_heads * num_q_blocks + h * num_q_blocks + q)
-
-    if num_valid_blocks <= id:
-        return
-    k = tl.load(q2k_block_sparse_index_ptr + b * num_heads * num_q_blocks * max_kv_blocks + h * num_q_blocks * max_kv_blocks + q * max_kv_blocks + id)
-
-    full_mask = (tl.arange(0, BLOCK_Q)[:, None] < BLOCK_Q) & (tl.arange(0, BLOCK_K)[None, :] < BLOCK_K)
-
-    q_lengths = num_q_blocks * BLOCK_Q
-    k_lengths = num_k_blocks * BLOCK_K
-    mask_ptr_base = mask_ptr + b * num_heads * q_lengths * k_lengths + h * q_lengths * k_lengths + q * BLOCK_Q * k_lengths + k * BLOCK_K
-
-    tl.store(mask_ptr_base + tl.arange(0, BLOCK_Q)[:, None] * k_lengths + tl.arange(0, BLOCK_K)[None, :], full_mask)
-
-def index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, BLOCK_Q, BLOCK_K, num_k_blocks):
-    """
-    Convert block sparse indices to a mask.
-    
-    Args:
-        q2k_block_sparse_index: Indices for query-to-key sparse blocks
-        q2k_block_sparse_num: Number of sparse blocks for each query block
-    
-    Returns:
-        mask: Block sparse mask tensor
-    """
-    batch_size, num_heads, num_q_blocks, max_kv_blocks = q2k_block_sparse_index.shape
-    assert q2k_block_sparse_num.shape == (batch_size, num_heads, num_q_blocks)
-
-    mask = torch.zeros((batch_size, num_heads, num_q_blocks * BLOCK_Q, num_k_blocks * BLOCK_K), dtype=torch.bool, device=q2k_block_sparse_index.device)
-
-    grid = (batch_size * num_heads, num_q_blocks, max_kv_blocks)
-    index_to_mask_kernel[grid](
-        q2k_block_sparse_index,
-        q2k_block_sparse_num,
-        mask,
-        batch_size,
-        num_heads,
-        num_q_blocks,
-        num_k_blocks,
-        max_kv_blocks,
-        BLOCK_Q=BLOCK_Q,
-        BLOCK_K=BLOCK_K,
-    )
-
-    return mask

@triton.jit
 def topk_index_to_map_kernel(
@@ -372,99 +276,3 @@ def map_to_index(block_map: torch.Tensor):

    return index, index_num

-class BlockSparseAttentionFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
-        o, lse = block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
-        ctx.save_for_backward(q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num)
-        return o
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num = ctx.saved_tensors
-        grad_q, grad_k, grad_v = block_sparse_attention_backward(
-            q, k, v, o, lse, grad_output, k2q_block_sparse_index, k2q_block_sparse_num
-        )
-        return grad_q, grad_k, grad_v, None, None, None, None
-
-
-class DummyOperator(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x):
-        return x
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        return grad_output
-
-class CheckpointSDPA(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, obj, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k):
-        """Forward pass."""
-        with torch.no_grad():
-            mask = index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k, k.shape[2] // block_k)
-            outputs = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-        ctx.save_for_backward(*detach_variable((q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)))
-        ctx.block_q = block_q
-        ctx.block_k = block_k
-        # the obj is passed in, then it can access the saved input
-        # tensors later for recomputation
-        obj.ctx = ctx
-        return outputs
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        """Backward pass."""
-        inputs = ctx.saved_tensors
-        output = ctx.output
-        torch.autograd.backward(output, grad_output)
-        ctx.output = None
-        grads = tuple(inp.grad for inp in inputs)
-        return (None, ) + grads + (None, None)
-
-
-class BlockSparseAttnTorch:
-    def __init__(self):
-        self.ctx = None
-    
-    def recompute_mask(self, _):
-        recomputed_mask = index_to_mask(self.q2k_block_sparse_index, self.q2k_block_sparse_num, self.block_q, self.block_k, self.num_kv_blocks)
-        mask_size = recomputed_mask.untyped_storage().size()
-        self.mask.untyped_storage().resize_(mask_size)
-        self.mask.untyped_storage().copy_(recomputed_mask.untyped_storage()) 
-    
-    def recompute(self, _):
-        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num = self.ctx.saved_tensors
-        block_q = self.ctx.block_q
-        block_k = self.ctx.block_k
-        mask = index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k, k.shape[2] // block_k)
-        with torch.enable_grad():
-            output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-        self.ctx.output = output
-        self.ctx = None
-
-    @torch._dynamo.disable
-    def forward(self, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k):
-        """
-        Differentiable block sparse attention function using PyTorch.
-        
-        Args:
-            q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
-            k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
-            v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
-            q2k_block_sparse_index: Indices for query-to-key sparse blocks
-            q2k_block_sparse_num: Number of sparse blocks for each query block
-            block_q: Block size for query
-            block_k: Block size for key-value
-
-        Returns:
-            output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
-        """
-
-        output = CheckpointSDPA.apply(
-            self, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k
-        )
-
-        o = DummyOperator.apply(output)
-        o.register_hook(self.recompute)
-        return o
@@ -0,0 +1,707 @@
+"""
+Fused Attention
+===============
+
+This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
+(https://tridao.me/publications/flash2/flash2.pdf)
+
+Credits: OpenAI kernel team
+"""
+
+import pytest
+import torch
+import triton
+import triton.language as tl
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+import math  # small utility needed by the sparse wrapper
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+@triton.jit
+def _attn_fwd_inner(acc, l_i, m_i, q,  #
+                    K_block_ptr, V_block_ptr,  #
+                    start_m, qk_scale,  #
+                    BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr,  #
+                    STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr,  #
+                    N_CTX: tl.constexpr, fp8_v: tl.constexpr):
+
+    # loop over k, v and update accumulator
+    for start_n in range(0, N_CTX, BLOCK_N):
+        # -- compute qk ----
+        k = tl.load(K_block_ptr)
+        qk = tl.dot(q, k)
+
+        m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
+        qk = qk * qk_scale - m_ij[:, None]
+        p = tl.math.exp2(qk)
+        l_ij = tl.sum(p, 1)
+        # -- update m_i and l_i
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        # -- update output accumulator --
+        acc = acc * alpha[:, None]
+        # update acc
+        v = tl.load(V_block_ptr)
+        if fp8_v:
+            p = p.to(tl.float8e5)
+        else:
+            p = p.to(tl.bfloat16)
+        acc = tl.dot(p, v, acc)
+        # update m_i and l_i
+        m_i = m_ij
+        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+    return acc, l_i, m_i
+
+
+# We don't run auto-tuning every time to keep the tutorial fast. Keeping
+# the code below and commenting out the equivalent parameters is convenient for
+# re-tuning.
+configs = [
+    triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+    for BM in [64]\
+    for BN in [64]\
+    for s in [3, 4, 7]\
+    for w in [4, 8]\
+]
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
+@triton.jit
+def _attn_fwd_sparse(Q, K, V, sm_scale,                         #
+                     q2k_index, q2k_num, max_kv_blks,           #
+                     M, Out,                                    #
+                     stride_qz, stride_qh, stride_qm, stride_qk,
+                     stride_kz, stride_kh, stride_kn, stride_kk,
+                     stride_vz, stride_vh, stride_vk, stride_vn,
+                     stride_oz, stride_oh, stride_om, stride_on,
+                     Z, H, N_CTX,                               #
+                     HEAD_DIM: tl.constexpr,                    #
+                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
+                     STAGE: tl.constexpr):
+    """
+    64×64 **block-sparse** forward kernel. Back-prop kernels remain dense
+    (32×64 and 64×32) – memory footprint unchanged.
+    """
+
+    # ----- program-id mapping -----
+    q_blk   = tl.program_id(0)          # Q-tile index
+    off_hz  = tl.program_id(1)          # fused (batch, head)
+    b       = off_hz // H
+    h       = off_hz %  H
+    q_tiles = N_CTX // BLOCK_M
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
+    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
+
+    # ----- base pointers -----
+    qvk_off = (b.to(tl.int64) * stride_qz +
+               h.to(tl.int64) * stride_qh)
+
+    Q_ptr = tl.make_block_ptr(
+        base=Q + qvk_off, shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(q_blk * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
+
+    K_base = tl.make_block_ptr(
+        base=K + qvk_off, shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N), order=(0, 1))
+
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
+    V_base = tl.make_block_ptr(
+        base=V + qvk_off, shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, HEAD_DIM), order=v_order)
+
+    O_ptr = tl.make_block_ptr(
+        base=Out + qvk_off, shape=(N_CTX, HEAD_DIM),
+        strides=(stride_om, stride_on),
+        offsets=(q_blk * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
+
+    # ----- accumulators -----
+    offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M)
+    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    qk_scale = sm_scale * 1.44269504  # 1/ln2
+    q = tl.load(Q_ptr)
+
+    # ----- sparse loop over valid K/V tiles -----
+    for i in range(0, kv_blocks):
+        kv_idx = tl.load(kv_ptr + i).to(tl.int32)
+
+        K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N))
+        V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0))
+
+        k = tl.load(K_ptr)
+        qk = tl.dot(q, k)
+
+        m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
+        p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
+        l_ij = tl.sum(p, 1)
+
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        acc = acc * alpha[:, None]
+
+        v = tl.load(V_ptr)
+        acc = tl.dot(p.to(tl.bfloat16), v, acc)
+        m_i = m_ij
+
+    # ----- epilogue -----
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    tl.store(M + off_hz * N_CTX + offs_m, m_i)
+    tl.store(O_ptr, acc.to(Out.type.element_ty))
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
+@triton.jit
+def _attn_fwd(Q, K, V, sm_scale, M, Out,  #
+              stride_qz, stride_qh, stride_qm, stride_qk,  #
+              stride_kz, stride_kh, stride_kn, stride_kk,  #
+              stride_vz, stride_vh, stride_vk, stride_vn,  #
+              stride_oz, stride_oh, stride_om, stride_on,  #
+              Z, H, N_CTX,  #
+              HEAD_DIM: tl.constexpr,  #
+              BLOCK_M: tl.constexpr,  #
+              BLOCK_N: tl.constexpr,  #
+              STAGE: tl.constexpr  #
+              ):
+    tl.static_assert(BLOCK_N <= HEAD_DIM)
+    start_m = tl.program_id(0)
+    off_hz = tl.program_id(1)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+
+    # block pointers
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
+    V_block_ptr = tl.make_block_ptr(
+        base=V + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, HEAD_DIM),
+        order=v_order,
+    )
+    K_block_ptr = tl.make_block_ptr(
+        base=K + qvk_offset,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N),
+        order=(0, 1),
+    )
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    # load scales
+    qk_scale = sm_scale
+    qk_scale *= 1.44269504  # 1/log(2)
+    # load q: it will stay in SRAM throughout
+    q = tl.load(Q_block_ptr)
+    acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr,  #
+                                    start_m, qk_scale,  #
+                                    BLOCK_M, HEAD_DIM, BLOCK_N,  #
+                                    3, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5  #
+                                    )
+
+    # epilogue
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    m_ptrs = M + off_hz * N_CTX + offs_m
+    tl.store(m_ptrs, m_i)
+    tl.store(O_block_ptr, acc.to(Out.type.element_ty))
+
+
+@triton.jit
+def _attn_bwd_preprocess(O, DO,  #
+                         Delta,  #
+                         Z, H, N_CTX,  #
+                         BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr  #
+                         ):
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_hz = tl.program_id(1)
+    off_n = tl.arange(0, HEAD_DIM)
+    # load
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    # write-back
+    tl.store(Delta + off_hz * N_CTX + off_m, delta)
+
+
+# The main inner-loop logic for computing dK and dV.
+@triton.jit
+def _attn_bwd_dkdv(dk, dv,  #
+                   Q, k, v, sm_scale,  #
+                   DO,  #
+                   M, D,  #
+                   k2q_index, k2q_num, max_q_blks,
+                   # shared by Q/K/V/DO.
+                   stride_tok, stride_d,  #
+                   H, N_CTX, BLOCK_M1: tl.constexpr,  #
+                   BLOCK_N1: tl.constexpr,  #
+                   HEAD_DIM: tl.constexpr,  #
+                   # Filled in by the wrapper.
+                   start_n, start_m, num_steps):
+    offs_m = start_m + tl.arange(0, BLOCK_M1)
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+    offs_k = tl.arange(0, HEAD_DIM)
+    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
+    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
+    step_m = BLOCK_M1
+    kv_blk   = tl.program_id(0)          # Q-tile index
+    off_hz  = tl.program_id(2)          # fused (batch, head)
+    b       = off_hz // H
+    h       = off_hz %  H
+    q_tiles = N_CTX // BLOCK_N1
+    meta_base = ((b * H + h) * q_tiles + kv_blk)
+
+    q_blocks = tl.load(k2q_num  + meta_base)                 # int32
+    q_ptr    = k2q_index + meta_base * max_q_blks           # ptr to list
+    
+        
+    for blk_idx in range(q_blocks*2):
+        block_sparse_offset = (tl.load(q_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_m
+        qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok)
+        # Load m before computing qk to reduce pipeline stall.
+        offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
+        m = tl.load(M + offs_m)
+        qkT = tl.dot(k, qT)
+        pT = tl.math.exp2(qkT - m[None, :])
+        do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
+        # Compute dV.
+        ppT = pT
+        ppT = ppT.to(tl.bfloat16)
+        dv += tl.dot(ppT, do)
+        # D (= delta) is pre-divided by ds_scale.
+        Di = tl.load(D + offs_m)
+        # Compute dP and dS.
+        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
+        dsT = pT * (dpT - Di[None, :])
+        dsT = dsT.to(tl.bfloat16)
+        dk += tl.dot(dsT, tl.trans(qT))
+        # Increment pointers.
+    return dk, dv
+
+
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _attn_bwd_dq(dq, q, K, V,  #
+                 do, m, D,
+                 # shared by Q/K/V/DO.
+                 q2k_index, q2k_num, max_kv_blks,
+                 stride_tok, stride_d,  #
+                 H, N_CTX,  #
+                 BLOCK_M2: tl.constexpr,  #
+                 BLOCK_N2: tl.constexpr,  #
+                 HEAD_DIM: tl.constexpr,
+                 # Filled in by the wrapper.
+                 start_m, start_n, num_steps):
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+    offs_n = start_n + tl.arange(0, BLOCK_N2)
+    offs_k = tl.arange(0, HEAD_DIM)
+    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    # D (= delta) is pre-divided by ds_scale.
+    Di = tl.load(D + offs_m)
+    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
+    step_n = BLOCK_N2
+    
+    q_blk   = tl.program_id(0)          # Q-tile index
+    off_hz  = tl.program_id(2)          # fused (batch, head)
+    b       = off_hz // H
+    h       = off_hz %  H
+    q_tiles = N_CTX // BLOCK_M2
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
+    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
+    
+    
+    for blk_idx in range(kv_blocks*2):
+        block_sparse_offset = (tl.load(kv_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_n * stride_tok
+        kT = tl.load(kT_ptrs + block_sparse_offset)
+        vT = tl.load(vT_ptrs + block_sparse_offset)
+        qk = tl.dot(q, kT)
+        p = tl.math.exp2(qk - m)
+        # Compute dP and dS.
+        dp = tl.dot(do, vT).to(tl.float32)
+        ds = p * (dp - Di[:, None])
+        ds = ds.to(tl.bfloat16)
+        # Compute dQ.
+        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        dq += tl.dot(ds, tl.trans(kT))
+        # Increment pointers.
+    return dq
+
+
+
+@triton.jit
+def _attn_bwd(Q, K, V, sm_scale,  #
+              DO,  #
+              DQ, DK, DV,  #
+              M, D,
+              q2k_index, q2k_num, max_kv_blks,
+              k2q_index, k2q_num, max_q_blks,
+              # shared by Q/K/V/DO.
+              stride_z, stride_h, stride_tok, stride_d,  #
+              H, N_CTX,  #
+              BLOCK_M1: tl.constexpr,  #
+              BLOCK_N1: tl.constexpr,  #
+              BLOCK_M2: tl.constexpr,  #
+              BLOCK_N2: tl.constexpr,  #
+              HEAD_DIM: tl.constexpr):
+    LN2 = 0.6931471824645996  # = ln(2)
+
+    bhid = tl.program_id(2)
+    off_chz = (bhid * N_CTX).to(tl.int64)
+    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
+    pid = tl.program_id(0)
+
+    # offset pointers for batch/head
+    Q += adj
+    K += adj
+    V += adj
+    DO += adj
+    DQ += adj
+    DK += adj
+    DV += adj
+    M += off_chz
+    D += off_chz
+
+    # load scales
+    offs_k = tl.arange(0, HEAD_DIM)
+
+    start_n = pid * BLOCK_N1
+    start_m = 0
+
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+
+    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+    # load K and V: they stay in SRAM throughout the inner loop.
+    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+
+    num_steps = N_CTX // BLOCK_M1
+
+    dk, dv = _attn_bwd_dkdv(  #
+        dk, dv,  #
+        Q, k, v, sm_scale,  #
+        DO,  #
+        M, D,  #
+        k2q_index, k2q_num, max_q_blks,
+        stride_tok, stride_d,  #
+        H, N_CTX,  #
+        BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
+        start_n, start_m, num_steps  #
+    )
+
+    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dv_ptrs, dv)
+
+    # Write back dK.
+    dk *= sm_scale
+    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dk_ptrs, dk)
+
+    # THIS BLOCK DOES DQ:
+    start_m = pid * BLOCK_M2
+    end_n = 0
+
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+
+    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    m = tl.load(M + offs_m)
+    m = m[:, None]
+
+    num_steps = N_CTX // BLOCK_N2
+    dq = _attn_bwd_dq(dq, q, K, V,  #
+                      do, m, D,  #
+                      q2k_index, q2k_num, max_kv_blks,
+                      stride_tok, stride_d,  #
+                      H, N_CTX,  #
+                      BLOCK_M2, BLOCK_N2, HEAD_DIM,  #
+                      start_m, end_n, num_steps  #
+                      )
+    # Write back dQ.
+    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    dq *= LN2
+    tl.store(dq_ptrs, dq)
+    
+    
+
+class _attention(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, k, v):
+        # shape constraints
+        sm_scale = 1.0 / (q.shape[-1] ** 0.5)
+        HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
+        # when v is in float8_e5m2 it is transposed.
+        HEAD_DIM_V = v.shape[-1]
+        assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
+        assert HEAD_DIM_K in {16, 32, 64, 128, 256}
+        o = torch.empty_like(q)
+        stage = 1
+        extra_kern_args = {}
+
+
+        grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
+        M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+        _attn_fwd[grid](
+            q, k, v, sm_scale, M, o,  #
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            k.stride(0), k.stride(1), k.stride(2), k.stride(3),  #
+            v.stride(0), v.stride(1), v.stride(2), v.stride(3),  #
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),  #
+            q.shape[0], q.shape[1],  #
+            N_CTX=q.shape[2],  #
+            HEAD_DIM=HEAD_DIM_K,  #
+            STAGE=stage,  #
+            **extra_kern_args)
+
+        ctx.save_for_backward(q, k, v, o, M)
+        ctx.grid = grid
+        ctx.sm_scale = sm_scale
+        ctx.HEAD_DIM = HEAD_DIM_K
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, M = ctx.saved_tensors
+        assert do.is_contiguous()
+        assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        BATCH, N_HEAD, N_CTX = q.shape[:3]
+        PRE_BLOCK = 128
+        BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
+        RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+        arg_k = k
+        arg_k = arg_k * (ctx.sm_scale * RCP_LN2)
+        PRE_BLOCK = 128
+        assert N_CTX % PRE_BLOCK == 0
+        pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
+        delta = torch.empty_like(M)
+        _attn_bwd_preprocess[pre_grid](
+            o, do,  #
+            delta,  #
+            BATCH, N_HEAD, N_CTX,  #
+            BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM  #
+        )
+        grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
+        _attn_bwd[grid](
+            q, arg_k, v, ctx.sm_scale, do, dq, dk, dv,  #
+            M, delta,  #
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            N_HEAD, N_CTX,  #
+            BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1,  #
+            BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2,  #
+            HEAD_DIM=ctx.HEAD_DIM  #
+        )
+
+        return dq, dk, dv, None, None
+
+
+attention = _attention.apply
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+class _attention_sparse(torch.autograd.Function):
+    """
+    Thin autograd wrapper that uses the sparse forward kernel above and the
+    standard dense backward kernels defined earlier (no extra memory use).
+    """
+
+    @staticmethod
+    def forward(ctx, q, k, v, q2k_index, q2k_num, k2q_index, k2q_num):
+        B, H, T, D = q.shape
+        sm_scale = 1.0 / math.sqrt(D)
+        max_kv_blks = q2k_index.shape[-1]
+        assert T % 64 == 0, f"T must be a multiple of 64, but got {T}"
+        assert T // 64 == q2k_num.shape[-1], f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
+        o = torch.empty_like(q)
+        M = torch.empty((B, H, T), dtype=torch.float32, device=q.device)
+
+        grid = lambda _: (triton.cdiv(T, 64), B * H, 1)
+        _attn_fwd_sparse[grid](
+            q, k, v, sm_scale,
+            q2k_index, q2k_num, max_kv_blks,
+            M, o,
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+            k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+            v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),
+            B, H, T,
+            HEAD_DIM=D, STAGE=3
+        )
+
+        ctx.save_for_backward(q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num)
+        ctx.grid = None
+        ctx.sm_scale = sm_scale
+        ctx.HEAD_DIM = D
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num = ctx.saved_tensors
+        assert do.is_contiguous()
+        assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        BATCH, N_HEAD, N_CTX = q.shape[:3]
+        PRE_BLOCK = 128
+        BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
+        RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+        arg_k = k
+        arg_k = arg_k * (ctx.sm_scale * RCP_LN2)
+        PRE_BLOCK = 128
+        assert N_CTX % PRE_BLOCK == 0
+        pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
+        delta = torch.empty_like(M)
+        _attn_bwd_preprocess[pre_grid](
+            o, do,  #
+            delta,  #
+            BATCH, N_HEAD, N_CTX,  #
+            BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM  #
+        )
+        
+        
+        max_q_blks = k2q_index.shape[-1]
+        max_kv_blks = q2k_index.shape[-1]
+        
+        
+        grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
+        _attn_bwd[grid](
+            q, arg_k, v, ctx.sm_scale, do, dq, dk, dv,  #
+            M, delta,  #
+            q2k_index, q2k_num, max_kv_blks,
+            k2q_index, k2q_num, max_q_blks,
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            N_HEAD, N_CTX,  #
+            BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1,  #
+            BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2,  #
+            HEAD_DIM=ctx.HEAD_DIM  #
+        )
+
+        return dq, dk, dv, None, None, None, None
+
+attention_sparse = _attention_sparse.apply
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+try:
+    from flash_attn.flash_attn_interface import \
+        flash_attn_qkvpacked_func as flash_attn_func
+    HAS_FLASH = True
+except BaseException:
+    HAS_FLASH = False
+
+TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2')
+BATCH, N_HEADS, HEAD_DIM = 4, 32, 128
+# vary seq length for fixed head and batch=4
+configs = []
+for mode in ["fwd", "bwd"]:
+        configs.append(
+            triton.testing.Benchmark(
+                x_names=["N_CTX"],
+                x_vals=[2**i for i in range(10, 15)],
+                line_arg="provider",
+                line_vals=["triton-fp16"] + (["triton-fp8"] if TORCH_HAS_FP8 else []) +
+                (["flash"] if HAS_FLASH else []),
+                line_names=["Triton [FP16]"] + (["Triton [FP8]"] if TORCH_HAS_FP8 else []) +
+                (["Flash-2"] if HAS_FLASH else []),
+                styles=[("red", "-"), ("blue", "-"), ("green", "-")],
+                ylabel="ms",
+                plot_name=f"fused-attention-batch{BATCH}-head{N_HEADS}-d{HEAD_DIM}-{mode}",
+                args={
+                    "H": N_HEADS,
+                    "BATCH": BATCH,
+                    "HEAD_DIM": HEAD_DIM,
+                    "mode": mode,
+                },
+            ))
+
+
+@triton.testing.perf_report(configs)
+def bench_flash_attention(BATCH, H, N_CTX, HEAD_DIM, mode, provider, device="cuda"):
+    assert mode in ["fwd", "bwd"]
+    warmup = 25
+    rep = 100
+    dtype = torch.bfloat16
+    if "triton" in provider:
+        q = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        k = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        v = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        if mode == "fwd" and "fp8" in provider:
+            q = q.to(torch.float8_e5m2)
+            k = k.to(torch.float8_e5m2)
+            v = v.permute(0, 1, 3, 2).contiguous()
+            v = v.permute(0, 1, 3, 2)
+            v = v.to(torch.float8_e5m2)
+        fn = lambda: attention(q, k, v)
+        if mode == "bwd":
+            o = fn()
+            do = torch.randn_like(o)
+            fn = lambda: o.backward(do, retain_graph=True)
+        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
+    if provider == "flash":
+        qkv = torch.randn((BATCH, N_CTX, 3, H, HEAD_DIM), dtype=dtype, device=device, requires_grad=True)
+        fn = lambda: flash_attn_func(qkv, causal=False)
+        if mode == "bwd":
+            o = fn()
+            do = torch.randn_like(o)
+            fn = lambda: o.backward(do, retain_graph=True)
+        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
+    flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * HEAD_DIM
+    total_flops = 2 * flops_per_matmul
+    if mode == "bwd":
+        total_flops *= 2.5  # 2.0(bwd) + 0.5(recompute)
+    return total_flops / ms * 1e-9
+
+
+if __name__ == "__main__":
+    # only works on post-Ampere GPUs right now
+    bench_flash_attention.run(save_path=".", print_data=True)
@@ -0,0 +1,47 @@
+import torch
+try:
+    from vsa_cuda import block_sparse_fwd, block_sparse_bwd
+except ImportError:
+    block_sparse_fwd = None
+    block_sparse_bwd = None
+from .block_sparse_attn_triton import attention_sparse as block_sparse_attn_triton
+
+class BlockSparseAttentionFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
+        o, lse = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
+        ctx.save_for_backward(q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num)
+        return o
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num = ctx.saved_tensors
+        grad_q, grad_k, grad_v = block_sparse_bwd(
+            q, k, v, o, lse, grad_output, k2q_block_sparse_index, k2q_block_sparse_num
+        )
+        return grad_q, grad_k, grad_v, None, None, None, None
+
+
+@torch._dynamo.disable
+def block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
+    """
+    Differentiable block sparse attention function.
+    
+    Args:
+        q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
+        k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
+        v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
+        q2k_block_sparse_index: Indices for query-to-key sparse blocks
+        q2k_block_sparse_num: Number of sparse blocks for each query block
+        k2q_block_sparse_index: Indices for key-to-query sparse blocks (for backward pass)
+        k2q_block_sparse_num: Number of sparse blocks for each key block (for backward pass)
+    
+    Returns:
+        output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
+    """
+    if block_sparse_fwd is not None:
+        return BlockSparseAttentionFunction.apply(
+            q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
+        )
+    else:
+        return block_sparse_attn_triton(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
@@ -23,3 +23,4 @@ clean:
 	@$(SPHINXBUILD) -M clean "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 	rm -rf "$(SOURCEDIR)/getting_started/examples"
 	rm -rf "$(SOURCEDIR)/inference/examples"
+	rm -rf "$(SOURCEDIR)/training/examples"
@@ -11,5 +11,5 @@ commonmark # Required by sphinx-argparse when using :markdownhelp:

 # packages to install to build the documentation
 cachetools
-f https://download.pytorch.org/whl/cpu
+# -f https://download.pytorch.org/whl/cpu
 torch
@@ -9,11 +9,11 @@
 ## Initialization Configuration

 ```{autodoc2-summary}
-    fastvideo.v1.configs.pipelines.PipelineConfig
+    fastvideo.configs.pipelines.PipelineConfig
 ```

 ## Sampling Configuration

 ```{autodoc2-summary}
-    fastvideo.v1.configs.sample.SamplingParam
+    fastvideo.configs.sample.SamplingParam
 ```
@@ -1,24 +1,24 @@
 # 🔍 FastVideo Overview

-This document outlines FastVideo's architecture for developers interested in framework internals or contributions. It serves as an onboarding guide for new contributors by providing an overview of the most important directories and files within the `fastvideo/v1/` codebase.
+This document outlines FastVideo's architecture for developers interested in framework internals or contributions. It serves as an onboarding guide for new contributors by providing an overview of the most important directories and files within the `fastvideo/` codebase.

-## Table of Contents - V1 Directory Structure and Files
+## Table of Contents - Directory Structure and Files

- [`fastvideo/v1/pipelines/`](#design-pipeline-system) - Core diffusion pipeline components
- [`fastvideo/v1/models/`](#design-model-components) - Model implementations
+- [`fastvideo/pipelines/`](#design-pipeline-system) - Core diffusion pipeline components
+- [`fastvideo/models/`](#design-model-components) - Model implementations
  - [`dits/`](#design-transformer-models) - Transformer-based diffusion models
  - [`vaes/`](#design-vae-variational-auto-encoder) - Variational autoencoders
  - [`encoders/`](#design-text-and-image-encoders) - Text and image encoders
  - [`schedulers/`](#design-schedulers) - Diffusion schedulers
- [`fastvideo/v1/attention/`](#design-optimized-attention) - Optimized attention implementations
- [`fastvideo/v1/distributed/`](#design-distributed-processing) - Distributed computing utilities
- [`fastvideo/v1/layers/`](#design-tensor-parallelism) - Custom neural network layers
- [`fastvideo/v1/platforms/`](#design-platforms) - Hardware platform abstractions
- [`fastvideo/v1/worker/`](#design-executor-and-worker-abstractions) - Multi-GPU process management
- [`fastvideo/v1/fastvideo_args.py`](#design-fastvideo-args) - Argument handling
- [`fastvideo/v1/forward_context.py`](#design-forwardcontext) - Forward pass context management
- `fastvideo/v1/utils.py` - Utility functions
- [`fastvideo/v1/logger.py`](#design-logger) - Logging infrastructure
+- [`fastvideo/attention/`](#design-optimized-attention) - Optimized attention implementations
+- [`fastvideo/distributed/`](#design-distributed-processing) - Distributed computing utilities
+- [`fastvideo/layers/`](#design-tensor-parallelism) - Custom neural network layers
+- [`fastvideo/platforms/`](#design-platforms) - Hardware platform abstractions
+- [`fastvideo/worker/`](#design-executor-and-worker-abstractions) - Multi-GPU process management
+- [`fastvideo/fastvideo_args.py`](#design-fastvideo-args) - Argument handling
+- [`fastvideo/forward_context.py`](#design-forwardcontext) - Forward pass context management
+- `fastvideo/utils.py` - Utility functions
+- [`fastvideo/logger.py`](#design-logger) - Logging infrastructure

 ## Core Architecture

@@ -32,7 +32,7 @@ FastVideo separates model components from execution logic with these principles:
 (design-fastvideo-args)=
 ## FastVideoArgs

-The `FastVideoArgs` class in `fastvideo/v1/fastvideo_args.py` serves as the central configuration system for FastVideo. It contains all parameters needed to control model loading, inference configuration, performance optimization settings, and more.
+The `FastVideoArgs` class in `fastvideo/fastvideo_args.py` serves as the central configuration system for FastVideo. It contains all parameters needed to control model loading, inference configuration, performance optimization settings, and more.

 Key features include:
 - **Command-line Interface**: Automatic conversion between CLI arguments and dataclass fields
@@ -111,7 +111,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
 (design-forwardbatch)=
 ### ForwardBatch

-Defined in `fastvideo/v1/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:
+Defined in `fastvideo/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:

 - **Input Data**: Prompts, images, generation parameters
 - **Intermediate State**: Embeddings, latents, timesteps, accumulated during stage execution
@@ -123,14 +123,14 @@ This structure facilitates clear state transitions between stages.
 (design-model-components)=
 ## Model Components

-The `fastvideo/v1/models/` directory contains implementations of the core neural network models used in video diffusion:
+The `fastvideo/models/` directory contains implementations of the core neural network models used in video diffusion:

 (design-transformer-models)=
 ### Transformer Models

 Transformer networks perform the actual denoising during diffusion:

- **Location**: `fastvideo/v1/models/dits/`
+- **Location**: `fastvideo/models/dits/`
 - **Examples**:
  - `WanTransformer3DModel`
  - `HunyuanVideoTransformer3DModel`
@@ -157,7 +157,7 @@ def forward(

 VAEs handle conversion between pixel space and latent space:

- **Location**: `fastvideo/v1/models/vaes/`
+- **Location**: `fastvideo/models/vaes/`
 - **Examples**:
  - `AutoencoderKLWan`
  - `AutoencoderKLHunyuanVideo`
@@ -175,7 +175,7 @@ FastVideo's VAE implementations include:

 Encoders process conditioning inputs into embeddings:

- **Location**: `fastvideo/v1/models/encoders/`
+- **Location**: `fastvideo/models/encoders/`
 - **Text Encoders**:
  - `CLIPTextModel`
  - `LlamaModel`
@@ -193,7 +193,7 @@ FastVideo implements optimizations such as:

 Schedulers manage the diffusion sampling process:

- **Location**: `fastvideo/v1/models/schedulers/`
+- **Location**: `fastvideo/models/schedulers/`
 - **Examples**:
  - `UniPCMultistepScheduler`
  - `FlowMatchEulerDiscreteScheduler`
@@ -219,7 +219,7 @@ def step(
 (design-optimized-attention)=
 ## Optimized Attention

-The `fastvideo/v1/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:
+The `fastvideo/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:

 ### Attention Backends
 Multiple implementations with automatic selection:
@@ -248,7 +248,7 @@ Supports various patterns with memory optimization techniques:
 (design-distributed-processing)=
 ## Distributed Processing

-The `fastvideo/v1/distributed/` directory contains implementations for distributed model execution:
+The `fastvideo/distributed/` directory contains implementations for distributed model execution:

 (design-tensor-parallelism)=
 ### Tensor Parallelism
@@ -260,7 +260,7 @@ Tensor parallelism splits model weights across devices:

 ```python
 # Tensor-parallel layers in a transformer block
-from fastvideo.v1.layers.linear import ColumnParallelLinear, RowParallelLinear
+from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear

 # Split along output dimension
 self.qkv_proj = ColumnParallelLinear(
@@ -288,7 +288,7 @@ Sequence parallelism splits sequences across devices:

 ```python
 # Distributed attention for long sequences
-from fastvideo.v1.attention import DistributedAttention
+from fastvideo.attention import DistributedAttention

 self.attn = DistributedAttention(
    num_heads=num_heads,
@@ -312,7 +312,7 @@ Efficient communication primitives minimize distributed overhead:

 ### ForwardContext

-Defined in `fastvideo/v1/forward_context.py`, `ForwardContext` manages execution-specific state *within* a forward pass, particularly for low-level optimizations. It is accessed via `get_forward_context()`.
+Defined in `fastvideo/forward_context.py`, `ForwardContext` manages execution-specific state *within* a forward pass, particularly for low-level optimizations. It is accessed via `get_forward_context()`.

 - **Attention Metadata**: Configuration for optimized attention kernels (`attn_metadata`)
 - **Profiling Data**: Potential hooks for performance metrics collection
@@ -333,7 +333,7 @@ with set_forward_context(current_timestep, attn_metadata, fastvideo_args):
 (design-executor-and-worker-abstractions)=
 ## Executor and Worker System

-The `fastvideo/v1/worker/` directory contains the distributed execution framework:
+The `fastvideo/worker/` directory contains the distributed execution framework:

 ### Executor Abstraction

@@ -360,7 +360,7 @@ This design allows FastVideo to efficiently utilize multiple GPUs while providin
 (design-platforms)=
 ## Platforms

-The `fastvideo/v1/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:
+The `fastvideo/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:

 ### Platform Abstraction

@@ -377,7 +377,7 @@ The primary components include:
 Usage example:

 ```python
-from fastvideo.v1.platforms import current_platform, _Backend
+from fastvideo.platforms import current_platform, _Backend

 # Check hardware capabilities
 if current_platform.supports_backend(_Backend.FLASH_ATTN):
@@ -398,9 +398,9 @@ See [PR](https://github.com/hao-ai-lab/FastVideo/pull/356)

 If you're a new contributor, here are some common areas to explore:

-1. **Adding a new model**: Implement new model types in the appropriate subdirectory of `fastvideo/v1/models/`
+1. **Adding a new model**: Implement new model types in the appropriate subdirectory of `fastvideo/models/`
 2. **Optimizing performance**: Look at attention implementations or memory management
-3. **Adding a new pipeline**: Create a new pipeline subclass in `fastvideo/v1/pipelines/`
+3. **Adding a new pipeline**: Create a new pipeline subclass in `fastvideo/pipelines/`
 4. **Hardware support**: Extend the `platforms` module for new hardware targets

 When adding code, follow these practices:
@@ -27,6 +27,15 @@ def fix_case(text: str) -> str:
        "openai": "OpenAI",
        "multilora": "MultiLoRA",
        "mlpspeculator": "MLPSpeculator",
+        "finetune": "Finetune",
+        "distillation": "Distillation",
+        "wan": "Wan",
+        "i2v": "I2V",
+        "t2v": "T2V",
+        "1.3b": "1.3B",
+        "14b": "14B",
+        "480p": "480P",
+        "720p": "720P",
        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
    }
@@ -161,31 +170,35 @@ class Example:
        return content


-def generate_examples(generate_main_index=False):
-    """
-    Generate example documentation.
-    
-    Args:
-        generate_main_index (bool): Whether to generate the main examples index.
-            If False, only category-specific indices will be generated.
-    """
-    # Create empty indices with dynamic paths
+@dataclass
+class NestedStructure:
+    """Helper class to manage nested documentation structures for training/distillation."""
+    category: str
+    method: str
+    model: str
+    dataset: str
+    example: Example
+
+    @property
+    def filename(self) -> str:
+        return f"{self.model}_{self.dataset}"
+
+    @property
+    def title(self) -> str:
+        return fix_case(self.dataset.replace('_', ' '))
+
+    @property
+    def description(self) -> str:
+        category_name = self.category.title()
+        return f"{category_name} example using the {self.dataset} dataset with the {self.model} model."
+
+
+def create_category_indices() -> dict[str, Index]:
+    """Create category indices with their respective configurations."""
    main_index_dir = ROOT_DIR / "docs/source/examples"
    if not main_index_dir.exists():
        main_index_dir.mkdir(parents=True)

-    # Create the main examples index only if requested
-    examples_index = None
-    if generate_main_index:
-        examples_index = Index(
-            path=main_index_dir / "examples_index.md",
-            title="💡 Examples",
-            description=
-            "A collection of examples demonstrating usage of FastVideo.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.",  # noqa: E501
-            caption="Examples",
-            maxdepth=2)
-
-    # Category indices with dynamic paths based on category names
    category_indices = {
        "inference":
        Index(
@@ -193,28 +206,54 @@ def generate_examples(generate_main_index=False):
            "docs/source/inference/examples/examples_inference_index.md",
            title="🚀 Examples",
            description=
-            "Inference examples demonstrate how to use FastVideo in an offline setting, where the model is queried for predictions in batches. We recommend starting with <project:basic.md>.",  # noqa: E501
+            "Inference examples demonstrate how to use FastVideo inference. We recommend starting with <project:basic.md>.",
            caption="Examples",
+            maxdepth=1,
+        ),
+        "training":
+        Index(
+            path=ROOT_DIR /
+            "docs/source/training/examples/examples_training_index.md",
+            title="🚀 Examples",
+            description=
+            "Training examples demonstrate how to use FastVideo training.",
+            caption="Examples",
+            maxdepth=3,
+        ),
+        "distillation":
+        Index(
+            path=ROOT_DIR /
+            "docs/source/distillation/examples/examples_distillation_index.md",
+            title="🚀 Examples",
+            description=
+            "Distillation examples demonstrate how to use FastVideo distillation.",
+            caption="Examples",
+            maxdepth=3,
        ),
    }

    # Ensure all category doc directories exist
-    for category, index in category_indices.items():
-        category_dir = index.path.parent
-        if not category_dir.exists():
-            category_dir.mkdir(parents=True)
+    for index in category_indices.values():
+        if not index.path.parent.exists():
+            index.path.parent.mkdir(parents=True)

+    return category_indices
+
+
+def find_examples(category_indices: dict[str, Index],
+                  generate_main_index: bool) -> list[Example]:
+    """Find all examples from the examples directory."""
    examples = []
    glob_patterns = ["*.py", "*.md", "*.sh"]
+
    # Find categorised examples
    for category in category_indices:
-        print(category)
        category_dir = EXAMPLE_DIR / category
        globs = [category_dir.glob(pattern) for pattern in glob_patterns]
        for path in itertools.chain(*globs):
            examples.append(Example(path, category))
-        # Find examples in subdirectories
-        for path in category_dir.glob("*/*.md"):
+        # Find examples in subdirectories (recursively)
+        for path in category_dir.glob("**/*.md"):
            examples.append(Example(path.parent, category))

    # Find uncategorised examples only if we're generating a main index
@@ -229,36 +268,173 @@ def generate_examples(generate_main_index=False):
                continue
            examples.append(Example(path.parent))

-    # Create document directories for each category based on category name and generate files
-    for example in sorted(examples, key=lambda e: e.path.stem):
-        print(example)
+    return examples
+
+
+def create_nested_structures(
+    examples: list[Example]
+) -> dict[str, dict[str, dict[str, dict[str, NestedStructure]]]]:
+    """Create nested structures for training and distillation categories."""
+    nested_structures: dict[str, dict[str, dict[str,
+                                                dict[str,
+                                                     NestedStructure]]]] = {}
+
+    for example in examples:
+        if example.category not in ["training", "distillation"]:
+            continue
+
+        category_dir = EXAMPLE_DIR / example.category
+        relative_path = example.path.relative_to(category_dir)
+        path_parts = relative_path.parts
+
+        # For nested examples like finetune/wan_i2v_14b_480p/crush_smol
+        if len(path_parts) >= 3:
+            method = path_parts[0]  # e.g., "finetune"
+            model = path_parts[1]  # e.g., "wan_i2v_14b_480p"
+            dataset = path_parts[2]  # e.g., "crush_smol"
+
+            # Initialize nested structure
+            if example.category not in nested_structures:
+                nested_structures[example.category] = {}
+            if method not in nested_structures[example.category]:
+                nested_structures[example.category][method] = {}
+            if model not in nested_structures[example.category][method]:
+                nested_structures[example.category][method][model] = {}
+
+            # Store the nested structure
+            nested_structures[
+                example.category][method][model][dataset] = NestedStructure(
+                    category=example.category,
+                    method=method,
+                    model=model,
+                    dataset=dataset,
+                    example=example)
+
+    return nested_structures
+
+
+def generate_flat_examples(examples: list[Example],
+                           category_indices: dict[str, Index],
+                           examples_index: Index | None,
+                           generate_main_index: bool) -> None:
+    """Generate documentation for flat structure examples (inference, etc.)."""
+    for example in examples:
+        if example.category in ["training", "distillation"]:
+            continue  # Skip nested structure examples

        # Determine which index to use for this example
        if example.category is not None and example.category in category_indices:
            index = category_indices[example.category]
        elif generate_main_index:
            assert examples_index is not None
-            index = examples_index  # Default to main index if available
+            index = examples_index
        else:
-            # Skip examples without a category if no main index
-            print(f"Skipping {example.path} (no category and no main index)")
            continue

-        # Place generated example markdown in the same directory as its index
+        # Generate the example documentation
        doc_path = index.path.parent / f"{example.path.stem}.md"
        with open(doc_path, "w+") as f:
            f.write(example.generate())
-        # Add the example to the index
        index.documents.append(example.path.stem)

+
+def generate_nested_examples(nested_structures: dict[str, dict[str, dict[
+    str, dict[str, NestedStructure]]]], category_indices: dict[str,
+                                                               Index]) -> None:
+    """Generate documentation for nested structure examples (training, distillation)."""
+    for category_name in ["training", "distillation"]:
+        if category_name not in category_indices or category_name not in nested_structures:
+            continue
+
+        category_index = category_indices[category_name]
+        category_base_dir = category_index.path.parent
+
+        for method, models in nested_structures[category_name].items():
+            # Create method-level index
+            method_index = Index(path=category_base_dir / f"{method}.md",
+                                 title=fix_case(method),
+                                 description=f"Examples using {method}.",
+                                 caption=f"{fix_case(method)} Examples",
+                                 maxdepth=2)
+
+            for model, datasets in models.items():
+                # Generate dataset examples using the Example class
+                for dataset, nested_struct in datasets.items():
+                    doc_path = category_base_dir / f"{nested_struct.filename}.md"
+                    with open(doc_path, "w+") as f:
+                        f.write(nested_struct.example.generate())
+
+                # Create model-level index
+                model_index = Index(
+                    path=category_base_dir / f"{model}.md",
+                    title=fix_case(model.replace('_', ' ')),
+                    description=f"Examples for the {model} model.",
+                    caption=f"{fix_case(model.replace('_', ' '))} Datasets",
+                    maxdepth=1)
+
+                # Add dataset indices to model index
+                for dataset, nested_struct in datasets.items():
+                    model_index.documents.append(nested_struct.filename)
+
+                # Write model index
+                with open(model_index.path, "w+") as f:
+                    f.write(model_index.generate())
+
+                # Add model to method index
+                method_index.documents.append(model)
+
+            # Write method index
+            with open(method_index.path, "w+") as f:
+                f.write(method_index.generate())
+
+            # Add method to main category index
+            category_index.documents.append(method)
+
+
+def generate_examples(generate_main_index=False):
+    """
+    Generate example documentation.
+    
+    Args:
+        generate_main_index (bool): Whether to generate the main examples index.
+            If False, only category-specific indices will be generated.
+    """
+    # Create category indices
+    category_indices = create_category_indices()
+
+    # Create the main examples index only if requested
+    examples_index = None
+    if generate_main_index:
+        main_index_dir = ROOT_DIR / "docs/source/examples"
+        examples_index = Index(
+            path=main_index_dir / "examples_index.md",
+            title="💡 Examples",
+            description=
+            "A collection of examples demonstrating usage of FastVideo.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.",
+            caption="Examples",
+            maxdepth=2)
+
+    # Find all examples
+    examples = find_examples(category_indices, generate_main_index)
+
+    # Create nested structures for training and distillation
+    nested_structures = create_nested_structures(examples)
+
+    # Generate flat structure examples (inference, etc.)
+    generate_flat_examples(examples, category_indices, examples_index,
+                           generate_main_index)
+
+    # Generate nested structure examples (training, distillation)
+    generate_nested_examples(nested_structures, category_indices)
+
    # Generate the index files for categories
    for category_index in category_indices.values():
        if category_index.documents:
            # Add to main index if it exists
-            if generate_main_index:
+            if generate_main_index and examples_index:
+                main_index_dir = examples_index.path.parent
                rel_path = category_index.path.relative_to(
                    main_index_dir.parent)
-                assert examples_index is not None
                examples_index.documents.insert(
                    0,
                    str(rel_path).replace(".md", ""))
@@ -1,120 +1,18 @@
-(fastvideo-installation)=
+(installation-index)=

 # 🔧 Installation

-FastVideo currently only supports Linux and NVIDIA CUDA GPUs.
+FastVideo supports the following hardware platforms:

-## Requirements
+:::{toctree}
+:maxdepth: 1
+:hidden:

- **OS: Linux**
- **Python: 3.10-3.12**
- **CUDA 12.4**
- **At least 1 NVIDIA GPU**
-
-## Set up using Python
-### Create a new Python environment
-
-#### Conda
-You can create a new python environment using [Conda](https://docs.conda.io/projects/conda/en/stable/user-guide/getting-started.html)
-##### 1. Install Miniconda (if not already installed)
-
-```bash
-wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
-bash Miniconda3-latest-Linux-x86_64.sh
-source ~/.bashrc
-```
-
-##### 2. Create and activate a Conda environment for FastVideo
-
-```bash
-# (Recommended) Create a new conda environment.
-conda create -n fastvideo python=3.12 -y
-conda activate fastvideo
-```
-
-:::{note}
-[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
+installation/gpu
+installation/mps
 :::

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

-`````{py:class} VideoGenerator(fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs, executor_class: type[fastvideo.v1.worker.executor.Executor], log_stats: bool)
-:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator
+`````{py:class} VideoGenerator(fastvideo_args: fastvideo.fastvideo_args.FastVideoArgs, executor_class: type[fastvideo.worker.executor.Executor], log_stats: bool)
+:canonical: fastvideo.entrypoints.video_generator.VideoGenerator

-```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator
+```{autodoc2-docstring} fastvideo.entrypoints.video_generator.VideoGenerator
 :parser: docs.source.autodoc2_docstring_parser
 ```

 `VideoGenerator.from_pretrained()` should be the primary way of creating a new video generator.

-````{py:method} from_pretrained(model_path: str, device: typing.Optional[str] = None, torch_dtype: typing.Optional[torch.dtype] = None, pipeline_config: typing.Optional[typing.Union[str | fastvideo.v1.configs.pipelines.PipelineConfig]] = None, **kwargs) -> fastvideo.v1.entrypoints.video_generator.VideoGenerator
-:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
+````{py:method} from_pretrained(model_path: str, device: typing.Optional[str] = None, torch_dtype: typing.Optional[torch.dtype] = None, pipeline_config: typing.Optional[typing.Union[str | fastvideo.configs.pipelines.PipelineConfig]] = None, **kwargs) -> fastvideo.entrypoints.video_generator.VideoGenerator
+:canonical: fastvideo.entrypoints.video_generator.VideoGenerator.from_pretrained
 :classmethod:

-```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
+```{autodoc2-docstring} fastvideo.entrypoints.video_generator.VideoGenerator.from_pretrained
 :parser: docs.source.autodoc2_docstring_parser
 ```

@@ -38,25 +38,25 @@ The follow two classes `PipelineConfig` and `SamplingParam` are used to configur
 ```

 `````{py:class} PipelineConfig
-:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig
+:canonical: fastvideo.configs.pipelines.base.PipelineConfig

-```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig
+```{autodoc2-docstring} fastvideo.configs.pipelines.base.PipelineConfig
 :parser: docs.source.autodoc2_docstring_parser
 ```

-````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.pipelines.base.PipelineConfig
-:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+````{py:method} from_pretrained(model_path: str) -> fastvideo.configs.pipelines.base.PipelineConfig
+:canonical: fastvideo.configs.pipelines.base.PipelineConfig.from_pretrained
 :classmethod:

-```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+```{autodoc2-docstring} fastvideo.configs.pipelines.base.PipelineConfig.from_pretrained
 :parser: docs.source.autodoc2_docstring_parser
 ```


 ````{py:method} dump_to_json(file_path: str)
-:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+:canonical: fastvideo.configs.pipelines.base.PipelineConfig.dump_to_json

-```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+```{autodoc2-docstring} fastvideo.configs.pipelines.base.PipelineConfig.dump_to_json
 :parser: docs.source.autodoc2_docstring_parser
 ```

@@ -68,16 +68,16 @@ The follow two classes `PipelineConfig` and `SamplingParam` are used to configur
 ```

 `````{py:class} SamplingParam
-:canonical: fastvideo.v1.configs.sample.base.SamplingParam
+:canonical: fastvideo.configs.sample.base.SamplingParam

-```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam
+```{autodoc2-docstring} fastvideo.configs.sample.base.SamplingParam
 :parser: docs.source.autodoc2_docstring_parser
 ```

-````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.sample.base.SamplingParam
-:canonical: fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+````{py:method} from_pretrained(model_path: str) -> fastvideo.configs.sample.base.SamplingParam
+:canonical: fastvideo.configs.sample.base.SamplingParam.from_pretrained
 :classmethod:

-```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+```{autodoc2-docstring} fastvideo.configs.sample.base.SamplingParam.from_pretrained
 :parser: docs.source.autodoc2_docstring_parser
 ```
@@ -63,22 +63,33 @@ getting_started/installation
 :maxdepth: 1

 inference/inference_quick_start
+inference/examples/examples_inference_index
 inference/configuration
 inference/optimizations
+inference/comfyui
 inference/support_matrix
-inference/examples/examples_inference_index
 inference/cli
 inference/add_pipeline
-inference/v0_inference
 :::

 :::{toctree}
 :caption: Training
 :maxdepth: 1

+training/examples/examples_training_index
 training/data_preprocess
 training/distillation
-training/finetune
+<!-- training/finetune -->
+:::
+
+<!-- :::{toctree}
+:caption: Distillation
+:maxdepth: 1
+
+distillation/examples/examples_distillation_index
+distillation/data_preprocess
+distillation/dmd -->
+<!-- training/finetune -->
 :::

 % What is STA Kernel?
@@ -46,25 +46,25 @@ FastVideo uses the Hugging Face Diffusers format for model organization:
 ### Implementing Modules

 Place new modules in the appropriate directories:
- Encoders: `fastvideo/v1/models/encoders/`
- VAEs: `fastvideo/v1/models/vaes/`
- Transformer models: `fastvideo/v1/models/dits/`
- Schedulers: `fastvideo/v1/models/schedulers/`
+- Encoders: `fastvideo/models/encoders/`
+- VAEs: `fastvideo/models/vaes/`
+- Transformer models: `fastvideo/models/dits/`
+- Schedulers: `fastvideo/models/schedulers/`

 ### Adapting Model Layers

 #### Layer Replacements
 Replace standard PyTorch layers with FastVideo optimized versions:
- nn.LayerNorm → fastvideo.v1.layers.layernorm.RMSNorm
- Embedding layers → fastvideo.v1.layers.vocab_parallel_embedding modules
- Activation functions → versions from fastvideo.v1.layers.activation
+- nn.LayerNorm → fastvideo.layers.layernorm.RMSNorm
+- Embedding layers → fastvideo.layers.vocab_parallel_embedding modules
+- Activation functions → versions from fastvideo.layers.activation

 #### Distributed Linear Layers
 Use appropriate parallel layers for distribution:

 ```python
 # Output dimension parallelism
-from fastvideo.v1.layers.linear import ColumnParallelLinear
+from fastvideo.layers.linear import ColumnParallelLinear
 self.q_proj = ColumnParallelLinear(
    input_size=hidden_size,
    output_size=head_size * num_heads,
@@ -73,7 +73,7 @@ self.q_proj = ColumnParallelLinear(
 )

 # Fused QKV projection
-from fastvideo.v1.layers.linear import QKVParallelLinear
+from fastvideo.layers.linear import QKVParallelLinear
 self.qkv_proj = QKVParallelLinear(
    hidden_size=hidden_size,
    head_size=attention_head_dim,
@@ -82,7 +82,7 @@ self.qkv_proj = QKVParallelLinear(
 )

 # Input dimension parallelism
-from fastvideo.v1.layers.linear import RowParallelLinear
+from fastvideo.layers.linear import RowParallelLinear
 self.out_proj = RowParallelLinear(
    input_size=head_size * num_heads,
    output_size=hidden_size,
@@ -96,8 +96,8 @@ Replace standard attention with FastVideo's optimized attention:

 ```python
 # Local attention patterns
-from fastvideo.v1.attention import LocalAttention
-from fastvideo.v1.attention.backends.abstract import _Backend
+from fastvideo.attention import LocalAttention
+from fastvideo.attention.backends.abstract import _Backend
 self.attn = LocalAttention(
    num_heads=num_heads,
    head_size=head_dim,
@@ -108,7 +108,7 @@ self.attn = LocalAttention(
 )

 # Distributed attention for long sequences
-from fastvideo.v1.attention import DistributedAttention
+from fastvideo.attention import DistributedAttention
 self.attn = DistributedAttention(
    num_heads=num_heads,
    head_size=head_dim,
@@ -130,7 +130,7 @@ self.attn = DistributedAttention(
 Register implemented modules in the model registry:

 ```python
-# In fastvideo/v1/models/registry.py
+# In fastvideo/models/registry.py
 _TEXT_TO_VIDEO_DIT_MODELS = {
    "YourTransformerModel": ("dits", "yourmodule", "YourTransformerClass"),
 }
@@ -145,7 +145,7 @@ _VAE_MODELS = {
 Create a new directory for your pipeline:

 ```
-fastvideo/v1/pipelines/
+fastvideo/pipelines/
 ├── your_pipeline/
 │   ├── __init__.py
 │   └── your_pipeline.py
@@ -167,13 +167,13 @@ Pipelines are composed of stages, each handling a specific part of the diffusion
 ### Creating Your Pipeline

 ```python
-from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
-from fastvideo.v1.pipelines.stages import (
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (
    InputValidationStage, CLIPTextEncodingStage, TimestepPreparationStage,
    LatentPreparationStage, DenoisingStage, DecodingStage
 )
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 import torch

 class MyCustomPipeline(ComposedPipelineBase):
@@ -246,7 +246,7 @@ EntryClass = MyCustomPipeline
 If existing stages don't meet your needs, create custom ones:

 ```python
-from fastvideo.v1.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.base import PipelineStage

 class MyCustomStage(PipelineStage):
    """Custom processing stage for the pipeline."""
@@ -305,7 +305,7 @@ EntryClass = [MyCustomPipeline, MyOtherPipeline]
 ```

 The registry will automatically:
-1. Scan all packages under `fastvideo/v1/pipelines/`
+1. Scan all packages under `fastvideo/pipelines/`
 2. Look for `EntryClass` variables
 3. Register pipelines using their class names as identifiers

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

 The folder `crush-smol_raw/` contains raw videos and captions for testing preprocessing, while `crush-smol_preprocessed/` contains latents prepared for testing training.
@@ -0,0 +1,3 @@
+# DMD Distillation Wan2.1-I2V-14B-480P Crush-Smol Example
+
+Coming soon!
@@ -18,6 +18,16 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 python examples/inference/basic/basic.py
 ```

+For an example on Apple silicon: 
+```
+python examples/inference/basic/basic_mps.py
+```
+
+For an example running DMD+VSA inference:
+```
+python examples/inference/basic/basic_dmd.py
+```
+
 ## Basic Walkthrough

 All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines!
@@ -1,6 +1,6 @@
 from fastvideo import VideoGenerator

-# from fastvideo.v1.configs.sample import SamplingParam
+# from fastvideo.configs.sample import SamplingParam

 OUTPUT_PATH = "video_samples"
 def main():
@@ -0,0 +1,46 @@
+import os
+import time
+from fastvideo import VideoGenerator
+
+from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_dmd2"
+def main():
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+
+    load_start_time = time.perf_counter()
+    model_name = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        # Adjust these offload parameters if you have < 32GB of VRAM
+        text_encoder_offload=False,
+        use_cpu_offload=False,
+        VSA_sparsity=0.8,
+    )
+    load_end_time = time.perf_counter()
+    load_time = load_end_time - load_start_time
+
+
+    sampling_param = SamplingParam.from_pretrained(model_name)
+
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+    start_time = time.perf_counter()
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+    end_time = time.perf_counter()
+    gen_time = end_time - start_time
+    e2e_gen_time = end_time - start_time
+
+    print(f"Time taken to load model: {load_time} seconds")
+    print(f"Time taken to generate video: {gen_time} seconds")
+    print(f"Time taken for e2e generation: {e2e_gen_time} seconds")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,40 @@
+from fastvideo import VideoGenerator, PipelineConfig
+from fastvideo.configs.sample import SamplingParam
+
+def main():
+    config = PipelineConfig.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    config.text_encoder_precisions = ["fp16"]
+    
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        pipeline_config=config,
+        use_fsdp_inference=False,      # Disable FSDP for MPS
+        use_cpu_offload=True,          
+        text_encoder_offload=True,    
+        pin_cpu_memory=True,           
+        disable_autocast=False,        
+        num_gpus=1,      
+    )
+
+    # Create sampling parameters with reduced number of frames
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    sampling_param.num_frames = 3  # Reduce from default 81 to 25 frames bc we have to use the SDPA attn backend for mps
+    sampling_param.height = 256
+    sampling_param.width = 256
+
+    prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+             "wide with interest. The playful yet serene atmosphere is complemented by soft "
+             "natural light filtering through the petals. Mid-shot, warm and cheerful tones.")
+    
+    video = generator.generate_video(prompt, sampling_param=sampling_param)
+
+    prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    
+    video2 = generator.generate_video(prompt2, sampling_param=sampling_param)
+
+if __name__ == "__main__":
+    main()
@@ -1,5 +1,5 @@
 from fastvideo import VideoGenerator
-from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.pipelines.base import PipelineConfig

 def main():

@@ -0,0 +1,39 @@
+from fastvideo import VideoGenerator, PipelineConfig, SamplingParam
+
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_fp16"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    model = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    pipeline_config = PipelineConfig.from_pretrained(model)
+    pipeline_config.text_encoder_precisions = ("bf16", )
+    generator = VideoGenerator.from_pretrained(
+        model,
+        # if num_gpus > 1, FastVideo will automatically handle distributed setup
+        pipeline_config=pipeline_config,
+        use_fsdp_inference=False,      # Disable FSDP for MPS
+        use_cpu_offload=True,          
+        text_encoder_offload=True,    
+        pin_cpu_memory=True,           
+        disable_autocast=False,        
+        num_gpus=1, 
+    )
+
+    sampling_param = SamplingParam.from_pretrained(model)
+    sampling_param.num_frames = 30
+    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    # Generate videos with the same simple API, regardless of GPU count
+    prompt = "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting."
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+    # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+
+
+if __name__ == "__main__":
+    main()
@@ -6,7 +6,7 @@ import gradio as gr
 import torch

 from fastvideo import VideoGenerator
-from fastvideo.v1.configs.sample.base import SamplingParam
+from fastvideo.configs.sample.base import SamplingParam

 if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="FastVideo Gradio Demo")
@@ -1,7 +1,7 @@
 from fastvideo import VideoGenerator
-from fastvideo.v1.configs.sample import SamplingParam
+from fastvideo.configs.sample import SamplingParam

-OUTPUT_PATH = "./lora"
+OUTPUT_PATH = "./lora_out"
 def main():
    # Initialize VideoGenerator with the Wan model
    generator = VideoGenerator.from_pretrained(
@@ -30,17 +30,8 @@ def main():
        negative_prompt=negative_prompt,
        **kwargs
    )
-    del generator
    
-    # Until FSDP resharding bug is fixed, multi-lora requires reloading the model
-    # see https://github.com/pytorch/pytorch/issues/157209
-    generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        num_gpus=1,
-        lora_path="motimalu/wan-flat-color-1.3b-v2",
-        lora_nickname="flat_color"
-    )
-    # generator.set_lora_adapter(lora_nickname="flat_color", lora_path="motimalu/wan-flat-color-1.3b-v2")
+    generator.set_lora_adapter(lora_nickname="flat_color", lora_path="motimalu/wan-flat-color-1.3b-v2")
    prompt = "flat color, no lineart, blending, negative space, artist:[john kafka|ponsuke kaikai|hara id 21|yoneyama mai|fuzichoco],  1girl, sakura miko, pink hair, cowboy shot, white shirt, floral print, off shoulder, outdoors, cherry blossom, tree shade, wariza, looking up, falling petals, half-closed eyes, white sky, clouds,  live2d animation, upper body, high quality cinematic video of a woman sitting under a sakura tree. Dreamy and lonely, the camera close-ups on the face of the woman as she turns towards the viewer. The Camera is steady, This is a cowboy shot. The animation is smooth and fluid."
    negative_prompt = "bad quality video,色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
    video = generator.generate_video(
@@ -0,0 +1,35 @@
+"""
+Inference using a LoRA checkpoint from FastVideo trainer.
+"""
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "./lora_out"
+def main():
+    # Initialize VideoGenerator with the Wan model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,
+        lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-1250/transformer",
+        lora_nickname="crush_smol"
+    )
+    kwargs = {
+        "height": 480,
+        "width": 832,
+        "num_frames": 77,
+        "guidance_scale": 5.0,
+        "num_inference_steps": 50,
+        "seed": 42,
+    }
+    # Generate video with LoRA style
+    prompt = "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table."
+
+    video = generator.generate_video(
+        prompt,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+        **kwargs
+    )    
+
+if __name__ == "__main__":
+    main()
@@ -1,10 +1,16 @@
-This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
+# Wan2.1-I2V-1.3B-InP Crush-Smol Example
+These are e2e example scripts for finetuning Wan2.1 T2V 1.3B InP on the crush-smol dataset.

-Execute the following commands from `FastVideo/` to run training:
+## Execute the following commands from `FastVideo/` to run training:
+
+### Download crush-smol dataset:

- Download crush-smol dataset:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
- Preprocess the videos and captions into latents:
+
+### Preprocess the videos and captions into latents:
+
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
- Edit the following file and run finetuning:
-`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`
+
+### Edit the following file and run finetuning:
+
+`bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`
@@ -84,7 +84,7 @@ miscellaneous_args=(
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    fastvideo/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -120,7 +120,7 @@ srun torchrun \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    fastvideo/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -5,21 +5,19 @@ MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_i2v_1_3b_inp/"
-VALIDATION_PATH="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"

 torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
    --model_path $MODEL_PATH \
    --data_merge_path $DATA_MERGE_PATH \
    --preprocess_video_batch_size 8 \
+    --seed 42 \
    --max_height 480 \
    --max_width 832 \
    --num_frames 77 \
    --dataloader_num_workers 0 \
    --output_dir=$OUTPUT_DIR \
-    --model_type $MODEL_TYPE \
    --train_fps 16 \
-    --validation_dataset_file $VALIDATION_PATH \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
@@ -1,10 +1,16 @@
-This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
+# Wan2.1-I2V-14B-480P Crush-Smol Example
+These are e2e examples scripts for finetuning Wan2.1 I2V 14B 480P on the crush-smol dataset.

-Execute the following commands from `FastVideo/` to run training:
+## Execute the following commands from `FastVideo/` to run training:
+
+### Download crush-smol dataset:

- Download crush-smol dataset:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
- Preprocess the videos and captions into latents:
+
+### Preprocess the videos and captions into latents:
+
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
- Edit the following file and run finetuning:
+
+### Edit the following file and run finetuning:
+
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`
@@ -15,7 +15,7 @@ NUM_GPUS=8
 # Training arguments
 training_args=(
  --tracker_project_name "wan_i2v_finetune"
-  --output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
+  --output_dir "checkpoints/wan_i2v_finetune"
  --max_train_steps 2000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -84,7 +84,7 @@ miscellaneous_args=(
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    fastvideo/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -49,7 +49,7 @@ VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/
 # Training arguments
 training_args=(
  --tracker_project_name wan_i2v_finetune
-  --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"
+  --output_dir="checkpoints/wan_i2v_finetune"
  --max_train_steps=2000
  --train_batch_size=2
  --train_sp_batch_size 1
@@ -120,7 +120,7 @@ srun torchrun \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_i2v_training_pipeline.py \
+    fastvideo/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -0,0 +1,93 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
+VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
+NUM_GPUS=4
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_i2v_finetune"
+  --output_dir "checkpoints/wan_i2v_finetune"
+  --max_train_steps 2000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 8
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --lora_rank 32
+  --lora_training True
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size $NUM_GPUS
+  --tp_size $NUM_GPUS
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_preprocessed_path "$VALIDATION_DIR"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "1.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --checkpointing_steps 1000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --allow_tf32
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_i2v_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -5,21 +5,19 @@ MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_i2v/"
-VALIDATION_PATH="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"

 torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
    --model_path $MODEL_PATH \
    --data_merge_path $DATA_MERGE_PATH \
    --preprocess_video_batch_size 8 \
+    --seed 42 \
    --max_height 480 \
    --max_width 832 \
    --num_frames 77 \
    --dataloader_num_workers 0 \
    --output_dir=$OUTPUT_DIR \
-    --model_type $MODEL_TYPE \
    --train_fps 16 \
-    --validation_dataset_file $VALIDATION_PATH \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
@@ -1,10 +1,16 @@
-This directory contain e2e examples scripts for finetuning Wan2.1 T2v. 
+# Wan2.1-T2V-1.3B Crush-Smol Example
+These are e2e example scripts for finetuning Wan2.1 T2V 1.3B on the crush-smol dataset.

-Execute the following commands from `FastVideo/` to run training:
+## Execute the following commands from `FastVideo/` to run training:
+
+### Download crush-smol dataset:

- Download crush-smol dataset:
 `bash examples/training/finetune/wan_t2v_1_3b/crush_smol/download_dataset.sh`
- Preprocess the videos and captions into latents:
+
+### Preprocess the videos and captions into latents:
+
 `bash examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh`
- Edit the following file and run finetuning:
+
+### Edit the following file and run finetuning:
+
 `bash examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh`
@@ -15,12 +15,12 @@ NUM_GPUS=4
 # Training arguments
 training_args=(
  --tracker_project_name "wan_t2v_finetune"
-  --output_dir "outputs/wan_t2v_finetune"
+  --output_dir "checkpoints/wan_t2v_finetune"
  --max_train_steps 5000
  --train_batch_size 1
  --train_sp_batch_size 1
  --gradient_accumulation_steps 8
-  --num_latent_t 8 
+  --num_latent_t 20
  --num_height 480
  --num_width 832
  --num_frames 77
@@ -61,7 +61,7 @@ validation_args=(
 optimizer_args=(
  --learning_rate 5e-5
  --mixed_precision "bf16"
-  --checkpointing_steps 6000
+  --checkpointing_steps 1000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
@@ -83,7 +83,7 @@ miscellaneous_args=(
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
-    fastvideo/v1/training/wan_training_pipeline.py \
+    fastvideo/training/wan_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -46,7 +46,7 @@ VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/vali
 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_finetune
-  --output_dir="outputs/wan_t2v_finetune"
+  --output_dir="checkpoints/wan_t2v_finetune"
  --max_train_steps=1000
  --train_batch_size=4
  --train_sp_batch_size 1
@@ -117,7 +117,7 @@ srun torchrun \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/v1/training/wan_training_pipeline.py \
+    fastvideo/training/wan_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
@@ -0,0 +1,93 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
+NUM_GPUS=2
+# export CUDA_VISIBLE_DEVICES=4,5
+
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_t2v_finetune"
+  --output_dir "checkpoints/wan_t2v_finetune_lora"
+  --max_train_steps 5000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 8
+  --num_latent_t 20
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --lora_rank 32
+  --lora_training True
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS 
+  --sp_size $NUM_GPUS 
+  --tp_size $NUM_GPUS
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path $DATA_DIR
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation 
+  --validation_dataset_file $VALIDATION_DATASET_FILE
+  --validation_steps 50
+  --validation_sampling_steps "50" 
+  --validation_guidance_scale "1.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 5e-5
+  --mixed_precision "bf16"
+  --checkpointing_steps 500
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --allow_tf32
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+  --master_port 29501 \
+    fastvideo/training/wan_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -5,21 +5,19 @@ MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_t2v/"
-VALIDATION_PATH="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"

 torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
    --model_path $MODEL_PATH \
    --data_merge_path $DATA_MERGE_PATH \
    --preprocess_video_batch_size 8 \
+    --seed 42 \
    --max_height 480 \
    --max_width 832 \
    --num_frames 77 \
    --dataloader_num_workers 0 \
    --output_dir=$OUTPUT_DIR \
-    --model_type $MODEL_TYPE \
    --train_fps 16 \
-    --validation_dataset_file $VALIDATION_PATH \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
@@ -6,7 +6,7 @@ from typing import Any

 import numpy as np

-from fastvideo.v1.utils import dict_to_3d_list
+from fastvideo.utils import dict_to_3d_list


 def configure_sta(mode: str = 'STA_searching',
@@ -1,6 +1,6 @@
-from fastvideo.v1.configs.pipelines import PipelineConfig
-from fastvideo.v1.configs.sample import SamplingParam
-from fastvideo.v1.entrypoints.video_generator import VideoGenerator
+from fastvideo.configs.pipelines import PipelineConfig
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.entrypoints.video_generator import VideoGenerator
 from fastvideo.version import __version__

-__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"]
+__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"]
@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.attention.layer import (DistributedAttention,
+                                       DistributedAttention_VSA, LocalAttention)
+from fastvideo.attention.selector import get_attn_backend
+
+__all__ = [
+    "DistributedAttention",
+    "LocalAttention",
+    "DistributedAttention_VSA",
+    "AttentionBackend",
+    "AttentionMetadata",
+    "AttentionMetadataBuilder",
+    # "AttentionState",
+    "get_attn_backend",
+]
@@ -6,8 +6,8 @@ from dataclasses import dataclass, fields
 from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar

 if TYPE_CHECKING:
-    from fastvideo.v1.fastvideo_args import FastVideoArgs
-    from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+    from fastvideo.fastvideo_args import FastVideoArgs
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch

 import torch

@@ -6,17 +6,17 @@ from flash_attn import flash_attn_func as flash_attn_2_func
 try:
    from flash_attn_interface import flash_attn_func as flash_attn_3_func

-    # flash_attn 3 has slightly different API: it returns lse by default
-    flash_attn_func = lambda q, k, v, softmax_scale, causal: flash_attn_3_func(
-        q, k, v, softmax_scale, causal)[0]
+    # flash_attn 3 no longer have a different API, see following commit:
+    # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7
+    flash_attn_func = flash_attn_3_func
 except ImportError:
    flash_attn_func = flash_attn_2_func

-from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
-                                                      AttentionImpl,
-                                                      AttentionMetadata,
-                                                      AttentionMetadataBuilder)
-from fastvideo.v1.logger import init_logger
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionImpl,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -3,11 +3,9 @@
 import torch
 from sageattention import sageattn

-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
+from fastvideo.attention.backends.abstract import (  # FlashAttentionMetadata,
+    AttentionBackend, AttentionImpl, AttentionMetadata)
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -2,11 +2,9 @@

 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
+from fastvideo.attention.backends.abstract import (  # FlashAttentionMetadata,
+    AttentionBackend, AttentionImpl, AttentionMetadata)
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -7,17 +7,17 @@ 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.fastvideo_args import FastVideoArgs
-from fastvideo.v1.forward_context import ForwardContext, get_forward_context
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.v1.utils import dict_to_3d_list
+import fastvideo.envs as envs
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionImpl,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.distributed import get_sp_group
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import ForwardContext, get_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.utils import dict_to_3d_list

 logger = init_logger(__name__)

@@ -10,14 +10,14 @@ try:
 except ImportError:
    video_sparse_attn = None

-from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
-                                                      AttentionImpl,
-                                                      AttentionMetadata,
-                                                      AttentionMetadataBuilder)
-from fastvideo.v1.distributed import get_sp_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.attention.backends.abstract import (AttentionBackend,
+                                                   AttentionImpl,
+                                                   AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.distributed import get_sp_group
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch

 logger = init_logger(__name__)

@@ -3,15 +3,14 @@
 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 (
+from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
+from fastvideo.distributed.communication_op import (
    sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D)
-from fastvideo.v1.distributed.parallel_state import (get_sp_parallel_rank,
-                                                     get_sp_world_size)
-from fastvideo.v1.forward_context import ForwardContext, get_forward_context
-from fastvideo.v1.platforms import AttentionBackendEnum
-from fastvideo.v1.utils import get_compute_dtype
+from fastvideo.distributed.parallel_state import (get_sp_parallel_rank,
+                                                  get_sp_world_size)
+from fastvideo.forward_context import ForwardContext, get_forward_context
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.utils import get_compute_dtype


 class DistributedAttention(nn.Module):
@@ -9,11 +9,11 @@ from typing import 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 AttentionBackendEnum, current_platform
-from fastvideo.v1.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
+import fastvideo.envs as envs
+from fastvideo.attention.backends.abstract import AttentionBackend
+from fastvideo.logger import init_logger
+from fastvideo.platforms import AttentionBackendEnum, current_platform
+from fastvideo.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname

 logger = init_logger(__name__)

@@ -0,0 +1,6 @@
+from fastvideo.configs.models.base import ModelConfig
+from fastvideo.configs.models.dits.base import DiTConfig
+from fastvideo.configs.models.encoders.base import EncoderConfig
+from fastvideo.configs.models.vaes.base import VAEConfig
+
+__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"]
@@ -2,7 +2,7 @@
 from dataclasses import dataclass, field, fields
 from typing import Any

-from fastvideo.v1.logger import init_logger
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -0,0 +1,5 @@
+from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
+from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
+from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
+
+__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig"]
@@ -2,9 +2,9 @@
 from dataclasses import dataclass, field
 from typing import Any

-from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.layers.quantization import QuantizationConfig
-from fastvideo.v1.platforms import AttentionBackendEnum
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.platforms import AttentionBackendEnum


@dataclass
@@ -3,7 +3,7 @@ from dataclasses import dataclass, field

 import torch

-from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig


 def is_double_block(n: str, m) -> bool:
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig


@dataclass
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig


 def is_blocks(n: str, m) -> bool:
@@ -0,0 +1,14 @@
+from fastvideo.configs.models.encoders.base import (BaseEncoderOutput,
+                                                    EncoderConfig,
+                                                    ImageEncoderConfig,
+                                                    TextEncoderConfig)
+from fastvideo.configs.models.encoders.clip import (CLIPTextConfig,
+                                                    CLIPVisionConfig)
+from fastvideo.configs.models.encoders.llama import LlamaConfig
+from fastvideo.configs.models.encoders.t5 import T5Config
+
+__all__ = [
+    "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
+    "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig",
+    "T5Config"
+]
@@ -4,9 +4,9 @@ from typing import Any

 import torch

-from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.layers.quantization import QuantizationConfig
-from fastvideo.v1.platforms import AttentionBackendEnum
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.platforms import AttentionBackendEnum


@dataclass
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
-                                                       ImageEncoderConfig,
-                                                       TextEncoderArchConfig,
-                                                       TextEncoderConfig)
+from fastvideo.configs.models.encoders.base import (ImageEncoderArchConfig,
+                                                    ImageEncoderConfig,
+                                                    TextEncoderArchConfig,
+                                                    TextEncoderConfig)


 def _is_transformer_layer(n: str, m) -> bool:
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
-                                                       TextEncoderConfig)
+from fastvideo.configs.models.encoders.base import (TextEncoderArchConfig,
+                                                    TextEncoderConfig)


 def _is_transformer_layer(n: str, m) -> bool:
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
-                                                       TextEncoderConfig)
+from fastvideo.configs.models.encoders.base import (TextEncoderArchConfig,
+                                                    TextEncoderConfig)


 def _is_transformer_layer(n: str, m) -> bool:
@@ -0,0 +1,9 @@
+from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
+from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig
+from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
+
+__all__ = [
+    "HunyuanVAEConfig",
+    "WanVAEConfig",
+    "StepVideoVAEConfig",
+]
@@ -6,8 +6,8 @@ from typing import Any

 import torch

-from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.utils import StoreBoolean
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.utils import StoreBoolean


@dataclass
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
+from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig


@dataclass
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
+from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig


@dataclass
@@ -3,7 +3,7 @@ from dataclasses import dataclass, field

 import torch

-from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
+from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig


@dataclass
@@ -0,0 +1,15 @@
+from fastvideo.configs.pipelines.base import (PipelineConfig,
+                                              SlidingTileAttnConfig)
+from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
+from fastvideo.configs.pipelines.registry import (
+    get_pipeline_config_cls_from_name)
+from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
+from fastvideo.configs.pipelines.wan import (WanI2V480PConfig, WanI2V720PConfig,
+                                             WanT2V480PConfig, WanT2V720PConfig)
+
+__all__ = [
+    "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
+    "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
+    "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig",
+    "get_pipeline_config_cls_from_name"
+]
@@ -7,13 +7,12 @@ from typing import Any, cast

 import torch

-from fastvideo.v1.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
-                                         VAEConfig)
-from fastvideo.v1.configs.models.encoders import BaseEncoderOutput
-from fastvideo.v1.configs.utils import update_config_from_args
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import (FlexibleArgumentParser, StoreBoolean,
-                                shallow_asdict)
+from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
+                                      VAEConfig)
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.utils import update_config_from_args
+from fastvideo.logger import init_logger
+from fastvideo.utils import FlexibleArgumentParser, StoreBoolean, shallow_asdict

 logger = init_logger(__name__)

@@ -73,12 +72,6 @@ class PipelineConfig:
                                  ...] = field(default_factory=lambda:
                                               (postprocess_text, ))

-    # LoRA parameters
-    lora_path: str | None = None
-    lora_nickname: str | None = "default"  # for swapping adapters in the pipeline
-    lora_target_names: list[
-        str] | None = None  # can restrict list of layers to adapt, e.g. ["q_proj"]
-
    # StepVideo specific parameters
    pos_magic: str | None = None
    neg_magic: str | None = None
@@ -89,6 +82,9 @@ class PipelineConfig:
    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
    skip_time_steps: int = 15

+    # DMD parameters
+    dmd_denoising_steps: list[int] | None = field(default=None)
+
    # Compilation
    # enable_torch_compile: bool = False

@@ -206,12 +202,21 @@ class PipelineConfig:
            "Bool for applying scheduler scale in set_timesteps, used in stepvideo",
        )

+        # DMD parameters
+        parser.add_argument(
+            f"--{prefix_with_dot}dmd-denoising-steps",
+            type=parse_int_list,
+            default=PipelineConfig.dmd_denoising_steps,
+            help=
+            "Comma-separated list of denoising steps (e.g., '1000,757,522')",
+        )
+
        # Add VAE configuration arguments
-        from fastvideo.v1.configs.models.vaes.base import VAEConfig
+        from fastvideo.configs.models.vaes.base import VAEConfig
        VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")

        # Add DiT configuration arguments
-        from fastvideo.v1.configs.models.dits.base import DiTConfig
+        from fastvideo.configs.models.dits.base import DiTConfig
        DiTConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}dit-config")

        return parser
@@ -235,7 +240,7 @@ class PipelineConfig:
        """
        use the pipeline class setting from model_path to match the pipeline config
        """
-        from fastvideo.v1.configs.pipelines.registry import (
+        from fastvideo.configs.pipelines.registry import (
            get_pipeline_config_cls_from_name)
        pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)

@@ -250,7 +255,7 @@ class PipelineConfig:
        kwargs: dictionary of kwargs
        config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance
        """
-        from fastvideo.v1.configs.pipelines.registry import (
+        from fastvideo.configs.pipelines.registry import (
            get_pipeline_config_cls_from_name)

        prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip()
@@ -387,3 +392,10 @@ class SlidingTileAttnConfig(PipelineConfig):
    # Additional configuration specific to sliding tile attention
    pad_to_square: bool = False
    use_overlap_optimization: bool = True
+
+
+def parse_int_list(value: str) -> list[int]:
+    """Parse a comma-separated string of integers into a list."""
+    if not value:
+        return []
+    return [int(x.strip()) for x in value.split(",")]
@@ -5,12 +5,12 @@ from typing import TypedDict

 import torch

-from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.v1.configs.models.dits import HunyuanVideoConfig
-from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
-                                                  CLIPTextConfig, LlamaConfig)
-from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
-from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import HunyuanVideoConfig
+from fastvideo.configs.models.encoders import (BaseEncoderOutput,
+                                               CLIPTextConfig, LlamaConfig)
+from fastvideo.configs.models.vaes import HunyuanVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig

 PROMPT_TEMPLATE_ENCODE_VIDEO = (
    "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
@@ -4,17 +4,15 @@
 import os
 from collections.abc import Callable

-from fastvideo.v1.configs.pipelines.base import PipelineConfig
-from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
-                                                    HunyuanConfig)
-from fastvideo.v1.configs.pipelines.stepvideo import StepVideoT2VConfig
-from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
-                                                WanI2V720PConfig,
-                                                WanT2V480PConfig,
-                                                WanT2V720PConfig)
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import (maybe_download_model_index,
-                                verify_model_config_and_directory)
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
+from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
+from fastvideo.configs.pipelines.wan import (FastWanT2V480PConfig,
+                                             WanI2V480PConfig, WanI2V720PConfig,
+                                             WanT2V480PConfig, WanT2V720PConfig)
+from fastvideo.logger import init_logger
+from fastvideo.utils import (maybe_download_model_index,
+                             verify_model_config_and_directory)

 logger = init_logger(__name__)

@@ -27,6 +25,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig,
    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
    # Add other specific weight variants
 }
@@ -36,6 +35,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
    "hunyuan": lambda id: "hunyuan" in id.lower(),
    "wanpipeline": lambda id: "wanpipeline" in id.lower(),
    "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
+    "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(),
    "stepvideo": lambda id: "stepvideo" in id.lower(),
    # Add other pipeline architecture detectors
 }
@@ -47,6 +47,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
    "wanpipeline":
    WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
    "wanimagetovideo": WanI2V480PConfig,
+    "wandmdpipeline": FastWanT2V480PConfig,
    "stepvideo": StepVideoT2VConfig
    # Other fallbacks by architecture
 }
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.models import DiTConfig, VAEConfig
-from fastvideo.v1.configs.models.dits import StepVideoConfig
-from fastvideo.v1.configs.models.vaes import StepVideoVAEConfig
-from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models import DiTConfig, VAEConfig
+from fastvideo.configs.models.dits import StepVideoConfig
+from fastvideo.configs.models.vaes import StepVideoVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig


@dataclass
@@ -4,12 +4,12 @@ from dataclasses import dataclass, field

 import torch

-from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.v1.configs.models.dits import WanVideoConfig
-from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
-                                                  CLIPVisionConfig, T5Config)
-from fastvideo.v1.configs.models.vaes import WanVAEConfig
-from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.configs.models.encoders import (BaseEncoderOutput,
+                                               CLIPVisionConfig, T5Config)
+from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig


 def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
@@ -95,3 +95,19 @@ class WanI2V720PConfig(WanI2V480PConfig):

    # Denoising stage
    flow_shift: int = 5
+
+
+@dataclass
+class FastWanT2V480PConfig(WanT2V480PConfig):
+    """Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD"""
+
+    # WanConfig-specific parameters with defaults
+
+    # Denoising stage
+    flow_shift: int = 8
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000, 757, 522])
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
@@ -0,0 +1,3 @@
+from fastvideo.configs.sample.base import SamplingParam
+
+__all__ = ["SamplingParam"]
@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from typing import Any

-from fastvideo.v1.logger import init_logger
+from fastvideo.logger import init_logger

 logger = init_logger(__name__)

@@ -66,7 +66,7 @@ class SamplingParam:

    @classmethod
    def from_pretrained(cls, model_path: str) -> "SamplingParam":
-        from fastvideo.v1.configs.sample.registry import (
+        from fastvideo.configs.sample.registry import (
            get_sampling_param_cls_for_name)
        sampling_cls = get_sampling_param_cls_for_name(model_path)
        if sampling_cls is not None:
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field

-from fastvideo.v1.configs.sample.base import SamplingParam
-from fastvideo.v1.configs.sample.teacache import TeaCacheParams
+from fastvideo.configs.sample.base import SamplingParam
+from fastvideo.configs.sample.teacache import TeaCacheParams


@dataclass
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
SolitaryThinker	350afecede	update template directions	2025-07-24 21:51:51 -07:00
SolitaryThinker	c868206c40	remove v1 folder	2025-07-24 21:47:57 -07:00
SolitaryThinker	5a4eebdf97	move collect env	2025-07-24 21:33:21 -07:00
SolitaryThinker	f859409844	remove .v1.	2025-07-24 21:31:42 -07:00
SolitaryThinker	ed5c2e0cc0	remove v1	2025-07-24 21:25:20 -07:00
SolitaryThinker	25a91bfb0f	remove v1	2025-07-24 21:22:48 -07:00
Wenxuan TanandWei Feng	48528f468c	[Feature] Multi-lora inference (#640 ) Co-authored-by: Wei (Will) Feng <134637289+weifengpy@users.noreply.github.com>	2025-07-24 21:01:28 -07:00
Yongqi ChenandSolitaryThinker	4264a80ca9	[Feature] Add DMD inference pipeline (#637 ) Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>	2025-07-24 21:00:47 -07:00
William Lin	8573d4f05e	[Docs] Docs update for Training and MPS (#641 )	2025-07-24 19:23:58 -07:00
Wenxuan Tan	210a733515	[Bugfix] Fix LoRA trainable params and training ckpt loading (#630 )	2025-07-23 20:01:40 -07:00
William Lin	0aef0e6f63	[bugfix] Fix preprocessing pipelines and nightly tests (#633 )	2025-07-22 22:44:25 -07:00
Kevin Lin	dd022ad9be	[CI] Fix CI for pull request targets other than main (#632 )	2025-07-22 21:00:03 -07:00
William Lin	832ad61e5b	[bugfix] fa3 no longer returns lse (#631 )	2025-07-22 18:31:34 -07:00
Wenxuan Tan	9419c04ee3	Fix lora train steps (#627 )	2025-07-21 23:29:30 -05:00
Zhang Peiyuanandroot	a37b39d83c	Py/add triton block sparse (#593 ) Co-authored-by: root <a1286225768@gmail,com>	2025-07-17 16:50:40 -07:00
Wenxuan Tan	bb8c769c8e	[LoRA] Support v1 LoRA training (#576 )	2025-07-17 15:52:09 -05:00
William Lin	576c214f28	[v0] Remove V0 code (#621 )	2025-07-15 22:05:16 -07:00
RandNMR73	b79d1fc15b	video gen working on apple silicon (addressed issues from prior pr) (#595 )	2025-07-15 22:04:27 -07:00
William Lin	eb66e1c18d	[chore] release 0.1.2 (#622 )	2025-07-15 14:42:54 -07:00