load_video

.github/workflows/pr-test.yml

@@ -14,9 +14,13 @@ on:
       - ".github/workflows/pr-test.yml"
       - "pyproject.toml"
       - "docker/Dockerfile.python3.12"
-      - "csrc/**"
   workflow_dispatch:
     inputs:
+      custom_image:
+        description: "Custom image from this repository (default: fastvideo-dev:py3.12-latest)"
+        required: false
+        default: "fastvideo-dev:py3.12-latest"
+        type: string
       run_encoder_test:
         description: "Run encoder-test"
         required: false
@@ -52,16 +56,6 @@ on:
         required: false
         default: false
         type: boolean
-      run_precision_test_STA:
-        description: "Run precision-test-STA"
-        required: false
-        default: false
-        type: boolean
-      run_precision_test_VSA:
-        description: "Run precision-test-VSA"
-        required: false
-        default: false
-        type: boolean
       run_nightly_test:
         description: "Run nightly-test"
         required: false
@@ -71,7 +65,6 @@ on:
 env:
   PYTHONUNBUFFERED: "1"
 
-
 concurrency:
   group: pr-test-${{ github.ref }}
   cancel-in-progress: true
@@ -91,69 +84,44 @@ jobs:
       training-test: ${{ steps.filter.outputs.training-test }}
       training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }}
       inference-test-STA: ${{ steps.filter.outputs.inference-test-STA }}
-      precision-test-STA: ${{ steps.filter.outputs.precision-test-STA }}
-      precision-test-VSA: ${{ steps.filter.outputs.precision-test-VSA }}
     steps:
       - uses: actions/checkout@v4
       - uses: dorny/paths-filter@v3
         id: filter
         with:
           filters: |
-            # Define reusable path patterns
-            common-paths: &common-paths
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
-            sta-kernel-paths: &sta-kernel-paths
-              - 'csrc/attn/st_attn/**'
-              - 'csrc/attn/setup_sta.py'
-              - 'csrc/attn/config_sta.py'
-              - 'csrc/attn/st_attn.cpp'
-            vsa-kernel-paths: &vsa-kernel-paths
-              - 'csrc/attn/vsa/**'
-              - 'csrc/attn/tk/**'
-              - 'csrc/attn/setup_vsa.py'
-              - 'csrc/attn/config_vsa.py'
-              - 'csrc/attn/vsa.cpp'
-            vsa-paths: &vsa-paths
-              - 'fastvideo/v1/**'
-              - *common-paths
-              - *vsa-kernel-paths
-            
-            # Actual tests
             encoder-test:
               - 'fastvideo/v1/models/encoders/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/encoders/**'
-              - *common-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
             vae-test:
               - 'fastvideo/v1/models/vaes/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/vaes/**'
-              - *common-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
             transformer-test:
               - 'fastvideo/v1/models/dits/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/transformers/**'
               - 'fastvideo/v1/layers/**'
               - 'fastvideo/v1/attention/**'
-              - *common-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
             training-test:
               - 'fastvideo/v1/**'
-              - *common-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
             training-test-VSA:
               - 'fastvideo/v1/**'
-              - *common-paths
-              - *vsa-kernel-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
             inference-test-STA:
               - 'fastvideo/v1/**'
-              - *common-paths
-              - *sta-kernel-paths
-            precision-test-STA:
-              - *common-paths
-              - *sta-kernel-paths
-            precision-test-VSA:
-              - *common-paths
-              - *vsa-kernel-paths
+              - 'pyproject.toml'
+              - 'docker/Dockerfile.python3.12'
 
   encoder-test:
     needs: change-filter
@@ -166,7 +134,7 @@ jobs:
       gpu_type: "NVIDIA A40"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
       test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
       timeout_minutes: 30
     secrets:
@@ -184,7 +152,7 @@ jobs:
       gpu_type: "NVIDIA A40"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
       test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
       timeout_minutes: 30
     secrets:
@@ -202,7 +170,7 @@ jobs:
       gpu_type: "NVIDIA L40S"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
       test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
       timeout_minutes: 30
     secrets:
@@ -212,7 +180,8 @@ jobs:
   ssim-test:
     needs: change-filter
     if: >-
-      github.event_name != 'workflow_dispatch' || (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
     strategy:
       fail-fast: false
       matrix:
@@ -238,7 +207,7 @@ jobs:
   training-test:
     needs: change-filter
     if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test == 'true') || 
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test == 'true')
     uses: ./.github/workflows/runpod-test.yml
     with:
@@ -247,7 +216,7 @@ jobs:
       gpu_count: 4
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
       test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/Vanilla -srP"
       timeout_minutes: 30
     secrets:
@@ -258,7 +227,7 @@ jobs:
   training-test-VSA:
     needs: change-filter
     if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test-VSA == 'true') || 
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test_VSA == 'true')
     uses: ./.github/workflows/runpod-test.yml
     with:
@@ -267,7 +236,7 @@ jobs:
       gpu_count: 1
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
       test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/VSA -srP"
       timeout_minutes: 30
     secrets:
@@ -278,7 +247,7 @@ jobs:
   inference-test-STA:
     needs: change-filter
     if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.inference-test-STA == 'true') || 
+      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_inference_test_STA == 'true')
     uses: ./.github/workflows/runpod-test.yml
     with:
@@ -287,51 +256,13 @@ jobs:
       gpu_count: 1
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
       test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/inference/STA -srP"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
 
-  precision-test-STA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-STA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_STA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "precision-test-STA"
-      gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_sta.py"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  precision-test-VSA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-VSA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_VSA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "precision-test-VSA"
-      gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_block_sparse.py"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
   nightly-test:
     if: >-
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
@@ -342,7 +273,7 @@ jobs:
       gpu_count: 4
       volume_size: 100
       disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
       test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs"
       timeout_minutes: 30
     secrets:
@@ -351,8 +282,7 @@ jobs:
       WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
 
   runpod-cleanup:
-    # Add other jobs to this list as you create them
-    needs: [encoder-test, vae-test, transformer-test, ssim-test, training-test, training-test-VSA, inference-test-STA, precision-test-STA, precision-test-VSA] 
+    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
     if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
     runs-on: ubuntu-latest
     steps:
@@ -369,7 +299,7 @@ jobs:
 
       - name: Cleanup all RunPod instances
         env:
-          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12", "training-test", "training-test-VSA", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]'
+          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12"]'
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
         run: python .github/scripts/runpod_cleanup.py

csrc/attn/README.md

@@ -4,7 +4,7 @@
 
 
 ## Installation
-We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only support H100/H200, because ThunderKittens uses TMA but doesn't support Blackwell yet.
+We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
 First, install C++20 for ThunderKittens:
 ```bash
 sudo apt update
@@ -53,14 +53,8 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)
 
 ## Test
 ```bash
-python tests/test_sta.py # test STA
-python tests/test_block_sparse.py # test VSA
+python test/test_sta.py
 ```
-## Benchmark
-```bash
-python benchmarks/bench_sta.py
-```
-
 
 ## How Does STA Work?
 We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 

csrc/attn/bench/bench_sta.py

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

csrc/attn/st_attn/st_attn_h100.cu

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

csrc/attn/test/test_block_sparse.py

@@ -7,7 +7,6 @@ from vsa import BLOCK_M, BLOCK_N
 
 import numpy as np
 import random
-import gc
 
 def set_seed(seed: int = 42):
     # Python random module
@@ -21,6 +20,15 @@ def set_seed(seed: int = 42):
     torch.cuda.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)  # if using multi-GPU
 
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
+    parser.add_argument('--batch_size', type=int, default=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=100, help='Number of test iterations to run')
+    return parser.parse_args()
 
 @torch.no_grad
 def precision_metric(quant_o, fa2_o): 
@@ -127,7 +135,9 @@ def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device=
                 
     return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
 
-def main(args):
+def main():
+    args = parse_arguments()
+
     set_seed(42)
     
     # Extract parameters
@@ -181,36 +191,23 @@ def main(args):
             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_sdpa = q.clone()
+            k_sdpa = k.clone()
+            v_sdpa = v.clone()
 
             q.requires_grad = True
             k.requires_grad = True
             v.requires_grad = True
-
-
-            # testing forward
-            o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
-            del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded
-            grad_o = torch.randn_like(o)
-            o.backward(grad_o)
-            # clear memory
-            q_sdpa = q.detach().clone()
-            k_sdpa = k.detach().clone()
-            v_sdpa = v.detach().clone()
             q_sdpa.requires_grad = True
             k_sdpa.requires_grad = True
             v_sdpa.requires_grad = True
-            q.data = torch.empty(0, device=q.device)
-            k.data = torch.empty(0, device=k.device)
-            v.data = torch.empty(0, device=v.device)
-            torch.cuda.empty_cache()
 
+            # testing forward
+            o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
             o_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)
@@ -218,72 +215,52 @@ def main(args):
             print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
 
             # test backward
+            grad_o = torch.randn_like(o)
+            o.backward(grad_o)
             o_sdpa.backward(grad_o)
 
             sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad)
-            # Error bounds collected on H100
-            assert sim > 0.9999, f"SSIM too low: {sim}"
-            assert l1 < 4e-3, f"l1 too large: {l1}"
-            assert rmse < 3e-4, f"RMSE too large: {rmse}"
             grad_q_metrics['sim'].append(sim)
             grad_q_metrics['l1'].append(l1)
             grad_q_metrics['rmse'].append(rmse)
             print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}")        
             
             sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad)
-            assert sim > 0.9999, f"SSIM too low: {sim}"
-            assert l1 < 4e-3, f"l1 too large: {l1}"
-            assert rmse < 2e-4, f"RMSE too large: {rmse}"
             grad_k_metrics['sim'].append(sim)
             grad_k_metrics['l1'].append(l1)
             grad_k_metrics['rmse'].append(rmse)
             print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
             
             sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad)
-            assert sim > 0.9999, f"SSIM too low: {sim}"
-            assert l1 < 1e-4, f"l1 too large: {l1}"
-            assert rmse < 2e-5, f"RMSE too large: {rmse}"
             grad_v_metrics['sim'].append(sim)
             grad_v_metrics['l1'].append(l1)
             grad_v_metrics['rmse'].append(rmse)
             print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
-            
-            del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa
-            gc.collect()
-            torch.cuda.empty_cache()
-
+        
         # Print summary statistics if multiple iterations were run
         if num_iterations > 1:
             print("\n" + "="*50)
             print(f"Summary Statistics (over {num_iterations} iterations):")
             
             print("\nForward metrics:")
-            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}")
+            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}")
+            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}")
+            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(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(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(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}")
+            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(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(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(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}")
+            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(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}")
+            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(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}")
+            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(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)
+    main()

csrc/attn/test/test_sta.py

@@ -81,7 +81,5 @@ std = 10
 
 # Run correctness check directly
 results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output')
-assert results['TK vs FLEX']['avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large"
-assert results['TK vs FLEX']['max_diff'] < 4e-2, f"Maximum difference: {results['TK vs FLEX']['max_diff']} is too large"
 print(f"Average difference: {results['TK vs FLEX']['avg_diff']}")
 print(f"Maximum difference: {results['TK vs FLEX']['max_diff']}")

csrc/attn/vsa/block_sparse_h100.cu

@@ -3,8 +3,6 @@
 #include "kittens.cuh"
 #include <cooperative_groups.h>
 #include <iostream>
-#include <c10/cuda/CUDAGuard.h>
-
 
 using namespace kittens;
 namespace cg = cooperative_groups;
@@ -942,9 +940,8 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
     float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
     float* d_l   = reinterpret_cast<float*>(l_ptr);
 
-    //cudadevicesynchronize();
-    const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
-    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
+    cudaDeviceSynchronize();
+    auto stream = at::cuda::getCurrentCUDAStream().stream(); 
 
     if (head_dim == 64) {
         using q_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>;
@@ -969,7 +966,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
 
         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;
+        auto mem_size = 54000;
 
         dim3 grid(seq_len/(64), qo_heads, batch);
 
@@ -982,7 +979,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
         fwd_attend_ker<64><<<grid, (128), mem_size, stream>>>(g);
 
         CHECK_CUDA_ERROR(cudaGetLastError());
-        // cudaStreamSynchronize(stream);
+        cudaStreamSynchronize(stream);
     }
 
     if (head_dim == 128) {
@@ -1008,7 +1005,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
 
         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;
+        auto mem_size = 54000;
 
         dim3 grid(seq_len/(64), qo_heads, batch);
 
@@ -1021,11 +1018,11 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
         fwd_attend_ker<128><<<grid, (128), mem_size, stream>>>(g);
 
         CHECK_CUDA_ERROR(cudaGetLastError());
-        // cudaStreamSynchronize(stream);
+        cudaStreamSynchronize(stream);
     }
 
     return {o, l_vec};
-    //cudadevicesynchronize();
+    cudaDeviceSynchronize();
 }
 
 std::vector<torch::Tensor> 
@@ -1135,14 +1132,13 @@ block_sparse_attention_backward(torch::Tensor q,
     float* d_kg = reinterpret_cast<float*>(kg_ptr);
     float* d_vg = reinterpret_cast<float*>(vg_ptr);
 
-    constexpr int mem_size = kittens::MAX_SHARED_MEMORY; 
-    int threads  = 4 * kittens::WARP_THREADS;
+    auto mem_size = kittens::MAX_SHARED_MEMORY; 
+    auto threads  = 4 * kittens::WARP_THREADS;
 
-    //cudadevicesynchronize();
-    const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
-    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
+    cudaDeviceSynchronize();
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
 
-    //  cudaStreamSynchronize(stream);
+    cudaStreamSynchronize(stream);
 
     // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256");
     dim3 grid_bwd(seq_len/(4*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
@@ -1226,7 +1222,7 @@ block_sparse_attention_backward(torch::Tensor q,
         dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
         threads = 128;
 
-        //cudadevicesynchronize();
+        cudaDeviceSynchronize();
 
         {
             cudaFuncSetAttribute(
@@ -1244,8 +1240,8 @@ block_sparse_attention_backward(torch::Tensor q,
         }
 
         // CHECK_CUDA_ERROR(cudaGetLastError());
-        // cudaStreamSynchronize(stream);
-        //cudadevicesynchronize();
+        cudaStreamSynchronize(stream);
+        cudaDeviceSynchronize();
         // const auto kernel_end = std::chrono::high_resolution_clock::now();
         // std::cout << "Kernel Time: " << std::chrono::duration_cast<std::chrono::microseconds>(kernel_end - start).count() << "us" << std::endl;
         // std::cout << "---" << std::endl;
@@ -1330,7 +1326,7 @@ block_sparse_attention_backward(torch::Tensor q,
         dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
         threads = 128;
 
-        //cudadevicesynchronize();
+        cudaDeviceSynchronize();
 
         {
             cudaFuncSetAttribute(
@@ -1342,10 +1338,10 @@ block_sparse_attention_backward(torch::Tensor q,
             bwd_attend_ker<128><<<grid_bwd_2, threads, 113000, stream>>>(bwd_global); 
         }
 
-        // cudaStreamSynchronize(stream);
-        //cudadevicesynchronize();
+        cudaStreamSynchronize(stream);
+        cudaDeviceSynchronize();
     }
 
     return {qg, kg, vg};
-    //cudadevicesynchronize();
+    cudaDeviceSynchronize();
 }

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

@@ -7,7 +7,7 @@
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
-#SBATCH --nodelist=fs-mbz-gpu-[100-850]
+#SBATCH --nodelist=fs-mbz-gpu-[400-550]
 #SBATCH --mem=1440G
 #SBATCH --output=4n_i2v/4n_i2v_%j.out
 #SBATCH --error=4n_i2v/4n_i2v_%j.err
@@ -62,8 +62,8 @@ srun torchrun \
     --cache_dir "/home/ray/.cache"\
     --data_path "$DATA_DIR"\
     --validation_preprocessed_path "$VALIDATION_DIR"\
-    --train_batch_size=2\
-    --num_latent_t 8 \
+    --train_batch_size=1\
+    --num_latent_t 16 \
     --num_gpus $NUM_GPUS \
     --sp_size $NUM_GPUS \
     --tp_size $NUM_GPUS \
@@ -71,12 +71,12 @@ srun torchrun \
     --hsdp_shard_dim $NUM_GPUS \
     --train_sp_batch_size 1\
     --dataloader_num_workers 10\
-    --gradient_accumulation_steps=1\
+    --gradient_accumulation_steps=2\
     --max_train_steps=10000 \
     --learning_rate=5e-5\
     --mixed_precision="bf16"\
     --checkpointing_steps=11000 \
-    --validation_steps 10\
+    --validation_steps 100\
     --validation_sampling_steps "40" \
     --log_validation \
     --checkpoints_total_limit 3\

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

@@ -16,7 +16,7 @@ training_args=(
   --max_train_steps 5000
   --train_batch_size 1
   --train_sp_batch_size 1
-  --gradient_accumulation_steps 8
+  --gradient_accumulation_steps 1
   --num_latent_t 8 
   --num_height 480
   --num_width 832
@@ -48,17 +48,17 @@ dataset_args=(
 validation_args=(
   --log_validation 
   --validation_preprocessed_path $VALIDATION_DIR
-  --validation_steps 50
+  --validation_steps 100
   --validation_sampling_steps "50" 
   --validation_guidance_scale "1.0"
 )
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 5e-5
+  --learning_rate 1e-5
   --mixed_precision "bf16"
   --checkpointing_steps 6000
-  --weight_decay 1e-4
+  --weight_decay 0.01
   --max_grad_norm 1.0
 )
 

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

@@ -1,5 +1,5 @@
 #!/bin/bash
-#SBATCH --job-name=t2v
+#SBATCH --job-name=FV_2N_14B
 #SBATCH --partition=main
 #SBATCH --qos=hao
 #SBATCH --nodes=1
@@ -9,8 +9,8 @@
 #SBATCH --cpus-per-task=128
 #SBATCH --nodelist=fs-mbz-gpu-[400-550]
 #SBATCH --mem=1440G
-#SBATCH --output=1n_t2v/1n_t2v_%j.out
-#SBATCH --error=1n_t2v/1n_t2v_%j.err
+#SBATCH --output=4n_i2v/4n_i2v_%j.out
+#SBATCH --error=4n_i2v/4n_i2v_%j.err
 #SBATCH --exclusive
 set -e -x
 
@@ -63,13 +63,13 @@ srun torchrun \
     --cache_dir "/home/ray/.cache"\
     --data_path "$DATA_DIR"\
     --validation_preprocessed_path "$VALIDATION_DIR"\
-    --train_batch_size=4\
+    --train_batch_size=1\
     --num_latent_t 8 \
     --num_gpus $NUM_GPUS \
-    --sp_size 4 \
-    --tp_size 4 \
-    --hsdp_replicate_dim 2 \
-    --hsdp_shard_dim 4 \
+    --sp_size $NUM_GPUS \
+    --tp_size $NUM_GPUS \
+    --hsdp_replicate_dim $SLURM_JOB_NUM_NODES \
+    --hsdp_shard_dim $NUM_GPUS \
     --train_sp_batch_size 1\
     --dataloader_num_workers 10\
     --gradient_accumulation_steps=1\
@@ -77,15 +77,15 @@ srun torchrun \
     --learning_rate=5e-5\
     --mixed_precision="bf16"\
     --checkpointing_steps=11000 \
-    --validation_steps 50\
-    --validation_sampling_steps "50" \
+    --validation_steps 100\
+    --validation_sampling_steps "40" \
     --log_validation \
     --checkpoints_total_limit 3\
     --allow_tf32\
     --ema_start_step 0\
     --cfg 0.0\
-    --output_dir="$DATA_DIR/outputs/wan_t2v_finetune_2n"\
-    --tracker_project_name wan_t2v_finetune \
+    --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"\
+    --tracker_project_name wan_i2v_finetune \
     --num_height 480 \
     --num_width 832 \
     --num_frames 77 \

fastvideo/v1/configs/models/dits/base.py

@@ -12,7 +12,6 @@ class DiTArchConfig(ArchConfig):
     _fsdp_shard_conditions: list = field(default_factory=list)
     _compile_conditions: list = field(default_factory=list)
     _param_names_mapping: dict = field(default_factory=dict)
-    _reverse_param_names_mapping: dict = field(default_factory=dict)
     _lora_param_names_mapping: dict = field(default_factory=dict)
     _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
         AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,

fastvideo/v1/configs/models/dits/hunyuanvideo.py

@@ -147,9 +147,6 @@ class HunyuanVideoArchConfig(DiTArchConfig):
             r"final_layer.linear.\1",
         })
 
-    # Reverse mapping for saving checkpoints: training -> diffusers
-    _reverse_param_names_mapping: dict = field(default_factory=lambda: {})
-
     patch_size: int = 2
     patch_size_t: int = 1
     in_channels: int = 16

fastvideo/v1/configs/models/dits/wanvideo.py

@@ -49,13 +49,9 @@ class WanVideoArchConfig(DiTArchConfig):
             r"blocks.\1.ffn.fc_in.\2",
             r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
             r"blocks.\1.ffn.fc_out.\2",
-            r"^blocks\.(\d+)\.norm2\.(.*)$":
+            r"blocks\.(\d+)\.norm2\.(.*)$":
             r"blocks.\1.self_attn_residual_norm.norm.\2",
         })
-
-    # Reverse mapping for saving checkpoints: training -> diffusers
-    _reverse_param_names_mapping: dict = field(default_factory=lambda: {})
-
     # Some LoRA adapters use the original official layer names instead of hf layer names,
     # so apply this before the param_names_mapping
     _lora_param_names_mapping: dict = field(

fastvideo/v1/dataset/preprocessing_datasets.py

@@ -411,29 +411,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
                                 torch.distributed.checkpoint.stateful.Stateful):
     """
     Merged dataset for video and caption data with stage-based processing.
-    Assumes that data_merge_path is a txt file with the following format:
-    <folder_path>,<json_file_path>
-
-        The folder should contain videos.
-
-        The json file should be a list of dictionaries with the following format:
-        [
-        {
-            "path": "1gGQy4nxyUo-Scene-016.mp4",
-            "resolution": {
-            "width": 1920,
-            "height": 1080
-            },
-            "size": 2439112,
-            "fps": 25.0,
-            "duration": 6.88,
-            "num_frames": 172,
-            "cap": [
-            "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open."
-            ]
-        },
-        ...
-        ]
     
     This dataset processes video and image data through a series of stages:
     - Data validation
@@ -488,23 +465,26 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
 
     def _load_raw_data(self) -> List[Dict]:
         """Load raw data from JSON files."""
+        all_data = []
+
         # Read folder-annotation pairs
         with open(self.data_merge_path) as f:
             folder_anno_pairs = [
                 line.strip().split(",") for line in f if line.strip()
             ]
-        assert len(folder_anno_pairs) == 1, "Only support one folder-annotation pair"
-        assert len(folder_anno_pairs[0]) == 2, "Folder-annotation pair should have two elements"
-        folder, annotation_file = folder_anno_pairs[0]
 
-        with open(annotation_file) as f:
-            data_items = json.load(f)
+        # Process each folder-annotation pair
+        for folder, annotation_file in folder_anno_pairs:
+            with open(annotation_file) as f:
+                data_items = json.load(f)
 
-        # Update paths with folder prefix
-        for item in data_items:
-            item["path"] = opj(folder, item["path"])
+            # Update paths with folder prefix
+            for item in data_items:
+                item["path"] = opj(folder, item["path"])
 
-        return data_items
+            all_data.extend(data_items)
+
+        return all_data[self.start_idx:]
 
     def _process_metadata(self) -> List[PreprocessBatch]:
         """Process the raw metadata through all filtering stages."""

fastvideo/v1/models/dits/base.py

@@ -14,7 +14,6 @@ class BaseDiT(nn.Module, ABC):
     _fsdp_shard_conditions: list = []
     _compile_conditions: list = []
     _param_names_mapping: dict
-    _reverse_param_names_mapping: dict
     hidden_size: int
     num_attention_heads: int
     num_channels_latents: int
@@ -79,7 +78,6 @@ class CachableDiT(BaseDiT):
     # These are required class attributes that should be overridden by concrete implementations
     _fsdp_shard_conditions = []
     _param_names_mapping = {}
-    _reverse_param_names_mapping = {}
     _lora_param_names_mapping: dict = {}
     # Ensure these instance attributes are properly defined in subclasses
     hidden_size: int

fastvideo/v1/models/dits/hunyuanvideo.py

@@ -442,8 +442,6 @@ class HunyuanVideoTransformer3DModel(CachableDiT):
     _supported_attention_backends = HunyuanVideoConfig(
     )._supported_attention_backends
     _param_names_mapping = HunyuanVideoConfig()._param_names_mapping
-    _reverse_param_names_mapping = HunyuanVideoConfig(
-    )._reverse_param_names_mapping
     _lora_param_names_mapping = HunyuanVideoConfig()._lora_param_names_mapping
 
     def __init__(self, config: HunyuanVideoConfig, hf_config: dict[str, Any]):

fastvideo/v1/models/dits/stepvideo.py

@@ -460,8 +460,6 @@ class StepVideoModel(BaseDiT):
         # lambda n, m: "pos_embed" in n  # If needed for the patch embedding.
     ]
     _param_names_mapping = StepVideoConfig()._param_names_mapping
-    _reverse_param_names_mapping = StepVideoConfig(
-    )._reverse_param_names_mapping
     _lora_param_names_mapping = StepVideoConfig()._lora_param_names_mapping
     _supported_attention_backends = StepVideoConfig(
     )._supported_attention_backends

fastvideo/v1/models/dits/wanvideo.py

@@ -521,7 +521,6 @@ class WanTransformer3DModel(CachableDiT):
     _supported_attention_backends = WanVideoConfig(
     )._supported_attention_backends
     _param_names_mapping = WanVideoConfig()._param_names_mapping
-    _reverse_param_names_mapping = WanVideoConfig()._reverse_param_names_mapping
     _lora_param_names_mapping = WanVideoConfig()._lora_param_names_mapping
 
     def __init__(self, config: WanVideoConfig, hf_config: dict[str,

fastvideo/v1/models/loader/fsdp_load.py

@@ -222,14 +222,10 @@ def load_model_from_full_model_state_dict(
     used_keys = set()
     sharded_sd = {}
     to_merge_params: DefaultDict[str, Dict[Any, Any]] = defaultdict(dict)
-    reverse_param_names_mapping = {}
-    assert param_names_mapping is not None
     for source_param_name, full_tensor in full_sd_iterator:
+        assert param_names_mapping is not None
         target_param_name, merge_index, num_params_to_merge = param_names_mapping(
             source_param_name)
-        reverse_param_names_mapping[target_param_name] = (source_param_name,
-                                                          merge_index,
-                                                          num_params_to_merge)
         used_keys.add(target_param_name)
         if merge_index is not None:
             to_merge_params[target_param_name][merge_index] = full_tensor
@@ -264,7 +260,6 @@ def load_model_from_full_model_state_dict(
                 sharded_tensor = sharded_tensor.cpu()
         sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)
 
-    model._reverse_param_names_mapping = reverse_param_names_mapping
     unused_keys = set(meta_sd.keys()) - used_keys
     if unused_keys:
         logger.warning("Found new parameters in meta state dict: %s",

fastvideo/v1/pipelines/pipeline_batch_info.py

@@ -149,7 +149,6 @@ class TrainingBatch:
 
     # Dataloader batch outputs
     latents: Optional[torch.Tensor] = None
-    # original_latents: Optional[torch.Tensor] = None
     encoder_hidden_states: Optional[torch.Tensor] = None
     encoder_attention_mask: Optional[torch.Tensor] = None
     # i2v

fastvideo/v1/pipelines/preprocess/v1_preprocess.py

@@ -59,6 +59,12 @@ if __name__ == "__main__":
         default=2,
         help="Batch size (per device) for the training dataloader.",
     )
+    parser.add_argument(
+        "--preprocess_text_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
     parser.add_argument("--samples_per_file", type=int, default=64)
     parser.add_argument("--flush_frequency",
                         type=int,

fastvideo/v1/pipelines/stages/encoding.py

@@ -68,8 +68,9 @@ class EncodingStage(PipelineStage):
                 width=batch.width).to(get_torch_device(), dtype=torch.float32)
 
             image = image.unsqueeze(2)
+        else:
+            image = image.transpose(1, 2)
         logger.info("image: %s", image.shape)
-        image = image.transpose(1, 2)
         video_condition = torch.cat([
             image,
             image.new_zeros(image.shape[0], image.shape[1],

fastvideo/v1/tests/training/Vanilla/test_training_loss.py

@@ -116,7 +116,7 @@ def test_distributed_training():
         'avg_step_time': 1.0,
         'grad_norm': 0.2,
         'step_time': 0.5,
-        'train_loss': 0.0025
+        'train_loss': 0.001
     }
 
     failures = []

fastvideo/v1/training/training_pipeline.py

@@ -25,8 +25,7 @@ from fastvideo.v1.dataset import build_parquet_map_style_dataloader
 from fastvideo.v1.dataset.dataloader.schema import (
     pyarrow_schema_t2v, pyarrow_schema_t2v_validation)
 from fastvideo.v1.distributed import (cleanup_dist_env_and_memory, get_sp_group,
-                                      get_torch_device, get_world_group,
-                                      sequence_model_parallel_all_gather)
+                                      get_torch_device, get_world_group)
 from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
@@ -321,10 +320,6 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
             model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
         target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
 
-        # if self.training_args.sp_size > 1:
-        #     model_pred = sequence_model_parallel_all_gather(model_pred, dim=2)
-
-        assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
         loss = (torch.mean((model_pred.float() - target.float())**2) /
                 self.training_args.gradient_accumulation_steps)
 
@@ -368,31 +363,14 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         for _ in range(self.training_args.gradient_accumulation_steps):
             training_batch = self._get_next_batch(training_batch)
 
-            training_batch.latents = training_batch.latents[:, :, :self.
-                                                            training_args.
-                                                            num_latent_t]
-
-            # Normalize DIT input
-            training_batch = self._normalize_dit_input(training_batch)
-
-            # Save original latents
-            # training_batch.original_latents = training_batch.latents.clone()
-
-
-            training_batch = self._prepare_dit_inputs(training_batch)
-
-            # # Shard latents across sp groups
+            # Shard latents across sp groups
             training_batch.latents = shard_latents_across_sp(
                 training_batch.latents,
                 num_latent_t=self.training_args.num_latent_t)
-            # shard noisy_model_input to match
-            training_batch.noisy_model_input = shard_latents_across_sp(
-                training_batch.noisy_model_input,
-                num_latent_t=self.training_args.num_latent_t)
-            # CRITICAL FIX: Also shard noise to match latents
-            training_batch.noise = shard_latents_across_sp(
-                training_batch.noise,
-                num_latent_t=self.training_args.num_latent_t)
+
+            # Normalize DIT input
+            training_batch = self._normalize_dit_input(training_batch)
+            training_batch = self._prepare_dit_inputs(training_batch)
             training_batch = self._build_attention_metadata(training_batch)
             training_batch = self._build_input_kwargs(training_batch)
             training_batch = self._transformer_forward_and_compute_loss(

fastvideo/v1/training/training_utils.py

@@ -8,6 +8,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dcp
+import torch.distributed.checkpoint.stateful
 from einops import rearrange
 from safetensors.torch import save_file
 
@@ -153,20 +154,13 @@ def save_checkpoint(transformer,
 
     if rank == 0:
         # Save model weights (consolidated)
-        transformer_save_dir = os.path.join(save_dir, "transformer")
-        os.makedirs(transformer_save_dir, exist_ok=True)
-        weight_path = os.path.join(transformer_save_dir,
+        weight_path = os.path.join(save_dir,
                                    "diffusion_pytorch_model.safetensors")
         logger.info("rank: %s, saving consolidated checkpoint to %s",
                     rank,
                     weight_path,
                     local_main_process_only=False)
-
-        # Convert training format to diffusers format and save
-        diffusers_state_dict = convert_training_to_diffusers_format(
-            cpu_state, transformer)
-        save_file(diffusers_state_dict, weight_path)
-
+        save_file(cpu_state, weight_path)
         logger.info("rank: %s, consolidated checkpoint saved to %s",
                     rank,
                     weight_path,
@@ -176,7 +170,7 @@ def save_checkpoint(transformer,
         config_dict = transformer.hf_config
         if "dtype" in config_dict:
             del config_dict["dtype"]  # TODO
-        config_path = os.path.join(transformer_save_dir, "config.json")
+        config_path = os.path.join(save_dir, "config.json")
         # save dict as json
         with open(config_path, "w") as f:
             json.dump(config_dict, f, indent=4)
@@ -485,68 +479,3 @@ def _has_foreach_support(tensors: List[torch.Tensor],
                          device: torch.device) -> bool:
     return _device_has_foreach_support(device) and all(
         t is None or type(t) in [torch.Tensor] for t in tensors)
-
-
-def convert_training_to_diffusers_format(state_dict: Dict[str, Any],
-                                         transformer) -> Dict[str, Any]:
-    """
-    Convert training format state dict to diffusers format using reverse_param_names_mapping.
-    
-    Args:
-        state_dict: State dict in training format
-        transformer: Transformer model object with _reverse_param_names_mapping
-        
-    Returns:
-        State dict in diffusers format
-    """
-    new_state_dict = {}
-
-    # Get the reverse mapping from the transformer
-    reverse_param_names_mapping = transformer._reverse_param_names_mapping
-    assert reverse_param_names_mapping != {}, "reverse_param_names_mapping is empty"
-
-    # Group parameters that need to be split (merged parameters)
-    merge_groups: Dict[str, List[Tuple[str, int, int]]] = {}
-
-    # First pass: collect all merge groups
-    for training_key, (
-            diffusers_key, merge_index,
-            num_params_to_merge) in reverse_param_names_mapping.items():
-        if merge_index is not None:
-            # This is a merged parameter that needs to be split
-            if training_key not in merge_groups:
-                merge_groups[training_key] = []
-            merge_groups[training_key].append(
-                (diffusers_key, merge_index, num_params_to_merge))
-
-    # Second pass: handle merged parameters by splitting them
-    used_keys = set()
-    for training_key, splits in merge_groups.items():
-        if training_key in state_dict:
-            v = state_dict[training_key]
-            # Sort by merge_index to ensure correct order
-            splits.sort(key=lambda x: x[1])
-            total = splits[0][2]
-            split_size = v.shape[0] // total
-            split_tensors = torch.split(v, split_size, dim=0)
-
-            for diffusers_key, split_index, _ in splits:
-                new_state_dict[diffusers_key] = split_tensors[split_index]
-            used_keys.add(training_key)
-
-    # Third pass: handle regular parameters (direct mappings)
-    for training_key, v in state_dict.items():
-        if training_key in used_keys:
-            continue
-
-        if training_key in reverse_param_names_mapping:
-            diffusers_key, merge_index, _ = reverse_param_names_mapping[
-                training_key]
-            if merge_index is None:
-                # Direct mapping
-                new_state_dict[diffusers_key] = v
-        else:
-            # No mapping found, keep as is
-            new_state_dict[training_key] = v
-
-    return new_state_dict

fastvideo/v1/training/wan_i2v_training_pipeline.py

@@ -105,30 +105,6 @@ class WanI2VTrainingPipeline(TrainingPipeline):
         # training_batch.extra_latents = extra_latents
         training_batch.infos = infos
 
-        assert training_batch.image_latents is not None
-        # if self.training_args and hasattr(self.training_args, 'num_latent_t'):
-        training_batch.image_latents = training_batch.image_latents[:, :, :self.training_args.num_latent_t]
-
-        return training_batch
-
-    def _prepare_dit_inputs(self,
-                            training_batch: TrainingBatch) -> TrainingBatch:
-        """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning."""
-        assert self.training_args is not None
-        assert training_batch.latents is not None
-        assert training_batch.encoder_hidden_states is not None
-        assert training_batch.encoder_attention_mask is not None
-        assert self.noise_random_generator is not None
-        assert training_batch.image_latents is not None
-
-        # First, call parent method to prepare noise, timesteps, etc. for video latents
-        training_batch = super()._prepare_dit_inputs(training_batch)
-        
-        image_latents = training_batch.image_latents.to(get_torch_device(), dtype=torch.bfloat16)
-        
-        training_batch.noisy_model_input = torch.cat(
-            [training_batch.noisy_model_input, image_latents], dim=1)
-        
         return training_batch
 
     def _build_input_kwargs(self,
@@ -138,15 +114,34 @@ class WanI2VTrainingPipeline(TrainingPipeline):
         assert training_batch.encoder_hidden_states is not None
         assert training_batch.encoder_attention_mask is not None
         assert training_batch.timesteps is not None
+        assert training_batch.preprocessed_image is not None
         assert training_batch.image_embeds is not None
+        assert training_batch.image_latents is not None
+        # assert training_batch.extra_latents is not None
 
-        # Image Embeds for conditioning
+        # extra_latents = training_batch.extra_latents
+        # if extra_latents:
+        #     image_embeds, image_latents = extra_latents[
+        #         "clip_feature"], extra_latents["first_frame_latent"]
+        # image_
+        # Image Embeds
         image_embeds = training_batch.image_embeds
+        image_latents = training_batch.image_latents
+        preprocessed_image = training_batch.preprocessed_image
         assert torch.isnan(image_embeds).sum() == 0
         image_embeds = image_embeds.to(get_torch_device(), dtype=torch.bfloat16)
         encoder_hidden_states_image = image_embeds
 
-        # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs
+        # Image Latents
+        assert torch.isnan(image_latents).sum() == 0
+        image_latents = image_latents.to(get_torch_device(),
+                                         dtype=torch.bfloat16)
+        image_latents = shard_latents_across_sp(
+            image_latents, num_latent_t=self.training_args.num_latent_t)
+
+        training_batch.noisy_model_input = torch.cat(
+            [training_batch.noisy_model_input, image_latents], dim=1)
+
         training_batch.input_kwargs = {
             "hidden_states":
             training_batch.noisy_model_input,
@@ -176,13 +171,25 @@ class WanI2VTrainingPipeline(TrainingPipeline):
         masks = validation_batch['text_attention_mask']
         clip_features = validation_batch['clip_feature']
         # extra_latents = validation_batch['extra_latents']
-        pil_image = validation_batch['pil_image']
+        preprocessed_image = validation_batch['pil_image']
         infos = validation_batch['info_list']
         prompt = infos[0]['prompt']
 
         prompt_embeds = embeddings.to(get_torch_device())
         prompt_attention_mask = masks.to(get_torch_device())
         clip_features = clip_features.to(get_torch_device())
+
+        if 'colorful candies' in prompt:
+            logger.info("colorful candies")
+            from fastvideo.v1.models.vision_utils import load_video
+            # video_path = 'validation_dataset/yYcK4nANZz4-Scene-030.mp4'
+            video_path = '/mnt/user_storage/fv/FastVideo/examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation_dataset/yYcK4nANZz4-Scene-030.mp4'
+            video = load_video(video_path)
+            pil_image = video[0]
+            preprocessed_image = None
+        else:
+            pil_image = None
+
         # clip_features = extra_latents.get("clip_feature")
         # first_frame_latent = extra_latents.get("first_frame_latent")
         # pil_image = extra_latents.get("pil_image")
@@ -220,7 +227,8 @@ class WanI2VTrainingPipeline(TrainingPipeline):
             negative_prompt_embeds=[negative_prompt_embeds],
             negative_attention_mask=[negative_prompt_attention_mask],
             image_embeds=[clip_features],
-            preprocessed_image=pil_image,
+            preprocessed_image=preprocessed_image,
+            pil_image=pil_image,
             height=training_args.num_height,
             width=training_args.num_width,
             num_frames=num_frames,