Release v0.2.0

Ready to release v0.2.0

scripts/build_docker.sh

+#!/bin/bash
+PYTHON_VERSION=$1
+TORCH_VERSION=$2
+CUDA_VERSION=$3
+NUNCHAKU_VERSION=$4
+
+# 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. Exit."
+  exit 2
+fi
+
+if [ "$CUDA_VERSION" == "12.8" ]; then
+  CUDA_IMAGE="12.8.1-devel-ubuntu24.04"
+  echo "CUDA_VERSION is 12.8, setting CUDA_IMAGE to $CUDA_IMAGE"
+elif [ "$CUDA_VERSION" == "12.4" ]; then
+  CUDA_IMAGE="12.4.1-devel-ubuntu22.04"
+  echo "CUDA_VERSION is 12.4, setting CUDA_IMAGE to $CUDA_IMAGE"
+else
+  echo "CUDA_VERSION is not 12.8 or 12.4. Exit."
+  exit 2
+fi
+
+
+docker build --no-cache \
+--build-arg PYTHON_VERSION=${PYTHON_VERSION} \
+--build-arg CUDA_SHORT_VERSION=${CUDA_VERSION//.} \
+--build-arg CUDA_IMAGE=${CUDA_IMAGE} \
+--build-arg TORCH_VERSION=${TORCH_VERSION} \
+--build-arg TORCHVISION_VERSION=${TORCHVISION_VERSION} \
+--build-arg TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION} \
+-t nunchaku:${NUNCHAKU_VERSION}-py${PYTHON_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION} .

scripts/build_linux_wheels.sh → scripts/build_linux_wheel.sh

@@ -7,6 +7,19 @@ 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/manylinux-builder:cuda${CUDA_VERSION} \
@@ -16,7 +29,7 @@ docker run --rm \
     yum install -y devtoolset-11 && \
     source scl_source enable devtoolset-11 && \
     gcc --version && g++ --version && \
-    ${PYTHON_ROOT_PATH}/bin/pip install --no-cache-dir torch==${TORCH_VERSION} numpy --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION//.} && \
+    ${PYTHON_ROOT_PATH}/bin/pip install --no-cache-dir torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} torchaudio==${TORCHAUDIO_VERSION} --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION//.} && \
     ${PYTHON_ROOT_PATH}/bin/pip install build ninja wheel setuptools && \
     export NUNCHAKU_INSTALL_MODE=ALL && \
     export NUNCHAKU_BUILD_WHEELS=1 && \

scripts/build_linux_wheel_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 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
+    "
\ No newline at end of file

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
+    "
\ No newline at end of file

scripts/build_windows_wheel.cmd

+@echo off
+setlocal enabledelayedexpansion
+
+:: get arguments
+set PYTHON_VERSION=%1
+set TORCH_VERSION=%2
+set CUDA_VERSION=%3
+
+set CUDA_SHORT_VERSION=%CUDA_VERSION:.=%
+echo %CUDA_SHORT_VERSION%
+
+:: setup some variables
+if "%TORCH_VERSION%"=="2.5" (
+    set TORCHVISION_VERSION=0.20
+    set TORCHAUDIO_VERSION=2.5
+) else if "%TORCH_VERSION%"=="2.6" (
+    set TORCHVISION_VERSION=0.21
+    set TORCHAUDIO_VERSION=2.6
+) else (
+    echo TORCH_VERSION is not 2.5 or 2.6, no changes to versions.
+)
+echo setting TORCHVISION_VERSION to %TORCHVISION_VERSION% and TORCHAUDIO_VERSION to %TORCHAUDIO_VERSION%
+
+:: conda environment name
+set ENV_NAME=build_env_%PYTHON_VERSION%_%TORCH_VERSION%
+echo Using conda environment: %ENV_NAME%
+
+:: create conda environment
+call conda create -y -n %ENV_NAME% python=%PYTHON_VERSION%
+call conda activate %ENV_NAME%
+
+:: install dependencies
+call pip install ninja setuptools wheel build
+call pip install --no-cache-dir torch==%TORCH_VERSION% torchvision==%TORCHVISION_VERSION% torchaudio==%TORCHAUDIO_VERSION% --index-url "https://download.pytorch.org/whl/cu%CUDA_SHORT_VERSION%/"
+
+:: set environment variables
+set NUNCHAKU_INSTALL_MODE=ALL
+set NUNCHAKU_BUILD_WHEELS=1
+
+:: cd to the parent directory
+cd /d "%~dp0.."
+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
+
+:: build wheels
+python -m build --wheel --no-isolation
+
+:: exit conda
+call conda deactivate
+call conda remove -y -n %ENV_NAME% --all
+
+echo Build complete!

scripts/build_windows_wheel_cu128.cmd

+@echo off
+setlocal enabledelayedexpansion
+
+:: get arguments
+set PYTHON_VERSION=%1
+set TORCH_VERSION=%2
+set CUDA_VERSION=%3
+
+set CUDA_SHORT_VERSION=%CUDA_VERSION:.=%
+echo %CUDA_SHORT_VERSION%
+
+:: conda environment name
+set ENV_NAME=build_env_%PYTHON_VERSION%_%TORCH_VERSION%
+echo Using conda environment: %ENV_NAME%
+
+:: create conda environment
+call conda create -y -n %ENV_NAME% python=%PYTHON_VERSION%
+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
+)
+
+:: set environment variables
+set NUNCHAKU_INSTALL_MODE=ALL
+set NUNCHAKU_BUILD_WHEELS=1
+
+:: cd to the parent directory
+cd /d "%~dp0.."
+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
+
+:: build wheels
+python -m build --wheel --no-isolation
+
+:: exit conda
+call conda deactivate
+call conda remove -y -n %ENV_NAME% --all
+
+echo Build complete!

scripts/linux_cleanup.sh

 #!/bin/bash
 set -ex
 
-#docker run --rm \
-#    -v "$(pwd)":/nunchaku \
-#    pytorch/manylinux-builder:cuda12.4 \
-#    bash -c "cd /nunchaku && rm -r *"
-docker run --rm -it \
+docker run --rm \
     -v "$(pwd)":/nunchaku \
     pytorch/manylinux-builder:cuda12.4 \
-    bash
\ No newline at end of file
+    bash -c "cd /nunchaku && rm -rf *"
\ No newline at end of file

setup.py

@@ -47,12 +47,12 @@ def get_sm_targets() -> list[str]:
             sm = f"{capability[0]}{capability[1]}"
             if sm == "120" and support_sm120:
                 sm = "120a"
-            assert sm in ["80", "86", "89", "120a"], f"Unsupported SM {sm}"
+            assert sm in ["75", "80", "86", "89", "120a"], f"Unsupported SM {sm}"
             if sm not in ret:
                 ret.append(sm)
     else:
         assert install_mode == "ALL"
-        ret = ["80", "86", "89"]
+        ret = ["75", "80", "86", "89"]
         if support_sm120:
             ret.append("120a")
     return ret
@@ -142,6 +142,7 @@ if __name__ == "__main__":
             *ncond("src/FluxModel.cpp"),
             *ncond("src/SanaModel.cpp"),
             "src/Serialization.cpp",
+            "src/Module.cpp",
             *ncond("third_party/Block-Sparse-Attention/csrc/block_sparse_attn/src/flash_fwd_hdim64_fp16_sm80.cu"),
             *ncond("third_party/Block-Sparse-Attention/csrc/block_sparse_attn/src/flash_fwd_hdim64_bf16_sm80.cu"),
             *ncond("third_party/Block-Sparse-Attention/csrc/block_sparse_attn/src/flash_fwd_hdim128_fp16_sm80.cu"),
@@ -158,9 +159,14 @@ if __name__ == "__main__":
             "src/kernels/layernorm_kernels.cu",
             "src/kernels/misc_kernels.cu",
             "src/kernels/zgemm/gemm_w4a4.cu",
-            "src/kernels/zgemm/gemm_w4a4_launch_fp16.cu",
-            "src/kernels/zgemm/gemm_w4a4_launch_bf16.cu",
+            "src/kernels/zgemm/gemm_w4a4_test.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_fp16_int4.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_fp16_int4_fasteri2f.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_fp16_fp4.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_bf16_int4.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_bf16_fp4.cu",
             "src/kernels/zgemm/gemm_w8a8.cu",
+            "src/kernels/zgemm/attention.cu",
             "src/kernels/dwconv.cu",
             "src/kernels/gemm_batched.cu",
             "src/kernels/gemm_f16.cu",

src/FluxModel.cpp

 #include "FluxModel.h"
 #include "kernels/misc_kernels.h"
 #include "kernels/gemm_batched.h"
+#include "kernels/zgemm/zgemm.h"
 #include "flash_api.h"
 #include "activation.h"
 
@@ -117,6 +118,33 @@ Attention::Attention(int num_heads, int dim_head, Device device) :
     headmask_type = headmask_type.copy(device);
 }
 
+Tensor Attention::forward(Tensor qkv) {
+    assert(qkv.ndims() == 3);
+
+    const Device device = qkv.device();
+    const int batch_size = qkv.shape[0];
+    const int num_tokens = qkv.shape[1];
+    assert(qkv.shape[2] == num_heads * dim_head * 3);
+
+    Tensor reshaped = qkv.view({batch_size, num_tokens, num_heads * 3, dim_head});
+    Tensor q = reshaped.slice(2, 0, num_heads);
+    Tensor k = reshaped.slice(2, num_heads, num_heads * 2);
+    Tensor v = reshaped.slice(2, num_heads * 2, num_heads * 3);
+
+    Tensor raw_attn_output = mha_fwd(q, k, v,
+        0.0f,
+        pow(q.shape[-1], (-0.5)),
+        false, -1, -1, false
+    ).front();
+
+    assert(raw_attn_output.shape[0] == batch_size);
+    assert(raw_attn_output.shape[1] == num_tokens);
+    assert(raw_attn_output.shape[2] == num_heads);
+    assert(raw_attn_output.shape[3] == dim_head);
+    
+    return raw_attn_output.view({batch_size * num_tokens, num_heads, dim_head});
+}
+
 Tensor Attention::forward(Tensor qkv, Tensor pool_qkv, float sparsityRatio) {
     const bool cast_fp16 = this->force_fp16 && qkv.scalar_type() != Tensor::FP16;
 
@@ -235,7 +263,7 @@ Tensor Attention::forward(Tensor qkv, Tensor pool_qkv, float sparsityRatio) {
     Tensor raw_attn_output = mha_varlen_fwd(
         q, k, v,
         cu_seqlens, cu_seqlens,
-        num_tokens_img + num_tokens_context, num_tokens_img + num_tokens_context,
+        num_tokens_img + num_tokens_txt, num_tokens_img + num_tokens_txt,
         0.0f,
         pow(q.shape[-1], (-0.5)),
         false, false, -1, -1, false
@@ -298,19 +326,50 @@ Tensor FluxSingleTransformerBlock::forward(Tensor hidden_states, Tensor temb, Te
 
     Tensor residual = hidden_states;
 
+    Tensor attn_output;
+
+    debug("rotary_emb", rotary_emb);
+
+    if (attnImpl == AttentionImpl::FlashAttention2) {
         Tensor qkv = Tensor::allocate({batch_size, num_tokens, dim * 3}, norm_hidden_states.scalar_type(), norm_hidden_states.device());
         // qkv_proj.forward(norm_hidden_states, qkv, {});
         // debug("qkv_raw", qkv);
 
-    debug("rotary_emb", rotary_emb);
         qkv_proj.forward(norm_hidden_states, qkv, {}, norm_q.weight, norm_k.weight, rotary_emb);
         debug("qkv", qkv);
         // Tensor qkv = forward_fc(qkv_proj, norm_hidden_states);
 
-    Tensor attn_output = attn.forward(qkv, {}, 0);
+        // attn_output = attn.forward(qkv, {}, 0);
+        attn_output = attn.forward(qkv);
         attn_output = attn_output.reshape({batch_size, num_tokens, num_heads * dim_head});
+    } else if (attnImpl == AttentionImpl::NunchakuFP16) {
+        assert(batch_size == 1);
+
+        const int num_tokens_pad = ceilDiv(num_tokens, 256) * 256;
+
+        Tensor q = Tensor::allocate({batch_size, num_heads, num_tokens_pad, dim_head}, Tensor::FP16, norm_hidden_states.device());
+        Tensor k = Tensor::allocate({batch_size, num_heads, num_tokens_pad, dim_head}, Tensor::FP16, norm_hidden_states.device());
+        Tensor v = Tensor::allocate({batch_size, num_heads, num_tokens_pad, dim_head}, Tensor::FP16, norm_hidden_states.device());
+
+        qkv_proj.forward(norm_hidden_states, {}, {}, norm_q.weight, norm_k.weight, rotary_emb, q, k, v, num_tokens);
+
+        debug("packed_q", q);
+        debug("packed_k", k);
+        debug("packed_v", v);
+
+        Tensor o = Tensor::allocate({batch_size, num_tokens_pad, num_heads * dim_head}, norm_hidden_states.scalar_type(), norm_hidden_states.device());
+
+        kernels::attention_fp16(q, k, v, o, pow(dim_head, (-0.5)));
+
+        attn_output = o.slice(1, 0, num_tokens);
+    } else {
+        assert(false);
+    }
+
     debug("raw_attn_output", attn_output);
 
+
+
     attn_output = forward_fc(out_proj, attn_output);
     debug("attn_output", attn_output);
 
@@ -384,13 +443,13 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
 
 
     int num_tokens_img = hidden_states.shape[1];
-    int num_tokens_context = encoder_hidden_states.shape[1];
+    int num_tokens_txt = encoder_hidden_states.shape[1];
 
     assert(hidden_states.shape[2] == dim);
     assert(encoder_hidden_states.shape[2] == dim);
 
     spdlog::debug("hidden_states={} encoder_hidden_states={} temb={}", hidden_states.shape.str(), encoder_hidden_states.shape.str(), temb.shape.str());
-    spdlog::debug("batch_size={} num_tokens_img={} num_tokens_context={}", batch_size, num_tokens_img, num_tokens_context);
+    spdlog::debug("batch_size={} num_tokens_img={} num_tokens_txt={}", batch_size, num_tokens_img, num_tokens_txt);
 
     auto norm1_output = norm1.forward(hidden_states, temb);
     auto norm1_context_output = norm1_context.forward(encoder_hidden_states, temb);
@@ -408,6 +467,14 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
     nvtxRangePop();
 
     auto stream = getCurrentCUDAStream();
+
+    int num_tokens_img_pad = 0, num_tokens_txt_pad = 0;
+    Tensor raw_attn_output;
+
+    if (attnImpl == AttentionImpl::FlashAttention2) {
+        num_tokens_img_pad = num_tokens_img;
+        num_tokens_txt_pad = num_tokens_txt;
+
         Tensor concat;
         Tensor pool;
 
@@ -416,8 +483,8 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
 
             const bool blockSparse = sparsityRatio > 0;
 
-        const int poolTokens = num_tokens_img / POOL_SIZE + num_tokens_context / POOL_SIZE;
-        concat = Tensor::allocate({batch_size, num_tokens_img + num_tokens_context, dim * 3}, norm1_output.x.scalar_type(), norm1_output.x.device());
+            const int poolTokens = num_tokens_img / POOL_SIZE + num_tokens_txt / POOL_SIZE;
+            concat = Tensor::allocate({batch_size, num_tokens_img + num_tokens_txt, dim * 3}, norm1_output.x.scalar_type(), norm1_output.x.device());
 
             pool = blockSparse
                 ? Tensor::allocate({batch_size, poolTokens, dim * 3}, norm1_output.x.scalar_type(), norm1_output.x.device())
@@ -426,13 +493,13 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
             for (int i = 0; i < batch_size; i++) {
                 // img first
                 Tensor qkv = concat.slice(0, i, i + 1).slice(1, 0, num_tokens_img);
-            Tensor qkv_context = concat.slice(0, i, i + 1).slice(1, num_tokens_img, num_tokens_img + num_tokens_context);
+                Tensor qkv_context = concat.slice(0, i, i + 1).slice(1, num_tokens_img, num_tokens_img + num_tokens_txt);
 
                 Tensor pool_qkv = pool.valid()
                     ? pool.slice(0, i, i + 1).slice(1, 0, num_tokens_img / POOL_SIZE)
                     : Tensor{};
                 Tensor pool_qkv_context = pool.valid()
-                ? concat.slice(0, i, i + 1).slice(1, num_tokens_img / POOL_SIZE, num_tokens_img / POOL_SIZE + num_tokens_context / POOL_SIZE)
+                    ? concat.slice(0, i, i + 1).slice(1, num_tokens_img / POOL_SIZE, num_tokens_img / POOL_SIZE + num_tokens_txt / POOL_SIZE)
                     : Tensor{};
 
                 // qkv_proj.forward(norm1_output.x.slice(0, i, i + 1), qkv);
@@ -462,22 +529,79 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
 
         nvtxRangePushA("Attention");
 
-    Tensor raw_attn_output = attn.forward(concat, pool, sparsityRatio);
+        if (pool.valid()) {
+            raw_attn_output = attn.forward(concat, pool, sparsityRatio);
+        } else {
+            raw_attn_output = attn.forward(concat);
+        }
 
         nvtxRangePop();
 
         spdlog::debug("raw_attn_output={}", raw_attn_output.shape.str());
 
-    raw_attn_output = raw_attn_output.view({batch_size, num_tokens_img + num_tokens_context, num_heads, dim_head});
-    debug("raw_attn_output", raw_attn_output);
+        raw_attn_output = raw_attn_output.view({batch_size, num_tokens_img + num_tokens_txt, num_heads, dim_head});
 
+    } else if (attnImpl == AttentionImpl::NunchakuFP16) {
+        num_tokens_img_pad = ceilDiv(num_tokens_img, 256) * 256;
+        num_tokens_txt_pad = ceilDiv(num_tokens_txt, 256) * 256;
+
+        Tensor concat_q, concat_k, concat_v;
+
+        {
+            nvtxRangePushA("qkv_proj");
+
+            concat_q = Tensor::allocate({batch_size, num_heads, num_tokens_img_pad + num_tokens_txt_pad, dim_head}, Tensor::FP16, norm1_output.x.device());
+            concat_k = Tensor::empty_like(concat_q);
+            concat_v = Tensor::empty_like(concat_q);
+
+            for (int i = 0; i < batch_size; i++) {
+                // img first
+                auto sliceImg = [&](Tensor x) {
+                    return x.slice(0, i, i+1).slice(2, 0, num_tokens_img_pad);
+                };
+                auto sliceTxt = [&](Tensor x) {
+                    return x.slice(0, i, i+1).slice(2, num_tokens_img_pad, num_tokens_img_pad + num_tokens_txt_pad);
+                };
+
+                qkv_proj.forward(
+                    norm1_output.x.slice(0, i, i + 1), {}, {}, norm_q.weight, norm_k.weight, rotary_emb,
+                    sliceImg(concat_q), sliceImg(concat_k), sliceImg(concat_v), num_tokens_img
+                );
+
+                qkv_proj_context.forward(
+                    norm1_context_output.x.slice(0, i, i + 1), {}, {}, norm_added_q.weight, norm_added_k.weight, rotary_emb_context,
+                    sliceTxt(concat_q), sliceTxt(concat_k), sliceTxt(concat_v), num_tokens_txt
+                );
+            }
+
+            debug("concat_q", concat_q);
+            debug("concat_k", concat_k);
+            debug("concat_v", concat_v);
+
+            nvtxRangePop();
+        }
+
+        raw_attn_output = Tensor::allocate({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads * dim_head}, norm1_output.x.scalar_type(), norm1_output.x.device());
+
+        nvtxRangePushA("Attention");
+
+        kernels::attention_fp16(concat_q, concat_k, concat_v, raw_attn_output, pow(dim_head, (-0.5)));
+
+        nvtxRangePop();
+
+        raw_attn_output = raw_attn_output.view({batch_size, num_tokens_img_pad + num_tokens_txt_pad, num_heads, dim_head});
+    } else {
+        assert(false);
+    }
+
+    debug("raw_attn_output", raw_attn_output);
 
     {
         nvtxRangePushA("o_proj");
 
         auto &&[_, gate_msa, shift_mlp, scale_mlp, gate_mlp] = norm1_output;
 
-        // raw_attn_output: [batch_size, num_tokens_img + num_tokens_context, num_heads * dim_head]
+        // raw_attn_output: [batch_size, num_tokens_img + num_tokens_txt, num_heads * dim_head]
 
         Tensor raw_attn_output_split;
         if (batch_size == 1) {
@@ -488,7 +612,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
                 raw_attn_output_split.data_ptr(),
                 num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(),
                 raw_attn_output.data_ptr(),
-                (num_tokens_img + num_tokens_context) * num_heads * dim_head * raw_attn_output.scalar_size(),
+                (num_tokens_img_pad + num_tokens_txt_pad) * num_heads * dim_head * raw_attn_output.scalar_size(),
                 num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(),
                 batch_size,
                 cudaMemcpyDeviceToDevice,
@@ -546,15 +670,15 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
 
         Tensor raw_attn_output_split;
         if (batch_size == 1) {
-            raw_attn_output_split = raw_attn_output.slice(1, num_tokens_img, num_tokens_img + num_tokens_context).reshape({batch_size, num_tokens_context, num_heads * dim_head});
+            raw_attn_output_split = raw_attn_output.slice(1, num_tokens_img_pad, num_tokens_img_pad + num_tokens_txt).reshape({batch_size, num_tokens_txt, num_heads * dim_head});
         } else {
-            raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_context, num_heads * dim_head}, raw_attn_output.scalar_type(), raw_attn_output.device());
+            raw_attn_output_split = Tensor::allocate({batch_size, num_tokens_txt, num_heads * dim_head}, raw_attn_output.scalar_type(), raw_attn_output.device());
             checkCUDA(cudaMemcpy2DAsync(
                 raw_attn_output_split.data_ptr(),
-                num_tokens_context * num_heads * dim_head * raw_attn_output_split.scalar_size(),
-                raw_attn_output.data_ptr<char>() + num_tokens_img * num_heads * dim_head * raw_attn_output_split.scalar_size(),
-                (num_tokens_img + num_tokens_context) * num_heads * dim_head * raw_attn_output.scalar_size(),
-                num_tokens_context * num_heads * dim_head * raw_attn_output_split.scalar_size(),
+                num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(),
+                raw_attn_output.data_ptr<char>() + num_tokens_img_pad * num_heads * dim_head * raw_attn_output_split.scalar_size(),
+                (num_tokens_img_pad + num_tokens_txt_pad) * num_heads * dim_head * raw_attn_output.scalar_size(),
+                num_tokens_txt * num_heads * dim_head * raw_attn_output_split.scalar_size(),
                 batch_size,
                 cudaMemcpyDeviceToDevice,
                 stream));
@@ -607,7 +731,7 @@ std::tuple<Tensor, Tensor> JointTransformerBlock::forward(Tensor hidden_states,
     return { hidden_states, encoder_hidden_states };
 }
 
-FluxModel::FluxModel(bool use_fp4, bool offload, Tensor::ScalarType dtype, Device device) : offload(offload) {
+FluxModel::FluxModel(bool use_fp4, bool offload, Tensor::ScalarType dtype, Device device) : dtype(dtype), offload(offload) {
     for (int i = 0; i < 19; i++) {
         transformer_blocks.push_back(std::make_unique<JointTransformerBlock>(3072, 24, 3072, false, use_fp4, dtype, device));
         registerChildren(*transformer_blocks.back(), format("transformer_blocks.{}", i));
@@ -626,7 +750,16 @@ FluxModel::FluxModel(bool use_fp4, bool offload, Tensor::ScalarType dtype, Devic
     }
 }
 
-Tensor FluxModel::forward(Tensor hidden_states, Tensor encoder_hidden_states, Tensor temb, Tensor rotary_emb_img, Tensor rotary_emb_context, Tensor rotary_emb_single) {
+Tensor FluxModel::forward(
+        Tensor hidden_states,
+        Tensor encoder_hidden_states,
+        Tensor temb,
+        Tensor rotary_emb_img,
+        Tensor rotary_emb_context,
+        Tensor rotary_emb_single,
+        Tensor controlnet_block_samples,
+        Tensor controlnet_single_block_samples,
+        bool skip_first_layer) {
     const int batch_size = hidden_states.shape[0];
     const Tensor::ScalarType dtype = hidden_states.dtype();
     const Device device = hidden_states.device();
@@ -639,9 +772,20 @@ Tensor FluxModel::forward(Tensor hidden_states, Tensor encoder_hidden_states, Te
     Tensor concat;
 
     auto compute = [&](int layer) {
+        if (skip_first_layer && size_t(layer) == 0) return;
         if (size_t(layer) < transformer_blocks.size()) {
             auto &block = transformer_blocks.at(layer);
             std::tie(hidden_states, encoder_hidden_states) = block->forward(hidden_states, encoder_hidden_states, temb, rotary_emb_img, rotary_emb_context, 0.0f);
+            if (controlnet_block_samples.valid()) {
+                const int num_controlnet_block_samples = controlnet_block_samples.shape[0];
+
+                int interval_control = ceilDiv(transformer_blocks.size(), static_cast<size_t>(num_controlnet_block_samples));
+                int block_index = layer / interval_control;
+                // Xlabs ControlNet
+                // block_index = layer % num_controlnet_block_samples;
+
+                hidden_states = kernels::add(hidden_states, controlnet_block_samples[block_index]);
+            }
         } else {
             if (size_t(layer) == transformer_blocks.size()) {
                 // txt first, same as diffusers
@@ -652,10 +796,23 @@ Tensor FluxModel::forward(Tensor hidden_states, Tensor encoder_hidden_states, Te
                 }
                 hidden_states = concat;
                 encoder_hidden_states = {};
+
             }
 
             auto &block = single_transformer_blocks.at(layer - transformer_blocks.size());
             hidden_states = block->forward(hidden_states, temb, rotary_emb_single);
+            if (controlnet_single_block_samples.valid()) {
+                const int num_controlnet_single_block_samples = controlnet_single_block_samples.shape[0];
+
+                int interval_control = ceilDiv(single_transformer_blocks.size(), static_cast<size_t>(num_controlnet_single_block_samples));
+                int block_index = (layer - transformer_blocks.size()) / interval_control;
+                // Xlabs ControlNet
+                // block_index = layer % num_controlnet_single_block_samples
+
+                auto slice = hidden_states.slice(1, txt_tokens, txt_tokens + img_tokens);
+                slice = kernels::add(slice, controlnet_single_block_samples[block_index]);
+                hidden_states.slice(1, txt_tokens, txt_tokens + img_tokens).copy_(slice);
+            }
         }
     };
     auto load = [&](int layer) {
@@ -682,3 +839,57 @@ Tensor FluxModel::forward(Tensor hidden_states, Tensor encoder_hidden_states, Te
 
     return hidden_states;
 }
+
+std::tuple<Tensor, Tensor> FluxModel::forward_layer(
+        size_t layer,
+        Tensor hidden_states,
+        Tensor encoder_hidden_states,
+        Tensor temb,
+        Tensor rotary_emb_img,
+        Tensor rotary_emb_context,
+        Tensor controlnet_block_samples,
+        Tensor controlnet_single_block_samples) {
+
+    std::tie(hidden_states, encoder_hidden_states) = transformer_blocks.at(layer)->forward(
+        hidden_states,
+        encoder_hidden_states,
+        temb,
+        rotary_emb_img,
+        rotary_emb_context, 0.0f);
+
+    const int txt_tokens = encoder_hidden_states.shape[1];
+    const int img_tokens = hidden_states.shape[1];
+
+    if (layer < transformer_blocks.size() && controlnet_block_samples.valid()) {
+        const int num_controlnet_block_samples = controlnet_block_samples.shape[0];
+
+        int interval_control = ceilDiv(transformer_blocks.size(), static_cast<size_t>(num_controlnet_block_samples));
+        int block_index = layer / interval_control;
+        // Xlabs ControlNet
+        // block_index = layer % num_controlnet_block_samples;
+
+        hidden_states = kernels::add(hidden_states, controlnet_block_samples[block_index]);
+    } else if (layer >= transformer_blocks.size() && controlnet_single_block_samples.valid()) {
+        const int num_controlnet_single_block_samples = controlnet_single_block_samples.shape[0];
+
+        int interval_control = ceilDiv(single_transformer_blocks.size(), static_cast<size_t>(num_controlnet_single_block_samples));
+        int block_index = (layer - transformer_blocks.size()) / interval_control;
+        // Xlabs ControlNet
+        // block_index = layer % num_controlnet_single_block_samples
+
+        auto slice = hidden_states.slice(1, txt_tokens, txt_tokens + img_tokens);
+        slice = kernels::add(slice, controlnet_single_block_samples[block_index]);
+        hidden_states.slice(1, txt_tokens, txt_tokens + img_tokens).copy_(slice);
+    }
+
+    return { hidden_states, encoder_hidden_states };
+}
+
+void FluxModel::setAttentionImpl(AttentionImpl impl) {
+    for (auto &&block : this->transformer_blocks) {
+        block->attnImpl = impl;
+    }
+    for (auto &&block : this->single_transformer_blocks) {
+        block->attnImpl = impl;
+    }
+}

src/FluxModel.h

@@ -6,6 +6,11 @@
 #include "Linear.h"
 #include "layernorm.h"
 
+enum class AttentionImpl {
+    FlashAttention2 = 0,
+    NunchakuFP16,
+};
+
 class AdaLayerNormZeroSingle : public Module {
 public:
     static constexpr bool USE_4BIT = true;
@@ -58,6 +63,7 @@ public:
     static constexpr int POOL_SIZE = 128;
 
     Attention(int num_heads, int dim_head, Device device);
+    Tensor forward(Tensor qkv);
     Tensor forward(Tensor qkv, Tensor pool_qkv, float sparsityRatio);
 
     static void setForceFP16(Module *module, bool value);
@@ -86,6 +92,8 @@ public:
     const int num_heads;
     const int mlp_hidden_dim;
 
+    AttentionImpl attnImpl = AttentionImpl::FlashAttention2;
+
 private:
     AdaLayerNormZeroSingle norm;
     GEMM mlp_fc1;
@@ -110,6 +118,8 @@ public:
     const int num_heads;
     const bool context_pre_only;
 
+    AttentionImpl attnImpl = AttentionImpl::FlashAttention2;
+
 private:
     AdaLayerNormZero norm1;
     AdaLayerNormZero norm1_context;
@@ -129,9 +139,30 @@ private:
 class FluxModel : public Module {
 public:
     FluxModel(bool use_fp4, bool offload, Tensor::ScalarType dtype, Device device);
-    Tensor forward(Tensor hidden_states, Tensor encoder_hidden_states, Tensor temb, Tensor rotary_emb_img, Tensor rotary_emb_context, Tensor rotary_emb_single);
+    Tensor forward(
+        Tensor hidden_states,
+        Tensor encoder_hidden_states,
+        Tensor temb,
+        Tensor rotary_emb_img,
+        Tensor rotary_emb_context,
+        Tensor rotary_emb_single,
+        Tensor controlnet_block_samples,
+        Tensor controlnet_single_block_samples,
+        bool skip_first_layer = false);
+    std::tuple<Tensor, Tensor> forward_layer(
+        size_t layer,
+        Tensor hidden_states,
+        Tensor encoder_hidden_states,
+        Tensor temb,
+        Tensor rotary_emb_img,
+        Tensor rotary_emb_context,
+        Tensor controlnet_block_samples,
+        Tensor controlnet_single_block_samples);
+    void setAttentionImpl(AttentionImpl impl);
 
 public:
+    const Tensor::ScalarType dtype;
+
     std::vector<std::unique_ptr<JointTransformerBlock>> transformer_blocks;
     std::vector<std::unique_ptr<FluxSingleTransformerBlock>> single_transformer_blocks;
 

src/Linear.cpp

@@ -52,13 +52,14 @@ void GEMV_AWQ::loadParam(std::string key, Tensor &dst, Tensor src) {
     if (key == "lora_down" || key == "lora_up") {
         assert(src.ndims() == 2);
         if (dst.shape.dataExtent != src.shape.dataExtent) {
-            dst = src.copy(this->device);
+            dst = Tensor::allocate(src.shape.dataExtent, dst.scalar_type(), this->device);
+            Module::loadParam(key, dst, src);
             if (key == "lora_down") {
                 const int new_rank = dst.shape[0];
                 this->lora_rank = new_rank;
             }
         } else {
-            dst.copy_(src);
+            Module::loadParam(key, dst, src);
         }
     } else {
         Module::loadParam(key, dst, src);
@@ -143,16 +144,18 @@ void GEMM_W4A4::loadParam(std::string key, Tensor &dst, Tensor src) {
     if (key == "lora_down" || key == "lora_up") {
         assert(src.ndims() == 2);
         if (dst.shape.dataExtent != src.shape.dataExtent) {
-            dst = src.copy(this->device);
+            dst = Tensor::allocate(src.shape.dataExtent, dst.scalar_type(), this->device);
+            Module::loadParam(key, dst, src);
             this->lora_rank = dst.shape[1];
             this->lora_scales.resize(ceilDiv(this->lora_rank, 16), 1.0f);
         } else {
-            dst.copy_(src);
+            Module::loadParam(key, dst, src);
         }
     } else if (key == "wcscales") {
         assert(src.ndims() == 1);
         assert(src.shape[0] == out_features_pad);
-        dst = src.copy(this->device);
+        dst = Tensor::allocate(src.shape.dataExtent, dst.scalar_type(), this->device);
+        Module::loadParam(key, dst, src);
     } else if (key == "wtscale") {
         assert(src.numel() == 1);
         if (src.dtype() == Tensor::BF16) {
@@ -160,7 +163,7 @@ void GEMM_W4A4::loadParam(std::string key, Tensor &dst, Tensor src) {
         } else if (src.dtype() == Tensor::FP16) {
             *dst.data_ptr<float>() = float(*src.data_ptr<half>());
         } else if (src.dtype() == Tensor::FP32) {
-            dst.copy_(src);
+            Module::loadParam(key, dst, src);
         } else {
             assert(false);
         }
@@ -181,7 +184,7 @@ std::variant<Tensor, GEMM_W4A4::QuantizedActivation> GEMM_W4A4::forward(Tensor x
     return forward_quant(quantize(x, false), fuse, nextGEMM);
 }
 
-void GEMM_W4A4::forward(Tensor x, Tensor out, Tensor pool, Tensor norm_q, Tensor norm_k, Tensor rotary_emb) {
+void GEMM_W4A4::forward(Tensor x, Tensor out, Tensor pool, Tensor norm_q, Tensor norm_k, Tensor rotary_emb, Tensor out_q, Tensor out_k, Tensor out_v, int numTokens) {
     QuantizedActivation qact = quantize(x, false);
 
 #if !NO_LORA_FUSION
@@ -196,7 +199,8 @@ void GEMM_W4A4::forward(Tensor x, Tensor out, Tensor pool, Tensor norm_q, Tensor
 
     kernels::gemm_w4a4(
         qact.act, qweight, out, {}, qact.ascales, wscales, {}, pool, qact.lora_act, this->lora_up, {}, {}, norm_q, norm_k, rotary_emb, this->bias, {}, {}, {}, qact.is_unsigned, this->lora_scales, false,
-        use_fp4, *this->wtscale.data_ptr<float>(), wcscales.numel() > 0 ? wcscales: Tensor{}
+        use_fp4, *this->wtscale.data_ptr<float>(), wcscales.numel() > 0 ? wcscales: Tensor{},
+        out_q, out_k, out_v, numTokens
     );
 
     debug("gemm.out", out);
@@ -277,7 +281,8 @@ std::variant<Tensor, GEMM_W4A4::QuantizedActivation> GEMM_W4A4::forward_quant(Qu
 
     kernels::gemm_w4a4(
         qact.act, qweight, out, qout.act, qact.ascales, wscales, qout.ascales, {}, qact.lora_act, this->lora_up, next_lora, qout.lora_act, {}, {}, {}, this->bias, next_smooth, {}, {}, qact.is_unsigned, this->lora_scales, fuse == FuseOptions::SILU,
-        use_fp4, *this->wtscale.data_ptr<float>(), wcscales.numel() > 0 ? wcscales: Tensor{}
+        use_fp4, *this->wtscale.data_ptr<float>(), wcscales.numel() > 0 ? wcscales: Tensor{},
+        {}, {}, {}, 0
     );
 
     if (fuse == FuseOptions::EMPTY || fuse == FuseOptions::SILU) {
@@ -446,9 +451,9 @@ GEMM_W8A8::QuantizedActivation GEMM_W8A8::quantize(Tensor x, bool fuse_glu) {
 }
 
 Tensor GEMM_W8A8::forward_quant(QuantizedActivation qact) {
-    auto oshape = qact.act.shape;
-    oshape[-1] = out_features;
-    Tensor out = Tensor::allocate(oshape, this->dtype, qact.act.device());
+    auto shape = TensorShape(qact.act.shape.dataExtent);
+    shape[-1] = out_features;
+    Tensor out = Tensor::allocate(shape, this->dtype, qact.act.device());
     kernels::gemm_w8a8(qact.act, this->qweight, out, qact.ascales, this->wscales, this->bias);
 
     debug("gemm.out", out);

src/Linear.h

@@ -69,7 +69,11 @@ public:
     Tensor forward(Tensor x);
     Tensor forward_silu(Tensor x);
     std::variant<Tensor, QuantizedActivation> forward(Tensor x, FuseOptions fuse, GEMM_W4A4 *nextGEMM = nullptr);
-    void forward(Tensor x, Tensor out, Tensor pool = {}, Tensor norm_q = {}, Tensor norm_k = {}, Tensor rotary_emb = {});
+    void forward(
+        Tensor x, Tensor out, 
+        Tensor pool = {}, Tensor norm_q = {}, Tensor norm_k = {}, Tensor rotary_emb = {}, 
+        Tensor out_q = {}, Tensor out_k = {}, Tensor out_v = {}, int numTokens = 0
+    );
     std::variant<Tensor, QuantizedActivation> forward_quant(QuantizedActivation qact, FuseOptions fuse, GEMM_W4A4 *nextGEMM = nullptr);
     Tensor forward_quant(QuantizedActivation qact);
 

src/Module.cpp

+#include "common.h"
+#include "Module.h"
+#include "kernels/misc_kernels.h"
+
+void Module::copyWithCast(Tensor dst, Tensor src) {
+    assert(dst.is_contiguous());
+    assert(dst.device().type == Device::CUDA);
+
+    if (src.device().type == Device::CUDA && src.device().idx == dst.device().idx) {
+        nunchaku::kernels::cast(src, dst);
+    } else {
+        Tensor tmp;
+        tmp.buffer = dst.buffer;
+        tmp.shape = dst.shape;
+        tmp.scalarType = src.scalarType;
+        tmp.copy_(src);
+        nunchaku::kernels::cast(tmp, dst);
+    }
+}

src/Module.h

@@ -131,11 +131,24 @@ public:
             m->enabledLazyLoad = val;
         });
     }
+    void setAutoCastFP16(bool val) {
+        traverse([val](Module *m) {
+            m->enabledAutoCastFP16 = val;
+        });
+    }
 
 protected:
     virtual void loadParam(std::string key, Tensor &dst, Tensor src) {
+        static const std::set<Tensor::ScalarType> whitelist = {
+            Tensor::FP16,
+            Tensor::BF16,
+        };
+        if (enabledAutoCastFP16 && dst.scalar_type() != src.scalar_type() && whitelist.contains(dst.scalar_type()) && whitelist.contains(src.scalar_type())) {
+            copyWithCast(dst, src);
+        } else {
             dst.copy_(src);
         }
+    }
 
     struct ChildrenRegisterHelper {
         ChildrenRegisterHelper(Module &self) : self(self) {}
@@ -174,7 +187,7 @@ protected:
     }
 
     void debug(std::string name, Tensor tensor) {
-        if (DebugContext::ctxs.empty()) {
+        if (DebugContext::ctxs.empty() || !tensor.valid()) {
             return;
         }
         std::string prefix = getFullName();
@@ -187,6 +200,9 @@ protected:
         }
     }
 
+private:
+    void copyWithCast(Tensor dst, Tensor src);
+
 public:
     Module *parent = nullptr;
     std::string name = "";
@@ -194,6 +210,7 @@ public:
     std::map<std::string, Param> params;
 
     bool enabledLazyLoad = false;
+    bool enabledAutoCastFP16 = true;
 };
 
 struct LayerOffloadHelper {

src/SanaModel.cpp

+#include <iostream>
+
 #include "SanaModel.h"
 #include "kernels/zgemm/zgemm.h"
 #include "flash_api.h"
@@ -8,6 +10,7 @@
 using spdlog::fmt_lib::format;
 using namespace nunchaku;
 
+
 SanaLinearAttention::SanaLinearAttention(int dim, bool bias, bool pag, bool use_fp4, Tensor::ScalarType dtype, Device device) :
     dim(dim),
     dim_pad(ceilDiv(dim, 128) * 128),
@@ -66,7 +69,8 @@ Tensor SanaLinearAttention::forward(Tensor x, Tensor out) {
         qact.is_unsigned, qkv_proj.lora_scales, false,
         qkv_proj.use_fp4,
         *qkv_proj.wtscale.data_ptr<float>(),
-        qkv_proj.wcscales.numel() > 0 ? qkv_proj.wcscales : Tensor{}
+        qkv_proj.wcscales.numel() > 0 ? qkv_proj.wcscales : Tensor{},
+        {}, {}, {}, 0
         );
 
     debug("vk", vk);
@@ -324,8 +328,8 @@ SanaModel::SanaModel(SanaConfig config, Tensor::ScalarType dtype, Device device)
     }
 }
 
-Tensor SanaModel::forward(Tensor hidden_states, Tensor encoder_hidden_states, Tensor timestep, Tensor cu_seqlens_img, Tensor cu_seqlens_txt, int H, int W, bool pag, bool cfg) {
-    for (int i = 0; i < config.num_layers; i++) {
+Tensor SanaModel::forward(Tensor hidden_states, Tensor encoder_hidden_states, Tensor timestep, Tensor cu_seqlens_img, Tensor cu_seqlens_txt, int H, int W, bool pag, bool cfg, bool skip_first_layer) {
+    for (int i = (skip_first_layer ? 1 : 0); i < config.num_layers; i++) {
         auto &&block = transformer_blocks[i];
         hidden_states = block->forward(
             hidden_states, encoder_hidden_states, timestep, cu_seqlens_img, cu_seqlens_txt, H, W,

src/SanaModel.h

@@ -89,7 +89,7 @@ struct SanaConfig {
 class SanaModel : public Module {
 public:
     SanaModel(SanaConfig config, Tensor::ScalarType dtype, Device device);
-    Tensor forward(Tensor hidden_states, Tensor encoder_hidden_states, Tensor timestep, Tensor cu_seqlens_img, Tensor cu_seqlens_txt, int H, int W, bool pag, bool cfg);
+    Tensor forward(Tensor hidden_states, Tensor encoder_hidden_states, Tensor timestep, Tensor cu_seqlens_img, Tensor cu_seqlens_txt, int H, int W, bool pag, bool cfg, bool skip_first_layer);
 
 public:
     const SanaConfig config;

src/Tensor.h

@@ -81,7 +81,8 @@ public:
     BufferCUDA(size_t size) {
         this->size = size;
         this->device.type = Device::CUDA;
-        checkCUDA(cudaGetDevice(&this->device.idx));
+        // checkCUDA(cudaGetDevice(&this->device.idx));
+        this->device.idx = CUDADeviceContext::getDevice();
         if (size == 0) {
             this->ptr = nullptr;
         }
@@ -418,6 +419,7 @@ public:
             result.buffer = std::make_shared<BufferMalloc>(shape.size() * scalarSize.at(scalarType));
         } else if (device.type == Device::CUDA) {
             // TODO: cross device allocate
+            CUDADeviceContext ctx(device.idx);
             result.buffer = std::make_shared<BufferCUDA>(shape.size() * scalarSize.at(scalarType));
         } else {
             assert(false);
@@ -429,6 +431,7 @@ public:
             if (device.type == Device::CPU) {
                 memset(result.buffer->getPtr(), 0xCC, result.buffer->getSize());
             } else if (device.type == Device::CUDA) {
+                CUDADeviceContext ctx(device.idx);
                 checkCUDA(cudaMemsetAsync(result.buffer->getPtr(), 0xCC, result.buffer->getSize(), getCurrentCUDAStream()));
             }
         }

src/common.h

@@ -107,16 +107,97 @@ struct CUDAEventWrapper {
     }
 };
 
+
+/**
+ * 1. hold one when entered from external code (set `device` to -1 to avoid device change)
+ * 2. hold one when switching device
+ * 3. hold one with `disableCache` when calling external code that may change the device
+ */
+class CUDADeviceContext {
+public:
+    CUDADeviceContext(int device = -1, bool disableCache = false) : disableCache(disableCache) {
+        if (cacheDisabled()) {
+            // no previous context => we might entered from external code, reset cache
+            // previous context is reset on => external code may be executed, reset
+            currentDeviceCache = -1;
+        }
+        
+        ctxs.push(this);
+        lastDevice = getDevice();
+        if (device >= 0) {
+            setDevice(device);
+        }
+
+        if (disableCache) {
+            // we are about to call external code, reset cache
+            currentDeviceCache = -1;
+        }
+    }
+    CUDADeviceContext(const CUDADeviceContext &) = delete;
+    CUDADeviceContext(CUDADeviceContext &&) = delete;
+
+    ~CUDADeviceContext() {
+        if (disableCache) {
+            // retured from external code, cache is not reliable, reset
+            currentDeviceCache = -1;
+        }
+
+        setDevice(lastDevice);
+        assert(ctxs.top() == this);
+        ctxs.pop();
+
+        if (cacheDisabled()) {
+            // ctxs.empty() => we are about to return to external code, reset cache
+            // otherwise => we are a nested context in a previous context with reset on, we might continue to execute external code, reset
+            currentDeviceCache = -1;
+        }
+    }
+
+    const bool disableCache;
+    int lastDevice;
+
+
+public:
+    static int getDevice() {
+        int idx = -1;
+        if (cacheDisabled() || currentDeviceCache < 0) {
+            checkCUDA(cudaGetDevice(&idx));
+        } else {
+            idx = currentDeviceCache;
+        }
+        currentDeviceCache = cacheDisabled() ? -1 : idx;
+        return idx;
+    }
+private:
+    static void setDevice(int idx) {
+        // TODO: deal with stream when switching device
+        assert(idx >= 0);
+        if (!cacheDisabled() && currentDeviceCache == idx) {
+            return;
+        }
+        checkCUDA(cudaSetDevice(idx));
+        currentDeviceCache = cacheDisabled() ? -1 : idx;
+    }
+
+private:
+    static inline thread_local std::stack<CUDADeviceContext *> ctxs;
+    static inline thread_local int currentDeviceCache = -1;
+
+    static bool cacheDisabled() {
+        return ctxs.empty() || ctxs.top()->disableCache;
+    }
+};
+
 inline cudaDeviceProp *getCurrentDeviceProperties() {
-    static thread_local cudaDeviceProp prop;
-    static thread_local bool propAvailable = false;
-    if (!propAvailable) {
-        int device;
-        checkCUDA(cudaGetDevice(&device));
-        checkCUDA(cudaGetDeviceProperties(&prop, device));
-        propAvailable = true;
-    }
-    return &prop;
+    static thread_local std::map<int, cudaDeviceProp> props;
+
+    int deviceId = CUDADeviceContext::getDevice();
+    if (!props.contains(deviceId)) {
+        cudaDeviceProp prop;
+        checkCUDA(cudaGetDeviceProperties(&prop, deviceId));
+        props[deviceId] = prop;
+    }
+    return &props.at(deviceId);
 }
 
 template<typename T>

src/interop/torch.cpp

@@ -22,6 +22,7 @@ Tensor from_torch(at::Tensor input) {
     }
 
     static const std::map<at::ScalarType, Tensor::ScalarType> mapType = {
+        { at::ScalarType::Char, Tensor::INT8 },
         { at::ScalarType::Byte, Tensor::INT8 },
         { at::ScalarType::Int, Tensor::INT32 },
         { at::ScalarType::Long, Tensor::INT64 },
@@ -36,7 +37,7 @@ Tensor from_torch(at::Tensor input) {
     result.scalarType = mapType.at(input.scalar_type());
     result.buffer = std::make_shared<BufferTorchTensor>(std::move(input));
 
-    // Tensor::lockBuffer(result.buffer, getCurrentCUDAStream());
+    Tensor::lockBuffer(result.buffer, getCurrentCUDAStream());
 
     return result;
 }

src/interop/torch.h

@@ -15,7 +15,7 @@ public:
     }
     virtual bool isAsyncBuffer() override {
         // TODO: figure out how torch manages memory
-        return true;
+        return this->device.type == Device::CUDA;
     }
 private:
     at::Tensor tensor;
@@ -31,3 +31,21 @@ public:
 
 Tensor from_torch(at::Tensor input);
 at::Tensor to_torch(Tensor input);
+
+class TensorsProviderTorch : public TensorsProvider {
+public:
+    TensorsProviderTorch(std::map<std::string, at::Tensor> dict) : storage(std::move(dict)) {}
+
+    virtual bool contains(const std::string &key) const override {
+        return storage.contains(key);
+    }
+    virtual Tensor getTensor(const std::string &key) override {
+        if (!storage.contains(key)) {
+            return Tensor{};
+        }
+        return from_torch(storage.at(key));
+    }
+
+private:
+    std::map<std::string, at::Tensor> storage;
+};
\ No newline at end of file