revert other stages

.buildkite/pipeline.yml

@@ -93,17 +93,6 @@ steps:
                   - TEST_TYPE=training
                 agents:
                   queue: "default"
-            - path:
-                - "fastvideo/training/*distillation_pipeline.py"
-                - "pyproject.toml"
-                - "docker/Dockerfile.python3.12"
-              config:
-                command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Distillation DMDTests"
-                env:
-                  - TEST_TYPE=distillation_dmd
-                agents:
-                  queue: "default"
             - path:
                 - "fastvideo/**"
                 - "pyproject.toml"

.buildkite/scripts/pr_test.sh

@@ -105,10 +105,6 @@ case "$TEST_TYPE" in
         log "Running LoRA tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_lora_tests"
         ;;
-    "distillation_dmd")
-        log "Running distillation DMD tests..."
-        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_distill_dmd_tests"
-        ;;
     *)
         log "Error: Unknown test type: $TEST_TYPE"
         exit 1

.github/workflows/pr-test.yml

@@ -326,7 +326,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] && python csrc/attn/tests/test_vsa.py"
+      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_block_sparse.py"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}

README.md

@@ -130,17 +130,9 @@ We learned and reused code from the following projects:
 We thank MBZUAI and [Anyscale](https://www.anyscale.com/) for their support throughout this project.
 
 ## Citation
-If you use FastVideo for your research, please cite our work:
+If you use FastVideo for your research, please cite our paper:
 
 ```bibtex
-@software{fastvideo2024,
-  title        = {FastVideo: A Unified Framework for Accelerated Video Generation},
-  author       = {The FastVideo Team},
-  url          = {https://github.com/hao-ai-lab/FastVideo},
-  month        = apr,
-  year         = {2024},
-}
-
 @misc{zhang2025vsafastervideodiffusion,
       title={VSA: Faster Video Diffusion with Trainable Sparse Attention}, 
       author={Peiyuan Zhang and Haofeng Huang and Yongqi Chen and Will Lin and Zhengzhong Liu and Ion Stoica and Eric Xing and Hao Zhang},

__init__.py

@@ -1,15 +0,0 @@
-try:
-    from .comfyui.video_generator.nodes import (NODE_CLASS_MAPPINGS,
-                                                NODE_DISPLAY_NAME_MAPPINGS)
-    WEB_DIRECTORY = "./web"
-    __all__ = [
-        'NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'
-    ]
-except ImportError:
-    # ComfyUI environment not available, skip comfyui imports
-    NODE_CLASS_MAPPINGS = {}
-    NODE_DISPLAY_NAME_MAPPINGS = {}
-    WEB_DIRECTORY = "./web"
-    __all__ = [
-        'NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'
-    ]

csrc/attn/README.md

@@ -2,18 +2,20 @@
 
 # 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
+```
 
 ## Video Sparse Attention (VSA)
-
-### Installation
-We support H100 (via TK) and any other GPU (via triton) for 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
@@ -32,35 +34,10 @@ export PATH=${CUDA_HOME}/bin:${PATH}
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 ```
 
-### Verify if you have successfully installed
-
-```bash
-# test numerical
-python tests/test_vsa.py
-# (For H100) test speed
-python benchmarks/bench_vsa_hopper.py
-```
-bench_vsa_hopper.py should print something like this:
-```bash
-Using topk=76 kv blocks per q block (out of 768 total kv blocks)
-
-=== BLOCK SPARSE ATTENTION BENCHMARK ===
-Block Sparse Forward  - TFLOPS: 5622.26
-Block Sparse Backward - TFLOPS: 3865.68
-```
 
 
-## Sliding Tile Attention (STA)
-We only support H100 for STA.
-```bash
-git submodule update --init --recursive
-python setup_sta.py install
-```
-
-
-
-
-###  Usage
+## Usage
+### STA
 End-2-end inference with FastVideo:
 ```bash
 bash scripts/inference/v1_inference_wan_STA.sh
@@ -80,19 +57,22 @@ out = sliding_tile_attention(q, k, v, window_size, text_length)
 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
+
+## Test
 ```bash
 python tests/test_sta.py # test STA
 python tests/test_block_sparse.py # test VSA
 ```
-### Benchmark
+## Benchmark
 ```bash
 python benchmarks/bench_sta.py
 ```
 
 
-### How Does STA Work?
+## How Does STA Work?
 We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 
 
 

csrc/attn/benchmarks/bench_vsa_hopper.py

@@ -1,9 +1,9 @@
 import torch
 import argparse
-from triton.testing import do_bench
+from flash_attn.utils.benchmark import benchmark_forward
 from vsa import block_sparse_fwd, block_sparse_bwd
 from vsa import BLOCK_M, BLOCK_N
-import  triton
+
 import numpy as np
 import random
 
@@ -23,7 +23,7 @@ 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=128, help='Head dimension')
+    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()
@@ -130,19 +130,19 @@ def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_
     
     # Forward pass
     # Warm-up run
-    variable_block_sizes = torch.ones(q2k_block_sparse_index.shape[2], device=q.device).int() * BLOCK_M
-    o, l_vec = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
+    o, l_vec = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
     torch.cuda.synchronize()
     
     # Benchmark forward
-    fwd_time = do_bench(
-        lambda: block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes),
-        warmup=5,
-        rep=20,
-        quantiles=None
+    _, fwd_time = benchmark_forward(
+        block_sparse_fwd,
+        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num,
+        repeats=20,
+        verbose=False,
+        desc='Block Sparse Forward'
     )
     
-    sparse_tflops = flops / fwd_time * 1e-12 * 1e3
+    sparse_tflops = flops / fwd_time.mean * 1e-12
     print(f"Block Sparse Forward - TFLOPS: {sparse_tflops:.2f}")
     
     # Backward pass
@@ -150,19 +150,20 @@ def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_
     
     # Warm-up runs
     for _ in range(5):
-        block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes)
+        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 = do_bench(
-        lambda: block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes),
-        warmup=5,
-        rep=20,
-        quantiles=None
+    _, bwd_time = benchmark_forward(
+        block_sparse_bwd,
+        q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num,
+        repeats=20,
+        verbose=False,
+        desc='Block Sparse Backward'
     )
     bwd_flops = 2.5 * flops  # Approximation
 
-    sparse_bwd_tflops = bwd_flops / bwd_time * 1e-12 * 1e3
+    sparse_bwd_tflops = bwd_flops / bwd_time.mean * 1e-12
     print(f"Block Sparse Backward - TFLOPS: {sparse_bwd_tflops:.2f}")
     
     return sparse_tflops, sparse_bwd_tflops

csrc/attn/tests/test_block_sparse.py

@@ -0,0 +1,289 @@
+import torch
+import argparse
+from flash_attn.utils.benchmark import benchmark_forward
+from flash_attn import flash_attn_func
+from vsa import block_sparse_attn
+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 = 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)
+            # clear memory
+            q_sdpa = q.detach().clone()
+            k_sdpa = k.detach().clone()
+            v_sdpa = v.detach().clone()
+            q_sdpa.requires_grad = True
+            k_sdpa.requires_grad = True
+            v_sdpa.requires_grad = True
+            q.data = torch.empty(0, device=q.device)
+            k.data = torch.empty(0, device=k.device)
+            v.data = torch.empty(0, device=v.device)
+            torch.cuda.empty_cache()
+
+            o_sdpa = torch.nn.functional.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=full_mask)
+
+
+            sim, l1, rmse = precision_metric(o, o_sdpa)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 8e-5, f"l1 too large: {l1}"
+            assert rmse < 2e-5, f"RMSE too large: {rmse}"
+            forward_metrics['sim'].append(sim)
+            forward_metrics['l1'].append(l1)
+            forward_metrics['rmse'].append(rmse)
+
+            print(f"block_sparse_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+
+            # test backward
+            o_sdpa.backward(grad_o)
+
+            sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad)
+            # Error bounds collected on H100
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 4e-3, f"l1 too large: {l1}"
+            assert rmse < 3e-4, f"RMSE too large: {rmse}"
+            grad_q_metrics['sim'].append(sim)
+            grad_q_metrics['l1'].append(l1)
+            grad_q_metrics['rmse'].append(rmse)
+            print(f"block_sparse_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}")        
+            
+            sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 4e-3, f"l1 too large: {l1}"
+            assert rmse < 2e-4, f"RMSE too large: {rmse}"
+            grad_k_metrics['sim'].append(sim)
+            grad_k_metrics['l1'].append(l1)
+            grad_k_metrics['rmse'].append(rmse)
+            print(f"block_sparse_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            
+            sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad)
+            assert sim > 0.9999, f"SSIM too low: {sim}"
+            assert l1 < 1e-4, f"l1 too large: {l1}"
+            assert rmse < 2e-5, f"RMSE too large: {rmse}"
+            grad_v_metrics['sim'].append(sim)
+            grad_v_metrics['l1'].append(l1)
+            grad_v_metrics['rmse'].append(rmse)
+            print(f"block_sparse_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
+            
+            del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        # Print summary statistics if multiple iterations were run
+        if num_iterations > 1:
+            print("\n" + "="*50)
+            print(f"Summary Statistics (over {num_iterations} iterations):")
+            
+            print("\nForward metrics:")
+            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}")
+            
+            print("\nGradient Q metrics:")
+            print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}, min={np.min(grad_q_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}, max={np.max(grad_q_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}, max={np.max(grad_q_metrics['rmse']):.6f}")
+            
+            print("\nGradient K metrics:")
+            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}, min={np.min(grad_k_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}, max={np.max(grad_k_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}, max={np.max(grad_k_metrics['rmse']):.6f}")
+            
+            print("\nGradient V metrics:")
+            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}, min={np.min(grad_v_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}, max={np.max(grad_v_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}, max={np.max(grad_v_metrics['rmse']):.6f}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
+    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
+    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
+    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
+    parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to')
+    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark')
+    parser.add_argument('--num_iterations', type=int, default=50, help='Number of test iterations to run')
+    args = parser.parse_args()
+    main(args)

csrc/attn/tests/test_block_sparse_triton.py

@@ -0,0 +1,289 @@
+import torch
+import argparse
+from flash_attn.utils.benchmark import benchmark_forward
+from flash_attn import flash_attn_func
+from vsa import 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)

csrc/attn/tests/test_full_attn.py

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

csrc/attn/tests/test_triton.py

@@ -0,0 +1,183 @@
+import torch
+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
+    QK = torch.matmul(q_, k_.transpose(-2, -1))
+    QK /= (q_.size(-1) ** 0.5)
+    
+    # Causal mask removed since causal is always false
+    
+    QK = torch.nn.functional.softmax(QK, dim=-1)
+    output = torch.matmul(QK, v_)
+    
+    output.backward(dO_)
+    
+    q_grad = q_.grad
+    k_grad = k_.grad
+    v_grad = v_.grad
+    
+    return output, q_grad, k_grad, v_grad
+
+def fa2_test(Q, K, V, dO):
+    Q.requires_grad = True
+    K.requires_grad = True
+    V.requires_grad = True
+    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)
+
+    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
+    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
+    
+    return scaled_tensor.contiguous()
+
+def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all', test_mode='forward_backward'):
+    results = {
+        'FA2 vs PT': {'sum_diff': 0.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):
+        torch.manual_seed(0)
+        
+        Q  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        K  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        V  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        dO = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
+        
+        pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, dO)
+        fa2_o, fa2_qg, fa2_kg, fa2_vg = fa2_test(Q, K, V, dO)
+        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
+            abs_diff = torch.abs(diff)
+            results['FA2 vs PT']['sum_diff'] += torch.sum(abs_diff).item()
+            results['FA2 vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
+            results['FA2 vs PT']['max_diff'] = max(results['FA2 vs PT']['max_diff'], torch.max(abs_diff).item())
+        
+        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()
+
+    # Calculate total elements based on test mode and error mode
+    if test_mode == 'forward_only':
+        total_elements = b * h * n * d * num_iterations
+    else:  # 'forward_backward'
+        total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4)
+    
+    for name, data in results.items():
+        avg_diff = data['sum_diff'] / total_elements
+        max_diff = data['max_diff']
+        results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff}
+
+    return results
+
+def generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward'):
+    seq_lengths = [768 * (2**i) for i in range(1)]
+
+    print(f"\n{'='*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"{'='*100}")
+    
+    # Print header
+    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 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"{'='*100}\n")
+
+# fix random seed
+torch.manual_seed(0)
+
+mean = 1e-1
+std = 10
+configs = [
+    (4, 1, 128),  # Larger batch, single head, larger dim
+    (2, 8, 64),   # Medium batch, many heads, medium dim
+]
+
+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.")

csrc/attn/tests/test_vsa.py

@@ -1,159 +0,0 @@
-import torch
-import sys
-import os
-import numpy as np
-from tqdm import tqdm
-
-# Add the parent directory to the path to import block_sparse_attn
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from tests.utils import generate_block_sparse_mask_for_function, create_full_mask_from_block_mask
-from vsa import block_sparse_attn
-
-BLOCK_M = 64
-BLOCK_N = 64
-
-def pytorch_test(Q, K, V, block_sparse_mask, dO):
-    q_ = Q.clone().requires_grad_()
-    k_ = K.clone().requires_grad_()
-    v_ = V.clone().requires_grad_()
-
-    QK = torch.matmul(q_, k_.transpose(-2, -1))
-    QK /= (q_.size(-1) ** 0.5)
-    QK = QK.masked_fill(~block_sparse_mask.unsqueeze(0), float('-inf'))
-
-    QK = torch.nn.functional.softmax(QK, dim=-1)
-    output = torch.matmul(QK, v_)
-
-    dO_ = dO
-    output.backward(dO_)
-    return (
-        output.to(torch.bfloat16),
-        q_.grad.to(torch.bfloat16),
-        k_.grad.to(torch.bfloat16),
-        v_.grad.to(torch.bfloat16),
-    )
-
-
-def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, non_pad_index, dO):
-    Q = Q.clone().requires_grad_()
-    K = K.clone().requires_grad_()
-    V = V.clone().requires_grad_()
-    
-    q_padded = vsa_pad(Q, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
-    k_padded = vsa_pad(K, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
-    v_padded = vsa_pad(V, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
-    output, _= block_sparse_attn(q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes)
-    output = output[:, :, non_pad_index, :]
-    output.backward(dO)
-    return output, Q.grad, K.grad, V.grad
-
-
-def get_non_pad_index(
-    vid_len: torch.LongTensor,
-    n_win: int,
-    win_size: int,
-):
-    device = vid_len.device
-    starts_pad = torch.arange(n_win, device=device) * win_size
-    index_pad = starts_pad[:, None] + torch.arange(win_size, device=device)[None, :]
-    index_mask = torch.arange(win_size, device=device)[None, :] < vid_len[:, None]
-
-    return index_pad[index_mask]
-
-def generate_tensor(shape, mean, std, dtype, device):
-    tensor = torch.randn(shape, dtype=dtype, device=device)
-    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
-    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
-    return scaled_tensor.contiguous()
-
-def generate_variable_block_sizes(num_blocks, min_size=32, max_size=64, device="cuda"):
-    return torch.randint(min_size, max_size + 1, (num_blocks,), device=device, dtype=torch.int32)
-
-
-def vsa_pad(x, non_pad_index, num_blocks, block_size):
-    padded_x = torch.zeros((1, x.shape[1], num_blocks * BLOCK_M, x.shape[3]), device=x.device, dtype=x.dtype)
-    padded_x[:, :, non_pad_index, :] = x
-    return padded_x
-
-def check_correctness(h, d, num_blocks, k, mean, std, num_iterations=20, error_mode='all'):
-    results = {
-        'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-        'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-        'gK': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-        'gV': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
-    }
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    variable_block_sizes = generate_variable_block_sizes(num_blocks, device=device)
-    S = int(variable_block_sizes.sum().item())
-    padded_S = num_blocks * BLOCK_M
-    non_pad_index = get_non_pad_index(variable_block_sizes, num_blocks, BLOCK_M)
-    block_mask = generate_block_sparse_mask_for_function(h, num_blocks, k, device)
-    full_mask = create_full_mask_from_block_mask(block_mask, variable_block_sizes, device)
-    for _ in range(num_iterations):
-        Q = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
-        K = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
-        V = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
-        dO = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
-
-        # dO_padded = torch.zeros_like(dO_padded)
-        # dO_padded[:, :, non_pad_index, :] = dO
-
-        pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO)
-        bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), variable_block_sizes,non_pad_index, dO)
-        for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), (pt_o, bs_o)]):
-            if bs is not None:
-                diff = pt - bs
-                abs_diff = torch.abs(diff)
-                results[name]['sum_diff'] += torch.sum(abs_diff).item()
-                results[name]['sum_abs'] += torch.sum(torch.abs(pt)).item()
-                results[name]['max_diff'] = max(results[name]['max_diff'], torch.max(abs_diff).item())
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    total_elements = h * S * d * num_iterations
-    for name, data in results.items():
-        avg_diff = data['sum_diff'] / total_elements
-        max_diff = data['max_diff']
-        results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff}
-
-    return results
-
-def generate_error_graphs(h, d, mean, std, error_mode='all'):
-    test_configs = [
-        {"num_blocks": 16, "k": 2, "description": "Small sequence"},
-        {"num_blocks": 32, "k": 4, "description": "Medium sequence"},
-        {"num_blocks": 53, "k": 6, "description": "Large sequence"},
-    ]
-    
-    print(f"\nError Analysis for h={h}, d={d}, mean={mean}, std={std}, mode={error_mode}")
-    print("=" * 150)
-    print(f"{'Config':<20} {'Blocks':<8} {'K':<4} "
-          f"{'gQ Avg':<12} {'gQ Max':<12} "
-          f"{'gK Avg':<12} {'gK Max':<12} "
-          f"{'gV Avg':<12} {'gV Max':<12} "
-          f"{'gO Avg':<12} {'gO Max':<12}")
-    print("-" * 150)
-
-    for config in test_configs:
-        num_blocks = config["num_blocks"]
-        k = config["k"]
-        description = config["description"]
-        results = check_correctness(h, d, num_blocks, k, mean, std, error_mode=error_mode)
-        print(f"{description:<20} {num_blocks:<8} {k:<4} "
-              f"{results['gQ']['avg_diff']:<12.6e} {results['gQ']['max_diff']:<12.6e} "
-              f"{results['gK']['avg_diff']:<12.6e} {results['gK']['max_diff']:<12.6e} "
-              f"{results['gV']['avg_diff']:<12.6e} {results['gV']['max_diff']:<12.6e} "
-              f"{results['gO']['avg_diff']:<12.6e} {results['gO']['max_diff']:<12.6e}")
-
-    print("-" * 150)
-
-if __name__ == "__main__":
-    h, d = 16, 128
-    mean = 0.0
-    std = 1
-    print("Block Sparse Attention with Variable Block Sizes Analysis")
-    print("=" * 60)
-    for mode in ['backward']:
-        generate_error_graphs(h, d, mean, std, error_mode=mode)
-    print("\nAnalysis completed for all modes.")

csrc/attn/tests/utils.py

@@ -1,54 +0,0 @@
-import torch
-
-def generate_block_sparse_mask_for_function(h, num_blocks, k, device="cuda"):
-    """
-    Generate block sparse mask of shape [h, num_blocks, num_blocks].
-    
-    Args:
-        h: number of heads
-        num_blocks: number of blocks
-        k: number of kv blocks each q block attends to
-        device: device to create tensors on
-
-    Returns:
-        block_sparse_mask: [h, num_blocks, num_blocks] bool tensor
-    """
-    k = min(k, num_blocks)
-    scores = torch.rand(h, num_blocks, num_blocks, device=device)
-    _, indices = torch.topk(scores, k, dim=-1)
-    block_sparse_mask = torch.zeros(h, num_blocks, num_blocks, dtype=torch.bool, device=device)
-
-    block_sparse_mask = block_sparse_mask.scatter_(2, indices, 1).bool()
-    return block_sparse_mask
-
-
-def create_full_mask_from_block_mask(block_sparse_mask, variable_block_sizes, device="cuda"):
-    """
-    Convert block-level sparse mask to full attention mask.
-    
-    Args:
-        block_sparse_mask: [h, num_blocks, num_blocks] bool tensor
-        variable_block_sizes: [num_blocks] tensor
-        device: device to create tensors on
-
-    Returns:
-        full_mask: [h, S, S] bool tensor where S = total sequence length
-    """
-    h, num_blocks, _ = block_sparse_mask.shape
-    total_seq_len = variable_block_sizes.sum().item()
-    cumsum = torch.cat([torch.tensor([0], device=device), variable_block_sizes.cumsum(dim=0)[:-1]])
-
-    full_mask = torch.zeros(h, total_seq_len, total_seq_len, dtype=torch.bool, device=device)
-
-    for head in range(h):
-        for q_block in range(num_blocks):
-            q_start = cumsum[q_block]
-            q_end = q_start + variable_block_sizes[q_block]
-
-            for kv_block in range(num_blocks):
-                if block_sparse_mask[head, q_block, kv_block]:
-                    kv_start = cumsum[kv_block]
-                    kv_end = kv_start + variable_block_sizes[kv_block]
-                    full_mask[head, q_start:q_end, kv_start:kv_end] = True
-
-    return full_mask

csrc/attn/vsa.cpp

@@ -9,10 +9,10 @@
 
 #ifdef TK_COMPILE_BLOCK_SPARSE
 extern std::vector<torch::Tensor> block_sparse_attention_forward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v,  torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num, torch::Tensor block_size
+    torch::Tensor q, torch::Tensor k, torch::Tensor v,  torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num
 ); 
 extern std::vector<torch::Tensor> block_sparse_attention_backward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og, torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num, torch::Tensor block_size
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og, torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num
 );
 #endif
 

csrc/attn/vsa/__init__.py

@@ -1,17 +1,12 @@
 import torch
 from typing import Tuple
-block_sparse_attn=None
-
+from vsa.vsa import block_sparse_attn
 try:
     from vsa_cuda import block_sparse_fwd, block_sparse_bwd
-    from vsa.block_sparse_wrapper import block_sparse_attn_SM90
-    block_sparse_attn = block_sparse_attn_SM90
 except ImportError:
-    from vsa.block_sparse_wrapper import block_sparse_attn_triton
     block_sparse_fwd = None
     block_sparse_bwd = None
-    block_sparse_attn = block_sparse_attn_triton
-
+from vsa.block_sparse_attn_triton import attention as triton_attention, attention_sparse as triton_attention_sparse
 BLOCK_M = 64
 BLOCK_N = 64
 
@@ -27,7 +22,7 @@ def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
     return output, QK
 
 
-def video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight=None):
+def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None):
     """
     q: [batch_size, num_heads, seq_len, head_dim]
     k: [batch_size, num_heads, seq_len, head_dim]
@@ -37,7 +32,7 @@ def video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_
     video_shape: tuple of (T, H, W)
     compress_attn_weight: [batch_size, num_heads, seq_len, head_dim]
     select_attn_weight: [batch_size, num_heads, seq_len, head_dim]
-    NOTE: We assume q, k, v is zero padded!!
+    
     V1 of sparse attention. Include compress attn and sparse attn branch, use average pooling to compress. 
     Assume q, k, v is flattened in this way: [batch_size, num_heads, T//block_size[0], H//block_size[1], W//block_size[2], block_size[0], block_size[1], block_size[2]]
     """
@@ -46,16 +41,16 @@ def video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_
         block_size = (block_size, block_size, block_size)
 
     block_elements = block_size[0] * block_size[1] * block_size[2]
-    assert block_elements == 64
+    assert block_elements % 64 == 0 and block_elements >= 64
     assert q.shape[2] % block_elements == 0
     batch_size, num_heads, seq_len, head_dim = q.shape
     # compress attn
-    q_compress = (q.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(q.dtype)
-    k_compress = (k.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype)
-    v_compress = (v.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype)
+    q_compress = q.view(batch_size, num_heads, seq_len // block_elements,
+                        block_elements, head_dim).mean(dim=3)
+    k_compress = k.view(batch_size, num_heads, seq_len // block_elements,
+                        block_elements, head_dim).mean(dim=3)
+    v_compress = v.view(batch_size, num_heads, seq_len // block_elements,
+                        block_elements, head_dim).mean(dim=3)
 
     output_compress, block_attn_score = torch_attention(q_compress, k_compress,
                                                         v_compress)
@@ -67,9 +62,13 @@ def video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_
                                              1).view(batch_size, num_heads,
                                                      seq_len, head_dim)
 
-    topK_indices = torch.topk(block_attn_score, topk, dim=-1).indices
-    block_mask = torch.zeros_like(block_attn_score, dtype=torch.bool).scatter_(-1, topK_indices, True)
-    output_select, _ = block_sparse_attn(q, k, v, block_mask, variable_block_sizes)
+    q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num = generate_topk_block_sparse_pattern(
+        block_attn_score, topk)
+
+    output_select = block_sparse_attn(q, k, v, q2k_block_sparse_index,
+                                      q2k_block_sparse_num,
+                                      k2q_block_sparse_index,
+                                      k2q_block_sparse_num)
 
     if compress_attn_weight is not None:
         final_output = output_compress * compress_attn_weight + output_select
@@ -77,3 +76,203 @@ def video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_
         final_output = output_compress + output_select
     return final_output
 
+
+def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor,
+                                       topk: int):
+    """
+    Generate a block sparse pattern where each q block attends to exactly topk kv blocks,
+    based on the provided attention scores.
+    
+    Args:
+        block_attn_score: [bs, h, num_q_blocks, num_kv_blocks]
+            Attention scores between query and key blocks
+        topk: int
+            Number of kv blocks each q block attends to
+        
+    Returns:
+        q2k_block_sparse_index: [bs, h, num_q_blocks, topk]
+            Contains the indices of kv blocks that each q block attends to.
+        q2k_block_sparse_num: [bs, h, num_q_blocks]
+            Contains the number of kv blocks that each q block attends to (all equal to topk).
+        k2q_block_sparse_index: [bs, h, num_kv_blocks, max_q_per_kv]
+            Contains the indices of q blocks that attend to each kv block.
+        k2q_block_sparse_num: [bs, h, num_kv_blocks]
+            Contains the number of q blocks that attend to each kv block.
+    """
+    device = block_attn_score.device
+    # Extract dimensions from block_attn_score
+    bs, h, num_q_blocks, num_kv_blocks = block_attn_score.shape
+
+    sorted_result = torch.sort(block_attn_score, dim=-1, descending=True)
+
+    sorted_indice = sorted_result.indices
+
+    q2k_block_sparse_index, _ = torch.sort(sorted_indice[:, :, :, :topk],
+                                           dim=-1)
+    q2k_block_sparse_index = q2k_block_sparse_index.to(dtype=torch.int32)
+    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks),
+                                      topk,
+                                      device=device,
+                                      dtype=torch.int32)
+
+    block_map = topk_index_to_map(q2k_block_sparse_index,
+                                  num_kv_blocks,
+                                  transpose_map=True)
+    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(
+        block_map.transpose(2, 3))
+
+    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
+
+
+## pytorch sdpa version of block sparse ##
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def topk_index_to_map_kernel(
+    map_ptr,
+    index_ptr,
+    map_bs_stride,
+    map_h_stride,
+    map_q_stride,
+    map_kv_stride,
+    index_bs_stride,
+    index_h_stride,
+    index_q_stride,
+    index_kv_stride,
+    topk: tl.constexpr,
+):
+    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
+    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
+
+    for i in tl.static_range(topk):
+        index = tl.load(index_ptr_base + i * index_kv_stride)
+        tl.store(map_ptr_base + index * map_kv_stride, 1.0)
+
+@triton.jit
+def map_to_index_kernel(
+    map_ptr,
+    index_ptr,
+    index_num_ptr,
+    map_bs_stride,
+    map_h_stride,
+    map_q_stride,
+    map_kv_stride,
+    index_bs_stride,
+    index_h_stride,
+    index_q_stride,
+    index_kv_stride,
+    index_num_bs_stride,
+    index_num_h_stride,
+    index_num_q_stride,
+    num_kv_blocks: tl.constexpr,
+):
+    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
+    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
+
+    num = 0
+    for i in tl.static_range(num_kv_blocks):
+        map_entry = tl.load(map_ptr_base + i * map_kv_stride)
+        if map_entry:
+            tl.store(index_ptr_base + num * index_kv_stride, i)
+            num += 1
+
+    tl.store(
+        index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride +
+        q * index_num_q_stride, num)
+
+def topk_index_to_map(index: torch.Tensor,
+                      num_kv_blocks: int,
+                      transpose_map: bool = False):
+    """
+    Convert topk indices to a map.
+    
+    Args:
+        index: [bs, h, num_q_blocks, topk]
+            The topk indices tensor.
+        num_kv_blocks: int
+            The number of key-value blocks in the block_map returned
+        transpose_map: bool
+            If True, the block_map will be transposed on the final two dimensions.
+    
+    Returns:
+        block_map: [bs, h, num_q_blocks, num_kv_blocks]
+            A binary map where 1 indicates that the q block attends to the kv block.
+    """
+    bs, h, num_q_blocks, topk = index.shape
+
+    if transpose_map is False:
+        block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks),
+                                dtype=torch.bool,
+                                device=index.device)
+    else:
+        block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks),
+                                dtype=torch.bool,
+                                device=index.device)
+        block_map = block_map.transpose(2, 3)
+
+    grid = (bs, h, num_q_blocks)
+    topk_index_to_map_kernel[grid](
+        block_map,
+        index,
+        block_map.stride(0),
+        block_map.stride(1),
+        block_map.stride(2),
+        block_map.stride(3),
+        index.stride(0),
+        index.stride(1),
+        index.stride(2),
+        index.stride(3),
+        topk=topk,
+    )
+
+    return block_map
+
+def map_to_index(block_map: torch.Tensor):
+    """
+    Convert a block map to indices and counts.
+    
+    Args:
+        block_map: [bs, h, num_q_blocks, num_kv_blocks]
+            The block map tensor.
+    
+    Returns:
+        index: [bs, h, num_q_blocks, num_kv_blocks]
+            The indices of the blocks.
+        index_num: [bs, h, num_q_blocks]
+            The number of blocks for each q block.
+    """
+    bs, h, num_q_blocks, num_kv_blocks = block_map.shape
+
+    index = torch.full((block_map.shape),
+                       -1,
+                       dtype=torch.int32,
+                       device=block_map.device)
+    index_num = torch.empty((bs, h, num_q_blocks),
+                            dtype=torch.int32,
+                            device=block_map.device)
+
+    grid = (bs, h, num_q_blocks)
+    map_to_index_kernel[grid](
+        block_map,
+        index,
+        index_num,
+        block_map.stride(0),
+        block_map.stride(1),
+        block_map.stride(2),
+        block_map.stride(3),
+        index.stride(0),
+        index.stride(1),
+        index.stride(2),
+        index.stride(3),
+        index_num.stride(0),
+        index_num.stride(1),
+        index_num.stride(2),
+        num_kv_blocks=num_kv_blocks,
+    )
+
+    return index, index_num
+

csrc/attn/vsa/block_sparse_attn_triton.py

@@ -18,6 +18,43 @@ 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.
@@ -34,7 +71,6 @@ configs = [
 @triton.jit
 def _attn_fwd_sparse(Q, K, V, sm_scale,                         #
                      q2k_index, q2k_num, max_kv_blks,           #
-                     variable_block_sizes,
                      M, Out,                                    #
                      stride_qz, stride_qh, stride_qm, stride_qk,
                      stride_kz, stride_kh, stride_kn, stride_kk,
@@ -100,15 +136,12 @@ def _attn_fwd_sparse(Q, K, V, sm_scale,                         #
     # ----- sparse loop over valid K/V tiles -----
     for i in range(0, kv_blocks):
         kv_idx = tl.load(kv_ptr + i).to(tl.int32)
-        block_size = tl.load(variable_block_sizes + kv_idx)
+
         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)
-        # mask out invalid columns
-        mask = tl.arange(0, BLOCK_N) < block_size
-        qk = tl.where(mask[None, :], qk, -float("inf"))
 
         m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
         p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
@@ -130,6 +163,84 @@ def _attn_fwd_sparse(Q, K, V, sm_scale,                         #
 # ──────────────────────────── 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
@@ -156,7 +267,6 @@ def _attn_bwd_dkdv(dk, dv,  #
                    DO,  #
                    M, D,  #
                    k2q_index, k2q_num, max_q_blks,
-                   variable_block_sizes,
                    # shared by Q/K/V/DO.
                    stride_tok, stride_d,  #
                    H, N_CTX, BLOCK_M1: tl.constexpr,  #
@@ -181,8 +291,6 @@ def _attn_bwd_dkdv(dk, dv,  #
 
     q_blocks = tl.load(k2q_num  + meta_base)                 # int32
     q_ptr    = k2q_index + meta_base * max_q_blks           # ptr to list
-    block_size = tl.load(variable_block_sizes + kv_blk)
-    
     
         
     for blk_idx in range(q_blocks*2):
@@ -193,9 +301,6 @@ def _attn_bwd_dkdv(dk, dv,  #
         m = tl.load(M + offs_m)
         qkT = tl.dot(k, qT)
         pT = tl.math.exp2(qkT - m[None, :])
-        mask = tl.arange(0, BLOCK_N1) < block_size
-        pT = tl.where(mask[:, None], pT, 0.0)
-
         do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
         # Compute dV.
         ppT = pT
@@ -219,7 +324,6 @@ def _attn_bwd_dq(dq, q, K, V,  #
                  do, m, D,
                  # shared by Q/K/V/DO.
                  q2k_index, q2k_num, max_kv_blks,
-                 variable_block_sizes,
                  stride_tok, stride_d,  #
                  H, N_CTX,  #
                  BLOCK_M2: tl.constexpr,  #
@@ -251,13 +355,10 @@ def _attn_bwd_dq(dq, q, K, V,  #
     
     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
-        block_size = tl.load(variable_block_sizes + blk_idx//2) - (blk_idx%2) * step_n
         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)
-        mask = tl.arange(0, BLOCK_N2) < block_size.to(tl.int32)
-        p = tl.where(mask[None, :], p , 0.0)
         # Compute dP and dS.
         dp = tl.dot(do, vT).to(tl.float32)
         ds = p * (dp - Di[:, None])
@@ -277,7 +378,6 @@ def _attn_bwd(Q, K, V, sm_scale,  #
               M, D,
               q2k_index, q2k_num, max_kv_blks,
               k2q_index, k2q_num, max_q_blks,
-              variable_block_sizes,
               # shared by Q/K/V/DO.
               stride_z, stride_h, stride_tok, stride_d,  #
               H, N_CTX,  #
@@ -328,7 +428,6 @@ def _attn_bwd(Q, K, V, sm_scale,  #
         DO,  #
         M, D,  #
         k2q_index, k2q_num, max_q_blks,
-        variable_block_sizes,
         stride_tok, stride_d,  #
         H, N_CTX,  #
         BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
@@ -360,7 +459,6 @@ def _attn_bwd(Q, K, V, sm_scale,  #
     dq = _attn_bwd_dq(dq, q, K, V,  #
                       do, m, D,  #
                       q2k_index, q2k_num, max_kv_blks,
-                      variable_block_sizes,
                       stride_tok, stride_d,  #
                       H, N_CTX,  #
                       BLOCK_M2, BLOCK_N2, HEAD_DIM,  #
@@ -373,77 +471,237 @@ def _attn_bwd(Q, K, V, sm_scale,  #
     
     
 
+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 ───────────────────────────
-def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes):
-    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)
+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).
+    """
 
-    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,
-        variable_block_sizes,
-        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
-    )
+    @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)
 
-    return o, M
+        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
+        )
 
-def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num, variable_block_sizes):
-    assert do.is_contiguous()
-    assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
-    
-    B, H, T, D = q.shape
-    sm_scale = 1.0 / math.sqrt(D)
-    dq = torch.empty_like(q)
-    dk = torch.empty_like(k)
-    dv = torch.empty_like(v)
-    BATCH, N_HEAD, N_CTX = q.shape[:3]
-    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 * (sm_scale * RCP_LN2)
-    PRE_BLOCK = 64
-    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=D  #
-    )
-    
-    
-    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, sm_scale, do, dq, dk, dv,  #
-        M, delta,  #
-        q2k_index, q2k_num, max_kv_blks,
-        k2q_index, k2q_num, max_q_blks,
-        variable_block_sizes,
-        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=D  #
-    )
+        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
 
-    return dq, dk, dv
+    @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)

csrc/attn/vsa/block_sparse_h100.cu

@@ -46,9 +46,225 @@ template<int D> struct fwd_globals {
 
     int32_t *__restrict__  q2k_block_sparse_index;
     int32_t *__restrict__  q2k_block_sparse_num;
-    int32_t *__restrict__  block_size;
 };
 
+template<int D>
+__global__  __launch_bounds__(128, 3) // encourage compiler to reduce register usage so that an SM can hold 3 CTAs. Performance will drop from 391T to 353T if not specified explicitly.
+void fwd_attend_ker_even(const __grid_constant__ fwd_globals<D> g) {
+    extern __shared__ int __shm[]; 
+    tma_swizzle_allocator al((int*)&__shm[0]);
+
+    using K = fwd_attend_ker_tile_dims<D>;
+
+    using q_tile    =         st_bf<64, K::tile_width>;
+    using k_tile    =         st_bf<128, K::tile_width>;
+    using v_tile    =         st_bf<128, K::tile_width>;
+    using k_tile_half    =         st_bf<64, K::tile_width>;
+    using v_tile_half    =         st_bf<64, K::tile_width>;
+    using l_col_vec = col_vec<st_fl<64, K::tile_width>>;
+    using o_tile    =         st_bf<64, K::tile_width>;
+    
+    q_tile    (&q_smem)[1] = al.allocate<q_tile, 1>();
+
+    k_tile_half    (&k_smem_0)[1]           = al.allocate<k_tile_half, 1          >();
+    k_tile_half    (&k_smem_1)[1]           = al.allocate<k_tile_half, 1          >();
+    k_tile    (*k_smem)           = reinterpret_cast<k_tile(*)>(k_smem_0);
+
+    v_tile_half    (&v_smem_0)[1]           = al.allocate<v_tile_half, 1          >();
+    v_tile_half    (&v_smem_1)[1]           = al.allocate<v_tile_half, 1          >();
+    v_tile    (*v_smem)           = reinterpret_cast<v_tile(*)>(v_smem_0);
+
+    l_col_vec (&l_smem)[1] = al.allocate<l_col_vec, 1>();
+
+    auto      (*o_smem)                      = reinterpret_cast<o_tile(*)>(q_smem);
+    
+    int kv_head_idx = blockIdx.y / g.hr;
+    int seq_idx     = blockIdx.x; 
+
+
+    int32_t* q2k_block_sparse_index_ptr = g.q2k_block_sparse_index + blockIdx.z * gridDim.y * gridDim.x * g.max_kv_blocks_per_q + blockIdx.y * gridDim.x * g.max_kv_blocks_per_q  + blockIdx.x * g.max_kv_blocks_per_q;
+    int32_t* q2k_block_sparse_num_ptr = g.q2k_block_sparse_num + blockIdx.z * gridDim.y * gridDim.x  + blockIdx.y * gridDim.x + blockIdx.x;
+    int32_t kv_blocks   = q2k_block_sparse_num_ptr[0] / 2; // each iter load 2 kv blocks
+    __shared__ kittens::semaphore qsmem_semaphore, k_smem_arrived, v_smem_arrived;
+    if (threadIdx.x == 0) {
+        int32_t kv_block_index[2];
+        reinterpret_cast<float2*>(kv_block_index)[0] = reinterpret_cast<float2*>(q2k_block_sparse_index_ptr)[0];
+        
+        init_semaphore(qsmem_semaphore, 0, 1);
+        init_semaphore(k_smem_arrived, 0, 1);
+        init_semaphore(v_smem_arrived, 0, 1);
+        
+        // preload q block
+        coord<q_tile> q_tile_idx = {blockIdx.z, blockIdx.y, seq_idx, 0};
+        tma::expect_bytes(qsmem_semaphore, sizeof(q_smem));
+        tma::load_async(q_smem[0], g.q, q_tile_idx, qsmem_semaphore);
+        
+        // preload the zeroth block of kv
+        tma::expect_bytes(k_smem_arrived, sizeof(k_tile));
+        coord<k_tile_half> k_tile_idx_0 = {blockIdx.z, kv_head_idx, kv_block_index[0], 0};
+        coord<k_tile_half> k_tile_idx_1 = {blockIdx.z, kv_head_idx, kv_block_index[1], 0};
+        tma::load_async(k_smem_0[0], g.k, k_tile_idx_0, k_smem_arrived);
+        tma::load_async(k_smem_1[0], g.k, k_tile_idx_1, k_smem_arrived);
+
+        tma::expect_bytes(v_smem_arrived, sizeof(v_tile));
+        coord<v_tile_half> v_tile_idx_0 = {blockIdx.z, kv_head_idx, kv_block_index[0], 0};
+        coord<v_tile_half> v_tile_idx_1 = {blockIdx.z, kv_head_idx, kv_block_index[1], 0};
+        tma::load_async(v_smem_0[0], g.v, v_tile_idx_0, v_smem_arrived);
+        tma::load_async(v_smem_1[0], g.v, v_tile_idx_1, v_smem_arrived);
+    }
+    __syncthreads(); 
+
+    rt_fl<16, 128>  att_block;
+    rt_bf<16, 128>  att_block_mma;
+    rt_fl<16, K::tile_width> o_reg;
+    
+    col_vec<rt_fl<16, 128>> max_vec, norm_vec, max_vec_last_scaled, max_vec_scaled;
+    
+    neg_infty(max_vec);
+    zero(norm_vec);
+    zero(o_reg);
+
+    // wait for q block
+    wait(qsmem_semaphore, 0);
+
+    for (int kv_idx = 0; kv_idx < kv_blocks - 1; kv_idx++) {
+        // preload kv index
+        int32_t kv_block_index[2];
+        reinterpret_cast<float2*>(kv_block_index)[0] = reinterpret_cast<float2*>(q2k_block_sparse_index_ptr)[kv_idx + 1];
+        
+        // wait k
+        wait(k_smem_arrived, kv_idx % 2);
+
+        // compute QK^T
+        warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]);
+
+        copy(max_vec_last_scaled, max_vec);
+        if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); }
+        else                   { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); }
+        
+        warpgroup::mma_async_wait();
+
+        // load K
+        if (threadIdx.x == 0) {
+            tma::expect_bytes(k_smem_arrived, sizeof(k_tile));
+            coord<k_tile_half> k_tile_idx_0 = {blockIdx.z, kv_head_idx, kv_block_index[0], 0};
+            coord<k_tile_half> k_tile_idx_1 = {blockIdx.z, kv_head_idx, kv_block_index[1], 0};
+            tma::load_async(k_smem_0[0], g.k, k_tile_idx_0, k_smem_arrived);
+            tma::load_async(k_smem_1[0], g.k, k_tile_idx_1, k_smem_arrived);
+        }
+
+        // exp
+        row_max(max_vec, att_block, max_vec);
+            
+        if constexpr (D == 64) { 
+            mul(att_block, att_block,    1.44269504089f*0.125f); 
+            mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f);
+        }
+        else                   { 
+            mul(att_block, att_block,    1.44269504089f*0.08838834764f); 
+            mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f);
+        }
+
+        sub_row(att_block, att_block, max_vec_scaled);
+        exp2(att_block, att_block);
+        sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled);
+        exp2(max_vec_last_scaled,       max_vec_last_scaled);
+        mul(norm_vec,            norm_vec,     max_vec_last_scaled);
+        row_sum(norm_vec,  att_block, norm_vec);
+        add(att_block, att_block, 0.f);
+        copy(att_block_mma, att_block);
+        mul_row(o_reg, o_reg, max_vec_last_scaled); 
+        
+        // wait v
+        wait(v_smem_arrived, kv_idx % 2);
+
+        // compute SV
+        warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]);
+        warpgroup::mma_async_wait();
+
+        // load V
+        if (threadIdx.x == 0) {
+            tma::expect_bytes(v_smem_arrived, sizeof(v_tile));
+            // coord<v_tile> v_tile_idx = {blockIdx.z, kv_head_idx, kv_idx + 1, 0};
+            // tma::load_async(v_smem[0], g.v, v_tile_idx, v_smem_arrived);
+            coord<v_tile_half> v_tile_idx_0 = {blockIdx.z, kv_head_idx, kv_block_index[0], 0};
+            coord<v_tile_half> v_tile_idx_1 = {blockIdx.z, kv_head_idx, kv_block_index[1], 0};
+            tma::load_async(v_smem_0[0], g.v, v_tile_idx_0, v_smem_arrived);
+            tma::load_async(v_smem_1[0], g.v, v_tile_idx_1, v_smem_arrived);
+        }
+    }
+
+    // last iter
+    {
+        int kv_idx = kv_blocks - 1;
+        // wait k
+        wait(k_smem_arrived, kv_idx % 2);
+
+        // compute QK^T
+        warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]);
+
+        copy(max_vec_last_scaled, max_vec);
+        if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); }
+        else                   { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); }
+        
+        warpgroup::mma_async_wait();      
+
+        // exp
+        row_max(max_vec, att_block, max_vec);
+            
+        if constexpr (D == 64) { 
+            mul(att_block, att_block,    1.44269504089f*0.125f); 
+            mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f);
+        }
+        else                   { 
+            mul(att_block, att_block,    1.44269504089f*0.08838834764f); 
+            mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f);
+        }
+
+        sub_row(att_block, att_block, max_vec_scaled);
+        exp2(att_block, att_block);
+        sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled);
+        exp2(max_vec_last_scaled,       max_vec_last_scaled);
+        mul(norm_vec,            norm_vec,     max_vec_last_scaled);
+        row_sum(norm_vec,  att_block, norm_vec);
+        add(att_block, att_block, 0.f);
+        copy(att_block_mma, att_block); 
+        mul_row(o_reg, o_reg, max_vec_last_scaled); 
+        
+        // wait v
+        wait(v_smem_arrived, kv_idx % 2);
+
+        // compute SV
+        warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]);
+        warpgroup::mma_async_wait();
+    }
+
+    div_row(o_reg, o_reg, norm_vec);
+    warpgroup::store(o_smem[0], o_reg); 
+    __syncthreads();
+
+    // TK store_async internally calls syncwarp so we need to route on warp level
+    if (threadIdx.x / 32 == 0) {
+        coord<o_tile> o_tile_idx = {blockIdx.z, blockIdx.y, seq_idx, 0};
+        tma::store_async(g.o, o_smem[0], o_tile_idx);
+    }
+
+    mul(max_vec_scaled,   max_vec_scaled, 0.69314718056f);
+    log(norm_vec, norm_vec);
+    add(norm_vec, norm_vec, max_vec_scaled);
+
+    if constexpr (D == 64) { mul(norm_vec, norm_vec, -8.0f); }
+    else                   { mul(norm_vec, norm_vec, -11.313708499f); }
+
+    warpgroup::store(l_smem[0], norm_vec);
+    __syncthreads();
+
+    if (threadIdx.x / 32 == 0) {
+        coord<l_col_vec> tile_idx = {blockIdx.z, blockIdx.y, 0, seq_idx};
+        tma::store_async(g.l, l_smem[0], tile_idx);
+    }
+    tma::store_async_wait();
+}
 
 template<int D>
 __global__  __launch_bounds__(128, 4)
@@ -141,7 +357,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) { // use block siz
         }
 
         // exp
-        right_fill(att_block, att_block, g.block_size[q2k_block_sparse_index_ptr[kv_idx]], base_types::constants<float>::neg_infty());
         row_max(max_vec, att_block, max_vec);
             
         if constexpr (D == 64) { 
@@ -194,8 +409,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) { // use block siz
         warpgroup::mma_async_wait();      
 
         // exp
-        right_fill(att_block, att_block, g.block_size[q2k_block_sparse_index_ptr[kv_idx]], base_types::constants<float>::neg_infty());
-
         row_max(max_vec, att_block, max_vec);
             
         if constexpr (D == 64) { 
@@ -271,9 +484,8 @@ struct bwd_prep_globals {
     d_gl  d;
 };
 
-constexpr int PREP_NUM_WARPS = (1);
 template<int D>
-__global__  __launch_bounds__(PREP_NUM_WARPS*kittens::WARP_THREADS, (D == 64) ? 6 / PREP_NUM_WARPS : 3 / PREP_NUM_WARPS)
+__global__  __launch_bounds__(4*kittens::WARP_THREADS, (D == 64) ? 2 : 1)
 void bwd_attend_prep_ker(const __grid_constant__ bwd_prep_globals<D> g) {
     extern __shared__ int __shm[]; 
     tma_swizzle_allocator al((int*)&__shm[0]);
@@ -284,9 +496,9 @@ void bwd_attend_prep_ker(const __grid_constant__ bwd_prep_globals<D> g) {
     using o_tile  = st_bf<4*16, D>;
     using d_tile  = col_vec<st_fl<4*16, D>>;
 
-    og_tile (&og_smem)[PREP_NUM_WARPS] = al.allocate<og_tile, PREP_NUM_WARPS>();
-    o_tile  (&o_smem) [PREP_NUM_WARPS] = al.allocate<o_tile , PREP_NUM_WARPS>();
-    d_tile  (&d_smem) [PREP_NUM_WARPS] = al.allocate<d_tile , PREP_NUM_WARPS>();
+    og_tile (&og_smem)[4] = al.allocate<og_tile, 4>();
+    o_tile  (&o_smem) [4] = al.allocate<o_tile , 4>();
+    d_tile  (&d_smem) [4] = al.allocate<d_tile , 4>();
     
     rt_fl<4*16, D> og_reg, o_reg; 
     col_vec<rt_fl<4*16, D>> d_reg;
@@ -295,13 +507,13 @@ void bwd_attend_prep_ker(const __grid_constant__ bwd_prep_globals<D> g) {
 
     if (threadIdx.x == 0) {
         init_semaphore(smem_semaphore, 0, 1);
-        tma::expect_bytes(smem_semaphore, sizeof(og_smem[0]) * PREP_NUM_WARPS * 2);
+        tma::expect_bytes(smem_semaphore, sizeof(og_smem[0]) * 4 * 2);
     }
     __syncthreads();
 
     if (warpid == 0) {
-        for (int w = 0; w < PREP_NUM_WARPS; w++) {
-            coord<o_tile> tile_idx = {blockIdx.z, blockIdx.y, (blockIdx.x * PREP_NUM_WARPS) + w, 0};
+        for (int w = 0; w < 4; w++) {
+            coord<o_tile> tile_idx = {blockIdx.z, blockIdx.y, (blockIdx.x * 4) + w, 0};
             tma::load_async(o_smem[w],  g.o,  tile_idx, smem_semaphore);
             tma::load_async(og_smem[w], g.og, tile_idx, smem_semaphore);
         }
@@ -316,8 +528,8 @@ void bwd_attend_prep_ker(const __grid_constant__ bwd_prep_globals<D> g) {
     __syncthreads(); 
 
     if (warpid == 0) {
-        for (int w = 0; w < PREP_NUM_WARPS; w++) {
-            coord<d_tile> tile_idx = {blockIdx.z, blockIdx.y, 0, (blockIdx.x * PREP_NUM_WARPS) + w};
+        for (int w = 0; w < 4; w++) {
+            coord<d_tile> tile_idx = {blockIdx.z, blockIdx.y, 0, (blockIdx.x * 4) + w};
             tma::store_async(g.d, d_smem[w], tile_idx);
         }
     }
@@ -379,7 +591,6 @@ struct bwd_globals {
 
     int32_t *__restrict__ k2q_block_sparse_index;
     int32_t *__restrict__ k2q_block_sparse_num;
-    int32_t *__restrict__ block_size;
 };
 
 __device__ static inline void
@@ -502,7 +713,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
     
     // wait for kv
     wait(kv_b, 0);
-    int fill_start = g.block_size[blockIdx.x] - 16 * kittens::warpid();
+
     for (int qo_idx = 0; qo_idx < qo_blocks - 1; qo_idx++) {
         // preload q index
         store_qg_block_index = load_q_block_index;
@@ -521,8 +732,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
 
         if constexpr (D == 64) { mul(s_block_t, s_block_t, 1.44269504089f*0.125f); }
         else                   { mul(s_block_t, s_block_t, 1.44269504089f*0.08838834764f); }
-
-        lower_fill(s_block_t, s_block_t, fill_start, base_types::constants<float>::neg_infty());
+    
         exp2(s_block_t, s_block_t); // P_i
         copy(p_block_t, s_block_t);
         copy(p_block_t_mma, s_block_t);
@@ -600,7 +810,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
 
         if constexpr (D == 64) { mul(s_block_t, s_block_t, 1.44269504089f*0.125f); }
         else                   { mul(s_block_t, s_block_t, 1.44269504089f*0.08838834764f); }
-        lower_fill(s_block_t, s_block_t, fill_start, base_types::constants<float>::neg_infty());
+    
         exp2(s_block_t, s_block_t); // P_i
         copy(p_block_t, s_block_t);
         copy(p_block_t_mma, s_block_t);
@@ -666,14 +876,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
 #include <iostream>
 
 std::vector<torch::Tensor> 
-block_sparse_attention_forward(
-    torch::Tensor q, 
-    torch::Tensor k, 
-    torch::Tensor v, 
-    torch::Tensor q2k_block_sparse_index, 
-    torch::Tensor q2k_block_sparse_num, 
-    torch::Tensor block_size
-)
+block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num)
 {
     CHECK_INPUT(q);
     CHECK_INPUT(k);
@@ -685,10 +888,6 @@ block_sparse_attention_forward(
     auto qo_heads = q.size(1);
     auto kv_heads = k.size(1);
     auto max_kv_blocks_per_q = q2k_block_sparse_index.size(3);
-    auto num_q_blocks = block_size.size(0);
-    TORCH_CHECK(batch==1, "Batch size dim will be removed in the future, please set batch to 1");
-    TORCH_CHECK(num_q_blocks * 64 == seq_len, "This kernel supports variable block size, but it assumes the input sequence is properly padded.");
-    TORCH_CHECK(num_q_blocks == q2k_block_sparse_index.size(2), "Number of Q blocks does not match between q2k_block_sparse_index and block_size");
     // check to see that these dimensions match for all inputs
     TORCH_CHECK(q.size(0) == batch, "Q batch dimension - idx 0 - must match for all inputs");
     TORCH_CHECK(k.size(0) == batch, "K batch dimension - idx 0 - must match for all inputs");
@@ -768,19 +967,7 @@ block_sparse_attention_forward(
         l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,   static_cast<unsigned int>(seq_len)};
         o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
 
-        globals g{
-            qg_arg, 
-            kg_arg, 
-            vg_arg, 
-            lg_arg, 
-            og_arg, 
-            static_cast<int>(seq_len), 
-            static_cast<int>(hr), 
-            static_cast<int>(max_kv_blocks_per_q), 
-            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), 
-            reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr()),
-            reinterpret_cast<int32_t*>(block_size.data_ptr())
-        };
+        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr())};
 
         constexpr int mem_size = 54000;
 
@@ -819,19 +1006,7 @@ block_sparse_attention_forward(
         l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,   static_cast<unsigned int>(seq_len)};
         o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
 
-        globals g{
-            qg_arg, 
-            kg_arg, 
-            vg_arg, 
-            lg_arg, 
-            og_arg, 
-            static_cast<int>(seq_len), 
-            static_cast<int>(hr), 
-            static_cast<int>(max_kv_blocks_per_q), 
-            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), 
-            reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr()),
-            reinterpret_cast<int32_t*>(block_size.data_ptr())
-        };
+        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr())};
 
         constexpr int mem_size = 54000;
 
@@ -861,8 +1036,7 @@ block_sparse_attention_backward(torch::Tensor q,
                    torch::Tensor l_vec,
                    torch::Tensor og,
                    torch::Tensor k2q_block_sparse_index,
-                   torch::Tensor k2q_block_sparse_num,
-                   torch::Tensor block_size)
+                   torch::Tensor k2q_block_sparse_num)
 {
     CHECK_INPUT(q);
     CHECK_INPUT(k);
@@ -875,7 +1049,7 @@ block_sparse_attention_backward(torch::Tensor q,
     auto seq_len  = q.size(2);
     auto head_dim = q.size(3);
     auto max_q_blocks_per_kv = k2q_block_sparse_index.size(3);
-    TORCH_CHECK(k2q_block_sparse_index.size(2) == block_size.size(0), "k2q_block_sparse_index.size(2) must match block_size.size(0)");
+
     // check to see that these dimensions match for all inputs
     TORCH_CHECK(q.size(0)     == batch, "Q  batch dimension - idx 0 - must match for all inputs");
     TORCH_CHECK(k.size(0)     == batch, "K  batch dimension - idx 0 - must match for all inputs");
@@ -962,7 +1136,7 @@ block_sparse_attention_backward(torch::Tensor q,
     float* d_vg = reinterpret_cast<float*>(vg_ptr);
 
     constexpr int mem_size = kittens::MAX_SHARED_MEMORY; 
-    int threads  = PREP_NUM_WARPS * kittens::WARP_THREADS;
+    int threads  = 4 * kittens::WARP_THREADS;
 
     //cudadevicesynchronize();
     const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
@@ -971,7 +1145,7 @@ block_sparse_attention_backward(torch::Tensor q,
     //  cudaStreamSynchronize(stream);
 
     // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256");
-    dim3 grid_bwd(seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
+    dim3 grid_bwd(seq_len/(4*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
 
     if (head_dim == 64)  {
         using og_tile = st_bf<4*16, 64>;
@@ -1046,8 +1220,8 @@ block_sparse_attention_backward(torch::Tensor q,
                         static_cast<int>(hr),
                         static_cast<int>(max_q_blocks_per_kv),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_index.data_ptr()),
-                        reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr()),
-                        reinterpret_cast<int32_t*>(block_size.data_ptr())};
+                        reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr())
+                        };
 
         dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
         threads = 128;
@@ -1150,8 +1324,8 @@ block_sparse_attention_backward(torch::Tensor q,
                         static_cast<int>(hr),
                         static_cast<int>(max_q_blocks_per_kv),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_index.data_ptr()),
-                        reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr()),
-                        reinterpret_cast<int32_t*>(block_size.data_ptr())};
+                        reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr())
+                        };
 
         dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
         threads = 128;

csrc/attn/vsa/block_sparse_wrapper.py

@@ -1,182 +0,0 @@
-import torch
-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 triton_block_sparse_attn_forward, triton_block_sparse_attn_backward
-assert torch.__version__ >= "2.4.0", "VSA requires PyTorch 2.4.0 or higher"
-from vsa.index import map_to_index
-from typing import Tuple, Optional
-
-
-
-@torch.library.custom_op("vsa::block_sparse_attn_SM90", mutates_args=(), device_types="cuda")
-def block_sparse_attn_SM90(
-    q_padded: torch.Tensor,
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor, 
-    )-> Tuple[torch.Tensor, torch.Tensor]:
-    q_padded = q_padded.contiguous()
-    k_padded = k_padded.contiguous()
-    v_padded = v_padded.contiguous()
-    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
-    variable_block_sizes = variable_block_sizes.int()
-    o_padded, lse_padded = block_sparse_fwd(q_padded, k_padded, v_padded, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
-    return o_padded, lse_padded
-
-
-
-
-@torch.library.register_fake("vsa::block_sparse_attn_SM90")
-def _block_sparse_attn_SM90_fake(
-    q_padded: torch.Tensor,
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor, 
-) -> Tuple[torch.Tensor,  torch.Tensor]:
-    q_padded, k_padded, v_padded = [x.contiguous() for x in (q_padded, k_padded, v_padded)]
-    B, H, S, D = q_padded.shape
-    o_padded = torch.empty_like(q_padded)
-    lse_padded = torch.empty((B, H, S, 1), device=q_padded.device, dtype=torch.float32)
-    return o_padded, lse_padded
-
-
-@torch.library.custom_op("vsa::block_sparse_attn_backward_SM90", mutates_args=(), device_types="cuda")
-def block_sparse_attn_backward_SM90(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    lse_padded: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor, 
-    )-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    grad_output_padded = grad_output_padded.contiguous()
-    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
-    grad_q_padded, grad_k_padded, grad_v_padded = block_sparse_bwd(
-        q_padded, k_padded, v_padded, o_padded, lse_padded, grad_output_padded, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes
-    )
-    grad_q_padded = grad_q_padded.to(grad_output_padded.dtype)
-    grad_k_padded = grad_k_padded.to(grad_output_padded.dtype)
-    grad_v_padded = grad_v_padded.to(grad_output_padded.dtype)
-    return grad_q_padded, grad_k_padded, grad_v_padded
-
-@torch.library.register_fake("vsa::block_sparse_attn_backward_SM90")
-def _block_sparse_attn_backward_SM90_fake(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    lse_padded: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor, 
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    torch._check(grad_output_padded.dtype == torch.bfloat16)
-    torch._check(lse_padded.dtype == torch.float32)
-    grad_output_padded = grad_output_padded.contiguous()
-    dq = torch.empty_like(grad_output_padded)
-    dk = torch.empty_like(grad_output_padded)
-    dv = torch.empty_like(grad_output_padded)
-    return dq, dk, dv
-
-
-def backward_SM90(ctx, grad_output1, grad_output2):
-    q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes= ctx.saved_tensors
-    dq, dk, dv = block_sparse_attn_backward_SM90(grad_output1, q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
-    return dq, dk, dv, None, None
-
-def setup_context_SM90(ctx, inputs, output):
-    q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
-    o_padded, lse_padded = output
-    ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
-
-
-block_sparse_attn_SM90.register_autograd(backward_SM90, setup_context=setup_context_SM90)
-
-
-@torch.library.custom_op("vsa::block_sparse_attn_triton", mutates_args=(), device_types="cuda")
-def block_sparse_attn_triton(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    q = q.contiguous()
-    k = k.contiguous()
-    v = v.contiguous()
-    block_map = block_map.int()
-    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
-    o, M = triton_block_sparse_attn_forward(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
-    return o, M
-
-
-@torch.library.register_fake("vsa::block_sparse_attn_triton")
-def _block_sparse_attn_triton_fake(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    q = q.contiguous()
-    k = k.contiguous()
-    v = v.contiguous()
-    o = torch.empty_like(q)
-    M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
-    return o, M
-
-
-@torch.library.custom_op("vsa::block_sparse_attn_backward_triton", mutates_args=(), device_types="cuda")
-def block_sparse_attn_backward_triton(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    M: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    grad_output_padded = grad_output_padded.contiguous()
-    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
-    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
-    dq, dk, dv = triton_block_sparse_attn_backward(grad_output_padded, q_padded, k_padded, v_padded, o_padded, M, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes)
-    return dq, dk, dv
-
-@torch.library.register_fake("vsa::block_sparse_attn_backward_triton")
-def _block_sparse_attn_backward_triton_fake(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    M: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    grad_output_padded = grad_output_padded.contiguous()
-    dq = torch.empty_like(grad_output_padded)
-    dk = torch.empty_like(grad_output_padded)
-    dv = torch.empty_like(grad_output_padded)
-    return dq, dk, dv
-
-
-def backward_triton(ctx, grad_output1, grad_output2):
-    q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes = ctx.saved_tensors
-    dq, dk, dv = block_sparse_attn_backward_triton(grad_output1, q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes)
-    return dq, dk, dv, None, None
-
-
-def setup_context_triton(ctx, inputs, output):
-    q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
-    o_padded, M = output
-    ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes)
-    
-block_sparse_attn_triton.register_autograd(backward_triton, setup_context=setup_context_triton)

csrc/attn/vsa/index.py

@@ -1,152 +0,0 @@
-
-## pytorch sdpa version of block sparse ##
-import triton
-import triton.language as tl
-import torch
-
-@triton.jit
-def topk_index_to_map_kernel(
-    map_ptr,
-    index_ptr,
-    map_bs_stride,
-    map_h_stride,
-    map_q_stride,
-    map_kv_stride,
-    index_bs_stride,
-    index_h_stride,
-    index_q_stride,
-    index_kv_stride,
-    topk,
-):
-    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
-    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
-
-    for i in tl.static_range(topk):
-        index = tl.load(index_ptr_base + i * index_kv_stride)
-        tl.store(map_ptr_base + index * map_kv_stride, 1.0)
-
-@triton.jit
-def map_to_index_kernel(
-    map_ptr,
-    index_ptr,
-    index_num_ptr,
-    map_bs_stride,
-    map_h_stride,
-    map_q_stride,
-    map_kv_stride,
-    index_bs_stride,
-    index_h_stride,
-    index_q_stride,
-    index_kv_stride,
-    index_num_bs_stride,
-    index_num_h_stride,
-    index_num_q_stride,
-    num_kv_blocks,
-):
-    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
-    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
-
-    num = 0
-    for i in tl.range(num_kv_blocks):
-        map_entry = tl.load(map_ptr_base + i * map_kv_stride)
-        if map_entry:
-            tl.store(index_ptr_base + num * index_kv_stride, i)
-            num += 1
-
-    tl.store(
-        index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride +
-        q * index_num_q_stride, num)
-
-def topk_index_to_map(index: torch.Tensor,
-                      num_kv_blocks: int,
-                      transpose_map: bool = False):
-    """
-    Convert topk indices to a map.
-    
-    Args:
-        index: [bs, h, num_q_blocks, topk]
-            The topk indices tensor.
-        num_kv_blocks: int
-            The number of key-value blocks in the block_map returned
-        transpose_map: bool
-            If True, the block_map will be transposed on the final two dimensions.
-    
-    Returns:
-        block_map: [bs, h, num_q_blocks, num_kv_blocks]
-            A binary map where 1 indicates that the q block attends to the kv block.
-    """
-    bs, h, num_q_blocks, topk = index.shape
-
-    if transpose_map is False:
-        block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks),
-                                dtype=torch.bool,
-                                device=index.device)
-    else:
-        block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks),
-                                dtype=torch.bool,
-                                device=index.device)
-        block_map = block_map.transpose(2, 3)
-
-    grid = (bs, h, num_q_blocks)
-    topk_index_to_map_kernel[grid](
-        block_map,
-        index,
-        block_map.stride(0),
-        block_map.stride(1),
-        block_map.stride(2),
-        block_map.stride(3),
-        index.stride(0),
-        index.stride(1),
-        index.stride(2),
-        index.stride(3),
-        topk=topk,
-    )
-
-    return block_map
-
-def map_to_index(block_map: torch.Tensor):
-    """
-    Convert a block map to indices and counts.
-    
-    Args:
-        block_map: [bs, h, num_q_blocks, num_kv_blocks]
-            The block map tensor.
-    
-    Returns:
-        index: [bs, h, num_q_blocks, num_kv_blocks]
-            The indices of the blocks.
-        index_num: [bs, h, num_q_blocks]
-            The number of blocks for each q block.
-    """
-    bs, h, num_q_blocks, num_kv_blocks = block_map.shape
-
-    index = torch.full((block_map.shape),
-                       -1,
-                       dtype=torch.int32,
-                       device=block_map.device)
-    index_num = torch.empty((bs, h, num_q_blocks),
-                            dtype=torch.int32,
-                            device=block_map.device)
-
-    grid = (bs, h, num_q_blocks)
-    map_to_index_kernel[grid](
-        block_map,
-        index,
-        index_num,
-        block_map.stride(0),
-        block_map.stride(1),
-        block_map.stride(2),
-        block_map.stride(3),
-        index.stride(0),
-        index.stride(1),
-        index.stride(2),
-        index.stride(3),
-        index_num.stride(0),
-        index_num.stride(1),
-        index_num.stride(2),
-        num_kv_blocks=num_kv_blocks,
-    )
-
-    return index, index_num

csrc/attn/vsa/vsa.py

@@ -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)

examples/distill/Wan-Syn-480P/README.md

@@ -1,14 +1,6 @@
 # Wan2.1-T2V-1.3B Distill Example
 These are end-to-end example scripts for distilling Wan2.1 T2V 1.3B model using DMD-only and DMD+VSA methods.
 
-### 0. Make sure you have installed VSA
-
-```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
-```
-
 ### 1. Download dataset:
 ```bash
 bash examples/distill/Wan-Syn-480P/download_dataset.sh

examples/distill/Wan-Syn-480P/distill_dmd_VSA_t2v.slurm

@@ -47,15 +47,15 @@ VALIDATION_DATASET_FILE=your_validation_dataset_file
 # Training arguments
 training_args=(
   --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir"checkpoints/wan_t2v_finetune"
-  --max_train_steps 4000
-  --train_batch_size 1
+  --output_dir="checkpoints/wan_t2v_finetune"
+  --max_train_steps=4000
+  --train_batch_size=1
   --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 480
+  --gradient_accumulation_steps=1
+  --num_latent_t 16
+  --num_height 448
   --num_width 832
-  --num_frames 81
+  --num_frames 61
   --enable_gradient_checkpointing_type "full"
 )
 
@@ -91,10 +91,9 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 500
-  --weight_only_checkpointing_steps 500
+  --learning_rate=1e-5
+  --mixed_precision="bf16"
+  --checkpointing_steps=500
   --weight_decay 0.01
   --max_grad_norm 1.0
 )
@@ -105,7 +104,7 @@ miscellaneous_args=(
   --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.0
-  --dit_precision "fp32"
+  --dit_precision "bf16"
   --ema_start_step 0
   --flow_shift 8
   --seed 1000
@@ -127,7 +126,7 @@ srun torchrun \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/wan_distillation_pipeline.py \
+    fastvideo/training/wan_training_pipeline.py \
     "${parallel_args[@]}" \
     "${model_args[@]}" \
     "${dataset_args[@]}" \

examples/distill/Wan-Syn-480P/distill_dmd_t2v.slurm

@@ -47,15 +47,15 @@ VALIDATION_DATASET_FILE=your_validation_dataset_file
 # Training arguments
 training_args=(
   --tracker_project_name wan_t2v_distill_dmd
-  --output_dir "checkpoints/wan_t2v_finetune"
-  --max_train_steps 4000
-  --train_batch_size 1
+  --output_dir="checkpoints/wan_t2v_finetune"
+  --max_train_steps=4000
+  --train_batch_size=1
   --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 480
+  --gradient_accumulation_steps=1
+  --num_latent_t 16
+  --num_height 448
   --num_width 832
-  --num_frames 81
+  --num_frames 61
   --enable_gradient_checkpointing_type "full"
 )
 
@@ -91,10 +91,9 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 500
-  --weight_only_checkpointing_steps 500
+  --learning_rate=1e-5
+  --mixed_precision="bf16"
+  --checkpointing_steps=500
   --weight_decay 0.01
   --max_grad_norm 1.0
 )
@@ -105,7 +104,7 @@ miscellaneous_args=(
   --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.0
-  --dit_precision "fp32"
+  --dit_precision "bf16"
   --ema_start_step 0
   --flow_shift 8
   --seed 1000
@@ -126,7 +125,7 @@ srun torchrun \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/wan_distillation_pipeline.py \
+    fastvideo/training/wan_training_pipeline.py \
     "${parallel_args[@]}" \
     "${model_args[@]}" \
     "${dataset_args[@]}" \

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm

@@ -1,138 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=t2v
-#SBATCH --partition=main
-#SBATCH --nodes=8
-#SBATCH --ntasks=8
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=dmd_t2v_output/t2v_%j.out
-#SBATCH --error=dmd_t2v_output/t2v_%j.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate your_env
-
-# Basic Info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
-DATA_DIR=your_data_dir
-VALIDATION_DATASET_FILE=your_validation_dataset_file
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir "checkpoints/wan_t2v_finetune"
-  --max_train_steps 4000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 480
-  --num_width 832
-  --num_frames 81
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus 64
-  --sp_size 4
-  --tp_size 1
-  --hsdp_replicate_dim 8
-  --hsdp_shard_dim 8
-)
-
-# 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 4
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 200
-  --validation_sampling_steps "3"
-  --validation_guidance_scale "1.0" # not used for dmd inference
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 500
-  --weight_only_checkpointing_steps 500
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 3
-  --seed 1000
-)
-
-# DMD arguments
-dmd_args=(
-  --dmd_denoising_steps '1000,757,522'
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.98
-  --generator_update_interval 5
-  --real_score_guidance_scale 3.5
-  --VSA_sparsity 0.9
-)
-
-srun torchrun \
---nnodes $SLURM_JOB_NUM_NODES \
---nproc_per_node $NUM_GPUS \
---node_rank $SLURM_PROCID \
---rdzv_backend=c10d \
---rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/wan_distillation_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/README.md

@@ -1,12 +0,0 @@
-# Wan2.2-5B Distill Example
-These are end-to-end example scripts for distilling Wan2.2 TI2V 5B model DMD+VSA methods.
-
-### 0. Make sure you have installed VSA
-
-```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
-```
-
-### TODO

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh

@@ -1,110 +0,0 @@
-#!/bin/bash
-
-# Basic Info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-export MASTER_PORT=29500
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-# Configs
-NUM_GPUS=1
-MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir="checkpoints/wan_t2v_finetune"
-  --max_train_steps=4000
-  --train_batch_size=1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps=1
-  --num_latent_t 31
-  --num_height 704
-  --num_width 1280
-  --num_frames 121
-  --enable_gradient_checkpointing_type "full"
-  --training_state_checkpointing_steps=500
-  --weight_only_checkpointing_steps=500
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus 1
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 1
-)
-
-# 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 4
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 200
-  --validation_sampling_steps "3"
-  --validation_guidance_scale "1.0" # not used for dmd inference
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate=1e-5
-  --mixed_precision="bf16"
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 8
-  --seed 1000
-)
-
-# DMD arguments
-dmd_args=(
-  --dmd_denoising_steps '1000,757,522'
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.98
-  --generator_update_interval 5
-  --real_score_guidance_scale 3.5
-  --VSA_sparsity 0.8
-)
-
-torchrun \
---nnodes 1 \
---nproc_per_node $NUM_GPUS \
-    fastvideo/training/wan_distillation_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/preprocess_wan_data_ti2v_5b.sh

@@ -1,24 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1 # 2,4,8
-MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-MODEL_TYPE="wan"
-DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_ti2v/"
-
-torchrun --nproc_per_node=$GPU_NUM \
-    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 704 \
-    --max_width 1280 \
-    --num_frames 121 \
-    --dataloader_num_workers 0 \
-    --output_dir=$OUTPUT_DIR \
-    --train_fps 24 \
-    --samples_per_file 8 \
-    --flush_frequency 8 \
-    --video_length_tolerance_range 5 \
-    --preprocess_task "t2v" 

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json

@@ -1,31 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 704,
-      "width": 1280,
-      "num_frames": 121
-    },
-    {
-      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 704,
-      "width": 1280,
-      "num_frames": 121
-    },
-    {
-      "caption": "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.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 704,
-      "width": 1280,
-      "num_frames": 121
-    }
-  ]
-}

examples/distillation/dmd/Wan2.1-Fun-1.3B-InP/crush_smol/README.md

@@ -0,0 +1,3 @@
+# DMD Distillation Wan2.1-I2V-14B-480P Crush-Smol Example
+
+Coming soon!

examples/inference/basic/basic.py

@@ -11,13 +11,11 @@ def main():
     generator = VideoGenerator.from_pretrained(
         "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         # FastVideo will automatically handle distributed setup
-        num_gpus=1,
+        num_gpus=2,
         use_fsdp_inference=True,
         dit_cpu_offload=False,
         vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
-        pin_cpu_memory=False,
+        text_encoder_cpu_offload=False,
         # image_encoder_cpu_offload=False,
     )
 

examples/inference/basic/basic_dmd.py

@@ -16,8 +16,7 @@ def main():
         num_gpus=1,
         use_fsdp_inference=True,
         # Adjust these offload parameters if you have < 32GB of VRAM
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
+        text_encoder_cpu_offload=False,
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         VSA_sparsity=0.8,
@@ -29,7 +28,9 @@ def main():
     sampling_param = SamplingParam.from_pretrained(model_name)
 
     prompt = (
-        "A neon-lit alley in futuristic Tokyo during a heavy rainstorm at night. The puddles reflect glowing signs in kanji, advertising ramen, karaoke, and VR arcades. A woman in a translucent raincoat walks briskly with an LED umbrella. Steam rises from a street food cart, and a cat darts across the screen. Raindrops are visible on the camera lens, creating a cinematic bokeh effect."
+        "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)

examples/inference/basic/basic_mps.py

@@ -18,7 +18,7 @@ def main():
 
     # 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 = 25  # Reduce from default 81 to 25 frames bc we have to use the SDPA attn backend for mps
+    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
 

examples/inference/basic/default_args.py

@@ -0,0 +1,62 @@
+from fastvideo import VideoGenerator
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+def main():
+
+    # This is the config class for the model initialization
+    config = PipelineConfig.from_pretrained("FastVideo/FastHunyuan-Diffusers")
+    # can be used to dump the config to a yaml file
+    config.dump_to_yaml("config.yaml")
+    print(config)
+    # {
+    #    'vae_config': {
+    #        'scale_factor': 8,
+    #        'sp': True,
+    #        'tiling': True,
+    #        'precision': 'fp16'
+    #    },
+    #    'text_encoder_config': {
+    #        'precision': 'fp16'
+    #    },
+    #    'dit_config': {
+    #        'precision': 'fp16'
+    #    },
+    #    'inference_args': {
+    #        'guidance_scale': 7.5,
+    #        'num_inference_steps': 5,
+    #        'seed': 1024,
+    #        'guidance_rescale': 0.0,
+    #        'flow_shift': 17,
+    #        'num_inference_steps': 5,
+    #    }
+    # }
+
+    config.vae_config.scale_factor = 16
+
+    # FastVideo will automatically used 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.
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/FastHunyuan-Diffusers",
+        num_gpus=4,
+        config=config,
+        # or
+        config_path="config.yaml",
+    )
+
+    sampling_param = SamplingParam.from_pretrained(
+        "FastVideo/FastHunyuan-Diffusers")
+    sampling_param.num_inference_steps = 5
+
+    # Generate videos with the same simple API, regardless of GPU count
+    prompt = "A beautiful woman in a red dress walking down a street"
+    video = generator.generate_video(prompt,
+                                     sampling_param=sampling_param,
+                                     num_inference_steps=6)
+
+    video2 = generator.generate_video(prompt2)
+    prompt2 = "A beautiful woman in a blue dress walking down a street"
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/test.py

@@ -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
+        dit_cpu_offload=True,          
+        text_encoder_cpu_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()

examples/inference/lora/wan_lora_inference.py

@@ -7,10 +7,6 @@ def main():
     generator = VideoGenerator.from_pretrained(
         "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         num_gpus=1,
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
         lora_path="benjamin-paine/steamboat-willie-1.3b",
         lora_nickname="steamboat"
     )

examples/inference/lora/wan_lora_inference_from_ckpt.py

@@ -10,10 +10,6 @@ def main():
     generator = VideoGenerator.from_pretrained(
         "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         num_gpus=1,
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
         lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-1250/transformer",
         lora_nickname="crush_smol"
     )

examples/inference/optimizations/attention_example.py

@@ -11,10 +11,6 @@ def main():
     gen = VideoGenerator.from_pretrained(
         model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         num_gpus=1,
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
     )
     load_time = time.perf_counter() - start_time
     print(f"Model loading time: {load_time:.2f} seconds")

examples/inference/optimizations/teacache_example.py

@@ -10,9 +10,6 @@ def main():
         model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         num_gpus=1,
         dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
     )
     load_time = time.perf_counter() - start_time
     print(f"Model loading time: {load_time:.2f} seconds")

examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm

@@ -49,11 +49,11 @@ VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_sm
 # Training arguments
 training_args=(
   --tracker_project_name wan_i2v_finetune
-  --output_dir "checkpoints/wan_i2v_finetune"
-  --max_train_steps 2000
-  --train_batch_size 2
+  --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"
+  --max_train_steps=2000
+  --train_batch_size=2
   --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
+  --gradient_accumulation_steps=1
   --num_latent_t 8
   --num_height 480
   --num_width 832
@@ -93,9 +93,9 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 1000
+  --learning_rate=1e-5
+  --mixed_precision="bf16"
+  --checkpointing_steps=1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm

@@ -1,134 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=i2v
-#SBATCH --partition=main
-#SBATCH --nodes=8
-#SBATCH --ntasks=8
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=VSA_i2v_output/i2v_%j.out
-#SBATCH --error=VSA_i2v_output/i2v_%j.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate your_env
-
-# Basic Info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
-DATA_DIR=your_data_dir
-VALIDATION_DATASET_FILE=your_validation_dataset_file
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_i2v_VSA
-  --output_dir "checkpoints/wan_i2v_finetune_VSA"
-  --max_train_steps 4000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 720
-  --num_width 1280
-  --num_frames 81
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus 64
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 8
-  --hsdp_shard_dim 8
-)
-
-# 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 4
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
-  --validation_sampling_steps "40"
-  --validation_guidance_scale "5.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-6
-  --mixed_precision "bf16"
-  --checkpointing_steps 1000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 3
-  --seed 1000
-)
-
-# VSA arguments
-vsa_args=(
-  --VSA_decay_rate 0.03 \
-  --VSA_decay_interval_steps 30 \
-  --VSA_sparsity 0.9 \
-)
-
-srun torchrun \
---nnodes $SLURM_JOB_NUM_NODES \
---nproc_per_node $NUM_GPUS \
---node_rank $SLURM_PROCID \
---rdzv_backend=c10d \
---rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/wan_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/README.md

@@ -1,26 +0,0 @@
-# Finetuning Wan2.1 to make it work with VSA to accelerate 
-
-These are e2e example scripts for finetuning Wan2.1 T2V with VSA to accelerate inference.
-
-## Execute the following commands from `FastVideo/` to run training:
-
-## Make sure you have installed VSA
-
-```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
-```
-
-### Download the synthetic dataset:
-
-```bash
-bash examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/download_dataset.sh
-```
-### Slurm script to train the model
-```bash
-sbatch examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm
-```
-```bash
-sbatch examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm
-```

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm

@@ -1,134 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=t2v
-#SBATCH --partition=main
-#SBATCH --nodes=8
-#SBATCH --ntasks=8
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=VSA_t2v_output/t2v_%j.out
-#SBATCH --error=VSA_t2v_output/t2v_%j.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate your_env
-
-# Basic Info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
-DATA_DIR=your_data_dir
-VALIDATION_DATASET_FILE=your_validation_dataset_file
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name wan_t2v_VSA
-  --output_dir "checkpoints/wan_t2v_finetune_VSA"
-  --max_train_steps 4000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 720
-  --num_width 1280
-  --num_frames 81
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus 64
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 8
-  --hsdp_shard_dim 8
-)
-
-# 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 4
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "5.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 1000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 5
-  --seed 1000
-)
-
-# VSA arguments
-vsa_args=(
-  --VSA_decay_rate 0.03 \
-  --VSA_decay_interval_steps 30 \
-  --VSA_sparsity 0.9 \
-)
-
-srun torchrun \
---nnodes $SLURM_JOB_NUM_NODES \
---nproc_per_node $NUM_GPUS \
---node_rank $SLURM_PROCID \
---rdzv_backend=c10d \
---rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/wan_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/download_dataset.sh

@@ -1,7 +0,0 @@
-#!/bin/bash
-
-# 480P dataset
-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
-
-# 720P dataset
-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x768x1280_250k" --local_dir "FastVideo/Wan-Syn_77x768x1280_250k" --repo_type "dataset"

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_64.json

@@ -1,516 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "In the video, a woman is elegantly showcasing her earrings, bringing attention to their intricate design with a gentle touch of her fingers. She is bathed in ambient purple and pink lighting, which casts a soft glow on her delicate features and enhances the vivid tones of her lipstick and eye makeup. Her hair is styled to frame her face smoothly, emphasizing the contours of her jawline and cheekbones. The background features a blurred neon light, adding an artistic and modern touch to the overall aesthetic.",
-      "video_path": "Fashion/mixkit-face-of-an-elegant-and-captivating-woman-41914_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.",
-      "video_path": "Man/mixkit-a-rancher-riding-a-horse-at-sunset-1143_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In a dimly lit, eerie setting, a mysterious pink bottle labeled \"Authentic 100% organic POISON\" sits prominently in the foreground, casting a menacing aura. The bottle is accentuated by green fog, which swirls lightly around it, enhancing its sinister allure. Behind it, a shadowy golden bottle adorned with a spider emblem subtly emerges, adding an extra layer of mystery to the scene. Dim candles provide faint, flickering light, which complements the dark atmosphere, making the setting ideal for an illusion of hidden dangers.",
-      "video_path": "smoke/mixkit-poison-in-halloween-ritual-33879_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "The video opens with a tranquil scene in the heart of a dense forest, emphasizing two large, textured tree trunks in the foreground framing the view. Sunlight filters through the canopy above, casting intricate patterns of light and shadow on the trees and the ground. Between the tree trunks, a clear view of a calm, muddy river unfolds, its surface shimmering under the gentle sunlight. The riverbank is decorated with a variety of small bushes and vibrant foliage, subtly transitioning into the deep greens of tall, leafy plants. In the background, the dense forest looms, filled with dark, towering trees, their branches intertwining to form an intricate canopy. The scene is bathed in the soft glow of the sun, creating a serene and picturesque setting. Occasional sunbeams pierce through the foliage, adding a magical aura to the landscape. The vibrant reds and oranges of the smaller plants add contrast, bringing warmth to the earthy tones of the scenery. Overall, this harmonious blend of natural elements creates a peaceful and idyllic forest setting.",
-      "video_path": "forest/mixkit-view-of-a-river-between-two-old-trees-560_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the video, a martial artist dressed in a traditional white uniform with a black belt demonstrates a series of precise movements against a stark black background. The individual gracefully transitions between stances, embodying a sense of focused discipline and control. Each motion is executed with a deliberate pace, showcasing the fluidity of martial arts techniques. The soft lighting creates subtle highlights on the uniform, adding depth to the figure as it moves. The practitioner begins with an open-hand pose, feet firmly grounded, gradually shifting to a powerful forward punch. The fluidity of the sequence displays a mastery of balance and poise. Every trajectory of the limbs is precise and deliberate, capturing the elegance and strength of martial arts. The serene, isolated setting enhances the intensity and concentration of the practitioner. This visual presentation is an elegant interplay of motion and stillness, displaying the art form's discipline and grace.",
-      "video_path": "Man/mixkit-a-young-man-practicing-his-karate-moves-49635_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A tranquil coastal scene unfolds with a drone's aerial view capturing a serene beach landscape. The camera glides over a quiet stretch of sandy shoreline, where gentle waves kiss the shore under a clear blue sky. Nestled amidst lush palm trees are a series of traditional thatched-roof huts, their earthy tones blending harmoniously with the natural surroundings. The sandy beach stretches endlessly, bordered by the rhythmic dance of ocean waves on one side and verdant greenery on the other. A pair of white umbrellas is set up on the sand, suggesting a place to relax and enjoy the sun. In the distance, two small human figures can be seen walking leisurely along the water's edge, leaving faint footprints behind them. The scene exudes a calm and inviting atmosphere, with the soft rustle of palm leaves and the whisper of the ocean breeze almost audible. The overall composition is a captivating blend of nature's tranquility and architectural simplicity. This picturesque setting invites viewers to imagine themselves steps away from this idyllic coastal escape.",
-      "video_path": "beach/mixkit-sunny-beach-in-a-dynamic-shot-from-a-drone-44383_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A lone figure stands on a large, moss-covered rock, surrounded by the soft rush of a nearby stream. The figure is wearing white sneakers and shorts, with a plaid shirt that hangs loosely in the breeze. The lighting creates dramatic shadows, enhancing the textures of the rock and the subtle movement of the water below. In the background, a waterfall cascades into the stream, completing this tranquil and serene nature scene.",
-      "video_path": "forest/mixkit-woman-standing-in-front-of-waterfall-559_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In an industrial setting, a person leans casually against a railing, exuding a sense of confidence and composure. They are wearing a striking outfit, consisting of a vibrant, patterned jacket over a simple white crop top, creating a bold contrast. The atmosphere is infused with warm, ambient lighting that casts soft shadows on the concrete walls and metallic surfaces. Intricate wiring and pipes form an intricate backdrop, enhancing the urban aesthetic. Their relaxed posture and direct, engaging gaze suggest a sense of ease in this industrial environment. This scene encapsulates a blend of modern fashion and gritty, urban architecture, creating a visually compelling narrative.",
-      "video_path": "Fashion/mixkit-portrait-of-a-hipster-woman-walking-down-a-stairs-1297_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A man is energetically stretching in an open-air setting, surrounded by rows of vibrant red seats that suggest an amphitheater or outdoor venue. He wears a sleeveless black shirt layered with a hooded vest, emphasizing his athletic build as he engages in a warm-up routine. Behind him, the striking modern architecture of the building features geometric panels, with large sections of glass and overlapping metallic beams creating a dynamic backdrop. The scene captures the contrast between his focused movements and the static, bold design of the structure, while the surrounding greenery adds a touch of nature to the environment. The overall atmosphere is one of preparation and anticipation, with the man appearing determined and ready for an upcoming event or performance.",
-      "video_path": "Sport/mixkit-man-doing-arm-stretches-595_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A young woman is seated on the floor in front of a plush, beige tufted couch, fully engrossed in sorting through a stack of papers. Her dark hair falls loosely past her shoulders, and she wears a green plaid shirt, contributing to the casual yet focused atmosphere. She gently places the papers onto a small round white table, occasionally lifting individual sheets to examine them more closely. Her expression shifts subtly, reflecting concentration and contemplation as she processes the information on the pages. Two small, round nested tables hold her documents, along with a small plant in a gray pot, adding a touch of greenery to the scene. The background features a dark paneled wall, creating a contrasting backdrop for the light-colored furniture. The setting is tranquil and organized, the couch and tables arranged symmetrically, conveying a sense of harmony. A calculator rests on the smaller table, hinting at a task involving calculations or budgeting.",
-      "video_path": "Woman/mixkit-frustrated-woman-throws-paperwork-on-the-floor-4526_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A heavily rusted metal gate stands firmly locked, with two vertical bars joined by a thick, old chain that loops elegantly around them. The chain's texture is coarse and rugged, its surface reflecting varying shades of orange and brown, indicative of years exposed to the elements. At the heart of the chain, a black iron padlock, slightly worn yet imposing, secures the gate, its curves and edges smooth against the aged links. The gate's metalwork is outlined by a backdrop of soft, blurred greenery, suggesting a serene and isolated location beyond the barrier. Tall trees rise in the distance, their trunks and leaves creating a lush, forest-like setting that contrasts with the gate's severe rust. A pathway leads away from the gate, its surface uneven with patches of moss and weathered stone visible in the soft focus, inviting yet inaccessible. The ambiance is quiet and mysterious, with a sense of abandonment hanging subtly in the air, evoking curiosity about what lies beyond. Shadows play across the gate, cast by branches swaying gently in the breeze, adding to the dynamic interaction of light and texture. This scene, rich in detail and atmosphere, captures the viewer's imagination, evoking both the allure of the forbidden and the beauty of decay.",
-      "video_path": "forest/mixkit-rusty-fence-with-a-chain-of-a-property-in-nature-5294_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In a serene and softly lit yoga studio, three individuals engage in a yoga session, each performing an upward-facing stretch. The central figure is a woman with shoulder-length brown hair, dressed in a light cropped top and green leggings, her posture reflecting grace and concentration. To her right, another participant, a woman in a purple outfit, mirrors the pose with equal poise. On her left, a person with a bun focuses intently, supported slightly by yoga blocks beneath their hands. The warm-colored wooden floor contrasts soothingly with the soft pastel mural on the back wall, featuring an abstract design and partial visage of a serene face. Natural light floods the space from a large window on the right, where lush greens peek through, adding an element of tranquility. In the corner of the room, a collection of meditation instruments, including a gong and a Buddha statue, subtly frame the peaceful setting. The mood is calm yet focused, as all three participants are deeply engaged in their practice. The scene combines elements of balance, harmony, and a shared journey towards mindfulness. This depiction captures the essence of a yoga session that blends personal growth with collective experience.",
-      "video_path": "People/mixkit-small-group-of-people-doing-yoga-together-43730_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the deep blue expanse of the ocean, two dolphins glide effortlessly, their sleek bodies reflecting the sunlight filtering through the water. The prominent shadows and caustics create a shimmering effect on their skin, capturing the beauty of their natural habitat. Each dolphin moves with a fluid grace, occasionally interacting with gentle nudges, showcasing their playful and social nature. The scene is vibrant and dynamic, with the clear blue background accentuating the dolphins' movements, making it an ideal subject for AI recreation.",
-      "video_path": "sea/mixkit-dolphins-underwater-4133_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the video, a young woman stands against a vibrant graffiti-covered wall, deeply engrossed in her smartphone. Her expression reflects a mix of focus and subtle satisfaction as she interacts with the screen. She wears a black floral-patterned top, which contrasts with the bright, abstract shapes and bold colors of the mural behind her. As she continues to engage with her phone, a series of like count notifications appear on the screen, indicating a growing online appreciation. The wall behind her features a striking mix of geometric and organic shapes, including swirls of teal, orange, and black, with large humanoid figures in a pop-art style. Her long, light-brown hair frames her face, adding a calm, composed aura amidst the lively backdrop. The video captures a blend of contemporary digital interaction and expressive urban art, creating a dynamic yet harmonious scene.",
-      "video_path": "Girl/mixkit-girl-looking-at-the-likes-in-her-post-4914_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A young mother and her baby sit comfortably on a bed, surrounded by an inviting, cozy atmosphere. The woman, wearing a sleeveless top and jeans, is gently engaging with the baby, who is dressed in an adorable animal-print onesie. The child is seated on the bed with colorful toys scattered around, including a plush toy and a board book. The warm glow from a hanging lamp casts a soft light on them, enhancing the serene environment. Pillows are propped up against the headboard, providing a cushioned backdrop as the mother leans slightly over to interact with the baby. A small bottle is visible beside her, suggesting a nurturing setting. Her hand gestures animatedly as she holds up a soft, white cushion with red and blue accents, likely stimulating the baby\u2019s curiosity. Their shared moment is filled with affection and joy, a perfect snapshot of familial bonding.",
-      "video_path": "Baby/mixkit-loving-mother-and-her-baby-playing-with-soft-toys-49966_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A young girl with long brown hair sits at a round wooden table, engrossed in working on her laptop. The laptop screen is a vivid green, suggesting a green screen effect is in use. To her left, a doll dressed in a yellow and white outfit is casually laid on top of some books, adding a playful and innocent touch to the scene. The setting is cozy, with sheer curtains in the background allowing soft natural light to spill into the room. The girl's posture and focused attention on the laptop suggest she is either playing a game or learning something new. This serene and domestic atmosphere is complemented by the slight blur of a dark couch in the foreground, framing the focused activity of the child.",
-      "video_path": "Girl/mixkit-little-girl-doing-homework-on-a-laptop-4757_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "An expansive view of a calm bay reveals a fleet of sailboats, each anchored in a regimented line stretching toward the horizon. The water is a serene blue, reflecting the soft hues of the early morning sky. A gentle breeze is indicated by the subtle ripples trailing behind the boats, while a single, larger vessel cuts a distinct path, leaving a graceful wake in its journey to the open sea. On one side, a cluster of modern high-rise buildings stands, contrasting against the natural simplicity of the water, suggesting a blend of urban and marine life. The distant shoreline is barely visible, softened by the atmospheric perspective, giving a sense of endless waters meeting the sky. The overall mood is peaceful and orderly, with the boats appearing almost as sentinels guarding the expanse of the tranquil bay.",
-      "video_path": "beach/mixkit-flying-backwards-over-the-sea-near-a-coast-50187_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the video, a person is standing in the center of a dark, featureless space, illuminated by a spotlight that emphasizes their presence. The individual is dressed in a traditional martial arts uniform, known as a gi, which is predominantly white with a black belt tied around the waist, indicating a high level of expertise. The background remains pitch black, creating a stark contrast with the brightly lit figure, ensuring complete focus on them. The person's expression is serious and focused, reflecting a deep sense of discipline and concentration. Their hands move gracefully, transitioning through various martial arts stances, demonstrating practiced skill and fluidity. The uniform's crisp fabric folds and subtly reflects the light, further highlighting each precise movement. Despite the simplicity of the environment, the scene is dynamic, with each motion capturing the essence of martial arts practice. The video effectively conveys a sense of calm strength and mastery, making it ideal for an AI to recreate with attention to posture, lighting, and attire.",
-      "video_path": "Sport/mixkit-karate-fighter-bowing-to-the-front-49706_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In a dimly lit room bathed in a mix of neon purple and blue lights, a focused individual is seated in a gaming chair. She wears a white hoodie and large headphones with cat ears that glow softly, creating a striking silhouette. Her hands rest on a keyboard, typing swiftly as she concentrates intently on the screen in front of her. The atmosphere exudes a sense of intensity and immersion, with the soft-colored lighting enhancing the futuristic vibe. Her long hair cascades down her shoulders, adding a touch of elegance to the otherwise tech-centric setting. The overall scene captures the essence of a dedicated gamer deeply engaged in her virtual world.",
-      "video_path": "earth/mixkit-a-young-woman-wearing-headphones-with-rgb-lights-suddenly-gets-51621_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "Inside a dimly-lit bus, five individuals are seated along the rows of worn seats, each subtly illuminated by the colorful lights emanating from overhead. On the left, a woman sits with a relaxed posture, her curly hair accented by a patterned scarf, wearing a plaid outfit paired with bright neon socks. Next to her, a person clad in a denim jacket appears deep in thought, resting their head on a hand. Further back, another figure in a bucket hat and oversized yellow attire gazes across the aisle, evoking a sense of introspection. The atmosphere is enriched by the soft glow of red and green lights, bathing the bus interior in an almost surreal ambiance, creating a compelling tableau of urban life.",
-      "video_path": "Music/mixkit-conceptual-urban-fashion-42581_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "An aerial view captures two tennis players on a court, with one dressed in white on the left and another in red on the right. They are mid-game, each poised for action with rackets in hand, accentuated by their strategic positioning at opposite baselines. The court itself is a stark, deep blue, bordered by the vibrant green of the surrounding area, with a dark central net dividing the space. Long shadows stretch dramatically across the ground, suggesting a late afternoon setting. The subtly textured surface of the court contrasts with the crisp, white lines marking its boundaries and sections. This scene creates a vivid, balanced composition, highlighting both the competitive tension and serene atmosphere of the game.",
-      "video_path": "People/mixkit-two-people-playing-tennis-aerial-view-880_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In a vibrant, dreamlike setting, a lone figure moves energetically against a backdrop of deep blue and purple hues, casting emotive shadows that ripple with dynamic motion. The figure, almost obscured by a smeared effect, suggests a rhythmic dance or a passionate performance, arms blurred as they sweep through colorful, streaked lighting. A neon glow accentuates their form, particularly highlighting the face which is abstractly illuminated in bursts of orange and red, suggesting intense emotional expression. The scene is dominated by two primary elements \u2013 the figure\u2019s motion and the dramatic lighting, creating a synergy of human emotion and visual spectacle. Swirling trails of light seem to intertwine with the figure, like a visual symphony of movement and color that floods the space. The lighting changes, casting intricate patterns on the figure and the surrounding space, giving the impression of a kaleidoscope in motion. Despite the blurred and abstract portrayal, there is a sense of focus conveyed through the figure\u2019s intent movements, akin to a conductor orchestrating a visual and auditory performance. The environment resonates with an electric energy, suggesting a seamless fusion of art and technology. As the visual drama unfolds, the scene invites viewers to lose themselves in the abstract dance and the play of vivid luminance.",
-      "video_path": "Music/mixkit-dancer-dancing-with-a-light-bar-in-his-hands-42221_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In a brightly lit studio, a photographer wearing a denim jacket focuses intently, capturing shots with a professional camera. Facing him, a model stands gracefully, adjusting her long, flowing hair with delicate movements. The scene is characterized by strong contrasts; the model's soft pink attire and gentle gestures complement the rugged, precise demeanor of the photographer. Positioned against a minimalist backdrop, the pair work seamlessly, with the camera\u2019s lens pointed directly at the model, capturing her elegance. The soft, diffused lighting casts a gentle glow on both subjects, creating an airy and ethereal atmosphere perfect for a high-fashion photo shoot.",
-      "video_path": "Fashion/mixkit-professional-photo-session-with-a-young-female-model-41621_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "The video showcases a serene, expansive landscape covered with a variety of trees dotting the hills. The hills gently slope across the frame, with patches of dry grass contrasting against the lush green foliage. Tall trees with dense canopies stand elegantly, casting soft shadows on the ground below. The sunlight bathes the entire scene, highlighting the varied textures of the leaves and terrain. Gaps between the trees reveal a narrow dirt path meandering through the hills, suggesting a sense of quiet solitude. The undulating hills extend into the distance, creating depth and a calming sense of vast space. The verdant hues of the leaves contrast with the earthy tones of the hills, enhancing the visual richness. In the background, a faint outline of distant hills can be seen, blurred softly by the atmospheric perspective. This tranquil setting could be efficiently recreated in a virtual environment by focusing on its layered composition, color palette, and natural textures.",
-      "video_path": "forest/mixkit-aerial-panorama-of-a-sunny-mountain-landscape-40846_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A bustling ski slope comes alive with skiers descending a pristine, snow-covered hill, surrounded by towering, snow-draped evergreens. Several figures stand atop the slope, silhouetted against a clear blue sky, preparing to embark on their ski run. The chair lift on the right continuously drops off eager adventurers, adding to the excitement at the hilltop. Each skier, clad in colorful winter gear, carves distinct paths into the textured snow as they weave their way down. The interplay of sunlight and shadows accentuates the myriad tracks etched into the slope, creating a dynamic visual rhythm. The scene captures a vibrant winter wonderland, full of action and the thrill of a perfect ski day.",
-      "video_path": "Car/mixkit-skiers-on-a-snowy-slope-3327_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "The scene unfolds within a dimly lit bus, where three young individuals are seated, each absorbed in their unique world. To the left, a person with tied-back hair rests their head on their hand, dressed casually in a jacket and jeans, projecting a relaxed demeanor. Central to the frame is another individual, sitting upright with intense focus, donning a plaid blazer and oversize hoops, enhancing their confident presence. The muted green and red lighting casts an atmospheric glow, adding depth and intrigue to the setting. On the right, a person in a bucket hat and striped shirt leans back, appearing contemplative as they adjust their hat with a nonchalant gesture. The interplay of light and shadow highlights their expressions, creating an intimate and cinematic ambiance. Together, these figures form a cohesive tableau, capturing a moment of introspection amid a bustling yet serene urban environment.",
-      "video_path": "City/mixkit-three-models-posing-to-the-lens-while-on-board-a-42575_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A determined climber is scaling a massive rock face, showcasing exceptional strength and skill. The person, clad in a teal shirt and dark pants, climbs with precision, their movements measured and deliberate. They are secured by climbing gear, which includes ropes and a harness, emphasizing their commitment to safety. The rugged texture of the sandy-colored rock provides an imposing backdrop, adding drama and scale to the climb. In the distance, other large rock formations and sparse vegetation can be seen under a bright, overcast sky, contributing to the natural and adventurous atmosphere. The scene captures a moment of focus and challenge, highlighting the climber's tenacity and the breathtaking environment.",
-      "video_path": "Sport/mixkit-alpinist-climbing-a-huge-rock-in-a-desert-43306_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A woman stands confidently in front of a large array of solar panels, her navy blue jumpsuit contrasting against the lush green grass beneath her feet. Her expression is calm and focused, eyes facing directly ahead, suggesting a deep connection to the subject matter\u2014renewable energy. The sunlight bathes the scene in warm hues, casting gentle shadows and highlighting the geometric precision of the solar panels' grid-like structure. The background reveals a blend of nature and technology, as the panels are anchored on a grassy slope with foliage on the left side of the frame. This composition captures a harmonious blend of human innovation and environmental consciousness, accentuated by the serene outdoor setting.",
-      "video_path": "Business/mixkit-woman-standing-in-front-of-a-solar-panel-4880_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the video, two people are working at a wooden desk, using an iMac computer. One person, wearing a white knit sweater, is using the apple wireless mouse with their right hand, while their left hand rests on the sleek white keyboard. Their movements are smooth yet intentional, suggesting they are focused on a task on the computer screen. The monitor displays a well-organized array of files and folders, hinting at a task that involves detailed organization or detailed data navigation. The second person, only subtly visible, sits closely by and appears to observe or assist, creating a collaborative atmosphere. Their presence adds a quiet dynamic to the scene, as if they are ready to provide input or guidance. Sticky notes with handwritten notes are attached to the monitor\u2019s stand, adding a touch of personal organization amidst the digital workspace. The focus on the keyboard and mouse emphasizes a streamlined workflow, indicative of a productive work environment. The overall ambiance is calm and focuses on teamwork, technology, and efficient workspace management.",
-      "video_path": "People/mixkit-person-with-glasses-working-on-a-desktop-computer-3248_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A man stands in front of a modern glass facade, taking off a dark hoodie to reveal his gray tank top underneath. His arms are lifted high as he maneuvers the hoodie over his head, showcasing a fluid motion that conveys a sense of calm and routine. The lighting highlights the contours of his muscles, emphasizing a combination of strength and quiet determination. Behind him, the reflective surface of the glass panels provides a subtle backdrop, enhancing the focus on his focused and serene demeanor.",
-      "video_path": "Sport/mixkit-man-puts-on-sleeveless-hoodie-603_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "The video displays a captivating dance of fiery orange flames against a stark black background, creating an intense visual contrast. The flames twist and intertwine, forming symmetrical, swirling patterns that expand and contract rhythmically across the frame. Each fiery tendril seems to be alive, moving with an almost hypnotic fluidity that captures the viewer's attention. The illumination from the flames casts subtle shadows, enhancing the depth and texture of the scene. Overall, the dynamic movement and vibrant color palette create an atmosphere of both beauty and power.",
-      "video_path": "fire/mixkit-two-orange-flames-on-black-background-685_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In this scene, a person is seated in a dimly lit room, possibly a recording studio, holding several drumsticks in their hands. The individual's face is partially obscured by sunglasses, adding a touch of mystery to their demeanor. They are wearing a colorful, patterned shirt with a mix of orange and blue tones that stands out against the darker background. The person appears focused and engaged with the drumsticks, their hands prominently displayed. The ambient light casts warm, soft shadows, emphasizing the texture and colors of their shirt and the wooden drumsticks. The room features wooden paneling, which complements the overall cozy, music-centric setting of the scene. The use of perspective centers on the drumsticks, highlighting the importance of rhythm and music in the captured moment.",
-      "video_path": "Music/mixkit-drummer-stretching-before-playing-42783_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A man is casually sitting on a sofa, engrossed in his meal and entertainment. He is holding a TV remote in one hand while reaching for food with the other, indicating a laid-back, comfortable evening. The table before him is filled with takeout containers, revealing a variety of appetizers and dishes, suggestive of a casual dining experience at home. The background is defined by colorful patterned cushions, adding a cozy, homey feel to the scene. Warm, ambient lighting highlights the relaxed atmosphere, casting soft shadows that contribute to the intimate setting. In this moment, he takes a bite of a sandwich, comfortably balancing his attention between food and whatever is playing on the screen.",
-      "video_path": "Man/mixkit-man-watching-tv-and-eating-fast-food-26089_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "The scene opens to a breathtaking view of a tranquil ocean horizon at dusk, displaying a vibrant tapestry of oranges, pinks, and purples as the sun sets. In the foreground, tall, swaying palm trees frame the scene, their silhouettes stark against the colorful sky. The ocean itself shimmers with reflections of the sunset, creating a peaceful, almost ethereal atmosphere. A small boat can be seen in the distance, centered on the horizon, adding a sense of scale and solitude to the scene. The waves gently lap the shore, creating faint patterns on the sandy beach, which stretches across the foreground. Above, the sky is dotted with scattered clouds that catch the last light of the day, enhancing the drama and beauty of the scene. The overall mood is serene and contemplative, capturing a perfect moment of nature\u2019s grandeur.",
-      "video_path": "beach/mixkit-sunset-with-sailing-boats-2166_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A man sits hunched on a couch, the weight of emotions clearly visible on his posture. He wears a simple, gray t-shirt, and his head is bowed, resting in his hands, which cover most of his face, obscuring his features. The gentle light filtering through sheer curtains in the background casts a soft glow upon him, emphasizing the contrast between his static form and the hazy brightness behind. His elbows rest upon his knees, suggesting a posture of deep contemplation or distress. The simplicity of the room, with its muted colors, highlights the focus on the man's internal struggle. Delicate detailing on the fabric of his shirt adds texture, enhancing the scene's realism. Subtle changes in the natural light indicate the passage of time, as the man remains unmoving, absorbed in thought. This intimate moment captures a profound vulnerability, making the scene universally relatable and poignant.",
-      "video_path": "Man/mixkit-worried-and-sad-man-with-his-head-down-4701_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A pair of hands, belonging to an unseen figure, carefully unrolls a large sheet of crisp, white paper on a dark wooden table. The lighting is warm, casting a gentle glow that highlights the textures of the paper and the wood grain of the table. As the paper unfurls, the edges reveal the faint beginnings of a colorful map printed on its surface. The arms, clad in a casual gray T-shirt, suggest a relaxed and focused task at hand. Each motion is deliberate, with fingers deftly guiding the paper, ensuring it lays flat without creases. In the background, a hint of a red curtain can be seen, adding a touch of color and depth to the setting. The composition of the scene emphasizes the contrast between the bright paper and the rich tones of the surroundings. This serene and methodical action evokes a sense of exploration and preparation.",
-      "video_path": "Man/mixkit-unrolling-a-world-map-on-a-table-21626_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A young woman sits on a vibrant green seat inside a bus, illuminated by the soft glow of pink and blue lights. Her outfit is a striking mix of colors: a neon pink top paired with a jacket featuring dark sleeves, and jeans that provide a neutral contrast. She wears large, hoop earrings that catch the light as she moves slightly, exuding an air of cool confidence. Her gaze is directed thoughtfully to the side, suggesting contemplation or daydreaming during her commute. The metallic pole beside her adds a geometric element to the composition, reflecting the kaleidoscope of neon hues. The background is a clean, futuristic white, serving as a blank canvas that amplifies the neon atmosphere. Her relaxed posture and the modern bus setting create a scene that captures a blend of urban life and personal introspection.",
-      "video_path": "City/mixkit-fashion-model-posing-on-a-bus-42578_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A silver SUV drives along a winding, snow-covered mountain road, with dense pine trees blanketed in snow lining both sides. The scene is serene, with the vehicle moving smoothly, possibly on a winter journey or vacation. As the SUV disappears around the bend, another, darker SUV follows, creating a sense of motion and perspective on the snow-dusted asphalt. The towering, snow-laden rock formation to the right contrasts with the dark green of the pines, highlighting the peacefulness of the wintry landscape.",
-      "video_path": "Car/mixkit-curve-on-a-snowy-forest-road-3317_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "The video showcases a vibrant urban skyline during twilight, with towering buildings reflecting the warm hues of the setting sun. A series of tall, cylindrical structures dominate the foreground, adjacent to a complex of industrial equipment and grids. The scene includes modern high-rise buildings with glass exteriors, capturing the evolving architecture of a bustling cityscape. A prominent structure labeled \"CITY OF AUSTIN POWER PLANT\" stands out, highlighting the industrial theme amidst the urban backdrop. The soft glow of city lights begins to pierce the approaching dusk, creating an inviting yet dynamic atmosphere. Shadows cast by the buildings add depth and contrast, emphasizing their massive scale and intricate designs. The overall composition is balanced between the natural light of the sunset and the artificial illumination of the city, offering a compelling visual narrative.",
-      "video_path": "Car/mixkit-slow-air-travel-in-reverse-over-a-big-city-49841_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the scene, a striking architectural structure dominates the view, bathed in a soft, ambient light. The enormous yellow arches serve as the centerpiece, drawing the eye upwards with their majestic curves and towering presence. The smooth, clean surfaces of the structure reflect the light, highlighting the texture and depth of the architecture. In the foreground, blurred streaks of headlights and taillights suggest the motion of vehicles passing by, adding dynamic energy to the otherwise still scene. The contrast between the fast-moving lights and the static arches creates a balanced composition. To the left, a lone streetlamp and a small tree provide a touch of nature and urban elements against the monumental backdrop. The night sky subtly peeks through the gaps in the structure, hinting at a clear, calm evening. Shadows from the arches create patterns on the ground, adding an intricate detail to the scene. Overall, the combination of light, shadow, and movement makes for a dramatic and visually captivating moment.",
-      "video_path": "Car/mixkit-a-fast-timelapse-of-the-street-with-a-monumental-yellow-50993_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A tranquil marina comes into full view under the golden hues of a setting sun. A collection of gleaming yachts and boats are neatly moored, their reflections shimmering softly on the gentle water. The sun's low position casts elongated shadows over the bustling harbor scene, while rolling hillsides surround the distant cityscape. The skyline is interspersed with modern buildings and clusters of residences, adding layers to the vibrant community. At the center, a broad wooden pier juts confidently into the harbor, extending an invitation for leisurely strolls. To the left, various shops and colorful structures line the waterfront, indicating a vibrant coastal economy. The entire atmosphere exudes a serene yet lively charm, balancing the hustle of maritime activity with the peacefulness of the encroaching dusk. It's a scene of calm anticipation, as if the whole place holds its breath before the night's events unfold.",
-      "video_path": "beach/mixkit-harbor-on-a-tourist-coast-with-many-boats-and-yachts-40077_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "The video features a confident individual standing atop a structure against a clear blue sky, exuding a sense of freedom and style. The person is clad in a striking yellow button-up shirt tied at the waist, and beneath it, they wear a simple white top that adds to their relaxed yet stylish appearance. Completing the ensemble are high-waisted white jeans paired with a black belt, adding a touch of contrast. Around their neck is a bold red scarf, providing a splash of color and an air of vintage flair. The person's sunglasses, tinted in yellow, reflect the sunlight and contribute to the overall cool and composed demeanor. Their hair is styled elegantly, pulled back with headphones resting over the ears, suggesting they are immersed in music. One hand casually grazes the headphones, while the other rests gently on the railing, grounding the individual in the moment. The scene is an effortless blend of fashion and tranquility, capturing the spirit of sunny, carefree days.",
-      "video_path": "Music/mixkit-standing-woman-listening-to-music-460_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A ballerina gracefully spins and moves across a pink-hued studio, her poised figure accentuated by a shimmering white tutu and bodice. The background, a continuous wash of soft pink, provides a serene and ethereal atmosphere, emphasizing her fluid movements. Her arms extend with elegance, highlighting the delicacy and precision of her ballet pose, while her focused expression adds intensity to the scene. The subtle details of her costume, combined with the pink monochromatic ambiance, create a dreamlike spectacle, ideal for an AI to envision a oneiric dance setting.",
-      "video_path": "Dance/mixkit-portrait-of-a-ballerina-spinning-with-pink-background-40163_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "The scene unfolds with two human figures in the distance, making their way through a serene meadow, thick with tall golden grass swaying gently in the breeze. The sun hangs low in the sky, casting a soft, diffused glow that illuminates the landscape with a warm, ethereal light. These figures, clad in hiking gear, move deliberately, suggesting they're either embarking on or concluding a journey. Their silhouettes contrast against the lush greenery of the surrounding trees, whose branches reach out, framing the horizon. The play of light and shadow among the trees creates a quilt of textures, with each leaf catching a hint of the sun's dying rays. This tranquil setting evokes a sense of calm and adventure, capturing the quintessential beauty of nature\u2019s landscape.",
-      "video_path": "People/mixkit-landscape-in-nature-while-two-people-are-jogging-44348_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A large cargo ship is docked at an industrial port, its white superstructure contrasting with the deep green and yellow of its deck. The foreground is dominated by the calm, deep blue waters of the harbor, which reflect the vessel\u2019s imposing presence. Surrounding the ship, a series of industrial buildings and storage facilities are visible, hinting at the bustling activity of the port. The deck is intricately detailed, featuring an array of pipes, equipment, and railings, showcasing the ship's functionality and purpose. In the background, a paved area with green patches and a few parked vehicles adds to the busy, industrious atmosphere of the scene.",
-      "video_path": "sea/mixkit-empty-cargo-ship-waiting-at-the-port-4209_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A lone climber ascends a towering rock face, clad in a pink shirt and gray pants, displaying a determined and focused expression. The climber navigates the rugged surface, where the texture of the rock is peppered with natural pockets and crevices that offer handholds and footholds. Sunlight casts soft shadows across the cliff, highlighting the intricate patterns and the climber\u2019s strategic movements. The cliff looms high, with sparse vegetation breaking the monotony of the stone, while distant rocky formations form a dramatic backdrop against the clear blue sky. The climber\u2019s gear, including a harness and chalk bag, underscores the adventure and challenge woven into this majestic, vertical journey.",
-      "video_path": "Sport/mixkit-mountaineer-girl-climbing-a-steep-rocky-mountain-41089_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A person is seen in a close-up shot, skillfully adjusting the tuning pegs of a guitar, showcasing a focused and practiced hand. The image is in black and white, highlighting the contrast between the textures of the instrument and the clothing. The individual's shirt, visible in the background, adds a soft, subtle texture, while the dark tones of the guitar neck create depth in the scene. This composition captures a moment of concentration and finesse, perfect for recreating an intimate musical setting.",
-      "video_path": "Music/mixkit-guitarist-playing-so-inspired-black-and-white-shot-44178_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A musician is playing a large brass instrument with the words \"Brass Band\" clearly visible on its bell. The scene is set against a vibrant yellow backdrop, casting a warm glow on the subject. The musician wears a dark cap and a matching suit, adding a formal touch to his attire. He is deeply focused on his performance, with the instrument's intricate tubing adding complexity to the visual composition. The lighting creates dramatic shadows and highlights, emphasizing the musician's expression and the instrument's metallic sheen. This harmonious blend of color and form captures the essence of a live brass band performance.",
-      "video_path": "Music/mixkit-musician-playing-the-trombone-while-dancing-43752_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the video, a lone musician stands gracefully in front of a grand cathedral, playing an accordion while surrounded by the lively water display of a central fountain. Dressed in a casual ensemble, he wears a light-colored shirt, dark pants, and a flat cap that gives him a vintage charm. His posture is relaxed, yet engaged, as he sways gently in rhythm with the music, casting soft shadows on the cobblestone steps beneath him. The backdrop features the cathedral's towering twin spires, with intricate stonework that casts a rich, historical aura around the scene. Sunlight bathes the entire setting, enhancing the golden hues of the cathedral facade and creating a halo-like effect around the musician. The fountain's water jets splash playfully, catching glimmers of light and adding a dynamic element to the tranquil atmosphere. The scene captures a harmonious blend of architectural majesty and human creativity, framed by the clear, azure sky that extends infinitely above. It's a vivid depiction of solitude and artistry, set against a timeless urban landscape.",
-      "video_path": "Music/mixkit-man-plays-an-accordion-in-front-of-a-fountain-630_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the tranquil video, a person sits in a meditative pose on a gentle hillside, silhouetted against the dawning sky. The person is facing the breathtaking sunrise, with their back slightly turned to the viewer, wearing a simple, light-colored shirt. Their right hand rests on their knee, fingers relaxed in a common meditation mudra, symbolizing calmness and peace. The sky, a stunning blend of soft oranges and deep purples, gradually brightens, casting a warm glow over the lush, green landscape. To the left, the outlines of distant urban buildings can be seen against the horizon, adding a contrast between nature and city life. A river reflecting the sky's colors meanders through the scene, lending a serene, flowing dynamic to the landscape. Trees rise and fall gently across the terrain, their leaves rustling only faintly in the morning breeze. The person remains still and focused, embodying a moment of mindfulness and connection with nature. This visual captures a harmonious balance, evoking a sense of tranquility and introspection.",
-      "video_path": "City/mixkit-girl-meditating-in-yoga-pose-at-sunset-4803_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A serene landscape video captures a breathtaking panoramic view of a vast valley covered in a gentle mist. The undulating hills are lush with dense greenery, their rich foliage creating a vibrant border on the left side of the frame. The mist weaves through the landscape like a soft, ethereal blanket, lending a dream-like quality to the scene. In the distance, several mountain peaks emerge, their dark outlines contrasting against the pale blue sky. A few faint, wispy clouds drift lazily across the horizon, complementing the tranquil atmosphere. The sunlight filters through the haze, casting a warm glow and highlighting different textures of the flora. The overall mood is calm and contemplative, inviting the viewer to pause and appreciate nature's untouched beauty. The composition emphasizes depth and expansiveness, drawing attention to the harmony between earth and sky. This captivating scene embodies tranquility, offering a perfect backdrop for meditation or relaxation.",
-      "video_path": "forest/mixkit-flying-over-a-hill-with-a-view-of-the-surrounding-49743_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In this scene, a bearded individual is intently focused on their smartphone, with the sun setting in the background, casting a warm glow across the cityscape. The person, partially visible, is wearing a dark, buttoned shirt that contrasts with the golden hue of the sunset. Their hands are holding the smartphone delicately but purposefully, reflecting a sense of engagement and focus on the screen. The sunlight creates a striking lens flare effect, enhancing the dramatic atmosphere of the moment as it glimmers off the phone\u2019s surface. The surrounding environment hints at an elevated vantage point, providing a panoramic view of the urban landscape below.",
-      "video_path": "City/mixkit-guy-texting-at-sunset-265_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In an expansive, industrial space defined by towering columns and high ceilings, a solitary figure takes center stage. The person, dressed in dark, fitted clothing, assumes a powerful, dynamic stance with one leg bent forward and both arms outstretched in a horizontal arc. Framing this pose are intense flames that engulf their arms, creating a striking visual contrast against the muted tones of the room. The fire forms a brilliant halo of orange and yellow, casting flickering shadows on the weathered walls and worn, tiled floor. This interplay between light and dark showcases the dancer's poise and agility, as they maintain balance amidst the intense heat. Windows line the background, their panes dimly illuminated by the daylight filtering in, adding depth and perspective to the scene. The entire performance evokes a sense of raw energy and elemental mastery, as the figure continues to manipulate the fire in a seamless, mesmerizing display.",
-      "video_path": "fire/mixkit-expert-juggler-doing-tricks-with-a-stick-with-fire-43663_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A man is playing the violin, focused intently on his music. His fingers gracefully dance along the strings, flawlessly executing each note. He holds the violin close to his chin with a sense of familiarity and expertise. The rich, warm tones of the violin reflect in the soft lighting of the room. He wears a dark shirt, and a subtle necklace rests against his chest, adding a personal touch to his attire. The bow moves smoothly across the strings, producing a melody that seems to fill the space with emotion. His expression is one of concentration and passion, immersing himself fully in the performance. The background is softly blurred, bringing the violin's intricate craftsmanship and his precise movements into sharp focus. This serene and intimate moment captures the essence of his musical artistry.",
-      "video_path": "Music/mixkit-fiddler-playing-a-song-639_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the dimly lit parking garage, two figures engage in an impromptu game of soccer. The first person, wearing a light grey shirt and black pants with three white stripes, skillfully maneuvers the ball with precise footwork. The ground is slick with patches of water, reflecting the vibrant neon lights above. A second figure, clad in dark clothing, stands poised in the background, ready to intercept. The space is defined by stark yellow lines and orange safety bollards, adding structure to the chaotic energy of the scene. The soccer ball glides smoothly across the wet floor, kicking up droplets as it passes. Despite the muted colors of the environment, the players' movements are dynamic and full of life. Their shadowy silhouettes dance with the reflecting light, creating a mesmerizing visual interplay. The atmosphere is charged with focus and camaraderie, encapsulating the essence of a late-night urban soccer experience.",
-      "video_path": "Sport/mixkit-player-making-skillful-play-in-a-street-soccer-game-43504_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A lone climber is seen scaling a towering vertical rock face, demonstrating remarkable strength and focus. Dressed in a light-colored shirt and jeans, the climber grips the stone tightly, navigating the rough textures and crevices with precision. The sheer cliff is massive, exhibiting a range of natural hues from light tan to deep gray, accentuating the climber's figure against the vast rocky backdrop. Surrounding the cliff, scattered greenery and rugged terrain provide a sense of wilderness and isolation. The scene portrays a daring ascension requiring concentration and skill, capturing the essence of human endeavor against nature's formidable beauty.",
-      "video_path": "Sport/mixkit-skilled-mountaineer-climbing-a-gigantic-mountain-41083_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In this serene landscape, a lush meadow stretches across the foreground, dotted with vibrant yellow wildflowers swaying gently in the breeze. A towering tree stands majestically on the right side, its branches reaching wide under the bright blue sky filled with fluffy white clouds. On the left, dense trees form a natural corridor leading to the horizon, suggesting a sense of journey and possibility. The richness of the green grass contrasts beautifully with the golden hue of the distant fields, creating a harmonious palette of nature\u2019s colors. The play of light and shadow adds depth and dimension, evoking a tranquil, inviting atmosphere. It's a scene where nature\u2019s beauty simply commands attention, offering a perfect escape into tranquility.",
-      "video_path": "sky/mixkit-countryside-meadow-4075_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A solitary boat glides across the expansive, tranquil expanse of a serene lake. The vessel leaves a gentle wake behind, creating delicate ripples across the mirror-like surface. The water appears a rich shade of teal, seamlessly blending with the sky at the horizon. Silhouettes of distant trees are faintly visible, creating a picturesque backdrop that enhances the solitary journey of the boat. The sky is a calm gradient, shifting from soft oranges near the shore to the pale blues above. In the distance, a few slender poles emerge from the water, remnants of an old structure or natural formation. The mood of the scene is one of peace and solitude, with the boat journeying steadily through the quiet landscape. There is a sense of endless possibilities as the boat moves toward the unseen beyond the frame. The simplicity and stillness of the scene invite contemplation and reflection, encapsulating a perfect moment of quietude on the water.",
-      "video_path": "mountain/mixkit-motorboat-on-a-large-lake-with-turquoise-blue-waters-4996_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In a cozy, dimly lit caf\u00e9, a woman sits alone at a rustic wooden table, fully engrossed in her reading. Her dark, wavy hair frames her face as she leans forward over an open book, suggesting deep focus and contemplation. The caf\u00e9\u2019s ambiance is warm, with hanging pendant lights casting a soft glow over the wooden shelves lined with jars and coffee paraphernalia in the background. A small cup of coffee rests just within her reach, alongside a glass dome encasing a solitary pastry, adding a touch of tranquility to the scene. Her casual attire, a denim jacket over a simple shirt, complements the laid-back, comfortable setting of the caf\u00e9. The contrast between her concentrated expression and the bustling, yet subdued caf\u00e9 atmosphere creates a harmonious, serene visual. The overall composition captures a quiet moment of introspection amidst the gentle hum of caf\u00e9 life.",
-      "video_path": "Woman/mixkit-woman-drinking-coffee-in-a-cafe-223_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In a vast, deserted landscape under the night sky, a solitary figure stands at a small music setup, illuminated by strategically placed lights. The person is engrossed in playing a keyboard, with various electronic equipment surrounding them, casting soft glows of orange and blue hues across the scene. To the left, a large circular light adds a dramatic focal point, highlighting the intense contrast between the darkness and the lit performance area. This setup, with its minimalistic design and strategic lighting, creates a captivating and easily recognizable scene that merges the serene, expansive backdrop with an intimate, focused music performance.",
-      "video_path": "Music/mixkit-talented-dj-playing-in-a-lonely-desert-42414_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In a bustling urban scene, cars zoom past a weathered building, their blurred motion a testament to the city\u2019s lively pace. The building, with its faded yellow and brown facade, boasts graffiti that speaks of both art and decay, framing the scene with an air of urban grit. A solitary figure stands slightly to the side, clad casually in a gray top and mustard trousers, gazing into the street, seemingly detached from the surrounding flurry. The motion of the traffic creates a dynamic contrast against the static backdrop, emphasizing the relentless movement of the city. As the video progresses, a bright yellow taxi appears, slowing down as it approaches the figure, adding a pop of color to the desaturated hues of the environment. The interaction suggests a routine, a possibly daily exchange between the driver and the pedestrian, hinting at the rhythms of city life. Overhead, a soft, overcast sky casts a diffused light, lending the scene a subdued, timeless quality. Small elements, like the vertical pole cutting through the frame and the distant chatter of urban sounds, complete this vivid tableau of urban existence.",
-      "video_path": "Car/mixkit-morning-in-the-street-time-lapse-1648_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "A young woman sits on a curb in a tranquil park, basking in the golden hue of the setting sun. Beside her, a collie dog rests calmly, its fur illuminated by the warm sunlight, creating a serene glow. The woman's hand gently strokes the dog's back, highlighting the bond and affection between them. Tall trees surround the pair, casting elongated shadows on the leaf-laden ground, adding to the peaceful and intimate ambiance of the scene.",
-      "video_path": "Pets/mixkit-a-woman-pets-a-dog-in-a-park-1562_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the video, a grand, majestic elephant stands in an open, sunlit field, its massive form dominating the scene. The elephant's skin is a tapestry of earthy tones, with rough, textured wrinkles that add character to its already imposing presence. Its trunk, a powerful and flexible appendage, moves gently, swaying as the elephant possibly enjoys the warmth of the day. The background is a blur of greenery, suggesting a lively environment filled with trees and shrubs that provide a natural habitat. Light plays on the elephant's skin, highlighting patches of dust and dirt that give it an authentic wilderness look. The scene captures the tranquility and majesty of this gentle giant in its natural surroundings.",
-      "video_path": "Zoo/mixkit-wet-elephant-in-the-savanna-3663_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the video, a fluffy dog with brown patches is intently engaged with a bright red toy shaped like a fire hydrant, which has a yellow and orange rope attached. The dog's body is relaxed as it lies on a plain white background, concentrating on nudging and playfully biting the toy. Its ears perk up slightly with curiosity, and its eyes are fixated on the toy, suggesting a scene of focused playfulness. The neutral tones of the dog's fur contrast starkly against the vivid red of the toy, creating a visually striking moment.",
-      "video_path": "Pets/mixkit-a-cute-border-collie-dog-play-with-a-fire-street-50662_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    }
-  ]
-}

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm

@@ -50,10 +50,10 @@ VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/
 training_args=(
   --tracker_project_name wan_i2v_finetune
   --output_dir="checkpoints/wan_i2v_finetune"
-  --max_train_steps 2000
-  --train_batch_size 2
+  --max_train_steps=2000
+  --train_batch_size=2
   --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
+  --gradient_accumulation_steps=1
   --num_latent_t 8
   --num_height 480
   --num_width 832
@@ -93,9 +93,9 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 1000
+  --learning_rate=1e-5
+  --mixed_precision="bf16"
+  --checkpointing_steps=1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

examples/training/finetune/wan_t2v_1.3B/crush_smol/download_dataset.sh

@@ -1,3 +0,0 @@
-#!/bin/bash
-
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm

@@ -46,11 +46,11 @@ 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 "checkpoints/wan_t2v_finetune"
-  --max_train_steps 1000
-  --train_batch_size 4
+  --output_dir="checkpoints/wan_t2v_finetune"
+  --max_train_steps=1000
+  --train_batch_size=4
   --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
+  --gradient_accumulation_steps=1
   --num_latent_t 8
   --num_height 480
   --num_width 832
@@ -90,9 +90,9 @@ validation_args=(
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 5e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 500
+  --learning_rate=5e-5
+  --mixed_precision="bf16"
+  --checkpointing_steps=500
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )

fastvideo/attention/backends/abstract.py

@@ -6,7 +6,8 @@ from dataclasses import dataclass, fields
 from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
 
 if TYPE_CHECKING:
-    pass
+    from fastvideo.fastvideo_args import FastVideoArgs
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 
 import torch
 
@@ -86,8 +87,10 @@ class AttentionMetadataBuilder(ABC, Generic[T]):
     @abstractmethod
     def build(
         self,
-        **kwargs: dict[str, Any],
-    ) -> AttentionMetadata:
+        current_timestep: int,
+        forward_batch: "ForwardBatch",
+        fastvideo_args: "FastVideoArgs",
+    ) -> T:
         """Build attention metadata with on-device tensors."""
         raise NotImplementedError
 

fastvideo/attention/backends/sliding_tile_attn.py

@@ -13,8 +13,10 @@ from fastvideo.attention.backends.abstract import (AttentionBackend,
                                                    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__)
@@ -74,13 +76,13 @@ class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder):
     def prepare(self):
         pass
 
-    def build(  # type: ignore
+    def build(
         self,
-        STA_param: list[list[Any]],
         current_timestep: int,
-        **kwargs: dict[str, Any],
+        forward_batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
     ) -> SlidingTileAttentionMetadata:
-        param = STA_param
+        param = forward_batch.STA_param
         if param is None:
             return SlidingTileAttentionMetadata(
                 current_timestep=current_timestep, STA_param=[])
@@ -128,6 +130,12 @@ class SlidingTileAttentionImpl(AttentionImpl):
         }
 
     def tile(self, x: torch.Tensor) -> torch.Tensor:
+        x = rearrange(x,
+                      "b (sp t h w) head d -> b (t sp h w) head d",
+                      sp=self.sp_size,
+                      t=self.dit_seq_shape_int[0] // self.sp_size,
+                      h=self.dit_seq_shape_int[1],
+                      w=self.dit_seq_shape_int[2])
         return rearrange(
             x,
             "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
@@ -148,7 +156,12 @@ class SlidingTileAttentionImpl(AttentionImpl):
             ts_t=self.STA_base_tile_size[0],
             ts_h=self.STA_base_tile_size[1],
             ts_w=self.STA_base_tile_size[2])
-        return x
+        return rearrange(x,
+                         "b (t sp h w) head d -> b (sp t h w) head d",
+                         sp=self.sp_size,
+                         t=self.dit_seq_shape_int[0] // self.sp_size,
+                         h=self.dit_seq_shape_int[1],
+                         w=self.dit_seq_shape_int[2])
 
     def preprocess_qkv(
         self,

fastvideo/attention/backends/video_sparse_attn.py

@@ -1,114 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
-import functools
 import math
 from dataclasses import dataclass
 
 import torch
+from einops import rearrange
 
 try:
     from vsa import video_sparse_attn
 except ImportError:
     video_sparse_attn = None
 
-from typing import Any
-
 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__)
-VSA_TILE_SIZE = (4, 4, 4)
-
-
-@functools.lru_cache(maxsize=10)
-def get_tile_partition_indices(
-    dit_seq_shape: tuple[int, int, int],
-    tile_size: tuple[int, int, int],
-    device: torch.device,
-) -> torch.LongTensor:
-    T, H, W = dit_seq_shape
-    ts, hs, ws = tile_size
-    indices = torch.arange(T * H * W, device=device,
-                           dtype=torch.long).reshape(T, H, W)
-    ls = []
-    for t in range(math.ceil(T / ts)):
-        for h in range(math.ceil(H / hs)):
-            for w in range(math.ceil(W / ws)):
-                ls.append(indices[t * ts:min(t * ts + ts, T),
-                                  h * hs:min(h * hs + hs, H),
-                                  w * ws:min(w * ws + ws, W)].flatten())
-    index = torch.cat(ls, dim=0)
-    return index
-
-
-@functools.lru_cache(maxsize=10)
-def get_reverse_tile_partition_indices(
-    dit_seq_shape: tuple[int, int, int],
-    tile_size: tuple[int, int, int],
-    device: torch.device,
-) -> torch.LongTensor:
-    return torch.argsort(
-        get_tile_partition_indices(dit_seq_shape, tile_size, device))
-
-
-@functools.lru_cache(maxsize=10)
-def construct_variable_block_sizes(
-    dit_seq_shape: tuple[int, int, int],
-    num_tiles: tuple[int, int, int],
-    device: torch.device,
-) -> torch.LongTensor:
-    """
-    Compute the number of valid (non‑padded) tokens inside every
-    (ts_t × ts_h × ts_w) tile after padding ‑‑ flattened in the order
-    (t‑tile, h‑tile, w‑tile) that `rearrange` uses.
-
-    Returns
-    -------
-    torch.LongTensor  # shape: [∏ full_window_size]
-    """
-    # unpack
-    t, h, w = dit_seq_shape
-    ts_t, ts_h, ts_w = VSA_TILE_SIZE
-    n_t, n_h, n_w = num_tiles
-
-    def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor:
-        """Vector with the size of each tile along one dimension."""
-        sizes = torch.full((n_tiles, ), tile, dtype=torch.int, device=device)
-        # size of last (possibly partial) tile
-        remainder = dim_len - (n_tiles - 1) * tile
-        sizes[-1] = remainder if remainder > 0 else tile
-        return sizes
-
-    t_sizes = _sizes(t, ts_t, n_t)  # [n_t]
-    h_sizes = _sizes(h, ts_h, n_h)  # [n_h]
-    w_sizes = _sizes(w, ts_w, n_w)  # [n_w]
-
-    # broadcast‑multiply to get voxels per tile, then flatten
-    block_sizes = (
-        t_sizes[:, None, None]  # [n_t, 1,   1]
-        * h_sizes[None, :, None]  # [1,   n_h, 1]
-        * w_sizes[None, None, :]  # [1,   1,   n_w]
-    ).reshape(-1)  # [n_t * n_h * n_w]
-
-    return block_sizes
-
-
-@functools.lru_cache(maxsize=10)
-def get_non_pad_index(
-    variable_block_sizes: torch.LongTensor,
-    max_block_size: int,
-):
-    n_win = variable_block_sizes.shape[0]
-    device = variable_block_sizes.device
-    starts_pad = torch.arange(n_win, device=device) * max_block_size
-    index_pad = starts_pad[:, None] + torch.arange(max_block_size,
-                                                   device=device)[None, :]
-    index_mask = torch.arange(
-        max_block_size, device=device)[None, :] < variable_block_sizes[:, None]
-    return index_pad[index_mask]
 
 
 class VideoSparseAttentionBackend(AttentionBackend):
@@ -141,12 +52,6 @@ class VideoSparseAttentionMetadata(AttentionMetadata):
     current_timestep: int
     dit_seq_shape: list[int]
     VSA_sparsity: float
-    num_tiles: list[int]
-    total_seq_length: int
-    tile_partition_indices: torch.LongTensor
-    reverse_tile_partition_indices: torch.LongTensor
-    variable_block_sizes: torch.LongTensor
-    non_pad_index: torch.LongTensor
 
 
 class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
@@ -157,44 +62,30 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
     def prepare(self):
         pass
 
-    def build(  # type: ignore
+    def build(
         self,
         current_timestep: int,
-        raw_latent_shape: tuple[int, int, int],
-        patch_size: tuple[int, int, int],
-        VSA_sparsity: float,
-        device: torch.device,
-        **kwargs: dict[str, Any],
+        forward_batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
     ) -> VideoSparseAttentionMetadata:
-        patch_size = patch_size
-        dit_seq_shape = (raw_latent_shape[0] // patch_size[0],
-                         raw_latent_shape[1] // patch_size[1],
-                         raw_latent_shape[2] // patch_size[2])
+        if forward_batch.latents is None:
+            raise ValueError("latents cannot be None")
 
-        num_tiles = (math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]),
-                     math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]),
-                     math.ceil(dit_seq_shape[2] / VSA_TILE_SIZE[2]))
-        total_seq_length = math.prod(dit_seq_shape)
+        raw_latent_shape = forward_batch.raw_latent_shape
+        if raw_latent_shape is None:
+            raise ValueError("raw_latent_shape cannot be None")
 
-        tile_partition_indices = get_tile_partition_indices(
-            dit_seq_shape, VSA_TILE_SIZE, device)
-        reverse_tile_partition_indices = get_reverse_tile_partition_indices(
-            dit_seq_shape, VSA_TILE_SIZE, device)
-        variable_block_sizes = construct_variable_block_sizes(
-            dit_seq_shape, num_tiles, device)
-        non_pad_index = get_non_pad_index(variable_block_sizes,
-                                          math.prod(VSA_TILE_SIZE))
+        patch_size = fastvideo_args.pipeline_config.dit_config.patch_size
+        dit_seq_shape = [
+            raw_latent_shape[2] // patch_size[0],
+            raw_latent_shape[3] // patch_size[1],
+            raw_latent_shape[4] // patch_size[2]
+        ]
+        VSA_sparsity = forward_batch.VSA_sparsity
 
-        return VideoSparseAttentionMetadata(
-            current_timestep=current_timestep,
-            dit_seq_shape=dit_seq_shape,  # type: ignore
-            VSA_sparsity=VSA_sparsity,  # type: ignore
-            num_tiles=num_tiles,  # type: ignore
-            total_seq_length=total_seq_length,  # type: ignore
-            tile_partition_indices=tile_partition_indices,  # type: ignore
-            reverse_tile_partition_indices=reverse_tile_partition_indices,
-            variable_block_sizes=variable_block_sizes,
-            non_pad_index=non_pad_index)
+        return VideoSparseAttentionMetadata(current_timestep=current_timestep,
+                                            dit_seq_shape=dit_seq_shape,
+                                            VSA_sparsity=VSA_sparsity)
 
 
 class VideoSparseAttentionImpl(AttentionImpl):
@@ -212,44 +103,66 @@ class VideoSparseAttentionImpl(AttentionImpl):
         self.prefix = prefix
         sp_group = get_sp_group()
         self.sp_size = sp_group.world_size
+        self.VSA_base_tile_size = [4, 4, 4]
+        self.dit_seq_shape: list[int]
+        self.full_window_size: list[int]
+        self.img_seq_length: int
 
-    def tile(self, x: torch.Tensor, num_tiles: list[int],
-             tile_partition_indices: torch.LongTensor,
-             non_pad_index: torch.LongTensor) -> torch.Tensor:
-        t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0]
-        h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1]
-        w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2]
+    def tile(self, x: torch.Tensor) -> torch.Tensor:
+        x = rearrange(x,
+                      "b (sp t h w) head d -> b (t sp h w) head d",
+                      sp=self.sp_size,
+                      t=self.dit_seq_shape[0] // self.sp_size,
+                      h=self.dit_seq_shape[1],
+                      w=self.dit_seq_shape[2])
 
-        x_padded = torch.zeros(
-            (x.shape[0], t_padded_size * h_padded_size * w_padded_size,
-             x.shape[-2], x.shape[-1]),
-            device=x.device,
-            dtype=x.dtype)
-        x_padded[:, non_pad_index] = x[:, tile_partition_indices]
-        return x_padded
+        return rearrange(
+            x,
+            "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.VSA_base_tile_size[0],
+            ts_h=self.VSA_base_tile_size[1],
+            ts_w=self.VSA_base_tile_size[2])
 
-    def untile(self, x: torch.Tensor,
-               reverse_tile_partition_indices: torch.LongTensor,
-               non_pad_index: torch.LongTensor) -> torch.Tensor:
-        x = x[:, non_pad_index][:, reverse_tile_partition_indices]
-        return x
+    def untile(self, x: torch.Tensor) -> torch.Tensor:
+        x = rearrange(
+            x,
+            "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.VSA_base_tile_size[0],
+            ts_h=self.VSA_base_tile_size[1],
+            ts_w=self.VSA_base_tile_size[2])
+        return rearrange(x,
+                         "b (t sp h w) head d -> b (sp t h w) head d",
+                         sp=self.sp_size,
+                         t=self.dit_seq_shape[0] // self.sp_size,
+                         h=self.dit_seq_shape[1],
+                         w=self.dit_seq_shape[2])
 
     def preprocess_qkv(
         self,
         qkv: torch.Tensor,
         attn_metadata: VideoSparseAttentionMetadata,
     ) -> torch.Tensor:
-        return self.tile(qkv, attn_metadata.num_tiles,
-                         attn_metadata.tile_partition_indices,
-                         attn_metadata.non_pad_index)
+        self.dit_seq_shape = attn_metadata.dit_seq_shape
+        self.full_window_size = [
+            self.dit_seq_shape[0] // self.VSA_base_tile_size[0],
+            self.dit_seq_shape[1] // self.VSA_base_tile_size[1],
+            self.dit_seq_shape[2] // self.VSA_base_tile_size[2]
+        ]
+        self.img_seq_length = math.prod(self.dit_seq_shape)
+        return self.tile(qkv)
 
     def postprocess_output(
         self,
         output: torch.Tensor,
         attn_metadata: VideoSparseAttentionMetadata,
     ) -> torch.Tensor:
-        return self.untile(output, attn_metadata.reverse_tile_partition_indices,
-                           attn_metadata.non_pad_index)
+        return self.untile(output)
 
     def forward(  # type: ignore[override]
         self,
@@ -268,17 +181,17 @@ class VideoSparseAttentionImpl(AttentionImpl):
 
         cur_topk = math.ceil(
             (1 - VSA_sparsity) *
-            (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE)))
+            (self.img_seq_length / math.prod(self.VSA_base_tile_size)))
 
         if video_sparse_attn is None:
             raise NotImplementedError("video_sparse_attn is not installed")
+
         hidden_states = video_sparse_attn(
             query,
             key,
             value,
-            variable_block_sizes=attn_metadata.variable_block_sizes,
             topk=cur_topk,
-            block_size=VSA_TILE_SIZE,
+            block_size=(4, 4, 4),
             compress_attn_weight=gate_compress).transpose(1, 2)
 
         return hidden_states

fastvideo/configs/models/dits/wanvideo.py

@@ -89,7 +89,6 @@ class WanVideoArchConfig(DiTArchConfig):
     image_dim: int | None = None
     added_kv_proj_dim: int | None = None
     rope_max_seq_len: int = 1024
-    pos_embed_seq_len: int | None = None
     exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
 
     def __post_init__(self):

fastvideo/configs/models/vaes/wanvae.py

@@ -9,7 +9,6 @@ from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
 @dataclass
 class WanVAEArchConfig(VAEArchConfig):
     base_dim: int = 96
-    decoder_base_dim: int | None = None
     z_dim: int = 16
     dim_mult: tuple[int, ...] = (1, 2, 4, 4)
     num_res_blocks: int = 2
@@ -52,21 +51,14 @@ class WanVAEArchConfig(VAEArchConfig):
         2.8251,
         1.9160,
     )
-    is_residual: bool = False
-    in_channels: int = 3
-    out_channels: int = 3
-    patch_size: int | None = None
-    scale_factor_temporal: int = 4
-    scale_factor_spatial: int = 8
-    clip_output: bool = True
+    temporal_compression_ratio = 4
+    spatial_compression_ratio = 8
 
     def __post_init__(self):
         self.scaling_factor: torch.tensor = 1.0 / torch.tensor(
             self.latents_std).view(1, self.z_dim, 1, 1, 1)
         self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view(
             1, self.z_dim, 1, 1, 1)
-        self.temporal_compression_ratio = self.scale_factor_temporal
-        self.spatial_compression_ratio = self.scale_factor_spatial
 
 
 @dataclass

fastvideo/configs/pipelines/base.py

@@ -85,6 +85,9 @@ class PipelineConfig:
     # DMD parameters
     dmd_denoising_steps: list[int] | None = field(default=None)
 
+    # Wan2.2 TI2V parameters
+    ti2v_task: bool = False
+
     # Compilation
     # enable_torch_compile: bool = False
 

fastvideo/configs/pipelines/registry.py

@@ -8,6 +8,7 @@ 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,
+                                             Wan2_2_TI2V_5B_Config,
                                              WanI2V480PConfig, WanI2V720PConfig,
                                              WanT2V480PConfig, WanT2V720PConfig)
 from fastvideo.logger import init_logger
@@ -26,9 +27,12 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "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/FastWan2.1-T2V-14B-480P-Diffusers": FastWanT2V480PConfig,
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
     "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
-    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig
+    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
+    # "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
+    # "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,
     # Add other specific weight variants
 }
 

fastvideo/configs/pipelines/wan.py

@@ -111,3 +111,26 @@ class FastWanT2V480PConfig(WanT2V480PConfig):
     def __post_init__(self) -> None:
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
+
+
+@dataclass
+class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
+    """Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD"""
+
+    # Denoising stage
+    flow_shift: int = 5
+    ti2v_task: bool = True
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
+    pass
+
+
+@dataclass
+class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
+    pass

fastvideo/configs/sample/registry.py

@@ -6,7 +6,7 @@ from typing import Any
 from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                               HunyuanSamplingParam)
 from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
-from fastvideo.configs.sample.wan import (FastWanT2V480PConfig,
+from fastvideo.configs.sample.wan import (Wan2_1_Fun_1_3B_InP_SamplingParam,
                                           Wan2_2_TI2V_5B_SamplingParam,
                                           WanI2V_14B_480P_SamplingParam,
                                           WanI2V_14B_720P_SamplingParam,
@@ -25,9 +25,18 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam,
     "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
     "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
+    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
+    Wan2_1_Fun_1_3B_InP_SamplingParam,
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
-    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers":
+    Wan2_1_Fun_1_3B_InP_SamplingParam,
+    "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers":
+    Wan2_1_Fun_1_3B_InP_SamplingParam,
     "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
+    # "Wan-AI/Wan2.2-T2V-A14B-Diffusers":
+    # Wan2_2_T2V_A14B_SamplingParam,
+    # "Wan-AI/Wan2.2-I2V-A14B-Diffusers":
+    # Wan2_2_I2V_A14B_SamplingParam,
     # Add other specific weight variants
 }
 

fastvideo/configs/sample/wan.py

@@ -107,6 +107,21 @@ class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
     fps: int = 16
 
 
+# =============================================
+# ============= Wan2.1 Fun Models =============
+# =============================================
+@dataclass
+class Wan2_1_Fun_1_3B_InP_SamplingParam(SamplingParam):
+    """Sampling parameters for Wan2.1 Fun 1.3B InP model."""
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+    negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+    guidance_scale: float = 6.0
+    num_inference_steps: int = 50
+
+
 # =============================================
 # ============= Wan2.2 TI2V Models =============
 # =============================================

fastvideo/distributed/device_communicators/base_device_communicator.py

@@ -120,43 +120,38 @@ class DistributedAutograd:
             ) == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
 
             if scatter_dim == 2 and gather_dim == 1:
-                bs, shard_seqlen, hn, hd = input_.shape
+                bs, shard_seqlen, hc, hs = input_.shape
                 seqlen = shard_seqlen * world_size
-                shard_hn = hn // world_size
+                shard_hc = hc // world_size
 
-                input_ = input_.transpose(
-                    0, 2).contiguous()  # hn, shard_seqlen, bs, hd
-                output = torch.empty_like(input_)
+                input_t = input_.reshape(bs, shard_seqlen, world_size, shard_hc,
+                                         hs).transpose(0, 2).contiguous()
+                output = torch.empty_like(input_t)
 
-                dist.all_to_all_single(output, input_,
-                                       group=group)  # hn, shard_seqlen, bs, hd
+                dist.all_to_all_single(output, input_t, group=group)
 
-                output = torch.cat(output.split(shard_hn),
-                                   dim=1)  # sharded hn, seqlen, bs, hd
-
-                output = output.transpose(
-                    0, 2).contiguous()  # bs, seqlen, sharded_hn, hd
+                output = output.reshape(seqlen, bs, shard_hc,
+                                        hs).transpose(0, 1).contiguous()
+                output = output.reshape(bs, seqlen, shard_hc, hs)
 
                 return output
             elif scatter_dim == 1 and gather_dim == 2:
-                bs, seqlen, shard_hn, hd = input_.shape
-                hn = shard_hn * world_size
+                bs, seqlen, shard_hc, hs = input_.shape
+                hc = shard_hc * world_size
                 shard_seqlen = seqlen // world_size
 
-                input_ = input_.transpose(
-                    0, 2).contiguous()  # shard_hn, seqlen, bs, hd
+                input_t = input_.reshape(bs, world_size, shard_seqlen, shard_hc,
+                                         hs)
+                input_t = input_t.transpose(0, 3).transpose(0, 1).contiguous()
+                input_t = input_t.reshape(world_size, shard_hc, shard_seqlen,
+                                          bs, hs)
 
-                input_ = input_.reshape(shard_hn, world_size, shard_seqlen, bs,
-                                        hd).transpose(0, 1).reshape(
-                                            shard_hn * world_size, shard_seqlen,
-                                            bs, hd).contiguous()
+                output = torch.empty_like(input_t)
+                dist.all_to_all_single(output, input_t, group=group)
 
-                output = torch.empty_like(input_)
-
-                dist.all_to_all_single(output, input_, group=group)
-
-                output = output.transpose(
-                    0, 2).contiguous()  # bs, seqlen, sharded_hn, hd
+                output = output.reshape(hc, shard_seqlen, bs, hs)
+                output = output.transpose(0, 2).contiguous()
+                output = output.reshape(bs, shard_seqlen, hc, hs)
 
                 return output
             else:

fastvideo/entrypoints/video_generator.py

@@ -9,6 +9,7 @@ diffusion models.
 import math
 import os
 import time
+from copy import deepcopy
 from typing import Any
 
 import imageio
@@ -202,6 +203,8 @@ class VideoGenerator:
         if sampling_param is None:
             sampling_param = SamplingParam.from_pretrained(
                 fastvideo_args.model_path)
+        else:
+            sampling_param = deepcopy(sampling_param)
 
         kwargs["prompt"] = prompt
         sampling_param.update(kwargs)
@@ -288,6 +291,10 @@ class VideoGenerator:
         """ # type: ignore[attr-defined]
         logger.info(debug_str)
 
+        # Use prompt[:100] for video name
+        if sampling_param.output_video_name is None:
+            sampling_param.output_video_name = prompt[:100]
+
         # Prepare batch
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
@@ -297,10 +304,6 @@ class VideoGenerator:
             extra={},
         )
 
-        # Use prompt[:100] for video name
-        if batch.output_video_name is None:
-            batch.output_video_name = prompt[:100]
-
         # Run inference
         start_time = time.perf_counter()
         output_batch = self.executor.execute_forward(batch, fastvideo_args)

fastvideo/fastvideo_args.py

@@ -615,7 +615,7 @@ class TrainingArgs(FastVideoArgs):
     output_dir: str = ""
     checkpoints_total_limit: int = 0
     checkpointing_steps: int = 0
-    resume_from_checkpoint: str = ""  # specify the checkpoint folder to resume from
+    resume_from_checkpoint: bool = False
 
     # optimizer & scheduler
     num_train_epochs: int = 0
@@ -670,8 +670,6 @@ class TrainingArgs(FastVideoArgs):
     min_timestep_ratio: float = 0.2
     max_timestep_ratio: float = 0.98
     real_score_guidance_scale: float = 3.5
-    training_state_checkpointing_steps: int = 0  # for resuming training
-    weight_only_checkpointing_steps: int = 0  # for inference
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -831,15 +829,6 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--checkpointing-steps",
                             type=int,
                             help="Steps between checkpoints")
-        parser.add_argument(
-            "--training-state-checkpointing-steps",
-            type=int,
-            help=
-            "Steps between training state checkpoints (for resuming training)")
-        parser.add_argument(
-            "--weight-only-checkpointing-steps",
-            type=int,
-            help="Steps between weight-only checkpoints (for inference)")
         parser.add_argument("--resume-from-checkpoint",
                             type=str,
                             help="Path to checkpoint to resume from")

fastvideo/models/schedulers/scheduling_unipc_multistep.py

@@ -1,4 +1,6 @@
-# Copyright 2025 TSAIL Team and The HuggingFace Team. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright 2024 TSAIL Team and The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,36 +14,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# DISCLAIMER: check https://huggingface.co/papers/2302.04867 and https://github.com/wl-zhao/UniPC for more info
+# DISCLAIMER: check https://arxiv.org/abs/2302.04867 and https://github.com/wl-zhao/UniPC for more info
 # The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
 # ==============================================================================
 #
-# Modified from diffusers==0.35.0.dev0
+# Modified from diffusers==0.33.0.dev0
 #
 # ==============================================================================
 
 import math
-from typing import List, Optional, Tuple, Union
 
 import numpy as np
+import scipy.stats
 import torch
-
 from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.utils import deprecate, is_scipy_available
-from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
+from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
+                                                   SchedulerMixin,
+                                                   SchedulerOutput)
+from diffusers.utils import deprecate
 
 from fastvideo.models.schedulers.base import BaseScheduler
 
-if is_scipy_available():
-    import scipy.stats
-
 
 # Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
 def betas_for_alpha_bar(
     num_diffusion_timesteps,
     max_beta=0.999,
     alpha_transform_type="cosine",
-):
+) -> torch.Tensor:
     """
     Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
     (1-beta) over time from t = [0,1].
@@ -62,16 +62,17 @@ def betas_for_alpha_bar(
     """
     if alpha_transform_type == "cosine":
 
-        def alpha_bar_fn(t):
-            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+        def alpha_bar_fn(t: float) -> float:
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2)**2
 
     elif alpha_transform_type == "exp":
 
-        def alpha_bar_fn(t):
+        def alpha_bar_fn(t: float) -> float:
             return math.exp(t * -12.0)
 
     else:
-        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+        raise ValueError(
+            f"Unsupported alpha_transform_type: {alpha_transform_type}")
 
     betas = []
     for i in range(num_diffusion_timesteps):
@@ -82,9 +83,9 @@ def betas_for_alpha_bar(
 
 
 # Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
-def rescale_zero_terminal_snr(betas):
+def rescale_zero_terminal_snr(betas: torch.Tensor) -> torch.Tensor:
     """
-    Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1)
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
 
 
     Args:
@@ -107,7 +108,8 @@ def rescale_zero_terminal_snr(betas):
     alphas_bar_sqrt -= alphas_bar_sqrt_T
 
     # Scale so the first timestep is back to the old value.
-    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 -
+                                            alphas_bar_sqrt_T)
 
     # Convert alphas_bar_sqrt to betas
     alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
@@ -174,8 +176,6 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         use_beta_sigmas (`bool`, *optional*, defaults to `False`):
             Whether to use beta sigmas for step sizes in the noise schedule during the sampling process. Refer to [Beta
             Sampling is All You Need](https://huggingface.co/papers/2407.12173) for more information.
-        use_flow_sigmas (`bool`, *optional*, defaults to `False`):
-            Whether to use flow sigmas for step sizes in the noise schedule during the sampling process.
         timestep_spacing (`str`, defaults to `"linspace"`):
             The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
             Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
@@ -200,7 +200,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         beta_start: float = 0.0001,
         beta_end: float = 0.02,
         beta_schedule: str = "linear",
-        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        trained_betas: np.ndarray | list[float] | None = None,
         solver_order: int = 2,
         prediction_type: str = "epsilon",
         thresholding: bool = False,
@@ -209,38 +209,44 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         predict_x0: bool = True,
         solver_type: str = "bh2",
         lower_order_final: bool = True,
-        disable_corrector: List[int] = [],
+        disable_corrector: tuple[int, ...] = (),
         solver_p: SchedulerMixin = None,
-        use_karras_sigmas: Optional[bool] = False,
-        use_exponential_sigmas: Optional[bool] = False,
-        use_beta_sigmas: Optional[bool] = False,
-        use_flow_sigmas: Optional[bool] = False,
-        flow_shift: Optional[float] = 1.0,
+        use_karras_sigmas: bool | None = False,
+        use_exponential_sigmas: bool | None = False,
+        use_beta_sigmas: bool | None = False,
+        use_flow_sigmas: bool | None = False,
+        flow_shift: float | None = 1.0,
         timestep_spacing: str = "linspace",
         steps_offset: int = 0,
-        final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+        final_sigmas_type: str | None = "zero",  # "zero", "sigma_min"
         rescale_betas_zero_snr: bool = False,
-        use_dynamic_shifting: bool = False,
-        time_shift_type: str = "exponential",
     ):
-        if self.config.use_beta_sigmas and not is_scipy_available():
-            raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
-        if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
+        if sum([
+                self.config.use_beta_sigmas, self.config.use_exponential_sigmas,
+                self.config.use_karras_sigmas
+        ]) > 1:
             raise ValueError(
                 "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
             )
         if trained_betas is not None:
             self.betas = torch.tensor(trained_betas, dtype=torch.float32)
         elif beta_schedule == "linear":
-            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+            self.betas = torch.linspace(beta_start,
+                                        beta_end,
+                                        num_train_timesteps,
+                                        dtype=torch.float32)
         elif beta_schedule == "scaled_linear":
             # this schedule is very specific to the latent diffusion model.
-            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+            self.betas = torch.linspace(beta_start**0.5,
+                                        beta_end**0.5,
+                                        num_train_timesteps,
+                                        dtype=torch.float32)**2
         elif beta_schedule == "squaredcos_cap_v2":
             # Glide cosine schedule
             self.betas = betas_for_alpha_bar(num_train_timesteps)
         else:
-            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
+            raise NotImplementedError(
+                f"{beta_schedule} is not implemented for {self.__class__}")
 
         if rescale_betas_zero_snr:
             self.betas = rescale_zero_terminal_snr(self.betas)
@@ -257,7 +263,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self.alpha_t = torch.sqrt(self.alphas_cumprod)
         self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)
         self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)
-        self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5
+        self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod)**0.5
 
         # standard deviation of the initial noise distribution
         self.init_noise_sigma = 1.0
@@ -266,22 +272,27 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             if solver_type in ["midpoint", "heun", "logrho"]:
                 self.register_to_config(solver_type="bh2")
             else:
-                raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")
+                raise NotImplementedError(
+                    f"{solver_type} is not implemented for {self.__class__}")
 
         self.predict_x0 = predict_x0
         # setable values
-        self.num_inference_steps = None
-        timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        self.num_inference_steps: int | None = None
+        timesteps = np.linspace(0,
+                                num_train_timesteps - 1,
+                                num_train_timesteps,
+                                dtype=np.float32)[::-1].copy()
         self.timesteps = torch.from_numpy(timesteps)
         self.model_outputs = [None] * solver_order
-        self.timestep_list = [None] * solver_order
+        self.timestep_list: list[int | torch.Tensor] = [None] * solver_order
         self.lower_order_nums = 0
-        self.disable_corrector = disable_corrector
+        self.disable_corrector = list(disable_corrector)
         self.solver_p = solver_p
         self.last_sample = None
-        self._step_index = None
-        self._begin_index = None
-        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self._step_index: int | None = None
+        self._begin_index: int | None = None
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
 
         BaseScheduler.__init__(self)
 
@@ -313,9 +324,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         """
         self._begin_index = begin_index
 
-    def set_timesteps(
-        self, num_inference_steps: int, device: Union[str, torch.device] = None, mu: Optional[float] = None
-    ):
+    def set_timesteps(self,
+                      num_inference_steps: int,
+                      device: str | torch.device = None):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
@@ -325,40 +336,41 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             device (`str` or `torch.device`, *optional*):
                 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
         """
-        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://huggingface.co/papers/2305.08891
-        if mu is not None:
-            assert self.config.use_dynamic_shifting and self.config.time_shift_type == "exponential"
-            self.config.flow_shift = np.exp(mu)
+        # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
         if self.config.timestep_spacing == "linspace":
-            timesteps = (
-                np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps + 1)
-                .round()[::-1][:-1]
-                .copy()
-                .astype(np.int64)
-            )
+            timesteps = (np.linspace(
+                0, self.config.num_train_timesteps - 1, num_inference_steps +
+                1).round()[::-1][:-1].copy().astype(np.int64))
         elif self.config.timestep_spacing == "leading":
-            step_ratio = self.config.num_train_timesteps // (num_inference_steps + 1)
+            step_ratio = self.config.num_train_timesteps // (
+                num_inference_steps + 1)
             # creates integer timesteps by multiplying by ratio
             # casting to int to avoid issues when num_inference_step is power of 3
-            timesteps = (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64)
+            timesteps = (np.arange(0, num_inference_steps + 1) *
+                         step_ratio).round()[::-1][:-1].copy().astype(np.int64)
             timesteps += self.config.steps_offset
         elif self.config.timestep_spacing == "trailing":
             step_ratio = self.config.num_train_timesteps / num_inference_steps
             # creates integer timesteps by multiplying by ratio
             # casting to int to avoid issues when num_inference_step is power of 3
-            timesteps = np.arange(self.config.num_train_timesteps, 0, -step_ratio).round().copy().astype(np.int64)
+            timesteps = np.arange(self.config.num_train_timesteps, 0,
+                                  -step_ratio).round().copy().astype(np.int64)
             timesteps -= 1
         else:
             raise ValueError(
                 f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
             )
 
-        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+        sigmas = np.array(
+            ((1 - self.alphas_cumprod) / self.alphas_cumprod)**0.5)
         if self.config.use_karras_sigmas:
             log_sigmas = np.log(sigmas)
             sigmas = np.flip(sigmas).copy()
-            sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
-            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
+            sigmas = self._convert_to_karras(
+                in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+            timesteps = np.array([
+                self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas
+            ]).round()
             if self.config.final_sigmas_type == "sigma_min":
                 sigma_last = sigmas[-1]
             elif self.config.final_sigmas_type == "zero":
@@ -371,8 +383,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         elif self.config.use_exponential_sigmas:
             log_sigmas = np.log(sigmas)
             sigmas = np.flip(sigmas).copy()
-            sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
-            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas])
+            sigmas = self._convert_to_exponential(
+                in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+            timesteps = np.array(
+                [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas])
             if self.config.final_sigmas_type == "sigma_min":
                 sigma_last = sigmas[-1]
             elif self.config.final_sigmas_type == "zero":
@@ -385,8 +399,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         elif self.config.use_beta_sigmas:
             log_sigmas = np.log(sigmas)
             sigmas = np.flip(sigmas).copy()
-            sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
-            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas])
+            sigmas = self._convert_to_beta(
+                in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+            timesteps = np.array(
+                [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas])
             if self.config.final_sigmas_type == "sigma_min":
                 sigma_last = sigmas[-1]
             elif self.config.final_sigmas_type == "zero":
@@ -397,9 +413,12 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
                 )
             sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)
         elif self.config.use_flow_sigmas:
-            alphas = np.linspace(1, 1 / self.config.num_train_timesteps, num_inference_steps + 1)
+            alphas = np.linspace(1, 1 / self.config.num_train_timesteps,
+                                 num_inference_steps + 1)
             sigmas = 1.0 - alphas
-            sigmas = np.flip(self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy()
+            sigmas = np.flip(
+                self.config.flow_shift * sigmas /
+                (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy()
             timesteps = (sigmas * self.config.num_train_timesteps).copy()
             if self.config.final_sigmas_type == "sigma_min":
                 sigma_last = sigmas[-1]
@@ -413,7 +432,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         else:
             sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
             if self.config.final_sigmas_type == "sigma_min":
-                sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
+                sigma_last = ((1 - self.alphas_cumprod[0]) /
+                              self.alphas_cumprod[0])**0.5
             elif self.config.final_sigmas_type == "zero":
                 sigma_last = 0
             else:
@@ -423,7 +443,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)
 
         self.sigmas = torch.from_numpy(sigmas)
-        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)
+        self.timesteps = torch.from_numpy(timesteps).to(device=device,
+                                                        dtype=torch.int64)
 
         self.num_inference_steps = len(timesteps)
 
@@ -438,7 +459,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         # add an index counter for schedulers that allow duplicated timesteps
         self._step_index = None
         self._begin_index = None
-        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
 
     # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
     def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
@@ -449,25 +471,31 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         pixels from saturation at each step. We find that dynamic thresholding results in significantly better
         photorealism as well as better image-text alignment, especially when using very large guidance weights."
 
-        https://huggingface.co/papers/2205.11487
+        https://arxiv.org/abs/2205.11487
         """
         dtype = sample.dtype
         batch_size, channels, *remaining_dims = sample.shape
 
         if dtype not in (torch.float32, torch.float64):
-            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+            sample = sample.float(
+            )  # upcast for quantile calculation, and clamp not implemented for cpu half
 
         # Flatten sample for doing quantile calculation along each image
         sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
 
         abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
 
-        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.quantile(abs_sample,
+                           self.config.dynamic_thresholding_ratio,
+                           dim=1)
         s = torch.clamp(
             s, min=1, max=self.config.sample_max_value
         )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
-        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
-        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+        s = s.unsqueeze(
+            1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(
+            sample, -s, s
+        ) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
 
         sample = sample.reshape(batch_size, channels, *remaining_dims)
         sample = sample.to(dtype)
@@ -475,7 +503,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return sample
 
     # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t
-    def _sigma_to_t(self, sigma, log_sigmas):
+    def _sigma_to_t(self, sigma, log_sigmas) -> torch.Tensor:
         # get log sigma
         log_sigma = np.log(np.maximum(sigma, 1e-10))
 
@@ -483,7 +511,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         dists = log_sigma - log_sigmas[:, np.newaxis]
 
         # get sigmas range
-        low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
+        low_idx = np.cumsum(
+            (dists >= 0),
+            axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
         high_idx = low_idx + 1
 
         low = log_sigmas[low_idx]
@@ -499,18 +529,20 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return t
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._sigma_to_alpha_sigma_t
-    def _sigma_to_alpha_sigma_t(self, sigma):
+    def _sigma_to_alpha_sigma_t(
+            self, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         if self.config.use_flow_sigmas:
             alpha_t = 1 - sigma
             sigma_t = sigma
         else:
-            alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
+            alpha_t = 1 / ((sigma**2 + 1)**0.5)
             sigma_t = sigma * alpha_t
 
         return alpha_t, sigma_t
 
     # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
-    def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
+    def _convert_to_karras(self, in_sigmas: torch.Tensor,
+                           num_inference_steps) -> torch.Tensor:
         """Constructs the noise schedule of Karras et al. (2022)."""
 
         # Hack to make sure that other schedulers which copy this function don't break
@@ -530,13 +562,14 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
         rho = 7.0  # 7.0 is the value used in the paper
         ramp = np.linspace(0, 1, num_inference_steps)
-        min_inv_rho = sigma_min ** (1 / rho)
-        max_inv_rho = sigma_max ** (1 / rho)
-        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        min_inv_rho = sigma_min**(1 / rho)
+        max_inv_rho = sigma_max**(1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho
         return sigmas
 
     # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
-    def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
+    def _convert_to_exponential(self, in_sigmas: torch.Tensor,
+                                num_inference_steps: int) -> torch.Tensor:
         """Constructs an exponential noise schedule."""
 
         # Hack to make sure that other schedulers which copy this function don't break
@@ -554,13 +587,17 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
         sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
 
-        sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps))
+        sigmas = np.exp(
+            np.linspace(math.log(sigma_max), math.log(sigma_min),
+                        num_inference_steps))
         return sigmas
 
     # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
-    def _convert_to_beta(
-        self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6
-    ) -> torch.Tensor:
+    def _convert_to_beta(self,
+                         in_sigmas: torch.Tensor,
+                         num_inference_steps: int,
+                         alpha: float = 0.6,
+                         beta: float = 0.6) -> torch.Tensor:
         """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)"""
 
         # Hack to make sure that other schedulers which copy this function don't break
@@ -578,15 +615,12 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
         sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
 
-        sigmas = np.array(
-            [
-                sigma_min + (ppf * (sigma_max - sigma_min))
-                for ppf in [
-                    scipy.stats.beta.ppf(timestep, alpha, beta)
-                    for timestep in 1 - np.linspace(0, 1, num_inference_steps)
-                ]
+        sigmas = np.array([
+            sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in [
+                scipy.stats.beta.ppf(timestep, alpha, beta)
+                for timestep in 1 - np.linspace(0, 1, num_inference_steps)
             ]
-        )
+        ])
         return sigmas
 
     def convert_model_output(
@@ -616,7 +650,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             if len(args) > 1:
                 sample = args[1]
             else:
-                raise ValueError("missing `sample` as a required keyword argument")
+                raise ValueError(
+                    "missing `sample` as a required keyword argument")
         if timestep is not None:
             deprecate(
                 "timesteps",
@@ -659,15 +694,14 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             else:
                 raise ValueError(
                     f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
-                    " `v_prediction` for the UniPCMultistepScheduler."
-                )
+                    " `v_prediction` for the UniPCMultistepScheduler.")
 
     def multistep_uni_p_bh_update(
         self,
         model_output: torch.Tensor,
         *args,
         sample: torch.Tensor = None,
-        order: int = None,
+        order: int | None = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -687,17 +721,20 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             `torch.Tensor`:
                 The sample tensor at the previous timestep.
         """
-        prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None)
+        prev_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "prev_timestep", None)
         if sample is None:
             if len(args) > 1:
                 sample = args[1]
             else:
-                raise ValueError("missing `sample` as a required keyword argument")
+                raise ValueError(
+                    " missing `sample` as a required keyword argument")
         if order is None:
             if len(args) > 2:
                 order = args[2]
             else:
-                raise ValueError("missing `order` as a required keyword argument")
+                raise ValueError(
+                    " missing `order` as a required keyword argument")
         if prev_timestep is not None:
             deprecate(
                 "prev_timestep",
@@ -714,7 +751,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             x_t = self.solver_p.step(model_output, s0, x).prev_sample
             return x_t
 
-        sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
+        sigma_t, sigma_s0 = self.sigmas[self.step_index +
+                                        1], self.sigmas[self.step_index]
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
 
@@ -728,7 +766,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         D1s = []
         for i in range(1, order):
             si = self.step_index - i
-            mi = model_output_list[-(i + 1)]
+            mi: torch.Tensor = model_output_list[-(i + 1)]
             alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
             lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
             rk = (lambda_si - lambda_s0) / h
@@ -760,30 +798,32 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             factorial_i *= i + 1
             h_phi_k = h_phi_k / hh - 1 / factorial_i
 
-        R = torch.stack(R)
+        R_tensor: torch.Tensor = torch.stack(R)
         b = torch.tensor(b, device=device)
 
+        D1s_tensor: torch.Tensor | None = None
         if len(D1s) > 0:
-            D1s = torch.stack(D1s, dim=1)  # (B, K)
+            D1s_tensor = torch.stack(D1s, dim=1)  # (B, K)
             # for order 2, we use a simplified version
             if order == 2:
                 rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
             else:
-                rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype)
+                rhos_p = torch.linalg.solve(R_tensor[:-1, :-1],
+                                            b[:-1]).to(device).to(x.dtype)
         else:
-            D1s = None
+            D1s_tensor = None
 
         if self.predict_x0:
             x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
-            if D1s is not None:
-                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s)
+            if D1s_tensor is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s_tensor)
             else:
                 pred_res = 0
             x_t = x_t_ - alpha_t * B_h * pred_res
         else:
             x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
-            if D1s is not None:
-                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s)
+            if D1s_tensor is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s_tensor)
             else:
                 pred_res = 0
             x_t = x_t_ - sigma_t * B_h * pred_res
@@ -795,9 +835,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self,
         this_model_output: torch.Tensor,
         *args,
-        last_sample: torch.Tensor = None,
-        this_sample: torch.Tensor = None,
-        order: int = None,
+        last_sample: torch.Tensor | None = None,
+        this_sample: torch.Tensor | None = None,
+        order: int | None = None,
         **kwargs,
     ) -> torch.Tensor:
         """
@@ -819,22 +859,26 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             `torch.Tensor`:
                 The corrected sample tensor at the current timestep.
         """
-        this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None)
+        this_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "this_timestep", None)
         if last_sample is None:
             if len(args) > 1:
                 last_sample = args[1]
             else:
-                raise ValueError("missing `last_sample` as a required keyword argument")
+                raise ValueError(
+                    " missing`last_sample` as a required keyword argument")
         if this_sample is None:
             if len(args) > 2:
                 this_sample = args[2]
             else:
-                raise ValueError("missing `this_sample` as a required keyword argument")
+                raise ValueError(
+                    " missing`this_sample` as a required keyword argument")
         if order is None:
             if len(args) > 3:
                 order = args[3]
             else:
-                raise ValueError("missing `order` as a required keyword argument")
+                raise ValueError(
+                    " missing`order` as a required keyword argument")
         if this_timestep is not None:
             deprecate(
                 "this_timestep",
@@ -849,7 +893,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         x_t = this_sample
         model_t = this_model_output
 
-        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1]
+        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[
+            self.step_index - 1]
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
 
@@ -863,7 +908,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         D1s = []
         for i in range(1, order):
             si = self.step_index - (i + 1)
-            mi = model_output_list[-(i + 1)]
+            mi: torch.Tensor = model_output_list[-(i + 1)]
             alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
             lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
             rk = (lambda_si - lambda_s0) / h
@@ -898,10 +943,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         R = torch.stack(R)
         b = torch.tensor(b, device=device)
 
-        if len(D1s) > 0:
-            D1s = torch.stack(D1s, dim=1)
-        else:
-            D1s = None
+        D1s_tensor: torch.Tensor | None = torch.stack(
+            D1s, dim=1) if len(D1s) > 0 else None
 
         # for order 1, we use a simplified version
         if order == 1:
@@ -911,16 +954,18 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
         if self.predict_x0:
             x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
-            if D1s is not None:
-                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            if D1s_tensor is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1],
+                                        D1s_tensor)
             else:
                 corr_res = 0
             D1_t = model_t - m0
             x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
         else:
             x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
-            if D1s is not None:
-                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            if D1s_tensor is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1],
+                                        D1s_tensor)
             else:
                 corr_res = 0
             D1_t = model_t - m0
@@ -929,7 +974,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return x_t
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.index_for_timestep
-    def index_for_timestep(self, timestep, schedule_timesteps=None):
+    def index_for_timestep(self, timestep, schedule_timesteps=None) -> int:
         if schedule_timesteps is None:
             schedule_timesteps = self.timesteps
 
@@ -949,7 +994,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return step_index
 
     # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
-    def _init_step_index(self, timestep):
+    def _init_step_index(self, timestep) -> None:
         """
         Initialize the step_index counter for the scheduler.
         """
@@ -964,10 +1009,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: Union[int, torch.Tensor],
+        timestep: int | torch.Tensor,
         sample: torch.Tensor,
         return_dict: bool = True,
-    ) -> Union[SchedulerOutput, Tuple]:
+    ) -> SchedulerOutput | tuple:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
         the multistep UniPC.
@@ -996,11 +1041,13 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         if self.step_index is None:
             self._init_step_index(timestep)
 
-        use_corrector = (
-            self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None
-        )
+        assert self.step_index is not None
+        use_corrector = (self.step_index > 0
+                         and self.step_index - 1 not in self.disable_corrector
+                         and self.last_sample is not None)
 
-        model_output_convert = self.convert_model_output(model_output, sample=sample)
+        model_output_convert = self.convert_model_output(model_output,
+                                                         sample=sample)
         if use_corrector:
             sample = self.multistep_uni_c_bh_update(
                 this_model_output=model_output_convert,
@@ -1017,16 +1064,19 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self.timestep_list[-1] = timestep
 
         if self.config.lower_order_final:
-            this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index)
+            this_order = min(self.config.solver_order,
+                             len(self.timesteps) - self.step_index)
         else:
             this_order = self.config.solver_order
 
-        self.this_order = min(this_order, self.lower_order_nums + 1)  # warmup for multistep
+        self.this_order: int = min(this_order, self.lower_order_nums +
+                                   1)  # warmup for multistep
         assert self.this_order > 0
 
         self.last_sample = sample
         prev_sample = self.multistep_uni_p_bh_update(
-            model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
+            model_output=
+            model_output,  # pass the original non-converted model output, in case solver-p is used
             sample=sample,
             order=self.this_order,
         )
@@ -1035,14 +1085,16 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             self.lower_order_nums += 1
 
         # upon completion increase step index by one
+        assert self._step_index is not None
         self._step_index += 1
 
         if not return_dict:
-            return (prev_sample,)
+            return (prev_sample, )
 
         return SchedulerOutput(prev_sample=prev_sample)
 
-    def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+    def scale_model_input(self, sample: torch.Tensor, *args,
+                          **kwargs) -> torch.Tensor:
         """
         Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
         current timestep.
@@ -1065,18 +1117,25 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         timesteps: torch.IntTensor,
     ) -> torch.Tensor:
         # Make sure sigmas and timesteps have the same device and dtype as original_samples
-        sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype)
-        if original_samples.device.type == "mps" and torch.is_floating_point(timesteps):
+        sigmas = self.sigmas.to(device=original_samples.device,
+                                dtype=original_samples.dtype)
+        if original_samples.device.type == "mps" and torch.is_floating_point(
+                timesteps):
             # mps does not support float64
-            schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32)
-            timesteps = timesteps.to(original_samples.device, dtype=torch.float32)
+            schedule_timesteps = self.timesteps.to(original_samples.device,
+                                                   dtype=torch.float32)
+            timesteps = timesteps.to(original_samples.device,
+                                     dtype=torch.float32)
         else:
             schedule_timesteps = self.timesteps.to(original_samples.device)
             timesteps = timesteps.to(original_samples.device)
 
         # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
         if self.begin_index is None:
-            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps]
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps)
+                for t in timesteps
+            ]
         elif self.step_index is not None:
             # add_noise is called after first denoising step (for inpainting)
             step_indices = [self.step_index] * timesteps.shape[0]
@@ -1093,4 +1152,4 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         return noisy_samples
 
     def __len__(self):
-        return self.config.num_train_timesteps
+        return self.config.num_train_timesteps

fastvideo/models/vaes/wanvae.py

@@ -20,7 +20,6 @@ from contextlib import contextmanager
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from einops import rearrange
 
 from fastvideo.configs.models.vaes import WanVAEConfig
 from fastvideo.layers.activation import get_act_fn
@@ -33,123 +32,22 @@ CACHE_T = 2
 is_first_frame = contextvars.ContextVar("is_first_frame", default=False)
 feat_cache = contextvars.ContextVar("feat_cache", default=None)
 feat_idx = contextvars.ContextVar("feat_idx", default=0)
-first_chunk = contextvars.ContextVar("first_chunk", default=None)
 
 
 @contextmanager
 def forward_context(first_frame_arg=False,
                     feat_cache_arg=None,
-                    feat_idx_arg=None,
-                    first_chunk_arg=None):
+                    feat_idx_arg=None):
     is_first_frame_token = is_first_frame.set(first_frame_arg)
     feat_cache_token = feat_cache.set(feat_cache_arg)
     feat_idx_token = feat_idx.set(feat_idx_arg)
-    first_chunk_token = first_chunk.set(first_chunk_arg)
     try:
         yield
     finally:
         is_first_frame.reset(is_first_frame_token)
         feat_cache.reset(feat_cache_token)
         feat_idx.reset(feat_idx_token)
-        first_chunk.reset(first_chunk_token)
 
-class AvgDown3D(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        factor_t,
-        factor_s=1,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.factor_t = factor_t
-        self.factor_s = factor_s
-        self.factor = self.factor_t * self.factor_s * self.factor_s
-
-        assert in_channels * self.factor % out_channels == 0
-        self.group_size = in_channels * self.factor // out_channels
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
-        pad = (0, 0, 0, 0, pad_t, 0)
-        x = F.pad(x, pad)
-        B, C, T, H, W = x.shape
-        x = x.view(
-            B,
-            C,
-            T // self.factor_t,
-            self.factor_t,
-            H // self.factor_s,
-            self.factor_s,
-            W // self.factor_s,
-            self.factor_s,
-        )
-        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
-        x = x.view(
-            B,
-            C * self.factor,
-            T // self.factor_t,
-            H // self.factor_s,
-            W // self.factor_s,
-        )
-        x = x.view(
-            B,
-            self.out_channels,
-            self.group_size,
-            T // self.factor_t,
-            H // self.factor_s,
-            W // self.factor_s,
-        )
-        x = x.mean(dim=2)
-        return x
-
-
-class DupUp3D(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        factor_t,
-        factor_s=1,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-
-        self.factor_t = factor_t
-        self.factor_s = factor_s
-        self.factor = self.factor_t * self.factor_s * self.factor_s
-
-        assert out_channels * self.factor % in_channels == 0
-        self.repeats = out_channels * self.factor // in_channels
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = x.repeat_interleave(self.repeats, dim=1)
-        x = x.view(
-            x.size(0),
-            self.out_channels,
-            self.factor_t,
-            self.factor_s,
-            self.factor_s,
-            x.size(2),
-            x.size(3),
-            x.size(4),
-        )
-        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
-        x = x.view(
-            x.size(0),
-            self.out_channels,
-            x.size(2) * self.factor_t,
-            x.size(4) * self.factor_s,
-            x.size(6) * self.factor_s,
-        )
-
-        _first_chunk = first_chunk.get()
-        if _first_chunk:
-            x = x[:, :, self.factor_t - 1 :, :, :]
-        return x
 
 class WanCausalConv3d(nn.Conv3d):
     r"""
@@ -260,26 +158,20 @@ class WanResample(nn.Module):
             - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution.
     """
 
-    def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None:
+    def __init__(self, dim: int, mode: str) -> None:
         super().__init__()
         self.dim = dim
         self.mode = mode
 
-        # default to dim //2
-        if upsample_out_dim is None:
-            upsample_out_dim = dim // 2
-
         # layers
         if mode == "upsample2d":
             self.resample = nn.Sequential(
                 WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
-                nn.Conv2d(dim, upsample_out_dim, 3, padding=1),
-            )
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
         elif mode == "upsample3d":
             self.resample = nn.Sequential(
                 WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
-                nn.Conv2d(dim, upsample_out_dim, 3, padding=1),
-            )
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
             self.time_conv = WanCausalConv3d(dim,
                                              dim * 2, (3, 1, 1),
                                              padding=(1, 0, 0))
@@ -552,40 +444,6 @@ class WanMidBlock(nn.Module):
 
         return x
 
-class WanResidualDownBlock(nn.Module):
-    def __init__(self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False):
-        super().__init__()
-
-        # Shortcut path with downsample
-        self.avg_shortcut = AvgDown3D(
-            in_dim,
-            out_dim,
-            factor_t=2 if temperal_downsample else 1,
-            factor_s=2 if down_flag else 1,
-        )
-
-        # Main path with residual blocks and downsample
-        resnets = []
-        for _ in range(num_res_blocks):
-            resnets.append(WanResidualBlock(in_dim, out_dim, dropout))
-            in_dim = out_dim
-        self.resnets = nn.ModuleList(resnets)
-
-        # Add the final downsample block
-        if down_flag:
-            mode = "downsample3d" if temperal_downsample else "downsample2d"
-            self.downsampler = WanResample(out_dim, mode=mode)
-        else:
-            self.downsampler = None
-
-    def forward(self, x):
-        x_copy = x.clone()
-        for resnet in self.resnets:
-            x = resnet(x)
-        if self.downsampler is not None:
-            x = self.downsampler(x)
-
-        return x + self.avg_shortcut(x_copy)
 
 class WanEncoder3d(nn.Module):
     r"""
@@ -604,7 +462,6 @@ class WanEncoder3d(nn.Module):
 
     def __init__(
         self,
-        in_channels: int = 3,
         dim=128,
         z_dim=4,
         dim_mult=(1, 2, 4, 4),
@@ -613,7 +470,6 @@ class WanEncoder3d(nn.Module):
         temperal_downsample=(True, True, False),
         dropout=0.0,
         non_linearity: str = "silu",
-        is_residual: bool = False,  # wan 2.2 vae use a residual downblock
     ):
         super().__init__()
         self.dim = dim
@@ -630,36 +486,26 @@ class WanEncoder3d(nn.Module):
         scale = 1.0
 
         # init block
-        self.conv_in = WanCausalConv3d(in_channels, dims[0], 3, padding=1)
+        self.conv_in = WanCausalConv3d(3, dims[0], 3, padding=1)
 
         # downsample blocks
         self.down_blocks = nn.ModuleList([])
         for i, (in_dim,
                 out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)):
             # residual (+attention) blocks
-            if is_residual:
+            for _ in range(num_res_blocks):
                 self.down_blocks.append(
-                    WanResidualDownBlock(
-                        in_dim,
-                        out_dim,
-                        dropout,
-                        num_res_blocks,
-                        temperal_downsample=temperal_downsample[i] if i != len(dim_mult) - 1 else False,
-                        down_flag=i != len(dim_mult) - 1,
-                    )
-                )
-            else:
-                for _ in range(num_res_blocks):
-                    self.down_blocks.append(WanResidualBlock(in_dim, out_dim, dropout))
-                    if scale in attn_scales:
-                        self.down_blocks.append(WanAttentionBlock(out_dim))
-                    in_dim = out_dim
+                    WanResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    self.down_blocks.append(WanAttentionBlock(out_dim))
+                in_dim = out_dim
 
-                # downsample block
-                if i != len(dim_mult) - 1:
-                    mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
-                    self.down_blocks.append(WanResample(out_dim, mode=mode))
-                    scale /= 2.0
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = "downsample3d" if temperal_downsample[
+                    i] else "downsample2d"
+                self.down_blocks.append(WanResample(out_dim, mode=mode))
+                scale /= 2.0
 
         # middle blocks
         self.mid_block = WanMidBlock(out_dim,
@@ -726,85 +572,6 @@ class WanEncoder3d(nn.Module):
             x = self.conv_out(x)
         return x
 
-# adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
-class WanResidualUpBlock(nn.Module):
-    """
-    A block that handles upsampling for the WanVAE decoder.
-    Args:
-        in_dim (int): Input dimension
-        out_dim (int): Output dimension
-        num_res_blocks (int): Number of residual blocks
-        dropout (float): Dropout rate
-        temperal_upsample (bool): Whether to upsample on temporal dimension
-        up_flag (bool): Whether to upsample or not
-        non_linearity (str): Type of non-linearity to use
-    """
-
-    def __init__(
-        self,
-        in_dim: int,
-        out_dim: int,
-        num_res_blocks: int,
-        dropout: float = 0.0,
-        temperal_upsample: bool = False,
-        up_flag: bool = False,
-        non_linearity: str = "silu",
-    ):
-        super().__init__()
-        self.in_dim = in_dim
-        self.out_dim = out_dim
-
-        if up_flag:
-            self.avg_shortcut = DupUp3D(
-                in_dim,
-                out_dim,
-                factor_t=2 if temperal_upsample else 1,
-                factor_s=2,
-            )
-        else:
-            self.avg_shortcut = None
-
-        # create residual blocks
-        resnets = []
-        current_dim = in_dim
-        for _ in range(num_res_blocks + 1):
-            resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity))
-            current_dim = out_dim
-
-        self.resnets = nn.ModuleList(resnets)
-
-        # Add upsampling layer if needed
-        if up_flag:
-            upsample_mode = "upsample3d" if temperal_upsample else "upsample2d"
-            self.upsampler = WanResample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim)
-        else:
-            self.upsampler = None
-
-        self.gradient_checkpointing = False
-
-    def forward(self, x):
-        """
-        Forward pass through the upsampling block.
-        Args:
-            x (torch.Tensor): Input tensor
-            feat_cache (list, optional): Feature cache for causal convolutions
-            feat_idx (list, optional): Feature index for cache management
-        Returns:
-            torch.Tensor: Output tensor
-        """
-        if self.avg_shortcut is not None:
-            x_copy = x.clone()
-
-        for resnet in self.resnets:
-            x = resnet(x)
-
-        if self.upsampler is not None:
-            x = self.upsampler(x)
-
-        if self.avg_shortcut is not None:
-            x = x + self.avg_shortcut(x_copy)
-
-        return x
 
 class WanUpBlock(nn.Module):
     """
@@ -896,8 +663,6 @@ class WanDecoder3d(nn.Module):
         temperal_upsample=(False, True, True),
         dropout=0.0,
         non_linearity: str = "silu",
-        out_channels: int = 3,
-        is_residual: bool = False,
     ):
         super().__init__()
         self.dim = dim
@@ -912,6 +677,7 @@ class WanDecoder3d(nn.Module):
 
         # dimensions
         dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
 
         # init block
         self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1)
@@ -927,44 +693,33 @@ class WanDecoder3d(nn.Module):
         for i, (in_dim,
                 out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)):
             # residual (+attention) blocks
-            if i > 0 and not is_residual:
-                # wan vae 2.1
+            if i > 0:
                 in_dim = in_dim // 2
 
-            # determine if we need upsampling
-            up_flag = i != len(dim_mult) - 1
-            # determine upsampling mode, if not upsampling, set to None
+            # Determine if we need upsampling
             upsample_mode = None
-            if up_flag and temperal_upsample[i]:
-                upsample_mode = "upsample3d"
-            elif up_flag:
-                upsample_mode = "upsample2d"
+            if i != len(dim_mult) - 1:
+                upsample_mode = "upsample3d" if temperal_upsample[
+                    i] else "upsample2d"
 
             # Create and add the upsampling block
-            if is_residual:
-                up_block = WanResidualUpBlock(
-                    in_dim=in_dim,
-                    out_dim=out_dim,
-                    num_res_blocks=num_res_blocks,
-                    dropout=dropout,
-                    temperal_upsample=temperal_upsample[i] if up_flag else False,
-                    up_flag=up_flag,
-                    non_linearity=non_linearity,
-                )
-            else:
-                up_block = WanUpBlock(
-                    in_dim=in_dim,
-                    out_dim=out_dim,
-                    num_res_blocks=num_res_blocks,
-                    dropout=dropout,
-                    upsample_mode=upsample_mode,
-                    non_linearity=non_linearity,
-                )
+            up_block = WanUpBlock(
+                in_dim=in_dim,
+                out_dim=out_dim,
+                num_res_blocks=num_res_blocks,
+                dropout=dropout,
+                upsample_mode=upsample_mode,
+                non_linearity=non_linearity,
+            )
             self.up_blocks.append(up_block)
 
+            # Update scale for next iteration
+            if upsample_mode is not None:
+                scale *= 2.0
+
         # output blocks
         self.norm_out = WanRMS_norm(out_dim, images=False)
-        self.conv_out = WanCausalConv3d(out_dim, out_channels, 3, padding=1)
+        self.conv_out = WanCausalConv3d(out_dim, 3, 3, padding=1)
 
         self.gradient_checkpointing = False
 
@@ -1021,41 +776,6 @@ class WanDecoder3d(nn.Module):
             x = self.conv_out(x)
         return x
 
-def patchify(x, patch_size):
-    if patch_size == 1:
-        return x
-
-    if x.dim() == 4:
-        x = rearrange(
-            x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
-    elif x.dim() == 5:
-        x = rearrange(
-            x,
-            "b c f (h q) (w r) -> b (c r q) f h w",
-            q=patch_size,
-            r=patch_size,
-        )
-    else:
-        raise ValueError(f"Invalid input shape: {x.shape}")
-
-    return x
-
-def unpatchify(x, patch_size):
-    if patch_size == 1:
-        return x
-
-    if x.dim() == 4:
-        x = rearrange(
-            x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
-    elif x.dim() == 5:
-        x = rearrange(
-            x,
-            "b (c r q) f h w -> b c f (h q) (w r)",
-            q=patch_size,
-            r=patch_size,
-        )
-
-    return x
 
 class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
     r"""
@@ -1075,43 +795,24 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
         self.z_dim = config.z_dim
         self.temperal_downsample = list(config.temperal_downsample)
         self.temperal_upsample = list(config.temperal_downsample)[::-1]
-
-        if config.decoder_base_dim is None:
-            decoder_base_dim = config.base_dim
-        else:
-            decoder_base_dim = config.decoder_base_dim
-
         self.latents_mean = list(config.latents_mean)
         self.latents_std = list(config.latents_std)
         self.shift_factor = config.shift_factor
 
         if config.load_encoder:
-            self.encoder = WanEncoder3d(
-                in_channels=config.in_channels,
-                dim=config.base_dim, 
-                z_dim=self.z_dim * 2,
-                dim_mult=config.dim_mult, 
-                num_res_blocks=config.num_res_blocks,
-                attn_scales=config.attn_scales,
-                temperal_downsample=self.temperal_downsample,
-                dropout=config.dropout,
-                is_residual=config.is_residual,
-            )
+            self.encoder = WanEncoder3d(config.base_dim, self.z_dim * 2,
+                                        config.dim_mult, config.num_res_blocks,
+                                        config.attn_scales,
+                                        self.temperal_downsample,
+                                        config.dropout)
         self.quant_conv = WanCausalConv3d(self.z_dim * 2, self.z_dim * 2, 1)
         self.post_quant_conv = WanCausalConv3d(self.z_dim, self.z_dim, 1)
 
         if config.load_decoder:
-            self.decoder = WanDecoder3d(
-                dim=decoder_base_dim, 
-                z_dim=self.z_dim,
-                dim_mult=config.dim_mult, 
-                num_res_blocks=config.num_res_blocks,
-                attn_scales=config.attn_scales,
-                temperal_upsample=self.temperal_upsample, 
-                dropout=config.dropout,
-                out_channels=config.out_channels,
-                is_residual=config.is_residual,
-            )
+            self.decoder = WanDecoder3d(config.base_dim, self.z_dim,
+                                        config.dim_mult, config.num_res_blocks,
+                                        config.attn_scales,
+                                        self.temperal_upsample, config.dropout)
 
         self.use_feature_cache = config.use_feature_cache
 
@@ -1137,8 +838,6 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
     def encode(self, x: torch.Tensor) -> torch.Tensor:
         if self.use_feature_cache:
             self.clear_cache()
-            if self.config.patch_size is not None:
-                x = patchify(x, patch_size=self.config.patch_size)
             with forward_context(feat_cache_arg=self._enc_feat_map,
                                  feat_idx_arg=self._enc_conv_idx):
                 t = x.shape[2]
@@ -1204,17 +903,11 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
                 for i in range(iter_):
                     feat_idx.set(0)
                     if i == 0:
-                        first_chunk.set(True)
                         out = self.decoder(x[:, :, i:i + 1, :, :])
                     else:
-                        first_chunk.set(False)
                         out_ = self.decoder(x[:, :, i:i + 1, :, :])
                         out = torch.cat([out, out_], 2)
 
-            if self.config.patch_size is not None:
-                out = unpatchify(out, patch_size=self.config.patch_size)
-                
-            out = out.float()
             out = torch.clamp(out, min=-1.0, max=1.0)
             self.clear_cache()
         else:

fastvideo/pipelines/composed_pipeline_base.py

@@ -234,9 +234,6 @@ class ComposedPipelineBase(ABC):
         # remove keys that are not pipeline modules
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
-        # @TODO(Wei): Temporary hack
-        model_index.pop("boundary_ratio", None)
-        model_index.pop("expand_timesteps", None)
 
         # some sanity checks
         assert len(
@@ -258,8 +255,6 @@ class ComposedPipelineBase(ABC):
         modules = {}
         for module_name, (transformers_or_diffusers,
                           architecture) in model_index.items():
-            if transformers_or_diffusers is None:
-                continue
             if module_name not in required_modules:
                 logger.info("Skipping module %s", module_name)
                 continue

fastvideo/pipelines/pipeline_registry.py

@@ -19,8 +19,8 @@ logger = init_logger(__name__)
 # map pipeline name to folder name
 _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "WanPipeline": "wan",
-    "WanDMDPipeline": "wan",
     "WanImageToVideoPipeline": "wan",
+    "WanDMDPipeline": "wan",
     "StepVideoPipeline": "stepvideo",
     "HunyuanVideoPipeline": "hunyuan",
 }

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -23,6 +23,7 @@ def main(args) -> None:
     assert num_gpus == 1, "Only support 1 GPU"
     pipeline_config = PipelineConfig.from_pretrained(args.model_path)
     kwargs = {
+        "dit_cpu_offload": False,
         "vae_precision": "fp32",
         "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
     }
@@ -30,9 +31,6 @@ def main(args) -> None:
     fastvideo_args = FastVideoArgs(
         model_path=args.model_path,
         num_gpus=get_world_size(),
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=False,
         pipeline_config=pipeline_config,
     )
     PreprocessPipeline = PreprocessPipeline_I2V if args.preprocess_task == "i2v" else PreprocessPipeline_T2V

fastvideo/pipelines/stages/denoising.py

@@ -12,7 +12,7 @@ from typing import Any
 import torch
 from einops import rearrange
 from tqdm.auto import tqdm
-import time
+
 from fastvideo.attention import get_attn_backend
 from fastvideo.configs.pipelines.base import STA_Mode
 from fastvideo.distributed import (get_local_torch_device, get_sp_parallel_rank,
@@ -30,7 +30,7 @@ from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.utils import dict_to_3d_list
+from fastvideo.utils import dict_to_3d_list, masks_like
 
 try:
     from fastvideo.attention.backends.sliding_tile_attn import (
@@ -88,11 +88,6 @@ class DenoisingStage(PipelineStage):
         Returns:
             The batch with denoised latents.
         """
-        # Start timing for the entire DenoisingStage
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        stage_start_time = time.time()
-        
         pipeline = self.pipeline() if self.pipeline else None
         if not fastvideo_args.model_loaded["transformer"]:
             loader = TransformerLoader()
@@ -199,9 +194,21 @@ class DenoisingStage(PipelineStage):
                 # Expand latents for I2V
                 latent_model_input = latents.to(target_dtype)
                 if batch.image_latent is not None:
+                    assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task"
                     latent_model_input = torch.cat(
                         [latent_model_input, batch.image_latent],
                         dim=1).to(target_dtype)
+                elif fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+                    # TI2V directly replaces the first frame of the latent with
+                    # the image latent instead of appending along the channel dim
+                    assert batch.image_latent is None, "TI2V task should not have image latents"
+                    z = self.get_module("vae").encode([batch.pil_image])
+                    mask1, mask2 = masks_like([latent_model_input],
+                                              zero=True,
+                                              generator=batch.generator)
+                    latent_model_input = (
+                        1. - mask2[0]) * z[0] + mask2[0] * latent_model_input
+
                 assert torch.isnan(latent_model_input).sum() == 0
                 latent_model_input = self.scheduler.scale_model_input(
                     latent_model_input, t)
@@ -233,17 +240,10 @@ class DenoisingStage(PipelineStage):
                             self.attn_metadata_builder = self.attn_metadata_builder_cls(
                             )
                             # TODO(will): clean this up
-                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
-                                current_timestep=i,  # type: ignore
-                                raw_latent_shape=batch.
-                                raw_latent_shape[2:5],  # type: ignore
-                                patch_size=fastvideo_args.
-                                pipeline_config.  # type: ignore
-                                dit_config.patch_size,  # type: ignore
-                                STA_param=batch.STA_param,  # type: ignore
-                                VSA_sparsity=fastvideo_args.
-                                VSA_sparsity,  # type: ignore
-                                device=get_local_torch_device(),
+                            attn_metadata = self.attn_metadata_builder.build(
+                                current_timestep=i,
+                                forward_batch=batch,
+                                fastvideo_args=fastvideo_args,
                             )
                             assert attn_metadata is not None, "attn_metadata cannot be None"
                         else:
@@ -606,11 +606,6 @@ class DmdDenoisingStage(DenoisingStage):
         Returns:
             The batch with denoised latents.
         """
-        # Start timing for the entire DmdDenoisingStage
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        dmd_stage_start_time = time.time()
-        
         # Setup precision and autocast settings
         # TODO(will): make the precision configurable for inference
         # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
@@ -691,13 +686,7 @@ class DmdDenoisingStage(DenoisingStage):
                 image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
                 batch.image_latent = image_latent
 
-
-        
-        # Synchronize CUDA before starting timer
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        denoise_start_time = time.time()
-        
+        # Run denoising loop
         with self.progress_bar(total=len(timesteps)) as progress_bar:
             for i, t in enumerate(timesteps):
                 # Skip if interrupted
@@ -731,7 +720,9 @@ class DmdDenoisingStage(DenoisingStage):
                 with torch.autocast(device_type="cuda",
                                     dtype=target_dtype,
                                     enabled=autocast_enabled):
-                    if (vsa_available and self.attn_backend
+                    if (st_attn_available
+                            and self.attn_backend == SlidingTileAttentionBackend
+                        ) or (vsa_available and self.attn_backend
                               == VideoSparseAttentionBackend):
                         self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
                         )
@@ -740,18 +731,11 @@ class DmdDenoisingStage(DenoisingStage):
                             self.attn_metadata_builder = self.attn_metadata_builder_cls(
                             )
                             # TODO(will): clean this up
-                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
-                                current_timestep=i,  # type: ignore
-                                raw_latent_shape=batch.
-                                raw_latent_shape[2:5],  # type: ignore
-                                patch_size=fastvideo_args.
-                                pipeline_config.  # type: ignore
-                                dit_config.patch_size,  # type: ignore
-                                STA_param=batch.STA_param,  # type: ignore
-                                VSA_sparsity=fastvideo_args.
-                                VSA_sparsity,  # type: ignore
-                                device=get_local_torch_device(),  # type: ignore
-                            )  # type: ignore
+                            attn_metadata = self.attn_metadata_builder.build(
+                                current_timestep=i,
+                                forward_batch=batch,
+                                fastvideo_args=fastvideo_args,
+                            )
                             assert attn_metadata is not None, "attn_metadata cannot be None"
                         else:
                             attn_metadata = None
@@ -812,20 +796,6 @@ class DmdDenoisingStage(DenoisingStage):
         if sp_group:
             latents = sequence_model_parallel_all_gather(latents, dim=1)
         latents = latents.permute(0, 2, 1, 3, 4)
-        
-        # Calculate total denoising time
-        # Synchronize CUDA before ending timer
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        denoise_end_time = time.time()
-        total_denoise_time = denoise_end_time - denoise_start_time
-        print(f"Total denoising time: {total_denoise_time:.2f} seconds")
-        
-        # Calculate total stage time
-        dmd_stage_end_time = time.time()
-        total_dmd_stage_time = dmd_stage_end_time - dmd_stage_start_time
-        print(f"Total DmdDenoisingStage time: {total_dmd_stage_time:.2f} seconds")
-        
         # Update batch with final latents
         batch.latents = latents
 

fastvideo/platforms/__init__.py

@@ -43,8 +43,6 @@ def cuda_platform_plugin() -> str | None:
 
         if cuda_is_jetson():
             is_cuda = True
-    if is_cuda:
-        logger.info("CUDA is available")
 
     return "fastvideo.platforms.cuda.CudaPlatform" if is_cuda else None
 
@@ -58,6 +56,8 @@ def mps_platform_plugin() -> str | None:
         if torch.backends.mps.is_available():
             is_mps = True
             logger.info("MPS (Metal Performance Shaders) is available")
+        else:
+            logger.info("MPS is not available")
     except Exception as e:
         logger.info("MPS detection failed: %s", e)
 
@@ -70,27 +70,8 @@ def cpu_platform_plugin() -> str | None:
     return "fastvideo.platforms.cpu.CpuPlatform"
 
 
-def rocm_platform_plugin() -> str | None:
-    is_rocm = False
-
-    try:
-        import amdsmi
-        amdsmi.amdsmi_init()
-        try:
-            if len(amdsmi.amdsmi_get_processor_handles()) > 0:
-                is_rocm = True
-                logger.info("ROCm platform is available")
-        finally:
-            amdsmi.amdsmi_shut_down()
-    except Exception as e:
-        logger.info("ROCm detection failed: %s", e)
-
-    return "fastvideo.platforms.rocm.RocmPlatform" if is_rocm else None
-
-
 builtin_platform_plugins = {
     'cuda': cuda_platform_plugin,
-    'rocm': rocm_platform_plugin,
     'mps': mps_platform_plugin,
     'cpu': cpu_platform_plugin,
 }
@@ -105,11 +86,6 @@ def resolve_current_platform_cls_qualname() -> str:
     if platform_cls_qualname is not None:
         return platform_cls_qualname
 
-    # Fall back to ROCm
-    platform_cls_qualname = rocm_platform_plugin()
-    if platform_cls_qualname is not None:
-        return platform_cls_qualname
-
     # Fall back to CUDA
     platform_cls_qualname = cuda_platform_plugin()
     if platform_cls_qualname is not None:

fastvideo/platforms/interface.py

@@ -91,7 +91,8 @@ class Platform:
 
     def is_cuda_alike(self) -> bool:
         """Stateless version of :func:`torch.cuda.is_available`."""
-        return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
+        # TODO(will): ROCM will be supported in the future here
+        return self._enum == PlatformEnum.CUDA
 
     def is_mps(self) -> bool:
         return self._enum == PlatformEnum.MPS

fastvideo/platforms/rocm.py

@@ -1,122 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from rocm/vllm: https://github.com/ROCm/vllm/blob/v0.7.3%2Brocm/vllm/platforms/rocm.py
-"""
-This file is a platform abstraction for ROCm GPUs,
-adjusted to match the structure and interface of `cuda.py`.
-"""
-
-import torch
-
-import fastvideo.envs as envs
-from fastvideo.logger import init_logger
-from fastvideo.platforms.interface import (AttentionBackendEnum,
-                                           DeviceCapability, Platform,
-                                           PlatformEnum)
-
-logger = init_logger(__name__)
-
-
-# ROCm uses the same torch.cuda interface
-class RocmPlatform(Platform):
-    _enum = PlatformEnum.ROCM
-    device_name: str = "rocm"
-    device_type: str = "cuda"  # torch uses 'cuda' backend string
-    dispatch_key: str = "CUDA"
-    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
-
-    @classmethod
-    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
-        major, minor = torch.cuda.get_device_capability(device_id)
-        return DeviceCapability(major=major, minor=minor)
-
-    @classmethod
-    def get_device_name(cls, device_id: int = 0) -> str:
-        return str(torch.cuda.get_device_name(device_id))
-
-    @classmethod
-    def get_device_total_memory(cls, device_id: int = 0) -> int:
-        return torch.cuda.get_device_properties(device_id).total_memory
-
-    @classmethod
-    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
-        if enforce_eager:
-            logger.warning(
-                "To see benefits of async output processing, enable CUDA graph. "
-                "Since enforce-eager is enabled, async output processor cannot be used"
-            )
-            return False
-        return True
-
-    @classmethod
-    def log_warnings(cls) -> None:
-        pass  # ROCm-specific warnings can be added here
-
-    @classmethod
-    def get_current_memory_usage(cls,
-                                 device: torch.device | None = None) -> float:
-        torch.cuda.reset_peak_memory_stats(device)
-        return float(torch.cuda.max_memory_allocated(device))
-
-    @classmethod
-    def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None,
-                             head_size: int, dtype: torch.dtype) -> str:
-        logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s",
-                    envs.FASTVIDEO_ATTENTION_BACKEND)
-
-        if selected_backend == AttentionBackendEnum.TORCH_SDPA:
-            logger.info("Using Torch SDPA backend.")
-            return "fastvideo.attention.backends.sdpa.SDPABackend"
-
-        elif selected_backend in (AttentionBackendEnum.FLASH_ATTN, None):
-            pass
-
-        elif selected_backend in (AttentionBackendEnum.SLIDING_TILE_ATTN,
-                                  AttentionBackendEnum.SAGE_ATTN):
-            raise ValueError(
-                f"{selected_backend.name} is not supported on {cls.device_name}."
-            )
-        elif selected_backend:
-            raise ValueError(
-                f"Invalid attention backend for {cls.device_name}: {selected_backend}"
-            )
-
-        target_backend = AttentionBackendEnum.FLASH_ATTN
-        if dtype not in (torch.float16, torch.bfloat16):
-            logger.info(
-                "Cannot use FlashAttention backend for dtype other than "
-                "torch.float16 or torch.bfloat16.")
-            target_backend = AttentionBackendEnum.TORCH_SDPA
-
-        if target_backend == AttentionBackendEnum.FLASH_ATTN:
-            try:
-                import flash_attn  # noqa: F401
-
-                from fastvideo.attention.backends.flash_attn import (  # noqa: F401
-                    FlashAttentionBackend)
-
-                supported_sizes = \
-                    FlashAttentionBackend.get_supported_head_sizes()
-                if head_size not in supported_sizes:
-                    logger.info(
-                        "Cannot use FlashAttention-2 backend for head size %d.",
-                        head_size)
-                    target_backend = AttentionBackendEnum.TORCH_SDPA
-            except ImportError:
-                logger.info("Cannot use FlashAttention backend because the "
-                            "flash_attn package is not found. "
-                            "Make sure that flash_attn was built and installed "
-                            "(on by default).")
-                target_backend = AttentionBackendEnum.TORCH_SDPA
-
-        if target_backend == AttentionBackendEnum.TORCH_SDPA:
-            logger.info("Using Torch SDPA backend.")
-
-            return "fastvideo.attention.backends.sdpa.SDPABackend"
-
-        logger.info("Using Flash Attention backend.")
-
-        return "fastvideo.attention.backends.flash_attn.FlashAttentionBackend"
-
-    @classmethod
-    def get_device_communicator_cls(cls) -> str:
-        return "fastvideo.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # works for ROCm too

fastvideo/tests/modal/pr_test.py

@@ -102,10 +102,7 @@ def run_precision_tests_STA():
 def run_precision_tests_VSA():
     run_test("python csrc/attn/tests/test_block_sparse.py")
 
+
 @app.function(gpu="L40S:1", image=image, timeout=3600)
 def run_inference_lora_tests():
-    run_test("pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs")
-
-@app.function(gpu="L40S:2", image=image, timeout=900)
-def run_distill_dmd_tests():
-    run_test("pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs")
+    run_test("pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs")

fastvideo/tests/training/VSA/test_training_loss_VSA.py

@@ -32,7 +32,7 @@ def run_worker():
         "--inference_mode", "False",
         "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         "--data_path", "data/mini_dataset_i2v_VSA/combined_parquet_dataset",
-        "--validation_dataset_file", "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json",
+        "--validation_dataset_file", "examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json",
         "--train_batch_size", "1",
         "--num_latent_t", "4",
         "--num_gpus", "2",

fastvideo/tests/training/Vanilla/test_training_loss.py

@@ -36,7 +36,7 @@ def run_worker():
         "--inference_mode", "False",
         "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset",
-        "--validation_dataset_file", "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json",
+        "--validation_dataset_file", "examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json",
         "--train_batch_size", "2",
         "--num_latent_t", "4",
         "--num_gpus", "2",

fastvideo/tests/training/distill/test_distill_dmd.py

@@ -1,130 +0,0 @@
-import os
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29512"
-import sys
-import subprocess
-from pathlib import Path
-import torch
-import json
-from huggingface_hub import snapshot_download
-from fastvideo.utils import logger
-# Import the training pipeline
-sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-from fastvideo.training.wan_training_pipeline import main
-from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.utils import FlexibleArgumentParser
-from fastvideo.training.wan_distillation_pipeline import WanDistillationPipeline
-
-wandb_name = "test_distill_dmd"
-
-NUM_NODES = "1"
-NUM_GPUS_PER_NODE = "2"
-
-
-def run_worker():
-    """Worker function that will be run on each GPU"""
-    # Create and populate args
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    
-    # Set the arguments as they are in finetune_v1_test.sh
-    args = parser.parse_args([
-        "--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        "--inference_mode", "False",
-        "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset",
-        "--validation_dataset_file", "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json",
-        "--train_batch_size", "1",
-        "--num_latent_t", "4",
-        "--num_gpus", "2",
-        "--sp_size", "2",
-        "--tp_size", "1",
-        "--hsdp_replicate_dim", "1",
-        "--hsdp_shard_dim", "2",
-        "--train_sp_batch_size", "1",
-        "--dataloader_num_workers", "1",
-        "--gradient_accumulation_steps", "2",
-        "--max_train_steps", "2",
-        "--learning_rate", "1e-5",
-        "--mixed_precision", "bf16",
-        "--training_state_checkpointing_steps", "30",
-        "--weight_only_checkpointing_steps", "30",
-        "--validation_steps", "10",
-        "--validation_sampling_steps", "3",
-        "--log_validation",
-        "--checkpoints_total_limit", "3",
-        "--allow_tf32",
-        "--ema_start_step", "0",
-        "--training_cfg_rate", "0.0",
-        "--output_dir", "data/wan_finetune_test",
-        "--tracker_project_name", "wan_finetune_ci",
-        "--wandb_run_name", wandb_name,
-        "--num_height", "480",
-        "--num_width", "832",
-        "--num_frames", "13",
-        "--flow_shift", "8",
-        "--validation_guidance_scale", "1.0",
-        "--weight_decay", "0.01",
-        "--dit_precision", "fp32",
-        "--max_grad_norm", "1.0",
-        "--generator_update_interval", "5",
-        "--dmd_denoising_steps", "1000,757,522",
-        "--min_timestep_ratio", "0.02",
-        "--max_timestep_ratio", "0.98",
-        "--real_score_guidance_scale", "3.5",
-        "--enable_gradient_checkpointing_type", "full"
-    ])
-    # Call the main training function
-    pipeline = WanDistillationPipeline.from_pretrained(
-        args.pretrained_model_name_or_path, args=args)
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Training pipeline done")
-
-def test_distributed_training():
-    """Test the distributed training setup"""
-    os.environ["WANDB_MODE"] = "offline"
-
-    data_dir = Path("data/crush-smol_processed_t2v")
-    
-    if not data_dir.exists():
-        print(f"Downloading test dataset to {data_dir}...")
-        snapshot_download(
-            repo_id="wlsaidhi/crush-smol_processed_t2v",
-            local_dir=str(data_dir),
-            repo_type="dataset",
-            local_dir_use_symlinks=False
-        )
-    
-    # Get the current file path
-    current_file = Path(__file__).resolve()
-    
-    # Run torchrun command
-    cmd = [
-        "torchrun",
-        "--nnodes", NUM_NODES,
-        "--nproc_per_node", NUM_GPUS_PER_NODE,
-        "--master_port", os.environ["MASTER_PORT"],
-        str(current_file)
-    ]
-    process = subprocess.run(cmd, capture_output=True, text=True)
-    
-    # Print stdout and stderr for debugging
-    if process.stdout:
-        print("STDOUT:", process.stdout)
-    if process.stderr:
-        print("STDERR:", process.stderr)
-    
-    # Check if the process failed
-    if process.returncode != 0:
-        print(f"Process failed with return code: {process.returncode}")
-        raise subprocess.CalledProcessError(process.returncode, cmd, process.stdout, process.stderr)
-
-if __name__ == "__main__":
-    if os.environ.get("LOCAL_RANK") is not None:
-        # We're being run by torchrun
-        run_worker()
-    else:
-        # We're being run directly
-        test_distributed_training()

fastvideo/tests/training/lora/test_lora_training.py

@@ -41,7 +41,7 @@ def test_lora_training():
         "--inference_mode", "False",
         "--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         "--data_path", "data/crush-smol_processed_t2v/combined_parquet_dataset",
-        "--validation_dataset_file", "examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json",
+        "--validation_dataset_file", "examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json",
         "--train_batch_size", "1",
         "--num_latent_t", "8",
         "--num_gpus", NUM_GPUS_PER_NODE,

fastvideo/tests/vaes/test_official_wan_vae.py

@@ -1,113 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-
-import numpy as np
-import pytest
-import torch
-from diffusers import AutoencoderKLWan
-
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import VAELoader
-from fastvideo.configs.models.vaes import WanVAEConfig
-from fastvideo.utils import maybe_download_model
-
-try:
-    from wan.modules.vae2_2 import Wan2_2_VAE
-except ImportError:
-    Wan2_2_VAE = None
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH)
-VAE_PATH = os.path.join(MODEL_PATH, "vae")
-
-
-@pytest.mark.skip(reason="disable test")
-@pytest.mark.usefixtures("distributed_setup")
-def test_official_wan_vae():
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str))
-    args.device = device
-    args.vae_cpu_offload = False
-
-    loader = VAELoader()
-    model2 = loader.load(VAE_PATH, args)
-    assert model2.use_feature_cache # Default to use the original WanVAE algorithm
-
-    model1 = Wan2_2_VAE(vae_pth="/mnt/weka/home/hao.zhang/wei/Wan2.2_VAE.pth", device=device, dtype=precision)
-
-    # Create identical inputs for both models
-    batch_size = 1
-
-    # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size,
-                               3,
-                               81,
-                               32,
-                               32,
-                               device=device,
-                               dtype=precision)
-    # latent_tensor = torch.randn(batch_size,
-    #                             16,
-    #                             21,
-    #                             32,
-    #                             32,
-    #                             device=device,
-    #                             dtype=precision)
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        # Test encoding
-        logger.info("Testing encoding...")
-        latent1 = model1.encode([input_tensor.squeeze(0)])[0].unsqueeze(0)
-        print("--------------------------------")
-        latent2 = model2.encode(input_tensor)
-        # Check if latents have the same shape
-        assert latent1.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
-        # Check if latents are similar
-        max_diff_encode = torch.max(torch.abs(latent1 - latent2.mean))
-        mean_diff_encode = torch.mean(torch.abs(latent1 - latent2.mean))
-        logger.info("Maximum difference between encoded latents: %s",
-                    max_diff_encode.item())
-        logger.info("Mean difference between encoded latents: %s",
-                    mean_diff_encode.item())
-        assert max_diff_encode < 1e-5, f"Encoded latents differ significantly: max diff = {mean_diff_encode.item()}"
-        # Test decoding
-        logger.info("Testing decoding...")
-        # latent1_tensor = latent1.mode()
-        # mean1 = (torch.tensor(model1.config.latents_mean).view(
-        #     1, model1.config.z_dim, 1, 1, 1).to(input_tensor.device,
-        #                                         input_tensor.dtype))
-        # std1 = (1.0 / torch.tensor(model1.config.latents_std).view(
-        #     1, model1.config.z_dim, 1, 1, 1)).to(input_tensor.device,
-        #                                         input_tensor.dtype)
-        # latent1_tensor = latent1_tensor / std1 + mean1
-        # output1 = model1.decode(latent1_tensor).sample
-
-        output1 = model1.decode([latent1.squeeze(0)])
-        output1 = output1[0].unsqueeze(0)
-
-        mean2 = model2.config.arch_config.shift_factor.to(input_tensor.device, input_tensor.dtype)
-        std2 = model2.config.arch_config.scaling_factor.to(input_tensor.device, input_tensor.dtype)
-        latent2_tensor = latent2.mode()
-        latent2_tensor = latent2_tensor / std2 + mean2
-        output2 = model2.decode(latent2_tensor)
-        # Check if outputs have the same shape
-        assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
-
-        # Check if outputs are similar
-        max_diff_decode = torch.max(torch.abs(output1.float() - output2.float()))
-        mean_diff_decode = torch.mean(torch.abs(output1.float() - output2.float()))
-        logger.info("Maximum difference between decoded outputs: %s",
-                    max_diff_decode.item())
-        logger.info("Mean difference between decoded outputs: %s",
-                    mean_diff_decode.item())
-        assert max_diff_decode < 1e-5, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}"

fastvideo/tests/vaes/test_wan_vae.py

@@ -104,4 +104,4 @@ def test_wan_vae():
                     max_diff_decode.item())
         logger.info("Mean difference between decoded outputs: %s",
                     mean_diff_decode.item())
-        assert max_diff_decode < 1e-5, f"Decoded outputs differ significantly: max diff = {mean_diff_decode.item()}"
+        assert max_diff_decode < 1e-5, f"Decoded outputs differ significantly: max diff = {mean_diff_decode.item()}"

fastvideo/training/distillation_pipeline.py

@@ -36,9 +36,8 @@ from fastvideo.training.activation_checkpoint import (
     apply_activation_checkpointing)
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
-    clip_grad_norm_while_handling_failing_dtensor_cases,
-    load_distillation_checkpoint, pred_noise_to_pred_video,
-    save_distillation_checkpoint, shift_timestep)
+    clip_grad_norm_while_handling_failing_dtensor_cases, load_checkpoint,
+    pred_noise_to_pred_video, save_checkpoint, shift_timestep)
 from fastvideo.utils import is_vsa_available, set_random_seed
 
 import wandb  # isort: skip
@@ -61,8 +60,6 @@ class DistillationPipeline(TrainingPipeline):
     train_dataloader: StatefulDataLoader
     train_loader_iter: Iterator[dict[str, Any]]
     current_epoch: int = 0
-    init_steps: int
-    current_trainstep: int
     video_latent_shape: tuple[int, ...]
     video_latent_shape_sp: tuple[int, ...]
 
@@ -391,9 +388,10 @@ class DistillationPipeline(TrainingPipeline):
 
         training_batch.conditional_dict = conditional_dict
         training_batch.unconditional_dict = unconditional_dict
-        training_batch.raw_latent_shape = training_batch.latents.shape
+
         training_batch.latents = training_batch.latents.permute(0, 2, 1, 3, 4)
         self.video_latent_shape = training_batch.latents.shape
+        training_batch.raw_latent_shape = training_batch.latents.shape
 
         if self.sp_world_size > 1:
             training_batch.latents = rearrange(
@@ -484,17 +482,19 @@ class DistillationPipeline(TrainingPipeline):
         training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss
         return training_batch
 
-    def _resume_from_checkpoint(self) -> None:
+    def _resume_from_checkpoint(self) -> None:  #TODO(yongqi)
         """Resume training from checkpoint with distillation models."""
 
         logger.info("Loading distillation checkpoint from %s",
                     self.training_args.resume_from_checkpoint)
 
-        resumed_step = load_distillation_checkpoint(
-            self.transformer, self.fake_score_transformer, self.global_rank,
+        resumed_step = load_checkpoint(
+            self.transformer, self.global_rank,
             self.training_args.resume_from_checkpoint, self.optimizer,
-            self.fake_score_optimizer, self.train_dataloader, self.lr_scheduler,
-            self.fake_score_lr_scheduler, self.noise_random_generator)
+            self.train_dataloader, self.lr_scheduler,
+            self.noise_random_generator)
+
+        # TODO: Add checkpoint loading for critic and teacher models
 
         if resumed_step > 0:
             self.init_steps = resumed_step
@@ -680,29 +680,24 @@ class DistillationPipeline(TrainingPipeline):
         else:
             set_random_seed(seed + self.global_rank)
 
-        # Set random seeds for deterministic training
-        self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
-            self.seed)
+        self.noise_random_generator = torch.Generator(
+            device="cpu").manual_seed(seed)
         self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed(
             self.seed)
         self.validation_random_generator = torch.Generator(
-            device="cpu").manual_seed(self.seed)
+            device="cpu").manual_seed(seed)
+
         logger.info("Initialized random seeds with seed: %s", seed)
 
-        # Resume from checkpoint if specified (this will restore random states)
         if self.training_args.resume_from_checkpoint:
             self._resume_from_checkpoint()
-            logger.info("Resumed from checkpoint, random states restored")
-        else:
-            logger.info("Starting training from scratch")
 
         self.train_loader_iter = iter(self.train_dataloader)
 
         step_times: deque[float] = deque(maxlen=100)
 
         self._log_training_info()
-        self._log_validation(self.transformer, self.training_args,
-                             self.init_steps)
+        self._log_validation(self.transformer, self.training_args, 0)
 
         progress_bar = tqdm(
             range(0, self.training_args.max_train_steps),
@@ -770,51 +765,30 @@ class DistillationPipeline(TrainingPipeline):
                     log_data["VSA_train_sparsity"] = current_vsa_sparsity
                 wandb.log(log_data, step=step)
 
-            # Save training state checkpoint (for resuming training)
-            if (self.training_args.training_state_checkpointing_steps > 0
-                    and step %
-                    self.training_args.training_state_checkpointing_steps == 0):
-                print("rank", self.global_rank,
-                      "save training state checkpoint at step", step)
-                save_distillation_checkpoint(
-                    self.transformer, self.fake_score_transformer,
-                    self.global_rank, self.training_args.output_dir, step,
-                    self.optimizer, self.fake_score_optimizer,
-                    self.train_dataloader, self.lr_scheduler,
-                    self.fake_score_lr_scheduler, self.noise_random_generator)
-
+            if step % self.training_args.checkpointing_steps == 0 and step > 0:
+                print("rank", self.global_rank, "save checkpoint at step", step)
+                save_checkpoint(
+                    self.transformer,
+                    self.global_rank,  #TODO(yongqi)
+                    self.training_args.output_dir,
+                    step,
+                    self.optimizer,
+                    self.train_dataloader,
+                    self.lr_scheduler,
+                    self.noise_random_generator)
                 if self.transformer:
                     self.transformer.train()
                 self.sp_group.barrier()
 
-            # Save weight-only checkpoint
-            if (self.training_args.weight_only_checkpointing_steps > 0
-                    and step %
-                    self.training_args.weight_only_checkpointing_steps == 0):
-                print("rank", self.global_rank,
-                      "save weight-only checkpoint at step", step)
-                save_distillation_checkpoint(self.transformer,
-                                             self.fake_score_transformer,
-                                             self.global_rank,
-                                             self.training_args.output_dir,
-                                             f"{step}_weight_only",
-                                             only_save_generator_weight=True)
-
             if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
                 self._log_validation(self.transformer, self.training_args, step)
 
         wandb.finish()
-
-        # Save final training state checkpoint
-        print("rank", self.global_rank,
-              "save final training state checkpoint at step",
-              self.training_args.max_train_steps)
-        save_distillation_checkpoint(
-            self.transformer, self.fake_score_transformer, self.global_rank,
-            self.training_args.output_dir, self.training_args.max_train_steps,
-            self.optimizer, self.fake_score_optimizer, self.train_dataloader,
-            self.lr_scheduler, self.fake_score_lr_scheduler,
-            self.noise_random_generator)
+        save_checkpoint(self.transformer, self.global_rank,
+                        self.training_args.output_dir,
+                        self.training_args.max_train_steps, self.optimizer,
+                        self.train_dataloader, self.lr_scheduler,
+                        self.noise_random_generator)
 
         if get_sp_group():
             cleanup_dist_env_and_memory()

fastvideo/training/training_pipeline.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-import dataclasses
 import math
 import os
 import time
@@ -22,7 +21,7 @@ from tqdm.auto import tqdm
 
 import fastvideo.envs as envs
 from fastvideo.attention.backends.video_sparse_attn import (
-    VideoSparseAttentionMetadataBuilder)
+    VideoSparseAttentionMetadata)
 from fastvideo.configs.sample import SamplingParam
 from fastvideo.dataset import build_parquet_map_style_dataloader
 from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
@@ -163,9 +162,8 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         if self.global_rank == 0:
             project = training_args.tracker_project_name or "fastvideo"
-            wandb_config = dataclasses.asdict(training_args)
             wandb.init(project=project,
-                       config=wandb_config,
+                       config=training_args,
                        name=training_args.wandb_run_name)
 
     @abstractmethod
@@ -209,9 +207,8 @@ class TrainingPipeline(LoRAPipeline, ABC):
     def _normalize_dit_input(self,
                              training_batch: TrainingBatch) -> TrainingBatch:
         # TODO(will): support other models
-        training_batch.latents = normalize_dit_input(
-            'wan', training_batch.latents,
-            self.validation_pipeline.get_module("vae"))
+        training_batch.latents = normalize_dit_input('wan',
+                                                     training_batch.latents)
         return training_batch
 
     def _prepare_dit_inputs(self,
@@ -263,13 +260,15 @@ class TrainingPipeline(LoRAPipeline, ABC):
         assert latents_shape is not None
         assert training_batch.timesteps is not None
         if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
-            training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder(  # type: ignore
-            ).build(  # type: ignore
-                raw_latent_shape=latents_shape[2:5],
+            dit_seq_shape = [
+                latents_shape[2] // patch_size[0],
+                latents_shape[3] // patch_size[1],
+                latents_shape[4] // patch_size[2]
+            ]
+            training_batch.attn_metadata = VideoSparseAttentionMetadata(
                 current_timestep=training_batch.timesteps,
-                patch_size=patch_size,
-                VSA_sparsity=current_vsa_sparsity,
-                device=get_local_torch_device())
+                dit_seq_shape=dit_seq_shape,
+                VSA_sparsity=current_vsa_sparsity)
         else:
             training_batch.attn_metadata = None
 
@@ -409,7 +408,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
     def train(self) -> None:
         assert self.seed is not None, "seed must be set"
-        set_random_seed(self.seed + self.global_rank)
+        set_random_seed(self.seed)
         logger.info('rank: %s: start training',
                     self.global_rank,
                     local_main_process_only=False)
@@ -438,8 +437,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         step_times: deque[float] = deque(maxlen=100)
 
         self._log_training_info()
-        self._log_validation(self.transformer, self.training_args,
-                             self.init_steps)
+        self._log_validation(self.transformer, self.training_args, 0)
 
         # Train!
         progress_bar = tqdm(

fastvideo/training/training_utils.py

@@ -33,8 +33,7 @@ def gather_state_dict_on_cpu_rank0(
     cpu_state_dict = {}
     sharded_sd = model.state_dict()
     param_requires_grad = set([
-        k.replace("._checkpoint_wrapped_module.", ".")
-        for k, v in dict(model.named_parameters()).items() if v.requires_grad
+        k for k, v in dict(model.named_parameters()).items() if v.requires_grad
     ])
     for param_name, param in sharded_sd.items():
         if param_name not in param_requires_grad:
@@ -58,7 +57,6 @@ def gather_state_dict_on_cpu_rank0(
 
         if rank == 0:
             cpu_state_dict[param_name] = param.cpu()
-
     return cpu_state_dict
 
 
@@ -189,148 +187,6 @@ def save_checkpoint(transformer,
         logger.info("--> checkpoint saved at step %s to %s", step, weight_path)
 
 
-def save_distillation_checkpoint(generator_transformer,
-                                 fake_score_transformer,
-                                 rank,
-                                 output_dir,
-                                 step,
-                                 generator_optimizer=None,
-                                 fake_score_optimizer=None,
-                                 dataloader=None,
-                                 generator_scheduler=None,
-                                 fake_score_scheduler=None,
-                                 noise_generator=None,
-                                 only_save_generator_weight=False) -> None:
-    """
-    Save distillation checkpoint with both generator and fake_score models.
-    Saves both distributed checkpoint and consolidated model weights.
-    Only saves the generator model for inference (consolidated weights).
-    
-    Args:
-        only_save_generator_weight: If True, only save the generator model weights for inference
-                                   without saving distributed checkpoint for training resume.
-    """
-    save_dir = os.path.join(output_dir, f"checkpoint-{step}")
-    os.makedirs(save_dir, exist_ok=True)
-
-    # Create directories for models
-    inference_save_dir = os.path.join(save_dir,
-                                      "generator_inference_transformer")
-
-    # Only save distributed checkpoint if not only saving generator weight
-    if not only_save_generator_weight:
-        # Save generator distributed checkpoint
-        generator_states = {
-            "model": ModelWrapper(generator_transformer),
-        }
-        if generator_optimizer is not None:
-            generator_states["optimizer"] = OptimizerWrapper(
-                generator_transformer, generator_optimizer)
-        if dataloader is not None:
-            generator_states["dataloader"] = dataloader
-        if generator_scheduler is not None:
-            generator_states["scheduler"] = SchedulerWrapper(
-                generator_scheduler)
-
-        generator_dcp_dir = os.path.join(save_dir, "distributed_checkpoint",
-                                         "generator")
-        logger.info("rank: %s, saving generator distributed checkpoint to %s",
-                    rank,
-                    generator_dcp_dir,
-                    local_main_process_only=False)
-
-        begin_time = time.perf_counter()
-        dcp.save(generator_states, checkpoint_id=generator_dcp_dir)
-        end_time = time.perf_counter()
-
-        logger.info(
-            "rank: %s, generator distributed checkpoint saved in %.2f seconds",
-            rank,
-            end_time - begin_time,
-            local_main_process_only=False)
-
-        # Save critic distributed checkpoint
-        critic_states = {
-            "model": ModelWrapper(fake_score_transformer),
-        }
-        if fake_score_optimizer is not None:
-            critic_states["optimizer"] = OptimizerWrapper(
-                fake_score_transformer, fake_score_optimizer)
-        if dataloader is not None:
-            critic_states["dataloader"] = dataloader
-        if fake_score_scheduler is not None:
-            critic_states["scheduler"] = SchedulerWrapper(fake_score_scheduler)
-
-        critic_dcp_dir = os.path.join(save_dir, "distributed_checkpoint",
-                                      "critic")
-        logger.info("rank: %s, saving critic distributed checkpoint to %s",
-                    rank,
-                    critic_dcp_dir,
-                    local_main_process_only=False)
-
-        begin_time = time.perf_counter()
-        dcp.save(critic_states, checkpoint_id=critic_dcp_dir)
-        end_time = time.perf_counter()
-
-        logger.info(
-            "rank: %s, critic distributed checkpoint saved in %.2f seconds",
-            rank,
-            end_time - begin_time,
-            local_main_process_only=False)
-
-        # Save shared random state separately
-        shared_states = {
-            "random_state": RandomStateWrapper(noise_generator),
-        }
-        shared_dcp_dir = os.path.join(save_dir, "distributed_checkpoint",
-                                      "shared")
-
-        dcp.save(shared_states, checkpoint_id=shared_dcp_dir)
-
-    else:
-        logger.info(
-            "rank: %s, skipping distributed checkpoint save (only_save_generator_weight=True)",
-            rank,
-            local_main_process_only=False)
-
-    # Save generator model weights (consolidated) for inference
-    cpu_state = gather_state_dict_on_cpu_rank0(generator_transformer,
-                                               device=None)
-
-    if rank == 0:
-        # Save generator model weights (consolidated) for inference
-        os.makedirs(inference_save_dir, exist_ok=True)
-        weight_path = os.path.join(inference_save_dir,
-                                   "diffusion_pytorch_model.safetensors")
-        logger.info(
-            "rank: %s, saving consolidated generator inference checkpoint to %s",
-            rank,
-            weight_path,
-            local_main_process_only=False)
-
-        # Convert training format to diffusers format and save
-        diffusers_state_dict = custom_to_hf_state_dict(
-            cpu_state, generator_transformer.reverse_param_names_mapping)
-        save_file(diffusers_state_dict, weight_path)
-
-        logger.info(
-            "rank: %s, consolidated generator inference checkpoint saved to %s",
-            rank,
-            weight_path,
-            local_main_process_only=False)
-
-        # Save model config
-        config_dict = generator_transformer.hf_config
-        if "dtype" in config_dict:
-            del config_dict["dtype"]  # TODO
-        config_path = os.path.join(inference_save_dir, "config.json")
-        # save dict as json
-        with open(config_path, "w") as f:
-            json.dump(config_dict, f, indent=4)
-        logger.info("--> distillation checkpoint saved at step %s to %s", step,
-                    weight_path)
-
-
 def load_checkpoint(transformer,
                     rank,
                     checkpoint_path,
@@ -391,137 +247,14 @@ def load_checkpoint(transformer,
     return step
 
 
-def load_distillation_checkpoint(generator_transformer,
-                                 fake_score_transformer,
-                                 rank,
-                                 checkpoint_path,
-                                 generator_optimizer=None,
-                                 fake_score_optimizer=None,
-                                 dataloader=None,
-                                 generator_scheduler=None,
-                                 fake_score_scheduler=None,
-                                 noise_generator=None) -> int:
-    """
-    Load distillation checkpoint with both generator and fake_score models.
-    Returns the step number from which training should resume.
-    """
-    if not os.path.exists(checkpoint_path):
-        logger.warning("Distillation checkpoint path %s does not exist",
-                       checkpoint_path)
-        return 0
-
-    # Extract step number from checkpoint path
-    step = int(os.path.basename(checkpoint_path).split('-')[-1])
-
-    if rank == 0:
-        logger.info("Loading distillation checkpoint from step %s", step)
-
-    # Load generator distributed checkpoint
-    generator_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
-                                     "generator")
-    if not os.path.exists(generator_dcp_dir):
-        logger.warning(
-            "Generator distributed checkpoint directory %s does not exist",
-            generator_dcp_dir)
-        return 0
-
-    generator_states = {
-        "model": ModelWrapper(generator_transformer),
-    }
-
-    if generator_optimizer is not None:
-        generator_states["optimizer"] = OptimizerWrapper(
-            generator_transformer, generator_optimizer)
-
-    if dataloader is not None:
-        generator_states["dataloader"] = dataloader
-
-    if generator_scheduler is not None:
-        generator_states["scheduler"] = SchedulerWrapper(generator_scheduler)
-
-    logger.info("rank: %s, loading generator distributed checkpoint from %s",
-                rank,
-                generator_dcp_dir,
-                local_main_process_only=False)
-
-    begin_time = time.perf_counter()
-    dcp.load(generator_states, checkpoint_id=generator_dcp_dir)
-    end_time = time.perf_counter()
-
-    logger.info(
-        "rank: %s, generator distributed checkpoint loaded in %.2f seconds",
-        rank,
-        end_time - begin_time,
-        local_main_process_only=False)
-
-    # Load critic distributed checkpoint
-    critic_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
-                                  "critic")
-    if not os.path.exists(critic_dcp_dir):
-        logger.warning(
-            "Critic distributed checkpoint directory %s does not exist",
-            critic_dcp_dir)
-        return 0
-
-    critic_states = {
-        "model": ModelWrapper(fake_score_transformer),
-    }
-
-    if fake_score_optimizer is not None:
-        critic_states["optimizer"] = OptimizerWrapper(fake_score_transformer,
-                                                      fake_score_optimizer)
-
-    if dataloader is not None:
-        critic_states["dataloader"] = dataloader
-
-    if fake_score_scheduler is not None:
-        critic_states["scheduler"] = SchedulerWrapper(fake_score_scheduler)
-
-    logger.info("rank: %s, loading critic distributed checkpoint from %s",
-                rank,
-                critic_dcp_dir,
-                local_main_process_only=False)
-
-    begin_time = time.perf_counter()
-    dcp.load(critic_states, checkpoint_id=critic_dcp_dir)
-    end_time = time.perf_counter()
-
-    logger.info(
-        "rank: %s, critic distributed checkpoint loaded in %.2f seconds",
-        rank,
-        end_time - begin_time,
-        local_main_process_only=False)
-
-    # Load shared random state
-    shared_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
-                                  "shared")
-    if not os.path.exists(shared_dcp_dir):
-        logger.warning("Shared random state directory %s does not exist",
-                       shared_dcp_dir)
-        return 0
-
-    shared_states = {
-        "random_state": RandomStateWrapper(noise_generator),
-    }
-
-    begin_time = time.perf_counter()
-    dcp.load(shared_states, checkpoint_id=shared_dcp_dir)
-    end_time = time.perf_counter()
-
-    logger.info("rank: %s, shared random state loaded in %.2f seconds",
-                rank,
-                end_time - begin_time,
-                local_main_process_only=False)
-    logger.info("--> distillation checkpoint loaded from step %s", step)
-    return step
-
-
-def normalize_dit_input(model_type, latents, vae) -> torch.Tensor:
+def normalize_dit_input(model_type, latents, args=None) -> torch.Tensor:
     if model_type == "hunyuan_hf" or model_type == "hunyuan":
         return latents * 0.476986
     elif model_type == "wan":
-        latents_mean = torch.tensor(vae.latents_mean)
-        latents_std = 1.0 / torch.tensor(vae.latents_std)
+        from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
+        vae_config = WanVAEConfig()
+        latents_mean = torch.tensor(vae_config.arch_config.latents_mean)
+        latents_std = 1.0 / torch.tensor(vae_config.arch_config.latents_std)
 
         latents_mean = latents_mean.view(1, -1, 1, 1,
                                          1).to(device=latents.device)