Merge pull request #388 from mit-han-lab/dev

Merge into to main for CI

.github/workflows/clean-nightly-releases.yaml

+#name: Clean Nightly Releases
+#
+#on:
+#  schedule:
+#    - cron: '0 4 * * *'  # Runs daily at 04:00 UTC
+#  workflow_dispatch:      # Allows manual triggering
+#
+#jobs:
+#  clean-nightly:
+#    runs-on: ubuntu-latest
+#    permissions:
+#      contents: write
+#      packages: write
+#
+#    steps:
+#      - name: Set up GitHub CLI
+#        uses: cli/cli-action@v2
+#
+#      - name: Authenticate GitHub CLI
+#        run: echo "${{ secrets.GITHUB_TOKEN }}" | gh auth login --with-token
+#
+#      - name: Get and clean old nightly releases
+#        env:
+#          GH_REPO: ${{ github.repository }}
+#        run: |
+#          # Get all release titles that match the nightly pattern
+#          mapfile -t releases < <(
+#            gh release list --limit 1000 --json tagName,name,createdAt --jq '
+#              sort_by(.createdAt) | reverse |
+#              map(select(.name | test("^Nunchaku Nightly v[0-9]+\\.[0-9]+\\.[0-9]+dev[0-9]{8}$"))) |
+#              map(.tagName)'
+#          )
+#
+#          # Keep only the first 5 (latest), delete the rest
+#          if (( ${#releases[@]} > 5 )); then
+#            for tag in "${releases[@]:5}"; do
+#              echo "Deleting old release: $tag"
+#              gh release delete "$tag" --yes
+#            done
+#          else
+#            echo "Less than or equal to 5 nightly releases found. No cleanup needed."
+#          fi

.github/workflows/nightly-build.yaml

+#name: Nightly Build
+#
+#on:
+#  schedule:
+#    - cron: '0 8 * * *'
+#  workflow_dispatch:
+#
+#jobs:
+#  tag:
+#    name: Tag Dev Branch if dev version
+#    runs-on: ubuntu-latest
+#    outputs:
+#      is_dev: ${{ steps.check.outputs.is_dev }}
+#      tag_name: ${{ steps.tag.outputs.tag_name }}
+#    steps:
+#      - name: Checkout dev branch
+#        uses: actions/checkout@v4
+#        with:
+#          fetch-depth: 0
+#          ref: dev
+#
+#      - name: Extract version from __version__.py
+#        id: version
+#        run: |
+#          version=$(grep '__version__' nunchaku/__version__.py | sed -E 's/.*"([^"]+)".*/\1/')
+#          echo "version=$version" >> "$GITHUB_OUTPUT"
+#          echo "Extracted version: $version"
+#
+#      - name: Check if version contains 'dev'
+#        id: check
+#        run: |
+#          if [[ "${{ steps.version.outputs.version }}" == *dev* ]]; then
+#            echo "is_dev=true" >> "$GITHUB_OUTPUT"
+#          else
+#            echo "is_dev=false" >> "$GITHUB_OUTPUT"
+#          fi
+#
+#      - name: Set tag name
+#        id: tag
+#        if: steps.check.outputs.is_dev == 'true'
+#        run: |
+#          today=$(date -u +"%Y%m%d")
+#          tag_name="v${{ steps.version.outputs.version }}-$today"
+#          echo "tag_name=$tag_name" >> "$GITHUB_OUTPUT"
+#
+#      - name: Create and push tag
+#        if: steps.check.outputs.is_dev == 'true'
+#        run: |
+#          git config user.name "github-actions"
+#          git config user.email "github-actions@users.noreply.github.com"
+#          git tag ${{ steps.tag.outputs.tag_name }}
+#          git push origin ${{ steps.tag.outputs.tag_name }}
+#
+#      - name: Skip tagging (version is not dev)
+#        if: steps.check.outputs.is_dev == 'false'
+#        run: echo "Version is not a dev version. Skipping tag."
+#
+#  windows-build:
+#    name: Build the windows nightly wheels
+#    runs-on: self-hosted, windows-build
+#    needs: tag
+#    if: needs.tag.outputs.is_dev == 'true'
+#    steps:
+#      - name: Checkout to the tag
+#        uses: actions/checkout@v4
+#        with:
+#          fetch-depth: 0
+#          ref: ${{ needs.tag.outputs.tag_name }}
+#          submodules: true
+#      - name: Show current commit
+#        run: git log -1 --oneline
+#      - name: setup Conda
+#        run: |
+#          pwd
+#      - name: Build all wheels
+#        run: .\scripts\build_all_windows_wheels.cmd
+#      - name: Upload wheels
+#        run: |
+#          echo "Uploading wheels"
+#          mkdir -p wheels
+#          cp dist/*.whl wheels/
+#
+#

.github/workflows/pr_test_ampere.yaml → .github/workflows/test_ampere.yaml

-name: pr_test_ampere
+name: test_ampere
 
 on:
   workflow_dispatch:

.github/workflows/pr_test_blackwell.yaml → .github/workflows/test_blackwell.yaml

-name: pr_test_blackwell
+name: test_blackwell
 
 on:
   workflow_dispatch:

MANIFEST.in

 recursive-include src *.cpp *.h *.hpp *.ipp *.cu *.cuh
-recursive-include nunchaku *.py
+recursive-include nunchaku *.py *.json
 recursive-include nunchaku/csrc *.cpp *.h *.hpp *.ipp *.cu *.cuh
 
 recursive-include third_party/Block-Sparse-Attention/csrc/block_sparse_attn *.cpp *.h *.hpp *.ipp *.cu *.cuh

nunchaku/__version__.py

-__version__ = "0.3.0dev1"
+__version__ = "0.3.0dev"

nunchaku/caching/utils.py

@@ -257,7 +257,7 @@ class SanaCachedTransformerBlocks(nn.Module):
         first_hidden_states_residual = hidden_states - original_hidden_states
         del original_hidden_states
 
-        can_use_cache = get_can_use_cache(
+        can_use_cache, _ = get_can_use_cache(
             first_hidden_states_residual,
             threshold=self.residual_diff_threshold,
             parallelized=self.transformer is not None and getattr(self.transformer, "_is_parallelized", False),
@@ -272,7 +272,7 @@ class SanaCachedTransformerBlocks(nn.Module):
         else:
             if self.verbose:
                 print("Cache miss!!!")
-            set_buffer("first_hidden_states_residual", first_hidden_states_residual)
+            set_buffer("first_multi_hidden_states_residual", first_hidden_states_residual)
             del first_hidden_states_residual
 
             hidden_states, hidden_states_residual = self.call_remaining_transformer_blocks(
@@ -284,7 +284,7 @@ class SanaCachedTransformerBlocks(nn.Module):
                 post_patch_height=post_patch_height,
                 post_patch_width=post_patch_width,
             )
-            set_buffer("hidden_states_residual", hidden_states_residual)
+            set_buffer("multi_hidden_states_residual", hidden_states_residual)
         torch._dynamo.graph_break()
 
         return hidden_states
@@ -399,7 +399,7 @@ class FluxCachedTransformerBlocks(nn.Module):
             controlnet_block_samples = (
                 torch.stack(controlnet_block_samples).to(original_device) if len(controlnet_block_samples) > 0 else None
             )
-        if controlnet_single_block_samples is not None and len(controlnet_single_block_samples) > 0:
+        if controlnet_single_block_samples is not None:
             controlnet_single_block_samples = (
                 torch.stack(controlnet_single_block_samples).to(original_device)
                 if len(controlnet_single_block_samples) > 0

scripts/build_all_linux_wheels.sh

 #!/bin/bash
 # Define the versions for Python, Torch, and CUDA
-python_versions=("3.10" "3.11" "3.12" "3.13")
+python_versions=("3.10" "3.11" "3.12")
 torch_versions=("2.5" "2.6")
 cuda_versions=("12.4")
 
@@ -18,12 +18,10 @@ for python_version in "${python_versions[@]}"; do
   done
 done
 
-bash scripts/build_linux_wheel_torch2.7_cu128.sh "3.10" "2.7" "12.8"
-bash scripts/build_linux_wheel_torch2.7_cu128.sh "3.11" "2.7" "12.8"
-bash scripts/build_linux_wheel_torch2.7_cu128.sh "3.12" "2.7" "12.8"
-bash scripts/build_linux_wheel_torch2.7_cu128.sh "3.13" "2.7" "12.8"
+bash scripts/build_linux_wheel.sh "3.10" "2.7" "12.8"
+bash scripts/build_linux_wheel.sh "3.11" "2.7" "12.8"
+bash scripts/build_linux_wheel.sh "3.12" "2.7" "12.8"
 
 bash scripts/build_linux_wheel_cu128.sh "3.10" "2.8" "12.8"
 bash scripts/build_linux_wheel_cu128.sh "3.11" "2.8" "12.8"
 bash scripts/build_linux_wheel_cu128.sh "3.12" "2.8" "12.8"
-bash scripts/build_linux_wheel_cu128.sh "3.13" "2.8" "12.8"

scripts/build_all_windows_wheels.cmd

@@ -2,7 +2,7 @@
 setlocal enabledelayedexpansion
 
 REM Define Python and Torch versions
-set "python_versions=3.10 3.11 3.12 3.13"
+set "python_versions=3.10 3.11 3.12"
 set "torch_versions=2.5 2.6"
 set "cuda_version=12.4"
 
@@ -23,12 +23,10 @@ for %%P in (%python_versions%) do (
 call scripts\build_windows_wheel_cu128.cmd 3.10 2.7 12.8
 call scripts\build_windows_wheel_cu128.cmd 3.11 2.7 12.8
 call scripts\build_windows_wheel_cu128.cmd 3.12 2.7 12.8
-call scripts\build_windows_wheel_cu128.cmd 3.13 2.7 12.8
 
-call scripts\build_windows_wheel_cu128.cmd 3.10 2.8 12.8
-call scripts\build_windows_wheel_cu128.cmd 3.11 2.8 12.8
-call scripts\build_windows_wheel_cu128.cmd 3.12 2.8 12.8
-call scripts\build_windows_wheel_cu128.cmd 3.13 2.8 12.8
+REM call scripts\build_windows_wheel_cu128.cmd 3.10 2.8 12.8
+REM call scripts\build_windows_wheel_cu128.cmd 3.11 2.8 12.8
+REM call scripts\build_windows_wheel_cu128.cmd 3.12 2.8 12.8
 
 echo All builds completed successfully!
 exit /b 0

scripts/build_linux_wheel.sh

@@ -16,8 +16,12 @@ elif [ "$TORCH_VERSION" == "2.6" ]; then
   TORCHVISION_VERSION="0.21"
   TORCHAUDIO_VERSION="2.6"
   echo "TORCH_VERSION is 2.6, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
+elif [ "$TORCH_VERSION" == "2.7" ]; then
+  TORCHVISION_VERSION="0.22"
+  TORCHAUDIO_VERSION="2.7"
+  echo "TORCH_VERSION is 2.7, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
 else
-  echo "TORCH_VERSION is not 2.5 or 2.6, no changes to versions."
+  echo "TORCH_VERSION is not 2.5, 2.6 or 2.7, no changes to versions."
 fi
 
 docker run --rm \

scripts/build_linux_wheel_torch2.7_cu128.sh

-#!/bin/bash
-# Modified from https://github.com/sgl-project/sglang/blob/main/sgl-kernel/build.sh
-set -ex
-PYTHON_VERSION=$1
-TORCH_VERSION=$2 # has no use for now
-CUDA_VERSION=$3
-MAX_JOBS=${4:-} # optional
-PYTHON_ROOT_PATH=/opt/python/cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}
-
-# Check if TORCH_VERSION is 2.5 or 2.6 and set the corresponding versions for TORCHVISION and TORCHAUDIO
-#if [ "$TORCH_VERSION" == "2.5" ]; then
-#  TORCHVISION_VERSION="0.20"
-#  TORCHAUDIO_VERSION="2.5"
-#  echo "TORCH_VERSION is 2.5, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
-#elif [ "$TORCH_VERSION" == "2.6" ]; then
-#  TORCHVISION_VERSION="0.21"
-#  TORCHAUDIO_VERSION="2.6"
-#  echo "TORCH_VERSION is 2.6, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
-#else
-#  echo "TORCH_VERSION is not 2.5 or 2.6, no changes to versions."
-#fi
-
-docker run --rm \
-    -v "$(pwd)":/nunchaku \
-    pytorch/manylinux2_28-builder:cuda${CUDA_VERSION} \
-    bash -c "
-    cd /nunchaku && \
-    rm -rf build && \
-    gcc --version && g++ --version && \
-    ${PYTHON_ROOT_PATH}/bin/pip install --pre torch==2.7.0.dev20250307+cu128 torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128 && \
-    ${PYTHON_ROOT_PATH}/bin/pip install build ninja wheel setuptools && \
-    export NUNCHAKU_INSTALL_MODE=ALL && \
-    export NUNCHAKU_BUILD_WHEELS=1 && \
-    export MAX_JOBS=${MAX_JOBS} && \
-    ${PYTHON_ROOT_PATH}/bin/python -m build --wheel --no-isolation
-    "

scripts/build_windows_wheel.cmd

@@ -16,8 +16,11 @@ if "%TORCH_VERSION%"=="2.5" (
 ) else if "%TORCH_VERSION%"=="2.6" (
     set TORCHVISION_VERSION=0.21
     set TORCHAUDIO_VERSION=2.6
+) else if "%TORCH_VERSION%"=="2.7" (
+    set TORCHVISION_VERSION=0.22
+    set TORCHAUDIO_VERSION=2.7
 ) else (
-    echo TORCH_VERSION is not 2.5 or 2.6, no changes to versions.
+    echo TORCH_VERSION is not 2.5, 2.6 or 2.7, no changes to versions.
 )
 echo setting TORCHVISION_VERSION to %TORCHVISION_VERSION% and TORCHAUDIO_VERSION to %TORCHAUDIO_VERSION%
 
@@ -44,6 +47,7 @@ if exist build rd /s /q build
 :: set up Visual Studio compilation environment
 call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
 set DISTUTILS_USE_SDK=1
+set MAX_JOBS=4
 
 :: build wheels
 python -m build --wheel --no-isolation

scripts/build_windows_wheel_cu128.cmd → scripts/build_windows_wheel_torch_nightly.cmd

@@ -20,11 +20,7 @@ call conda activate %ENV_NAME%
 :: install dependencies
 call pip install ninja setuptools wheel build
 
-if "%TORCH_VERSION%"=="2.7" (
-    call pip install --pre torch==2.7.0.dev20250307+cu128 torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
-) else (
-    call pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
-)
+call pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
 
 :: set environment variables
 set NUNCHAKU_INSTALL_MODE=ALL
@@ -37,6 +33,7 @@ if exist build rd /s /q build
 :: set up Visual Studio compilation environment
 call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
 set DISTUTILS_USE_SDK=1
+set MAX_JOBS=4
 
 :: build wheels
 python -m build --wheel --no-isolation

setup.py

@@ -2,6 +2,7 @@ import os
 import re
 import subprocess
 import sys
+from datetime import date
 
 import setuptools
 import torch
@@ -64,6 +65,8 @@ if __name__ == "__main__":
 
     torch_version = torch.__version__.split("+")[0]
     torch_major_minor_version = ".".join(torch_version.split(".")[:2])
+    if "dev" in version:
+        version = version + date.today().strftime("%Y%m%d")  # data
     version = version + "+torch" + torch_major_minor_version
 
     ROOT_DIR = os.path.dirname(__file__)