update docker

.buildkite/pipeline.yml

@@ -22,7 +22,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 20m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                 label: "Encoder Tests"
                 env:
                   - TEST_TYPE=encoder
@@ -35,7 +35,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 20m .buildkite/scripts/pr_test.sh"
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                 label: "VAE Tests"
                 env:
                   - TEST_TYPE=vae
@@ -129,7 +129,11 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
-                - "fastvideo-kernel/**"
+                - "csrc/attn/video_sparse_attn/**"
+                - "csrc/attn/video_sparse_attn/tk/**"
+                - "csrc/attn/video_sparse_attn/setup.py"
+                - "csrc/attn/video_sparse_attn/config_vsa.py"
+                - "csrc/attn/video_sparse_attn/vsa.cpp"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -141,7 +145,10 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
-                - "fastvideo-kernel/**"
+                - "csrc/attn/sliding_tile_attn/**"
+                - "csrc/attn/sliding_tile_attn/setup.py"
+                - "csrc/attn/sliding_tile_attn/config_sta.py"
+                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -152,16 +159,48 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "fastvideo-kernel/**"
+                - "csrc/attn/sliding_tile_attn/**"
+                - "csrc/attn/sliding_tile_attn/setup.py"
+                - "csrc/attn/sliding_tile_attn/config_sta.py"
+                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
                 command: "timeout 15m .buildkite/scripts/pr_test.sh"
-                label: "Kernel Tests"
+                label: "Precision Tests STA"
                 env:
-                  - TEST_TYPE=kernel_tests
+                  - TEST_TYPE=precision_sta
+                agents:
+                  queue: "default"
             - path:
-                - "fastvideo-kernel/**"
+                - "csrc/attn/video_sparse_attn/**"
+                - "csrc/attn/video_sparse_attn/tk/**"
+                - "csrc/attn/tests/test_vsa.py"
+                - "csrc/attn/video_sparse_attn/setup.py"
+                - "csrc/attn/video_sparse_attn/config_vsa.py"
+                - "csrc/attn/video_sparse_attn/vsa.cpp"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Precision Tests VSA"
+                env:
+                  - TEST_TYPE=precision_vsa
+                agents:
+                  queue: "default"
+            - path:
+                - "csrc/attn/vmoba_attn/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Precision Tests VMoBA"
+                env:
+                  - TEST_TYPE=precision_vmoba
+                agents:
+                  queue: "default"
+            - path:
+                - "csrc/attn/vmoba_attn/vmoba/**"
                 - "fastvideo/attention/backends/vmoba.py"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
@@ -183,14 +222,14 @@ steps:
                   - TEST_TYPE=unit_test
                 agents:
                   queue: "default"
-            # - path:
+            - path:
-            #     - "scripts/lora_extraction/**"
+                - "scripts/lora_extraction/**"
-            #     - "pyproject.toml"
+                - "pyproject.toml"
-            #     - "docker/Dockerfile.python3.12"
+                - "docker/Dockerfile.python3.12"
-            #   config:
+              config:
-            #     command: "timeout 90m .buildkite/scripts/pr_test.sh"
+                command: "timeout 90m .buildkite/scripts/pr_test.sh"
-            #     label: "LoRA Extraction Tests"
+                label: "LoRA Extraction Tests"
-            #     env:
+                env:
-            #       - TEST_TYPE=lora_extraction
+                  - TEST_TYPE=lora_extraction
-            #     agents:
+                agents:
-            #       queue: "default"
+                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -93,9 +93,13 @@ case "$TEST_TYPE" in
         log "Running inference STA tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_STA"
         ;;
-    "kernel_tests")
+    "precision_sta")
-        log "Running kernel tests..."
+        log "Running precision STA tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_kernel_tests"
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_STA"
+        ;;
+    "precision_vsa")
+        log "Running precision VSA tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_VSA"
         ;;
     "inference_lora")
         log "Running LoRA tests..."
@@ -114,6 +118,10 @@ case "$TEST_TYPE" in
         log "Running V-MoBA inference tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_vmoba"
         ;;
+    "precision_vmoba")
+        log "Running V-MoBA precision tests..."
+        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_vmoba"
+        ;;
     "unit_test")
         log "Running unit tests..."
         MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_unit_test"

.github/workflows/fastvideo-kernel-publish.yml

@@ -5,7 +5,7 @@ on:
     branches:
       - main
     paths:
-      - "fastvideo-kernel/pyproject.toml"
+      - "csrc/fastvideo_kernel/pyproject.toml"
   workflow_dispatch:
 
 jobs:
@@ -23,15 +23,13 @@ jobs:
       - name: Check if version changed
         id: check-version
         run: |
-          cd fastvideo-kernel
+          cd csrc/fastvideo_kernel
           # Get current commit's version from pyproject.toml
-          # Use ^ to match start of line to avoid matching minimum-version
+          NEW_VERSION=$(grep -oP 'version\s*=\s*"\K[^"]+' pyproject.toml)
-          NEW_VERSION=$(grep -oP '^version\s*=\s*"\K[^"]+' pyproject.toml)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          # Note: git show expects path relative to repo root
+          OLD_VERSION=$(git show HEAD~1:./pyproject.toml | grep -oP 'version\s*=\s*"\K[^"]+' || echo "0.0.0")
-          OLD_VERSION=$(git show HEAD~1:fastvideo-kernel/pyproject.toml | grep -oP '^version\s*=\s*"\K[^"]+' || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -53,18 +51,15 @@ jobs:
       fail-fast: false
       matrix:
           os: [ubuntu-22.04]
-          python-version: ['3.10', '3.11', '3.12']
+          python-version: ['3.10', '3.11', '3.12', '3.13']
           torch-cuda:
-            # - torch-version: '2.5.1'
+            - torch-version: '2.5.1'
-            #   cuda-version: '12.4.1'
+              cuda-version: '12.4.1'
-            #   torch-cuda-short: 'cu124'
+              torch-cuda-short: 'cu124'
-            # - torch-version: '2.6.0'
+            - torch-version: '2.6.0'
-            #   cuda-version: '12.6.3'
+              cuda-version: '12.6.3'
-            #   torch-cuda-short: 'cu126'
+              torch-cuda-short: 'cu126'
-            # - torch-version: '2.7.1'
+            - torch-version: '2.7.1'
-            #   cuda-version: '12.8.0'
-            #   torch-cuda-short: 'cu128'
-            - torch-version: '2.9.1'
               cuda-version: '12.8.0'
               torch-cuda-short: 'cu128'
 
@@ -143,32 +138,30 @@ jobs:
         run: |
           export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
           
-          pip install setuptools ninja packaging wheel triton scikit-build-core cmake build
+          pip install setuptools ninja packaging wheel triton
           
-          cd fastvideo-kernel
+          cd csrc/fastvideo_kernel
           git submodule update --init --recursive  # Ensure ThunderKittens submodule is initialized
-          # Release builds are produced on GPU-less runners, so force-enable TK and target Hopper.
+          python setup.py bdist_wheel --dist-dir=dist
-          export TORCH_CUDA_ARCH_LIST="9.0a"
+
-          export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
+      - name: Rename wheel file
-          
+        run: |
-          # Build standard wheel (no local version suffix) for PyPI
+          cd csrc/fastvideo_kernel
-          python -m build --wheel --outdir dist
+
-          
+          CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
-          # Fix the wheel to be manylinux compliant
+          TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2)
-          pip install auditwheel
+          # Get the correct version format
-          # Target manylinux_2_35 (Ubuntu 22.04 native)
+          tmpname=cu${CUDA_SHORT_VERSION}torch${TORCH_SHORT_VERSION}
-          auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist
+          wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
-          # Move fixed wheels back to dist for upload consistency
+          # Rename with version information
-          rm dist/*.whl
+          ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
-          mv fixed_dist/*.whl dist/
+          echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
 
       - name: Upload wheel artifact
-        # Only upload if it's the "main" CUDA version we want on PyPI
-        # We upload all to artifacts for inspection/GH releases, but give them distinct artifact names
         uses: actions/upload-artifact@v4
         with:
-          name: fastvideo_kernel-py${{ matrix.python-version }}-${{ matrix.torch-cuda.torch-cuda-short }}-torch${{ matrix.torch-cuda.torch-version }}
+          name: ${{ env.wheel_name }}-py${{ matrix.python-version }}
-          path: fastvideo-kernel/dist/*.whl
+          path: csrc/fastvideo_kernel/dist/*.whl
           retention-days: 90
 
   publish_package:
@@ -186,22 +179,58 @@ jobs:
         with:
           python-version: '3.10'
 
-      - name: Download PyPI wheels
+      - name: Install CUDA 12.4.1
-        uses: actions/download-artifact@v4
+        uses: Jimver/cuda-toolkit@v0.2.21
+        id: cuda-toolkit
         with:
-          path: fastvideo-kernel/dist/
+          cuda: 12.4.1
-          pattern: 'fastvideo_kernel-py*'
+          linux-local-args: '["--toolkit"]'
-          merge-multiple: true
+          method: 'network'
+          sub-packages: '["nvcc"]'
+
+      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
+        run: |
+          sudo apt update
+          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
+          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+          # Allow Git to Access Safe Directory
+          git config --global --add safe.directory /__w/FastVideo/FastVideo
+
+          # Set CUDA environment variables
+          export CUDA_HOME=/usr/local/cuda-12.4.1
+          export PATH=${CUDA_HOME}/bin:${PATH}
+          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+          
+          # Verify installation
+          gcc --version
+          g++ --version
+          clang-11 --version
+          nvcc --version
+
+      - name: Install PyTorch 2.5.1+cu12.4.1
+        run: |
+          pip install --upgrade pip
+          pip install typing-extensions==4.12.2
+          export TORCH_CUDA_VERSION=124
+          pip install --no-cache-dir torch==2.5.1 --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
+          nvcc --version
+          python --version
+          python -c "import torch; print('PyTorch:', torch.__version__)"
+          python -c "import torch; print('CUDA:', torch.version.cuda)"
+          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
 
       - name: Build source distribution
         run: |
-          pip install build scikit-build-core cmake ninja
+          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
           
-          cd fastvideo-kernel
+          pip install setuptools ninja packaging wheel triton
-          # We don't need full CUDA/Torch to just package the source (sdist)
+          
-          python -m build --sdist --outdir dist
+          cd csrc/fastvideo_kernel
+          git submodule update --init --recursive
+          python setup.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: fastvideo-kernel/dist/
+          packages-dir: csrc/fastvideo_kernel/dist/

.gitmodules

@@ -1,6 +1,7 @@
-[submodule "fastvideo-kernel/include/tk"]
+[submodule "csrc/attn/video_sparse_attn/tk"]
-	path = fastvideo-kernel/include/tk
+	path = csrc/attn/video_sparse_attn/tk
+	url = https://github.com/HazyResearch/ThunderKittens.git
+
+[submodule "csrc/attn/sliding_tile_attn/tk"]
+	path = csrc/attn/sliding_tile_attn/tk
 	url = https://github.com/HazyResearch/ThunderKittens.git
-[submodule "fastvideo-kernel/include/cutlass"]
-	path = fastvideo-kernel/include/cutlass
-	url = https://github.com/NVIDIA/cutlass.git

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ default_stages:
 exclude: |
   (?x)(
     fastvideo/third_party/.*|
-    fastvideo-kernel/.*|
+    csrc/.*|
     assets/.*|
     tests/.*|
     demo/.*|

benchmarks/fvd/README.md

@@ -38,8 +38,7 @@ python -m benchmarks.fvd.cli \
     --num-frames 32 \
     --clip-strategy random \
     --batch-size 32 \
-    --seed 42 \
+    --seed 42
-    --extractor clip
 ```
 
 **Standard protocols:**
@@ -52,8 +51,6 @@ python -m benchmarks.fvd.cli \
     --protocol fvd2048_16f  # or fvd2048_128f, quick_test, etc.
 ```
 
-This would use i3d model by default as the feature extractor
-
 **Feature caching** (speed up repeated evaluations):
 
 ```bash
@@ -61,7 +58,7 @@ python -m benchmarks.fvd.cli \
     --real-path data/real/ \
     --gen-path outputs/gen/ \
     --protocol fvd2048_16f \
-    --cache-real-features fvd-cache/extractor_name  # Directory path (will save/load fvd-cache/extractor_name/extractor-name_real_features.pkl)
+    --cache-real-features cache/real  # Directory path (will save/load cache/real/real_features.pkl)
 ```
 
 Run `python -m benchmarks.fvd.cli --help` for all options.
@@ -86,7 +83,6 @@ batch_size=32,                # GPU batch size
 device='cuda',                # cuda|cpu
 cache_real_features=None,     # Cache path for speed
 seed=42,                      # Reproducibility
-extractor='i3d',              # i3d|clip|videomae
 ```
 
 ## Programmatic Usage
@@ -101,6 +97,7 @@ print(f"FVD: {results['fvd']:.2f}")
 
 ## Notes
 
+- I3D model auto-downloads from Hugging Face on first run
 - Requires minimum 10 frames per clip
 - Supports both video files (.mp4, .avi, etc.) and frame directories
 - `--cache-real-features` expects a **directory path** (e.g., `cache/real`), it will automatically create/load `real_features.pkl` inside that directory

benchmarks/fvd/__init__.py

@@ -13,8 +13,7 @@ from .fvd import (
     compute_statistics,
     FVDConfig,
 )
-from .feature_extractors import (BaseFeatureExtractor, I3DFeatureExtractor,
+from .i3d_model import I3DFeatureExtractor
-                                 load_extractor)
 from .video_utils import (
     load_video_auto,
     sample_clips_from_video,
@@ -28,9 +27,7 @@ __all__ = [
     'compute_frechet_distance',
     'compute_statistics',
     'FVDConfig',
-    'BaseFeatureExtractor',
     'I3DFeatureExtractor',
-    'load_extractor',
     'load_video_auto',
     'sample_clips_from_video',
     'load_video_clips_streaming',

benchmarks/fvd/cli.py

@@ -1,64 +1,122 @@
 import argparse
+import json
 import sys
-import traceback
+from pathlib import Path
+
 from .fvd import compute_fvd_with_config, FVDConfig
 
 
 def main() -> int:
     parser = argparse.ArgumentParser(
-        description='Compute Fréchet Video Distance (FVD)')
+        description='Compute Fréchet Video Distance (FVD)',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Standard FVD2048_16f protocol
+  python -m fastvideo.benchmarks.fvd.cli \\
+      --real-path data/real/ \\
+      --gen-path outputs/gen/ \\
+      --protocol fvd2048_16f
+
+  # Custom configuration
+  python -m fastvideo.benchmarks.fvd.cli \\
+      --real-path data/real/ \\
+      --gen-path outputs/gen/ \\
+      --num-videos 1024 \\
+      --num-frames 32 \\
+      --clip-strategy random \\
+      --frame-stride 2
+        """)
 
     # Required arguments
     parser.add_argument('--real-path',
                         type=str,
                         required=True,
-                        help='Path to real videos')
+                        help='Path to real videos directory')
     parser.add_argument('--gen-path',
                         type=str,
                         required=True,
-                        help='Path to generated videos')
+                        help='Path to generated videos directory')
 
-    # Extractor selection
+    # Reproducibility
-    parser.add_argument('--extractor',
+    parser.add_argument(
-                        type=str,
+        '--seed',
-                        default='i3d',
+        type=int,
-                        choices=['i3d', 'clip', 'videomae'],
+        default=None,
-                        help='Feature extractor model to use (default: i3d)')
+        help='Random seed for reproducibility (np.random, random, torch)')
 
-    # Standard args
+    # Protocol presets
-    parser.add_argument('--seed',
-                        type=int,
-                        default=None,
-                        help='Random seed for reproducibility')
     parser.add_argument('--protocol',
                         type=str,
                         default=None,
-                        choices=['fvd2048_16f', 'fvd2048_128f', 'quick_test'],
+                        choices=[
+                            'fvd2048_16f', 'fvd2048_128f',
+                            'fvd2048_128f_subsample8', 'quick_test'
+                        ],
                         help='Use standard protocol (overrides other settings)')
+
+    # Video selection
     parser.add_argument('--num-videos',
                         type=int,
                         default=2048,
-                        help='Number of videos to use')
+                        help='Number of videos to use (default: 2048)')
+
+    # Clip sampling
     parser.add_argument('--num-frames',
                         type=int,
                         default=16,
-                        help='Number of frames per clip')
+                        help='Number of frames per clip (default: 16)')
-    parser.add_argument('--clip-strategy',
+    parser.add_argument('--num-clips',
-                        type=str,
+                        type=int,
-                        default='beginning',
+                        default=1,
-                        help='Clip sampling strategy')
+                        help='Number of clips per video (default: 1)')
+    parser.add_argument(
+        '--clip-strategy',
+        type=str,
+        default='beginning',
+        choices=['beginning', 'random', 'uniform', 'middle', 'sliding', 'all'],
+        help='Clip sampling strategy (default: beginning)')
+    parser.add_argument(
+        '--frame-stride',
+        type=int,
+        default=1,
+        help='Frame stride for FPS subsampling (default: 1, no subsampling)')
+    parser.add_argument('--temporal-stride',
+                        type=int,
+                        default=1,
+                        help='Temporal stride for sliding window (default: 1)')
+
+    # Data processing
+    parser.add_argument('--no-frame-dirs',
+                        action='store_true',
+                        help='Disable frame directory support')
+
+    # Computation
     parser.add_argument('--batch-size',
                         type=int,
                         default=32,
-                        help='Batch size for feature extraction')
+                        help='Batch size for feature extraction (default: 32)')
     parser.add_argument('--device',
                         type=str,
                         default='cuda',
-                        help='Device to use (cuda or cpu)')
+                        choices=['cuda', 'cpu'],
+                        help='Device to use (default: cuda)')
+
+    # Caching
     parser.add_argument('--cache-real-features',
                         type=str,
                         default=None,
                         help='Path to cache real video features')
+    parser.add_argument('--i3d-model-path',
+                        type=str,
+                        default=None,
+                        help='Custom cache path for I3D model')
+
+    # Output
+    parser.add_argument('--output',
+                        type=str,
+                        default='fvd_results.json',
+                        help='Output JSON file (default: fvd_results.json)')
     parser.add_argument('--quiet',
                         action='store_true',
                         help='Suppress progress output')
@@ -70,36 +128,56 @@ def main() -> int:
         protocol_map = {
             'fvd2048_16f': FVDConfig.fvd2048_16f,
             'fvd2048_128f': FVDConfig.fvd2048_128f,
+            'fvd2048_128f_subsample8': FVDConfig.fvd2048_128f_subsample8,
             'quick_test': FVDConfig.quick_test,
         }
         config = protocol_map[args.protocol]()
-        # Apply overrides
+
+        # Override device and caching from args
         config.device = args.device
         config.cache_real_features = args.cache_real_features
-        config.extractor_model = args.extractor  # Apply extractor arg
+        config.i3d_model_path = args.i3d_model_path
+        config.batch_size = args.batch_size
+        config.seed = args.seed
     else:
-        config = FVDConfig(
+        # Custom config from args
-            num_videos=args.num_videos,
+        config = FVDConfig(num_videos=args.num_videos,
-            num_frames_per_clip=args.num_frames,
+                           num_frames_per_clip=args.num_frames,
-            extractor_model=args.extractor,  # Apply extractor arg
+                           num_clips_per_video=args.num_clips,
-            clip_strategy=args.clip_strategy,
+                           clip_strategy=args.clip_strategy,
-            batch_size=args.batch_size,
+                           frame_stride=args.frame_stride,
-            device=args.device,
+                           temporal_stride=args.temporal_stride,
-            cache_real_features=args.cache_real_features,
+                           support_frame_dirs=not args.no_frame_dirs,
-            seed=args.seed)
+                           batch_size=args.batch_size,
+                           device=args.device,
+                           cache_real_features=args.cache_real_features,
+                           i3d_model_path=args.i3d_model_path,
+                           seed=args.seed)
 
+    # Compute FVD
     try:
-        _ = compute_fvd_with_config(
+        results = compute_fvd_with_config(real_videos=args.real_path,
-            args.real_path,  # noqa: F841
+                                          gen_videos=args.gen_path,
-            args.gen_path,
+                                          config=config,
-            config,
+                                          verbose=not args.quiet)
-            verbose=not args.quiet)
+
+        # Save results
+        output_path = Path(args.output)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(output_path, 'w') as f:
+            json.dump(results, f, indent=2)
+
+        print(f"\nResults saved to {output_path}")
+        print(f"FVD: {results['fvd']:.2f}")
+        print(f"Protocol: {results['protocol']}")
 
         return 0
 
     except Exception as e:
         print(f"Error: {e}", file=sys.stderr)
-        traceback.print_exc(file=sys.stderr)
+        import traceback
+        traceback.print_exc()
         return 1
 
 

benchmarks/fvd/feature_extractors.py

@@ -1,264 +0,0 @@
-"""
-Pluggable Feature Extractors for FVD Computation.
-Supports I3D (standard), CLIP, and VideoMAE via a common interface.
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from abc import ABC, abstractmethod
-from huggingface_hub import hf_hub_download
-from tqdm import tqdm
-
-try:
-    from transformers import CLIPModel, CLIPProcessor, VideoMAEModel
-    TRANSFORMERS_AVAILABLE = True
-except ImportError:
-    TRANSFORMERS_AVAILABLE = False
-
-
-class BaseFeatureExtractor(ABC, nn.Module):
-    """Abstract base class for all video feature extractors."""
-
-    def __init__(self, device: str = 'cuda'):
-        super().__init__()
-        self.device = torch.device(
-            device if torch.cuda.is_available() else 'cpu')
-
-    @property
-    @abstractmethod
-    def feature_dim(self) -> int:
-        """Dimension of the output feature vector."""
-        pass
-
-    @abstractmethod
-    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            videos: [B, T, C, H, W] in [0, 255] range.
-        Returns:
-            Preprocessed tensor ready for the model.
-        """
-        pass
-
-    @abstractmethod
-    def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
-        """
-        Extract features for a single batch.
-        Args:
-            videos: [B, T, C, H, W] (raw input)
-        Returns:
-            Features: [B, feature_dim]
-        """
-        pass
-
-    @torch.no_grad()
-    def extract_features(self,
-                         videos: torch.Tensor,
-                         batch_size: int = 32,
-                         verbose: bool = True) -> torch.Tensor:
-        """
-        Extract features for a large tensor of videos by batching.
-        """
-        N = len(videos)
-        all_features = []
-
-        iterator = range(0, N, batch_size)
-        if verbose:
-            iterator = tqdm(
-                iterator,
-                desc=f"Extracting features ({self.__class__.__name__})")
-
-        for i in iterator:
-            batch = videos[i:i + batch_size].to(self.device)
-            features = self.extract_features_batch(batch)
-            all_features.append(features.cpu())
-
-        return torch.cat(all_features, dim=0)
-
-
-# 1. I3D Extractor (The Standard FVD Metric)
-class I3DFeatureExtractor(BaseFeatureExtractor):
-    REPO_ID = 'flateon/FVD-I3D-torchscript'
-    MODEL_FILENAME = 'i3d_torchscript.pt'
-
-    def __init__(self, device: str = 'cuda', cache_dir: str | None = None):
-        super().__init__(device)
-        self.cache_dir = cache_dir
-        self.model = self._load_model()
-        self.model.eval()
-        self.model.to(self.device)
-
-    @property
-    def feature_dim(self) -> int:
-        return 400
-
-    def _load_model(self) -> torch.nn.Module:
-        try:
-            model_path = hf_hub_download(repo_id=self.REPO_ID,
-                                         filename=self.MODEL_FILENAME,
-                                         cache_dir=self.cache_dir)
-            return torch.jit.load(model_path, map_location=self.device)
-        except Exception as e:
-            raise RuntimeError(f"Failed to load I3D model: {e}") from e
-
-    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
-        """Standard I3D preprocessing: Resize to 224, Norm to [-1, 1]."""
-        B, T, C, H, W = videos.shape
-
-        if T < 10:
-            raise ValueError(f"I3D requires at least 10 frames, got {T}")
-
-        # Normalize to [0, 1]
-        if videos.max() > 1.0:
-            videos = videos / 255.0
-
-        # Scale to [-1, 1]
-        videos = videos * 2.0 - 1.0
-
-        # Resize to 224x224
-        if H != 224 or W != 224:
-            videos = videos.reshape(B * T, C, H, W)
-            videos = F.interpolate(videos,
-                                   size=(224, 224),
-                                   mode='bilinear',
-                                   align_corners=False)
-            videos = videos.reshape(B, T, C, 224, 224)
-
-        # [B, T, C, H, W] -> [B, C, T, H, W]
-        return videos.permute(0, 2, 1, 3, 4).contiguous()
-
-    def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
-        batch = self.preprocess(videos)
-        # TorchScript I3D returns raw logits when return_features=True
-        return self.model(batch,
-                          rescale=False,
-                          resize=False,
-                          return_features=True)
-
-
-# 2. CLIP Extractor (Semantic/Content Quality)
-class CLIPFeatureExtractor(BaseFeatureExtractor):
-
-    def __init__(self,
-                 device: str = 'cuda',
-                 model_name: str = "openai/clip-vit-base-patch32"):
-        if not TRANSFORMERS_AVAILABLE:
-            raise ImportError(
-                "Please install transformers: pip install transformers")
-        super().__init__(device)
-        self.processor = CLIPProcessor.from_pretrained(model_name)
-        self.model = CLIPModel.from_pretrained(model_name).to(self.device)
-        self.model.eval()
-        self._feature_dim = self.model.config.projection_dim
-
-    @property
-    def feature_dim(self) -> int:
-        return self._feature_dim
-
-    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
-        # Ensure values are [0, 255]
-        if videos.max() <= 1.0:
-            videos = videos * 255.0
-
-        return videos.to(torch.uint8)
-
-    def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
-        # Input: [B, T, C, H, W]
-        B, T, C, H, W = videos.shape
-        videos = self.preprocess(videos)
-
-        # Flatten B*T to treat frames as images
-        images = videos.view(B * T, C, H, W)
-
-        # HF Processor
-        inputs = self.processor(images=images,
-                                return_tensors="pt",
-                                padding=True)
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}
-
-        # Extract features [B*T, Dim]
-        outputs = self.model.get_image_features(**inputs)
-
-        # Reshape [B, T, Dim] and Average Pooling over time
-        outputs = outputs.view(B, T, -1)
-        return outputs.mean(dim=1)
-
-
-# 3. VideoMAE Extractor (Structure/Motion Quality)
-class VideoMAEFeatureExtractor(BaseFeatureExtractor):
-
-    def __init__(self,
-                 device: str = 'cuda',
-                 model_name: str = "MCG-NJU/videomae-base"):
-        if not TRANSFORMERS_AVAILABLE:
-            raise ImportError(
-                "Please install transformers: pip install transformers")
-        super().__init__(device)
-        self.model = VideoMAEModel.from_pretrained(model_name).to(self.device)
-        self.model.eval()
-
-        self.register_buffer(
-            'mean',
-            torch.tensor([0.485, 0.456, 0.406],
-                         device=self.device).view(1, 1, 3, 1, 1))
-        self.register_buffer(
-            'std',
-            torch.tensor([0.229, 0.224, 0.225],
-                         device=self.device).view(1, 1, 3, 1, 1))
-
-    @property
-    def feature_dim(self) -> int:
-        return self.model.config.hidden_size
-
-    def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
-        """
-        Efficient GPU-based preprocessing.
-        Input: [B, T, C, H, W] in range [0, 255]
-        """
-        B, T, C, H, W = videos.shape
-
-        # 1. Resize to 224x224
-        if H != 224 or W != 224:
-            videos = videos.view(B * T, C, H, W)
-            videos = F.interpolate(videos,
-                                   size=(224, 224),
-                                   mode='bilinear',
-                                   align_corners=False)
-            videos = videos.view(B, T, C, 224, 224)
-
-        # 2. Normalize to [0, 1]
-        if videos.dtype != torch.float32:
-            videos = videos.float()
-
-        if videos.max() > 1.0:
-            videos = videos / 255.0
-
-        # 3. Apply ImageNet Mean/Std
-        return (videos - self.mean) / self.std
-
-    def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
-        # Input: [B, T, C, H, W]
-
-        # Fast GPU Preprocessing
-        pixel_values = self.preprocess(videos)
-
-        # Forward pass
-        outputs = self.model(pixel_values)
-
-        # Global Average Pooling of last hidden state [B, T_patches, 768] -> [B, 768]
-        return outputs.last_hidden_state.mean(dim=1)
-
-
-# Factory
-def load_extractor(name: str, device: str = 'cuda') -> BaseFeatureExtractor:
-    name = name.lower()
-    if name == 'i3d':
-        return I3DFeatureExtractor(device)
-    elif name == 'clip':
-        return CLIPFeatureExtractor(device)
-    elif name == 'videomae':
-        return VideoMAEFeatureExtractor(device)
-    else:
-        raise ValueError(
-            f"Unknown extractor: {name}. Options: i3d, clip, videomae")

benchmarks/fvd/fvd.py

@@ -5,7 +5,8 @@ from pathlib import Path
 from collections.abc import Iterator
 import pickle
 from dataclasses import dataclass, field
-from .feature_extractors import BaseFeatureExtractor, load_extractor
+
+from .i3d_model import I3DFeatureExtractor
 from .video_utils import ClipSamplingStrategy, load_video_clips_streaming
 
 
@@ -54,9 +55,6 @@ class FVDConfig:
     # Video selection
     num_videos: int = 2048
 
-    # Feature Extractor Selection
-    extractor_model: str = 'i3d'  # Options: 'i3d', 'clip', 'videomae'
-
     # Clip sampling
     num_frames_per_clip: int = 16
     num_clips_per_video: int = 1
@@ -87,7 +85,11 @@ class FVDConfig:
 
     @classmethod
     def fvd2048_16f(cls) -> 'FVDConfig':
-        """Standard FVD protocol: 2048 videos, 16 frames, beginning clip."""
+        """
+        Standard FVD protocol: 2048 videos, 16 frames, beginning clip.
+        
+        most common FVD configuration used in papers
+        """
         return cls(num_videos=2048,
                    num_frames_per_clip=16,
                    clip_strategy='beginning',
@@ -101,6 +103,18 @@ class FVDConfig:
                    clip_strategy='beginning',
                    use_streaming=True)
 
+    @classmethod
+    def fvd2048_128f_subsample8(cls) -> 'FVDConfig':
+        """
+        Long video with FPS subsampling: 2048 videos, 128 frames (every 8th).
+        Used for very long videos - samples every 8th frame
+        """
+        return cls(num_videos=2048,
+                   num_frames_per_clip=16,
+                   frame_stride=8,
+                   clip_strategy='beginning',
+                   use_streaming=True)
+
     @classmethod
     def quick_test(cls) -> 'FVDConfig':
         """Quick test config: 100 videos, 16 frames."""
@@ -110,13 +124,22 @@ class FVDConfig:
 
     def to_dict(self) -> dict:
         """Export config to dict for logging"""
-        d = self.__dict__.copy()
+        return {
-        d['clip_strategy'] = str(self.clip_strategy)
+            'num_videos': self.num_videos,
-        return d
+            'num_frames_per_clip': self.num_frames_per_clip,
+            'num_clips_per_video': self.num_clips_per_video,
+            'clip_strategy': str(self.clip_strategy),
+            'frame_stride': self.frame_stride,
+            'temporal_stride': self.temporal_stride,
+            'batch_size': self.batch_size,
+            'device': self.device,
+            'seed': self.seed,
+            'use_streaming': self.use_streaming,
+        }
 
     def __str__(self) -> str:
         """Human-readable protocol name"""
-        desc = f"FVD_{self.extractor_model.upper()}_{self.num_videos}_{self.num_frames_per_clip}f"
+        desc = f"FVD{self.num_videos}_{self.num_frames_per_clip}f"
         if self.frame_stride > 1:
             desc += f"_subsample{self.frame_stride}"
         if self.num_clips_per_video > 1:
@@ -127,42 +150,57 @@ class FVDConfig:
 
 
 def extract_features_streaming(video_generator: Iterator[torch.Tensor],
-                               extractor: BaseFeatureExtractor,
+                               extractor: I3DFeatureExtractor,
                                batch_size: int = 32,
                                max_clips: int | None = None,
                                verbose: bool = True) -> np.ndarray:
     """
     Extract features from a video clip generator using streaming.
+
+    Args:
+        video_generator: Iterator yielding clips [T, C, H, W]
+        extractor: I3D feature extractor
+        batch_size: Batch size for processing
+        max_clips: Maximum clips to process (for validation)
+        verbose: Show progress
+    
+    Returns:
+        features: [N, 400] numpy array
     """
     all_features = []
     batch = []
+    clip_count = 0
 
     if verbose:
         print(f"Extracting features with batch_size={batch_size}...")
 
-    with torch.no_grad():
+    for clip_count, clip in enumerate(video_generator):
-        for clip_count, clip in enumerate(video_generator):
+        batch.append(clip)
-            batch.append(clip)
 
-            # Process batch when full
+        # Process batch when full
-            if len(batch) == batch_size:
+        if len(batch) == batch_size:
-                batch_tensor = torch.stack(batch).to(extractor.device)
-                features = extractor.extract_features_batch(batch_tensor)
-
-                all_features.append(features.detach().cpu().numpy())
-                batch = []
-
-                if verbose and clip_count % (batch_size * 10) == 0:
-                    print(f"Processed {clip_count} clips...")
-
-            if max_clips is not None and clip_count >= max_clips:
-                break
-
-        # Process remaining clips
-        if len(batch) > 0:
             batch_tensor = torch.stack(batch).to(extractor.device)
-            features = extractor.extract_features_batch(batch_tensor)
+            features = extractor.extract_features(batch_tensor,
-            all_features.append(features.detach().cpu().numpy())
+                                                  batch_size=batch_size,
+                                                  verbose=False)
+            all_features.append(features.cpu().numpy())
+
+            batch = []  # Clear batch
+
+            if verbose and clip_count % (batch_size * 10) == 0:
+                print(f"Processed {clip_count} clips...")
+
+        # Stop if we've reached max_clips
+        if max_clips is not None and clip_count >= max_clips:
+            break
+
+    # Process remaining clips
+    if len(batch) > 0:
+        batch_tensor = torch.stack(batch).to(extractor.device)
+        features = extractor.extract_features(batch_tensor,
+                                              batch_size=len(batch),
+                                              verbose=False)
+        all_features.append(features.cpu().numpy())
 
     if len(all_features) == 0:
         raise RuntimeError("No features extracted - check video loading")
@@ -176,17 +214,14 @@ def extract_features_streaming(video_generator: Iterator[torch.Tensor],
 
 
 def load_or_compute_features(videos: str | Path | torch.Tensor,
-                             extractor: BaseFeatureExtractor,
+                             extractor: I3DFeatureExtractor,
                              config: FVDConfig,
                              cache_path: str | None = None,
                              cache_name: str = "real_features") -> np.ndarray:
     """Load features from cache or compute (with streaming support)"""
 
     if cache_path is not None:
-        script_dir = Path(__file__).parent
+        cache_file = Path(cache_path) / f"{cache_name}.pkl"
-        cache_dir = script_dir / cache_path
-        cache_file = cache_dir / f"{config.extractor_model}_{cache_name}.pkl"
-
         if cache_file.exists():
             print(f"Loading cached features from {cache_file}")
             with open(cache_file, 'rb') as f:
@@ -194,25 +229,19 @@ def load_or_compute_features(videos: str | Path | torch.Tensor,
 
             # Validate and limit based on config
             max_features = config.num_videos * config.num_clips_per_video
-
             if len(features) < max_features:
                 print(
                     f"WARNING: Cache has {len(features)} features but need {max_features}"
                 )
-                print("Cached features insufficient - will recompute...")
+                print("Recomputing features...")
             elif len(features) > max_features:
-                print(
-                    f"Using {max_features} features from cache (truncated from {len(features)})"
-                )
                 features = features[:max_features]
                 return features
             else:
-                print(f"Using all {len(features)} cached features")
                 return features
 
-    print("Computing features from scratch...")
+    # Compute features
-
+    if isinstance(videos, str | Path):
-    if isinstance(videos, (str | Path)):
         target_size = (224, 224) if config.resize_before_extraction else None
 
         video_generator = load_video_clips_streaming(
@@ -233,7 +262,9 @@ def load_or_compute_features(videos: str | Path | torch.Tensor,
                                               batch_size=config.batch_size,
                                               max_clips=max_clips,
                                               verbose=True)
+
     else:
+        # Already a tensor
         print(f"Extracting features from {len(videos)} video tensors...")
         features = extractor.extract_features(videos,
                                               batch_size=config.batch_size,
@@ -252,10 +283,9 @@ def load_or_compute_features(videos: str | Path | torch.Tensor,
 
     # Cache features if requested
     if cache_path is not None:
-        script_dir = Path(__file__).parent
+        cache_dir = Path(cache_path)
-        cache_dir = script_dir / cache_path
         cache_dir.mkdir(parents=True, exist_ok=True)
-        cache_file = cache_dir / f"{config.extractor_model}_{cache_name}.pkl"
+        cache_file = cache_dir / f"{cache_name}.pkl"
         print(f"Caching features to {cache_file}")
         with open(cache_file, 'wb') as f:
             pickle.dump(features, f)
@@ -263,34 +293,79 @@ def load_or_compute_features(videos: str | Path | torch.Tensor,
     return features
 
 
+def compute_fvd(real_videos: str | Path | torch.Tensor,
+                gen_videos: str | Path | torch.Tensor,
+                num_frames: int = 16,
+                batch_size: int = 32,
+                device: str = 'cuda',
+                num_videos: int | None = 2048,
+                cache_real_features: str | None = None,
+                i3d_model_path: str | None = None,
+                seed: int | None = None,
+                verbose: bool = True) -> float:
+    """
+    Compute Fréchet Video Distance (FVD)
+    
+    For advanced control, use compute_fvd_with_config() instead.
+    
+    Args:
+        real_videos: Path to real videos or tensor [N, T, C, H, W]
+        gen_videos: Path to generated videos or tensor [N, T, C, H, W]
+        num_frames: Frames per video (default: 16)
+        batch_size: Batch size (default: 32)
+        device: 'cuda' or 'cpu' (default: 'cuda')
+        num_videos: Max videos (default: 2048)
+        cache_real_features: Cache path for real features
+        i3d_model_path: Custom I3D model cache path
+        seed: Random seed for reproducibility
+        verbose: Print progress
+    
+    Returns:
+        FVD score (float). Lower is better.
+    """
+    num_videos = num_videos if num_videos is not None else 2048
+
+    config = FVDConfig(
+        num_videos=num_videos,
+        num_frames_per_clip=num_frames,
+        batch_size=batch_size,
+        device=device,
+        cache_real_features=cache_real_features,
+        i3d_model_path=i3d_model_path,
+        seed=seed,
+    )
+
+    result = compute_fvd_with_config(real_videos, gen_videos, config, verbose)
+    return result['fvd']
+
+
 def compute_fvd_with_config(real_videos: str | Path | torch.Tensor,
                             gen_videos: str | Path | torch.Tensor,
                             config: FVDConfig,
                             verbose: bool = True) -> dict:
     """
-        Compute FVD using a standardized configuration.
+    Compute FVD using a standardized configuration.
-
+    
-        This is the recommended way to compute FVD for reproducibility.
+    This is the recommended way to compute FVD for reproducibility.
-
+    
-        Args:
+    Args:
-            real_videos: Path or tensors
+        real_videos: Path or tensors
-            gen_videos: Path or tensors
+        gen_videos: Path or tensors
-            config: FVDConfig specifying protocol
+        config: FVDConfig specifying protocol
-            verbose: Print progress
+        verbose: Print progress
-
+    
-        Returns:
+    Returns:
-            results: Dictionary with:
+        results: Dictionary with:
-                - 'fvd': FVD score (float)
+            - 'fvd': FVD score (float)
-                - 'protocol': Protocol name (str)
+            - 'protocol': Protocol name (str)
-                - 'model': Feature extractor model name (str)
+            - 'config': Configuration dict
-                - 'config': Configuration dict
+    
-
+    Example:
-        Example:
+        >>> config = FVDConfig.fvd2048_16f()
-            >>> config = FVDConfig.fvd2048_16f()
+        >>> results = compute_fvd_with_config('data/real/', 'outputs/gen/', config)
-            >>> results = compute_fvd_with_config('data/real/', 'outputs/gen/', config)
+        >>> print(f"FVD: {results['fvd']:.2f}")
-            >>> print(f"FVD: {results['fvd']:.2f}")
+        >>> print(f"Protocol: {results['protocol']}")  # "FVD2048_16f"
     """
-
     # Seed for reproducibility
     if config.seed is not None:
         import random as _rnd
@@ -303,20 +378,18 @@ def compute_fvd_with_config(real_videos: str | Path | torch.Tensor,
     if verbose:
         print("=" * 70)
         print(f"Computing FVD with protocol: {config}")
-        print(f"Model: {config.extractor_model.upper()}")
         print("=" * 70)
         print("\nConfiguration:")
         for key, value in config.to_dict().items():
             print(f"  {key}: {value}")
         print()
 
-    # Initialize Extractor using Factory
+    # Initialize I3D
     if verbose:
-        print(
+        print(f"\nInitializing I3D model on {config.device}...")
-            f"\nInitializing {config.extractor_model.upper()} model on {config.device}..."
-        )
 
-    extractor = load_extractor(config.extractor_model, device=config.device)
+    extractor = I3DFeatureExtractor(device=config.device,
+                                    cache_dir=config.i3d_model_path)
 
     # Extract features
     if verbose:
@@ -361,45 +434,14 @@ def compute_fvd_with_config(real_videos: str | Path | torch.Tensor,
 
     if verbose:
         print(f"\n{'='*70}")
-        print(f"FVD Score ({config.extractor_model.upper()}): {fvd:.4f}")
+        print(f"FVD Score: {fvd:.4f}")
         print(f"Protocol: {config}")
         print(f"{'='*70}\n")
 
     results = {
         'fvd': fvd,
         'protocol': str(config),
-        'model': config.extractor_model,
         'config': config.to_dict(),
     }
 
     return results
-
-
-def compute_fvd(real_videos: str | Path | torch.Tensor,
-                gen_videos: str | Path | torch.Tensor,
-                num_frames: int = 16,
-                batch_size: int = 32,
-                device: str = 'cuda',
-                num_videos: int | None = 2048,
-                cache_real_features: str | None = None,
-                i3d_model_path: str | None = None,
-                seed: int | None = None,
-                verbose: bool = True) -> float:
-    """
-    Backward compatibility wrapper for computing FVD (defaults to I3D).
-    """
-    num_videos = num_videos if num_videos is not None else 2048
-
-    config = FVDConfig(
-        num_videos=num_videos,
-        num_frames_per_clip=num_frames,
-        extractor_model='i3d',  # Default to I3D
-        batch_size=batch_size,
-        device=device,
-        cache_real_features=cache_real_features,
-        i3d_model_path=i3d_model_path,
-        seed=seed,
-    )
-
-    result = compute_fvd_with_config(real_videos, gen_videos, config, verbose)
-    return result['fvd']

benchmarks/fvd/run_fvd.py

@@ -1,53 +1,33 @@
 import sys
 from pathlib import Path
+from benchmarks.fvd.fvd import FVDConfig, compute_fvd_with_config
 
 root_dir = Path(__file__).parent.parent.parent
 sys.path.insert(0, str(root_dir))
 
-from benchmarks.fvd.fvd import FVDConfig, compute_fvd_with_config  # noqa: E402
-
 
 def main() -> None:
+    # Get script directory
     script_dir = Path(__file__).parent.resolve()
 
-    # Define directories
+    clip_strategy = 'beginning'  # Options: 'uniform', 'random', 'beginning', 'end', 'all'
+    cfg = FVDConfig(
+        num_videos=650,
+        num_frames_per_clip=16,
+        num_clips_per_video=1,
+        clip_strategy=clip_strategy,
+        frame_stride=1,
+        batch_size=32,
+        device='cuda',
+        seed=42,
+        cache_real_features=str(script_dir / f'fvd-cache/{clip_strategy}'),
+    )
+
     real_dir = "benchmarks/data/real_videos"
     gen_dir = "benchmarks/data/generated_videos"
 
-    # Compare all 3 models
+    results = compute_fvd_with_config(real_dir, gen_dir, cfg, verbose=True)
-    models_to_test = ['i3d', 'clip', 'videomae']
+    print(f"FVD = {results['fvd']:.2f}")
-
-    print(f"\n{'='*60}")
-    print("STARTING COMPARISON BENCHMARK")
-    print(f"{'='*60}")
-
-    for model_name in models_to_test:
-        print(f"\n>>> Running evaluation with {model_name.upper()}...")
-
-        try:
-            cfg = FVDConfig(
-                num_videos=650,
-                num_frames_per_clip=16,
-                extractor_model=model_name,
-                clip_strategy='beginning',
-                device='cuda',
-                seed=42,
-                # Use separate cache folders for each model to avoid conflicts
-                cache_real_features=str(script_dir / f'fvd-cache/{model_name}'),
-            )
-
-            results = compute_fvd_with_config(real_dir,
-                                              gen_dir,
-                                              cfg,
-                                              verbose=False)
-            print(f"FVD: {results['fvd']}\nModel: {results['model']}")
-
-        except Exception as e:
-            print(f"{model_name.upper()} Failed: {e}")
-
-    print(f"\n{'='*60}")
-    print("BENCHMARK COMPLETE")
-    print(f"{'='*60}")
 
 
 if __name__ == '__main__':

benchmarks/scripts/run.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 # 1. Install missing dependency
-pip install -q opencv-python-headless transformers huggingface_hub
+pip install -q opencv-python-headless
 
 # 2. Run FVD script
 python benchmarks/fvd/run_fvd.py

csrc/attn/README.md

@@ -0,0 +1,113 @@
+
+
+# Attention Kernel Used in FastVideo
+
+
+
+## Video Sparse Attention (VSA)
+
+### Installation
+We support H100 (via TK) and any other GPU (via triton) for VSA.
+```bash
+git submodule update --init --recursive
+python setup_vsa.py install
+```
+
+
+If you encounter error during installation, try below:
+Install C++20 for ThunderKittens:
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+(If you use CUDA12.8)
+```bash
+export CUDA_HOME=/usr/local/cuda-12.8
+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
+End-2-end inference with FastVideo:
+```bash
+bash scripts/inference/v1_inference_wan_STA.sh
+```
+
+If you want to use sliding tile attention in your custom model:
+```python
+from st_attn import sliding_tile_attention
+# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
+# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256
+# a tile is a cube of size (6, 8, 8)
+# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video.
+# text_length: int ranging from 0 to 256
+# If your attention contains text token (Hunyuan)
+out = sliding_tile_attention(q, k, v, window_size, text_length)
+# If your attention does not contain text token (StepVideo)
+out = sliding_tile_attention(q, k, v, window_size, 0, False)
+```
+
+
+### Test
+```bash
+python tests/test_sta.py # test STA
+python tests/test_vsa.py # test VSA
+```
+### Benchmark
+```bash
+python benchmarks/bench_sta.py
+```
+
+
+### How Does STA Work?
+We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 
+
+
+https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747
+
+## Why is STA Fast?
+2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. 
+
+STA removes mixed blocks.
+
+
+<div align="center">
+<img src=../../assets/sliding_tile_attn_map.png width="80%"/>
+</div>
+
+## Acknowledgement
+
+We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/attn/benchmarks/bench_sta.py

@@ -0,0 +1,145 @@
+import os
+from collections import defaultdict
+
+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"):
+    assert mode in ["fwd", "bwd", "fwd_bwd"]
+    f = 4 * batch * seqlen**2 * nheads * headdim // (2 if causal else 1)
+    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 benchmark_attention(configurations):
+    results = {'fwd': defaultdict(list), 'bwd': defaultdict(list)}
+
+    for B, H, N, D, causal, dit_seq_shape, window_size in configurations:
+        print("=" * 60)
+        print(f"Timing forward and backward pass for B={B}, H={H}, N={N}, D={D}, causal={causal}")
+
+        q = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
+        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()
+
+
+        # # Warmup for forward pass
+        # for _ in range(10):
+        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
+
+        # # 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))
+
+        # 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 = compute_TFLOPS(flops(B, N, H, D, causal, 'fwd'), ms)
+        results['fwd'][(D, causal)].append((N, tflops_fwd))
+
+        print(f"Average time for forward pass (ms): {ms:.2f}")
+        print(f"Average TFLOPS: {tflops_fwd}")
+        print("-" * 60)
+
+        # torch.cuda.empty_cache()
+        # torch.cuda.synchronize()
+
+        # # Prepare for timing backward pass
+        # start_events_bwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+        # end_events_bwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+
+        # # Warmup for backward pass
+        # for _ in range(10):
+        #     qg, kg, vg = tk.mha_backward(q, k, v, o, l_vec, grad_output, causal)
+
+        # # Time the backward pass
+        # for i in range(10):
+        #     start_events_bwd[i].record()
+        #     qg, kg, vg = tk.mha_backward(q, k, v, o, l_vec, grad_output, causal)
+        #     end_events_bwd[i].record()
+
+        # torch.cuda.synchronize()
+        # 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)
+        # 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)
+
+    return results
+
+
+def plot_results(results):
+    os.makedirs('benchmark_results', exist_ok=True)
+    for mode in ['fwd', 'bwd']:
+        for (D, causal), values in results[mode].items():
+            seq_lens = [x[0] for x in values]
+            tflops = [x[1] for x in values]
+
+            plt.figure(figsize=(10, 6))
+            bars = plt.bar(range(len(seq_lens)), tflops, tick_label=seq_lens)
+            plt.xlabel('Sequence Length')
+            plt.ylabel('TFLOPS')
+            plt.title(f'{mode.upper()} Pass - Head Dim: {D}, Causal: {causal}')
+            plt.grid(True)
+
+            # Adding the numerical y value on top of each bar
+            for bar in bars:
+                yval = bar.get_height()
+                plt.text(bar.get_x() + bar.get_width() / 2, yval, round(yval, 2), ha='center', va='bottom')
+
+            filename = f'benchmark_results/{mode}_D{D}_causal{causal}.png'
+            plt.savefig(filename)
+            plt.close()
+
+
+# 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]),
+    # (16, 16, 768*16,  128, False),
+    # (16, 16, 768*2,  128, False),
+    # (16, 16, 768*4,  128, False),
+    # (16, 16, 768*8,  128, False),
+    # (16, 16, 768*16, 128, False),
+    # (16, 16, 768,    128, True),
+    # (16, 16, 768*2,  128, True),
+    # (16, 16, 768*4,  128, True),
+    # (16, 16, 768*8,  128, True),
+    # (16, 16, 768*16, 128, True),
+    # (16, 32, 768,    64,  False),
+    # (16, 32, 768*2,  64,  False),
+    # (16, 32, 768*4,  64,  False),
+    # (16, 32, 768*8,  64,  False),
+    # (16, 32, 768*16, 64,  False),
+    # (16, 32, 768,    64,  True),
+    # (16, 32, 768*2,  64,  True),
+    # (16, 32, 768*4,  64,  True),
+    # (16, 32, 768*8,  64,  True),
+    # (16, 32, 768*16, 64,  True),
+]
+
+results = benchmark_attention(configurations)
+# plot_results(results)

csrc/attn/benchmarks/bench_vsa_hopper.py

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

csrc/attn/benchmarks/bench_vsa_triton.py

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

csrc/attn/sliding_tile_attn/MANIFEST.in

@@ -0,0 +1,2 @@
+recursive-include tk *
+include config_sta.py

csrc/attn/sliding_tile_attn/README.md

@@ -0,0 +1,103 @@
+
+# Attention Kernel Used in FastVideo
+
+## Sliding Tile Attention (STA)
+We support H100 (via TK) and any other GPU (via triton) for STA.
+
+### Installation
+```bash
+pip install st_attn
+```
+
+Install from source:
+
+```bash
+git submodule update --init --recursive
+python setup.py install
+```
+
+If you want to skip the compilation of the TK kernel and only use the Triton version, try below:
+
+```bash
+SKIP_SM90_EXT=1 python setup.py install
+or
+SKIP_SM90_EXT=1 pip install --no-build-isolation .
+```
+
+If you encounter error during installation, try below:
+Install C++20 for ThunderKittens:
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+(If you use CUDA12.8)
+```bash
+export CUDA_HOME=/usr/local/cuda-12.8
+export PATH=${CUDA_HOME}/bin:${PATH}
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+###  Usage
+End-2-end inference with FastVideo:
+```bash
+bash scripts/inference/v1_inference_wan_STA.sh
+```
+
+If you want to use sliding tile attention in your custom model:
+```python
+from st_attn import sliding_tile_attention
+# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding.
+# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256
+# a tile is a cube of size (6, 8, 8)
+# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video.
+# text_length: int ranging from 0 to 256
+# If your attention contains text token (Hunyuan)
+out = sliding_tile_attention(q, k, v, window_size, text_length)
+# If your attention does not contain text token (StepVideo)
+out = sliding_tile_attention(q, k, v, window_size, 0, False)
+```
+
+
+### Test
+```bash
+python ../tests/test_sta.py # test STA
+```
+### 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.
+
+
+https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747
+
+
+## STA Configuration Logic
+Here is a diagram of how the window is configured and passed through the FastVideo pipeline:
+
+<div align="center">
+<img src="../../../docs/assets/images/STA_configuration.png" width="80%"/>
+</div>
+
+
+## Why is STA Fast?
+2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation.
+
+STA removes mixed blocks.
+
+
+<div align="center">
+<img src=../../../assets/sliding_tile_attn_map.png width="80%"/>
+</div>
+
+## Acknowledgement
+
+We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/attn/sliding_tile_attn/config_sta.py

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

csrc/attn/sliding_tile_attn/setup.py

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

csrc/attn/sliding_tile_attn/st_attn.cpp

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

csrc/attn/sliding_tile_attn/st_attn/__init__.py

@@ -0,0 +1,63 @@
+import math
+
+import torch
+from torch.utils.checkpoint import detach_variable
+try:
+    from st_attn_cuda import sta_fwd
+except ImportError:
+    sta_fwd = None
+
+try:
+    from st_attn.st_attn_triton import sliding_tile_attention_triton
+except ImportError:
+    sliding_tile_attention_triton = None
+
+def sliding_tile_attention_SM90(q_all, k_all, v_all, window_size, text_length, has_text=True, dit_seq_shape='30x48x80'):
+    seq_length = q_all.shape[2]
+    dit_seq_shape_mapping = {
+        '30x48x80':1,
+        '36x48x48':2,
+        '18x48x80':3,
+    }
+    if has_text:
+        assert q_all.shape[
+            2] >= 115200 and q_all.shape[2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '30x48x80' for HunyuanVideo"
+        assert q_all.shape[1] == len(window_size), "Number of heads must match the number of window sizes"
+        target_size = math.ceil(seq_length / 384) * 384
+        pad_size = target_size - seq_length
+        if pad_size > 0:
+            q_all = torch.cat([q_all, q_all[:, :, -pad_size:]], dim=2)
+            k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2)
+            v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2)
+    else:
+        if dit_seq_shape == '36x48x48': # Stepvideo 204x768x68
+            assert q_all.shape[2] == 82944
+        elif dit_seq_shape == '18x48x80': # Wan 69x768x1280
+            assert q_all.shape[2] == 69120
+        else:
+            raise ValueError(f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
+
+    kernel_aspect_ratio_flag = dit_seq_shape_mapping[dit_seq_shape]
+    hidden_states = torch.empty_like(q_all)
+    # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
+    for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
+        for batch in range(q_all.shape[0]):
+            q_head, k_head, v_head, o_head = (q_all[batch:batch + 1, head_index:head_index + 1],
+                                              k_all[batch:batch + 1,
+                                                    head_index:head_index + 1], v_all[batch:batch + 1,
+                                                                                      head_index:head_index + 1],
+                                              hidden_states[batch:batch + 1, head_index:head_index + 1])
+
+            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text, kernel_aspect_ratio_flag)
+    if has_text:
+        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True, kernel_aspect_ratio_flag)
+    return hidden_states[:, :, :seq_length]
+
+def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, dit_seq_shape='30x48x80'):
+    major, minor = torch.cuda.get_device_capability(q_all.device)
+    if major == 9 and minor == 0 and sta_fwd is not None:
+        return sliding_tile_attention_SM90(q_all, k_all, v_all, window_size, text_length, has_text, dit_seq_shape)
+    elif sliding_tile_attention_triton is not None:
+        return sliding_tile_attention_triton(q_all, k_all, v_all, window_size, text_length, has_text, dit_seq_shape)
+    else:
+        raise ImportError("No suitable sliding tile attention implementation found.")

csrc/attn/sliding_tile_attn/st_attn/st_attn_h100.cu

@@ -0,0 +1,841 @@
+// # Define TORCH_COMPILE macro
+
+#include "kittens.cuh"
+#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;
+}
+
+
+constexpr int CONSUMER_WARPGROUPS = (3); 
+constexpr int PRODUCER_WARPGROUPS = (1); 
+constexpr int NUM_WARPGROUPS      = (CONSUMER_WARPGROUPS+PRODUCER_WARPGROUPS); 
+constexpr int NUM_WORKERS         = (NUM_WARPGROUPS*kittens::WARPGROUP_WARPS); 
+
+using namespace kittens;
+namespace cg = cooperative_groups;
+
+template<int D> struct fwd_attend_ker_tile_dims {};
+template<> struct fwd_attend_ker_tile_dims<64> {
+    constexpr static int tile_width = (64);
+    constexpr static int qo_height  = (4*16);
+    constexpr static int kv_height  = (8*16);
+    constexpr static int stages     = (4); 
+};
+template<> struct fwd_attend_ker_tile_dims<128> {
+    constexpr static int tile_width = (128);
+    constexpr static int qo_height  = (4*16);
+    constexpr static int kv_height  = (8*16);
+    constexpr static int stages     = (2); 
+};
+
+template<int D> struct fwd_globals {
+    using q_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::qo_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+    using k_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::kv_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+    using v_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::kv_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+    using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<D>::qo_height, fwd_attend_ker_tile_dims<D>::tile_width>>;
+    using o_tile    =         st_bf<fwd_attend_ker_tile_dims<D>::qo_height, fwd_attend_ker_tile_dims<D>::tile_width>;
+
+    using q_gl = gl<bf16,  -1, -1, -1, -1, q_tile>;
+    using k_gl = gl<bf16,  -1, -1, -1, -1, k_tile>;
+    using v_gl = gl<bf16,  -1, -1, -1, -1, v_tile>;
+    using l_gl = gl<float, -1, -1, -1, -1, l_col_vec>;
+    using o_gl = gl<bf16,  -1, -1, -1, -1, o_tile>;
+
+    q_gl q;
+    k_gl k;
+    v_gl v;
+    l_gl l;
+    o_gl o;
+
+    const int N; 
+    const int text_L;
+    const int hr;
+};
+
+
+template<int D, bool is_causal, bool text_q, bool text_kv, int DT, int DH, int DW, int CT, int CH, int CW>
+__global__  __launch_bounds__((NUM_WORKERS)*kittens::WARP_THREADS, 1)
+void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
+    extern __shared__ int __shm[]; 
+    tma_swizzle_allocator al((int*)&__shm[0]);
+    int warpid = kittens::warpid(), warpgroupid = warpid/kittens::WARPGROUP_WARPS;
+
+    using K = fwd_attend_ker_tile_dims<D>;
+
+    using q_tile    =         st_bf<K::qo_height, K::tile_width>;
+    using k_tile    =         st_bf<K::kv_height, K::tile_width>;
+    using v_tile    =         st_bf<K::kv_height, K::tile_width>;
+    using l_col_vec = col_vec<st_fl<K::qo_height, K::tile_width>>;
+    using o_tile    =         st_bf<K::qo_height, K::tile_width>;
+    
+    q_tile    (&q_smem)[CONSUMER_WARPGROUPS] = al.allocate<q_tile, CONSUMER_WARPGROUPS>();
+    k_tile    (&k_smem)[K::stages]           = al.allocate<k_tile, K::stages          >();
+    v_tile    (&v_smem)[K::stages]           = al.allocate<v_tile, K::stages          >();
+    l_col_vec (&l_smem)[CONSUMER_WARPGROUPS] = al.allocate<l_col_vec, CONSUMER_WARPGROUPS>();
+    auto      (*o_smem)                      = reinterpret_cast<o_tile(*)>(q_smem);
+    int img_kv_blocks;
+    int kv_blocks   = g.N / (K::kv_height);
+    if constexpr (text_kv) {
+        img_kv_blocks = kv_blocks - 3;
+    } else {
+        img_kv_blocks = kv_blocks;
+    }
+    int kv_head_idx = blockIdx.y / g.hr;
+    int seq_idx;
+    if constexpr (text_q) {
+        seq_idx = CT * CH * CW * 6.0 + blockIdx.x * CONSUMER_WARPGROUPS;
+    } else {
+        seq_idx = blockIdx.x * CONSUMER_WARPGROUPS; 
+    }
+    __shared__ kittens::semaphore qsmem_semaphore, k_smem_arrived[K::stages], v_smem_arrived[K::stages], compute_done[K::stages];
+    if (threadIdx.x == 0) { 
+        init_semaphore(qsmem_semaphore, 0, 1); 
+        for(int j = 0; j < K::stages; j++) {
+            init_semaphore(k_smem_arrived[j], 0, 1); 
+            init_semaphore(v_smem_arrived[j], 0, 1); 
+            init_semaphore(compute_done[j], CONSUMER_WARPGROUPS, 0); 
+        }
+
+        tma::expect_bytes(qsmem_semaphore, sizeof(q_smem));
+
+        for (int wg = 0; wg < CONSUMER_WARPGROUPS; wg++) {
+            coord<q_tile> q_tile_idx = {blockIdx.z, blockIdx.y, (seq_idx) + wg, 0};
+            tma::load_async(q_smem[wg], g.q, q_tile_idx, qsmem_semaphore);
+        }
+
+        if constexpr (text_q){
+            for (int j = 0; j < K::stages - 1; j++) {
+                coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0};
+                tma::expect_bytes(k_smem_arrived[j], sizeof(k_tile));
+                tma::load_async(k_smem[j], g.k, kv_tile_idx, k_smem_arrived[j]);
+                tma::expect_bytes(v_smem_arrived[j], sizeof(v_tile));
+                tma::load_async(v_smem[j], g.v, kv_tile_idx, v_smem_arrived[j]);
+            }
+        } else {
+            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 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);
+                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));
+                    tma::load_async(k_smem[count], g.k, kv_tile_idx, k_smem_arrived[count]);
+                    tma::expect_bytes(v_smem_arrived[count], sizeof(v_tile));
+                    tma::load_async(v_smem[count], g.v, kv_tile_idx, v_smem_arrived[count]);
+                    count += 1;
+                }
+                j += 1;
+            }
+        }
+    }
+    __syncthreads(); 
+
+    int pipe_idx = K::stages - 1; 
+    
+    if(warpgroupid == NUM_WARPGROUPS-1) {
+        warpgroup::decrease_registers<32>();      
+        
+        int kv_iters; 
+        if constexpr (is_causal) {
+            kv_iters = (seq_idx * (K::qo_height/kittens::TILE_ROW_DIM<bf16>)) - 1 + (CONSUMER_WARPGROUPS * (K::qo_height/kittens::TILE_ROW_DIM<bf16>)); 
+            kv_iters = ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM<bf16>)) == 0) ? (0) : ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM<bf16>)) - 1);
+        }
+        else { kv_iters = kv_blocks-2;}
+
+        if(warpid == NUM_WORKERS-4) {
+            if constexpr (text_q){
+                for (auto kv_idx = pipe_idx - 1; kv_idx <= kv_iters; kv_idx++) {
+                    coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, kv_idx + 1, 0};
+                    tma::expect_bytes(k_smem_arrived[(kv_idx+1)%K::stages], sizeof(k_tile));
+                    tma::load_async(k_smem[(kv_idx+1)%K::stages], g.k, kv_tile_idx, k_smem_arrived[(kv_idx+1)%K::stages]);
+                    tma::expect_bytes(v_smem_arrived[(kv_idx+1)%K::stages], sizeof(v_tile));
+                    tma::load_async(v_smem[(kv_idx+1)%K::stages], g.v, kv_tile_idx, v_smem_arrived[(kv_idx+1)%K::stages]);
+                    kittens::wait(compute_done[(kv_idx)%K::stages], (kv_idx/K::stages)%2);
+                }
+            } else {
+                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);
+                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++) {
+                        for (int kw = k_w_min; kw <= k_w_max; kw++) {
+                            for (int j = 0; j <= 2; j++){
+                                if (count >= K::stages - 1) {
+                                    int index = ((kt * (CH * CW)) + (kh * CW) + kw) * 3 + j;
+                                    coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, index, 0};
+                                    tma::expect_bytes(k_smem_arrived[count%K::stages], sizeof(k_tile));
+                                    tma::load_async(k_smem[count%K::stages], g.k, kv_tile_idx, k_smem_arrived[count%K::stages]);
+                                    tma::expect_bytes(v_smem_arrived[count%K::stages], sizeof(v_tile));
+                                    tma::load_async(v_smem[count%K::stages], g.v, kv_tile_idx, v_smem_arrived[count%K::stages]);
+                                    kittens::wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2);
+                                    count += 1;
+                                } else {
+                                    count += 1;
+                                }
+                            }
+                        }
+                    }
+                }
+                // for text 
+                for (int index = img_kv_blocks; index < kv_blocks; index++) {
+                    coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, index, 0};
+                    tma::expect_bytes(k_smem_arrived[count%K::stages], sizeof(k_tile));
+                    tma::load_async(k_smem[count%K::stages], g.k, kv_tile_idx, k_smem_arrived[count%K::stages]);
+                    tma::expect_bytes(v_smem_arrived[count%K::stages], sizeof(v_tile));
+                    tma::load_async(v_smem[count%K::stages], g.v, kv_tile_idx, v_smem_arrived[count%K::stages]);
+                    kittens::wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2);
+                    count += 1;
+                }
+            }
+
+
+        }
+    }
+    else {
+        warpgroup::increase_registers<160>();
+
+        rt_fl<16, K::kv_height>  att_block;
+        rt_bf<16, K::kv_height>  att_block_mma;
+        rt_fl<16, K::tile_width> o_reg;
+        
+        col_vec<rt_fl<16, K::kv_height>> max_vec, norm_vec, max_vec_last_scaled, max_vec_scaled;
+        
+        neg_infty(max_vec);
+        zero(norm_vec);
+        zero(o_reg);
+
+        int kv_iters; 
+        if constexpr (is_causal) {
+            kv_iters = (seq_idx * 4) - 1 + (CONSUMER_WARPGROUPS * 4);
+            kv_iters = (kv_iters/8);
+        }
+        else if constexpr (text_q){ 
+            // 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 ; 
+        }
+
+        kittens::wait(qsmem_semaphore, 0);
+        for (auto kv_idx = 0; kv_idx <= kv_iters; kv_idx++) {
+
+            kittens::wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2);
+            warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[(kv_idx)%K::stages]);
+            
+            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();
+
+            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); 
+
+            kittens::wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); 
+
+            warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]);
+            warpgroup::mma_async_wait();
+
+            if(warpgroup::laneid() == 0) arrive(compute_done[(kv_idx)%K::stages], 1);
+        }
+        // the last three kv blocks are for text, we process them separately
+        if constexpr(text_kv) {
+            for (auto kv_idx = kv_iters + 1; kv_idx <= kv_iters + 3; kv_idx++) {
+
+                kittens::wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2);
+                warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[(kv_idx)%K::stages]);
+                
+                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();
+                // apply non-pad mask
+                int offset = g.text_L - (kv_idx - (kv_iters + 1)) * K::kv_height;
+                // printf("k_idx_start: %d, k_idx_end: %d, text_end: %d, offset: %d\n", k_idx_start, k_idx_end, text_end, offset);
+                right_fill(att_block, att_block, offset, base_types::constants<float>::neg_infty());
+
+
+                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); 
+
+                kittens::wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); 
+
+                warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]);
+                warpgroup::mma_async_wait();
+
+                if(warpgroup::laneid() == 0) arrive(compute_done[(kv_idx)%K::stages], 1);
+            }
+        }
+
+        div_row(o_reg, o_reg, norm_vec);
+        warpgroup::store(o_smem[warpgroupid], o_reg); 
+        warpgroup::sync(warpgroupid+4);
+
+        if (warpid % 4 == 0) {
+            coord<o_tile> o_tile_idx = {blockIdx.z, blockIdx.y, (seq_idx) + warpgroupid, 0};
+            tma::store_async(g.o, o_smem[warpgroupid], 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[warpgroupid], norm_vec);
+        warpgroup::sync(warpgroupid+4);
+
+        if (warpid % 4 == 0) {
+            coord<l_col_vec> tile_idx = {blockIdx.z, blockIdx.y, 0, (seq_idx) + warpgroupid};
+            tma::store_async(g.l, l_smem[warpgroupid], tile_idx);
+        }
+        tma::store_async_wait();
+    }
+}
+
+
+
+#include "pyutils/torch_helpers.cuh"
+#include <ATen/cuda/CUDAContext.h>
+#include <iostream>
+
+torch::Tensor 
+sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag)
+{
+    CHECK_INPUT(q);
+    CHECK_INPUT(k);
+    CHECK_INPUT(v);
+
+    auto batch    = q.size(0);
+    auto seq_len  = q.size(2); 
+    auto head_dim = q.size(3);  
+    auto qo_heads = q.size(1);
+    auto kv_heads = k.size(1);
+
+    // 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");
+    TORCH_CHECK(v.size(0) == batch, "V batch dimension - idx 0 - must match for all inputs");
+
+    TORCH_CHECK(q.size(2) == seq_len, "Q sequence length dimension - idx 2 - must match for all inputs");
+    TORCH_CHECK(k.size(2) == seq_len, "K sequence length dimension - idx 2 - must match for all inputs");
+    TORCH_CHECK(v.size(2) == seq_len, "V sequence length dimension - idx 2 - must match for all inputs");
+
+    TORCH_CHECK(q.size(3) == head_dim, "Q head dimension - idx 3 - must match for all non-vector inputs");
+    TORCH_CHECK(k.size(3) == head_dim, "K head dimension - idx 3 - must match for all non-vector inputs");
+    TORCH_CHECK(v.size(3) == head_dim, "V head dimension - idx 3 - must match for all non-vector inputs");
+
+    TORCH_CHECK(qo_heads >= kv_heads, "QO heads must be greater than or equal to KV heads");
+    TORCH_CHECK(qo_heads % kv_heads == 0, "QO heads must be divisible by KV heads");
+    TORCH_CHECK(q.size(1) == qo_heads, "QO head dimension - idx 1 - must match for all inputs");
+    TORCH_CHECK(k.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs");
+    TORCH_CHECK(v.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs");  
+
+    auto hr = qo_heads / kv_heads;
+
+    c10::BFloat16* q_ptr = q.data_ptr<c10::BFloat16>();
+    c10::BFloat16* k_ptr = k.data_ptr<c10::BFloat16>();
+    c10::BFloat16* v_ptr = v.data_ptr<c10::BFloat16>();
+
+    bf16*  d_q = reinterpret_cast<bf16*>(q_ptr);
+    bf16*  d_k = reinterpret_cast<bf16*>(k_ptr);
+    bf16*  d_v = reinterpret_cast<bf16*>(v_ptr);
+    
+
+    
+    torch::Tensor l_vec = torch::empty({static_cast<const uint>(batch), 
+                                        static_cast<const uint>(qo_heads), 
+                                        static_cast<const uint>(seq_len), 
+                                        static_cast<const uint>(1)}, 
+                                        torch::TensorOptions().dtype(torch::kFloat).device(q.device()).memory_format(at::MemoryFormat::Contiguous));
+        
+
+    bf16*  o_ptr = reinterpret_cast<bf16*>(o.data_ptr<c10::BFloat16>());
+    bf16*  d_o   = reinterpret_cast<bf16*>(o_ptr);
+
+    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(); 
+
+
+    if (head_dim == 128) {
+        using q_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using k_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using v_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>>;
+        using o_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+
+        using q_global = gl<bf16,  -1, -1, -1, -1, q_tile>;
+        using k_global = gl<bf16,  -1, -1, -1, -1, k_tile>;
+        using v_global = gl<bf16,  -1, -1, -1, -1, v_tile>;
+        using l_global = gl<float, -1, -1, -1, -1, l_col_vec>;
+        using o_global = gl<bf16,  -1, -1, -1, -1, o_tile>;
+
+        using globals      = fwd_globals<128>;
+
+        q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
+        k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
+        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>(text_length), static_cast<int>(hr)};
+
+        constexpr int mem_size = kittens::MAX_SHARED_MEMORY;
+        int threads = NUM_WORKERS * kittens::WARP_THREADS;
+        if (has_text) {
+            // TORCH_CHECK(seq_len % (CONSUMER_WARPGROUPS*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 192");
+            dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4)-2, qo_heads, batch);
+            dim3 grid_text(2, qo_heads, batch);
+            if (!process_text) {
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 1, 1, 1, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 1, 1, 1, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 1, 1, 2, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true,1, 1, 2, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 5 && kernel_h_size == 3 && kernel_w_size == 3) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 1, 1, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 1, 1, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                }else if (kernel_t_size ==3 && kernel_h_size == 5 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 1, 2, 2, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 1, 2, 2, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size ==5 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 3, 0, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 3, 0, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size ==5 && kernel_h_size == 3 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 1, 2, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 1, 2, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 2, 2, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 2, 2, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 2, 3, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 2, 3, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 5 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 3, 5, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 3, 5, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 0, 0, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 2, 0, 0, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 0, 3, 5, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true, 0, 3, 5, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, true, 2, 0, 5, 5, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, true,2, 0, 5, 5, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else {
+                    // print error
+                    std::cout << "Invalid kernel size" << std::endl;
+                    //print kernel size
+                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                }
+            } else {
+                cudaFuncSetAttribute(
+                    fwd_attend_ker<128, false, true, true, 1, 1, 1, 5, 6, 10>,
+                    cudaFuncAttributeMaxDynamicSharedMemorySize,
+                    mem_size
+                );
+                fwd_attend_ker<128, false, true, true, 1, 1, 1, 5, 6, 10><<<grid_text, (32*NUM_WORKERS), mem_size, stream>>>(g);
+            }
+
+        } else {
+            dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
+            if (kernel_aspect_ratio_flag == 2){
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }  else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }  else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else {
+                    // print error
+                    std::cout << "Invalid kernel size" << std::endl;
+                    //print kernel size
+                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                }
+            }
+            else if (kernel_aspect_ratio_flag == 3) {
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 2, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 1, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 9){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 9){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 0, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 0, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 3 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 1, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,0, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 2, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,0, 3, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else {
+                    // print error
+                    std::cout << "Invalid kernel size" << std::endl;
+                    //print kernel size
+                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                }
+            }
+
+            else {
+                std::cout << "Unsupported kernel_aspect_ratio_flag: " << kernel_aspect_ratio_flag << std::endl;
+            }
+
+        }
+        CHECK_CUDA_ERROR(cudaGetLastError());
+        // cudaStreamSynchronize(stream);
+    }
+
+    return o;
+    //cudadevicesynchronize();
+}
+

csrc/attn/sliding_tile_attn/st_attn/st_attn_triton.py

@@ -11,20 +11,12 @@ def is_hip():
     target = triton.runtime.driver.active.get_current_target()
     return target.backend == 'hip'
 
-
-def is_cdna3_cdna4():
-    target = triton.runtime.driver.active.get_current_target()
-    return (target.arch == 'gfx1201' or target.arch == 'gfx1101' or target.arch == 'gfx1100' or target.arch == 'gfx1030')
-
-
 def get_common_autotune_config():
-    # cdna arch does not support a 4-stage software pipeline, see https://github.com/ROCm/triton/issues/916
-    supported_num_staged = [1, 2] if is_cdna3_cdna4() else [1, 2, 3, 4]
     configs = [
         triton.Config({'BLOCK_Q': BLOCK_Q, 'BLOCK_KV': BLOCK_KV}, num_stages=s, num_warps=w) \
         for BLOCK_Q in [32, 64, 128]\
         for BLOCK_KV in [32, 64, 128]\
-        for s in supported_num_staged\
+        for s in [1, 2, 3, 4]\
         for w in [4, 8]\
     ]
     return configs

csrc/attn/tests/test_sta.py

@@ -1,6 +1,6 @@
 import torch
-from .support_flex_sta import get_sliding_tile_attention_mask
+from flex_sta_ref import get_sliding_tile_attention_mask
-from fastvideo_kernel import sliding_tile_attention
+from st_attn import sliding_tile_attention
 from torch.nn.attention.flex_attention import flex_attention
 # from flash_attn_interface import flash_attn_func
 from tqdm import tqdm
@@ -73,23 +73,15 @@ def check_correctness(b, h, n, d, causal, mean, std, num_iterations=50, error_mo
 
 
 # Example usage
-def test_sliding_tile_attention():
+b, h, d = 2, 24, 128
-    if not torch.cuda.is_available():
+n = 69120  # Sequence length
-        return
+causal = False
-
+mean = 1e-1
-    b, h, d = 2, 24, 128
+std = 10
-    n = 69120  # Sequence length
-    causal = False
-    mean = 1e-1
-    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']}")
-
-if __name__ == "__main__":
-    test_sliding_tile_attention()
 
+# 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/tests/test_vsa.py

@@ -1,13 +1,13 @@
 import torch
-import pytest
 import sys
 import os
 import numpy as np
 from tqdm import tqdm
 
-from .utils import generate_block_sparse_mask_for_function, create_full_mask_from_block_mask
+# Add the parent directory to the path to import block_sparse_attn
-# Use installed package
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-# from fastvideo_kernel import video_sparse_attn as block_sparse_attn
+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
@@ -42,36 +42,7 @@ def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, q
     q_padded = vsa_pad(Q, q_non_pad_index, q_num_blocks, BLOCK_M)
     k_padded = vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M)
     v_padded = vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M)
-    # Use raw kernel or triton
+    output, _= block_sparse_attn(q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes)
-    try:
-        from fastvideo_kernel._C import fastvideo_kernel_ops
-        raw_kernel = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None)
-    except ImportError:
-        raw_kernel = None
-
-    from fastvideo_kernel.triton_kernels.index import map_to_index
-
-    
-    # Convert mask to indices
-    # block_sparse_mask is [H, M, N] bool
-    # We need to map it to index.
-    # block_sparse_mask needs to be expanded/reshaped? 
-    # generate_block_sparse_mask_for_function returns [H, NumBlocksQ, NumBlocksKV]
-    
-    # Ops.py logic:
-    # mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
-    # idx, num = map_to_index(mask)
-    
-    idx, num = map_to_index(block_sparse_mask.unsqueeze(0)) # Add batch dim [1, H, M, N]
-    
-    if raw_kernel:
-        out_s = raw_kernel(q_padded, k_padded, v_padded, idx, num, variable_block_sizes.int())
-        output = out_s[0]
-    else:
-        # Fallback to triton testing if C++ not available
-        from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward
-        output, _ = triton_block_sparse_attn_forward(q_padded, k_padded, v_padded, idx, num, variable_block_sizes)
-
     output = output[:, :, q_non_pad_index, :]
     output.backward(dO)
     return output, Q.grad, K.grad, V.grad
@@ -264,19 +235,11 @@ def generate_error_graphs_qkdiff(h, d, error_mode='all'):
 
     print("-" * 150)
 
-@pytest.mark.skip()
+if __name__ == "__main__":
-def test_video_sparse_attention_backward():
-    if not torch.cuda.is_available():
-        return
-
     h, d = 16, 128
     print("Block Sparse Attention with Variable Block Sizes Analysis")
     print("=" * 60)
     for mode in ['backward']:
         generate_error_graphs(h, d, error_mode=mode)
         generate_error_graphs_qkdiff(h, d, error_mode=mode)
-    print("\nAnalysis completed for all modes.")
+    print("\nAnalysis completed for all modes.")
-
-if __name__ == "__main__":
-    test_video_sparse_attention_backward()
-

csrc/attn/tests/test_vsa_forward.py

@@ -4,12 +4,21 @@ from typing import Tuple
 
 import torch
 
-from .utils import (
+# Make sure we can import from the project root (`vsa`, `tests.utils`, etc.)
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
+PROJECT_ROOT = os.path.dirname(CURRENT_DIR)
+if PROJECT_ROOT not in sys.path:
+    sys.path.append(PROJECT_ROOT)
+if CURRENT_DIR not in sys.path:
+    sys.path.append(CURRENT_DIR)
+
+from tests.utils import (
     generate_block_sparse_mask_for_function,
     create_full_mask_from_block_mask,
 )
-from .test_vsa import BLOCK_M  # Import from local test_vsa
+from vsa import block_sparse_attn, BLOCK_M
-from . import test_vsa as ref
+import test_vsa as ref  # reuse helper functions from backward test
+
 
 def pytorch_forward(
     Q: torch.Tensor,
@@ -47,7 +56,8 @@ def block_sparse_forward_test(
     kv_num_blocks: int,
 ) -> torch.Tensor:
     """
-    Forward-only wrapper
+    Forward-only wrapper around `block_sparse_attn`, mirroring `block_sparse_kernel_test`
+    but without any backward / grad logic.
     """
     Q = Q.detach()
     K = K.detach()
@@ -57,24 +67,9 @@ def block_sparse_forward_test(
     k_padded = ref.vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M)
     v_padded = ref.vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M)
 
-    # Use raw kernel or triton
+    out_padded, _ = block_sparse_attn(
-    try:
+        q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes
-        from fastvideo_kernel._C import fastvideo_kernel_ops
+    )
-        raw_kernel = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None)
-    except ImportError:
-        raw_kernel = None
-        
-    from fastvideo_kernel.triton_kernels.index import map_to_index
-    idx, num = map_to_index(block_sparse_mask)
-    
-    if raw_kernel:
-        out_padded = raw_kernel(q_padded, k_padded, v_padded, idx, num, variable_block_sizes.int())[0]
-    else:
-        from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward
-        out_padded, _ = triton_block_sparse_attn_forward(
-            q_padded, k_padded, v_padded, idx, num, variable_block_sizes
-        )
-
     # Remove padding on the query side
     out = out_padded[:, :, q_non_pad_index, :]
     return out
@@ -223,10 +218,7 @@ def run_forward_qk_diff(
     return avg_abs_err, max_rel_diff
 
 
-def test_video_sparse_attention_forward():
+if __name__ == "__main__":
-    if not torch.cuda.is_available():
-        return
-
     h, d = 16, 128
     print("Forward Block Sparse Attention Check (QK Equal)")
     print("=" * 80)
@@ -242,6 +234,3 @@ def test_video_sparse_attention_forward():
         f"QK diff:  avg |ΔO| = {avg_err_diff:.6e}, max rel ΔO = {max_rel_diff:.6e}"
     )
 
-if __name__ == "__main__":
-    test_video_sparse_attention_forward()
-

csrc/attn/video_sparse_attn/MANIFEST.in

@@ -0,0 +1,2 @@
+recursive-include tk *
+include config_vsa.py

csrc/attn/video_sparse_attn/README.md

@@ -0,0 +1,61 @@
+
+
+# Attention Kernel Used in FastVideo
+
+## Video Sparse Attention (VSA)
+
+### Installation
+We support H100 (via TK) and any other GPU (via triton) for VSA.
+
+```bash
+pip install vsa
+``` 
+
+Install from source:
+
+```bash
+git submodule update --init --recursive
+python setup.py install
+```
+
+
+If you encounter error during installation, try below:
+Install C++20 for ThunderKittens:
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+(If you use CUDA12.8)
+```bash
+export CUDA_HOME=/usr/local/cuda-12.8
+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
+```
+
+## Acknowledgement
+
+We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/attn/video_sparse_attn/config_vsa.py

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

csrc/attn/video_sparse_attn/setup.py

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

csrc/attn/video_sparse_attn/vsa.cpp

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

csrc/attn/video_sparse_attn/vsa/__init__.py

@@ -0,0 +1,80 @@
+import torch
+from typing import Tuple
+block_sparse_attn=None
+import torch
+major, minor = torch.cuda.get_device_capability(0)
+if major == 9 and minor == 0:# check if H100
+    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
+else:
+    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
+
+BLOCK_M = 64
+BLOCK_N = 64
+
+
+def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
+    QK = torch.matmul(q, k.transpose(-2, -1))
+    QK /= (q.size(-1)**0.5)
+
+    # Causal mask removed since causal is always false
+
+    QK = torch.nn.functional.softmax(QK, dim=-1)
+    output = torch.matmul(QK, v)
+    return output, QK
+
+
+def video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight=None):
+    """
+    q: [batch_size, num_heads, seq_len, head_dim]
+    k: [batch_size, num_heads, seq_len, head_dim]
+    v: [batch_size, num_heads, seq_len, head_dim]
+    topk: int
+    block_size: int or tuple of 3 ints
+    video_shape: tuple of (T, H, W)
+    compress_attn_weight: [batch_size, num_heads, seq_len, head_dim]
+    select_attn_weight: [batch_size, num_heads, seq_len, head_dim]
+    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]]
+    """
+
+    if isinstance(block_size, int):
+        block_size = (block_size, block_size, block_size)
+
+    block_elements = block_size[0] * block_size[1] * block_size[2]
+    assert block_elements == 64
+    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)
+
+    output_compress, block_attn_score = torch_attention(q_compress, k_compress,
+                                                        v_compress)
+
+    output_compress = output_compress.view(batch_size, num_heads,
+                                           seq_len // block_elements, 1,
+                                           head_dim)
+    output_compress = output_compress.repeat(1, 1, 1, block_elements,
+                                             1).view(batch_size, num_heads,
+                                                     seq_len, head_dim)
+
+    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)
+
+    if compress_attn_weight is not None:
+        final_output = output_compress * compress_attn_weight + output_select
+    else:
+        final_output = output_compress + output_select
+    return final_output
+

csrc/attn/video_sparse_attn/vsa/block_sparse_attn_triton.py

@@ -0,0 +1,450 @@
+"""
+Fused Attention
+===============
+
+This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
+(https://tridao.me/publications/flash2/flash2.pdf)
+
+Credits: OpenAI kernel team
+"""
+
+import pytest
+import torch
+import triton
+import triton.language as tl
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+import math  # small utility needed by the sparse wrapper
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+# We don't run auto-tuning every time to keep the tutorial fast. Keeping
+# the code below and commenting out the equivalent parameters is convenient for
+# re-tuning.
+configs = [
+    triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+    for BM in [64]\
+    for BN in [64]\
+    for s in [3, 4, 7]\
+    for w in [4, 8]\
+]
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
+@triton.jit
+def _attn_fwd_sparse(Q, K, V, sm_scale,                         #
+                     q2k_index, q2k_num, max_kv_blks,           #
+                     variable_block_sizes,
+                     M, Out,                                    #
+                     stride_qz, stride_qh, stride_qm, stride_qk,
+                     stride_kz, stride_kh, stride_kn, stride_kk,
+                     stride_vz, stride_vh, stride_vk, stride_vn,
+                     stride_oz, stride_oh, stride_om, stride_on,
+                     Z, H, N_CTX,                               #
+                     HEAD_DIM: tl.constexpr,                    #
+                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
+                     STAGE: tl.constexpr):
+    """
+    64×64 **block-sparse** forward kernel. Back-prop kernels remain dense
+    (32×64 and 64×32) – memory footprint unchanged.
+    """
+
+    # ----- program-id mapping -----
+    q_blk   = tl.program_id(0)          # Q-tile index
+    off_hz  = tl.program_id(1)          # fused (batch, head)
+    b       = off_hz // H
+    h       = off_hz %  H
+    q_tiles = N_CTX // BLOCK_M
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
+    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
+
+    # ----- base pointers -----
+    qvk_off = (b.to(tl.int64) * stride_qz +
+               h.to(tl.int64) * stride_qh)
+
+    Q_ptr = tl.make_block_ptr(
+        base=Q + qvk_off, shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(q_blk * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
+
+    K_base = tl.make_block_ptr(
+        base=K + qvk_off, shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N), order=(0, 1))
+
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
+    V_base = tl.make_block_ptr(
+        base=V + qvk_off, shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, HEAD_DIM), order=v_order)
+
+    O_ptr = tl.make_block_ptr(
+        base=Out + qvk_off, shape=(N_CTX, HEAD_DIM),
+        strides=(stride_om, stride_on),
+        offsets=(q_blk * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
+
+    # ----- accumulators -----
+    offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M)
+    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    qk_scale = sm_scale * 1.44269504  # 1/ln2
+    q = tl.load(Q_ptr)
+
+    # ----- sparse loop over valid K/V tiles -----
+    for i in range(0, kv_blocks):
+        kv_idx = tl.load(kv_ptr + i).to(tl.int32)
+        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])
+        l_ij = tl.sum(p, 1)
+
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        acc = acc * alpha[:, None]
+
+        v = tl.load(V_ptr)
+        acc = tl.dot(p.to(tl.bfloat16), v, acc)
+        m_i = m_ij
+
+    # ----- epilogue -----
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    tl.store(M + off_hz * N_CTX + offs_m, m_i)
+    tl.store(O_ptr, acc.to(Out.type.element_ty))
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+
+
+@triton.jit
+def _attn_bwd_preprocess(O, DO,  #
+                         Delta,  #
+                         Z, H, N_CTX,  #
+                         BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr  #
+                         ):
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_hz = tl.program_id(1)
+    off_n = tl.arange(0, HEAD_DIM)
+    # load
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    # write-back
+    tl.store(Delta + off_hz * N_CTX + off_m, delta)
+
+
+# The main inner-loop logic for computing dK and dV.
+@triton.jit
+def _attn_bwd_dkdv(dk, dv,  #
+                   Q, k, v, sm_scale,  #
+                   DO,  #
+                   M, D,  #
+                   k2q_index, k2q_num, max_q_blks,
+                   variable_block_sizes,
+                   # shared by Q/K/V/DO.
+                   stride_tok, stride_d,  #
+                   H, N_CTX, BLOCK_M1: tl.constexpr,  #
+                   BLOCK_N1: tl.constexpr,  #
+                   HEAD_DIM: tl.constexpr,  #
+                   # Filled in by the wrapper.
+                   start_n, start_m, num_steps):
+    offs_m = start_m + tl.arange(0, BLOCK_M1)
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+    offs_k = tl.arange(0, HEAD_DIM)
+    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
+    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
+    step_m = BLOCK_M1
+    kv_blk   = tl.program_id(0)          # Q-tile index
+    off_hz  = tl.program_id(2)          # fused (batch, head)
+    b       = off_hz // H
+    h       = off_hz %  H
+    q_tiles = N_CTX // BLOCK_N1
+    meta_base = ((b * H + h) * q_tiles + kv_blk)
+
+    q_blocks = tl.load(k2q_num  + meta_base)                 # int32
+    q_ptr    = k2q_index + meta_base * max_q_blks           # ptr to list
+    block_size = tl.load(variable_block_sizes + kv_blk)
+    
+    
+        
+    for blk_idx in range(q_blocks*2):
+        block_sparse_offset = (tl.load(q_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_m
+        qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok)
+        # Load m before computing qk to reduce pipeline stall.
+        offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
+        m = tl.load(M + offs_m)
+        qkT = tl.dot(k, qT)
+        pT = tl.math.exp2(qkT - m[None, :])
+        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
+        ppT = ppT.to(tl.bfloat16)
+        dv += tl.dot(ppT, do)
+        # D (= delta) is pre-divided by ds_scale.
+        Di = tl.load(D + offs_m)
+        # Compute dP and dS.
+        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
+        dsT = pT * (dpT - Di[None, :])
+        dsT = dsT.to(tl.bfloat16)
+        dk += tl.dot(dsT, tl.trans(qT))
+        # Increment pointers.
+    return dk, dv
+
+
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _attn_bwd_dq(dq, q, K, V,  #
+                 do, m, D,
+                 # shared by Q/K/V/DO.
+                 q2k_index, q2k_num, max_kv_blks,
+                 variable_block_sizes,
+                 stride_tok, stride_d,  #
+                 H, N_CTX,  #
+                 BLOCK_M2: tl.constexpr,  #
+                 BLOCK_N2: tl.constexpr,  #
+                 HEAD_DIM: tl.constexpr,
+                 # Filled in by the wrapper.
+                 start_m, start_n, num_steps):
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+    offs_n = start_n + tl.arange(0, BLOCK_N2)
+    offs_k = tl.arange(0, HEAD_DIM)
+    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    # D (= delta) is pre-divided by ds_scale.
+    Di = tl.load(D + offs_m)
+    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
+    step_n = BLOCK_N2
+    
+    q_blk   = tl.program_id(0)          # Q-tile index
+    off_hz  = tl.program_id(2)          # fused (batch, head)
+    b       = off_hz // H
+    h       = off_hz %  H
+    q_tiles = N_CTX // BLOCK_M2
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
+    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
+    
+    
+    for blk_idx in range(kv_blocks*2):
+        kv_idx = tl.load(kv_ptr + blk_idx//2).to(tl.int32)
+        block_size = tl.load(variable_block_sizes + kv_idx) -  (blk_idx % 2) * step_n
+        block_sparse_offset = (kv_idx*2 + blk_idx%2) * step_n * stride_tok
+        kT = tl.load(kT_ptrs + block_sparse_offset)
+        vT = tl.load(vT_ptrs + block_sparse_offset)
+        qk = tl.dot(q, kT)
+        p = tl.math.exp2(qk - m)
+        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])
+        ds = ds.to(tl.bfloat16)
+        # Compute dQ.
+        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        dq += tl.dot(ds, tl.trans(kT))
+        # Increment pointers.
+    return dq
+
+
+
+@triton.jit
+def _attn_bwd(Q, K, V, sm_scale,  #
+              DO,  #
+              DQ, DK, DV,  #
+              M, D,
+              q2k_index, q2k_num, max_kv_blks,
+              k2q_index, k2q_num, max_q_blks,
+              variable_block_sizes,
+              # shared by Q/K/V/DO.
+              stride_z, stride_h, stride_tok, stride_d,  #
+              H, N_CTX,  #
+              BLOCK_M1: tl.constexpr,  #
+              BLOCK_N1: tl.constexpr,  #
+              BLOCK_M2: tl.constexpr,  #
+              BLOCK_N2: tl.constexpr,  #
+              HEAD_DIM: tl.constexpr):
+    LN2 = 0.6931471824645996  # = ln(2)
+
+    bhid = tl.program_id(2)
+    off_chz = (bhid * N_CTX).to(tl.int64)
+    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
+    pid = tl.program_id(0)
+
+    # offset pointers for batch/head
+    Q += adj
+    K += adj
+    V += adj
+    DO += adj
+    DQ += adj
+    DK += adj
+    DV += adj
+    M += off_chz
+    D += off_chz
+
+    # load scales
+    offs_k = tl.arange(0, HEAD_DIM)
+
+    start_n = pid * BLOCK_N1
+    start_m = 0
+
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+
+    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+    # load K and V: they stay in SRAM throughout the inner loop.
+    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+
+    num_steps = N_CTX // BLOCK_M1
+
+    dk, dv = _attn_bwd_dkdv(  #
+        dk, dv,  #
+        Q, k, v, sm_scale,  #
+        DO,  #
+        M, D,  #
+        k2q_index, k2q_num, max_q_blks,
+        variable_block_sizes,
+        stride_tok, stride_d,  #
+        H, N_CTX,  #
+        BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
+        start_n, start_m, num_steps  #
+    )
+
+    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dv_ptrs, dv)
+
+    # Write back dK.
+    dk *= sm_scale
+    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dk_ptrs, dk)
+
+    # THIS BLOCK DOES DQ:
+    start_m = pid * BLOCK_M2
+    end_n = 0
+
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+
+    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    m = tl.load(M + offs_m)
+    m = m[:, None]
+
+    num_steps = N_CTX // BLOCK_N2
+    dq = _attn_bwd_dq(dq, q, K, V,  #
+                      do, m, D,  #
+                      q2k_index, q2k_num, max_kv_blks,
+                      variable_block_sizes,
+                      stride_tok, stride_d,  #
+                      H, N_CTX,  #
+                      BLOCK_M2, BLOCK_N2, HEAD_DIM,  #
+                      start_m, end_n, num_steps  #
+                      )
+    # Write back dQ.
+    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    dq *= LN2
+    tl.store(dq_ptrs, dq)
+    
+    
+
+
+
+# ──────────────────────────── 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)
+
+    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
+    )
+
+    return o, M
+
+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  #
+    )
+
+    return dq, dk, dv
+
+

csrc/attn/video_sparse_attn/vsa/block_sparse_wrapper.py

@@ -0,0 +1,185 @@
+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_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)
+
+
+major, minor = torch.cuda.get_device_capability(0)
+
+if major == 9 and minor == 0:# check if H100
+    @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)

csrc/attn/video_sparse_attn/vsa/index.py

@@ -1,9 +1,9 @@
+
 ## 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,
@@ -26,7 +26,6 @@ def topk_index_to_map_kernel(
         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,
@@ -60,7 +59,6 @@ def map_to_index_kernel(
         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):
@@ -108,7 +106,6 @@ def topk_index_to_map(index: torch.Tensor,
 
     return block_map
 
-
 def map_to_index(block_map: torch.Tensor):
     """
     Convert a block map to indices and counts.

csrc/attn/vmoba_attn/README.md

@@ -0,0 +1,32 @@
+# Attention Kernel Used in FastVideo
+
+## VMoBA: Mixture-of-Block Attention for Video Diffusion Models (VMoBA)
+
+### Installation
+Please ensure that you have installed FlashAttention version **2.7.1 or higher**, as some interfaces have changed in recent releases.
+
+### Usage
+
+You can use `moba_attn_varlen` in the following ways:
+
+**Install from source:**
+```bash
+python setup.py install
+```
+
+**Import after installation:**
+```python
+from vmoba import moba_attn_varlen
+```
+
+**Or import directly from the project root:**
+```python
+from csrc.attn.vmoba_attn.vmoba import moba_attn_varlen
+```
+
+### Verify if you have successfully installed
+
+```bash
+python csrc/attn/vmoba_attn/vmoba/vmoba.py
+```
+

csrc/attn/vmoba_attn/setup.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from setuptools import find_packages, setup
+
+PACKAGE_NAME = "vmoba"
+VERSION = "0.0.0"
+AUTHOR = "JianzongWu"
+DESCRIPTION = "VMoBA: Mixture-of-Block Attention for Video Diffusion Models"
+URL = "https://github.com/KwaiVGI/VMoBA"
+
+setup(
+    name=PACKAGE_NAME,
+    version=VERSION,
+    author=AUTHOR,
+    description=DESCRIPTION,
+    url=URL,
+    packages=find_packages(),
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: Apache Software License",
+    ],
+    python_requires='>=3.12',
+    install_requires=[
+        "flash-attn >= 2.7.1",
+    ]
+)

csrc/attn/vmoba_attn/tests/test_vmoba_attn.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import pytest
+import random
+from csrc.attn.vmoba_attn.vmoba  import moba_attn_varlen
+
+def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda"):
+    """
+    Generates random data for testing the variable-length attention function.
+    """
+    torch.manual_seed(42)
+    random.seed(42)
+    torch.cuda.manual_seed_all(42)
+
+    # Generate sequence lengths for each item in the batch
+    if batch_size > 1:
+        # Ensure sequence lengths are reasonably distributed
+        avg_seqlen = total_seqlen // batch_size
+        seqlens = [random.randint(avg_seqlen // 2, avg_seqlen + avg_seqlen // 2) for _ in range(batch_size - 1)]
+        remaining_len = total_seqlen - sum(seqlens)
+        if remaining_len > 0:
+            seqlens.append(remaining_len)
+        else: # Adjust if sum exceeds total_seqlen
+            seqlens.append(avg_seqlen)
+            current_sum = sum(seqlens)
+            seqlens[-1] -= (current_sum - total_seqlen)
+        # Ensure all lengths are positive
+        seqlens = [max(1, s) for s in seqlens]
+        # Final adjustment to match total_seqlen
+        seqlens[-1] += total_seqlen - sum(seqlens)
+
+    else:
+        seqlens = [total_seqlen]
+
+    cu_seqlens = torch.tensor([0] + list(torch.cumsum(torch.tensor(seqlens), 0)), device=device, dtype=torch.int32)
+    max_seqlen = max(seqlens) if seqlens else 0
+
+    q = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False)
+    k = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False)
+    v = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False)
+
+    return q, k, v, cu_seqlens, max_seqlen
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("total_seqlen", [512, 1024])
+@pytest.mark.parametrize("num_heads", [8])
+@pytest.mark.parametrize("head_dim", [64])
+@pytest.mark.parametrize("moba_chunk_size", [64])
+@pytest.mark.parametrize("moba_topk", [2, 4])
+@pytest.mark.parametrize("select_mode", ["topk", "threshold"])
+@pytest.mark.parametrize("threshold_type", ["query_head", "head_global", "overall"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_moba_attn_varlen_forward(
+    batch_size, total_seqlen, num_heads, head_dim, moba_chunk_size, moba_topk, select_mode, threshold_type, dtype
+):
+    """
+    Tests the forward pass of moba_attn_varlen for basic correctness.
+    It checks output shape, dtype, and for the presence of NaNs/Infs.
+    """
+    if dtype == torch.float32:
+        pytest.skip("float32 is not supported in flash attention")
+
+    q, k, v, cu_seqlens, max_seqlen = generate_test_data(
+        batch_size, total_seqlen, num_heads, head_dim, dtype
+    )
+
+    # Ensure chunk size is not larger than the smallest sequence length
+    min_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).min().item()
+    if moba_chunk_size > min_seqlen:
+        pytest.skip("moba_chunk_size is larger than the minimum sequence length in the batch")
+
+    try:
+        output = moba_attn_varlen(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            moba_chunk_size=moba_chunk_size,
+            moba_topk=moba_topk,
+            select_mode=select_mode,
+            threshold_type=threshold_type,
+            simsum_threshold=0.5, # A reasonable default for threshold mode
+        )
+    except Exception as e:
+        pytest.fail(f"moba_attn_varlen forward pass failed with exception: {e}")
+
+    # 1. Check output shape
+    assert output.shape == q.shape, f"Expected output shape {q.shape}, but got {output.shape}"
+
+    # 2. Check output dtype
+    assert output.dtype == q.dtype, f"Expected output dtype {q.dtype}, but got {output.dtype}"
+
+    # 3. Check for NaNs or Infs in the output
+    assert torch.all(torch.isfinite(output)), "Output contains NaN or Inf values"

csrc/attn/vmoba_attn/vmoba/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+from .vmoba import moba_attn_varlen, process_moba_input, process_moba_output

csrc/fastvideo_kernel/.gitignore

@@ -0,0 +1,7 @@
+build/
+dist/
+*.egg-info/
+__pycache__/
+*.so
+*.pyc
+.ipynb_checkpoints/

csrc/fastvideo_kernel/MANIFEST.in

@@ -0,0 +1,6 @@
+include LICENSE
+include README.md
+include pyproject.toml
+recursive-include src/fastvideo_kernel *.cu *.cuh *.cpp *.h
+recursive-include csrc *.cu *.cuh *.cpp *.h
+recursive-include tk *.cu *.cuh *.cpp *.h

csrc/fastvideo_kernel/README.md

@@ -0,0 +1,31 @@
+# FastVideo Kernel
+
+CUDA kernels for FastVideo video generation.
+
+## Installation
+
+```bash
+git submodule update --init --recursive
+cd csrc/fastvideo_kernel
+pip install .
+```
+
+## Usage
+
+```python
+from fastvideo_kernel import sliding_tile_attention, video_sparse_attn, moba_attn_varlen
+
+# Example: Sliding Tile Attention
+out = sliding_tile_attention(q, k, v, window_sizes, text_len)
+
+# Example: Video Sparse Attention (with Triton fallback)
+out = video_sparse_attn(q, k, v, block_sizes, topk=5)
+
+# Example: VMoBA
+out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)
+```
+
+## Requirements
+
+- H100 GPU (sm_90a) for CUDA kernels
+- Triton for non-H100 fallback

csrc/fastvideo_kernel/csrc/st_attn.cpp

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

csrc/fastvideo_kernel/csrc/vsa.cpp

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

csrc/fastvideo_kernel/pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["setuptools>=61.0", "torch>=2.5.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "fastvideo-kernel"
+version = "0.1.0"
+description = "CUDA kernels for FastVideo"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "Apache-2.0"}
+authors = [{name = "Hao AI Lab"}]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Environment :: GPU :: NVIDIA CUDA :: 12",
+    "License :: OSI Approved :: Apache Software License",
+]
+dependencies = [
+    "torch>=2.5.0",
+    "triton>=2.0.0"
+]
+
+[project.urls]
+Repository = "https://github.com/hao-ai-lab/FastVideo"
+
+[tool.setuptools.packages.find]
+where = ["src"]

csrc/fastvideo_kernel/setup.py

@@ -0,0 +1,132 @@
+import os
+import subprocess
+import sys
+from pathlib import Path
+from setuptools import find_packages, setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+ROOT = Path(__file__).parent.absolute()
+CSRC_DIR = ROOT / "csrc"
+
+# Path to ThunderKittens (TK)
+def get_tk_dir():
+    tk_env = os.getenv("THUNDERKITTENS_ROOT")
+    if tk_env:
+        return tk_env
+    
+    # Check common locations
+    possible_paths = [
+        ROOT / "tk",
+        ROOT / "csrc" / "tk",
+        ROOT.parent / "attn" / "sliding_tile_attn" / "tk",
+        ROOT.parent / "attn" / "video_sparse_attn" / "tk",
+    ]
+    for p in possible_paths:
+        if (p / "include" / "kittens.cuh").exists():
+            return str(p)
+    
+    # Default fallback
+    return str(ROOT.parent / "attn" / "sliding_tile_attn" / "tk")
+
+TK_DIR = get_tk_dir()
+
+def get_cuda_flags(tk_root: str) -> list:
+    python_include = subprocess.check_output(
+        ["python", "-c", "import sysconfig; print(sysconfig.get_path('include'))"]
+    ).decode().strip()
+    
+    torch_includes = subprocess.check_output([
+        "python", "-c",
+        "import torch; from torch.utils.cpp_extension import include_paths; "
+        "print(' '.join(['-I' + p for p in include_paths()]))"
+    ]).decode().strip().split()
+    
+    return [
+        "-DNDEBUG",
+        "-Xcompiler=-Wno-psabi",
+        "-Xcompiler=-fno-strict-aliasing",
+        "--expt-extended-lambda",
+        "--expt-relaxed-constexpr",
+        "-forward-unknown-to-host-compiler",
+        "--use_fast_math",
+        "-std=c++20",
+        "-O3",
+        "-Xnvlink=--verbose",
+        "-Xptxas=--verbose",
+        "-Xptxas=--warn-on-spills",
+        f"-I{tk_root}/include",
+        f"-I{tk_root}/prototype",
+        f"-I{python_include}",
+        "-DTORCH_COMPILE",
+        "-DKITTENS_HOPPER",
+        "-arch=sm_90a",
+    ] + torch_includes
+
+def get_extensions():
+    if not torch.cuda.is_available():
+        return []
+
+    extensions = []
+    cpp_flags = ["-std=c++20", "-O3"]
+    
+    # Check if TK is available
+    if not os.path.exists(os.path.join(TK_DIR, "include", "kittens.cuh")):
+        print(f"Warning: ThunderKittens not found at {TK_DIR}. CUDA kernels will not be built.")
+        return []
+
+    cuda_flags = get_cuda_flags(TK_DIR)
+    
+    # STA Extension
+    extensions.append(CUDAExtension(
+        "fastvideo_kernel._C.st_attn",
+        sources=[
+            "csrc/st_attn.cpp",
+            "csrc/st_attn_h100.cu",
+        ],
+        extra_compile_args={
+            "cxx": cpp_flags + ["-DTK_COMPILE_ST_ATTN"], 
+            "nvcc": cuda_flags + ["-DTK_COMPILE_ST_ATTN"]
+        },
+        libraries=["cuda"],
+    ))
+    
+    # VSA Extension
+    extensions.append(CUDAExtension(
+        "fastvideo_kernel._C.vsa",
+        sources=[
+            "csrc/vsa.cpp",
+            "csrc/block_sparse_h100.cu",
+        ],
+        extra_compile_args={
+            "cxx": cpp_flags + ["-DTK_COMPILE_BLOCK_SPARSE"], 
+            "nvcc": cuda_flags + ["-DTK_COMPILE_BLOCK_SPARSE"]
+        },
+        libraries=["cuda"],
+    ))
+    
+    return extensions
+
+ext_modules = []
+if not any(arg in sys.argv for arg in ["clean", "egg_info", "--version"]):
+    try:
+        import torch
+        ext_modules = get_extensions()
+    except Exception as e:
+        print(f"Warning: Failed to configure CUDA extensions: {e}")
+
+setup(
+    name="fastvideo-kernel",
+    version="0.1.0",
+    description="Unified CUDA kernels for FastVideo",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    license="Apache-2.0",
+    author="Hao AI Lab",
+    url="https://github.com/hao-ai-lab/FastVideo",
+    package_dir={"": "src"},
+    packages=find_packages(where="src"),
+    ext_modules=ext_modules,
+    cmdclass={"build_ext": BuildExtension} if ext_modules else {},
+    python_requires=">=3.10",
+    install_requires=["torch>=2.5.0", "triton>=2.0.0"],
+)

csrc/fastvideo_kernel/src/fastvideo_kernel/__init__.py

@@ -1,4 +1,4 @@
-from .version import __version__
+__version__ = "0.1.0"
 
 from fastvideo_kernel.ops import (
     sliding_tile_attention,
@@ -11,24 +11,11 @@ from fastvideo_kernel.vmoba import (
     process_moba_output,
 )
 
-from fastvideo_kernel.turbodiffusion_ops import (
-    Int8Linear,
-    FastRMSNorm,
-    FastLayerNorm,
-    int8_linear,
-    int8_quant,
-)
-
 __all__ = [
     "sliding_tile_attention",
     "video_sparse_attn",
     "moba_attn_varlen",
     "process_moba_input",
     "process_moba_output",
-    "Int8Linear",
-    "FastRMSNorm",
-    "FastLayerNorm",
-    "int8_linear",
-    "int8_quant",
     "__version__",
 ]

csrc/fastvideo_kernel/src/fastvideo_kernel/ops.py

@@ -1,20 +1,20 @@
 import math
 import torch
 from .triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward
-from .triton_kernels.st_attn_triton import sliding_tile_attention_triton
 from .triton_kernels.index import map_to_index
 
-# Try to load the C++ extension
 try:
-    from fastvideo_kernel._C import fastvideo_kernel_ops
+    from fastvideo_kernel._C.st_attn import sta_fwd
-    sta_fwd = getattr(fastvideo_kernel_ops, "sta_fwd", None)
-    block_sparse_fwd = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None)
-    block_sparse_bwd = getattr(fastvideo_kernel_ops, "block_sparse_bwd", None)
 except ImportError:
     sta_fwd = None
+
+try:
+    from fastvideo_kernel._C.vsa import block_sparse_fwd, block_sparse_bwd
+except ImportError:
     block_sparse_fwd = None
     block_sparse_bwd = None
 
+
 def sliding_tile_attention(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -24,15 +24,12 @@ def sliding_tile_attention(
     has_text: bool = True,
     seq_shape: str = "30x48x80",
 ) -> torch.Tensor:
-    # Check if the specific op is available
     if sta_fwd is None:
-        return sliding_tile_attention_triton(
+        raise RuntimeError("STA kernel not compiled. Requires H100 and ThunderKittens at build time.")
-            q, k, v, window_size, text_length, has_text, seq_shape
+    
-        )
-
     seq_length = q.shape[2]
     shape_map = {"30x48x80": 1, "36x48x48": 2, "18x48x80": 3}
-
+    
     if has_text:
         target_size = math.ceil(seq_length / 384) * 384
         pad_size = target_size - seq_length
@@ -40,20 +37,22 @@ def sliding_tile_attention(
             q = torch.cat([q, q[:, :, -pad_size:]], dim=2)
             k = torch.cat([k, k[:, :, -pad_size:]], dim=2)
             v = torch.cat([v, v[:, :, -pad_size:]], dim=2)
-
+    
     output = torch.empty_like(q)
     flag = shape_map[seq_shape]
-
+    
     for head_idx, (t, h, w) in enumerate(window_size):
         sta_fwd(
-            q[:, head_idx:head_idx + 1], k[:, head_idx:head_idx + 1],
+            q[:, head_idx:head_idx+1],
-            v[:, head_idx:head_idx + 1], output[:, head_idx:head_idx + 1],
+            k[:, head_idx:head_idx+1],
+            v[:, head_idx:head_idx+1],
+            output[:, head_idx:head_idx+1],
             t, h, w, text_length, False, has_text, flag
         )
-
+    
     if has_text:
         sta_fwd(q, k, v, output, 3, 3, 3, text_length, True, True, flag)
-
+    
     return output[:, :, :seq_length]
 
 
@@ -68,45 +67,37 @@ def video_sparse_attn(
 ) -> torch.Tensor:
     if isinstance(block_size, int):
         block_size = (block_size, block_size, block_size)
-
+    
     block_elements = block_size[0] * block_size[1] * block_size[2]
     batch, heads, seq_len, dim = q.shape
-
+    
     # Compression branch
     q_c = q.view(batch, heads, seq_len // block_elements, block_elements, dim)
     k_c = k.view(batch, heads, seq_len // block_elements, block_elements, dim)
     v_c = v.view(batch, heads, seq_len // block_elements, block_elements, dim)
-
+    
-    q_c = (q_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
+    q_c = (q_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(q.dtype)
-        q.dtype)
+    k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype)
-    k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
+    v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype)
-        k.dtype)
+    
-    v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
+    scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim ** 0.5)
-        v.dtype)
-
-    scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim**0.5)
     attn = torch.softmax(scores, dim=-1)
     out_c = torch.matmul(attn, v_c)
-
+    
     out_c = out_c.view(batch, heads, seq_len // block_elements, 1, dim)
-    out_c = out_c.repeat(1, 1, 1, block_elements,
+    out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, seq_len, dim)
-                         1).view(batch, heads, seq_len, dim)
+    
-
     # Sparse branch
     topk_idx = torch.topk(scores, topk, dim=-1).indices
-    mask = torch.zeros_like(scores,
+    mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
-                            dtype=torch.bool).scatter_(-1, topk_idx, True)
-
-    idx, num = map_to_index(mask)
     
     if block_sparse_fwd is not None:
-        out_s = block_sparse_fwd(
+        idx, num = map_to_index(mask)
-            q, k, v, idx, num, variable_block_sizes.int()
+        out_s, _ = block_sparse_fwd(q, k, v, idx, num, variable_block_sizes.int())
-        )[0] # block_sparse_fwd returns vector<Tensor>
     else:
-        out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num,
+        idx, num = map_to_index(mask)
-                                                    variable_block_sizes)
+        out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num, variable_block_sizes)
-
+    
     if compress_attn_weight is not None:
         return out_c * compress_attn_weight + out_s
     return out_c + out_s

csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py

@@ -8,6 +8,7 @@ This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
 Credits: OpenAI kernel team
 """
 
+
 import torch
 import triton
 import triton.language as tl
@@ -16,6 +17,7 @@ import triton.language as tl
 import math  # small utility needed by the sparse wrapper
 # ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
 
+
 # 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.
@@ -27,92 +29,65 @@ configs = [
     for w in [4, 8]\
 ]
 
-
 # ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
 @triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
 @triton.jit
-def _attn_fwd_sparse(
+def _attn_fwd_sparse(Q, K, V, sm_scale,                         #
-        Q,
+                     q2k_index, q2k_num, max_kv_blks,           #
-        K,
+                     variable_block_sizes,
-        V,
+                     M, Out,                                    #
-        sm_scale,  #
+                     stride_qz, stride_qh, stride_qm, stride_qk,
-        q2k_index,
+                     stride_kz, stride_kh, stride_kn, stride_kk,
-        q2k_num,
+                     stride_vz, stride_vh, stride_vk, stride_vn,
-        max_kv_blks,  #
+                     stride_oz, stride_oh, stride_om, stride_on,
-        variable_block_sizes,
+                     Z, H, N_CTX,                               #
-        M,
+                     HEAD_DIM: tl.constexpr,                    #
-        Out,  #
+                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
-        stride_qz,
+                     STAGE: tl.constexpr):
-        stride_qh,
-        stride_qm,
-        stride_qk,
-        stride_kz,
-        stride_kh,
-        stride_kn,
-        stride_kk,
-        stride_vz,
-        stride_vh,
-        stride_vk,
-        stride_vn,
-        stride_oz,
-        stride_oh,
-        stride_om,
-        stride_on,
-        Z,
-        H,
-        N_CTX,  #
-        HEAD_DIM: tl.constexpr,  #
-        BLOCK_M: tl.constexpr,
-        BLOCK_N: tl.constexpr,
-        STAGE: tl.constexpr):
     """
     64×64 **block-sparse** forward kernel. Back-prop kernels remain dense
     (32×64 and 64×32) – memory footprint unchanged.
     """
 
     # ----- program-id mapping -----
-    q_blk = tl.program_id(0)  # Q-tile index
+    q_blk   = tl.program_id(0)          # Q-tile index
-    off_hz = tl.program_id(1)  # fused (batch, head)
+    off_hz  = tl.program_id(1)          # fused (batch, head)
-    b = off_hz // H
+    b       = off_hz // H
-    h = off_hz % H
+    h       = off_hz %  H
     q_tiles = N_CTX // BLOCK_M
     meta_base = ((b * H + h) * q_tiles + q_blk)
 
-    kv_blocks = tl.load(q2k_num + meta_base)  # int32
+    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
-    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
 
     # ----- base pointers -----
-    qvk_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
+    qvk_off = (b.to(tl.int64) * stride_qz +
+               h.to(tl.int64) * stride_qh)
 
-    Q_ptr = tl.make_block_ptr(base=Q + qvk_off,
+    Q_ptr = tl.make_block_ptr(
-                              shape=(N_CTX, HEAD_DIM),
+        base=Q + qvk_off, shape=(N_CTX, HEAD_DIM),
-                              strides=(stride_qm, stride_qk),
+        strides=(stride_qm, stride_qk),
-                              offsets=(q_blk * BLOCK_M, 0),
+        offsets=(q_blk * BLOCK_M, 0),
-                              block_shape=(BLOCK_M, HEAD_DIM),
+        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
-                              order=(1, 0))
 
-    K_base = tl.make_block_ptr(base=K + qvk_off,
+    K_base = tl.make_block_ptr(
-                               shape=(HEAD_DIM, N_CTX),
+        base=K + qvk_off, shape=(HEAD_DIM, N_CTX),
-                               strides=(stride_kk, stride_kn),
+        strides=(stride_kk, stride_kn),
-                               offsets=(0, 0),
+        offsets=(0, 0),
-                               block_shape=(HEAD_DIM, BLOCK_N),
+        block_shape=(HEAD_DIM, BLOCK_N), order=(0, 1))
-                               order=(0, 1))
 
-    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1,
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
-                                                                              0)
+    V_base = tl.make_block_ptr(
-    V_base = tl.make_block_ptr(base=V + qvk_off,
+        base=V + qvk_off, shape=(N_CTX, HEAD_DIM),
-                               shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vk, stride_vn),
-                               strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
-                               offsets=(0, 0),
+        block_shape=(BLOCK_N, HEAD_DIM), order=v_order)
-                               block_shape=(BLOCK_N, HEAD_DIM),
-                               order=v_order)
 
-    O_ptr = tl.make_block_ptr(base=Out + qvk_off,
+    O_ptr = tl.make_block_ptr(
-                              shape=(N_CTX, HEAD_DIM),
+        base=Out + qvk_off, shape=(N_CTX, HEAD_DIM),
-                              strides=(stride_om, stride_on),
+        strides=(stride_om, stride_on),
-                              offsets=(q_blk * BLOCK_M, 0),
+        offsets=(q_blk * BLOCK_M, 0),
-                              block_shape=(BLOCK_M, HEAD_DIM),
+        block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0))
-                              order=(1, 0))
 
     # ----- accumulators -----
     offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M)
@@ -152,30 +127,23 @@ def _attn_fwd_sparse(
     acc = acc / l_i[:, None]
     tl.store(M + off_hz * N_CTX + offs_m, m_i)
     tl.store(O_ptr, acc.to(Out.type.element_ty))
-
-
 # ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
 
 
+
+
 @triton.jit
-def _attn_bwd_preprocess(
+def _attn_bwd_preprocess(O, DO,  #
-        O,
+                         Delta,  #
-        DO,  #
+                         Z, H, N_CTX,  #
-        Delta,  #
+                         BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr  #
-        Z,
+                         ):
-        H,
-        N_CTX,  #
-        BLOCK_M: tl.constexpr,
-        HEAD_DIM: tl.constexpr  #
-):
     off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
     off_hz = tl.program_id(1)
     off_n = tl.arange(0, HEAD_DIM)
     # load
-    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
-                off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
-    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
-                 off_n[None, :]).to(tl.float32)
     delta = tl.sum(o * do, axis=1)
     # write-back
     tl.store(Delta + off_hz * N_CTX + off_m, delta)
@@ -183,32 +151,19 @@ def _attn_bwd_preprocess(
 
 # The main inner-loop logic for computing dK and dV.
 @triton.jit
-def _attn_bwd_dkdv(
+def _attn_bwd_dkdv(dk, dv,  #
-        dk,
+                   Q, k, v, sm_scale,  #
-        dv,  #
+                   DO,  #
-        Q,
+                   M, D,  #
-        k,
+                   k2q_index, k2q_num, max_q_blks,
-        v,
+                   variable_block_sizes,
-        sm_scale,  #
+                   # shared by Q/K/V/DO.
-        DO,  #
+                   stride_tok, stride_d,  #
-        M,
+                   H, N_CTX, BLOCK_M1: tl.constexpr,  #
-        D,  #
+                   BLOCK_N1: tl.constexpr,  #
-        k2q_index,
+                   HEAD_DIM: tl.constexpr,  #
-        k2q_num,
+                   # Filled in by the wrapper.
-        max_q_blks,
+                   start_n, start_m, num_steps):
-        variable_block_sizes,
-        # shared by Q/K/V/DO.
-        stride_tok,
-        stride_d,  #
-        H,
-        N_CTX,
-        BLOCK_M1: tl.constexpr,  #
-        BLOCK_N1: tl.constexpr,  #
-        HEAD_DIM: tl.constexpr,  #
-        # Filled in by the wrapper.
-    start_n,
-        start_m,
-        num_steps):
     offs_m = start_m + tl.arange(0, BLOCK_M1)
     offs_n = start_n + tl.arange(0, BLOCK_N1)
     offs_k = tl.arange(0, HEAD_DIM)
@@ -217,20 +172,21 @@ def _attn_bwd_dkdv(
     # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
     tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
     step_m = BLOCK_M1
-    kv_blk = tl.program_id(0)  # Q-tile index
+    kv_blk   = tl.program_id(0)          # Q-tile index
-    off_hz = tl.program_id(2)  # fused (batch, head)
+    off_hz  = tl.program_id(2)          # fused (batch, head)
-    b = off_hz // H
+    b       = off_hz // H
-    h = off_hz % H
+    h       = off_hz %  H
     q_tiles = N_CTX // BLOCK_N1
     meta_base = ((b * H + h) * q_tiles + kv_blk)
 
-    q_blocks = tl.load(k2q_num + meta_base)  # int32
+    q_blocks = tl.load(k2q_num  + meta_base)                 # int32
-    q_ptr = k2q_index + meta_base * max_q_blks  # ptr to list
+    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):
+    
-        block_sparse_offset = (tl.load(q_ptr + blk_idx // 2).to(tl.int32) * 2 +
+        
-                               blk_idx % 2) * step_m
+    for blk_idx in range(q_blocks*2):
+        block_sparse_offset = (tl.load(q_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_m
         qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok)
         # Load m before computing qk to reduce pipeline stall.
         offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
@@ -256,32 +212,21 @@ def _attn_bwd_dkdv(
     return dk, dv
 
 
+
 # the main inner-loop logic for computing dQ
 @triton.jit
-def _attn_bwd_dq(
+def _attn_bwd_dq(dq, q, K, V,  #
-        dq,
+                 do, m, D,
-        q,
+                 # shared by Q/K/V/DO.
-        K,
+                 q2k_index, q2k_num, max_kv_blks,
-        V,  #
+                 variable_block_sizes,
-        do,
+                 stride_tok, stride_d,  #
-        m,
+                 H, N_CTX,  #
-        D,
+                 BLOCK_M2: tl.constexpr,  #
-        # shared by Q/K/V/DO.
+                 BLOCK_N2: tl.constexpr,  #
-        q2k_index,
+                 HEAD_DIM: tl.constexpr,
-        q2k_num,
+                 # Filled in by the wrapper.
-        max_kv_blks,
+                 start_m, start_n, num_steps):
-        variable_block_sizes,
-        stride_tok,
-        stride_d,  #
-        H,
-        N_CTX,  #
-        BLOCK_M2: tl.constexpr,  #
-        BLOCK_N2: tl.constexpr,  #
-        HEAD_DIM: tl.constexpr,
-        # Filled in by the wrapper.
-        start_m,
-        start_n,
-        num_steps):
     offs_m = start_m + tl.arange(0, BLOCK_M2)
     offs_n = start_n + tl.arange(0, BLOCK_N2)
     offs_k = tl.arange(0, HEAD_DIM)
@@ -292,27 +237,27 @@ def _attn_bwd_dq(
     # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
     tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
     step_n = BLOCK_N2
-
+    
-    q_blk = tl.program_id(0)  # Q-tile index
+    q_blk   = tl.program_id(0)          # Q-tile index
-    off_hz = tl.program_id(2)  # fused (batch, head)
+    off_hz  = tl.program_id(2)          # fused (batch, head)
-    b = off_hz // H
+    b       = off_hz // H
-    h = off_hz % H
+    h       = off_hz %  H
     q_tiles = N_CTX // BLOCK_M2
     meta_base = ((b * H + h) * q_tiles + q_blk)
 
-    kv_blocks = tl.load(q2k_num + meta_base)  # int32
+    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
-    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
     block_size = tl.load(variable_block_sizes + q_blk)
-
+    
-    for blk_idx in range(kv_blocks * 2):
+    
-        block_sparse_offset = (tl.load(kv_ptr + blk_idx // 2).to(tl.int32) * 2 +
+    for blk_idx in range(kv_blocks*2):
-                               blk_idx % 2) * step_n * stride_tok
+        block_sparse_offset = (tl.load(kv_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_n * stride_tok
         kT = tl.load(kT_ptrs + block_sparse_offset)
         vT = tl.load(vT_ptrs + block_sparse_offset)
         qk = tl.dot(q, kT)
         p = tl.math.exp2(qk - m)
         mask = tl.arange(0, BLOCK_N2) < block_size.to(tl.int32)
-        p = tl.where(mask[None, :], p, 0.0)
+        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])
@@ -324,37 +269,23 @@ def _attn_bwd_dq(
     return dq
 
 
+
 @triton.jit
-def _attn_bwd(
+def _attn_bwd(Q, K, V, sm_scale,  #
-        Q,
+              DO,  #
-        K,
+              DQ, DK, DV,  #
-        V,
+              M, D,
-        sm_scale,  #
+              q2k_index, q2k_num, max_kv_blks,
-        DO,  #
+              k2q_index, k2q_num, max_q_blks,
-        DQ,
+              variable_block_sizes,
-        DK,
+              # shared by Q/K/V/DO.
-        DV,  #
+              stride_z, stride_h, stride_tok, stride_d,  #
-        M,
+              H, N_CTX,  #
-        D,
+              BLOCK_M1: tl.constexpr,  #
-        q2k_index,
+              BLOCK_N1: tl.constexpr,  #
-        q2k_num,
+              BLOCK_M2: tl.constexpr,  #
-        max_kv_blks,
+              BLOCK_N2: tl.constexpr,  #
-        k2q_index,
+              HEAD_DIM: tl.constexpr):
-        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,  #
-        BLOCK_M1: tl.constexpr,  #
-        BLOCK_N1: tl.constexpr,  #
-        BLOCK_M2: tl.constexpr,  #
-        BLOCK_N2: tl.constexpr,  #
-        HEAD_DIM: tl.constexpr):
     LN2 = 0.6931471824645996  # = ln(2)
 
     bhid = tl.program_id(2)
@@ -388,32 +319,20 @@ def _attn_bwd(
     k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
     v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
 
+
     num_steps = N_CTX // BLOCK_M1
 
     dk, dv = _attn_bwd_dkdv(  #
-        dk,
+        dk, dv,  #
-        dv,  #
+        Q, k, v, sm_scale,  #
-        Q,
-        k,
-        v,
-        sm_scale,  #
         DO,  #
-        M,
+        M, D,  #
-        D,  #
+        k2q_index, k2q_num, max_q_blks,
-        k2q_index,
-        k2q_num,
-        max_q_blks,
         variable_block_sizes,
-        stride_tok,
+        stride_tok, stride_d,  #
-        stride_d,  #
+        H, N_CTX,  #
-        H,
+        BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
-        N_CTX,  #
+        start_n, start_m, num_steps  #
-        BLOCK_M1,
-        BLOCK_N1,
-        HEAD_DIM,  #
-        start_n,
-        start_m,
-        num_steps  #
     )
 
     dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
@@ -438,88 +357,54 @@ def _attn_bwd(
     m = m[:, None]
 
     num_steps = N_CTX // BLOCK_N2
-    dq = _attn_bwd_dq(
+    dq = _attn_bwd_dq(dq, q, K, V,  #
-        dq,
+                      do, m, D,  #
-        q,
+                      q2k_index, q2k_num, max_kv_blks,
-        K,
+                      variable_block_sizes,
-        V,  #
+                      stride_tok, stride_d,  #
-        do,
+                      H, N_CTX,  #
-        m,
+                      BLOCK_M2, BLOCK_N2, HEAD_DIM,  #
-        D,  #
+                      start_m, end_n, num_steps  #
-        q2k_index,
+                      )
-        q2k_num,
-        max_kv_blks,
-        variable_block_sizes,
-        stride_tok,
-        stride_d,  #
-        H,
-        N_CTX,  #
-        BLOCK_M2,
-        BLOCK_N2,
-        HEAD_DIM,  #
-        start_m,
-        end_n,
-        num_steps  #
-    )
     # Write back dQ.
     dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
     dq *= LN2
     tl.store(dq_ptrs, dq)
+    
+    
+
 
 
 # ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
-def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
+def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes):
-                                     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 q2k_num.shape[
+    assert T // 64 == q2k_num.shape[-1], f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
-        -1] == T // 64, f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
     o = torch.empty_like(q)
     M = torch.empty((B, H, T), dtype=torch.float32, device=q.device)
 
     grid = lambda _: (triton.cdiv(T, 64), B * H, 1)
-    _attn_fwd_sparse[grid](q,
+    _attn_fwd_sparse[grid](
-                           k,
+        q, k, v, sm_scale,
-                           v,
+        q2k_index, q2k_num, max_kv_blks,
-                           sm_scale,
+        variable_block_sizes,
-                           q2k_index,
+        M, o,
-                           q2k_num,
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
-                           max_kv_blks,
+        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
-                           variable_block_sizes,
+        v.stride(0), v.stride(1), v.stride(2), v.stride(3),
-                           M,
+        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
-                           o,
+        B, H, T,
-                           q.stride(0),
+        HEAD_DIM=D, STAGE=3
-                           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)
 
     return o, M
 
-
+def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num, variable_block_sizes):
-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)
@@ -535,49 +420,30 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
     pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
     delta = torch.empty_like(M)
     _attn_bwd_preprocess[pre_grid](
-        o,
+        o, do,  #
-        do,  #
         delta,  #
-        BATCH,
+        BATCH, N_HEAD, N_CTX,  #
-        N_HEAD,
+        BLOCK_M=PRE_BLOCK, HEAD_DIM=D  #
-        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,
+        q, arg_k, v, sm_scale, do, dq, dk, dv,  #
-        arg_k,
+        M, delta,  #
-        v,
+        q2k_index, q2k_num, max_kv_blks,
-        sm_scale,
+        k2q_index, k2q_num, max_q_blks,
-        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(0), q.stride(1), q.stride(2), q.stride(3),  #
-        q.stride(1),
+        N_HEAD, N_CTX,  #
-        q.stride(2),
+        BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1,  #
-        q.stride(3),  #
+        BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2,  #
-        N_HEAD,
-        N_CTX,  #
-        BLOCK_M1=BLOCK_M1,
-        BLOCK_N1=BLOCK_N1,  #
-        BLOCK_M2=BLOCK_M2,
-        BLOCK_N2=BLOCK_N2,  #
         HEAD_DIM=D  #
     )
 
     return dq, dk, dv
+
+

csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/index.py

@@ -0,0 +1,152 @@
+
+## 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/fastvideo_kernel/src/fastvideo_kernel/vmoba.py

@@ -0,0 +1,868 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapt from https://github.com/KwaiVGI/VMoBA/blob/main/src/vmoba.py
+
+import random
+import time
+import os
+import torch
+from typing import Tuple
+try:
+    from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
+    from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
+except ImportError:
+    def _unsupported(*args, **kwargs):
+        raise ImportError("flash-attn is not installed. Please install it, e.g., `pip install flash-attn`.")
+    _flash_attn_varlen_forward = _unsupported
+    _flash_attn_varlen_backward = _unsupported
+    flash_attn_varlen_func = _unsupported
+    
+from functools import lru_cache
+from einops import rearrange
+
+
+@lru_cache(maxsize=16)
+def calc_chunks(cu_seqlen, moba_chunk_size):
+    """
+    Calculate chunk boundaries.
+
+    For vision tasks we include all chunks (even the last one which might be shorter)
+    so that every chunk can be selected.
+    """
+    batch_sizes = cu_seqlen[1:] - cu_seqlen[:-1]
+    batch_num_chunk = (batch_sizes + (moba_chunk_size - 1)) // moba_chunk_size
+    cu_num_chunk = torch.ones(
+        batch_num_chunk.numel() + 1,
+        device=cu_seqlen.device,
+        dtype=batch_num_chunk.dtype,
+    )
+    cu_num_chunk[1:] = batch_num_chunk.cumsum(dim=0)
+    num_chunk = cu_num_chunk[-1]
+    chunk_sizes = torch.full(
+        (num_chunk + 1,), moba_chunk_size, dtype=torch.int32, device=cu_seqlen.device
+    )
+    chunk_sizes[0] = 0
+    batch_last_chunk_size = batch_sizes - (batch_num_chunk - 1) * moba_chunk_size
+    chunk_sizes[cu_num_chunk[1:]] = batch_last_chunk_size
+    cu_chunk = chunk_sizes.cumsum(dim=-1, dtype=torch.int32)
+    chunk_to_batch = torch.zeros(
+        (num_chunk,), dtype=torch.int32, device=cu_seqlen.device
+    )
+    chunk_to_batch[cu_num_chunk[1:-1]] = 1
+    chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32)
+
+    # Do not filter out any chunk
+    filtered_chunk_indices = torch.arange(
+        num_chunk, device=cu_seqlen.device, dtype=torch.int32
+    )
+    num_filtered_chunk = num_chunk
+
+    return cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch
+
+
+# --- Threshold Selection Helper Functions ---
+
+def _select_threshold_query_head(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects chunks for each <query, head> pair based on threshold.
+    Normalization and sorting happen along the chunk dimension (dim=0).
+    """
+    C, H, S = gate.shape
+    eps = 1e-6
+
+    # LSE‐style normalization per <head, query> (across chunks)
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
+
+    row_min = gate_min_val.amin(dim=0)        # (H, S)
+    row_max = gate_masked.amax(dim=0)         # (H, S)
+    denom = row_max - row_min
+    denom = torch.where(denom <= eps, torch.ones_like(denom), denom) # avoid divide‑by‑zero
+
+    gate_norm = (gate - row_min.unsqueeze(0)) / denom.unsqueeze(0)
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
+
+    # 1) pull out the self‐chunk’s normalized weight for each <head,seq>
+    self_norm = (gate_norm * gate_self_chunk_mask).sum(dim=0)  # (H, S)
+
+    # 2) compute how much more normalized weight we need beyond self
+    total_norm_sum = gate_norm.sum(dim=0) # (H, S)
+    remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps)  # (H, S)
+    remain_ratio = torch.clamp(remain_ratio, min=0.0) # if already ≥ thresh, no extra needed
+
+    # 3) zero out the self‐chunk in a copy, so we only sort “others”
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0
+
+    # 4) sort the other chunks by descending norm, per <head,seq>
+    sorted_norm, sorted_idx = torch.sort(others_norm, descending=True, dim=0) # (C, H, S)
+
+    # 5) cumulative‑sum the sorted norms per <head,seq>
+    cumsum_others = sorted_norm.cumsum(dim=0) # (C, H, S)
+
+    # 6) for each <head,seq>, find the smallest k where cumsum_ratio ≥ remain_ratio
+    ratio = cumsum_others / (total_norm_sum.unsqueeze(0) + eps) # (C, H, S)
+    cond = ratio >= remain_ratio.unsqueeze(0) # (C, H, S) boolean mask
+    any_cond = cond.any(dim=0) # (H, S)
+    # Find the index of the first True value along dim 0. If none, use C-1.
+    cutoff = torch.where(any_cond, cond.float().argmax(dim=0), torch.full_like(any_cond, fill_value=C - 1)) # (H, S)
+
+    # 7) build a mask in sorted order up to that cutoff
+    idx_range = torch.arange(C, device=gate.device).view(-1, 1, 1) # (C, 1, 1)
+    sorted_mask = idx_range <= cutoff.unsqueeze(0) # (C, H, S)
+
+    # 8) scatter it back to original chunk order
+    others_mask = torch.zeros_like(gate, dtype=torch.bool)
+    others_mask.scatter_(0, sorted_idx, sorted_mask)
+
+    # 9) finally, include every self‐chunk plus all selected others
+    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)
+
+    return final_gate_mask
+
+
+def _select_threshold_block(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects <query, head> pairs for each block based on threshold.
+    Normalization and sorting happen across the head and sequence dimensions (dim=1, 2).
+    """
+    C, H, S = gate.shape
+    HS = H * S
+    eps = 1e-6
+
+    # LSE‐style normalization per block (across heads and queries)
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
+
+    block_max = gate_masked.amax(dim=(1, 2), keepdim=True) # (C, 1, 1)
+    block_min = gate_min_val.amin(dim=(1, 2), keepdim=True) # (C, 1, 1)
+    block_denom = block_max - block_min
+    block_denom = torch.where(block_denom <= eps, torch.ones_like(block_denom), block_denom) # (C, 1, 1)
+
+    gate_norm = (gate - block_min) / block_denom # (C, H, S)
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
+
+    # 1) identify normalized weights of entries that *are* self-chunks (from query perspective)
+    self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S)
+    # Sum these weights *per block*
+    self_norm_sum_per_block = self_norm_entries.sum(dim=(1, 2)) # (C,)
+
+    # 2) compute how much more normalized weight each block needs beyond its self-chunk contributions
+    total_norm_sum_per_block = gate_norm.sum(dim=(1, 2)) # (C,)
+    remain_ratio = simsum_threshold - self_norm_sum_per_block / (total_norm_sum_per_block + eps) # (C,)
+    remain_ratio = torch.clamp(remain_ratio, min=0.0) # (C,)
+
+    # 3) zero out the self‐chunk entries in a copy, so we only sort “others”
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries
+
+    # 4) sort the other <head, seq> pairs by descending norm, per block
+    others_flat = others_norm.contiguous().view(C, HS) # (C, H*S)
+    sorted_others_flat, sorted_indices_flat = torch.sort(others_flat, dim=1, descending=True) # (C, H*S)
+
+    # 5) cumulative‑sum the sorted norms per block
+    cumsum_others_flat = sorted_others_flat.cumsum(dim=1) # (C, H*S)
+
+    # 6) for each block, find the smallest k where cumsum_ratio ≥ remain_ratio
+    ratio_flat = cumsum_others_flat / (total_norm_sum_per_block.unsqueeze(1) + eps) # (C, H*S)
+    cond_flat = ratio_flat >= remain_ratio.unsqueeze(1) # (C, H*S) boolean mask
+    any_cond = cond_flat.any(dim=1) # (C,)
+    # Find the index of the first True value along dim 1. If none, use HS-1.
+    cutoff_flat = torch.where(any_cond, cond_flat.float().argmax(dim=1), torch.full_like(any_cond, fill_value=HS - 1)) # (C,)
+
+    # 7) build a mask in sorted order up to that cutoff per block
+    idx_range_flat = torch.arange(HS, device=gate.device).unsqueeze(0) # (1, H*S)
+    sorted_mask_flat = idx_range_flat <= cutoff_flat.unsqueeze(1) # (C, H*S)
+
+    # 8) scatter it back to original <head, seq> order per block
+    others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C, H*S)
+    others_mask_flat.scatter_(1, sorted_indices_flat, sorted_mask_flat)
+    others_mask = others_mask_flat.view(C, H, S) # (C, H, S)
+
+    # 9) finally, include every self‐chunk entry plus all selected others
+    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)
+
+    return final_gate_mask
+
+
+def _select_threshold_overall(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects <chunk, query, head> triplets globally based on threshold.
+    Normalization and sorting happen across all valid entries.
+    """
+    C, H, S = gate.shape
+    CHS = C * H * S
+    eps = 1e-6
+
+    # LSE‐style normalization globally across all valid entries
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min
+
+    overall_max = gate_masked.max() # scalar
+    overall_min = gate_min_val.min() # scalar
+    overall_denom = overall_max - overall_min
+    overall_denom = torch.where(overall_denom <= eps, torch.tensor(1.0, device=gate.device, dtype=gate.dtype), overall_denom)
+
+    gate_norm = (gate - overall_min) / overall_denom # (C, H, S)
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S)
+
+    # 1) identify normalized weights of entries that *are* self-chunks
+    self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S)
+    # Sum these weights globally
+    self_norm_sum_overall = self_norm_entries.sum() # scalar
+
+    # 2) compute how much more normalized weight is needed globally beyond self-chunk contributions
+    total_norm_sum_overall = gate_norm.sum() # scalar
+    remain_ratio = simsum_threshold - self_norm_sum_overall / (total_norm_sum_overall + eps) # scalar
+    remain_ratio = torch.clamp(remain_ratio, min=0.0) # scalar
+
+    # 3) zero out the self‐chunk entries in a copy, so we only sort “others”
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries
+
+    # 4) sort all other entries by descending norm, globally
+    others_flat = others_norm.flatten() # (C*H*S,)
+    valid_others_mask_flat = valid_gate_mask.flatten() & ~gate_self_chunk_mask.flatten() # Mask for valid, non-self entries
+    
+    # Only sort the valid 'other' entries
+    valid_others_indices = torch.where(valid_others_mask_flat)[0]
+    valid_others_values = others_flat[valid_others_indices]
+
+    sorted_others_values, sort_perm = torch.sort(valid_others_values, descending=True) # (N_valid_others,)
+    sorted_original_indices = valid_others_indices[sort_perm] # Original indices in C*H*S space, sorted by value
+
+    # 5) cumulative‑sum the sorted valid 'other' norms globally
+    cumsum_others_values = sorted_others_values.cumsum(dim=0) # (N_valid_others,)
+
+    # 6) find the smallest k where cumsum_ratio ≥ remain_ratio globally
+    ratio_values = cumsum_others_values / (total_norm_sum_overall + eps) # (N_valid_others,)
+    cond_values = ratio_values >= remain_ratio # (N_valid_others,) boolean mask
+    any_cond = cond_values.any() # scalar
+    
+    # Find the index of the first True value in the *sorted* list. If none, use all valid others.
+    cutoff_idx_in_sorted = torch.where(
+        any_cond, 
+        cond_values.float().argmax(dim=0), 
+        torch.tensor(len(sorted_others_values) - 1, device=gate.device, dtype=torch.long)
+    )
+
+    # 7) build a mask selecting the top-k others based on the cutoff
+    # Select the original indices corresponding to the top entries in the sorted list
+    selected_other_indices = sorted_original_indices[:cutoff_idx_in_sorted + 1]
+
+    # 8) create the mask in the original flat shape
+    others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C*H*S,)
+    if selected_other_indices.numel() > 0: # Check if any 'other' indices were selected
+        others_mask_flat[selected_other_indices] = True
+    others_mask = others_mask_flat.view(C, H, S) # (C, H, S)
+
+    # 9) finally, include every self‐chunk entry plus all selected others
+    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)
+
+    return final_gate_mask
+
+
+def _select_threshold_head_global(
+    gate: torch.Tensor,
+    valid_gate_mask: torch.Tensor,
+    gate_self_chunk_mask: torch.Tensor,
+    simsum_threshold: float
+) -> torch.Tensor:
+    """
+    Selects <chunk, query> globally for each head based on threshold.
+    """
+    C, H, S = gate.shape
+    eps = 1e-6
+
+    # 1) LSE‐style normalization per head (across chunks and sequence dims)
+    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
+    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)
+
+    max_per_head = gate_masked.amax(dim=(0, 2), keepdim=True)  # (1, H, 1)
+    min_per_head = gate_min_val.amin(dim=(0, 2), keepdim=True)  # (1, H, 1)
+    denom = max_per_head - min_per_head
+    denom = torch.where(denom <= eps, torch.ones_like(denom), denom)
+
+    gate_norm = (gate - min_per_head) / denom
+    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0)  # (C, H, S)
+
+    # 2) sum normalized self‐chunk contributions per head
+    self_norm_sum = (gate_norm * gate_self_chunk_mask).sum(dim=(0, 2))  # (H,)
+
+    # 3) total normalized sum per head
+    total_norm_sum = gate_norm.sum(dim=(0, 2))  # (H,)
+
+    # 4) how much more normalized weight needed per head
+    remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps)  # (H,)
+    remain_ratio = torch.clamp(remain_ratio, min=0.0)
+
+    # 5) zero out self‐chunk entries to focus on "others"
+    others_norm = gate_norm.clone()
+    others_norm[gate_self_chunk_mask] = 0.0  # (C, H, S)
+
+    # 6) flatten chunk and sequence dims, per head
+    CS = C * S
+    others_flat = others_norm.permute(1, 0, 2).reshape(H, CS)  # (H, C*S)
+    valid_flat  = (valid_gate_mask & ~gate_self_chunk_mask) \
+                   .permute(1, 0, 2).reshape(H, CS)           # (H, C*S)
+
+    # 7) vectorized selection of “others” per head
+    masked_flat = torch.where(valid_flat, others_flat, torch.zeros_like(others_flat))
+    sorted_vals, sorted_idx = torch.sort(masked_flat, dim=1, descending=True)  # (H, C*S)
+
+    cumsum_vals = sorted_vals.cumsum(dim=1)                                    # (H, C*S)
+    ratio_vals  = cumsum_vals / (total_norm_sum.unsqueeze(1) + eps)            # (H, C*S)
+    cond        = ratio_vals >= remain_ratio.unsqueeze(1)                      # (H, C*S)
+
+    has_cutoff = cond.any(dim=1)                                               # (H,)
+    default    = torch.full((H,), CS - 1, device=gate.device, dtype=torch.long)
+    cutoff     = torch.where(has_cutoff, cond.float().argmax(dim=1), default)  # (H,)
+
+    idx_range   = torch.arange(CS, device=gate.device).unsqueeze(0)            # (1, C*S)
+    sorted_mask = idx_range <= cutoff.unsqueeze(1)                             # (H, C*S)
+
+    selected_flat = torch.zeros_like(valid_flat)                               # (H, C*S)
+    selected_flat.scatter_(1, sorted_idx, sorted_mask)                         # (H, C*S)
+
+    # 8) reshape selection mask back to (C, H, S)
+    others_mask = selected_flat.reshape(H, C, S).permute(1, 0, 2)               # (C, H, S)
+
+    # 9) include self‐chunks plus selected others, and obey valid mask
+    final_gate_mask = valid_gate_mask & (gate_self_chunk_mask | others_mask)
+
+    return final_gate_mask
+
+
+class MixedAttention(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        self_attn_cu_seqlen,
+        moba_q,
+        moba_kv,
+        moba_cu_seqlen_q,
+        moba_cu_seqlen_kv,
+        max_seqlen,
+        moba_chunk_size,
+        moba_q_sh_indices,
+    ):
+        ctx.max_seqlen = max_seqlen
+        ctx.moba_chunk_size = moba_chunk_size
+        ctx.softmax_scale = softmax_scale = q.shape[-1] ** (-0.5)
+
+        # Non-causal self-attention branch
+        # return out, softmax_lse, S_dmask, rng_state
+        self_attn_out_sh, self_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=self_attn_cu_seqlen,
+            cu_seqlens_k=self_attn_cu_seqlen,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+        )
+        # MOBA attention branch (non-causal)
+        moba_attn_out, moba_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
+            q=moba_q,
+            k=moba_kv[:, 0],
+            v=moba_kv[:, 1],
+            cu_seqlens_q=moba_cu_seqlen_q,
+            cu_seqlens_k=moba_cu_seqlen_kv,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=moba_chunk_size,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+        )
+
+        self_attn_lse_sh = self_attn_lse_hs.t().contiguous()
+        moba_attn_lse = moba_attn_lse_hs.t().contiguous()
+
+        output = torch.zeros((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+        output_2d = output.view(-1, q.shape[2])
+
+        max_lse_1d = self_attn_lse_sh.view(-1)
+        max_lse_1d = max_lse_1d.index_reduce(
+            0, moba_q_sh_indices, moba_attn_lse.view(-1), "amax"
+        )
+        self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh)
+        moba_attn_lse = (
+            moba_attn_lse.view(-1)
+            .sub(max_lse_1d.index_select(0, moba_q_sh_indices))
+            .reshape_as(moba_attn_lse)
+        )
+
+        mixed_attn_se_sh = self_attn_lse_sh.exp()
+        moba_attn_se = moba_attn_lse.exp()
+
+        mixed_attn_se_sh.view(-1).index_add_(
+            0, moba_q_sh_indices, moba_attn_se.view(-1)
+        )
+        mixed_attn_lse_sh = mixed_attn_se_sh.log()
+
+        # Combine self-attention output
+        factor = (self_attn_lse_sh - mixed_attn_lse_sh).exp()  # [S, H]
+        self_attn_out_sh = self_attn_out_sh * factor.unsqueeze(-1)
+        output_2d += self_attn_out_sh.reshape_as(output_2d)
+
+        # Combine MOBA attention output
+        mixed_attn_lse = (
+            mixed_attn_lse_sh.view(-1)
+            .index_select(0, moba_q_sh_indices)
+            .view_as(moba_attn_lse)
+        )
+        factor = (moba_attn_lse - mixed_attn_lse).exp()  # [S, H]
+        moba_attn_out = moba_attn_out * factor.unsqueeze(-1)
+        raw_attn_out = moba_attn_out.view(-1, moba_attn_out.shape[-1])
+        output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out)
+        output = output.to(q.dtype)
+        mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(mixed_attn_se_sh)
+        ctx.save_for_backward(
+            output,
+            mixed_attn_lse_sh,
+            q,
+            k,
+            v,
+            self_attn_cu_seqlen,
+            moba_q,
+            moba_kv,
+            moba_cu_seqlen_q,
+            moba_cu_seqlen_kv,
+            moba_q_sh_indices,
+        )
+
+        return output
+
+    @staticmethod
+    def backward(ctx, d_output):
+
+        max_seqlen = ctx.max_seqlen
+        moba_chunk_size = ctx.moba_chunk_size
+        softmax_scale = ctx.softmax_scale
+
+        (
+            output,
+            mixed_attn_vlse_sh,
+            q,
+            k,
+            v,
+            self_attn_cu_seqlen,
+            moba_q,
+            moba_kv,
+            moba_cu_seqlen_q,
+            moba_cu_seqlen_kv,
+            moba_q_sh_indices,
+        ) = ctx.saved_tensors
+
+        d_output = d_output.contiguous()
+
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        _ = _flash_attn_varlen_backward(
+            dout=d_output,
+            q=q,
+            k=k,
+            v=v,
+            out=output,
+            softmax_lse=mixed_attn_vlse_sh.t().contiguous(),
+            dq=dq,
+            dk=dk,
+            dv=dv,
+            cu_seqlens_q=self_attn_cu_seqlen,
+            cu_seqlens_k=self_attn_cu_seqlen,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=True,
+            window_size_left=-1,
+            window_size_right=-1
+        )
+
+        headdim = q.shape[-1]
+        d_moba_output = (
+            d_output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)
+        )
+        moba_output = (
+            output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)
+        )
+
+        mixed_attn_vlse = (
+            mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1)
+        )
+
+        dmq = torch.empty_like(moba_q)
+        dmkv = torch.empty_like(moba_kv)
+        _ = _flash_attn_varlen_backward(
+            dout=d_moba_output,
+            q=moba_q,
+            k=moba_kv[:, 0],
+            v=moba_kv[:, 1],
+            out=moba_output,
+            softmax_lse=mixed_attn_vlse,
+            dq=dmq,
+            dk=dmkv[:,0],
+            dv=dmkv[:,1],
+            cu_seqlens_q=moba_cu_seqlen_q,
+            cu_seqlens_k=moba_cu_seqlen_kv,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=moba_chunk_size,
+            softmax_scale=softmax_scale,
+            causal=False,
+            dropout_p=0.0,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=True,
+            window_size_left=-1,
+            window_size_right=-1
+        )
+
+        return dq, dk, dv, None, dmq, dmkv, None, None, None, None, None
+
+
+def moba_attn_varlen(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    max_seqlen: int,
+    moba_chunk_size: int,
+    moba_topk: int,
+    select_mode: str = 'threshold', # "topk" or "threshold"
+    simsum_threshold: float = 0.25,
+    threshold_type: str = 'query_head',
+) -> torch.Tensor:
+    """
+    Accelerated MOBA attention for vision tasks with proper LSE normalization.
+
+    This version:
+      - Splits KV into chunks.
+      - For each query head, selects the top-k relevant KV chunks (including the self chunk)
+        by amplifying the diagonal (self-chunk) logits.
+      - Aggregates the attention outputs from the selected chunks using a log-sum-exp
+        reduction so that attending to each query over the selected chunks is equivalent
+        to the original algorithm.
+    """
+    # Stack keys and values.
+    kv = torch.stack((k, v), dim=1)
+    seqlen, num_head, head_dim = q.shape
+
+    # Compute chunk boundaries.
+    cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(
+        cu_seqlens, moba_chunk_size
+    )
+
+    self_attn_cu_seqlen = cu_chunk
+
+    # Update top-k selection to include the self chunk.
+    moba_topk = min(moba_topk, num_filtered_chunk)
+
+    # --- Build filtered KV from chunks ---
+    chunk_starts = cu_chunk[filtered_chunk_indices]          # [num_filtered_chunk]
+    chunk_ends = cu_chunk[filtered_chunk_indices + 1]          # [num_filtered_chunk]
+    chunk_lengths = chunk_ends - chunk_starts                  # [num_filtered_chunk]
+    max_chunk_len = int(chunk_lengths.max().item())
+
+    range_tensor = torch.arange(max_chunk_len, device=kv.device, dtype=chunk_starts.dtype).unsqueeze(0)
+    indices = chunk_starts.unsqueeze(1) + range_tensor
+    indices = torch.clamp(indices, max=kv.shape[0] - 1)
+    valid_mask = range_tensor < chunk_lengths.unsqueeze(1)
+    gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, *kv.shape[1:])
+    gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(gathered)
+
+    # Compute key_gate_weight over valid tokens.
+    key_values = gathered[:, :, 0].float()  # [num_filtered_chunk, max_chunk_len, num_head, head_dim]
+    valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1)
+    key_sum = (key_values * valid_mask_exp).sum(dim=1)
+    divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1)
+    key_gate_weight = key_sum / divisor  # [num_filtered_chunk, num_head, head_dim]
+
+    # Compute gate logits between key_gate_weight and queries.
+    q_float = q.float()
+    # gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float)  # [num_filtered_chunk, num_head, seqlen]
+    gate = torch.bmm(key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2)
+
+    # Amplify the diagonal (self chunk) contributions.
+    gate_seq_idx = torch.arange(seqlen, device=q.device, dtype=torch.int32).unsqueeze(0).expand(num_filtered_chunk, seqlen)
+    chunk_start = cu_chunk[filtered_chunk_indices]  # [num_filtered_chunk]
+    chunk_end = cu_chunk[filtered_chunk_indices + 1]  # [num_filtered_chunk]
+    gate_self_chunk_mask = ((gate_seq_idx >= chunk_start.unsqueeze(1)) &
+                            (gate_seq_idx < chunk_end.unsqueeze(1))).unsqueeze(1).expand(-1, num_head, -1)
+    amplification_factor = 1e9  # Example factor; adjust as needed.
+    origin_gate = gate.clone()
+    gate = gate.clone()
+    if select_mode == "topk":
+        gate[gate_self_chunk_mask] += amplification_factor
+
+    # Exclude positions that are outside the valid batch boundaries.
+    batch_starts = cu_seqlens[chunk_to_batch[filtered_chunk_indices]]
+    batch_ends = cu_seqlens[chunk_to_batch[filtered_chunk_indices] + 1]
+    gate_batch_start_mask = gate_seq_idx < batch_starts.unsqueeze(1)
+    gate_batch_end_mask = gate_seq_idx >= batch_ends.unsqueeze(1)
+    gate_inf_mask = gate_batch_start_mask | gate_batch_end_mask
+    gate.masked_fill_(gate_inf_mask.unsqueeze(1), -float("inf"))
+
+    if select_mode == 'topk':
+        # We amplify self‐chunk in gate already, so self entries will rank highest.
+        valid_gate_mask = gate != -float("inf")
+        if threshold_type == 'query_head':
+            # === per‐<head,seq> top-k across chunks (original behavior) ===
+            # gate: (C, H, S)
+            _, gate_topk_idx = torch.topk(gate, k=moba_topk, dim=0, largest=True, sorted=False)
+            gate_idx_mask = torch.zeros_like(gate, dtype=torch.bool)
+            gate_idx_mask.scatter_(0, gate_topk_idx, True)
+            gate_mask = valid_gate_mask & gate_idx_mask
+        elif threshold_type == 'overall':
+            # === global top-k across all (chunk, head, seq) entries ===
+            C, H, S = gate.shape
+            flat_gate = gate.flatten()
+            flat_mask = valid_gate_mask.flatten()
+            flat_gate_masked = torch.where(flat_mask, flat_gate, -float("inf"))
+            # pick topk global entries
+            vals, idx = torch.topk(flat_gate_masked, k=moba_topk * H * S, largest=True, sorted=False)
+            others_mask_flat = torch.zeros_like(flat_mask, dtype=torch.bool)
+            others_mask_flat[idx] = True
+            gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view(gate.shape)
+        elif threshold_type == 'head_global':
+            # per-head top-k across all chunks and sequence positions
+            C, H, S = gate.shape
+            CS = C * S
+            flat_gate  = gate.permute(1, 0, 2).reshape(H, CS)
+            flat_valid = valid_gate_mask.permute(1, 0, 2).reshape(H, CS)
+            flat_gate_masked = torch.where(flat_valid, flat_gate, torch.full_like(flat_gate, -float('inf')))
+            # pick top-k indices per head
+            _, topk_idx = torch.topk(flat_gate_masked, k=moba_topk * S, dim=1, largest=True, sorted=False)
+            gate_idx_flat = torch.zeros_like(flat_valid, dtype=torch.bool)
+            gate_idx_flat.scatter_(1, topk_idx, True)
+            gate_mask = gate_idx_flat.reshape(H, C, S).permute(1, 0, 2)
+        else:
+            raise ValueError(
+                f"Invalid threshold_type for topk: {threshold_type}. "
+                "Choose 'query_head', 'block', or 'overall'."
+            )
+    elif select_mode == 'threshold':
+        # Delegate to the specific thresholding function
+        valid_gate_mask = gate != -float("inf")  # (num_chunk, num_head, seqlen)
+        if threshold_type == 'query_head':
+            gate_mask = _select_threshold_query_head(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        elif threshold_type == 'block':
+            gate_mask = _select_threshold_block(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        elif threshold_type == 'overall':
+            gate_mask = _select_threshold_overall(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        elif threshold_type == 'head_global':
+            gate_mask = _select_threshold_head_global(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold)
+        else:
+            raise ValueError(f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'.")
+    else:
+        raise ValueError(f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'.")
+    
+    # eliminate self_chunk in MoBA branch
+    gate_mask = gate_mask & ~gate_self_chunk_mask
+    # if gate_mask is all false, perform flash_attn instead
+    if gate_mask.sum() == 0:
+        return flash_attn_varlen_func(
+            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=False
+        )
+
+    # Determine which query positions are selected.
+    # nonzero_indices has shape [N, 3] where each row is [chunk_index, head_index, seq_index].
+    moba_q_indices = gate_mask.reshape(gate_mask.shape[0], -1).nonzero(as_tuple=True)[-1] # [(h s k)]
+    moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (moba_q_indices // seqlen)
+    moba_q = rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1)
+
+    # Build cumulative sequence lengths for the selected queries.
+    moba_seqlen_q = gate_mask.sum(dim=-1).flatten()
+    q_zero_mask = moba_seqlen_q == 0
+    valid_expert_mask = ~q_zero_mask
+    if q_zero_mask.sum() > 0:
+        moba_seqlen_q = moba_seqlen_q[valid_expert_mask]
+    moba_cu_seqlen_q = torch.cat(
+        (
+            torch.tensor([0], device=q.device, dtype=moba_seqlen_q.dtype),
+            moba_seqlen_q.cumsum(dim=0),
+        ),
+        dim=0,
+    ).to(torch.int32)
+
+    # Rearrange gathered KV for the MOBA branch.
+    experts_tensor = rearrange(gathered, "nc cl two h d -> (nc h) cl two d")
+    valid_expert_lengths = chunk_lengths.unsqueeze(1).expand(num_filtered_chunk, num_head).reshape(-1).to(torch.int32)
+    if q_zero_mask.sum() > 0:
+        experts_tensor = experts_tensor[valid_expert_mask]
+        valid_expert_lengths = valid_expert_lengths[valid_expert_mask]
+
+    seq_range = torch.arange(experts_tensor.shape[1], device=experts_tensor.device).unsqueeze(0)
+    mask = seq_range < valid_expert_lengths.unsqueeze(1)
+    moba_kv = experts_tensor[mask]  # Shape: ((nc h cl_valid) two d)
+    moba_kv = moba_kv.unsqueeze(2)   # Shape: ((nc h cl_valid) two 1 d)
+
+    moba_cu_seqlen_kv = torch.cat(
+        [torch.zeros(1, device=experts_tensor.device, dtype=torch.int32),
+         valid_expert_lengths.cumsum(dim=0)],
+        dim=0,
+    ).to(torch.int32)
+
+    assert (
+        moba_cu_seqlen_kv.shape == moba_cu_seqlen_q.shape
+    ), f"Mismatch between moba_cu_seqlen_kv.shape and moba_cu_seqlen_q.shape: {moba_cu_seqlen_kv.shape} vs {moba_cu_seqlen_q.shape}"
+
+    return MixedAttention.apply(
+        q,
+        k,
+        v,
+        self_attn_cu_seqlen,
+        moba_q,
+        moba_kv,
+        moba_cu_seqlen_q,
+        moba_cu_seqlen_kv,
+        max_seqlen,
+        moba_chunk_size,
+        moba_q_sh_indices,
+    )
+
+
+def process_moba_input(
+    x,
+    patch_resolution,
+    chunk_size,
+):
+    """
+    Process inputs for the attention function.
+
+    Args:
+        x (torch.Tensor): Input tensor with shape [batch_size, num_patches, num_heads, head_dim].
+        patch_resolution (tuple): Tuple containing the patch resolution (t, h, w).
+        chunk_size (int): Size of the chunk. (maybe tuple or int, according to chunk type)
+
+    Returns:
+        torch.Tensor: Processed input tensor.
+    """
+    if isinstance(chunk_size, float) or isinstance(chunk_size, int):
+        moba_chunk_size = int(chunk_size * patch_resolution[1] * patch_resolution[2])
+    else:
+        assert isinstance(chunk_size, (Tuple, list)), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}"
+        if len(chunk_size) == 2:
+            assert patch_resolution[1] % chunk_size[0] == 0 and patch_resolution[2] % chunk_size[1] == 0, f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}"
+            nch, ncw = patch_resolution[1] // chunk_size[0], patch_resolution[2] // chunk_size[1]
+            x = rearrange(x, "b (t nch ch ncw cw) n d -> b (nch ncw t ch cw) n d", t=patch_resolution[0], nch=nch, ncw=ncw, ch=chunk_size[0], cw=chunk_size[1])
+            moba_chunk_size = patch_resolution[0] * chunk_size[0] * chunk_size[1]
+        elif len(chunk_size) == 3:
+            assert patch_resolution[0] % chunk_size[0] == 0 and patch_resolution[1] % chunk_size[1] == 0 and patch_resolution[2] % chunk_size[2] == 0, f"patch_resolution {patch_resolution} should be divisible by 3d chunk_size {chunk_size}"
+            nct, nch, ncw = patch_resolution[0] // chunk_size[0], patch_resolution[1] // chunk_size[1], patch_resolution[2] // chunk_size[2]
+            x = rearrange(x, "b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d", nct=nct, nch=nch, ncw=ncw, ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2])
+            moba_chunk_size = chunk_size[0] * chunk_size[1] * chunk_size[2]
+        else:
+            raise ValueError(f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}")
+        
+    return x, moba_chunk_size
+
+
+def process_moba_output(
+    x,
+    patch_resolution,
+    chunk_size,
+):
+    if isinstance(chunk_size, float) or isinstance(chunk_size, int):
+        pass
+    elif len(chunk_size) == 2:
+        x = rearrange(x, "b (nch ncw t ch cw) n d -> b (t nch ch ncw cw) n d", nch=patch_resolution[1] // chunk_size[0], ncw=patch_resolution[2] // chunk_size[1], t=patch_resolution[0], ch=chunk_size[0], cw=chunk_size[1])
+    elif len(chunk_size) == 3:
+        x = rearrange(x, "b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d", nct=patch_resolution[0] // chunk_size[0], nch=patch_resolution[1] // chunk_size[1], ncw=patch_resolution[2] // chunk_size[2], ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2])
+
+    return x
+
+
+# TEST
+def generate_data(batch_size, seqlen, num_head, head_dim, dtype):
+    random.seed(0)
+    torch.manual_seed(0)
+    torch.cuda.manual_seed(0)
+    device = torch.cuda.current_device()
+
+    q = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
+    k = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
+    v = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda')
+    print(f"q.shape: {q.shape}, k.shape: {k.shape}, v.shape: {v.shape}")
+    cu_seqlens = torch.arange(0, q.shape[0] * q.shape[1] + 1, q.shape[1], dtype=torch.int32, device='cuda')
+    max_seqlen = q.shape[1]
+    q = rearrange(q, "b s ... -> (b s) ...")
+    k = rearrange(k, "b s ... -> (b s) ...")
+    v = rearrange(v, "b s ... -> (b s) ...")
+
+    return q, k, v, cu_seqlens, max_seqlen
+
+
+def test_attn_varlen_moba_speed(batch, head, seqlen, head_dim, moba_chunk_size, moba_topk, dtype=torch.bfloat16, select_mode='threshold', simsum_threshold=0.25, threshold_type='query_head'):
+    """Speed test comparing flash_attn vs moba_attention"""
+    # Get data
+    q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype)
+    print(f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}")
+    vo_grad = torch.randn_like(q)
+    
+    # Warmup
+    warmup_iters = 3
+    perf_test_iters = 10
+
+    # Warmup
+    for _ in range(warmup_iters):
+        o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False)
+        torch.autograd.backward(o, vo_grad)
+    
+    torch.cuda.synchronize()
+    start_flash = time.perf_counter()
+    for _ in range(perf_test_iters):
+        o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False)
+        torch.autograd.backward(o, vo_grad)
+        
+    torch.cuda.synchronize()
+    time_flash = (time.perf_counter() - start_flash) / perf_test_iters * 1000
+
+    # Warmup
+    for _ in range(warmup_iters):
+        om = moba_attn_varlen(q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type)
+        torch.autograd.backward(om, vo_grad)
+
+    torch.cuda.synchronize()
+    start_moba = time.perf_counter()
+    for _ in range(perf_test_iters):
+        om = moba_attn_varlen(q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type)
+        torch.autograd.backward(om, vo_grad)
+    
+    torch.cuda.synchronize()
+    time_moba = (time.perf_counter() - start_moba) / perf_test_iters * 1000
+    
+    print(f"Flash: {time_flash:.2f}ms, MoBA: {time_moba:.2f}ms")
+    print(f"Speedup:  {time_flash / time_moba:.2f}x")
+
+
+
+if __name__ == "__main__":
+    """
+    CUDA_VISIBLE_DEVICES=1 \
+    python -u csrc/attn/vmoba_attn/vmoba/vmoba.py
+    """
+    test_attn_varlen_moba_speed(batch=1, head=12, seqlen=32760, head_dim=128, moba_chunk_size=32760 // 3 // 6 // 4, moba_topk=3, select_mode='threshold', simsum_threshold=0.3, threshold_type='query_head')

csrc/fastvideo_kernel/tests/support_flex_sta.py

@@ -0,0 +1,71 @@
+from typing import Tuple
+
+import torch
+from torch import BoolTensor, IntTensor
+from torch.nn.attention.flex_attention import create_block_mask
+
+# Peiyuan: This is neccesay. Dont know why. see https://github.com/pytorch/pytorch/issues/135028
+torch._inductor.config.realize_opcount_threshold = 100
+
+
+def generate_sta_mask(canvas_twh, kernel_twh, tile_twh, text_length):
+    """Generates a 3D NATTEN attention mask with a given kernel size.
+    
+    Args:
+        canvas_t: The time dimension of the canvas.
+        canvas_h: The height of the canvas.
+        canvas_w: The width of the canvas.
+        kernel_t: The time dimension of the kernel.
+        kernel_h: The height of the kernel.
+        kernel_w: The width of the kernel.
+    """
+    canvas_t, canvas_h, canvas_w = canvas_twh
+    kernel_t, kernel_h, kernel_w = kernel_twh
+    tile_t_size, tile_h_size, tile_w_size = tile_twh
+    total_tile_size = tile_t_size * tile_h_size * tile_w_size
+    canvas_tile_t, canvas_tile_h, canvas_tile_w = canvas_t // tile_t_size, canvas_h // tile_h_size, canvas_w // tile_w_size
+    img_seq_len = canvas_t * canvas_h * canvas_w
+
+    def get_tile_t_x_y(idx: IntTensor) -> Tuple[IntTensor, IntTensor, IntTensor]:
+        tile_id = idx // total_tile_size
+        tile_t = tile_id // (canvas_tile_h * canvas_tile_w)
+        tile_h = (tile_id % (canvas_tile_h * canvas_tile_w)) // canvas_tile_w
+        tile_w = tile_id % canvas_tile_w
+        return tile_t, tile_h, tile_w
+
+    def sta_mask_3d(
+        b: IntTensor,
+        h: IntTensor,
+        q_idx: IntTensor,
+        kv_idx: IntTensor,
+    ) -> BoolTensor:
+        q_t_tile, q_x_tile, q_y_tile = get_tile_t_x_y(q_idx)
+        kv_t_tile, kv_x_tile, kv_y_tile = get_tile_t_x_y(kv_idx)
+        # kernel nominally attempts to center itself on the query, but kernel center
+        # is clamped to a fixed distance (kernel half-length) from the canvas edge
+        kernel_center_t = q_t_tile.clamp(kernel_t // 2, (canvas_tile_t - 1) - kernel_t // 2)
+        kernel_center_x = q_x_tile.clamp(kernel_h // 2, (canvas_tile_h - 1) - kernel_h // 2)
+        kernel_center_y = q_y_tile.clamp(kernel_w // 2, (canvas_tile_w - 1) - kernel_w // 2)
+        time_mask = (kernel_center_t - kv_t_tile).abs() <= kernel_t // 2
+        hori_mask = (kernel_center_x - kv_x_tile).abs() <= kernel_h // 2
+        vert_mask = (kernel_center_y - kv_y_tile).abs() <= kernel_w // 2
+        image_mask = (q_idx < img_seq_len) & (kv_idx < img_seq_len)
+        image_to_text_mask = (q_idx < img_seq_len) & (kv_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length)
+        text_to_all_mask = (q_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length)
+        return (image_mask & time_mask & hori_mask & vert_mask) | image_to_text_mask | text_to_all_mask
+
+    sta_mask_3d.__name__ = f"natten_3d_c{canvas_t}x{canvas_w}x{canvas_h}_k{kernel_t}x{kernel_w}x{kernel_h}"
+    return sta_mask_3d
+
+
+def get_sliding_tile_attention_mask(kernel_size, tile_size, img_size, text_length, device, text_max_len=256):
+    img_seq_len = img_size[0] * img_size[1] * img_size[2]
+    image_mask = generate_sta_mask(img_size, kernel_size, tile_size, text_length)
+    mask = create_block_mask(image_mask,
+                             B=None,
+                             H=None,
+                             Q_LEN=img_seq_len + text_max_len,
+                             KV_LEN=img_seq_len + text_max_len,
+                             device=device,
+                             _compile=True)
+    return mask

csrc/fastvideo_kernel/tests/test_correctness.py

@@ -0,0 +1,63 @@
+import torch
+import sys
+import os
+from tqdm import tqdm
+
+# Local support import
+from .support_flex_sta import get_sliding_tile_attention_mask
+
+# USE OUR NEW PACKAGE!
+from fastvideo_kernel import sliding_tile_attention
+from torch.nn.attention.flex_attention import flex_attention
+
+flex_attention = torch.compile(flex_attention, dynamic=False)
+
+def flex_test(Q, K, V, kernel_size):
+    mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (18, 48, 80), 0, 'cuda', 0)
+    output = flex_attention(Q, K, V, block_mask=mask)
+    return output
+
+def h100_fwd_kernel_test(Q, K, V, kernel_size):
+    # Using the same parameters as the original test
+    o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 0, False, '18x48x80')
+    return o
+
+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, causal, mean, std, num_iterations=2):
+    print(f"Running correctness check: batch={b}, heads={h}, seq_len={n}, dim={d}")
+    kernel_size_ls = [(3, 3, 5), (3, 1, 10)]
+    
+    for kernel_size in kernel_size_ls:
+        print(f"Testing kernel_size: {kernel_size}")
+        for xi in tqdm(range(num_iterations)):
+            torch.manual_seed(xi)
+            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')
+            
+            tk_o = h100_fwd_kernel_test(Q, K, V, kernel_size)
+            pt_o = flex_test(Q, K, V, kernel_size)
+
+            diff = pt_o - tk_o
+            abs_diff = torch.abs(diff)
+            max_d = torch.max(abs_diff).item()
+            avg_d = torch.sum(abs_diff).item() / (b * h * n * d)
+            
+            if max_d > 0.1:
+                print(f"Warning: Large diff detected! max={max_d}, avg={avg_d}")
+
+    print("\n✅ TEST COMPLETE: New package matches FlexAttention behavior.")
+
+if __name__ == "__main__":
+    b, h, d = 2, 24, 128
+    n = 69120 
+    causal = False
+    mean = 1e-1
+    std = 10
+    
+    check_correctness(b, h, n, d, causal, mean, std, num_iterations=2)

csrc/fastvideo_kernel/tests/test_vmoba_correctness.py

@@ -3,7 +3,7 @@
 import torch
 import pytest
 import random
-from fastvideo_kernel import moba_attn_varlen
+from fastvideo_kernel.vmoba import moba_attn_varlen
 
 def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda"):
     """

docker/Dockerfile.python3.10

@@ -55,12 +55,18 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd fastvideo-kernel && \
+    cd csrc/attn/sliding_tile_attn && \
     git submodule update --init --recursive && \
-    ./build.sh
+    python setup.py install
 
+# Install VSA 
+RUN source $HOME/.local/bin/env && \
+    source /opt/venv/bin/activate && \
+    cd csrc/attn/video_sparse_attn && \
+    git submodule update --init --recursive && \
+    python setup.py install
 
 EXPOSE 22

docker/Dockerfile.python3.11

@@ -55,12 +55,11 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Kernels
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd fastvideo-kernel && \
+    cd csrc/fastvideo_kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    python setup.py install
-
 
 EXPOSE 22

docker/Dockerfile.python3.12

@@ -55,11 +55,11 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Kernels
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd fastvideo-kernel && \
+    cd csrc/fastvideo_kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    python setup.py install
 
 EXPOSE 22

docker/Dockerfile.python3.12.cuda12.9.1

@@ -55,12 +55,11 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Kernels
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd fastvideo-kernel && \
+    cd csrc/fastvideo_kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    python setup.py install
-
 
 EXPOSE 22

docker/rocm7_1.Dockerfile

@@ -36,7 +36,7 @@ RUN echo "# Placeholder" > README.md
 RUN source $HOME/.local/bin/env && \
     uv venv --python 3.10 --seed /opt/venv && \
     source /opt/venv/bin/activate && \
-    uv pip install --no-cache-dir --upgrade pip
+    uv pip install --no-cache-dir --upgrade pip 
 
 COPY . .
 
@@ -48,11 +48,11 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Kernels
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd fastvideo-kernel && \
+    cd csrc/fastvideo_kernel && \
     git submodule update --init --recursive && \
-    ./build.sh --rocm
+    python setup.py install
 
 EXPOSE 22

docs/assets/custom.css

@@ -2,11 +2,9 @@
     writing-mode: sideways-lr;
     white-space: nowrap;
     max-width: 0;
-}
+    p {
-
+       margin: 0;
-/* Keep header cell paragraph content tight (avoid CSS nesting for compatibility) */
+    }
-.vertical-table-header th.head:not(.stub) p {
-    margin: 0;
 }
 
 /* Image sizing classes */

docs/attention/developer/index.md

@@ -1,179 +0,0 @@
-# Adding a New Attention Backend
-
-FastVideo allows integrating new attention mechanisms easily. This guide walks you through adding a new backend (e.g., `MyNewAttn`).
-
-## 1. Implement the Backend (Python)
-
-Create a new file in `fastvideo/attention/backends/` (e.g., `mynew_attn.py`).
-
-Your implementation should inherit from `AttentionBackend` defined in `abstract.py`.
-
-```python
-# fastvideo/attention/backends/mynew_attn.py
-import torch
-from .abstract import AttentionBackend
-# Import the context manager to access metadata (optional)
-from fastvideo.forward_context import get_forward_context
-
-# Import compiled kernel if applicable (see Section 2)
-try:
-    # Import from the top-level package
-    from fastvideo_kernel import my_compiled_attn_func
-except ImportError:
-    my_compiled_attn_func = None
-
-class MyNewAttnBackend(AttentionBackend):
-    def process_inputs(self, q, k, v, **kwargs):
-        # Pre-process inputs if necessary
-        return q, k, v
-
-    def forward(self, q, k, v, **kwargs):
-        # Optional: Access extra metadata passed via ForwardContext
-        # Only needed if your backend requires global state (e.g. window_size)
-        try:
-            context = get_forward_context()
-            metadata = context.attn_metadata
-            # Example: window_size = metadata.window_size
-        except (AssertionError, AttributeError):
-            # Handle case where context is not set (e.g. standard inference)
-            pass
-
-        if my_compiled_attn_func is not None:
-            return my_compiled_attn_func(q, k, v)
-        else:
-            # Fallback implementation (e.g., Triton or pure PyTorch)
-            return self.fallback_impl(q, k, v)
-```
-
-## 2. Passing Extra Information via ForwardContext (Optional)
-
-FastVideo uses a `ForwardContext` to pass global metadata (like current timestep, batch info, or custom attention configurations) to attention backends without changing the `forward` signature of every layer. **This is optional and only required if your backend needs dynamic per-step information.**
-
-To use this:
-1. **Set Context**: In your pipeline or generation loop, use the `set_forward_context` context manager.
-2. **Access Context**: Inside your attention backend, use `get_forward_context()`.
-
-See `docs/attention/sta/index.md` (Sliding Tile Attention) for an example of how complex configuration (window sizes) is passed this way.
-
-## 3. Adding Compiled Kernels (C++/CUDA)
-
-If your backend requires custom CUDA kernels, you need to add them to the `fastvideo-kernel` package.
-
-### A. Add Source Files
-Place your kernel implementation files in `fastvideo-kernel/csrc/attention/`.
-* `mynew_attn.cu` (CUDA implementation)
-* `mynew_attn.h` (Optional headers)
-
-### B. Register in Extension
-Update `fastvideo-kernel/csrc/common_extension.cpp` to expose your function to Python.
-
-```cpp
-// 1. Declare external function
-#ifdef COMPILE_MYNEW_ATTN
-extern torch::Tensor mynew_attn_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v);
-#endif
-
-// 2. Register in module
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    // ... other kernels ...
-
-#ifdef COMPILE_MYNEW_ATTN
-    m.def("mynew_attn_fwd", torch::wrap_pybind_function(mynew_attn_forward), "My New Attention Forward");
-#endif
-}
-```
-
-### C. Update CMakeLists.txt
-Update `fastvideo-kernel/CMakeLists.txt` to compile your new files.
-
-**Case 1: General CUDA Kernel (Runs on all GPUs)**
-Add your source file directly to `EXTENSION_SOURCES` and define the compilation flag.
-
-```cmake
-# Add to EXTENSION_SOURCES
-list(APPEND EXTENSION_SOURCES csrc/attention/mynew_attn.cu)
-
-# Add compilation definition for common_extension.cpp
-list(APPEND COMPILE_DEFS COMPILE_MYNEW_ATTN)
-```
-
-**Case 2: ThunderKittens Kernel (Hopper H100 Only)**
-If your kernel uses ThunderKittens (TK), it requires specific architecture flags (`sm_90a`). Add it inside the `ENABLE_TK_KERNELS` block.
-
-```cmake
-if(ENABLE_TK_KERNELS)
-    # Add source only if TK is enabled
-    list(APPEND EXTENSION_SOURCES csrc/attention/mynew_attn_tk.cu)
-    
-    # Add definition to guard registration
-    list(APPEND COMPILE_DEFS TK_COMPILE_MYNEW_ATTN)
-endif()
-```
-
-### D. Expose in Python Ops
-Update `fastvideo-kernel/python/fastvideo_kernel/ops.py` to make the function importable and handle fallbacks gracefully.
-
-```python
-# fastvideo-kernel/python/fastvideo_kernel/ops.py
-
-# Try to load C++ extension symbols
-try:
-    from fastvideo_kernel._C import fastvideo_kernel_ops
-    mynew_attn_fwd = getattr(fastvideo_kernel_ops, "mynew_attn_fwd", None)
-except ImportError:
-    mynew_attn_fwd = None
-
-def my_compiled_attn_func(q, k, v):
-    # Runtime check: use C++ kernel if available, else fallback
-    if mynew_attn_fwd is not None:
-        return mynew_attn_fwd(q, k, v)
-    else:
-        # Call Triton/Python fallback
-        return mynew_attn_triton(q, k, v)
-```
-
-### E. Expose in Package Init
-Update `fastvideo-kernel/python/fastvideo_kernel/__init__.py` to export the function.
-
-```python
-from fastvideo_kernel.ops import (
-    my_compiled_attn_func,
-    # ...
-)
-
-__all__ = [
-    "my_compiled_attn_func",
-    # ...
-]
-```
-
-## 4. Register the Backend
-
-Update `fastvideo/attention/backends/__init__.py` to export your new class.
-
-```python
-from .mynew_attn import MyNewAttnBackend
-```
-
-## 5. Platform Integration
-
-If your backend requires specific platform checks (e.g., checking for H100 support), handle that in `fastvideo/platforms/cuda.py` or within your backend's `__init__`.
-
-## 6. Add Documentation
-
-Create a new documentation page for your backend to explain its usage, installation (if custom kernels are needed), and features.
-
-1. **Create Directory**: `docs/attention/mynew_attn/`
-2. **Create Index**: `docs/attention/mynew_attn/index.md`
-3. **Update Navigation**: Add an entry to `mkdocs.yml` under the "Attention" tab.
-
-## Checklist
-
-* [ ] Created `fastvideo/attention/backends/mynew_attn.py`.
-* [ ] (Optional) Added CUDA kernels in `fastvideo-kernel/csrc/attention/`.
-* [ ] (Optional) Updated `common_extension.cpp` and `CMakeLists.txt`.
-* [ ] (Optional) Exposed kernel in `fastvideo-kernel/python/fastvideo_kernel/ops.py`.
-* [ ] (Optional) Exported kernel in `fastvideo-kernel/python/fastvideo_kernel/__init__.py`.
-* [ ] Implemented `forward` method respecting the standard signature.
-* [ ] Added unit tests in `tests/`.
-* [ ] Added documentation in `docs/attention/` and updated `mkdocs.yml`.

docs/attention/index.md

@@ -1,53 +0,0 @@
-# FastVideo Attention Kernels
-
-FastVideo provides highly optimized custom attention kernels to accelerate video generation.
-
-## Supported Kernels
-
-* **[Video Sparse Attention (VSA)](vsa/index.md)**: Sparse attention mechanism selecting top-k blocks.
-* **[Sliding Tile Attention (STA)](sta/index.md)**: Optimized attention for window-based video generation.
-
-## General Build Instructions
-
-These instructions apply to building the `fastvideo-kernel` package from source, which includes both STA and VSA kernels.
-
-### Prerequisites
-
-* **PyTorch**: 2.5.0+
-* **CUDA**: 12.4+ (12.8 recommended for best performance)
-* **C++ Compiler**: GCC 11+ (C++20 support required for ThunderKittens)
-
-Install system dependencies:
-
-```bash
-sudo apt update
-sudo apt install -y gcc-11 g++-11 clang-11 ninja-build
-
-# Set gcc-11 as default
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-```
-
-Set up your CUDA environment variables (adjust version as needed):
-
-```bash
-export CUDA_HOME=/usr/local/cuda-12.8
-export PATH=${CUDA_HOME}/bin:${PATH}
-export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-```
-
-### Compile and Install
-
-Clone the repository and build the kernel:
-
-```bash
-# Clone recursively to get ThunderKittens submodule
-git clone --recursive https://github.com/hao-ai-lab/FastVideo.git
-cd FastVideo/fastvideo-kernel
-
-# Build and install
-./build.sh
-```
-
-The build script automatically detects your GPU architecture:
-* **H100 (sm_90a)**: Compiles optimized C++ ThunderKittens kernels.
-* **Other (A100, etc.)**: Skips C++ compilation; installs Python package with Triton kernels.

docs/attention/sta/index.md

@@ -1,36 +0,0 @@
-# Sliding Tile Attention (STA)
-
-Optimized attention for window-based video generation (e.g., HunyuanVideo).
-
-## Installation
-
-STA is included in the `fastvideo-kernel` package. See the [main Attention page](../index.md) for build instructions.
-
-## Usage
-
-```python
-from fastvideo_kernel import sliding_tile_attention
-
-# q, k, v: [batch_size, num_heads, seq_length, head_dim]
-# window_size: List of (t, h, w) tiles. Tile size is (6, 8, 8).
-# text_length: Number of text tokens (0-256)
-
-out = sliding_tile_attention(
-    q, k, v, 
-    window_size=[(3, 3, 3)],  # Example window
-    text_length=256
-)
-```
-
-## Citation
-
-If you use Sliding Tile Attention in your research, please cite:
-
-```bibtex
-@article{zhang2025fast,
-  title={Fast video generation with sliding tile attention},
-  author={Zhang, Peiyuan and Chen, Yongqi and Su, Runlong and Ding, Hangliang and Stoica, Ion and Liu, Zhengzhong and Zhang, Hao},
-  journal={arXiv preprint arXiv:2502.04507},
-  year={2025}
-}
-```

docs/attention/vsa/index.md

@@ -1,36 +0,0 @@
-# Video Sparse Attention (VSA)
-
-Sparse attention mechanism selecting top-k blocks.
-
-## Installation
-
-VSA is included in the `fastvideo-kernel` package. See the [main Attention page](../index.md) for build instructions.
-
-## Usage
-
-```python
-from fastvideo_kernel import video_sparse_attn
-
-# q, k, v: [batch_size, num_heads, seq_len, head_dim]
-# variable_block_sizes: Number of valid tokens per block
-# topk: Number of blocks to attend
-
-output = video_sparse_attn(
-    q, k, v, 
-    variable_block_sizes=block_sizes, 
-    topk=32
-)
-```
-
-## Citation
-
-If you use Video Sparse Attention in your research, please cite:
-
-```bibtex
-@article{zhang2025vsa,
-  title={Vsa: Faster video diffusion with trainable sparse attention},
-  author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
-  journal={arXiv preprint arXiv:2505.13389},
-  year={2025}
-}
-```

docs/getting_started/quick_start.md

@@ -75,4 +75,3 @@ if __name__ == '__main__':
 - [Configuration](../inference/configuration.md) - Learn about configuration options
 - [Examples](../inference/examples/) - Explore more examples
 - [Optimizations](../inference/optimizations.md) - Performance optimization tips
-- [Low VRAM Inference](../inference/low_vram_inference.md) - Memory-saving settings (CPU offload, sharded loading, etc.)

docs/inference/inference_quick_start.md

@@ -98,7 +98,7 @@ Common issues and their solutions:
 ### Out of Memory Errors
 If you encounter CUDA out of memory errors:
 - Reduce `num_frames` or video resolution
-- Enable memory optimization with CPU-offload and sharded loading flags (see [Low VRAM Inference](low_vram_inference.md))
+- Enable memory optimization with `enable_model_cpu_offload`
 - Try a smaller model or use distilled versions
 - Use `num_gpus` > 1 if multiple GPUs are available
 

docs/sliding_tile_attention/demo.md

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

docs/sliding_tile_attention/installation.md

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

docs/video_sparse_attention/installation.md

@@ -0,0 +1,65 @@
+
+# 🔧 Installation
+You can install the Video Sparse Attention package using
+
+```bash
+pip install vsa
+```
+
+# Building from Source
+We support H100s (via ThunderKittens) and any other GPU (via Triton) for VSA.
+
+First, install C++20 for ThunderKittens (if using an H100):
+
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+
+Set up CUDA environment (if using CUDA 12.8):
+
+```bash
+export CUDA_HOME=/usr/local/cuda-12.8
+export PATH=${CUDA_HOME}/bin:${PATH} 
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+Install VSA:
+
+```bash
+cd csrc/attn/video_sparse_attn/
+git submodule update --init --recursive
+python setup.py install
+```
+
+# 🧪 Test
+
+```bash
+python csrc/attn/tests/test_vsa.py
+```
+
+# 📋 Usage
+
+```python
+from vsa import video_sparse_attn
+
+# q, k, v: [batch_size, num_heads, seq_len, head_dim]
+# variable_block_sizes: [num_blocks] - number of valid tokens in each block
+# topk: int - number of top-k blocks to attend to
+# block_size: int or tuple of 3 ints - size of each block (default: 64 tokens)
+# compress_attn_weight: optional weight for compressed attention branch
+
+output = video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight)
+
+``` 
+
+# 🚀Inference
+
+```bash
+bash scripts/inference/v1_inference_wan_VSA.sh
+```

fastvideo-kernel/CMakeLists.txt

@@ -1,159 +0,0 @@
-cmake_minimum_required(VERSION 3.26 FATAL_ERROR)
-project(fastvideo-kernel LANGUAGES CXX)
-
-# Prefer environment variable (used by CI or pip install git+repo_addr) if CMake var is not explicitly set.
-if(NOT DEFINED GPU_BACKEND AND DEFINED ENV{GPU_BACKEND})
-    set(GPU_BACKEND "$ENV{GPU_BACKEND}")
-endif()
-
-if(GPU_BACKEND STREQUAL "ROCM")
-    enable_language(HIP)
-else()
-    enable_language(CUDA)
-endif()
-
-# Import common utils if needed, but we keep it simple for now
-
-# Find Python and Torch
-find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)
-
-# Robustly find Torch include paths using Python
-execute_process(
-    COMMAND "${Python_EXECUTABLE}" -c "import torch; from torch.utils.cpp_extension import include_paths; print(';'.join(include_paths()))"
-    OUTPUT_VARIABLE TORCH_INCLUDE_PATHS
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-)
-list(APPEND TORCH_INCLUDE_DIRS ${TORCH_INCLUDE_PATHS})
-
-# Find Torch package (still useful for libraries)
-find_package(Torch REQUIRED)
-
-# Include directories
-include_directories(
-    ${CMAKE_SOURCE_DIR}/include
-    ${CMAKE_SOURCE_DIR}/include/cutlass/include
-    ${CMAKE_SOURCE_DIR}/include/tk/include
-    ${CMAKE_SOURCE_DIR}/include/tk/prototype
-    ${CMAKE_SOURCE_DIR}/csrc
-    ${CMAKE_SOURCE_DIR}/csrc/turbodiffusion
-    ${TORCH_INCLUDE_DIRS}
-)
-
-# ---------------------------
-# ThunderKittens (TK) toggles
-# ---------------------------
-# AUTO: enable TK only when we can confidently target Hopper (sm_90a).
-# ON:   force-enable TK kernels (intended for release wheels/images; does NOT require a GPU).
-# OFF:  never build TK kernels.
-set(FASTVIDEO_KERNEL_BUILD_TK "AUTO" CACHE STRING "Build ThunderKittens kernels: AUTO/ON/OFF")
-set_property(CACHE FASTVIDEO_KERNEL_BUILD_TK PROPERTY STRINGS AUTO ON OFF)
-
-# Prefer environment variable (used by CI) if CMake var is not explicitly set.
-if(NOT DEFINED TORCH_CUDA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST})
-    set(TORCH_CUDA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}")
-endif()
-
-message(STATUS "TORCH_CUDA_ARCH_LIST (cmake/env): ${TORCH_CUDA_ARCH_LIST}")
-message(STATUS "FASTVIDEO_KERNEL_BUILD_TK: ${FASTVIDEO_KERNEL_BUILD_TK}")
-
-set(ENABLE_TK_KERNELS OFF)
-if(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "ON")
-    set(ENABLE_TK_KERNELS ON)
-elseif(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "OFF")
-    set(ENABLE_TK_KERNELS OFF)
-else()
-    # AUTO: detect Hopper if possible.
-    if(TORCH_CUDA_ARCH_LIST)
-        # Accept common spellings: 9.0a, 90a, sm_90a.
-        string(REGEX MATCH "(^|[; ,])((9\\.0a)|(90a)|(sm_90a))([; ,]|$)" _HAS_90A "${TORCH_CUDA_ARCH_LIST}")
-        if(_HAS_90A)
-            set(ENABLE_TK_KERNELS ON)
-        endif()
-    else()
-        # Best-effort local detection (works when a CUDA device is visible).
-        execute_process(
-            COMMAND "${Python_EXECUTABLE}" -c "import torch; import sys; \nprint('1' if (torch.cuda.is_available() and torch.version.cuda and torch.cuda.get_device_capability()[0] >= 9) else '0')"
-            OUTPUT_VARIABLE _LOCAL_HAS_HOPPER
-            OUTPUT_STRIP_TRAILING_WHITESPACE
-            ERROR_QUIET
-        )
-        if(_LOCAL_HAS_HOPPER STREQUAL "1")
-            set(ENABLE_TK_KERNELS ON)
-        endif()
-    endif()
-endif()
-
-if(ENABLE_TK_KERNELS)
-    message(STATUS "ThunderKittens kernels: ENABLED")
-else()
-    message(STATUS "ThunderKittens kernels: DISABLED (will use Triton fallbacks at runtime)")
-endif()
-
-# Always try to build the extension if CUDA is available, but conditionally add sources/flags
-set(BUILD_CXX_KERNELS ON)
-
-# Compiler flags
-set(CUDA_FLAGS
-    "-DNDEBUG"
-    "-O3"
-    "-std=c++20"
-    "--use_fast_math"
-    "--expt-extended-lambda"
-    "--expt-relaxed-constexpr"
-    "-Xcompiler=-fno-strict-aliasing"
-    "-Xcompiler=-fPIC"
-    "-DTORCH_COMPILE"
-    "-Xnvlink=--verbose"
-    "-Xptxas=--verbose"
-    "-Xptxas=--warn-on-spills"
-)
-
-# If TK is enabled, ensure we target Hopper. This is required even on GPU-less builders (CI).
-if(ENABLE_TK_KERNELS)
-    if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES OR CMAKE_CUDA_ARCHITECTURES STREQUAL "")
-        set(CMAKE_CUDA_ARCHITECTURES "90a" CACHE STRING "CUDA architectures" FORCE)
-    endif()
-    list(APPEND CUDA_FLAGS "-DKITTENS_HOPPER")
-    message(STATUS "CMAKE_CUDA_ARCHITECTURES: ${CMAKE_CUDA_ARCHITECTURES}")
-endif()
-
-if(BUILD_CXX_KERNELS)
-    # Source files
-    set(EXTENSION_SOURCES 
-        csrc/common_extension.cpp
-        csrc/turbodiffusion/gemm/gemm.cu
-        csrc/turbodiffusion/norm/rmsnorm.cu
-        csrc/turbodiffusion/norm/layernorm.cu
-        csrc/turbodiffusion/quant/quant.cu
-    )
-    
-    # Conditionally add TK kernels
-    if(ENABLE_TK_KERNELS)
-        list(APPEND EXTENSION_SOURCES
-            csrc/attention/st_attn_h100.cu
-            csrc/attention/block_sparse_h100.cu
-        )
-    endif()
-
-    # Combined FastVideo Extension
-    # Using name 'fastvideo_kernel_ops' to distinguish from the python package namespace
-    Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI
-        ${EXTENSION_SOURCES}
-    )
-
-    # Build compile definitions list
-    set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops)
-    if(ENABLE_TK_KERNELS)
-        list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE)
-    endif()
-
-    target_compile_definitions(fastvideo_kernel_ops PRIVATE ${COMPILE_DEFS})
-
-    target_compile_options(fastvideo_kernel_ops PRIVATE
-        $<$<COMPILE_LANGUAGE:CUDA>:${CUDA_FLAGS}>
-    )
-
-    # We install it to fastvideo_kernel/_C so we can load it to register the ops
-    install(TARGETS fastvideo_kernel_ops LIBRARY DESTINATION fastvideo_kernel/_C)
-endif()
-

fastvideo-kernel/MANIFEST.in

@@ -1,6 +0,0 @@
-include LICENSE
-include README.md
-include pyproject.toml
-recursive-include python/fastvideo_kernel *.py
-recursive-include csrc *.cu *.cuh *.cpp *.h
-recursive-include include/tk *.cu *.cuh *.cpp *.h *.src

fastvideo-kernel/README.md

@@ -1,69 +0,0 @@
-# FastVideo Kernel
-
-CUDA kernels for FastVideo video generation.
-
-## Installation
-
-### Standard Installation (Local Development)
-This will automatically detect your GPU architecture. If an NVIDIA Hopper (H100/sm_90a) GPU is detected, ThunderKittens kernels will be enabled. Otherwise, they will be skipped, and the package will use Triton fallbacks at runtime.
-
-```bash
-git submodule update --init --recursive
-cd fastvideo-kernel
-./build.sh
-```
-
-### Rocm Build
-If you are in a rocm environment without the compilation toolchaine of CUDA.
-
-```bash
-cd fastvideo-kernel
-./build.sh --rocm
-```
-
-## Usage
-
-### Sliding Tile Attention (STA) & Video Sparse Attention (VSA)
-
-For detailed usage, please check the [Attention Documentation](../docs/attention/index.md).
-
-```python
-from fastvideo_kernel import sliding_tile_attention, video_sparse_attn, moba_attn_varlen
-
-# Example: Sliding Tile Attention
-out = sliding_tile_attention(q, k, v, window_sizes, text_len)
-
-# Example: Video Sparse Attention (with Triton fallback)
-out = video_sparse_attn(q, k, v, block_sizes, topk=5)
-
-# Example: VMoBA
-out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)
-```
-
-### TurboDiffusion Kernels
-
-This package also includes kernels from [TurboDiffusion](https://github.com/thu-ml/TurboDiffusion), including INT8 GEMM, Quantization, RMSNorm and LayerNorm.
-
-## Requirements
-
-- **Runtime**:
-  - NVIDIA H100 (sm_90a) for C++ optimized kernels.
-  - Any CUDA GPU for Triton-based fallbacks.
-- **Build**:
-  - CUDA Toolkit 12.3+
-  - C++20 compatible compiler (GCC 10+, Clang 11+)
-
-## Acknowledgement
-
-This package structure and build system are based on [sgl-kernel](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) from the SGLang project.
-
-The implementation of `turbodiffusion` kernels is adapted from [TurboDiffusion](https://github.com/thu-ml/TurboDiffusion). If you use these kernels, please cite:
-
-```bibtex
-@article{zhang2025turbodiffusion,
-  title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times},
-  author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun},
-  journal={arXiv preprint arXiv:2512.16093},
-  year={2025}
-}
-```

fastvideo-kernel/build.sh

@@ -1,38 +0,0 @@
-#!/bin/bash
-set -ex
-
-# Simple build script wrapping uv/pip
-# Usage:
-#   ./build.sh                # local dev build (auto-detect / skip TK kernels when not available)
-#   ./build.sh --release      # force-enable Hopper/TK kernels for release builds (no GPU required)
-
-echo "Building fastvideo-kernel..."
-
-# Ensure submodules are initialized if needed (tk)
-git submodule update --init --recursive
-
-# Install build dependencies
-pip install scikit-build-core cmake ninja
-
-RELEASE=0
-GPU_BACKEND=CUDA
-for arg in "$@"; do
-    case "$arg" in
-        --rocm)
-            GPU_BACKEND=ROCM
-            ;;
-    esac
-done
-
-# Force-enable ThunderKittens kernels and compile for Hopper.
-export TORCH_CUDA_ARCH_LIST="9.0a"
-export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
-
-export CMAKE_ARGS="${CMAKE_ARGS:-} -DGPU_BACKEND=${GPU_BACKEND}"
-
-echo "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST:-<unset>}"
-echo "CMAKE_ARGS: ${CMAKE_ARGS:-<unset>}"
-echo "GPU_BACKEND: ${GPU_BACKEND:-<unset>}"
-# Build and install
-# Use -v for verbose output
-pip install . -v --no-build-isolation

fastvideo-kernel/csrc/common_extension.cpp

@@ -1,48 +0,0 @@
-#include <torch/extension.h>
-#include <ATen/ATen.h>
-#include <vector>
-
-// Forward declarations
-#ifdef TK_COMPILE_ST_ATTN
-extern torch::Tensor sta_forward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, 
-    int kernel_t_size, int kernel_w_size, int kernel_h_size, 
-    int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag
-); 
-#endif
-
-#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
-); 
-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
-);
-#endif
-
-// TurboDiffusion kernels
-void register_quant(pybind11::module_ &);
-void register_rms_norm(pybind11::module_ &);
-void register_layer_norm(pybind11::module_ &);
-void register_gemm(pybind11::module_ &);
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.doc() = "FastVideo CUDA Kernels";
-
-#ifdef TK_COMPILE_ST_ATTN
-    m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention (Hopper)");
-#endif
-
-#ifdef TK_COMPILE_BLOCK_SPARSE
-    m.def("block_sparse_fwd", torch::wrap_pybind_function(block_sparse_attention_forward), "block sparse attention forward (Hopper)");
-    m.def("block_sparse_bwd", torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward (Hopper)");
-#endif
-
-    // TurboDiffusion
-    register_quant(m);
-    register_rms_norm(m);
-    register_layer_norm(m);
-    register_gemm(m);
-}

fastvideo-kernel/csrc/turbodiffusion/common/common.hpp

@@ -1,85 +0,0 @@
-#pragma once
-
-#include <torch/torch.h>
-#include <torch/extension.h>
-
-CUTLASS_HOST_DEVICE int64_t cdiv(int64_t const& a, int64_t const &b) {
-  return (a + b - 1) / b;
-}
-
-template <class T>
-CUTLASS_HOST_DEVICE T max(T a, T b) { return a > b ? a : b; }
-
-template <class T>
-CUTLASS_HOST_DEVICE T min(T a, T b) { return a > b ? b : a; }
-
-#define MIN(a, b) ((a) > (b) ? (b) : (a))
-#define MAX(a, b) ((a) > (b) ? (a) : (b))
-
-#define BOOL_SWITCH(COND, CONST_NAME, ...)    \
-[&] {                                         \
-  if (COND) {                                 \
-    static constexpr bool CONST_NAME = true;  \
-    return (__VA_ARGS__)();                   \
-  } else {                                    \
-    static constexpr bool CONST_NAME = false; \
-    return (__VA_ARGS__)();                   \
-  }                                           \
-}() 
-
-#define CUDA_CHECK(call)    \
-{                           \
-  cudaError_t err = call;   \
-  if (err != cudaSuccess) { \
-    fprintf(stderr, "CUDA Error at %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
-    exit(err);              \
-  }                         \
-}
-
-#define CONFIG_SWITCH(N, ...)                                   \
-[&] {                                                           \
-  if (N <= 1024) {                                              \
-    constexpr int NUM_THR_PER_CTA = 128;                        \
-    constexpr int MAX_HIDDEN_SIZE = 1024;                       \
-    return (__VA_ARGS__)();                                     \
-  } else if (N <= 2048) {                                       \
-    constexpr int NUM_THR_PER_CTA = 128;                        \
-    constexpr int MAX_HIDDEN_SIZE = 2048;                       \
-    return (__VA_ARGS__)();                                     \
-  } else if (N <= 4096) {                                       \
-    constexpr int NUM_THR_PER_CTA = 128;                        \
-    constexpr int MAX_HIDDEN_SIZE = 4096;                       \
-    return (__VA_ARGS__)();                                     \
-  } else if (N <= 8192) {                                       \
-    constexpr int NUM_THR_PER_CTA = 256;                        \
-    constexpr int MAX_HIDDEN_SIZE = 8192;                       \
-    return (__VA_ARGS__)();                                     \
-  } else {                                                      \
-    constexpr int NUM_THR_PER_CTA = 256;                        \
-    constexpr int MAX_HIDDEN_SIZE = 16384;                      \
-    return (__VA_ARGS__)();                                     \
-  }                                                             \
-}()
-
-
-template <int BlockSize>
-  void create_tensor(
-    torch::Device const &device,
-    std::optional<at::Tensor> &output,
-    std::optional<at::Tensor> &scale,
-    int m, int n
-  ) {
-    int num_block_m = cdiv(m, BlockSize);
-    int num_block_n = cdiv(n, BlockSize);
-    if (!output.has_value()) {
-      output.emplace(torch::empty(
-        {m, n},
-        torch::TensorOptions().device(device).dtype(torch::kInt8)
-      ));
-      scale.emplace(torch::empty(
-        {num_block_m, num_block_n},
-        torch::TensorOptions().device(device).dtype(torch::kFloat32)
-      ));
-    }
-  }
-

fastvideo-kernel/csrc/turbodiffusion/common/launch.hpp

@@ -1,44 +0,0 @@
-#pragma once
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-
-template <class Kernel>
-__global__ void device_kernel(
-  __grid_constant__ typename Kernel::Params const params
-) {
-  extern __shared__ char smem[];
-  Kernel op;
-  op(params, smem);
-}
-
-template <class Kernel>
-__global__ __launch_bounds__(Kernel::MaxThreadsPerBlock, Kernel::MinBlocksPerMultiprocessor)
-void device_kernel_with_launch_bounds(
-  __grid_constant__ typename Kernel::Params const params
-) {
-  extern __shared__ char smem[];
-  Kernel op;
-  op(params, smem);
-}
-
-template <class Kernel>
-void launch_kernel(
-  typename Kernel::Params const &params,
-  dim3 grid_shape,
-  dim3 cta_shape,
-  size_t ShmSize,
-  cudaStream_t stream = nullptr
-) {
-  auto func = device_kernel<Kernel>;
-  if (ShmSize >= 48 * 1024) {
-    CUDA_CHECK(cudaFuncSetAttribute(
-      func,
-      cudaFuncAttributeMaxDynamicSharedMemorySize,
-      ShmSize
-    ));
-  }
-  func<<<grid_shape, cta_shape, ShmSize, stream>>>(params);
-  CUDA_CHECK(cudaGetLastError());
-}

fastvideo-kernel/csrc/turbodiffusion/common/load.hpp

@@ -1,75 +0,0 @@
-#pragma once
-
-template <
-  class InputDtype_,
-  int TileM_,
-  int TileN_,
-  int NumThrPerCta_,
-  bool IsEvenM,
-  bool IsEvenN
->
-class Loader {
-public: 
-  using InputDtype = InputDtype_;
-  static constexpr int TileM = TileM_;
-  static constexpr int TileN = TileN_;
-  static constexpr int NumThrPerCta = NumThrPerCta_;
-  static constexpr int NumElementPerThread = TileM * TileN / NumThrPerCta;
-  static constexpr int NumThrPerRow = TileN / NumElementPerThread;
-
-  static_assert(NumThrPerCta % TileM == 0);
-  static_assert(TileM * TileN % NumThrPerCta == 0);
-
-
-  CUTLASS_DEVICE void
-  load(void const *input_ptr, void *thr_output_reg, int64_t m, int64_t n, int blk_m, int blk_n, int tid) {
-    int n_alignment = (n & 31) * sizeof(InputDtype);
-    int thr_m_offset = tid / NumThrPerRow;
-    int thr_n_offset = (tid % NumThrPerRow) * NumElementPerThread;
-    void const *cta_input_ptr = (void*)((InputDtype*)input_ptr + blk_m * TileM * n + blk_n * TileN);
-    void const *thr_input_ptr = (void*)((InputDtype*)cta_input_ptr + thr_m_offset * n + thr_n_offset);
-    InputDtype tmp_reg[NumElementPerThread];
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      tmp_reg[i] = InputDtype(0.f);
-    bool pred = IsEvenM ? true : thr_m_offset + blk_m * TileM < m;
-    int limit = IsEvenN ? NumElementPerThread : MIN(NumElementPerThread, n - (blk_n * TileN + thr_n_offset));
-    if (n_alignment % 128 == 0)
-      _load<int4, IsEvenN>(thr_input_ptr, (void*)tmp_reg, limit, pred);
-    else if (n_alignment % 64 == 0)
-      _load<int2, IsEvenN>(thr_input_ptr, (void*)tmp_reg, limit, pred);
-    else
-      _load<InputDtype, IsEvenN>(thr_input_ptr, (void*)tmp_reg, limit, pred);
-
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      *((float*)thr_output_reg + i) = static_cast<float>(reinterpret_cast<InputDtype const&>(tmp_reg[i]));
-  }
-
-private:
-  template <class LoadDataType, bool IsEven>
-  CUTLASS_DEVICE void
-  _load(void const *thr_input_ptr, void *thr_output_reg, int limit, bool pred) {
-    static constexpr int NumElementPerLoad = sizeof(LoadDataType) / sizeof(InputDtype);
-    if (pred) {
-      if constexpr (IsEven) {
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < NumElementPerThread; i += NumElementPerLoad) {
-          *(LoadDataType*)((InputDtype*)thr_output_reg + i) = *(LoadDataType*)((InputDtype*)thr_input_ptr + i);
-        }
-      } else {
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < limit; i += NumElementPerLoad) {
-          if (limit - i > NumElementPerLoad)
-            *(LoadDataType*)((InputDtype*)thr_output_reg + i) = *(LoadDataType*)((InputDtype*)thr_input_ptr + i);
-          else {
-            for (int j = 0; j < NumElementPerLoad; ++j) {
-              if (i + j < limit)
-                *((InputDtype*)thr_output_reg + i + j) = *((InputDtype*)thr_input_ptr + i + j);
-            }
-          }
-        }
-      }
-    }
-  }
-};

fastvideo-kernel/csrc/turbodiffusion/common/store.hpp

@@ -1,77 +0,0 @@
-#pragma once
-#include "common/common.hpp"
-
-
-template <
-  class OutputDtype_,
-  int TileM_,
-  int TileN_,
-  int NumThrPerCta_,
-  bool IsEvenM,
-  bool IsEvenN,
-  bool Round = true,
-  bool SaveScale = true
->
-class Saver {
-public:
-  using OutputDtype = OutputDtype_;
-
-  static constexpr int TileM = TileM_;
-  static constexpr int TileN = TileN_;
-  static constexpr int NumThrPerCta = NumThrPerCta_;
-  static constexpr int NumElementPerThread = TileM * TileN / NumThrPerCta;
-  static constexpr int NumThrPerRow = TileN / NumElementPerThread;
-
-  static_assert(TileM * TileN % NumThrPerCta == 0);
-  static_assert(NumThrPerCta % TileM == 0);
-
-  CUTLASS_DEVICE void
-  store(void *Optr, void *OSptr, void *reg, float scale_inv, int64_t m, int64_t n, int blk_m, int blk_n, int tid) {
-    int n_alignment = (n & 31) * sizeof(OutputDtype);
-    int thr_m_offset = tid / NumThrPerRow;
-    int thr_n_offset = (tid % NumThrPerRow) * NumElementPerThread;
-    void *cta_output_ptr = (void*)((OutputDtype*)Optr + blk_m * TileM * (Round ? cdiv(n, TileN) * TileN : n) + blk_n * TileN);
-    void *thr_output_ptr = (void*)((OutputDtype*)cta_output_ptr + thr_m_offset * (Round ? cdiv(n, TileN) * TileN : n) + thr_n_offset);
-    bool pred = IsEvenM ? true : thr_m_offset + blk_m * TileM < m;
-    int limit = IsEvenN ? NumElementPerThread : MIN(NumElementPerThread, n - (blk_n * TileN + thr_n_offset));
-    if (n_alignment % 128 == 0)
-      _store<int4, IsEvenN>(thr_output_ptr, reg, limit, pred);
-    else if (n_alignment % 64 == 0)
-      _store<int2, IsEvenN>(thr_output_ptr, reg, limit, pred);
-    else
-      _store<OutputDtype, IsEvenN>(thr_output_ptr, reg, limit, pred);
-
-    if constexpr (SaveScale) {
-      if (tid == 0) {
-        *((float*)OSptr + blk_m * cdiv(n, TileN)+ blk_n) = scale_inv;
-      }
-    }
-  }
-
-private:
-  template <class StoreDataType, bool IsEven>
-  CUTLASS_DEVICE void
-  _store(void *thr_output_ptr, void *reg, int limit, bool pred) {
-    static constexpr int NumElementPerStore = sizeof(StoreDataType) / sizeof(OutputDtype);
-    if (pred) {
-      if constexpr (IsEven) {
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < NumElementPerThread; i += NumElementPerStore) {
-          *(StoreDataType*)((OutputDtype*)thr_output_ptr + i) = *(StoreDataType*)((OutputDtype*)reg + i);
-        }
-      } else {
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < limit; i += NumElementPerStore) {
-          if (limit - i > NumElementPerStore)
-            *(StoreDataType*)((OutputDtype*)thr_output_ptr + i) = *(StoreDataType*)((OutputDtype*)reg + i);
-          else {
-            for (int j = 0; j < limit - i; ++j) {
-              *((OutputDtype*)thr_output_ptr + i + j) = *((OutputDtype*)reg + i + j);
-            }
-          }
-        }
-      }
-    }
-  }
-
-};

fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu

@@ -1,77 +0,0 @@
-/* 
- * Copyright (c) 2025 by TurboDiffusion team.
- * 
- * Licensed under the Apache License, Version 2.0 (the "License");
- * 
- * Citation (please cite if you use this code):
- * 
- * @article{zhang2025turbodiffusion,
- *   title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times},
- *   author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun},
- *   journal={arXiv preprint arXiv:2512.16093},
- *   year={2025}
- * }
- */
-
-#include <pybind11/pybind11.h>
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <torch/python.h>
-#include <cutlass/cutlass.h>
-#include <cutlass/numeric_types.h>
-
-#include "common/common.hpp"
-#include "gemm/launch.hpp"
-
-void int8_gemm(
-  at::Tensor const& A, at::Tensor const& A_S, 
-  at::Tensor const& B, at::Tensor const& B_S, 
-  torch::Tensor& C
-) {
-
-  
-  static constexpr int swizzle_dir = 1;
-  static constexpr int swizzle_size_log = 5;
-
-  int k = B.size(1);
-  int m = A.size(0);
-  int n = B.size(0);
-
-  switch (C.scalar_type()) {
-    case torch::kHalf:{
-        int8_gemm_<cutlass::half_t> (
-            (int8_t*)A.data_ptr(), A_S.data_ptr<float>(),
-            (int8_t*)B.data_ptr(), B_S.data_ptr<float>(), 
-            (cutlass::half_t*)C.data_ptr(), 
-            m, n, k, swizzle_dir, swizzle_size_log, at::cuda::getCurrentCUDAStream().stream()
-        );
-        break;
-    }
-
-    case torch::kBFloat16:{
-        int8_gemm_<cutlass::bfloat16_t> (
-                (int8_t*)A.data_ptr(), A_S.data_ptr<float>(),
-                (int8_t*)B.data_ptr(), B_S.data_ptr<float>(), 
-                (cutlass::bfloat16_t*)C.data_ptr(), 
-                m, n, k, swizzle_dir, swizzle_size_log, at::cuda::getCurrentCUDAStream().stream()
-            );
-        break;
-    }
-
-    default: {
-      std::cerr << "Observing: " << C.scalar_type() << " for the output datatype which is invalid";
-      throw std::runtime_error("Unsupported output data type for int8 gemm.");
-    }
-  }
-
-}
-
-void register_gemm(pybind11::module_ &m) {
-    m.def("gemm_cuda", &int8_gemm);
-}
-
-
-  
-
-

fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp

@@ -1,522 +0,0 @@
-/* 
- * Copyright (c) 2025 by TurboDiffusion team.
- * 
- * Licensed under the Apache License, Version 2.0 (the "License");
- * 
- * Citation (please cite if you use this code):
- * 
- * @article{zhang2025turbodiffusion,
- *   title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times},
- *   author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun},
- *   journal={arXiv preprint arXiv:2512.16093},
- *   year={2025}
- * }
- */
-
-#pragma once
-
-#include <cuda.h>
-#include "cute/tensor.hpp"
-
-#include "common/common.hpp"
-#include "gemm/utils.hpp"
-
-using namespace cute;
-
-template <
-  class OutputDtype_,
-  bool IsEvenM,
-  bool IsEvenN
->
-struct GemmKernel {
-  using ElementA = int8_t;
-  using ElementB = int8_t;
-  using OutputDtype = OutputDtype_;
-  using AccumulatorDtype = int32_t;
-  static constexpr int BlockSize = 128;
-  static constexpr int TileM = 128;
-  static constexpr int TileN = 128;
-  static constexpr int TileK = 128;
-  static constexpr int Stage = 3;
-  static constexpr int EpiStage = 2;
-  
-  static_assert(
-       BlockSize % TileM == 0
-    && BlockSize % TileN == 0
-    && BlockSize % TileK == 0
-  );
-
-  static constexpr int NumTilePerBlock = BlockSize / TileK;
-
-  using SmemLayoutAtom = decltype(
-    composition(
-      Swizzle<3, 4, 3>{},
-      make_layout(
-        make_shape(Int<8>{}, Int<TileK>{}),
-        make_stride(Int<TileK>{}, Int<1>{})
-      )
-    )
-  );
-
-  using SmemLayoutA = decltype(
-    tile_to_shape(
-      SmemLayoutAtom{},
-      make_shape(Int<TileM>{}, Int<TileK>{}, Int<Stage>{})
-    )
-  );
-
-  using SmemLayoutB = decltype(
-    tile_to_shape(
-      SmemLayoutAtom{},
-      make_shape(Int<TileN>{}, Int<TileK>{}, Int<Stage>{})
-    )
-  );
-
-  using MmaOP = cute::SM80_16x8x32_S32S8S8S32_TN;
-  using TiledMma = decltype(
-    make_tiled_mma(
-      MMA_Atom<MMA_Traits<MmaOP>>{},
-      make_layout(make_shape(
-        _4{}, _2{}, _1{}
-      )),
-      make_tile(Int<64>{}, Int<32>{}, Int<32>{})
-    )
-  );
-
-  using G2SCopyAtomA = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>>, ElementA>;
-  using G2SCopyAtomB = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>>, ElementB>;
-  using G2STiledCopyA = decltype(
-    make_tiled_copy(
-      G2SCopyAtomA{},
-      make_layout(
-        make_shape(Int<64>{}, Int<4>{}),
-        make_stride(Int<4>{}, Int<1>{})
-      ),
-      make_layout(make_shape(Int<1>{}, Int<16>{}))
-    )
-  );
-  using G2STiledCopyB = decltype(
-    make_tiled_copy(
-      G2SCopyAtomB{},
-      make_layout(
-        make_shape(Int<64>{}, Int<4>{}),
-        make_stride(Int<4>{}, Int<1>{})
-      ),
-      make_layout(make_shape(Int<1>{}, Int<16>{}))
-    )
-  );
-
-  using S2RCopyAtomA = Copy_Atom<Copy_Traits<SM75_U32x4_LDSM_N>, ElementA>;
-  using S2RCopyAtomB = Copy_Atom<Copy_Traits<SM75_U32x4_LDSM_N>, ElementB>;
-  using S2RTiledCopyA = decltype(make_tiled_copy_A(S2RCopyAtomA{}, TiledMma{}));
-  using S2RTiledCopyB = decltype(make_tiled_copy_B(S2RCopyAtomB{}, TiledMma{}));
-
-  // epilogue
-  using SmemLayoutAtomD = decltype(
-    composition(
-      Swizzle<2, 3, 3>{},
-      make_layout(
-        make_shape(Int<32>{}, Int<32>{}),
-        LayoutRight{}
-      )
-    )
-  );
-
-  using SmemLayoutD = decltype(
-    tile_to_shape(
-      SmemLayoutAtomD{},
-      make_shape(Int<64>{}, Int<32>{}, Int<EpiStage>{})
-    )
-  );
-
-  using R2SCopyAtomD = Copy_Atom<UniversalCopy<std::conditional_t<sizeof(OutputDtype) == 4, int32_t, int16_t>>, OutputDtype>;
-  using R2STiledCopyD = decltype(make_tiled_copy_C(R2SCopyAtomD{}, TiledMma{}));
-
-  using S2GCopyAtomD = Copy_Atom<UniversalCopy<uint128_t>, OutputDtype>;
-  using S2GCopyD = decltype(make_tiled_copy(
-    S2GCopyAtomD{},
-    make_layout(Shape<_64, _4>{}),
-    make_layout(Shape<_1, _8>{})
-  ));
-
-  using TileShape = decltype(make_shape(Int<TileM>{}, Int<TileN>{}, Int<TileK>{}));
-
-  struct SharedStorageAB: cute::aligned_struct<128> {
-    array_aligned<typename TiledMma::ValTypeA, cosize_v<SmemLayoutA>, 128> smem_A;
-    array_aligned<typename TiledMma::ValTypeB, cosize_v<SmemLayoutB>, 128> smem_B;
-    array_aligned<float, 1> smem_AS;
-    array_aligned<float, 1> smem_BS;
-    array_aligned<int32_t, 1> smem_AF;
-  };
-
-  struct SharedStorageD: cute::aligned_struct<128> {
-    array_aligned<OutputDtype, cosize_v<SmemLayoutD>> smem_D;
-  };
-
-  union SharedStorage {
-    SharedStorageAB storage_AB;
-    SharedStorageD storage_D;
-  };
-
-
-  struct Params {
-    void const* Aptr;
-    void const* ASptr;
-    void const* Bptr;
-    void const* BSptr;
-    void* Dptr;
-    int64_t const m;
-    int64_t const n;
-    int64_t const k;
-    int const swizzle_dir;
-    int const swizzle_size;
-  };
-
-  using Arguments = Params;
-
-  static constexpr int ThreadNum = size(TiledMma{});
-  static constexpr int ShmSize = sizeof(SharedStorage);
-  static constexpr bool FastInt2Float = false;
-
-  static bool can_implement(int64_t m, int64_t n, int64_t k) {
-    if (k % BlockSize != 0) return false;
-    if ((n * sizeof(OutputDtype)) % 16 != 0)
-      return false;
-    return true;
-  }
-
-  static Params to_underlying_arguments(Arguments const& args) {
-    return args;
-  }
-
-  static dim3 get_grid_size(int64_t m, int64_t n) {
-    return dim3(cdiv(m, TileM) * cdiv(n, TileN));
-  }
-
-  CUTLASS_HOST_DEVICE 
-  static auto get_block_coord(
-    int64_t m_blocks, 
-    int64_t n_blocks, 
-    int const swizzle_dir,
-    int64_t const swizzle_size_log
-  ) {
-    int64_t blk_m;
-    int64_t blk_n;
-
-    if (swizzle_dir == 1)
-      std::swap(m_blocks, n_blocks);
-
-    if (swizzle_size_log == 0) {
-      blk_m = blockIdx.x % m_blocks;
-      blk_n = blockIdx.x / m_blocks;
-    } else {
-      int64_t group_size = n_blocks << swizzle_size_log;
-      int64_t num_groups = m_blocks >> swizzle_size_log;
-      int64_t group_idx = blockIdx.x / group_size;
-      int64_t local_idx = blockIdx.x % group_size;
-      if (group_idx == num_groups) {
-        blk_m = (num_groups << swizzle_size_log) + local_idx % (m_blocks - (num_groups << swizzle_size_log));
-        blk_n = local_idx / (m_blocks - (num_groups << swizzle_size_log));
-      } else {
-        blk_m = (local_idx & ((1LL << swizzle_size_log) - 1)) + (group_idx << swizzle_size_log);
-        blk_n = local_idx >> swizzle_size_log;
-      }
-    }
-
-    if (swizzle_dir == 1)
-      std::swap(blk_m, blk_n);
-
-    return make_coord(blk_m, blk_n);
-  }
-
-  CUTLASS_DEVICE
-  void operator()(
-    Params const& params, char* smem_data
-  ) {
-
-    SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_data);
-
-    auto t_idx = threadIdx.x;
-
-    int64_t const m = params.m;
-    int64_t const n = params.n;
-    int64_t const k = params.k;
-    int const swizzle_dir = params.swizzle_dir;
-    int const swizzle_size = params.swizzle_size;
-
-    Tensor A = make_tensor(
-      make_gmem_ptr<ElementA>(params.Aptr),
-      make_shape(m, k),
-      make_stride(k, _1{})
-    );
-    Tensor B = make_tensor(
-      make_gmem_ptr<ElementB>(params.Bptr),
-      make_shape(m, k),
-      make_stride(k, _1{})
-    );
-    Tensor AS = make_tensor(
-      make_gmem_ptr<float>(params.ASptr),
-      make_shape(cdiv(m, BlockSize), cdiv(k, BlockSize)),
-      make_stride(cdiv(k, BlockSize), _1{})
-    );
-    Tensor BS = make_tensor(
-      make_gmem_ptr<float>(params.BSptr),
-      make_shape(cdiv(n, BlockSize), cdiv(k, BlockSize)),
-      make_stride(cdiv(k, BlockSize), _1{})
-    );
-    Tensor D = make_tensor(
-      make_gmem_ptr<OutputDtype>(params.Dptr),
-      make_shape(m, n),
-      LayoutRight{}
-    );
-
-    auto [m_coord, n_coord] = get_block_coord(
-      cdiv(m, size<0>(TileShape{})),
-      cdiv(n, size<1>(TileShape{})),
-      swizzle_dir, swizzle_size
-    );
-
-    int32_t blk_m_coord = m_coord / (BlockSize / TileM);
-    int32_t blk_n_coord = n_coord / (BlockSize / TileN);
-
-    // local tile
-    auto gA = local_tile(A, TileShape{}, make_coord(m_coord, n_coord, _), Step<_1, X, _1>{});
-    auto gB = local_tile(B, TileShape{}, make_coord(m_coord, n_coord, _), Step<X, _1, _1>{});
-    auto gD = local_tile(D, TileShape{}, make_coord(m_coord, n_coord, _), Step<_1, _1, X>{});
-
-    // shared memory
-    Tensor sA = make_tensor(
-      make_smem_ptr<ElementA>(shared_storage.storage_AB.smem_A.data()),
-      SmemLayoutA{}
-    );
-
-    Tensor sB = make_tensor(
-      make_smem_ptr<ElementB>(shared_storage.storage_AB.smem_B.data()),
-      SmemLayoutB{}
-    );
-
-    // register
-    TiledMma tiled_mma;
-    auto thr_mma = tiled_mma.get_slice(t_idx);
-    auto tCrA = thr_mma.partition_fragment_A(gA(_, _, 0));
-    auto tCrB = thr_mma.partition_fragment_B(gB(_, _, 0));
-    auto tDrC = thr_mma.partition_fragment_C(gD); // mma accumulator
-    auto tDrD = make_tensor_like<float>(tDrC); // float accumulator
-
-    if constexpr (FastInt2Float) {
-      CUTLASS_PRAGMA_UNROLL
-      for (int i = 0; i < size(tDrC); ++i)
-        tDrC(i) = 0x4B400000;
-    } else {
-      clear(tDrC);
-    }
-
-    clear(tDrD);
-
-
-    // global to shared copy
-    G2STiledCopyA g2s_tiled_copy_a;
-    auto g2s_thr_copy_a = g2s_tiled_copy_a.get_slice(t_idx);
-    auto tAgA = g2s_thr_copy_a.partition_S(gA);
-    auto tAsA = g2s_thr_copy_a.partition_D(sA);
-    auto cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA)));
-    auto tAcA = g2s_thr_copy_a.partition_S(cA);
-    int const m_limit = m - TileM * m_coord;
-    int const n_limit = n - TileN * n_coord;
-
-    G2STiledCopyB g2s_tiled_copy_b;
-    auto g2s_thr_copy_b = g2s_tiled_copy_b.get_slice(t_idx);
-    auto tBgB = g2s_thr_copy_b.partition_S(gB);
-    auto tBsB = g2s_thr_copy_b.partition_D(sB);
-    auto cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB)));
-    auto tBcB = g2s_thr_copy_a.partition_S(cB);
-
-
-    // shared to register copy
-    S2RTiledCopyA s2r_tiled_copy_a;
-    auto s2r_thr_copy_a = s2r_tiled_copy_a.get_slice(t_idx);
-    auto tCsA = s2r_thr_copy_a.partition_S(sA);
-    auto tCrA_view = s2r_thr_copy_a.retile_D(tCrA);
-
-    S2RTiledCopyB s2r_tiled_copy_b;
-    auto s2r_thr_copy_b = s2r_tiled_copy_b.get_slice(t_idx);
-    auto tCsB = s2r_thr_copy_b.partition_S(sB);
-    auto tCrB_view = s2r_thr_copy_b.retile_D(tCrB);
-
-    // pipeline status
-    int64_t g2s_a_tile = 0;
-    int64_t g2s_b_tile = 0;
-    int g2s_a_smem = 0;
-    int g2s_b_smem = 0;
-
-    int g2s_tile_in_block = 0;
-    int g2s_block = 0; // b block idx
-
-    int s2r_a_smem = 0;
-    int s2r_b_smem = 0;
-    int s2r_tile_in_block = 0;
-
-    int mma_block_a = 0;
-    int mma_block_b = 0;
-
-    int ntile = k / TileK;
-    // load scale and fallback
-    // we assume all ptrs are 128bit aligned
-    // auto smem_fallback_A = raw_pointer_cast(make_smem_ptr<int32_t>(shared_storage.storage_AB.smem_AF.data()));
-    // auto smem_scale_A = raw_pointer_cast(make_smem_ptr<float>(shared_storage.storage_AB.smem_AS.data()));
-    // auto smem_scale_B = raw_pointer_cast(make_smem_ptr<float>(shared_storage.storage_AB.smem_BS.data()));
-    __syncthreads();
-
-
-    int32_t fallbackA_load = 0;
-    int32_t fallbackA_mma = 0;
-
-    // copy first Stage - 1 tile
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0, _i = min(Stage - 1, ntile); i < _i; ++i) {
-      if (g2s_b_tile < ntile) {
-        g2s_tile_in_block = (g2s_tile_in_block + 1) % NumTilePerBlock;
-        copy_AB<IsEvenM>(g2s_tiled_copy_a, tAgA, tAsA, tAcA, g2s_a_tile, g2s_a_smem, m_limit);
-        copy_AB<IsEvenN>(g2s_tiled_copy_b, tBgB, tBsB, tBcB, g2s_b_tile, g2s_b_smem, n_limit);
-        ++g2s_b_tile;
-        ++g2s_b_smem;
-        ++g2s_block;
-        g2s_a_tile = g2s_block * NumTilePerBlock;
-        ++g2s_a_smem;
-      }
-      cp_async_fence();
-    }
-
-    constexpr int nk = size<2>(tCrA);
-    float scale_a = AS(blk_m_coord, 0);
-    float scale_b = BS(blk_n_coord, 0);
-
-    CUTLASS_PRAGMA_NO_UNROLL
-    for (int64_t mma_b_tile = 0; mma_b_tile < ntile; ++mma_b_tile) {
-      s2r_tile_in_block = (s2r_tile_in_block + 1) % NumTilePerBlock;
-      cp_async_wait<Stage - 2>();
-      __syncthreads();
-
-      // do mma first
-      CUTLASS_PRAGMA_UNROLL
-      for (int ik = 0; ik < nk; ++ik) {
-        cute::copy(s2r_tiled_copy_a, tCsA(_, _, ik, s2r_a_smem),
-                    tCrA_view(_, _, ik));
-        cute::copy(s2r_tiled_copy_b, tCsB(_, _, ik, s2r_b_smem),
-                    tCrB_view(_, _, ik));
-        cute::gemm(tiled_mma, tDrC, tCrA(_, _, ik), tCrB(_, _, ik), tDrC);
-      }
-
-      // a s2r increase anyway
-      s2r_a_smem = (s2r_a_smem + 1) % Stage;
-
-      // get next s2r b tile int64_t
-        // end of a block
-
-        // dequant first
-      dequant<AccumulatorDtype, TileM * TileN / ThreadNum, FastInt2Float>(
-        tDrC.data(), tDrD.data(), scale_a * scale_b
-      );
-
-      s2r_b_smem = (s2r_b_smem + 1) % Stage;
-      // b advance
-      ++mma_block_b;
-      if (mma_block_b < size<1>(BS)) scale_b = BS(blk_n_coord, mma_block_b);
-      mma_block_a = mma_block_b;
-      if (mma_block_a < size<1>(AS)) scale_a = AS(blk_m_coord, mma_block_a);
-
-      // load next stage
-      if (g2s_b_tile < ntile) {
-        g2s_tile_in_block = (g2s_tile_in_block + 1) % NumTilePerBlock;
-        copy_AB<IsEvenM>(g2s_tiled_copy_a, tAgA, tAsA, tAcA, g2s_a_tile, g2s_a_smem, m_limit);
-        copy_AB<IsEvenN>(g2s_tiled_copy_b, tBgB, tBsB, tBcB, g2s_b_tile, g2s_b_smem, n_limit);
-        ++g2s_b_tile;
-        g2s_b_smem = (g2s_b_smem + 1) % Stage;
-        ++g2s_block;
-        g2s_a_tile = g2s_block * NumTilePerBlock;
-        g2s_a_smem = (g2s_a_smem + 1) % Stage;
-      }
-      cp_async_fence();
-    }
-
-
-    // epilogue
-
-    Tensor sD = make_tensor(
-      make_smem_ptr<OutputDtype>(shared_storage.storage_D.smem_D.data()),
-      SmemLayoutD{}
-    );
-
-    R2STiledCopyD r2s_tiled_copy_d;
-    auto r2s_thr_copy_d = r2s_tiled_copy_d.get_slice(t_idx);
-    auto tDrD_r2s = r2s_thr_copy_d.retile_S(tDrD);
-    auto tDsD_r2s = r2s_thr_copy_d.partition_D(sD);
-
-    S2GCopyD s2g_tiled_copy_d;
-    auto s2g_thr_copy_d = s2g_tiled_copy_d.get_slice(t_idx);
-    auto tDsD_s2g = s2g_thr_copy_d.partition_S(sD);
-    auto tDgD_s2g = s2g_thr_copy_d.partition_D(gD);
-    Tensor cD = make_identity_tensor(make_shape(Int<TileM>{}, Int<TileN>{}));
-    auto tDcD_s2g = s2g_thr_copy_d.partition_D(cD);
-
-    auto tDgD_s2gx = group_modes<1, 3>(tDgD_s2g);  // (CPY_, CPY_MN)
-    auto tDrD_r2sx = group_modes<1, 3>(tDrD_r2s);  // (CPY_, CPY_MN)
-    auto tDcD_s2gx = group_modes<1, 3>(tDcD_s2g);
-    
-    int32_t step = size<3>(tDsD_r2s);  // pipe
-    CUTLASS_PRAGMA_UNROLL
-    for (int32_t i = 0; i < size<1>(tDrD_r2sx); i += step) {
-      CUTLASS_PRAGMA_UNROLL
-      for (int32_t j = 0; j < step; ++j) {
-        if constexpr (std::is_same<OutputDtype, float>::value) {
-            cute::copy(r2s_tiled_copy_d, tDrD_r2sx(_, i + j), tDsD_r2s(_, 0, 0, j));
-        } else {
-            auto t = make_tensor_like<OutputDtype>(tDrD_r2sx(_, i + j));
-            cute::copy(tDrD_r2sx(_, i + j), t);
-            cute::copy(r2s_tiled_copy_d, t, tDsD_r2s(_, 0, 0, j));
-        }
-      }
-
-      __syncthreads();
-
-      // shm -> global
-      if constexpr (IsEvenM && IsEvenN) {
-        CUTLASS_PRAGMA_UNROLL
-        for (int32_t j = 0; j < step; ++j)
-          cute::copy(s2g_tiled_copy_d, tDsD_s2g(_, 0, 0, j), tDgD_s2gx(_, i + j));
-      } else if constexpr (IsEvenN) {
-          CUTLASS_PRAGMA_UNROLL
-          for (int32_t j = 0; j < step; ++j) {
-            if (get<0>(tDcD_s2gx(0, i + j)) < m_limit)
-              cute::copy(s2g_tiled_copy_d, tDsD_s2g(_, 0, 0, j), tDgD_s2gx(_, i + j));
-          }
-      } else if constexpr (IsEvenM) {
-          CUTLASS_PRAGMA_UNROLL
-          for (int32_t j = 0; j < step; ++j)
-            if (get<1>(tDcD_s2gx(size<0>(tDsD_s2g) - 1, i + j)) < n_limit) {
-              cute::copy(s2g_tiled_copy_d, tDsD_s2g(_, 0, 0, j), tDgD_s2gx(_, i + j));
-            } else {
-              CUTLASS_PRAGMA_UNROLL
-              for (int k = 0; k < size<0>(tDsD_s2g); ++k)
-                if (get<1>(tDcD_s2gx(k, i + j)) < n_limit)
-                  tDgD_s2gx(k, i + j) = tDsD_s2g(k, 0, 0, j);
-            }
-      } else {
-          CUTLASS_PRAGMA_UNROLL
-          for (int32_t j = 0; j < step; ++j)
-            if (get<0>(tDcD_s2gx(0, i + j)) < m_limit) {
-              if (get<1>(tDcD_s2gx(size<0>(tDsD_s2g) - 1, i + j)) < n_limit) {
-                cute::copy(s2g_tiled_copy_d, tDsD_s2g(_, 0, 0, j), tDgD_s2gx(_, i + j));
-              } else {
-                for (int32_t k = 0; k < size<0>(tDsD_s2g); ++k)
-                  if (get<1>(tDcD_s2gx(k, i + j)) < n_limit)
-                    tDgD_s2gx(k, i + j) = tDsD_s2g(k, 0, 0, j);
-              }
-            }
-      }
-      __syncthreads();
-    }
-  }
-
-};
-  

fastvideo-kernel/csrc/turbodiffusion/gemm/launch.hpp

@@ -1,64 +0,0 @@
-/* 
- * Copyright (c) 2025 by TurboDiffusion team.
- * 
- * Licensed under the Apache License, Version 2.0 (the "License");
- * 
- * Citation (please cite if you use this code):
- * 
- * @article{zhang2025turbodiffusion,
- *   title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times},
- *   author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun},
- *   journal={arXiv preprint arXiv:2512.16093},
- *   year={2025}
- * }
- */
-
-#pragma once
-
-#include "common/common.hpp"
-#include "common/launch.hpp"
-#include "gemm/kernel.hpp"
-
-
-template <class OutputDtype>
-bool int8_gemm_(
-  int8_t const *Aptr, float const *ASptr,
-  int8_t const *Bptr, float const *BSptr,
-  OutputDtype* Dptr, int64_t m, int64_t n, int64_t k,
-  int swizzle_dir = 1, int swizzle_size_log = 0,
-  cudaStream_t stream = nullptr
-) {
-  BOOL_SWITCH(m % 128 == 0, IsEvenM, [&] {
-    BOOL_SWITCH(n % 128 == 0, IsEvenN, [&] {
-      using Kernel = GemmKernel<OutputDtype, IsEvenM, IsEvenN>;
-      if (!Kernel::can_implement(m, n, k))
-        return false;
-      using Args = typename Kernel::Arguments;
-      Args args {
-        (void*)Aptr, (void*)ASptr,
-        (void*)Bptr, (void*)BSptr, (void*)Dptr,
-        m, n, k, swizzle_dir,
-        swizzle_size_log
-      };
-
-      auto params = Kernel::to_underlying_arguments(args);
-
-      static constexpr size_t ShmSize = Kernel::ShmSize;
-      dim3 grid_shape = Kernel::get_grid_size(m, n);
-      dim3 block_shape = dim3(Kernel::ThreadNum);
-      auto func = device_kernel<Kernel>;
-      if (ShmSize >= 48 * 1024) {
-        cudaFuncSetAttribute(
-          func,
-          cudaFuncAttributeMaxDynamicSharedMemorySize,
-          ShmSize
-        );
-      }
-      func<<<grid_shape, block_shape, ShmSize, stream>>>(
-        params
-      );
-      return true;
-    });
-  });
-  return true;
-}

fastvideo-kernel/csrc/turbodiffusion/gemm/utils.hpp

@@ -1,129 +0,0 @@
-/* 
- * Copyright (c) 2025 by TurboDiffusion team.
- * 
- * Licensed under the Apache License, Version 2.0 (the "License");
- * 
- * Citation (please cite if you use this code):
- * 
- * @article{zhang2025turbodiffusion,
- *   title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times},
- *   author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun},
- *   journal={arXiv preprint arXiv:2512.16093},
- *   year={2025}
- * }
- */
-
-#pragma once
-
-#include <cuda.h>
-#include "cute/tensor.hpp"
-
-template <
-  bool IsEven, 
-  class TiledCopy,
-  class SrcTensor,
-  class DstTensor,
-  class PrdTensor
->
-CUTLASS_DEVICE void
-copy_AB(
-    TiledCopy const& _copy, 
-    SrcTensor const &S, 
-    DstTensor &D, 
-    PrdTensor const &ID, 
-    const int64_t &i_read, 
-    const int64_t &i_write, 
-    const int64_t &limit
-) {
-  using namespace cute;
-  if constexpr (IsEven)
-    cute::copy(_copy, S(_, _, _, i_read), D(_, _, _, i_write));
-  else {
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < size<1>(ID); ++i)
-      if (get<0>(ID(0, i, 0)) < limit)
-        cute::copy(_copy, S(_, i, _, i_read), D(_, i, _, i_write));
-  }
-}
-
-template <int N>
-CUTLASS_DEVICE void copy_async(
-  void const* gmem_src,
-  void* smem_dst
-) {
-  uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst));;
-  asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], %2;\n"
-      :: "r"(smem_int_ptr),
-          "l"(gmem_src),
-          "n"(N));
-}
-
-template<class LoadType, class T, int NumThreads>
-CUTLASS_DEVICE void copy_aligned(const void* src, void* dst, size_t N, int64_t thread_idx) {
-  static constexpr int NumElementPerLoad = sizeof(LoadType) / sizeof(T);
-  for (int64_t i = thread_idx * NumElementPerLoad; i < N; i += NumElementPerLoad * NumThreads) {
-    if (i + NumElementPerLoad <= N) {
-      copy_async<sizeof(LoadType)>(
-        (void*)((T*)src + i),
-        (void*)((T*)dst + i)
-      );
-    } else {
-      for (int64_t j = 0; j < N - i; ++j)
-        copy_async<sizeof(T)>(
-          (void*)((T*)src + i + j),
-          (void*)((T*)dst + i + j)
-        );
-    }
-  }
-}
-
-template<class T, int NumThreads, bool Wait = true, bool Commit = true>
-CUTLASS_DEVICE void g2s_vector_copy(const void* src, void* dst, size_t N, int64_t thread_idx) {
-  
-  uintptr_t src_addr = reinterpret_cast<uintptr_t>(src);
-
-  if (src_addr % 16 == 0) {
-      copy_aligned<int4, T, NumThreads>(src, dst, N, thread_idx);
-  } else if (src_addr % 8 == 0) {
-      copy_aligned<int2, T, NumThreads>(src, dst, N, thread_idx);
-  } else if (src_addr % 4 == 0) {
-      copy_aligned<int, T, NumThreads>(src, dst, N, thread_idx);
-  } else {
-    assert(0);
-  }
-  if constexpr (Commit) {
-    asm volatile("cp.async.commit_group;\n" ::);
-  }
-  if constexpr (Wait) {
-    asm volatile("cp.async.wait_all;\n" ::);
-  }
-}
-
-template<class T, int N, bool FastInt2Float>
-CUTLASS_DEVICE
-static void dequant(
-  T* mma_accum_ptr,
-  float* float_accum_ptr,
-  float scale
-) {
-  static int const ic = 0x4B400000;
-  if constexpr (FastInt2Float && std::is_same_v<T, int32_t>) {
-    CUTLASS_PRAGMA_UNROLL
-    for (size_t i = 0; i < N; ++i) {
-      *(float_accum_ptr + i) += (__int_as_float(*(mma_accum_ptr + i)) - __int_as_float(ic)) * scale;
-      *(mma_accum_ptr + i) = ic;
-    }
-  } else if constexpr (std::is_same_v<T, int32_t>) {
-    CUTLASS_PRAGMA_UNROLL
-    for (size_t i = 0; i < N; ++i) {
-      *(float_accum_ptr + i) += __int2float_rn(*(mma_accum_ptr + i)) * scale;
-      *(mma_accum_ptr + i) = 0;
-    }
-  } else {
-    CUTLASS_PRAGMA_UNROLL
-    for (size_t i = 0; i < N; ++i) {
-      *(float_accum_ptr + i) += (*(mma_accum_ptr + i)) * scale;
-      *(mma_accum_ptr + i) = 0;
-    }
-  }
-}

fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu

@@ -1,63 +0,0 @@
-#include <pybind11/pybind11.h>
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <torch/python.h>
-#include <cutlass/cutlass.h>
-#include "common/common.hpp"
-#include "norm/layernorm.hpp"
-
-auto layer_norm(
-  at::Tensor const Input, 
-  float eps,
-  std::optional<at::Tensor const> W,
-  std::optional<at::Tensor const> const B,
-  std::optional<at::Tensor> Output
-) {
-  using ElementIn = float;
-  using ElementOut = float;
-  using ElementWeight = float;
-
-  int64_t const m = Input.size(0);
-  int64_t const n = Input.size(1);
-  torch::Device const input_device = Input.device();
-
-  if (!Output.has_value()) {
-    Output.emplace(
-      torch::empty(
-        {m, n},
-        torch::TensorOptions().device(input_device).dtype(torch::kFloat32)
-      )
-    );
-  }
-
-
-  void *Iptr = Input.data_ptr();
-  void *Wptr = W.has_value() ? W.value().data_ptr() : nullptr;
-  void *Bptr = B.has_value() ? B.value().data_ptr() : nullptr;
-  void *Optr = Output.value().data_ptr();
-
-  BOOL_SWITCH(B.has_value(), BIAS, [&]{
-    BOOL_SWITCH(W.has_value(), AFFINE, [&]{
-    CONFIG_SWITCH(n, [&]{
-    layernorm<
-    ElementIn, ElementOut, ElementWeight,
-    AFFINE, BIAS,
-    MAX_HIDDEN_SIZE, NUM_THR_PER_CTA> (
-      Iptr, Wptr, Bptr,
-      Optr, eps, m, n,
-      at::cuda::getCurrentCUDAStream().stream()
-        );
-      });
-    });
-  });
-  
-    
-
-  return Output;
-}
-
-void register_layer_norm(pybind11::module_ &m) {
-    m.def("layer_norm_cuda", &layer_norm);
-}
-

fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp

@@ -1,202 +0,0 @@
-#pragma once
-
-#include "common/load.hpp"
-#include "common/store.hpp"
-#include "common/launch.hpp"
-
-
-template <
-  class InputDtype_,
-  class OutputDtype_,
-  class WeightDtype_,
-  bool Affine_,
-  bool Bias_,
-  int MaxHiddenSize_,
-  int NumThrPerCta_,
-  bool IsEven
->
-class LayerNorm {
-public:
-  using InputDtype = InputDtype_;
-  using OutputDtype = OutputDtype_;
-  using WeightDtype = WeightDtype_;
-  static constexpr int NumThrPerCta = NumThrPerCta_;
-  static constexpr int MaxHiddenSize = MaxHiddenSize_;
-  static constexpr bool Affine = Affine_;
-  static constexpr bool Bias = Bias_;
-
-  static constexpr size_t ShmSize = 32;
-  static constexpr int NumElementPerThread = MaxHiddenSize / NumThrPerCta;
-
-  static_assert(MaxHiddenSize % NumThrPerCta == 0);
-
-  struct Params {
-    void const *Iptr;
-    void const *Wptr;
-    void const *Bptr;
-    void *Optr;
-    float eps;
-    int64_t m;
-    int64_t n;
-  };
-
-  using Arguments = Params;
-
-  static Params to_underlying_arguments(Arguments const& args) {
-    return args;
-  }
-
-  static dim3 get_grid_size(int64_t m, int64_t n) {
-    return dim3(m);
-  }
-
-  static dim3 get_cta_size(int64_t m, int64_t n) {
-    return dim3(NumThrPerCta, 1, 1);
-  }
-
-  CUTLASS_DEVICE
-  void operator()(Params const& params, char *shared_data) {
-    int const blk_m = blockIdx.x;
-    int const blk_n = 1;
-    int tidx = threadIdx.x;
-    float x[NumElementPerThread];
-
-    // load
-    Loader<InputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> loader;
-    loader.load(params.Iptr, x, params.m, params.n, blk_m, 0, tidx);
-
-    // mean reduction
-    float u = _reduce_sum(x, shared_data) / params.n;
-
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      x[i] -= u;
-
-    __syncthreads();
-    // var reduction
-    float v = sqrtf(_reduce_square(x, shared_data) / params.n + params.eps);
-
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      x[i] /= v;
-
-    if constexpr (Affine) {
-      // load weight
-      Loader<WeightDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> weight_loader;
-      float w[NumElementPerThread];
-      weight_loader.load(params.Wptr, w, 1, params.n, 0, 0, tidx);
-      if constexpr (Bias) {
-        float b[NumElementPerThread];
-        weight_loader.load(params.Bptr, b, 1, params.n, 0, 0, tidx);
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < NumElementPerThread; ++i)
-          x[i] = x[i] * w[i] + b[i];
-      } else {
-        CUTLASS_PRAGMA_UNROLL
-        for (int i = 0; i < NumElementPerThread; ++i)
-          x[i] = x[i] * w[i];
-      }
-    }
-
-    // save y
-    {
-      Saver<OutputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven, false, false> saver;
-      if constexpr (std::is_same_v<OutputDtype, float>) {
-        saver.store(params.Optr, nullptr, x, 0, params.m, params.n, blk_m, 0, tidx);
-      } else {
-        OutputDtype tmp[NumElementPerThread];
-        for (int i = 0; i < NumElementPerThread; ++i)
-          tmp[i] = OutputDtype(x[i]);
-        saver.store(params.Optr, nullptr, tmp, 0, params.m, params.n, blk_m, 0, tidx);
-      }
-    }
-    
-  }
-
-private:
-  CUTLASS_DEVICE
-  float _reduce_square(float *reg, char *shared_data) {
-    // thread
-    float sum_square = 0;
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      sum_square += reg[i] * reg[i];
-
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 16; i >= 1; i >>= 1) {
-      sum_square += __shfl_down_sync(0xFFFFFFFF, sum_square, i);
-    }
-    if (threadIdx.x == 0) {
-      *(float*)shared_data = 0;
-    }
-    __syncthreads();
-
-    if (threadIdx.x % 32 == 0) {
-      atomicAdd((float*)shared_data, sum_square);
-    }
-
-    __syncthreads();
-    sum_square = *(float*)shared_data;
-    return sum_square;
-  }
-
-  CUTLASS_DEVICE
-  float _reduce_sum(float *reg, char *shared_data) {
-    // thread
-    float sum = 0;
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      sum += reg[i];
-
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 16; i >= 1; i >>= 1) {
-      sum += __shfl_down_sync(0xFFFFFFFF, sum, i);
-    }
-    if (threadIdx.x == 0) {
-      *(float*)shared_data = 0;
-    }
-    __syncthreads();
-
-    if (threadIdx.x % 32 == 0) {
-      atomicAdd((float*)shared_data, sum);
-    }
-
-    __syncthreads();
-    sum = *(float*)shared_data;
-    return sum;
-  }
-};
-
-template <
-  class InputDtype,
-  class OutputDtype,
-  class WeightDtype,
-  bool Affine,
-  bool Bias,
-  int MaxHiddenSize,
-  int NumThrPerCta
->
-bool layernorm(
-  void const *Iptr, void const *Wptr, void const *Bptr,
-  void *Optr, float eps, int64_t m, int64_t n, 
-  cudaStream_t stream = nullptr
-) {
-  BOOL_SWITCH(n % MaxHiddenSize == 0, IsEven, [&] {
-    using Kernel = LayerNorm<
-      InputDtype, OutputDtype, WeightDtype,
-      Affine, Bias,
-      MaxHiddenSize, NumThrPerCta,
-      IsEven>;
-    using Arguments = typename Kernel::Arguments;
-    Arguments args = {
-      Iptr, Wptr, Bptr, Optr, 
-      eps, m, n
-    };
-    auto params = Kernel::to_underlying_arguments(args);
-    auto grid_shape = Kernel::get_grid_size(m, n);
-    auto cta_shape = Kernel::get_cta_size(m, n);
-    static constexpr size_t ShmSize = Kernel::ShmSize;
-    launch_kernel<Kernel>(params, grid_shape, cta_shape, ShmSize, stream);
-  });
-  return true;
-}

fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu

@@ -1,59 +0,0 @@
-#include <pybind11/pybind11.h>
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <torch/python.h>
-#include <cutlass/cutlass.h>
-#include <pybind11/pybind11.h>
-
-#include "common/common.hpp"
-#include "norm/rmsnorm.hpp"
-
-auto rms_norm(
-  at::Tensor const& Input, 
-  float eps,
-  const std::optional<at::Tensor>& Weight,
-  std::optional<at::Tensor>& Output
-) {
-  
-  using ElementIn = float;
-  using ElementOut = float;
-  using ElementWeight = float;
-  
-  int64_t const m = Input.size(0);
-  int64_t const n = Input.size(1);
-  torch::Device const input_device = Input.device();
-  
-  if (!Output.has_value()) {
-    Output.emplace(
-      torch::empty(
-        {m, n},
-        torch::TensorOptions().device(input_device).dtype(torch::kFloat32)
-      )
-    );
-  }
-
-  void *Iptr = Input.data_ptr();
-  void *Wptr = Weight.has_value() ? Weight.value().data_ptr() : nullptr;
-  void *Optr = Output.value().data_ptr();
-
-
-  CONFIG_SWITCH(n, [&]{
-    rmsnorm<
-    ElementIn, ElementOut, ElementWeight,
-    MAX_HIDDEN_SIZE, NUM_THR_PER_CTA
-    > (
-      Iptr, Wptr,
-      Optr, 
-      eps, m, n, 
-      at::cuda::getCurrentCUDAStream().stream()
-    );
-});
-  
-
-  return Output;
-}
-
-void register_rms_norm(pybind11::module_ &m) {
-    m.def("rms_norm_cuda", &rms_norm);
-}

fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp

@@ -1,147 +0,0 @@
-#pragma once
-
-#include "common/load.hpp"
-#include "common/store.hpp"
-#include "common/launch.hpp"
-
-
-template <
-  class InputDtype_,
-  class OutputDtype_,
-  class WeightDtype_,
-  int MaxHiddenSize_,
-  int NumThrPerCta_,
-  bool IsEven
->
-class RMSNorm {
-public:
-  using InputDtype = InputDtype_;
-  using OutputDtype = OutputDtype_;
-  using WeightDtype = WeightDtype_;
-  static constexpr int NumThrPerCta = NumThrPerCta_;
-  static constexpr int MaxHiddenSize = MaxHiddenSize_;
-
-  static constexpr size_t ShmSize = 32;
-  static constexpr int NumElementPerThread = MaxHiddenSize / NumThrPerCta;
-
-  static_assert(MaxHiddenSize % NumThrPerCta == 0);
-
-  struct Params {
-    void const *Iptr;
-    void const *Wptr;
-    void *Optr;
-    float eps;
-    int64_t m;
-    int64_t n;
-  };
-
-  using Arguments = Params;
-
-  static Params to_underlying_arguments(Arguments const& args) {
-    return args;
-  }
-
-  static dim3 get_grid_size(int64_t m, int64_t n) {
-    return dim3(m);
-  }
-
-  static dim3 get_cta_size(int64_t m, int64_t n) {
-    return dim3(NumThrPerCta, 1, 1);
-  }
-
-  CUTLASS_DEVICE
-  void operator()(Params const& params, char *shared_data) {
-    int const blk_m = blockIdx.x;
-    int const blk_n = 1;
-    int tidx = threadIdx.x;
-    float x[NumElementPerThread];
-
-    // load
-    Loader<InputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> loader;
-    loader.load(params.Iptr, x, params.m, params.n, blk_m, 0, tidx);
-
-    // rms reduction
-    float rms = sqrtf(_reduce_square(x, shared_data) / params.n + params.eps);
-
-    // load weight
-    Loader<WeightDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> weight_loader;
-    float w[NumElementPerThread];
-    loader.load(params.Wptr, w, 1, params.n, 0, 0, tidx);
-
-    // norm
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      x[i] = w[i] * x[i] / rms ;
-
-    // save y
-    OutputDtype *output_reg = (OutputDtype*)x;
-    if constexpr (!std::is_same_v<OutputDtype, float>) {
-      output_reg = (OutputDtype*)w;
-      CUTLASS_PRAGMA_UNROLL
-      for (int i = 0; i < NumElementPerThread; ++i)
-        output_reg[i] = OutputDtype(x[i]);
-    }
-    Saver<OutputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven, false, false> saver;
-    saver.store(params.Optr, nullptr, output_reg, 0, params.m, params.n, blk_m, 0, tidx);
-    
-  }
-
-private:
-  CUTLASS_DEVICE
-  float _reduce_square(float *reg, char *shared_data) {
-    // thread
-    float sum_square = 0;
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      sum_square += reg[i] * reg[i];
-
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 16; i >= 1; i >>= 1) {
-      sum_square += __shfl_down_sync(0xFFFFFFFF, sum_square, i);
-    }
-    if (threadIdx.x == 0) {
-      *(float*)shared_data = 0;
-    }
-    __syncthreads();
-
-    if (threadIdx.x % 32 == 0) {
-      atomicAdd((float*)shared_data, sum_square);
-    }
-
-    __syncthreads();
-    sum_square = *(float*)shared_data;
-    return sum_square;
-  }
-};
-
-
-template <
-  class InputDtype,
-  class OutputDtype,
-  class WeightDtype,
-  int MaxHiddenSize,
-  int NumThrPerCta
->
-bool rmsnorm(
-  void const *Iptr, void const *Wptr, 
-  void *Optr, float eps, 
-  int64_t m, int64_t n, 
-  cudaStream_t stream = nullptr
-) {
-  BOOL_SWITCH(n % MaxHiddenSize == 0, IsEven, [&] {
-    using Kernel = RMSNorm<
-      InputDtype, OutputDtype, WeightDtype,
-      MaxHiddenSize, NumThrPerCta,
-      IsEven>;
-    using Arguments = typename Kernel::Arguments;
-    Arguments args = {
-      Iptr, Wptr, Optr, eps, m, n
-    };
-    auto params = Kernel::to_underlying_arguments(args);
-    auto grid_shape = Kernel::get_grid_size(m, n);
-    auto cta_shape = Kernel::get_cta_size(m, n);
-    static constexpr size_t ShmSize = Kernel::ShmSize;
-    launch_kernel<Kernel>(params, grid_shape, cta_shape, ShmSize, stream);
-  });
-  return true;
-}

fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu

@@ -1,75 +0,0 @@
-/* 
- * Copyright (c) 2025 by TurboDiffusion team.
- * 
- * Licensed under the Apache License, Version 2.0 (the "License");
- * 
- * Citation (please cite if you use this code):
- * 
- * @article{zhang2025turbodiffusion,
- *   title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times},
- *   author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun},
- *   journal={arXiv preprint arXiv:2512.16093},
- *   year={2025}
- * }
- */
-
-#include <pybind11/pybind11.h>
-#include <torch/extension.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <torch/python.h>
-#include <cutlass/cutlass.h>
-#include <cutlass/numeric_types.h>
-#include <pybind11/pybind11.h>
-
-#include "common/common.hpp"
-#include "quant/quant.hpp"
-
-auto quant(
-  torch::Tensor const& Input, 
-  std::optional<torch::Tensor>& Output, 
-  std::optional<torch::Tensor>& Output_S
-) {
-
-  using ElementOut = int8_t;
-  static constexpr int BlockSize = 128;
-  static constexpr int NumThrPerCta = 256;
-
-  int64_t m = Input.size(0);
-  int64_t n = Input.size(1);
-  torch::Device const input_device = Input.device();
-
-  create_tensor<BlockSize>(input_device, Output, Output_S, m, n);
-  
-  ElementOut *Optr = (ElementOut*)Output.value().data_ptr();
-  float *OSptr = Output_S.value().data_ptr<float>();
-
-  switch (Input.scalar_type()) {
-    case torch::kHalf:{
-        cutlass::half_t *Iptr = (cutlass::half_t*)Input.data_ptr();
-        quantization<cutlass::half_t, BlockSize, NumThrPerCta> (
-                    Iptr, Optr, OSptr, m, n, at::cuda::getCurrentCUDAStream().stream()
-                  );
-        break;
-    }
-
-    case torch::kBFloat16:{
-      cutlass::bfloat16_t *Iptr = (cutlass::bfloat16_t*)Input.data_ptr();
-      quantization<cutlass::bfloat16_t, BlockSize, NumThrPerCta> (
-                    Iptr, Optr, OSptr, m, n, at::cuda::getCurrentCUDAStream().stream()
-                  );
-      break;
-    }
-
-    default: {
-      std::cerr << "Observing: " << Input.scalar_type() << " for the input datatype which is invalid";
-      throw std::runtime_error("Unsupported input data type for quantize_to_fp4.");
-    }
-  }
-  
-  return std::make_tuple(Output, Output_S);
-}
-
-void register_quant(pybind11::module_ &m) {
-    m.def("quant_cuda", &quant);
-}

fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp

@@ -1,195 +0,0 @@
-/* 
- * Copyright (c) 2025 by TurboDiffusion team.
- * 
- * Licensed under the Apache License, Version 2.0 (the "License");
- * 
- * Citation (please cite if you use this code):
- * 
- * @article{zhang2025turbodiffusion,
- *   title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times},
- *   author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun},
- *   journal={arXiv preprint arXiv:2512.16093},
- *   year={2025}
- * }
- */
-
-#pragma once
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include "cutlass/numeric_conversion.h"
-
-#include "common/load.hpp"
-#include "common/store.hpp"
-#include "common/launch.hpp"
-
-template <
-  class InputDtype_, 
-  int NumThrPerCta_,
-  bool IsEvenM,
-  bool IsEvenN
->
-class Quantization {
-public:
-  using InputDtype = InputDtype_;
-  using OutputDtype = int8_t;
-  using FPConverter = cutlass::NumericConverter<int8_t, float, cutlass::FloatRoundStyle::round_to_nearest>;
-
-  static constexpr int BlockSize = 128;
-  static constexpr int NumThrPerCta = NumThrPerCta_;
-  static constexpr int NumElementPerThread = BlockSize * BlockSize / NumThrPerCta;
-  static constexpr int NumThrPerRow = BlockSize / NumElementPerThread;
-
-  static_assert(BlockSize * BlockSize % NumThrPerCta == 0);
-  static_assert(NumThrPerCta % BlockSize == 0);
-
-  static constexpr size_t ShmSize = 32;
-
-  static constexpr float int8_max = 128.f;
-
-  struct Params {
-    void const *Iptr;
-    void *Optr;
-    void *OSptr;
-    int64_t const m;
-    int64_t const n;
-  };
-
-  using Arguments = Params;
-
-  static Params to_underlying_arguments(Arguments const& args) {
-    return args;
-  }
-
-  static dim3 get_grid_size(int64_t m, int64_t n) {
-    return dim3(
-      cdiv(n, BlockSize),
-      cdiv(m, BlockSize) 
-    );
-  }
-
-  static dim3 get_cta_size(int64_t m, int64_t n) {
-    return dim3(
-      NumThrPerCta, 1, 1
-    );
-  }
-  
-  CUTLASS_DEVICE
-  void quantization(
-    float *float_reg, 
-    void *Optr, void *OSptr, 
-    int64_t const m, int64_t const n, 
-    int blk_m, int blk_n, int tidx,
-    char *shared_data
-  ) {
-
-    OutputDtype output_reg[NumElementPerThread];
-
-
-    Saver<OutputDtype, BlockSize, BlockSize, NumThrPerCta, IsEvenM, IsEvenN> saver;
-
-    float amax = _reduce_amax(float_reg, (float*)shared_data);
-    
-    
-    _quantization(float_reg, output_reg, int8_max / amax);
-    
-    float scale_inv = amax / int8_max;
-
-    saver.store(Optr, OSptr, output_reg, scale_inv, m, n, blk_m, blk_n, tidx);
-
-    __syncthreads();
-    }
-  
-
-  CUTLASS_DEVICE 
-  void operator()(Params const& params, char *shared_data) {
-    int blk_m = blockIdx.y;
-    int blk_n = blockIdx.x;
-    int tidx = threadIdx.x;
-
-    float float_reg[NumElementPerThread];
-
-    // load float32 data
-    Loader<InputDtype, BlockSize, BlockSize, NumThrPerCta, IsEvenM, IsEvenN> loader;
-    loader.load(params.Iptr, float_reg, params.m, params.n, blk_m, blk_n, tidx);
-    quantization(
-      float_reg, params.Optr, params.OSptr, params.m, params.n, blk_m, blk_n, tidx, shared_data
-    );
-  }
-
-private: 
-
-  CUTLASS_DEVICE float 
-  _reduce_amax(float *reg, float *smem_ptr) {
-    float amax = 1e-8;
-    // thread reduction
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      amax = max(amax, fabs(reg[i]));
-
-    __syncwarp();
-
-    // warp reduction
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 16; i >= 1; i /= 2) {
-      amax = max(
-        __shfl_xor_sync(0xffffffff, amax, i, 32),
-        amax
-      );
-    }
-
-    // cta reduction
-    if (threadIdx.x == 0) {
-      *smem_ptr = 0;
-    }
-    __syncthreads();
-
-    atomicMax((uint32_t*)smem_ptr, reinterpret_cast<const uint32_t&>(amax));
-
-    __syncthreads();
-
-    amax = *smem_ptr;
-
-    return amax;
-  }
-
-  CUTLASS_DEVICE void
-  _quantization(float *float_reg, OutputDtype *out_reg, float scale) {
-    FPConverter converter;
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i) {
-        out_reg[i] = converter(float_reg[i] * scale);
-    }
-
-  }
-};
-
-template <
-  class InputDtype,
-  int BlockSize,
-  int NumThrPerCta
->
-bool quantization(
-  void const *Iptr, void *Optr, void *OSptr,
-  int64_t m, int64_t n, 
-  cudaStream_t stream = nullptr
-) {
-  BOOL_SWITCH(m % BlockSize == 0, IsEvenM, [&] {
-    BOOL_SWITCH(n % BlockSize == 0, IsEvenN, [&] {
-      using Kernel = Quantization<
-        InputDtype, NumThrPerCta, IsEvenM, IsEvenN>;
-      using Arguments = typename Kernel::Arguments;
-      Arguments args = {
-        Iptr, Optr, OSptr,
-        m, n
-      };
-      auto params = Kernel::to_underlying_arguments(args);
-      auto grid_shape = Kernel::get_grid_size(m, n);
-      auto cta_shape = Kernel::get_cta_size(m, n);
-      static constexpr size_t ShmSize = Kernel::ShmSize;
-      launch_kernel<Kernel>(params, grid_shape, cta_shape, ShmSize, stream);
-    });
-  });
-
-  return true;
-}