refactor: unify ggml operation helpers around runner context (#2314)

Co-authored-by: Maxime Dréan <maximedrean@proton.me>

.dockerignore

@@ -1,7 +1,8 @@
 build*/
 docs/
 test/
+tests/
 
 .cache/
 *.swp
-models/
+models/

.github/workflows/build.yml

@@ -16,6 +16,7 @@ on:
         ".github/workflows/**",
         ".dockerignore",
         "docker/Dockerfile*",
+        "scripts/ci-cuda-cache.sh",
         "**/CMakeLists.txt",
         "**/Makefile",
         "**/*.h",
@@ -23,6 +24,7 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
+        "**/*.cuh",
         "ggml",
         "examples/server/frontend",
         "examples/server/frontend/**",
@@ -37,6 +39,7 @@ on:
         ".github/workflows/**",
         ".dockerignore",
         "docker/Dockerfile*",
+        "scripts/ci-cuda-cache.sh",
         "**/CMakeLists.txt",
         "**/Makefile",
         "**/*.h",
@@ -44,6 +47,7 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
+        "**/*.cuh",
         "ggml",
         "examples/server/frontend",
         "examples/server/frontend/**",
@@ -55,6 +59,9 @@ on:
 env:
   BRANCH_NAME: ${{ github.head_ref || github.ref_name }}
 
+permissions:
+  contents: read
+
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
   cancel-in-progress: true
@@ -66,17 +73,17 @@ jobs:
     steps:
       - name: Clone
         id: checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v6
         with:
           submodules: recursive
 
       - name: Setup Node
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@v5
         with:
-          node-version: 20
+          node-version: 24
 
       - name: Setup pnpm
-        uses: pnpm/action-setup@v4
+        uses: pnpm/action-setup@v5
         with:
           version: 10.15.1
 
@@ -97,7 +104,8 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        shell: bash
+        run: echo "short=${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT"
 
       - name: Fetch system info
         id: system-info
@@ -118,7 +126,7 @@ jobs:
 
       - name: Upload artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}.zip
           path: |
@@ -130,17 +138,17 @@ jobs:
     steps:
       - name: Clone
         id: checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v6
         with:
           submodules: recursive
 
       - name: Setup Node
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@v5
         with:
-          node-version: 20
+          node-version: 24
 
       - name: Setup pnpm
-        uses: pnpm/action-setup@v4
+        uses: pnpm/action-setup@v5
         with:
           version: 10.15.1
 
@@ -161,7 +169,8 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        shell: bash
+        run: echo "short=${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT"
 
       - name: Fetch system info
         id: system-info
@@ -182,7 +191,7 @@ jobs:
 
       - name: Upload artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-vulkan.zip
           path: |
@@ -230,25 +239,27 @@ jobs:
           submodules: recursive
 
       - name: Setup Node
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@v5
         with:
-          node-version: 20
+          node-version: 24
 
       - name: Setup pnpm
-        uses: pnpm/action-setup@v4
+        uses: pnpm/action-setup@v5
         with:
           version: 10.15.1
 
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        shell: bash
+        run: echo "short=${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT"
 
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
+        id: buildx
+        uses: docker/setup-buildx-action@v4
 
       - name: Log in to the container registry
-        uses: docker/login-action@v3
+        uses: docker/login-action@v4
         with:
           registry: ${{ env.REGISTRY }}
           username: ${{ github.actor }}
@@ -256,7 +267,7 @@ jobs:
 
       - name: Extract metadata for Docker
         id: meta
-        uses: docker/metadata-action@v5
+        uses: docker/metadata-action@v6
         with:
           images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
 
@@ -267,9 +278,24 @@ jobs:
           # if set to "true" but frees about 6 GB
           tool-cache: false
 
+      - name: Cache CUDA compilation
+        if: matrix.variant == 'cuda'
+        uses: actions/cache@v5
+        with:
+          path: ${{ runner.temp }}/cuda-ccache
+          key: cuda-ccache-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('docker/Dockerfile.cuda') }}-${{ github.run_id }}-${{ github.run_attempt }}
+          restore-keys: |
+            cuda-ccache-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('docker/Dockerfile.cuda') }}-
+
+      - name: Restore CUDA compiler cache mount
+        if: matrix.variant == 'cuda'
+        env:
+          CUDA_CACHE_BUILDER: ${{ steps.buildx.outputs.name }}
+        run: bash scripts/ci-cuda-cache.sh restore
+
       - name: Build and push Docker image
         id: build-push
-        uses: docker/build-push-action@v6
+        uses: docker/build-push-action@v7
         with:
           context: .
           platforms: ${{ matrix.platform }}
@@ -280,23 +306,29 @@ jobs:
           annotations: ${{ steps.meta.outputs.annotations }}
           build-args: ${{ matrix.build-args }}
 
+      - name: Export CUDA compiler cache mount
+        if: matrix.variant == 'cuda'
+        env:
+          CUDA_CACHE_BUILDER: ${{ steps.buildx.outputs.name }}
+        run: bash scripts/ci-cuda-cache.sh save
+
   macOS-latest-cmake:
     runs-on: macos-latest
 
     steps:
       - name: Clone
         id: checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v6
         with:
           submodules: recursive
 
       - name: Setup Node
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@v5
         with:
-          node-version: 20
+          node-version: 24
 
       - name: Setup pnpm
-        uses: pnpm/action-setup@v4
+        uses: pnpm/action-setup@v5
         with:
           version: 10.15.1
 
@@ -317,7 +349,8 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        shell: bash
+        run: echo "short=${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT"
 
       - name: Fetch system info
         id: system-info
@@ -338,7 +371,7 @@ jobs:
 
       - name: Upload artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-${{ steps.system-info.outputs.OS_NAME }}-${{ steps.system-info.outputs.OS_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}.zip
           path: |
@@ -356,35 +389,44 @@ jobs:
           - build: "cpu"
             defines: "-DGGML_NATIVE=OFF -DSD_BUILD_SHARED_LIBS=ON -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON"
           - build: "cuda12"
-            defines: "-DSD_CUDA=ON -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"' -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON"
+            defines: "-DSD_CUDA=ON -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"' -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON"
           - build: "vulkan"
             defines: "-DSD_VULKAN=ON -DSD_BUILD_SHARED_LIBS=ON -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON"
     steps:
       - name: Clone
         id: checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v6
         with:
           submodules: recursive
 
       - name: Setup Node
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@v5
         with:
-          node-version: 20
+          node-version: 24
 
       - name: Setup pnpm
-        uses: pnpm/action-setup@v4
+        uses: pnpm/action-setup@v5
         with:
           version: 10.15.1
 
       - name: Install cuda-toolkit
         id: cuda-toolkit
         if: ${{ matrix.build == 'cuda12' }}
-        uses: Jimver/cuda-toolkit@v0.2.22
+        uses: Jimver/cuda-toolkit@v0.2.35
         with:
           cuda: "12.8.1"
           method: "network"
           sub-packages: '["nvcc", "cudart", "cublas", "cublas_dev", "thrust", "visual_studio_integration"]'
 
+      - name: Cache CUDA compilation
+        if: matrix.build == 'cuda12'
+        uses: ggml-org/ccache-action@v1.2.21
+        with:
+          variant: ccache
+          install: binary
+          key: windows-cuda-12.8.1-x64
+          max-size: "2G"
+
       - name: Install Vulkan SDK
         id: get_vulkan
         if: ${{ matrix.build == 'vulkan' }}
@@ -394,13 +436,15 @@ jobs:
           Add-Content $env:GITHUB_ENV "VULKAN_SDK=C:\VulkanSDK\${env:VULKAN_VERSION}"
           Add-Content $env:GITHUB_PATH "C:\VulkanSDK\${env:VULKAN_VERSION}\bin"
 
-      - name: Activate MSVC environment
-        id: msvc_dev_cmd
-        uses: ilammy/msvc-dev-cmd@v1
-
       - name: Build
         id: cmake_build
+        shell: pwsh
         run: |
+          $vswhere = "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe"
+          $vsPath = & $vswhere -latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -property installationPath
+          if (-not $vsPath) { throw "Visual Studio C++ tools not found" }
+          & "$vsPath\Common7\Tools\Launch-VsDevShell.ps1" -Arch amd64 -HostArch amd64 -SkipAutomaticLocation
+
           mkdir build
           cd build
           cmake .. -DCMAKE_CXX_FLAGS='/bigobj' -G Ninja -DCMAKE_C_COMPILER=cl.exe -DCMAKE_CXX_COMPILER=cl.exe -DCMAKE_BUILD_TYPE=Release ${{ matrix.defines }}
@@ -409,7 +453,8 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        shell: bash
+        run: echo "short=${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT"
 
       - name: Pack artifacts
         id: pack_artifacts
@@ -429,6 +474,9 @@ jobs:
             ls .\build\bin
             throw "Can't find stable-diffusion.dll"
           }
+          if ("${{ matrix.build }}" -eq "cuda12") {
+            Copy-Item thirdparty/cutlass/LICENSE.txt (Join-Path (Split-Path -Parent $filePath) 'cutlass.txt')
+          }
           7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-${{ matrix.build }}-x64.zip $filePath .\start-server.sh .\start-server.bat
 
       - name: Copy and pack Cuda runtime
@@ -442,7 +490,7 @@ jobs:
 
       - name: Upload Cuda runtime
         if: ${{ matrix.build == 'cuda12' && (github.event_name == 'push' && github.ref == 'refs/heads/master' || github.event.inputs.create_release == 'true') }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: sd-cudart-sd-bin-win-cu12-x64.zip
           path: |
@@ -450,7 +498,7 @@ jobs:
 
       - name: Upload artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-${{ matrix.build }}-x64.zip
           path: |
@@ -464,29 +512,29 @@ jobs:
       gpu_targets: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v6
         with:
           submodules: recursive
 
       - name: Setup Node
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@v5
         with:
-          node-version: 20
+          node-version: 24
 
       - name: Setup pnpm
-        uses: pnpm/action-setup@v4
+        uses: pnpm/action-setup@v5
         with:
           version: 10.15.1
 
       - name: Cache ROCm Installation
         id: cache-rocm
-        uses: actions/cache@v4
+        uses: actions/cache@v5
         with:
           path: C:\TheRock\build
           key: rocm-wheels-${{ env.ROCM_VERSION }}-${{ runner.os }}
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.16
+        uses: ggml-org/ccache-action@v1.2.21
         with:
           key: windows-rocm-${{ env.ROCM_VERSION }}-x64
           evict-old-files: 1d
@@ -567,7 +615,8 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: pr-mpt/actions-commit-hash@v2
+        shell: bash
+        run: echo "short=${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT"
 
       - name: Bundle HIP runtime libraries
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
@@ -588,7 +637,7 @@ jobs:
 
       - name: Upload artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip
           path: |
@@ -615,7 +664,7 @@ jobs:
           submodules: recursive
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.16
+        uses: ggml-org/ccache-action@v1.2.21
         with:
           key: ubuntu-rocm-cmake-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
           evict-old-files: 1d
@@ -688,12 +737,12 @@ jobs:
       # setup-node installs into /opt/hostedtoolcache, which is removed above.
       # Keep Node/pnpm setup after disk cleanup so the server frontend can be embedded.
       - name: Setup Node
-        uses: actions/setup-node@v4
+        uses: actions/setup-node@v5
         with:
-          node-version: 20
+          node-version: 24
 
       - name: Setup pnpm
-        uses: pnpm/action-setup@v4
+        uses: pnpm/action-setup@v5
         with:
           version: 10.15.1
 
@@ -722,7 +771,8 @@ jobs:
       - name: Get commit hash
         id: commit
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
+        shell: bash
+        run: echo "short=${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT"
 
       - name: Prepare artifacts
         id: prepare_artifacts
@@ -751,7 +801,7 @@ jobs:
 
       - name: Upload artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v6
         with:
           name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm-${{ matrix.ROCM_VERSION }}.zip
           path: |
@@ -762,6 +812,9 @@ jobs:
 
     runs-on: ubuntu-latest
 
+    permissions:
+      contents: write
+
     needs:
       - ubuntu-latest-cmake
       - ubuntu-latest-cmake-vulkan
@@ -773,13 +826,13 @@ jobs:
 
     steps:
       - name: Clone
-        uses: actions/checkout@v3
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
 
       - name: Download artifacts
         id: download-artifact
-        uses: actions/download-artifact@v4
+        uses: actions/download-artifact@v7
         with:
           path: ./artifact
           pattern: sd-*
@@ -792,36 +845,47 @@ jobs:
 
       - name: Get commit hash
         id: commit
-        uses: prompt/actions-commit-hash@v2
+        shell: bash
+        run: echo "short=${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT"
 
       - name: Create release
         id: create_release
         if: ${{ github.event_name == 'workflow_dispatch' || github.ref_name == 'master' }}
-        uses: anzz1/action-create-release@v1
+        uses: actions/github-script@v9
         env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          RELEASE_TAG: ${{ format('{0}-{1}-{2}', env.BRANCH_NAME, steps.commit_count.outputs.count, steps.commit.outputs.short) }}
         with:
-          tag_name: ${{ format('{0}-{1}-{2}', env.BRANCH_NAME, steps.commit_count.outputs.count, steps.commit.outputs.short) }}
+          script: |
+            const { data: release } = await github.rest.repos.createRelease({
+              ...context.repo,
+              tag_name: process.env.RELEASE_TAG,
+              name: process.env.RELEASE_TAG,
+              target_commitish: context.sha,
+              draft: false,
+              prerelease: false,
+            });
+            core.setOutput('id', release.id);
 
       - name: Upload release
         id: upload_release
         if: ${{ github.event_name == 'workflow_dispatch' || github.ref_name == 'master' }}
-        uses: actions/github-script@v3
+        uses: actions/github-script@v9
+        env:
+          RELEASE_ID: ${{ steps.create_release.outputs.id }}
         with:
-          github-token: ${{secrets.GITHUB_TOKEN}}
           script: |
             const path = require('path');
-            const fs = require('fs');
-            const release_id = '${{ steps.create_release.outputs.id }}';
-            for (let file of await fs.readdirSync('./artifact')) {
+            const fs = require('fs').promises;
+            const release_id = Number(process.env.RELEASE_ID);
+            for (const file of await fs.readdir('./artifact')) {
               if (path.extname(file) === '.zip') {
                 console.log('uploadReleaseAsset', file);
-                await github.repos.uploadReleaseAsset({
+                await github.rest.repos.uploadReleaseAsset({
                   owner: context.repo.owner,
                   repo: context.repo.repo,
                   release_id: release_id,
                   name: file,
-                  data: await fs.readFileSync(`./artifact/${file}`)
+                  data: await fs.readFile(`./artifact/${file}`)
                 });
               }
             }

.gitignore

@@ -2,6 +2,7 @@ build*/
 cmake-build-*/
 test/
 tests/
+logs/
 .vscode/
 .idea/
 .cache/

.gitmodules

@@ -10,3 +10,6 @@
 [submodule "thirdparty/libwebm"]
 	path = thirdparty/libwebm
 	url = https://github.com/webmproject/libwebm.git
+[submodule "thirdparty/cutlass"]
+	path = thirdparty/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

CMakeLists.txt

@@ -9,13 +9,28 @@ if (NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE)
 endif()
 
 if (MSVC)
+    cmake_host_system_information(RESULT SD_LOGICAL_CORES QUERY NUMBER_OF_LOGICAL_CORES)
+    if (SD_LOGICAL_CORES GREATER 8)
+        math(EXPR SD_DEFAULT_MSVC_MP_COUNT "${SD_LOGICAL_CORES} - 2")
+    elseif (SD_LOGICAL_CORES GREATER 2)
+        math(EXPR SD_DEFAULT_MSVC_MP_COUNT "${SD_LOGICAL_CORES} - 1")
+    else()
+        set(SD_DEFAULT_MSVC_MP_COUNT "${SD_LOGICAL_CORES}")
+    endif()
+    if (NOT SD_DEFAULT_MSVC_MP_COUNT GREATER 0)
+        set(SD_DEFAULT_MSVC_MP_COUNT 1)
+    endif()
+    set(SD_MSVC_MP_COUNT "${SD_DEFAULT_MSVC_MP_COUNT}" CACHE STRING "Maximum parallel MSVC compilations per project")
+    if (NOT SD_MSVC_MP_COUNT MATCHES "^[1-9][0-9]*$" OR SD_MSVC_MP_COUNT GREATER 65536)
+        message(FATAL_ERROR "SD_MSVC_MP_COUNT must be an integer between 1 and 65536")
+    endif()
     add_compile_definitions(_CRT_SECURE_NO_WARNINGS)
     add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING)
     # /MP is MSVC-only: icx rejects it outright once offloading is enabled.
     add_compile_options(
-        $<$<AND:$<COMPILE_LANGUAGE:C>,$<C_COMPILER_ID:MSVC>>:/MP>
+        $<$<AND:$<COMPILE_LANGUAGE:C>,$<C_COMPILER_ID:MSVC>>:/MP${SD_MSVC_MP_COUNT}>
         $<$<COMPILE_LANGUAGE:C>:/utf-8>
-        $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/MP>
+        $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/MP${SD_MSVC_MP_COUNT}>
         $<$<COMPILE_LANGUAGE:CXX>:/utf-8>
     )
 endif()
@@ -111,6 +126,12 @@ if(SD_CUDA)
 endif()
 
 if(SD_METAL)
+    if(NOT SD_USE_UPSTREAM_GGML)
+        if(CMAKE_VERSION VERSION_LESS 3.16)
+            message(FATAL_ERROR "SD_METAL with SD_USE_UPSTREAM_GGML=OFF requires CMake 3.16 or later for OBJC and OBJCXX support")
+        endif()
+        enable_language(OBJC OBJCXX)
+    endif()
     message("-- Use Metal as backend stable-diffusion")
     set(GGML_METAL ON)
 endif()
@@ -225,6 +246,8 @@ file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS
     "src/core/*.h"
     "src/core/*.cpp"
     "src/core/*.hpp"
+    "src/core/cpu/*.h"
+    "src/core/cpu/*.cpp"
     "src/core/probe/*.h"
     "src/core/probe/*.cpp"
     "src/extensions/*.h"
@@ -343,6 +366,51 @@ if (SD_CUDA)
     # Keep the driver stub on downstream link lines when no driver is installed.
     target_link_libraries(${SD_LIB} PUBLIC CUDA::cuda_driver)
     set_property(SOURCE src/core/ggml_extend_backend.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_USE_CUDA)
+    if (NOT SD_USE_UPSTREAM_GGML)
+        include(cmake/cutlass.cmake)
+        enable_language(CUDA)
+        set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
+        if (SD_CUDA_CUTLASS_PATH)
+            target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu src/core/cuda/bf16_gemm.cu)
+            set_source_files_properties(src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu src/core/cuda/bf16_gemm.cu PROPERTIES
+                INCLUDE_DIRECTORIES "${SD_CUDA_CUTLASS_PATH}/include"
+                COMPILE_OPTIONS "--use_fast_math;--extended-lambda;--expt-relaxed-constexpr")
+            if (MSVC)
+                set_property(SOURCE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu src/core/cuda/bf16_gemm.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus)
+            endif()
+            if (GGML_CUDA_FORCE_CUBLAS)
+                set_property(SOURCE src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu src/core/cuda/bf16_gemm.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS)
+            endif()
+        endif()
+    endif()
+endif()
+set(SD_BACKEND_REGISTRATIONS)
+if (SD_METAL AND NOT SD_USE_UPSTREAM_GGML)
+    add_subdirectory(src/core/metal)
+    target_link_libraries(${SD_LIB} PRIVATE "-framework Foundation" "-framework Metal")
+    list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_metal_i8)
+endif()
+if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA))
+    if(SD_HIPBLAS)
+        include(cmake/hip.cmake)
+    endif()
+    add_subdirectory(src/core/cuda)
+    foreach(VENDOR CUDA HIP MUSA)
+        if(SD_${VENDOR} OR (VENDOR STREQUAL "HIP" AND SD_HIPBLAS))
+            string(TOLOWER "${VENDOR}" NAMESPACE)
+            list(APPEND SD_BACKEND_REGISTRATIONS "sd_${NAMESPACE}::register_extensions")
+        endif()
+    endforeach()
+endif()
+if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML)
+    add_subdirectory(src/core/vulkan)
+    list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope sd_ggml_register_vulkan_i8 sd_ggml_register_vulkan_modulate sd_ggml_register_vulkan_rms_affine)
+endif()
+if(SD_BACKEND_REGISTRATIONS)
+    string(REPLACE ";" ",\n    " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}")
+    configure_file(cmake/backend_extensions.h.in generated/sd_backend_extensions.h @ONLY)
+    set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY INCLUDE_DIRECTORIES "${CMAKE_CURRENT_BINARY_DIR}/generated")
+    set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_BACKEND_EXTENSIONS)
 endif()
 if (SD_HIPBLAS)
     find_package(hip REQUIRED CONFIG HINTS "$ENV{ROCM_PATH}" /opt/rocm)

CONTRIBUTING.md

@@ -29,6 +29,18 @@ PRs should include:
 - Linked issue/PR context when applicable.
 - Screenshots or sample outputs for UI/visual behavior changes.
 
+## Documentation
+
+Write model usage guides in `docs/` for users. Focus on supported weights,
+setup, runnable commands, parameter meanings, output behavior, and practical
+limits such as memory requirements. Include technical details only when they
+help users choose settings or troubleshoot a problem.
+
+Keep tensor layouts, internal class names, kernel implementation details,
+algorithm derivations, and development history out of usage guides. Keep task-specific
+profiling logs, numerical comparison procedures, and test artifacts local, and
+summarize relevant verification results in the PR description.
+
 ## README UI Listings
 
 The [UIs list in the README](README.md#uis) only includes open-source projects that use `stable-diffusion.cpp` as a backend. Do not submit PRs to add closed-source UIs.

README.md

@@ -59,6 +59,7 @@ API and command-line option may change frequently.***
     - [Z-Image](./docs/z_image.md)
     - [Z-Image L2P](./docs/z_image_l2p.md)
     - [MiniT2I](./docs/minit2i.md)
+    - [P-H3-Image](./docs/p_h3_image.md)
     - [SenseNova U1.5](./docs/sensenova_u1.md)
     - [Ovis-Image](./docs/ovis_image.md)
     - [Anima](./docs/anima.md)
@@ -87,6 +88,9 @@ API and command-line option may change frequently.***
     - [LTX-2.3/LTX-2.5](./docs/ltx2.md)
     - [HunyuanVideo 1.5](./docs/hunyuan_video.md)
     - [LingBot-Video](./docs/lingbot_video.md)
+  - Video Upscale Models
+    - [RelayVSR](./docs/relay_vsr.md)
+    - [SeedVR2](./docs/seedvr2.md)
   - [PhotoMaker](./docs/photo_maker.md) support.
   - [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL, including Plus)
   - Control Net support with SD 1.5
@@ -197,6 +201,9 @@ These projects use `stable-diffusion.cpp` as a backend for their image generatio
 - [Neural-Pixel](https://github.com/Luiz-Alcantara/Neural-Pixel)
 - [KoboldCpp](https://github.com/LostRuins/koboldcpp)
 - [Kokon](https://github.com/GabrieleLozupone/kokon)
+- [NodeTool](https://github.com/nodetool-ai/nodetool)
+- [Quo](https://github.com/quo-ai/quo)
+- [Prisme](https://github.com/d3Ita/prisme)
 
 ## Contributors
 

cmake/backend_extensions.h.in

@@ -0,0 +1,17 @@
+#ifndef __SD_BACKEND_EXTENSIONS_H__
+#define __SD_BACKEND_EXTENSIONS_H__
+
+#include "core/cuda/registration.h"
+
+void sd_ggml_register_metal_i8(ggml_backend_dev_t device);
+
+void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device);
+void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device);
+void sd_ggml_register_vulkan_rms_affine(ggml_backend_dev_t device);
+void sd_ggml_register_vulkan_modulate(ggml_backend_dev_t device);
+
+static void (*const sd_backend_extensions[])(ggml_backend_dev_t) = {
+    @SD_BACKEND_REGISTRATIONS@
+};
+
+#endif

cmake/cutlass.cmake

@@ -0,0 +1,20 @@
+set(sd_cutlass_default "${PROJECT_SOURCE_DIR}/thirdparty/cutlass")
+if(NOT DEFINED SD_CUDA_CUTLASS_PATH AND DEFINED GGML_CUDA_CUTLASS_PATH)
+    get_filename_component(sd_old_cutlass_default "${SD_GGML_SOURCE_DIR}/src/ggml-cuda/vendors/cutlass" ABSOLUTE)
+    if(NOT GGML_CUDA_CUTLASS_PATH STREQUAL sd_old_cutlass_default)
+        set(sd_cutlass_default "${GGML_CUDA_CUTLASS_PATH}")
+    endif()
+    message(STATUS "Migrating GGML_CUDA_CUTLASS_PATH to SD_CUDA_CUTLASS_PATH")
+endif()
+set(SD_CUDA_CUTLASS_PATH "${sd_cutlass_default}" CACHE PATH "sd: CUTLASS source directory for INT8 GEMM and CuTe attention, empty to disable")
+unset(GGML_CUDA_CUTLASS_PATH CACHE)
+if(SD_CUDA_CUTLASS_PATH)
+    if(NOT EXISTS "${SD_CUDA_CUTLASS_PATH}/include/cutlass/gemm/kernel/default_gemm_universal_with_visitor.h" OR
+       NOT EXISTS "${SD_CUDA_CUTLASS_PATH}/include/cute/tensor.hpp")
+        message(FATAL_ERROR
+            "CUTLASS headers not found at ${SD_CUDA_CUTLASS_PATH}. "
+            "Run git submodule update --init --recursive from the sd.cpp repository root, "
+            "set SD_CUDA_CUTLASS_PATH to a CUTLASS source tree (tested with 3.9.2), "
+            "or set SD_CUDA_CUTLASS_PATH to an empty string to disable CUTLASS.")
+    endif()
+endif()

cmake/embed_spirv.cmake

@@ -0,0 +1,3 @@
+file(READ "${INPUT}" SPIRV HEX)
+string(REGEX REPLACE "(..)(..)(..)(..)" "0x\\4\\3\\2\\1," SPIRV_WORDS "${SPIRV}")
+file(WRITE "${OUTPUT}" "#include <cstdint>\nnamespace sd_vulkan_shaders {\nstatic const uint32_t ${NAME}[] = {${SPIRV_WORDS}};\n}\n")

cmake/ggml.cmake

@@ -9,17 +9,19 @@ if(NOT TARGET ggml)
     add_library(ggml ALIAS ggml::ggml)
 endif()
 
-get_target_property(sd_ggml_imported ggml IMPORTED)
-if(sd_ggml_imported)
-    set(sd_ggml_private_include "${SD_GGML_SOURCE_DIR}/src")
-else()
-    get_target_property(sd_ggml_private_include ggml SOURCE_DIR)
+if(SD_USE_UPSTREAM_GGML)
+    get_target_property(sd_ggml_imported ggml IMPORTED)
+    if(sd_ggml_imported)
+        set(sd_ggml_private_include "${SD_GGML_SOURCE_DIR}/src")
+    else()
+        get_target_property(sd_ggml_private_include ggml SOURCE_DIR)
+    endif()
+    if(NOT EXISTS "${sd_ggml_private_include}/ggml-impl.h")
+        message(FATAL_ERROR "Set SD_GGML_SOURCE_DIR to the source tree matching the selected ggml library (ggml-impl.h is required).")
+    endif()
+    target_include_directories(${SD_LIB} PRIVATE "${sd_ggml_private_include}")
+    set_property(TARGET ${SD_LIB} PROPERTY SD_GGML_PRIVATE_INCLUDE_DIR "${sd_ggml_private_include}")
 endif()
-if(NOT EXISTS "${sd_ggml_private_include}/ggml-impl.h")
-    message(FATAL_ERROR "Set SD_GGML_SOURCE_DIR to the source tree matching the selected ggml library (ggml-impl.h is required).")
-endif()
-target_include_directories(${SD_LIB} PRIVATE "${sd_ggml_private_include}")
-set_property(TARGET ${SD_LIB} PROPERTY SD_GGML_PRIVATE_INCLUDE_DIR "${sd_ggml_private_include}")
 
 if(SD_USE_UPSTREAM_GGML)
     target_compile_definitions(${SD_LIB} PUBLIC SD_USE_UPSTREAM_GGML)

cmake/hip.cmake

@@ -0,0 +1,12 @@
+if(NOT WIN32 AND NOT CMAKE_CXX_COMPILER MATCHES "hipcc")
+    if(NOT CMAKE_HIP_ARCHITECTURES)
+        if(GPU_TARGETS)
+            set(CMAKE_HIP_ARCHITECTURES ${GPU_TARGETS})
+        elseif(AMDGPU_TARGETS)
+            set(CMAKE_HIP_ARCHITECTURES ${AMDGPU_TARGETS})
+        endif()
+    endif()
+    # HIP objects are linked by the parent target, so initialize its link rules here.
+    enable_language(HIP)
+    set_property(TARGET ${SD_LIB} PROPERTY HIP_ARCHITECTURES "${CMAKE_HIP_ARCHITECTURES}")
+endif()

docker/Dockerfile

@@ -8,7 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends build-essential
     curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
     gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
     rm /tmp/nodesource-repo.gpg.key && \
-    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_24.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
     apt-get update && \
     apt-get install -y --no-install-recommends nodejs && \
     npm install -g pnpm@10.15.1 && \

docker/Dockerfile.cuda

@@ -1,7 +1,9 @@
+# syntax=docker/dockerfile:1
+
 ARG CUDA_VERSION=12.6.3
 ARG UBUNTU_VERSION=24.04
 
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION} AS build
+FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION} AS build
 
 # sd-server embeds the web UI at build time, so the build image needs Node/pnpm.
 RUN apt-get update && apt-get install -y --no-install-recommends build-essential git ccache cmake ca-certificates curl gnupg && \
@@ -9,7 +11,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends build-essential
     curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
     gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
     rm /tmp/nodesource-repo.gpg.key && \
-    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_24.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
     apt-get update && \
     apt-get install -y --no-install-recommends nodejs && \
     npm install -g pnpm@10.15.1 && \
@@ -25,8 +27,15 @@ ARG CUDA_ARCHITECTURES=""
 ARG GGML_CUDA_FA_ALL_QUANTS=""
 ARG GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS=ON
 
+ENV CCACHE_DIR=/root/.cache/ccache \
+    CCACHE_MAXSIZE=2G \
+    CCACHE_COMPILERCHECK=content
+
 RUN set -- \
-        -DSD_CUDA=ON; \
+        -DSD_CUDA=ON \
+        -DCMAKE_C_COMPILER_LAUNCHER=ccache \
+        -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
+        -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache; \
     if [ "${GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS}" = "ON" ]; then \
         set -- "$@" \
             -DSD_BUILD_SHARED_LIBS=ON \
@@ -44,15 +53,20 @@ RUN set -- \
         set -- "$@" "-DGGML_CUDA_FA_ALL_QUANTS=${GGML_CUDA_FA_ALL_QUANTS}"; \
     fi; \
     cmake . -B ./build "$@"
-RUN cmake --build ./build --config Release -j$(nproc)
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    ccache --zero-stats && \
+    cmake --build ./build --config Release -j$(nproc) && \
+    ccache --cleanup && \
+    ccache --show-stats
 
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu${UBUNTU_VERSION} AS runtime
+FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu${UBUNTU_VERSION} AS runtime
 
 RUN apt-get update && \
     apt-get install --yes --no-install-recommends libgomp1 && \
     apt-get clean
 
 COPY --from=build /sd.cpp/build/bin /sd.cpp/bin
+COPY --from=build /sd.cpp/thirdparty/cutlass/LICENSE.txt /sd.cpp/bin/cutlass.txt
 RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
     printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
     chmod +x /sd-cli /sd-server

docker/Dockerfile.musa

@@ -8,7 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends ccache cmake gi
     curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
     gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
     rm /tmp/nodesource-repo.gpg.key && \
-    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_24.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
     apt-get update && \
     apt-get install -y --no-install-recommends nodejs && \
     npm install -g pnpm@10.15.1 && \

docker/Dockerfile.rocm

@@ -7,7 +7,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
     gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
     rm /tmp/nodesource-repo.gpg.key && \
-    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_24.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
     apt-get update && apt-get install -y --no-install-recommends nodejs && \
     npm install -g pnpm@10.15.1 && \
     apt-get clean && rm -rf /var/lib/apt/lists/*

docker/Dockerfile.sycl

@@ -9,7 +9,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends cmake ca-certif
     curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
     gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
     rm /tmp/nodesource-repo.gpg.key && \
-    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_24.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
     apt-get update && \
     apt-get install -y --no-install-recommends nodejs && \
     npm install -g pnpm@10.15.1 && \

docker/Dockerfile.vulkan

@@ -8,7 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends build-essential
     curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
     gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
     rm /tmp/nodesource-repo.gpg.key && \
-    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_24.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
     apt-get update && \
     apt-get install -y --no-install-recommends nodejs && \
     npm install -g pnpm@10.15.1 && \

docs/backend.md

@@ -135,7 +135,10 @@ regardless of argument order, even with `--auto-fit on`.
 
 Auto-fit preserves explicit `--backend` assignments, including per-module
 assignments and device lists. For modules without a runtime assignment, it chooses
-the GPU with the largest available memory budget (the first device on a tie).
+the discrete GPU with the largest positive memory budget (the first device on a
+tie), falling back to integrated GPUs when no discrete GPU has a positive budget.
+An integrated GPU's shared system memory does not give it priority over a discrete
+GPU. Explicit assignments such as `--backend Vulkan1` still select that device.
 It then derives parameter placements from the model metadata, each module's
 compute devices, and the remaining memory budgets. The chosen backend
 specifications are printed.
@@ -247,6 +250,9 @@ These are backend-level defaults. GGML still chooses the actual kernel based on
 weight type, shape, and hardware; enabling direct convolution does not guarantee
 that every operation uses implicit GEMM. The CUDA-family implicit GEMM paths
 require F16 weights and can use ordinary direct kernels for other cases.
+For 3D convolution, a supported direct kernel is also used as a fallback when
+the im2col decomposition is unsupported for the actual tensor sizes, including
+when an intermediate reshape or matrix multiplication exceeds backend limits.
 This option does not change the separately created image upscaler context.
 
 Repeating `--conv-direct` replaces the previous assignment string. Existing

docs/build.md

@@ -2,20 +2,44 @@
 
 ## Get the Code
 
-```
+```shell
 git clone --recursive https://github.com/leejet/stable-diffusion.cpp
 cd stable-diffusion.cpp
 ```
 
-- If you have already cloned the repository, you can use the following command to update the repository to the latest code.
+If you already cloned the repository without submodules, initialize them from the repository root:
 
+```shell
+git submodule update --init --recursive
 ```
+
+To update an existing checkout and its submodules:
+
+```shell
 cd stable-diffusion.cpp
 git pull origin master
-git submodule init
-git submodule update
+git submodule update --init --recursive
 ```
 
+## MSVC compilation parallelism
+
+MSVC builds limit parallel C/C++ compilation within each project by default.
+For `n` logical processors, the limit is `n` when `n <= 2`, `n - 1` when
+`3 <= n <= 8`, and `n - 2` when `n > 8`. The default is calculated when the
+build directory is first configured with this setting. Override it with
+`-DSD_MSVC_MP_COUNT=4`, or pass `-USD_MSVC_MP_COUNT` to recalculate it.
+
+For Visual Studio builds, keep project-level parallelism at one to avoid
+multiplying the per-project limit:
+
+```shell
+cmake --build build --config Release --parallel 1
+```
+
+This setting controls MSVC's `/MP` option, not CUDA compiler parallelism or a
+hard CPU usage limit. Other build tools, such as Ninja, schedule compilation
+jobs separately; use `cmake --build build --parallel N` to limit those jobs.
+
 ## Selecting a GGML source tree
 
 By default, sd.cpp builds the patched GGML submodule in `ggml/`. To build with
@@ -27,8 +51,10 @@ cmake -S . -B build-upstream -DSD_USE_UPSTREAM_GGML=ON -DSD_GGML_SOURCE_DIR=../g
 cmake --build build-upstream --config Release
 ```
 
-The selected source tree supplies both the library and its private headers.
-Backend options such as `-DSD_CUDA=ON` apply to the selected tree as usual.
+Upstream compatibility mode uses the selected library and matching private
+headers. The patched GGML build uses public headers, including `ggml-ext.h` and
+`ggml-graph.h`. Backend options such as `-DSD_CUDA=ON` apply to the selected tree
+as usual.
 
 `SD_USE_UPSTREAM_GGML` defaults to `OFF`, which enables the patched GGML
 extensions. Set it to `ON` when using upstream GGML; it selects the compatibility
@@ -45,10 +71,12 @@ creating an inference context. Ordinary floating-point and shared GGML
 quantization types remain available, subject to backend operator support.
 
 `SD_USE_SYSTEM_GGML=ON` instead links an installed GGML CMake package, located
-with `ggml_DIR` or `CMAKE_PREFIX_PATH`. In that mode, `SD_GGML_SOURCE_DIR` must
-point to the matching source tree for private headers. The installed library
-must use the same ABI settings as sd.cpp, including `GGML_MAX_NAME`.
-Set `SD_USE_UPSTREAM_GGML=ON` as well if the installed package is upstream GGML.
+with `ggml_DIR` or `CMAKE_PREFIX_PATH`. An installed patched package must include
+the extension and graph APIs; its source tree is not required for private
+headers. The installed library must use the same ABI settings as sd.cpp,
+including `GGML_MAX_NAME`. For an upstream package, also set
+`SD_USE_UPSTREAM_GGML=ON` and point `SD_GGML_SOURCE_DIR` to its matching source
+tree for the compatibility layer's private headers.
 
 ## WebP and WebM Support in Examples
 
@@ -90,8 +118,8 @@ cmake --build . --config Release
 
 ## Build with CUDA
 
-Native SageAttention is included when using CUDA with patched GGML
-(`SD_USE_UPSTREAM_GGML=OFF`).
+The sd.cpp SageAttention extension is included when using CUDA with patched
+GGML (`SD_USE_UPSTREAM_GGML=OFF`).
 See [SageAttention](sage_attention.md) for GPU requirements and `--sage-attn` usage.
 
 This provides GPU acceleration using NVIDIA GPU. Make sure to have the CUDA toolkit installed. You can download it from your Linux distro's package manager (e.g. `apt install nvidia-cuda-toolkit`) or from here: [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads). Recommended to have at least 4 GB of VRAM.
@@ -102,6 +130,28 @@ cmake .. -DSD_CUDA=ON
 cmake --build . --config Release
 ```
 
+CUDA extensions use sd.cpp's `thirdparty/cutlass` submodule by default
+for eligible INT8 GEMM and CuTe attention operations. If its headers are missing,
+CMake stops with an initialization hint; it does not download them automatically.
+Run `git submodule update --init --recursive` from the repository root before
+configuring.
+
+To build with the existing CUDA implementations without CUTLASS, configure from
+the repository root with:
+
+```shell
+cmake -S . -B build -DSD_CUDA=ON -DSD_CUDA_CUTLASS_PATH:PATH=
+cmake --build build --config Release
+```
+
+Alternatively, set `SD_CUDA_CUTLASS_PATH` to another local CUTLASS source
+directory (tested with 3.9.2). Existing CMake caches retain their configured path;
+pass `-USD_CUDA_CUTLASS_PATH` when reconfiguring to restore the bundled default.
+Existing `GGML_CUDA_CUTLASS_PATH` caches are migrated on first configuration:
+the old bundled path becomes `thirdparty/cutlass`, while custom paths and an
+empty disabled value are preserved. An explicitly set `SD_CUDA_CUTLASS_PATH`
+takes precedence. Non-CUDA builds do not require CUTLASS.
+
 ## Build with HipBLAS
 
 This provides GPU acceleration using AMD GPU. Make sure to have the ROCm toolkit installed.

docs/docker.md

@@ -18,12 +18,26 @@ docker run --rm --init -v /path/to/models:/models -v /path/to/output/:/output -p
 
 ## Building using Docker
 
+Initialize the submodules from the repository root before building. Docker copies
+the local checkout, so its build context must include the submodule contents:
+
 ```shell
+git submodule update --init --recursive
 docker build -f docker/Dockerfile -t sd .
 ```
 
 ## Building variants using Docker
 
+### CUDA (NVIDIA GPU):
+
+```shell
+docker build -f docker/Dockerfile.cuda -t sd-cuda .
+```
+
+To select GPU architectures, pass `--build-arg "CUDA_ARCHITECTURES=89"` (for
+example, RTX 4090). Run the image on a host configured with NVIDIA Container
+Toolkit, using `docker run --rm --gpus all ... sd-cuda [args...]`.
+
 ### Vulkan:
 
 ```shell
@@ -67,4 +81,4 @@ docker run --rm --device=/dev/dri --device=/dev/kfd --group-add video \
 #   -v ./models:/models -v ./build:/output \
 #   sd -m /models/sd-v1-4.ckpt -v -o /output/output.png \
 #   -p "a lovely cat"
-```
+```

docs/ggml_extensions.md

@@ -0,0 +1,550 @@
+# ggml operator extensions
+
+The patched ggml exposes `ggml-ext.h` for application-owned operators, kernel
+overrides, and graph fusion rules. Registrations are in-process and scoped to a
+backend device. Register them before capability queries or graph allocation;
+callback code and user data must remain alive until all dependent graphs and
+backends have been destroyed. Registration names are unique, and duplicate
+registrations fail rather than replacing live implementations.
+
+## Operator definitions
+
+`ggml_ext_register_op` registers a versioned name, attribute size, and an output
+metadata inference callback. `ggml_ext_op_create` creates an ordinary ggml tensor
+with explicit input dependencies. The initial interface supports one contiguous,
+out-of-place output with an existing `ggml_type`.
+
+Attributes are copied into the unused portion of `op_params`. They are small,
+plain values; use `memcpy` when reading attributes that need stricter alignment.
+Auxiliary tensors must appear in `src`, not in an attribute pointer or user data.
+Registered operators use `GGML_OP_CUSTOM` internally, with an identifiable
+callback marker. Unsupported devices do not mistake them for CPU callbacks.
+These process-local descriptors are not a graph serialization format.
+
+## Kernel selection
+
+`ggml_ext_register_kernel(device, &desc)` associates a kernel with a built-in
+`ggml_op`, or with a registered custom operator. `supports` describes correctness
+constraints; `preferred` optionally decides whether to override an available
+native implementation. Higher priorities are considered first. When native
+execution is unsupported, a supported extension can run regardless of its
+preference callback.
+
+Capability queries include registered single-node implementations, so the
+scheduler can select a device even when its native implementation rejects the
+operator. Callbacks inspect metadata only and must not recursively call the
+public capability query for the same device and operator. A fusion-only
+registration does not claim single-node support.
+
+`ggml_ext_dev_supports_native_op` queries the original device implementation
+without extension dispatch. Fusion policies can use it to select a supported
+extension even when their performance preference normally favors native kernels.
+
+`prepare` creates optional host-side plan state and reports workspace bytes.
+It can run during memory measurement, before tensor addresses are assigned, and
+must not launch work or allocate device buffers. Workspace is allocated by ggml
+before execution and included in scheduler memory estimates and buffer-size
+reports. `execute` receives the backend, workspace buffer, and current tensor
+addresses. `destroy` releases prepared state. A preparation or execution error
+is propagated; execution never retries a native implementation after launch.
+
+Workspace is scratch storage shared by sequential extension steps in a plan;
+its capacity is the maximum request, not the sum across nodes. Contents are not
+preserved between calls. Queue work in backend execution order, and finish any
+private-stream use before subsequent backend work can access the workspace.
+Persistent state belongs in `prepare`/`destroy` state or explicit tensor inputs.
+
+Overriding a built-in operator preserves its numerical and aliasing contract,
+including allocator reuse of input storage. Explicit views/in-place nodes are
+not intercepted by this first interface. Rebuild the scheduler allocation after
+changing graph metadata, selection policy, or registrations.
+
+## Fusion
+
+A `match` callback returns the number of consecutive nodes it can replace.
+The scheduler freezes this selection before allocating tensors. ggml checks
+compute flags, external uses, output flags, and views; the initial implementation
+accepts a conservative set of pure operators with only the final node exposed
+as the fused output. The rule remains responsible for mathematical equivalence,
+dtype, shape, layout, and hardware constraints. Internal views are accepted when
+their storage belongs to the fused subgraph. In-place intermediates are accepted
+only when their storage owner is an earlier node in the same match. External
+consumers and output flags still prevent eliminating those results. The final
+output must own storage, use a terminal `CPY`, or write back to an earlier owning
+node within the match. Internal writeback must cover exactly the owner's full
+shape, type, strides and storage offset, without external consumers of that
+owner. Output flags propagated to that owner are retained. Terminal `CPY` keeps
+its original destination and write semantics. The extension must handle input/output storage overlap, including
+allocator reuse, with scratch space when necessary. Other in-place writes are
+not eligible for extension fusion when they target storage outside the match.
+
+An optional `optimize` callback runs before matching and allocation. Use
+`ggml_ext_graph_delay` to move a producer range after independent nodes. The
+helper rejects dependencies, unknown side effects, and writes that alias moved
+inputs or outputs. It can cross independent copies and known in-place operators
+when their storage does not alias the moved inputs or outputs. Owning dtype
+conversions can also be moved or crossed.
+Reordering must preserve edges and tensor metadata. Fusion legality is checked
+again after reordering.
+
+External inputs stay allocated through the final node. Native graph optimization
+and execution operate on the remaining contiguous ranges, so native fusion cannot
+consume nodes selected by an extension. Fused intermediates are excluded from
+allocation through the backend's `skip_alloc` callback and
+`ggml_gallocr_set_skipped_tensors`. Their metadata and graph edges remain available,
+but execution must not access their data or buffers. Storage backing the final
+output is retained, including an internal owner used by an in-place final write.
+The skipped set is replaced for each allocation graph; changing which tensors are skipped
+invalidates cached allocation assignments even when tensor shapes are unchanged.
+
+Automatic extension fusion is enabled in scheduler execution with one pipeline
+copy and without evaluation callbacks. Multiple pipeline copies and evaluation
+callbacks retain individual operations. Direct backend graph computation supports
+single-node extensions; it does not fuse an already allocated graph. Direct calls
+use temporary plans and synchronize before releasing them; scheduler plans and
+workspaces are reused across executions.
+
+Use the backend graph APIs or the scheduler. The legacy CPU `ggml_graph_compute`
+and backend `graph_plan_*` interfaces do not dispatch these extensions.
+
+## Backend services
+
+Packed ComfyUI INT4 linear layers use `sd.linear_w4_convrot.v1`, built by
+`ggml_ext_mul_mat_w4_convrot`. The operator consumes byte-packed I8 weights and the
+explicit output of `sd.quantize_i8_convrot.v1`; F32 callers get a quantization
+node during graph construction. Logical output dimensions are stored in the
+operator attributes, so backend copies of the packed activation do not need to
+retain the producer's op marker. W4A4 uses row scales; W4A8 additionally exposes
+the codebook, channel scales, and FP8 relative scales as tensor dependencies.
+CPU, CUDA/HIP/MUSA, and Vulkan register kernels through the same extension API
+as INT8. No ggml submodule changes are required.
+
+Services are discovered through `ggml_backend_reg_get_proc_address`:
+
+| Service | Purpose |
+| --- | --- |
+| `ggml_ext_cpu_parallel_v1` | Run an application callback on the CPU backend's configured thread pool, with its abort callback. |
+| `ggml_ext_cuda_context_v1` | Borrow the current CUDA device and stream, ordered with native work. |
+| `ggml_ext_cuda_device_v1` | Get the physical CUDA ordinal for capability checks before backend creation, including virtual ggml devices. |
+| `ggml_ext_cuda_blas_v1` | Borrow the current stream's cuBLAS handle during execution, with host scalar pointers. |
+| `ggml_ext_gpu_device_v1` | Query CUDA/HIP/MUSA runtime kind, physical ordinal, shared-memory limit, warp size, and INT8 BLAS capability. |
+| `ggml_ext_gpu_context_v1` | Borrow the matching runtime's device and execution stream. |
+| `ggml_ext_gpu_blas_v1` | Borrow the matching cuBLAS, hipBLAS, or muBLAS handle on that stream. |
+| `ggml_ext_metal_device_v1` | Borrow the Metal device for application pipeline compilation. |
+| `ggml_ext_metal_buffer_v1` | Resolve a backend buffer range to its Metal buffer and offset, including private storage and split allocations. |
+| `ggml_ext_metal_submit_v1` | Encode application commands on a backend-owned command buffer, ordered with native Metal work and tracked by backend synchronization. |
+
+The borrowed cuBLAS handle remains owned by the backend. Extensions must not retain
+or destroy it, or change its stream, pointer mode, math mode, or workspace. This
+CUDA-specific BLAS service is retained for existing callers; the runtime-aware
+services support CUDA, HIP, and MUSA with the same ownership rules. Extensions
+call vendor runtime and BLAS APIs directly. These services do not execute GEMM,
+quantization, or model fusion on the application's behalf.
+
+Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; callers
+must use the representation appropriate to their backend. Vulkan extensions use
+the command recording services described below.
+
+CUDA INT8 Linear, RoPE, ConvRot, RMS affine and INT8 gated residual kernels omit output
+scratch when allocation guarantees that a fresh owning output cannot overlap
+live inputs. Existing outputs and views retain overlap protection. Internal
+GEMM and quantization workspace is still reserved when required; kernels also
+handle a null workspace when the entire step needs no scratch.
+
+Compute workspace measurement preserves the storage identities and offsets of
+external tensors, including views. Separate allocations must not share a synthetic
+buffer identity: that can invent overlap, overestimate scratch, and repeatedly
+rebuild a workspace whose actual allocation is smaller than the estimate.
+
+## sd.cpp integration
+
+`src/core/ggml_extensions.cpp` and `.h` contain the shared operator schemas,
+validation, and backend registration entry point. Graph construction helpers live
+in `src/core/ggml_extend.cpp` and `.h`; they take `GGMLRunnerContext*` so backend
+selection, compatibility workarounds, and shared ConvRot nodes use the same
+context as the calling model. Per-call precision, scaling, attention, and padding
+options remain explicit. Model blocks retain weight loading and adapter
+orchestration, while the shared helpers own FP8 fallback, embedding lookup,
+full/partial RoPE, grouped convolution, and backend-dependent repeat handling.
+LoRA/LoKr construction, n-mode products, and Kronecker products stay in
+`src/model/adapter/lora_ops.cpp` and `.h`, with the same runner-context interface.
+CPU kernels and fusion rules live in `src/core/cpu/`, with registration collected
+in `cpu/registration.cpp`. CUDA-style and Vulkan implementations live in their
+respective backend directories.
+
+`src/core/metal/int8.mm` registers Metal INT8 linear and ConvRot kernels. Their
+Metal source is embedded at build time and compiled for the selected device
+during registration. The public `ggml-ext-metal.h` services expose borrowed
+objects without private backend headers. The submit callback only encodes work;
+GGML owns submission, cancellation checks, command-buffer lifetime, and
+asynchronous error tracking. Rotation scratch and any overlapping output use
+the scheduler workspace. No full floating-point weight copy is allocated.
+When runtime compilation supports Metal 4 TensorOps, INT8 multiplication uses
+`matmul2d` with INT32 cooperative results and fused scale/bias restoration.
+The SIMD-group implementation remains the fallback for older systems and
+devices where the TensorOps pipeline is unavailable.
+TensorOps multiplies at most 1024 pairs per partial result and combines those
+results with wrapping integer addition. This preserves exact INT32 semantics
+even when TensorOps uses floating-point arithmetic internally.
+
+`ggml_ext_modulate(ctx, x, shift, scale)` emits `sd.modulate.v1` for F32 inputs
+when the runner context uses CPU, CUDA/HIP/MUSA, or Vulkan.
+Its explicit inputs are `x`, `scale`, and optional `shift`; it computes
+`(x + x * scale) + shift` with separate FP32 multiply/add rounding. The output
+is contiguous and out of place. Inputs may be strided and use ggml's
+dimension-wise repetition rules. Callers prepare token, batch, or frame layouts
+before constructing the operation; normalization and residual gates remain
+separate operations. CPU, CUDA/HIP/MUSA, and Vulkan have standalone kernels.
+Other backends (including Metal, SYCL, and OpenCL), upstream ggml builds, and
+other input types retain the original MUL/ADD graph.
+
+ConvRot fusion matches the Modulate node directly, including scale-only
+modulation, instead of recognizing its former MUL/ADD/ADD decomposition.
+CUDA accepts contiguous per-channel parameters for `Modulate -> ConvRot` and
+`RMSNorm -> Mul(weight) -> Modulate -> ConvRot`. Vulkan limits the first fusion
+to widths 4096/6144/5376 and the RMSNorm fusion to width 5376.
+Other layouts execute the standalone operation followed by ConvRot. The
+gated-residual MUL/ADD fusion remains available for residual branches.
+
+`src/core/cpu/rope.cpp` registers the CPU implementation of
+`sd.rope_apply.v1` and a
+`RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
+vector. Both interleaved and split-half RoPE, partial rotation, and strided input
+rows are supported. Model construction uses the custom operator on CPU, CUDA,
+HIP, MUSA, and Vulkan. Other backends can use the scheduler's CPU fallback.
+The former `ggml_rope_apply` API, builtin opcode, and CPU implementation have
+been removed from ggml.
+
+`src/core/cuda/rope.cu` supplies standalone CUDA RoPE
+and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares
+its RMSNorm/RoPE kernel. Partial rotation and arbitrary float-aligned input and
+position-encoding strides are supported by the standalone kernel; RMSNorm fusion
+requires contiguous position encoding and weights. Copy fusion uses workspace
+when the destination overlaps an input. These registrations do not require
+CUTLASS. The old CUDA/HIP/MUSA RoPE dispatch is removed.
+
+Shared CUDA-style kernels live in `src/core/cuda/` and are compiled separately
+for each enabled runtime. `SD_CUDA_NAMESPACE` expands to `sd_cuda`, `sd_hip`, or
+`sd_musa`, so implementations can coexist. `runtime.h` handles the small runtime
+differences, including logical 32-lane shuffles on HIP. The application defines
+and launches its kernels directly; it includes only public ggml headers.
+CMake generates the registration table from enabled backend options. The shared
+entry point uses one `SD_BACKEND_EXTENSIONS` macro; backend registrars verify the
+device kind, and disabled backends add no symbol references.
+
+`SD_USE_UPSTREAM_GGML` excludes these registrations. The stable-diffusion public C
+API is unchanged. INT8 ConvRot retains its existing packed activation layout and
+separate weight scale tensor.
+
+`src/core/cuda/cute_attention.cu` supplies CuTe Flash Attention on CUDA when
+CUTLASS headers are configured. It registers an override of `FLASH_ATTN_EXT`,
+with the existing sequence-length preference and native fallback. It also
+registers Q preparation (`RMSNorm -> Mul -> RoPE`), output scaling, and packed
+ConvRot output fusion. These kernels use only the public extension API, borrowed
+CUDA stream, and scheduler-owned workspace; they do not use ggml's CUDA pool or
+private conversion helpers. The existing CUDA fusion and CuTe disable environment
+variables are honored.
+Aligned contiguous FP16 K/V inputs are read directly without reserving conversion
+buffers. Fresh F32 outputs are written by the attention epilogue, retaining the
+original FP16 rounding before output scaling. Packed ConvRot and preallocated or
+view outputs retain the intermediate half buffer.
+
+`src/core/cuda/i8_projection.cu` supplies the INT8 GEMM projection epilogues:
+`GEMM -> Reshape -> RMSNorm -> Mul -> RoPE`, optional scaling and FP16 copy,
+and K/V prefix concatenation with FP16 output. CuTe can consume a directly
+produced FP16 Q without materializing the FP32 GEMM or RoPE intermediates.
+Independent affine operands, including broadcasted constants, are computed before
+their normalization/projection
+chain when graph dependencies allow it. The attention optimizer can then move
+the complete five-node projection when safe.
+These rules accept custom RoPE nodes. They preserve the original
+large-matrix selection thresholds, full interleaved head dimension 128, packed
+activation layout, and separate weight scales. External consumers prevent fusion;
+aliasing terminal copies use scheduler workspace. The CUDA projection kernels
+and matchers now live in sd.cpp. `GGML_CUDA_FORCE_CUBLAS` disables these epilogues.
+
+`src/core/ggml_extensions.cpp` defines `sd.linear_i8.v1` and
+`sd.quantize_i8_convrot.v1`. `src/core/cpu/int8.cpp` supplies their CPU
+implementations. The CPU kernels use the
+configured thread pool and scheduler workspace; the old CPU INT8 GEMM and ConvRot
+implementations have been removed from ggml.
+Model code keeps calling `ggml_ext_linear_i8_tensorwise`,
+`ggml_ext_mul_mat_i8_tensorwise`, and `ggml_ext_quantize_i8_convrot`. The latter two
+constructors are declared and implemented in sd.cpp and create custom nodes.
+Their former ggml declarations, implementations, ConvRot opcode, and native
+backend paths have been removed. CPU, CUDA/HIP/MUSA, and Vulkan register kernels
+and fusion rules through the extension interfaces.
+
+For already quantized activations, `ggml_ext_mul_mat_i8_tensorwise` accepts ordinary
+I8 weight and activation tensors. Its optional final `input_scale` parameter is a
+separate F32 tensor with one value per activation row. Weight scales have one
+value per output channel; weight scales and bias can be null. Existing calls
+omit `input_scale` and retain their original argument order. It computes `float(int32_dot(W, X)) * X_scale * W_scale + bias`, in that
+order. All scales are explicit graph dependencies; there is no quantized-tensor
+wrapper. Shapes use ggml's low-to-high dimension order. Call
+`sd_ggml_register_extensions` before scheduling these custom nodes.
+
+`ggml_ext_mul_mat_i8_tensorwise` also accepts F32 activations (quantized into temporary
+workspace) or the output of `ggml_ext_quantize_i8_convrot`. ConvRot keeps the existing
+packed I8 activation rows, padded to a multiple of four rows, followed by F32 row
+scales. Weight scales remain separate. CPU ConvRot groups are 1, 4, 16, 64, or 256.
+Unsupported layouts and malformed scale shapes are rejected. Separate CPU quantize
+and GEMM phases propagate cancellation, and overlapping outputs use workspace.
+
+`src/core/cuda/int8.cu` implements the same custom schemas on CUDA. It uses the
+borrowed cuBLAS handle for I8-by-I8-to-I32 GEMM and application-owned quantization,
+ConvRot, and rescaling kernels. F32 inputs are quantized into workspace; ordinary
+I8 inputs retain separate scale tensors and use a padded temporary only when
+needed by cuBLAS. Packed ConvRot inputs are consumed directly. Large ConvRot rows
+use multiple kernels when they exceed shared memory. Aliasing outputs use scratch.
+CUDA GEMM requires Turing or newer, contiguous tensors, K/N divisible by four,
+float-aligned inputs, and row counts that fit a padded signed 32-bit dimension;
+CUDA ConvRot uses groups of 256.
+Unsupported custom nodes can fall back to CPU through the scheduler.
+
+When CUTLASS is configured, `src/core/cuda/i8_gemm.cu` supplies dense and SwiGLU
+INT8 GEMM kernels on Ampere or newer. It preserves the native tile selection,
+small-row policy, Stream-K scheduling, Ada column bands, and separate scaling
+order. The same extension matches `GEMM -> SwiGLU [-> ConvRot]` and row
+concatenation for custom GEMM nodes. Weight scales remain explicit
+inputs, including the gate/up scale permutation recorded on every SwiGLU launch.
+Two independent INT8 GEMMs sharing a quantized input can also fuse with split
+SwiGLU. A paired-column iterator reads the separate gate/up weight allocations,
+and the GEMM epilogue writes only the activated output. Both intermediate
+activations are omitted through the generic fusion allocation interface; weights
+are neither copied nor concatenated. This path requires bias-free branches,
+compatible layouts and alignments, and no external intermediate consumers.
+Only interleaved weight scales and any Stream-K workspace remain temporary.
+Plans reserve bounded Stream-K workspace, cache occupancy on first execution,
+and initialize barriers on every launch, so CUDA Graph replay observes updated
+inputs and scales. External consumers retain ordinary execution; overlapping
+outputs use workspace. For fresh outputs, the CUTLASS SwiGLU epilogue writes
+an allocator-managed tensor that the subsequent ConvRot launch consumes. This
+keeps the same GPU launches while allowing its storage to be reused by other
+operations. The combined SwiGLU/ConvRot path retains its workspace intermediate
+for other cases. Shapes or alignments unsuitable for CUTLASS use cuBLAS.
+Custom I8 nodes also retain cuBLAS when the reduction could overflow I32, to
+preserve their previous accumulation behavior. Fusion disable and force-cuBLAS
+settings are honored. No CUTLASS headers are needed for the cuBLAS path.
+
+The native CUDA/HIP INT8 GEMM, standalone ConvRot, and associated fusion paths
+have been removed from ggml. The shared application implementation calls cuBLAS
+or hipBLAS directly. HIP retains the existing architecture eligibility policy;
+MUSA does not advertise INT8 BLAS support and uses scheduler fallback for GEMM.
+Standalone ConvRot kernels are compiled for all three runtimes. CUTLASS
+implementations remain CUDA-specific.
+
+`src/core/cuda/i8_gated_residual.cu` implements
+`INT8 GEMM -> [Scale] -> Mul(gate) -> Add(residual) -> [Concat]` as a registered
+CUDA fusion, for packed ConvRot inputs and ordinary I8 nodes with separate activation
+scales. The corresponding native CUDA fusion rules and CUTLASS residual epilogue
+have been removed from ggml. Its optimizer moves independent gate producers ahead
+of GEMM and optional scaling through `ggml_ext_graph_delay`, before fusion matching.
+The former native INT8 affine reorder rule has also moved into this extension.
+It preserves the original tile selection, small-row
+threshold, Stream-K scheduling, Ada column bands, and separate FP32 rounding of
+gate multiplication and residual addition. The plan bounds Stream-K workspace
+over possible occupancy values using host metadata; execution verifies the
+actual requirement against that bound. Device occupancy is cached on first
+execution; each launch rebuilds parameters for the current tensor addresses and
+initializes its Stream-K barriers. Overlapping outputs use scratch, except
+for an exact in-place residual without concatenation. External consumers and
+unsupported layouts retain ordinary operator execution. Fusion and force-cuBLAS
+disable settings are honored.
+
+`src/core/cuda/convrot.cu` registers CUDA fusion for `Scale -> ConvRot`,
+`Norm -> Mul -> ConvRot`, `Modulate -> ConvRot`,
+`RMSNorm -> Mul(weight) -> Modulate -> ConvRot`, `SwiGLU -> ConvRot`, and
+`SiLU/Sigmoid -> Mul -> ConvRot`.
+ConvRot outputs are custom nodes. SwiGLU accepts packed,
+swapped, or separate gate/up inputs. Independent modulation and up-branch
+producers are moved ahead of their consumers through the safe graph-delay API.
+Reordering keeps an owning producer together with in-place consumers, such as
+Krea2's RMSNorm weight multiplication, so the alias checks can preserve their
+storage relationship.
+The application kernel preserves the original reduction order, thread selection,
+ConvRot row padding, and separate scale tail. Prepared plans cache operation
+metadata, use the borrowed CUDA stream, and reserve scheduler workspace for
+overlapping outputs. Long rows or external intermediate consumers retain
+ordinary execution. These fusions do not require CUTLASS and honor
+`GGML_CUDA_DISABLE_FUSION`.
+
+The corresponding scale, normalization, and unary-activation fusion rules have
+been removed from ggml for CUDA and HIP. GEMM/SwiGLU/ConvRot chains are selected by the
+application INT8 extension before the standalone activation fusion.
+
+`src/core/cuda/affine.cu` registers FP32 `Mul(gate) -> Add(residual)`,
+`Concat -> [Permute] -> FP16 Copy`, and `Contiguous -> Scale -> Permute -> FP16
+Copy` fusions. Gate multiplication and residual addition retain separate FP32
+rounding. Copy layouts and alignment requirements match the original CUDA
+kernels; unsupported layouts, external intermediate consumers, self-residual
+Add inputs, and fusion disabling retain ordinary operators. Internal in-place
+Add outputs are accepted when they overwrite the complete fused Mul allocation. Exact in-place residual
+or input addresses are safe for the pointwise kernel; partial overlap, broadcast
+gate overlap, and FP16 copy overlap use scheduler workspace. These registrations
+require no CUTLASS, compile for CUDA/HIP/MUSA, and support reusable CUDA Graph
+replay. Their former native fusion implementations have been removed from ggml.
+
+CUDA kernels can opt into device replay with `ggml_ext_register_kernel_ex` and
+`GGML_EXT_KERNEL_CAPTURE_SAFE`. The original registration function retains its
+per-execution callback behavior, and the descriptor ABI is unchanged. Opting in
+means the callback enqueues replayable device work: it has no required
+per-execution host side effects or decisions based on tensor contents, and any
+private streams rejoin the borrowed stream. Lazy setup must complete during the
+first ordinary execution. Tensor metadata, attributes, and host plan parameters
+remain frozen until scheduler reallocation; tensor contents may change normally.
+
+When every selected extension opts in, the CUDA backend can capture the whole
+scheduled plan, including its native fragments, into one CUDA Graph. The first
+execution warms up kernels and library state while retaining native fragment
+caches. The next stable execution captures
+and launches the graph; later calls replay it. Changing a tensor address or the
+execution stream invalidates the capture and repeats warmup, even when the graph
+UID is unchanged. Capture failures propagate without retrying partially executed
+work. An enclosing stream capture records the plan directly instead of starting
+a nested capture.
+
+`GGML_EXT_KERNEL_REPLAY_REUSABLE` additionally permits reuse across scheduler
+plan rebuilds and requires `CAPTURE_SAFE`. Recorded work must depend only on
+complete tensor metadata, addresses, the device and immutable registration data.
+Captured resources must belong to the backend or the supplied workspace, never
+per-plan host state. Workspace contents must be initialized by recorded device
+work on each execution. Prepare/destroy still run for every rebuilt host plan.
+The backend cache retains only the executor and workspace; it compares graph
+structure, tensor metadata/addresses, selected kernels and workspace requirements
+exactly, without dereferencing old graph objects. A new metadata arena can reuse
+an equivalent capture. Each backend retains at most 16 retired plans and 256 MiB
+of their workspaces, evicting the oldest entries first. Larger workspaces retain
+plan-local capture only. Backend destruction synchronizes and releases the cache.
+This allows sd.cpp's per-step scheduler rebuilds to retain device replay.
+
+Mixed plans containing an extension without the flag retain ordinary dispatch.
+Evaluation callbacks, multiple pipeline copies, and direct backend graph calls
+also retain their previous behavior. `GGML_CUDA_DISABLE_GRAPHS` disables whole-plan
+capture as well as native CUDA Graph caching. Native operations that cannot be
+captured retain the existing fragment execution path. Whole-plan replay requires
+CUDA VMM so native scratch addresses remain valid for the backend lifetime;
+legacy scratch pools retain the existing fragment path. HIP and MUSA do not yet
+provide this plan executor. sd.cpp's CUDA registrations opt in to reusable replay.
+
+## Vulkan execution services
+
+`ggml-ext-vulkan.h` exposes `ggml_ext_vulkan_v1` through the Vulkan backend
+registry. The versioned function table reports device limits and subgroup
+capabilities, resolves tensor/workspace byte ranges, dispatches application
+SPIR-V, and copies nonoverlapping buffer ranges. Extensions do not include
+private ggml Vulkan headers or own backend command buffers and descriptors.
+
+Shaders use the `main` entry point, consecutive storage-buffer bindings starting
+at zero, and consecutive 32-bit specialization constants. Dispatch accepts actual
+workgroup counts. The backend validates ranges and dispatch limits, caches the
+complete shader/configuration per device, and releases pipelines with that
+device. Push constants and specialization values are copied; the backend does
+not retain application execution state. Pipeline creation occurs on execution,
+never during host-only plan preparation.
+
+Bindings identify either a tensor or a backend buffer, with a byte offset and
+optional byte count. `buffer_info` returns an opaque storage identity and byte
+range for alias checks. Vulkan descriptor offsets are rounded down to the
+device's required alignment; shaders must add the returned `shader_offset`.
+Workspace bindings follow the same rules. Invalid or foreign-device ranges are
+rejected before recording commands. Kernel bounds and shader capability
+requirements remain the extension's responsibility.
+
+Dispatch and copy use the backend's existing compute command context, descriptor
+pool, transfer waits, and memory barriers. They preserve ordering with native
+fragments and invalidate cached converted operands after writes. Normal backend
+synchronization completes queued work before host access or resource release.
+Vulkan copy ranges must be equal in size, nonoverlapping, and aligned to four
+bytes. There is no per-dispatch host wait or extension-owned queue.
+
+Scheduled Vulkan plans share submission counters and FLOP thresholds across
+native fragments and extension callbacks. Native scratch is initialized once
+per plan invocation, with its offsets retained across fragments. Fragment
+boundaries do not force a
+queue submission, and the final pending commands are submitted when the plan
+returns. Extension dispatches and copies count toward the device's node limit.
+Every callback still executes on every invocation; this batching neither
+captures device replay nor retains plan state across scheduler rebuilds.
+Callback errors propagate without retry. Evaluation callbacks, direct backend
+calls, synchronous devices, serialized diagnostics, and timestamp profiling
+retain fragment execution.
+
+`src/core/vulkan/rope.cpp` registers the custom RoPE implementation,
+plus `RMSNorm -> Mul -> RoPE` and `RMSNorm -> Mul -> RoPE -> Scale -> FP16 Copy`
+fusion. These kernels retain the original strides, partial-rotation behavior,
+128-channel subgroup reduction, and separate arithmetic rounding. Their SPIR-V
+is generated from application shaders with the native RoPE compilation policy;
+SPIR-V optimization is disabled because it changes the reduction results.
+Safe graph-delay rules bring independent producers ahead of fusion candidates.
+External intermediate consumers retain ordinary operators, and overlapping
+outputs use plan workspace. `GGML_VK_DISABLE_FUSION` is honored.
+
+The corresponding Vulkan shader sources and native fusion rules have been
+removed from ggml. Unsupported custom nodes retain scheduler fallback to
+another supported backend.
+
+RoPE Apply, INT8/ConvRot, CuTe attention, SageAttention, and Sol-Attn use
+application-owned implementations. This interface does not add dynamic dtypes.
+
+`src/core/vulkan/int8.cpp` owns Vulkan INT8 GEMM, ConvRot, and their fusion rules.
+The Vulkan compute metadata service exposes shared-memory limits, subgroup
+arithmetic, and cooperative-matrix properties; shape selection stays in sd.cpp.
+Shader variants preserve the former tiling and compilation policy, including
+SwiGLU, residual, row-concatenation, and modulation paths. Kernels use public
+Vulkan dispatch, buffer, copy, and workspace services without private backend
+headers. Vulkan currently handles packed group-256 activations with separate
+weight scales; unsupported custom nodes fall back to CPU.
+
+## SageAttention
+
+The `ggml_ext_sage_attn` entry point is declared and implemented in sd.cpp's
+graph construction helpers. It constructs `sd.sage_attn.v1` with explicit Q/K/V inputs
+and scale/mode attributes. ggml has no Sage-specific opcode, allocation rule, or
+backend dispatch branch.
+
+`src/core/cuda/sage_attention.cu` and its `sageattention/` device headers own the
+CUDA implementation. Quantization and temporary attention output use scheduler
+workspace, rather than an oversized output tensor. The implementation obtains
+the backend stream through the public GPU service and converts the final FP16
+output with a local CUDA kernel. It retains the SM80 compatibility path and the
+SageAttention2/2++ toolkit and architecture requirements. Capability checks use
+the application's compiled CUDA target as well as the physical device.
+
+CPU, HIP, MUSA, and Vulkan do not implement this NVIDIA-specific kernel. Existing
+model-level capability checks retain flash/default attention for unsupported
+layers; a custom Sage node itself does not silently change algorithms.
+
+## Sol-Attn
+
+The `ggml_ext_sol_attn` entry point lives in the common sd.cpp graph construction
+files and constructs `sd.sol_attn.v1`. It accepts contiguous F32 Q/K/V tensors
+with shape `[128, tokens, heads, batch]`, a positive finite scale, and finite tau.
+The output remains F32 `[128, heads, tokens, batch]`. No current sd.cpp model or
+CLI path calls this entry point; the migration preserves its low-level behavior.
+
+`src/core/cuda/sol_attention.cu` owns packing, block summaries, threshold
+selection, and the BF16 tensor-core kernel. Temporary storage is separate
+scheduler workspace, so the output allocation contains only the output tensor.
+The kernel uses the public backend stream service and direct CUDA/PTX operations.
+Registration requires CUDA 12 or newer, an SM80-or-newer device, and an application
+kernel compiled for SM80 or newer. CPU, HIP, MUSA, and Vulkan do not advertise
+support. ggml no longer contains a Sol-Attn opcode, constructor, allocation rule,
+or backend kernel.
+
+## Public graph access
+
+`ggml-graph.h` exposes leaf enumeration, append-only leaf construction, recorded
+use counts, mapped inference-graph copying, use-count copying for segments, and
+backend execution of a node range. Mapped copies retain node order and consumer
+counts from the original graph, including consumers outside a graph view. The
+caller owns tensor copies and remaps their storage and edges; invalid mappings
+leave the destination graph unchanged. Range execution uses already allocated
+tensors and retains the original graph's consumer counts.
+
+Workspace measurement, graph cuts, and evaluation callbacks use these APIs.
+Normal sd.cpp builds require only ggml's public headers. Compatibility with an
+upstream ggml lacking this header is isolated in `src/core/ggml_compat.cpp`; that
+optional build still requires matching upstream private headers. The compatibility
+layer also retains the workaround for old Metal backends wrapping failed host
+mappings. The patched Metal backend returns null directly when mapping fails.

docs/int8_convrot.md

@@ -2,7 +2,11 @@
 
 sd.cpp can load and execute ComfyUI `int8_tensorwise` safetensors with `convrot` metadata directly. The stored INT8 weights are not converted to another weight type at load time.
 
-This requires the INT8 tensorwise/convrot extensions in the patched GGML.
+This uses sd.cpp's INT8 tensorwise/convrot operators through the patched GGML
+extension API. Schemas and graph builders live in sd.cpp, alongside the
+backend kernels and fusion rules; GGML supplies generic dispatch and workspace
+services. See [operator extensions](ggml_extensions.md) for the implementation
+and the preserved INT8 API entry points.
 Builds with `SD_USE_UPSTREAM_GGML=ON` reject these files during loading.
 
 ## Checkpoint format
@@ -89,13 +93,76 @@ The packed runtime activation tensor contains the I8 activation rows and their f
 ## Backend support
 
 - CPU provides the portable regular Hadamard, activation quantization, INT8 matrix multiplication, and scale restoration implementations.
-- NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. It uses cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition.
+- Metal supports INT8 tensorwise and convrot on Apple silicon (Apple GPU family 7 or newer). It keeps weights in INT8 and rotates and quantizes activations on the GPU. On macOS/iOS 26 or newer, it probes Metal TensorOps `matmul2d` with INT8 inputs and INT32 cooperative results. It combines 1024-product partial sums using wrapping integer addition to preserve exact accumulation on older GPUs, then applies scales and bias in registers. Older systems or unavailable TensorOps pipelines use SIMD-group matrix multiplication: INT8 tiles are converted to FP16, and FP32 partial sums cover at most 256 products before being combined as integers, preserving the I32 accumulation contract. Convrot groups 1, 4, 16, 64, and 256 are supported. These kernels require the patched GGML Metal extension services and are enabled by `SD_METAL=ON`.
+- NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. Eligible shapes use the CUTLASS INT8 implementation when enabled; other supported shapes use cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition.
 - AMD HIP devices in the CDNA, RDNA3 (including RDNA3.5), and RDNA4 families use the same INT8/H256 kernels with hipBLAS for I8 x I8 to I32 GEMM. Other AMD architectures fall back to CPU for INT8 matrix multiplication.
-- Vulkan provides native H256 activation quantization and INT8 matrix multiplication when the build and device support accelerated packed INT8 dot products. Unsupported configurations and GPU backends without these kernels use the backend scheduler to fall back to CPU.
+- MUSA shares the CUDA-style H256 activation kernels. INT8 GEMM falls back to CPU because the MUSA backend does not advertise INT8 BLAS support.
+- Vulkan provides H256 activation quantization and INT8 matrix multiplication when the build and device support accelerated packed INT8 dot products. Unsupported configurations and GPU backends without these kernels use the backend scheduler to fall back to CPU.
 
 LoRA adapters are applied at runtime without modifying the INT8 weights. The INT8 convrot path computes the base linear output, while LoRA, LoHa, LoKr, and raw weight-difference adapters compute their output corrections from the original, unrotated activation and add them to the base output. `--lora-apply-mode auto` selects this path for models containing INT8 tensorwise weights. If `immediately` is requested, sd.cpp falls back to runtime application because merging an adapter would require dequantizing and rotating its weight update, then recalculating the per-row scales and requantizing the result.
 
-The dedicated CUDA and HIP convrot activation paths currently require a group size of `256`; other supported group sizes use CPU execution.
+The shared CUDA/HIP/MUSA convrot activation paths support H64 and H256. Vulkan also supports H64 activation quantization; its INT8 matrix multiplication path currently requires H256.
+
+## Packed INT4 convrot
+
+ComfyUI `convrot_w4a4` and `asym_w4a8_int8` safetensors load directly, including
+checkpoints that mix packed INT4, INT8 tensorwise, and floating-point layers.
+Packed weights retain their checkpoint byte layout; they are not expanded or
+converted to GGUF Q4 blocks at load time.
+
+- `convrot_w4a4`: signed 4-bit weights, low nibble first, with F32 per-output
+  `weight_scale`. Activations use the same INT8 convrot quantization as above;
+  the format name does not select a native INT4 activation kernel.
+- `asym_w4a8_int8`: 4-bit indices into `weight_codebook` (16 F32 values), F32
+  `weight_s_channel`, and F8_E4M3 `weight_s_rel` per 16 input features. The
+  metadata must specify `group_size: 16`. CPU evaluates the floating-point
+  codebook; GPU kernels requantize it to INT8 with `amax / 127`, introducing a
+  small additional rounding error.
+
+Both formats imply convrot rotation even without a `convrot` flag; missing
+`convrot_groupsize` defaults to 256.
+
+Both formats reuse the convrot activation cache when the input scale is `1` and
+produce F32 outputs with optional bias. CPU provides a portable implementation.
+CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS.
+With CUTLASS enabled, NVIDIA
+SM80 and newer devices use mixed-input S8 x S4 GEMM for eligible W4A4 shapes,
+reusing the INT8 scale epilogue without expanding the full weight tensor. W4A8
+uses CUTLASS integer MMA per 16-value group with FP32 scale restoration.
+Other CUDA shapes and builds retain the packed byte-dot path. Set
+`GGML_CUDA_W4_FORCE_DP4A=1` to compare against that path. Supported AMD targets
+use packed byte-dot instructions, with portable byte arithmetic for other targets.
+Vulkan requires integer dot product support and at least 13312 bytes of shared
+memory for W4A4 or 16452 bytes for W4A8. GPU W4 kernels require the logical input
+width to be divisible by 64; H64 and H256 are supported. When supported by the
+compiler and device, Vulkan uses NV cooperative matrix integer GEMM. W4A4
+decodes packed weights directly into cooperative matrices for aligned input
+and output widths. Complete activation tiles and the final partial tile use
+separate shader specializations. Other W4A4 shapes use double-buffered weight
+and activation tiles (20480 bytes of shared memory, or 30720 for larger tiles).
+W4A8 stages both weights and activations in shared memory in K=64 tiles,
+reusing them for each 16-value group's integer
+MMA and FP32 scale restoration. FP8 scales use a shared lookup table, and
+expanding projections visit activation tiles in groups of weight panels for
+cache reuse. Adjacent groups share a K=32 activation fragment, with
+complementary zero padding in the weights to preserve each group's scale.
+This path uses 17476 bytes of shared memory, or 22596 for larger tiles.
+Both paths use 256-thread workgroups and keep the model's
+weights packed; neither expands a full layer in device memory. Other Vulkan
+devices retain the packed dot-product path; set `GGML_VK_W4_FORCE_DP4A=1` to compare
+against it. FP32 accumulation order can differ between these paths.
+Other shapes use CPU fallback. HIP W4A8 selects wave64 WMMA by the active
+compilation target's intrinsic support, preferring the GFX12 fragment layout
+over the GFX11 layout.
+There is no GPU architecture allowlist: newer targets supporting either
+intrinsic automatically use its matching layout. Selection happens per device
+image, including multi-architecture builds; targets supporting neither intrinsic
+retain the integer-dot fallback. W4A4 retains integer dot products on HIP.
+
+LoRA application remains at runtime, using the original unrotated activations
+and logical weight dimensions. Packed weights and their codebook/scale tensors
+are excluded from weight type conversion. `SD_USE_UPSTREAM_GGML` builds reject
+these formats because they require the application extension API.
 
 ## Example
 

docs/minimax_h3.md

@@ -99,10 +99,88 @@ The C API exposes the same inputs through `ref_images`, `ref_videos`, and
 `ref_audios` in `sd_vid_gen_params_t`. Each `sd_ref_video_t` supplies its own
 frame rate and optional soundtrack; non-24-fps inputs are resampled internally.
 
+## Latent hires second pass (two-stage sampling)
+
+MiniMax-H3 supports `--hires` with two interchangeable second-stage upscalers:
+
+- an interpolating latent upscaler (`Latent`, `Latent (nearest)`,
+  `Latent (bicubic)`, ...): after the first pass finishes, the packed AV latent
+  is resampled in latent space to `--hires-width` x `--hires-height` and
+  refined by a second sampling pass;
+- a model-backed upscaler (`--hires-upscaler <name>` + `--hires-upscalers-dir`):
+  a H3-shaped learned latent upscaler checkpoint runs between the passes,
+  like the LTX spatial upsampler. This is the sd.cpp equivalent of the
+  two-sampler workflows used for H3 in node-based frontends (e.g. the
+  community 3D-conv latent upscaler from
+  [LBH-123-AI/Minimax_h3_latent_Upscaler](https://huggingface.co/LBH-123-AI/Minimax_h3_latent_Upscaler)).
+
+Draft at a low megapixel count (e.g. 0.4 MP) and refine at a higher one
+(e.g. 0.8 MP) instead of paying for the full resolution in every step:
+
+```sh
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model  ..\models\diffusion_models\minimax_h3_fl2va-Q4_K_M.gguf --vae ..\models\vae\minimax_h3_video_vae_fp16.safetensors --audio-vae ..\models\vae\minimax_h3_audio_vae_fp32.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "a lovely cat" --cfg-scale 1.0 -v -W 864 -H 480 --diffusion-fa --offload-to-cpu --rng cpu --video-frames 56 --hires --hires-upscaler "Latent (bicubic)" --hires-width 1216 --hires-height 672 --hires-steps 4 --hires-denoising-strength 0.4
+```
+
+or with the learned upscaler:
+
+```sh
+.\bin\Release\sd-cli.exe -M vid_gen ... --hires --hires-upscalers-dir ..\models\latent_upscale_models --hires-upscaler h3-latent-upscaler-3d --hires-width 1216 --hires-height 672 --hires-steps 4 --hires-denoising-strength 0.4
+```
+
+Notes for this mode:
+
+- The learned upscaler expects a second normalization of the diffusion latent,
+  matching its training convention. The runner applies and reverses this extra
+  transform internally; callers should pass the usual H3 diffusion latents.
+- The frame count is never changed by the second stage, so the packed audio
+  latent length stays valid and the generated soundtrack is preserved.
+- Width and height still have to be multiples of 32; the hires target is
+  aligned like any other resolution. The learned upscaler additionally wants
+  an even latent size on both axes, which 32-pixel targets provide.
+- `--hires-sigmas` can replace `--hires-denoising-strength` when the second
+  stage should resume from an explicit sigma schedule, which is how most
+  two-stage H3 workflows split the sampler.
+- Keyframe conditioning (`--init-img` / `--end-img`) is not available together
+  with hires yet: those latents are encoded at the first stage resolution and
+  would have to be re-encoded for the hires canvas. Use T2VA or Ref2VA.
+- The second stage runs at the larger latent, so it is the memory-hungry half;
+  combine it with quantization, `--vae-tiling` and `--offload-to-cpu` on
+  small cards.
+- The learned upscaler is scale-conditional: the target-to-source ratio is
+  fed into the network as an embedding, 1.2-2.0x is the trained domain, and
+  downscaling is rejected.
+
 ## Shape and runtime notes
 
 - Width and height are aligned upward to a multiple of 32.
 - Frame count is aligned upward to the `17k + 5` grid, with a minimum of 5.
 - MiniMax-H3 runs at 24 fps; another requested value is overridden.
+- Attention head grouping is controlled by
+  `--model-args minimax_h3_attn_chunk_budget=2`, in GiB (default: `2`).
+  The budget estimates the size of a single FP32 Q/K tensor per group:
+  `head_dim * tokens * batch * 4 bytes * heads_per_group`. The group size adapts
+  to the sequence length, with a minimum of one head; inputs that fit use all
+  heads together. Set `0` to disable grouping, or a smaller budget such as
+  `0.25` to reduce peak activation memory for long sequences at the cost of
+  extra launches and concatenation. Each head still attends to the full sequence.
+  This is a grouping target, not a cap on total attention memory or VRAM usage;
+  other intermediates, backend workspace, and weights require additional memory.
+  `--max-vram` and `--offload-to-cpu` control weight residency, but the active
+  block's working memory must also fit within the VRAM budget.
 - The default video flow shift is 12. The audio stream is mapped internally to
   its shift of 3, so the regular samplers can operate on the packed AV latent.
+
+## HyperFlow H3 Flash
+
+[HyperFlow H3 Flash](https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash)
+uses the shared components and input options described in this guide. Replace
+`--diffusion-model` with the matching HyperFlow FL2VA or Ref2VA checkpoint and
+use these sampling options:
+
+```sh
+--sampling-method euler --cfg-scale 1 --sigmas hyperflow-h3-flash-6
+```
+
+Use `hyperflow-h3-flash-4` or `hyperflow-h3-flash-8` for 4 or 8 steps. These
+presets also work with `--hires-sigmas`; their length determines the step count.
+They already include video flow shift 12, so keep the default flow shift.

docs/model_config.md

@@ -60,6 +60,8 @@ static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage
 ```
 
 Use `TensorStorage` metadata, especially `n_dims` and `ne`, to infer shapes.
+Use `logical_ne0()` for Linear input widths; packed INT4 weights store half of
+the logical width in `ne[0]`. Keep raw `ne` values for storage layout calculations.
 Do not load or parse tensor data for config detection.
 
 Detection should respect `prefix`. For nested weights, construct full names from

docs/p_h3_image.md

@@ -0,0 +1,34 @@
+# How to Use
+
+P-H3-Image supports text-to-image and reference-image editing. It uses a MiniMax-H3 diffusion transformer, a single-frame VAE, and Qwen3-VL-4B with a text connector. Prompt enhancement is not included.
+
+## Download weights
+
+- Download P-H3-Image
+    - safetensors: https://huggingface.co/PrunaAI/P-H3-Image-Comfy/tree/main/diffusion_models
+- Download P-H3-Image VAE
+    - safetensors: https://huggingface.co/PrunaAI/P-H3-Image-Comfy/tree/main/vae
+- Download Qwen3-VL-4B with the text connector
+    - https://huggingface.co/PrunaAI/P-H3-Image-Comfy/tree/main/text_encoders
+
+Alternatively, reuse Qwen3-VL-4B, including its vision tower, with the [separate text connector](https://huggingface.co/PrunaAI/P-H3-Image/blob/main/text_connector/diffusion_pytorch_model.safetensors). Save the connector as `p-h3-image_text_connector.safetensors` and pass it with `--embeddings-connectors`.
+
+Use the P-H3 image VAE; the MiniMax-H3 video VAE is not compatible. No audio VAE is needed.
+
+## Examples
+
+### Text-to-image
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\p-h3-image_int8_convrot.safetensors --llm ..\models\text_encoders\qwen3vl_4b_bf16.safetensors --embeddings-connectors ..\models\text_encoders\p-h3-image_text_connector.safetensors --vae ..\models\vae\p-h3-image_vae_bf16.safetensors --cfg-scale 1.0 --steps 8 --sampling-method euler --sigmas p-h3-image -p "a lovely cat holding a sign says 'p-h3-image.cpp'" -v --offload-to-cpu --rng cpu -H 2912 -W 2912
+```
+
+### Image editing
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\p-h3-image_int8_convrot.safetensors --llm ..\models\text_encoders\qwen3vl_4b_bf16.safetensors --embeddings-connectors ..\models\text_encoders\p-h3-image_text_connector.safetensors --vae ..\models\vae\p-h3-image_vae_bf16.safetensors --cfg-scale 1.0 --steps 8 --sampling-method euler --sigmas p-h3-image -r source.png -r reference.png -p "Edit <Picture 1>. Give the cat the scarf from <Picture 2>, preserving its face and the background." -v --offload-to-cpu --rng cpu -H 2912 -W 2912
+```
+
+Reference images are numbered in argument order as `<Picture 1>`, `<Picture 2>`, and so on. A plain edit instruction gets `Edit <Picture 1>. ` prepended automatically. Prompts starting with `Edit <Picture N>` or `subject_definitions:` are used directly.
+
+Use dimensions divisible by 32. The upstream default is 2912 × 2912; smaller canvases reduce memory use. Use CFG 1 and `--sigmas p-h3-image` for the upstream 8-step schedule. Keep the default VAE tiling size of 256 pixels.

docs/performance.md

@@ -159,6 +159,25 @@ allocator/scheduler capacity. Compatible gallocr reservations are reused across
 graphs; idle workspaces can be reclaimed under pressure and are freed at runner
 completion. Cross-graph caches survive individual graphs, but cut buffers do not.
 
+On discrete Vulkan devices, runners reuse released device buffers of the same
+size to reduce allocation overhead during segmented execution. This idle buffer
+cache is limited to 1 GiB per device and is reclaimed before allocating a new
+size or at runner completion. It is outside the managed `--max-vram` budget,
+so physical VRAM usage can be higher. Set `GGML_VK_REUSE_BUFFERS=0` to disable
+it. Unified-memory and host-memory modes do not use this cache.
+
+Discrete NVIDIA and non-GCN AMD Vulkan devices use a separate transfer queue by
+default when available and graphics queue use is disabled, so weight prefetch can overlap compute. Set
+`GGML_VK_ASYNC_USE_TRANSFER_QUEUE=0` to disable it.
+
+On discrete NVIDIA devices, allocation chunks can grow up to 4 GiB, bounded by the device's buffer and allocation
+limits, to reduce fragmentation within large graph workspaces. Other devices
+retain the 1 GiB default. `GGML_VK_SUBALLOCATION_BLOCK_SIZE` overrides this limit
+in bytes on any vendor; for example, `1073741824` selects 1 GiB and `4294967296`
+selects 4 GiB, subject to device limits. This is a soft chunk limit, not a reservation
+or a change to `--max-vram`: a single oversized tensor or exhausting the allocator's
+chunk count can require a larger allocation, which may fail.
+
 The recommended shape for "biggest model my card can host" is:
 
 ```shell

docs/relay_vsr.md

@@ -0,0 +1,50 @@
+# RelayVSR image and video super-resolution
+
+[RelayVSR](https://github.com/kopperx/RelayVSR) upscales images and videos by
+2x or 4x with `-M vid_upscale`.
+
+## Weights
+
+[ComfyUI safetensors weights](https://huggingface.co/Comfy-Org/RelayVSR) load
+directly. Select the matching pair for the requested scale:
+
+| `--upscale-factor` | `--diffusion-model` | `--vae` |
+| --- | --- | --- |
+| 2 | `relayvsr_x2_lightx2v_int8_convrot.safetensors` | `relayvsr_flash_vae_x2_fp32.safetensors` |
+| 4 | `relayvsr_x4_int8_convrot.safetensors` | `relayvsr_flash_vae_x4_fp32.safetensors` |
+
+No separate text encoder or ordinary Wan VAE is needed.
+
+## Usage
+
+For a video, extract numbered frames into a directory, then run:
+
+```sh
+sd-cli -M vid_upscale \
+  --diffusion-model models/diffusion_models/relayvsr_x2_lightx2v_int8_convrot.safetensors \
+  --vae models/vae/relayvsr_flash_vae_x2_fp32.safetensors \
+  --input-frames input_frames --upscale-factor 2 --fps 24 \
+  --vsr-window-size 16 -s 42 -o output.avi
+```
+
+Frames are read in filename order. Set `--fps` to the source video's frame rate.
+Video outputs support AVI, WebM, and animated WebP. To save individual frames,
+use an output pattern such as `-o output/%04d.png`.
+
+For a single image, replace `--input-frames input_frames` with `-i input.png`
+and use `-o output.png` (or JPEG/WebP). No FPS setting is needed for an image.
+`-i` / `--init-img` and `--input-frames` are mutually exclusive.
+
+## Settings and limits
+
+- `--upscale-factor` is required. Use 2 or 4 with the corresponding weights.
+  Output width and height are exactly the input dimensions times the scale.
+- `--vsr-window-size` sets the reference-frame window, including both endpoints.
+  The default 16 selects reference frames 0, 15, 30, and so on, and always
+  includes the final frame. Every input frame appears once in the output.
+- `--flash-attn auto` and `--conv-direct auto` are enabled by default. See
+  [backend selection](backend.md) to choose a device or weight placement.
+- Sampler, scheduler, and step settings do not configure RelayVSR.
+  `--vsr-temporal-chunk-size` is a SeedVR2 option.
+- Long videos require enough system memory to hold the input and output frames.
+  Audio is not copied to the output.

docs/sage_attention.md

@@ -4,8 +4,9 @@
 the high-noise diffusion model when present. Python, PyTorch, and Triton are not
 required at build time or runtime.
 
-The CUDA backend automatically selects a kernel supported by both the GPU and
-the compiled CUDA toolkit:
+The sd.cpp CUDA extension registers SageAttention through GGML's generic
+operator API. It selects a kernel supported by the GPU, compiled CUDA target,
+and toolkit:
 
 | GPU / toolkit | Implementation |
 | --- | --- |

docs/seedvr2.md

@@ -0,0 +1,69 @@
+# SeedVR2 image and video super-resolution
+
+SeedVR2 upscales images and videos with `-M vid_upscale`. Support is experimental.
+
+## Weights
+
+These [ComfyUI safetensors weights](https://huggingface.co/Comfy-Org/SeedVR2)
+load directly, without conversion or a separate text encoder:
+
+| Model | `--diffusion-model` | `--vae` |
+| --- | --- | --- |
+| 3B | `seedvr2_3b_int8_convrot.safetensors` | `seedvr2_ema_vae_fp16.safetensors` |
+| 7B Sharp | `seedvr2_7b_sharp_int8_convrot.safetensors` | `seedvr2_ema_vae_fp16.safetensors` |
+
+Use the matching SeedVR2 VAE for either model. The 7B model requires more memory.
+
+## Upscale an image
+
+```sh
+sd-cli -M vid_upscale \
+  --diffusion-model models/diffusion_models/seedvr2_3b_int8_convrot.safetensors \
+  --vae models/vae/seedvr2_ema_vae_fp16.safetensors \
+  -i input.png --upscale-factor 2 -s 42 -o restored.png
+```
+
+`-i` / `--init-img` accepts a single image. Outputs can be PNG, JPEG, or WebP.
+No FPS or temporal chunk setting is needed for an image.
+
+## Upscale a video
+
+Extract the video into numbered frames, then pass the frame directory:
+
+```sh
+mkdir frames
+ffmpeg -i input.mp4 -fps_mode passthrough frames/%06d.png
+sd-cli -M vid_upscale \
+  --diffusion-model models/diffusion_models/seedvr2_3b_int8_convrot.safetensors \
+  --vae models/vae/seedvr2_ema_vae_fp16.safetensors \
+  --input-frames frames --upscale-factor 2 --fps 24 \
+  --vsr-temporal-chunk-size 1 -s 42 -o restored.avi
+```
+
+Set `--fps` to the source video's frame rate; 24 is only the example above.
+Frames are read in filename order. Use either `-i` or `--input-frames`, not both.
+
+Video outputs support AVI, WebM, and animated WebP. To save individual frames,
+use an output pattern such as `-o restored/%06d.png`.
+
+## Settings and limits
+
+- `--upscale-factor` is required and must be a positive integer. Output width
+  and height are exactly the input dimensions multiplied by this value.
+- `--vsr-temporal-chunk-size` controls how many video frames are processed
+  together. Valid values are 1, 5, 9, 13, and so on; the default is 5. A value of
+  1 processes each frame independently and uses less memory. Larger values use
+  neighboring frames but require more memory. Chunk boundaries may be visible.
+- Frame count and order are preserved, including a shorter final chunk. Audio
+  is not copied to the output.
+- If you run out of GPU memory, try `--vsr-temporal-chunk-size 1`. A five-frame
+  chunk at 1728x960 output resolution exceeded 24 GB of GPU memory in testing.
+- `--flash-attn auto` and `--conv-direct auto` are enabled by default. See
+  [backend selection](backend.md) to choose a device or weight placement.
+- SeedVR2 does not use sampler, scheduler, or step settings.
+  `--vsr-window-size` is a RelayVSR option and does not affect SeedVR2.
+
+3B has been tested on CUDA and CPU; 7B has been tested on CUDA with small inputs.
+Other backends have not been tested with this integration.
+
+Model reference: [ByteDance-Seed/SeedVR](https://github.com/ByteDance-Seed/SeedVR).

examples/cli/main.cpp

@@ -3,6 +3,7 @@
 #include <time.h>
 #include <algorithm>
 #include <cctype>
+#include <climits>
 #include <filesystem>
 #include <functional>
 #include <iostream>
@@ -33,10 +34,14 @@ const char* previews_str[] = {
 std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)");
 
 struct SDCliParams {
-    SDMode mode             = IMG_GEN;
-    std::string output_path = "output.png";
-    int output_begin_idx    = -1;
-    int compression_quality = 90;
+    std::string input_frames;
+    int video_upscale_factor    = 0;
+    int vsr_window_size         = 16;
+    int vsr_temporal_chunk_size = 5;
+    SDMode mode                 = IMG_GEN;
+    std::string output_path     = "output.png";
+    int output_begin_idx        = -1;
+    int compression_quality     = 90;
     std::string image_path;
     std::string metadata_format = "text";
 
@@ -62,6 +67,7 @@ struct SDCliParams {
         ArgOptions options;
 
         options.string_options = {
+            {"", "--input-frames", "directory of ordered RGB frames for vid_upscale (use -i for a single image)", 0, &input_frames},
             {"-o",
              "--output",
              "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs support .avi, .webm, and animated .webp",
@@ -90,6 +96,9 @@ struct SDCliParams {
         };
 
         options.int_options = {
+            {"", "--upscale-factor", "required positive video upscale factor (RelayVSR: 2 or 4; SeedVR2: integer)", &video_upscale_factor},
+            {"", "--vsr-window-size", "RelayVSR window including both endpoints (default: 16)", &vsr_window_size},
+            {"", "--vsr-temporal-chunk-size", "SeedVR2 input frames per independent chunk: 1 or 4n+1 (default: 5)", &vsr_temporal_chunk_size},
             {"",
              "--preview-interval",
              "preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1",
@@ -190,7 +199,7 @@ struct SDCliParams {
         options.manual_options = {
             {"-M",
              "--mode",
-             "run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen",
+             "run mode, one of [img_gen, adetailer, vid_gen, vid_upscale, upscale, convert, metadata], default: img_gen",
              on_mode_arg},
             {"",
              "--preview",
@@ -220,6 +229,10 @@ struct SDCliParams {
     }
 
     bool validate() {
+        if (mode == VID_UPSCALE && video_upscale_factor <= 0) {
+            LOG_ERROR("vid_upscale requires a positive --upscale-factor");
+            return false;
+        }
         if (mode != METADATA) {
             if (output_path.length() == 0) {
                 LOG_ERROR("error: the following arguments are required: output_path");
@@ -529,7 +542,9 @@ bool save_results(const SDCliParams& cli_params,
         return sucessful_reults != 0;
     }
 
-    if (cli_params.mode == VID_GEN && num_results > 1) {
+    const bool video_upscale_output = cli_params.mode == VID_UPSCALE &&
+                                      (gen_params.init_image_path.empty() || ext_lower == ".avi" || ext_lower == ".webm");
+    if ((cli_params.mode == VID_GEN && num_results > 1) || video_upscale_output) {
         if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm")
             ext = ".avi";
         std::string params = gen_params.embed_image_metadata
@@ -664,6 +679,72 @@ int main(int argc, const char* argv[]) {
         return 0;
     }
 
+    if (cli_params.mode == VID_UPSCALE) {
+        if (cli_params.input_frames.empty() == gen_params.init_image_path.empty()) {
+            LOG_ERROR("vid_upscale requires exactly one input: -i/--init-img or --input-frames");
+            return 1;
+        }
+        ctx_params.prepare_backend_assignments();
+        if (ctx_params.diffusion_model_path.empty() || ctx_params.vae_path.empty() || gen_params.fps <= 0) {
+            LOG_ERROR("vid_upscale requires --diffusion-model, --vae and positive --fps");
+            return 1;
+        }
+        std::vector<SDImageOwner> owned_frames;
+        if (!gen_params.init_image_path.empty()) {
+            owned_frames.emplace_back();
+            if (!load_sd_image_from_file(owned_frames.back().put(), gen_params.init_image_path.c_str(), 0, 0, 3)) {
+                LOG_ERROR("load image from '%s' failed", gen_params.init_image_path.c_str());
+                return 1;
+            }
+        } else if (!load_images_from_dir(cli_params.input_frames, owned_frames) || owned_frames.empty()) {
+            LOG_ERROR("vid_upscale input contains no readable frames");
+            return 1;
+        }
+        if (owned_frames.size() > size_t(INT_MAX))
+            return 1;
+        std::vector<sd_image_t> frames;
+        for (const auto& frame : owned_frames)
+            frames.push_back(frame.get());
+        sd_video_upscaler_ctx_params_t params;
+        sd_video_upscaler_ctx_params_init(&params);
+        params.diffusion_model_path  = ctx_params.diffusion_model_path.c_str();
+        params.vae_path              = ctx_params.vae_path.c_str();
+        params.upscale_factor        = cli_params.video_upscale_factor;
+        params.n_threads             = ctx_params.n_threads;
+        params.backend               = ctx_params.effective_backend.c_str();
+        params.params_backend        = ctx_params.effective_params_backend.c_str();
+        params.flash_attn_mode       = ctx_params.flash_attn_mode.c_str();
+        std::string conv_direct_mode = ctx_params.conv_direct_mode;
+        if (conv_direct_mode.empty()) {
+            conv_direct_mode = "auto";
+            if (ctx_params.diffusion_conv_direct)
+                conv_direct_mode += ",diffusion=on";
+            if (ctx_params.vae_conv_direct)
+                conv_direct_mode += ",vae=on";
+        }
+        params.conv_direct_mode = conv_direct_mode.c_str();
+        std::unique_ptr<video_upscaler_ctx_t, decltype(&free_video_upscaler_ctx)> upscaler(
+            new_video_upscaler_ctx(&params),
+            free_video_upscaler_ctx);
+        if (!upscaler)
+            return 1;
+        sd_video_generation_result_t* output = nullptr;
+        sd_video_upscale_params_t upscale_params;
+        sd_video_upscale_params_init(&upscale_params);
+        upscale_params.frames              = frames.data();
+        upscale_params.frame_count         = int(frames.size());
+        upscale_params.fps                 = gen_params.fps;
+        upscale_params.window_size         = cli_params.vsr_window_size;
+        upscale_params.temporal_chunk_size = cli_params.vsr_temporal_chunk_size;
+        upscale_params.seed                = gen_params.seed;
+        if (!upscale_video(upscaler.get(), &upscale_params, &output))
+            return 1;
+        std::string video_info = output->generation_info;
+        SDImageVec results;
+        results.adopt_video(output);
+        return save_results(cli_params, ctx_params, gen_params, results, int(results.size()), nullptr, video_info) ? 0 : 1;
+    }
+
     if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) {
         LOG_ERROR("--ad-model is only supported in image generation and adetailer modes");
         return 1;

examples/common/common.cpp

@@ -35,6 +35,7 @@ const char* const modes_str[] = {
     "convert",
     "upscale",
     "metadata",
+    "vid_upscale",
 };
 
 static sd_vae_format_t str_to_vae_format(const std::string& value) {
@@ -464,7 +465,7 @@ ArgOptions SDContextParams::get_options() {
          &uncond_diffusion_model_path},
         {"",
          "--embeddings-connectors",
-         "path to LTXAV embeddings connectors",
+         "path to embeddings connectors (LTXAV or P-H3-Image)",
          0,
          &embeddings_connectors_path},
         {"",
@@ -1135,7 +1136,7 @@ ArgOptions SDGenerationParams::get_options() {
          &extra_ad_args},
         {"-i",
          "--init-img",
-         "path to the init image",
+         "path to the init image, or a single input image for vid_upscale",
          0,
          &init_image_path},
         {"",
@@ -1187,6 +1188,13 @@ ArgOptions SDGenerationParams::get_options() {
          "under --hires-upscalers-dir (default: Latent)",
          0,
          &hires_upscaler},
+        {"",
+         "--control-type",
+         "kind of the control image, one of [openpose, depth, scribble, canny, normal, segment, tile, repaint]. "
+         "Required by an SDXL Control Net Union model (tile and repaint only exist in its ProMax variant), ignored by any other Control Net. "
+         "Also accepted: hed, softedge, pidi, ted (scribble); lineart, anime_lineart, mlsd (canny); seg (segment); inpaint, outpaint (repaint)",
+         0,
+         &control_type},
         {"",
          "--extra-sample-args",
          "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
@@ -1530,6 +1538,14 @@ ArgOptions SDGenerationParams::get_options() {
         std::string sigmas_str = value;
         if (sigmas_str == "qwen-image-2.1-turbo") {
             sigmas_str = "1.0,0.978453,0.95418,0.926626,0.89508,0.845148,0.704534,0.414568,0.0";
+        } else if (sigmas_str == "p-h3-image") {
+            sigmas_str = "1.0,0.984753,0.967905,0.951688,0.925372,0.881923,0.810083,0.528718,0.0";
+        } else if (sigmas_str == "hyperflow-h3-flash-4") {
+            sigmas_str = "1.0,0.962575853,0.755781889,0.312497854,0.0";
+        } else if (sigmas_str == "hyperflow-h3-flash-6") {
+            sigmas_str = "1.0,0.982655942,0.914642572,0.674755692,0.385223746,0.159055337,0.0";
+        } else if (sigmas_str == "hyperflow-h3-flash-8") {
+            sigmas_str = "1.0,0.993909776,0.9842875,0.966063678,0.923076928,0.834942341,0.696852028,0.468753338,0.0";
         }
         if (!sigmas_str.empty() && sigmas_str.front() == '[') {
             sigmas_str.erase(0, 1);
@@ -1777,11 +1793,11 @@ ArgOptions SDGenerationParams::get_options() {
          on_scheduler_arg},
         {"",
          "--sigmas",
-         "custom sigma values for the sampler, comma-separated (e.g., \"14.61,7.8,3.5,0.0\") or preset: qwen-image-2.1-turbo.",
+         "custom sigma values for the sampler, comma-separated (e.g., \"14.61,7.8,3.5,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.",
          on_sigmas_arg},
         {"",
          "--hires-sigmas",
-         "custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\") or preset: qwen-image-2.1-turbo.",
+         "custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.",
          on_hires_sigmas_arg},
         {"",
          "--skip-layers",
@@ -2247,6 +2263,7 @@ bool SDGenerationParams::from_json_str(
 
     load_if_exists("strength", strength);
     load_if_exists("control_strength", control_strength);
+    load_if_exists("control_type", control_type);
     load_if_exists("ip_adapter_strength", ip_adapter_strength);
     load_if_exists("moe_boundary", moe_boundary);
     load_if_exists("vace_strength", vace_strength);
@@ -2636,6 +2653,15 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s
         sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100);
     }
 
+    resolved_control_type = SD_CONTROL_TYPE_NONE;
+    if (!control_type.empty()) {
+        resolved_control_type = str_to_sd_control_type(control_type.c_str());
+        if (resolved_control_type == SD_CONTROL_TYPE_COUNT) {
+            LOG_ERROR("invalid control type '%s' (openpose, depth, scribble, canny, normal, segment, tile, repaint)", control_type.c_str());
+            return false;
+        }
+    }
+
     hires_upscaler_model_path.clear();
     if (hires_enabled) {
         if (hires_upscaler.empty()) {
@@ -2875,6 +2901,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.qwen_image_layers   = qwen_image_layers;
     params.control_image       = control_image.get();
     params.control_strength    = control_strength;
+    params.control_type        = resolved_control_type;
     params.ip_adapter_image    = ip_adapter_image.get();
     params.ip_adapter_strength = ip_adapter_strength;
     params.pm_params           = pm_params;
@@ -3080,6 +3107,7 @@ std::string SDGenerationParams::to_string() const {
         << "  vace_strength: " << vace_strength << ",\n"
         << "  strength: " << strength << ",\n"
         << "  control_strength: " << control_strength << ",\n"
+        << "  control_type: \"" << control_type << "\",\n"
         << "  seed: " << seed << ",\n"
         << "  upscale_repeats: " << upscale_repeats << ",\n"
         << "  upscale_tile_size: " << upscale_tile_size << ",\n"

examples/common/common.h

@@ -22,7 +22,7 @@ std::string argv_to_utf8(int index, const char** argv);
 #define BOOL_STR(b) ((b) ? "true" : "false")
 
 extern const char* const modes_str[];
-#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata"
+#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata, vid_upscale"
 
 enum SDMode {
     IMG_GEN,
@@ -31,6 +31,7 @@ enum SDMode {
     CONVERT,
     UPSCALE,
     METADATA,
+    VID_UPSCALE,
     MODE_COUNT
 };
 
@@ -227,6 +228,7 @@ struct SDGenerationParams {
     std::string end_image_path;
     std::string mask_image_path;
     std::string control_image_path;
+    std::string control_type;
     std::string ip_adapter_image_path;
     std::vector<std::string> ref_image_paths;
     std::vector<std::string> ref_video_paths;
@@ -291,6 +293,7 @@ struct SDGenerationParams {
     std::string prompt_with_lora;  // for metadata record only
     std::vector<sd_lora_t> lora_vec;
     sd_hires_upscaler_t resolved_hires_upscaler;
+    sd_control_type_t resolved_control_type = SD_CONTROL_TYPE_NONE;
 
     // Owned execution payload.
     SDImageOwner init_image;

examples/common/generation_metadata.cpp

@@ -1,5 +1,6 @@
 #include "generation_metadata.h"
 
+#include <cctype>
 #include <filesystem>
 #include <map>
 
@@ -70,6 +71,7 @@ std::string generation_info_to_webui(const std::string& text) {
     const auto& sampling = info.at("sample_params");
     std::string result   = info.value("prompt", "");
     if (info.contains("loras")) {
+        bool prompt_trimmed = false;
         for (const auto& lora : info.at("loras")) {
             // WebUI tags cannot restrict a LoRA to the high-noise model.
             if (lora.value("is_high_noise", false))
@@ -83,8 +85,15 @@ std::string generation_info_to_webui(const std::string& text) {
                 name.erase(extension);
             if (name.empty() || name.find_first_of(":<>\r\n") != std::string::npos)
                 continue;
-            if (!result.empty())
-                result += " ";
+            if (!result.empty()) {
+                if (!prompt_trimmed) {
+                    while (!result.empty() && std::isspace(static_cast<unsigned char>(result.back())))
+                        result.pop_back();
+                    prompt_trimmed = true;
+                }
+                if (!result.empty())
+                    result += " ";
+            }
             result += "<lora:" + name + ":" + lora.at("multiplier").dump() + ">";
         }
     }

examples/server/README.md

@@ -76,6 +76,16 @@ Recipes declaring `video_frames` also switch the form to Video Generation;
 other recipes switch to Image Generation. Generation is enabled after the
 selected recipe is loaded. Changing a selection does not reload the model
 until you click Load.
+
+The generation status bar shows the current stage and elapsed time even when
+Preview Mode is `none`. During sampling it shows completed steps for the current
+pass; prompt encoding and VAE decoding show an activity indicator when their
+total work is unknown. Sampling reaching 100% does not mean the job is finished:
+decoding, ADetailer, and output encoding may still follow.
+The canvas keeps the submitted output aspect ratio, and stage changes keep the
+same layout. Editing the next request's resolution does not resize the current
+preview or result.
+
 The server can also start with the existing model arguments shown below;
 existing API clients then work without selecting a recipe. The UI automatically
 selects **Command-line model (all features)** and uses the startup generation

examples/server/api.md

@@ -85,7 +85,7 @@ available during loading. Model lists are empty when no model is loaded;
 | GET | `/sdcpp/v1/model-files` | `{files: [{id, name, category, root, relative_path}], errors: [...]}` |
 | POST | `/sdcpp/v1/models/refresh` | Rescan recipes and files; return `{errors: [...]}` |
 | GET | `/sdcpp/v1/models/status` | Current state, recipe, selection, error, loading time and optional revision |
-| POST | `/sdcpp/v1/models/load` | Load the submitted recipe and selection; return status after completion |
+| POST | `/sdcpp/v1/models/load` | Load the submitted recipe and selection; return status after completion. A load is a no-op returning the current status with `already_loaded: true` only when the state is ready and the fully resolved selection, the complete recipe object, and the resolved component paths all match the active model; a catalog refresh with modified recipe metadata or a retargeted model path therefore reloads even with an unchanged recipe ID and selection |
 | POST | `/sdcpp/v1/models/unload` | Release the active context; return empty status |
 
 Load request:
@@ -269,6 +269,14 @@ Legacy `category` fields remain accepted, but must match one of the slot's
 supported directories; they cannot redefine or restrict its search directories.
 At least one required `model` or `diffusion_model` slot is necessary.
 
+An optional top-level `models` object maps slot IDs to file IDs, providing
+recipe-local defaults for component selection. A load request may omit
+`models` entirely and load with the recipe's defaults. Selections sent in the
+request always override recipe defaults, slot by slot. An empty string
+disables a slot's recipe default (omitting the slot selects it); the empty
+key is then normalized out of the reported selection, so for slots without a
+recipe default an empty selection loads identically to an omitted one.
+
 `size_multiple` is a positive integer from 1 to 4096 that declares the width and
 height increment used by the web UI. It defaults to `8` when omitted; the UI also
 uses `8` before a recipe is selected. For example, `"size_multiple": 16` makes
@@ -767,6 +775,7 @@ Field types:
 | `queue_position` | `integer` |
 | `result` | `object \| null` |
 | `error` | `object \| null` |
+| `progress` | `object` |
 | `preview` | `object \| null` |
 
 `created`, `started`, and `completed` are integer Unix timestamps in
@@ -774,6 +783,57 @@ Field types:
 the elapsed generation time in seconds, including output encoding and excluding
 queue time.
 
+`progress` is available independently of the preview setting, for both image
+and video jobs. Continue polling until `status` is terminal: a sampling pass
+can finish before decoding, ADetailer, or output encoding finishes.
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `progress.stage` | `string` | Current execution stage; see below |
+| `progress.detail` | `string` | Stage detail, or an empty string |
+| `progress.pass` | `integer` | Sampling pass number within the job; zero before sampling |
+| `progress.completed` | `integer \| null` | Completed units in the current stage/pass |
+| `progress.total` | `integer \| null` | Actual total units, or null when unknown |
+| `progress.unit` | `string \| null` | `steps` for sampling, `images` for image output encoding; null when uncounted |
+| `progress.stage_elapsed_ms` | `integer` | Elapsed time in the current stage/detail/pass |
+| `progress.elapsed_ms` | `integer` | Job execution time, excluding queue time; zero before execution |
+
+Stages include `queued`, `waiting` (waiting for the execution lock), `preparing`,
+`conditioning`, `sampling`, `vae_encode`, `vae_decode`, `audio_decode`,
+`upscaling`, `adetailer`, and `encoding`. Stages may repeat or be skipped.
+Details include `positive_prompt` / `negative_prompt` / `image_condition` for
+conditioning, `high_noise` / `low_noise` for two-stage sampling, and `image` /
+`video` for output encoding. During ADetailer, `detail` identifies the inner
+stage, such as `conditioning` or `sampling`. Clients should tolerate unknown
+stage, detail, and unit values.
+
+Sampling counts completed logical steps using the actual schedule, including
+strength adjustments. Each new sampling pass increments `pass` and starts at
+zero, including batch items, highres passes, high/low-noise passes, and ADetailer
+passes. The ratio is **current-pass progress**, not overall job completion.
+The progress snapshot is published atomically. Preview metadata still describes
+the preview frame, which may be older than the current progress.
+
+Unknown totals are represented by null counts, not an estimated percentage.
+For example, prompt encoding can report:
+
+```json
+"progress": {
+  "stage": "conditioning",
+  "detail": "positive_prompt",
+  "pass": 0,
+  "completed": null,
+  "total": null,
+  "unit": null,
+  "stage_elapsed_ms": 3200,
+  "elapsed_ms": 4100
+}
+```
+
+Elapsed times freeze when a job completes, fails, or is cancelled; the last
+stage is retained for inspection. Use `status` to determine the final outcome.
+Elapsed time alone does not indicate that computation is advancing.
+
 `preview` sub-fields:
 
 | Field | Type | Notes |
@@ -928,6 +988,7 @@ LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEH
 | `ref_image_args` | `string` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `control_type` | `string` |
 | `ip_adapter_strength` | `number` |
 | `hires` | `object` |
 | `hires.enabled` | `boolean` |
@@ -1108,6 +1169,7 @@ Example:
   "ref_image_args": "",
   "increase_ref_index": false,
   "control_strength": 0.9,
+  "control_type": "",
   "ip_adapter_strength": 1.0,
   "embed_image_metadata": true,
 
@@ -1233,6 +1295,7 @@ Top-level scalar fields:
 | `image_preprocess` | `string \| array<string>` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `control_type` | `string` |
 | `ip_adapter_strength` | `number` |
 | `embed_image_metadata` | `boolean` |
 | `preview` | `string` |

examples/server/async_jobs.cpp

@@ -92,6 +92,63 @@ static void set_preview_callback(AsyncGenerationJob& job) {
     sd_set_preview_callback(preview_callback, mode, interval, true, false, &job);
 }
 
+static void update_job_progress(AsyncGenerationJob& job, const char* stage, const char* detail = "", int sample_pass = -1, int64_t completed = 0, int64_t total = 0, const char* unit = "") {
+    std::lock_guard<std::mutex> lock(job.progress_mutex);
+    auto& progress = job.progress;
+    if (sample_pass < 0) {
+        sample_pass = progress.sample_pass;
+    }
+    if (progress.stage != stage || progress.detail != detail || progress.sample_pass != sample_pass) {
+        progress.stage_started = std::chrono::steady_clock::now();
+    }
+    progress.stage       = stage;
+    progress.detail      = detail;
+    progress.sample_pass = sample_pass;
+    progress.completed   = completed;
+    progress.total       = total;
+    progress.unit        = unit;
+}
+
+static void finish_job_progress(AsyncGenerationJob& job) {
+    std::lock_guard<std::mutex> lock(job.progress_mutex);
+    job.progress.finished = std::chrono::steady_clock::now();
+}
+
+class JobCallbacksGuard {
+    sd_ctx_t* ctx_;
+    AsyncGenerationJob& job_;
+    bool detailing_ = false;
+
+    static void progress_callback(const sd_generation_progress_t* progress, void* data) {
+        auto& guard = *static_cast<JobCallbacksGuard*>(data);
+        update_job_progress(guard.job_, guard.detailing_ ? "adetailer" : progress->stage,
+                            guard.detailing_ ? progress->stage : progress->detail,
+                            progress->sample_pass, progress->completed, progress->total, progress->unit);
+    }
+
+public:
+    JobCallbacksGuard(sd_ctx_t* ctx, AsyncGenerationJob& job)
+        : ctx_(ctx), job_(job) {
+        clear_preview_callback();
+        set_preview_callback(job);
+        sd_set_generation_progress_callback(ctx_, progress_callback, this);
+    }
+
+    ~JobCallbacksGuard() {
+        sd_set_generation_progress_callback(ctx_, nullptr, nullptr);
+        clear_preview_callback();
+    }
+
+    void begin_detailing() {
+        clear_preview_callback();
+        detailing_ = true;
+        update_job_progress(job_, "adetailer");
+    }
+
+    JobCallbacksGuard(const JobCallbacksGuard&)            = delete;
+    JobCallbacksGuard& operator=(const JobCallbacksGuard&) = delete;
+};
+
 const char* async_job_kind_name(AsyncJobKind kind) {
     switch (kind) {
         case AsyncJobKind::ImgGen:
@@ -186,6 +243,7 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) {
     job.result_fps         = 0;
     job.error_code         = "cancelled";
     job.error_message      = "job cancelled by client";
+    finish_job_progress(job);
     return true;
 }
 
@@ -200,6 +258,24 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo
     result["completed"]        = job.completed_at == 0 ? json(nullptr) : json(job.completed_at);
     result["queue_position"]   = 0;
 
+    {
+        std::lock_guard<std::mutex> lock(job.progress_mutex);
+        const auto& progress = job.progress;
+        const auto end       = progress.finished == std::chrono::steady_clock::time_point{}
+                                   ? std::chrono::steady_clock::now()
+                                   : progress.finished;
+        result["progress"]   = {
+            {"stage", progress.stage},
+            {"detail", progress.detail},
+            {"pass", progress.sample_pass},
+            {"completed", progress.total > 0 ? json(progress.completed) : json(nullptr)},
+            {"total", progress.total > 0 ? json(progress.total) : json(nullptr)},
+            {"unit", progress.unit.empty() ? json(nullptr) : json(progress.unit)},
+            {"stage_elapsed_ms", std::chrono::duration_cast<std::chrono::milliseconds>(end - progress.stage_started).count()},
+            {"elapsed_ms", job.started_at == 0 ? 0 : std::max<int64_t>(0, (job.completed_at == 0 ? unix_timestamp_now_ms() : job.completed_at) - job.started_at)},
+        };
+    }
+
     if (job.status == AsyncJobStatus::Queued) {
         size_t position = 1;
         for (const auto& queued_id : manager.queue) {
@@ -278,17 +354,19 @@ bool execute_img_gen_job(ServerRuntime& runtime,
             return false;
         }
         JobCancellationMonitor cancellation(runtime.sd_ctx, job);
-        set_preview_callback(job);
+        JobCallbacksGuard callbacks(runtime.sd_ctx, job);
         sd_image_generation_result_t* generated = nullptr;
         generate_image(runtime.sd_ctx, &params, &generated);
         results.adopt(generated, runtime.ctx_params->lora_model_dir);
-        clear_preview_callback();
         if (job.cancel_requested.load()) {
             return false;
         }
-        if (!results.empty() && !apply_server_adetailer(runtime, job.img_gen, params, results)) {
-            error_message = "ADetailer failed";
-            return false;
+        if (!results.empty() && job.img_gen.adetailer_enabled && !runtime.default_gen_params->ad_model_path.empty()) {
+            callbacks.begin_detailing();
+            if (!apply_server_adetailer(runtime, job.img_gen, params, results)) {
+                error_message = "ADetailer failed";
+                return false;
+            }
         }
     }
 
@@ -305,6 +383,7 @@ bool execute_img_gen_job(ServerRuntime& runtime,
         encoded_format = EncodedImageFormat::WEBP;
     }
 
+    update_job_progress(job, "encoding", "image", -1, 0, num_results, "images");
     for (int i = 0; i < num_results; ++i) {
         if (results[i].data == nullptr) {
             continue;
@@ -325,6 +404,7 @@ bool execute_img_gen_job(ServerRuntime& runtime,
         }
         output_images.push_back(base64_encode(image_bytes));
         output_info.push_back(json::parse(results.generation_info(i)));
+        update_job_progress(job, "encoding", "image", -1, i + 1, num_results, "images");
     }
 
     if (output_images.empty()) {
@@ -356,7 +436,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
             return false;
         }
         JobCancellationMonitor cancellation(runtime.sd_ctx, job);
-        set_preview_callback(job);
+        JobCallbacksGuard callbacks(runtime.sd_ctx, job);
         sd_video_generation_result_t* generated = nullptr;
         if (generate_video(runtime.sd_ctx, &params, &generated)) {
             output_fps       = generated->fps;
@@ -367,7 +447,6 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
             if (job.vid_gen.gen_params.embed_image_metadata)
                 str_params = results.generation_info(0);
         }
-        clear_preview_callback();
     }
 
     num_results = results.count();
@@ -377,6 +456,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
         return false;
     }
 
+    update_job_progress(job, "encoding", "video");
     std::vector<uint8_t> video_bytes = create_video_from_sd_images_to_vector(job.vid_gen.output_format,
                                                                              results.data(),
                                                                              num_results,
@@ -425,6 +505,7 @@ void async_job_worker(ServerRuntime& runtime) {
             job             = it->second;
             job->status     = AsyncJobStatus::Generating;
             job->started_at = unix_timestamp_now_ms();
+            update_job_progress(*job, "waiting");
         }
 
         std::vector<std::string> output_images;
@@ -459,6 +540,7 @@ void async_job_worker(ServerRuntime& runtime) {
             }
 
             job->completed_at = unix_timestamp_now_ms();
+            finish_job_progress(*job);
             if (job->cancel_requested.load()) {
                 job->status        = AsyncJobStatus::Cancelled;
                 job->error_code    = "cancelled";

examples/server/async_jobs.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <atomic>
+#include <chrono>
 #include <condition_variable>
 #include <cstdint>
 #include <deque>
@@ -29,6 +30,17 @@ enum class AsyncJobStatus {
 const char* async_job_kind_name(AsyncJobKind kind);
 const char* async_job_status_name(AsyncJobStatus status);
 
+struct AsyncJobProgress {
+    std::string stage = "queued";
+    std::string detail;
+    std::string unit;
+    int sample_pass                                     = 0;
+    int64_t completed                                   = 0;
+    int64_t total                                       = 0;
+    std::chrono::steady_clock::time_point stage_started = std::chrono::steady_clock::now();
+    std::chrono::steady_clock::time_point finished{};
+};
+
 struct AsyncGenerationJob {
     std::string id;
     AsyncJobKind kind     = AsyncJobKind::ImgGen;
@@ -47,6 +59,8 @@ struct AsyncGenerationJob {
     int result_fps         = 0;
     std::string error_code;
     std::string error_message;
+    mutable std::mutex progress_mutex;
+    AsyncJobProgress progress;
     mutable std::mutex preview_mutex;
     std::string preview_b64;
     int preview_pass          = 0;

examples/server/model_catalog.cpp

@@ -66,6 +66,14 @@ static void clear_model_paths(SDContextParams& params) {
     }
 }
 
+bool component_paths_match(const SDContextParams& a, const SDContextParams& b) {
+    for (const auto& slot : k_model_slots) {
+        if ((a.*slot.field) != (b.*slot.field))
+            return false;
+    }
+    return true;
+}
+
 std::string model_path_key(const fs::path& path) {
 #ifdef _WIN32
     const auto value = path.generic_wstring();
@@ -406,7 +414,7 @@ static void validate_video_constraints(const json& constraints) {
 }
 
 static void validate_recipe(const json& recipe) {
-    check_keys(recipe, {"schema_version", "id", "name", "description", "model_url", "slots", "required_any", "defaults", "context", "features", "disabled_features", "size_multiple", "video_constraints"});
+    check_keys(recipe, {"schema_version", "id", "name", "description", "model_url", "slots", "required_any", "models", "defaults", "context", "features", "disabled_features", "size_multiple", "video_constraints"});
     if (!recipe.at("schema_version").is_number_integer() || recipe.at("schema_version") != 1 ||
         !valid_identifier(recipe.at("id").get<std::string>()) || recipe.at("name").get<std::string>().empty()) {
         throw std::runtime_error("invalid schema_version, id or name");
@@ -449,6 +457,19 @@ static void validate_recipe(const json& recipe) {
     if (!has_required_model) {
         throw std::runtime_error("a required model or diffusion_model slot is needed");
     }
+    // Optional recipe-local model defaults: slot ID to file ID mappings that
+    // prefill slots a load request does not select itself.
+    if (recipe.contains("models")) {
+        const auto& recipe_models = recipe.at("models");
+        if (!recipe_models.is_object() || recipe_models.size() > 32)
+            throw std::runtime_error("models must be an object of at most 32 slot entries");
+        std::set<std::string> default_ids;
+        for (auto it = recipe_models.begin(); it != recipe_models.end(); ++it) {
+            if (!ids.count(it.key() == "taesd" ? "tae" : it.key()) || !it.value().is_string() || it.value().get<std::string>().empty() ||
+                !default_ids.insert(it.key() == "taesd" ? "tae" : it.key()).second)
+                throw std::runtime_error("unknown or invalid default model: " + it.key());
+        }
+    }
     if (recipe.contains("required_any")) {
         const auto& groups = recipe.at("required_any");
         if (!groups.is_array() || groups.size() > 32)
@@ -646,10 +667,11 @@ void initialize_server_models(ServerRuntime& runtime) {
     refresh_model_catalog(runtime);
 }
 
-void prepare_model_load(ServerRuntime& runtime, const json& body, SDContextParams& context, SDGenerationParams& defaults, json* loaded_recipe) {
+void prepare_model_load(ServerRuntime& runtime, json& body, SDContextParams& context, SDGenerationParams& defaults, json* loaded_recipe) {
     check_keys(body, {"recipe", "models"});
     const auto recipe_id = body.at("recipe").get<std::string>();
-    auto selection       = body.at("models");
+    // `models` is optional; unselected slots may fall back to recipe defaults below.
+    auto selection = body.value("models", json::object());
     if (!selection.is_object())
         throw std::runtime_error("models must be an object");
     auto& models = *runtime.models;
@@ -664,6 +686,20 @@ void prepare_model_load(ServerRuntime& runtime, const json& body, SDContextParam
     defaults = models.base_defaults;
     apply_recipe_context(recipe.value("context", json::object()), context);
     apply_recipe_defaults(recipe.value("defaults", json::object()), defaults);
+    // Recipe-provided default models fill any slot the request did not select.
+    // An explicit selection blocks the default for that slot, including an
+    // empty string (a deliberate opt-out of the default). The tae/taesd pair
+    // is one slot: an explicit request selection of either spelling overrides
+    // the recipe default for that slot.
+    if (recipe.contains("models")) {
+        for (auto it = recipe.at("models").begin(); it != recipe.at("models").end(); ++it) {
+            if (selection.contains(it.key()))
+                continue;
+            if ((it.key() == "tae" || it.key() == "taesd") && (selection.contains("tae") || selection.contains("taesd")))
+                continue;
+            selection[it.key()] = it.value();
+        }
+    }
     std::set<std::string> allowed;
     for (const auto& slot : recipe.at("slots")) {
         const auto id = slot.at("id").get<std::string>();
@@ -675,6 +711,9 @@ void prepare_model_load(ServerRuntime& runtime, const json& body, SDContextParam
                 if (selection.contains(id))
                     throw std::runtime_error("specify only one of tae and taesd");
                 selection[id] = selection.at(alias);
+                // Keep only the declared spelling: equivalent loads must resolve
+                // to an identical selection for the idempotent-load comparison.
+                selection.erase(alias);
             }
         }
         const std::string file_id = selection.value(id, std::string());
@@ -698,6 +737,20 @@ void prepare_model_load(ServerRuntime& runtime, const json& body, SDContextParam
         if (!allowed.count(it.key()))
             throw std::runtime_error("unknown model slot: " + it.key());
     }
+    // Omitting an optional slot and sending it as an empty string resolve to
+    // the same context; drop the empty key so equivalent loads compare
+    // identically instead of rebuilding the context. The explicit empty
+    // string already overrode any recipe default during the merge, so its
+    // intent is preserved.
+    for (auto it = selection.begin(); it != selection.end();) {
+        if (!it->is_string() || !it->get<std::string>().empty()) {
+            ++it;
+            continue;
+        }
+        it = selection.erase(it);
+    }
+    // Echo the fully resolved selection (recipe defaults merged in) for status reporting.
+    body["models"] = selection;
     for (const auto& group : recipe.value("required_any", json::array())) {
         if (std::none_of(group.begin(), group.end(), [&](const json& value) {
                 return !selection.value(value.get<std::string>(), std::string()).empty();

examples/server/model_catalog.h

@@ -51,6 +51,14 @@ void initialize_server_models(ServerRuntime& runtime);
 json resolve_recipe_features(const json& recipe);
 json model_generation_features(const ServerRuntime& runtime);
 
-void prepare_model_load(ServerRuntime& runtime, const json& body, SDContextParams& context, SDGenerationParams& defaults, json* loaded_recipe = nullptr);
+// Resolves the load request into context/generation params. `body` is updated
+// so `models` reflects the fully resolved selection (recipe defaults merged in),
+// letting the caller report exactly which components got loaded.
+void prepare_model_load(ServerRuntime& runtime, json& body, SDContextParams& context, SDGenerationParams& defaults, json* loaded_recipe = nullptr);
+
+// True when both parameter sets resolve to the same component files. File IDs
+// stay stable if a model directory symlink is retargeted inside the configured
+// root, so the idempotent-load check compares the resolved paths too.
+bool component_paths_match(const SDContextParams& a, const SDContextParams& b);
 
 #endif

examples/server/recipes/defaults.md

@@ -42,6 +42,8 @@ For community edit adapters and other optional extensions, these are the base-mo
 | `flux3-klein` | 1 | 16 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/main/README.md)  |
 | `flux3-klein-base` | 4 | 20 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/main/src/flux3/pipeline.py)  |
 | `hidream-o1-dev` | 1 | 20 | 1024 × 1024 | [Source](https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev)  |
+| `hyperflow-h3-flash-fl2va` | 1 | 6 | 1344 × 768 | [Source](https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash) Two-time FL2VA checkpoint; native text and first/last-frame generation. Fixed 6-step video sigma grid with shift 12, 124 frames at 24 FPS. Reuses MiniMax-H3 components. |
+| `hyperflow-h3-flash-ref2va` | 1 | 6 | 1344 × 768 | [Source](https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash) Separate two-time Ref2VA checkpoint; image references. Same schedule and components as FL2VA; video/audio references remain unverified. |
 | `hunyuan-video-15-i2v` | 6 | 20 | 1280 × 720 | [Source](https://huggingface.co/tencent/HunyuanVideo-1.5) 720p base I2V defaults; CFG/step-distilled and 480p checkpoints use different settings. |
 | `hunyuan-video-15-t2v` | 6 | 20 | 1280 × 720 | [Source](https://huggingface.co/tencent/HunyuanVideo-1.5) 720p base T2V defaults; CFG-distilled / 480p checkpoints need their published settings. |
 | `ideogram4` | 7 | 20 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/ideogram-oss/ideogram4/main/docs/inference.md) Constant-guidance pipeline defaults; this recipe does not reproduce the CLI’s per-step guidance schedule. |
@@ -71,6 +73,7 @@ For community edit adapters and other optional extensions, these are the base-mo
 | `minit2i` | 2.5 | 20 | 512 × 512 | [Source](https://huggingface.co/MiniT2I/MiniT2I/raw/main/pipeline.py) B/16 example and native model size; L/16 uses CFG 6. |
 | `moss-image-base` | 3 | 20 | 1024 × 1024 | [Source](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base)  |
 | `moss-image-sft` | 3 | 20 | 1024 × 1024 | [Source](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT)  |
+| `p-h3-image` | 1 | 8 | 1024 × 1024 | [Runtime guide](../../../docs/p_h3_image.md) Fixed P-H3 sigma schedule; upstream 2912 × 2912 reduced to the catalog default. |
 | `nitrosd-realism` | 1 | 1 | 1024 × 1024 | [Source](https://huggingface.co/ChenDY/NitroFusion)  |
 | `ovis-image` | 5 | 20 | 1024 × 1024 | [Source](https://huggingface.co/ATH-MaaS/Ovis-Image-7B)  |
 | `pid` | 1 | 4 | 2048 × 2048 | [Source](https://huggingface.co/nvidia/PiD) Standalone PiD decoding at 2K; guidance is handled by the upstream generator. |

examples/server/recipes/hyperflow-h3-flash-fl2va.json

@@ -0,0 +1,75 @@
+{
+  "schema_version": 1,
+  "id": "hyperflow-h3-flash-fl2va",
+  "name": "HyperFlow H3 Flash FL2VA",
+  "size_multiple": 32,
+  "description": "Generate video and audio from text, optionally with start/end images. Use the HyperFlow H3 Flash FL2VA transformer. Reuses MiniMax-H3 components with the HyperFlow 6-step schedule. Audio VAE is optional for silent video.",
+  "model_url": "https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash",
+  "slots": [
+    {
+      "id": "diffusion_model",
+      "required": true,
+      "label": "HyperFlow H3 Flash FL2VA Transformer"
+    },
+    {
+      "id": "llm",
+      "required": true,
+      "label": "MiniMax-H3 Qwen 3 VL 32B Text Encoder"
+    },
+    {
+      "id": "vae",
+      "required": false,
+      "label": "MiniMax H3 Video VAE"
+    },
+    {
+      "id": "tae",
+      "required": false
+    },
+    {
+      "id": "audio_vae",
+      "required": false,
+      "label": "MiniMax H3 Audio VAE"
+    }
+  ],
+  "required_any": [
+    [
+      "vae",
+      "tae"
+    ]
+  ],
+  "video_constraints": {
+    "frame_step": 17,
+    "frame_offset": 5,
+    "frame_rounding": "ceil",
+    "min_frames": 5,
+    "fixed_fps": 24
+  },
+  "defaults": {
+    "video_frames": 124,
+    "fps": 24,
+    "cfg_scale": 1,
+    "width": 1344,
+    "height": 768,
+    "sampling_method": "euler",
+    "steps": 6,
+    "flow_shift": 12,
+    "custom_sigmas": [
+      1.0,
+      0.982655942,
+      0.914642572,
+      0.674755692,
+      0.385223746,
+      0.159055337,
+      0.0
+    ]
+  },
+  "features": [
+    "end_image",
+    "flow_shift",
+    "init_image"
+  ],
+  "disabled_features": [
+    "img2img",
+    "mask_image"
+  ]
+}

examples/server/recipes/hyperflow-h3-flash-ref2va.json

@@ -0,0 +1,74 @@
+{
+  "schema_version": 1,
+  "id": "hyperflow-h3-flash-ref2va",
+  "name": "HyperFlow H3 Flash Ref2VA",
+  "size_multiple": 32,
+  "description": "Generate video and audio with reference images. Use the HyperFlow H3 Flash Ref2VA transformer; video/audio references are not validated for these weights. Reuses MiniMax-H3 components with the HyperFlow 6-step schedule. Audio VAE is optional for silent video.",
+  "model_url": "https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash",
+  "slots": [
+    {
+      "id": "diffusion_model",
+      "required": true,
+      "label": "HyperFlow H3 Flash REF2VA Transformer"
+    },
+    {
+      "id": "llm",
+      "required": true,
+      "label": "MiniMax-H3 Qwen 3 VL 32B Text Encoder"
+    },
+    {
+      "id": "vae",
+      "required": false,
+      "label": "MiniMax H3 Video VAE"
+    },
+    {
+      "id": "tae",
+      "required": false
+    },
+    {
+      "id": "audio_vae",
+      "required": false,
+      "label": "MiniMax H3 Audio VAE"
+    }
+  ],
+  "required_any": [
+    [
+      "vae",
+      "tae"
+    ]
+  ],
+  "video_constraints": {
+    "frame_step": 17,
+    "frame_offset": 5,
+    "frame_rounding": "ceil",
+    "min_frames": 5,
+    "fixed_fps": 24
+  },
+  "defaults": {
+    "video_frames": 124,
+    "fps": 24,
+    "cfg_scale": 1,
+    "width": 1344,
+    "height": 768,
+    "sampling_method": "euler",
+    "steps": 6,
+    "flow_shift": 12,
+    "custom_sigmas": [
+      1.0,
+      0.982655942,
+      0.914642572,
+      0.674755692,
+      0.385223746,
+      0.159055337,
+      0.0
+    ]
+  },
+  "features": [
+    "flow_shift",
+    "ref_images"
+  ],
+  "disabled_features": [
+    "img2img",
+    "mask_image"
+  ]
+}

examples/server/recipes/p-h3-image.json

@@ -0,0 +1,55 @@
+{
+  "schema_version": 1,
+  "id": "p-h3-image",
+  "name": "P-H3-Image",
+  "size_multiple": 32,
+  "description": "Supports text-to-image and reference-image editing with the fixed 8-step schedule. Use the P-H3 single-frame VAE and Qwen3-VL-4B with its vision tower. Supply the P-H3 text connector separately unless it is included in the text encoder.",
+  "model_url": "https://huggingface.co/PrunaAI/P-H3-Image",
+  "slots": [
+    {
+      "id": "diffusion_model",
+      "required": true
+    },
+    {
+      "id": "llm",
+      "required": true,
+      "label": "Qwen 3 VL 4B Text Encoder"
+    },
+    {
+      "id": "llm_vision",
+      "required": false,
+      "label": "Vision projector (required when not included in the text encoder)"
+    },
+    {
+      "id": "embeddings_connectors",
+      "required": false,
+      "label": "P-H3 Text Connector (required when not included in the text encoder)"
+    },
+    {
+      "id": "vae",
+      "required": true,
+      "label": "P-H3 Image VAE"
+    }
+  ],
+  "defaults": {
+    "sampling_method": "euler",
+    "cfg_scale": 1,
+    "steps": 8,
+    "width": 1024,
+    "height": 1024,
+    "custom_sigmas": [
+      1.0,
+      0.984753,
+      0.967905,
+      0.951688,
+      0.925372,
+      0.881923,
+      0.810083,
+      0.528718,
+      0.0
+    ]
+  },
+  "features": [
+    "ref_images"
+  ]
+}

examples/server/routes_models.cpp

@@ -84,8 +84,26 @@ static void change_model(ServerRuntime& runtime, const httplib::Request& req, ht
         }
     }
 
+    // Idempotent load: if this exact recipe and resolved selection are already
+    // active, report the current status instead of tearing down the context.
+    // Comparing full recipe objects (not just the id) also reloads when the
+    // catalog was refreshed with modified recipe metadata for the same id.
+    // Comparing resolved component paths catches the same file IDs resolving
+    // to different files (e.g. a model directory symlink retargeted inside
+    // the configured root since the catalog was last scanned).
+    if (!unload) {
+        json current = model_status(runtime);
+        if (current.value("state", "") == "ready" && current.value("recipe", json(nullptr)) == body.at("recipe") &&
+            current.value("models", json::object()) == body.at("models") && models.active_recipe == loaded_recipe &&
+            component_paths_match(*runtime.ctx_params, context)) {
+            current["already_loaded"] = true;
+            res.set_content(current.dump(), "application/json");
+            return;
+        }
+    }
+
     const json recipe    = unload ? json(nullptr) : body.at("recipe");
-    const json selection = unload ? json::object() : body.at("models");
+    const json selection = unload ? json::object() : body.value("models", json::object());
     set_model_status(runtime, "loading", recipe, selection);
     try {
         runtime.sd_ctx = nullptr;

examples/server/routes_sdcpp.cpp

@@ -154,6 +154,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"image_preprocess", defaults.image_preprocess},
         {"increase_ref_index", defaults.increase_ref_index},
         {"control_strength", defaults.control_strength},
+        {"control_type", defaults.control_type},
         {"ip_adapter_strength", defaults.ip_adapter_strength},
         {"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers, defaults.custom_sigmas)},
         {"hires", make_hires_json(defaults)},

include/stable-diffusion.h

@@ -400,6 +400,19 @@ enum sd_hires_upscaler_t {
     SD_HIRES_UPSCALER_COUNT,
 };
 
+enum sd_control_type_t {
+    SD_CONTROL_TYPE_NONE,
+    SD_CONTROL_TYPE_OPENPOSE,
+    SD_CONTROL_TYPE_DEPTH,
+    SD_CONTROL_TYPE_SCRIBBLE,  // thick lines: scribble, hed, softedge, ted
+    SD_CONTROL_TYPE_CANNY,     // thin lines: canny, lineart, anime lineart, mlsd
+    SD_CONTROL_TYPE_NORMAL,
+    SD_CONTROL_TYPE_SEGMENT,
+    SD_CONTROL_TYPE_TILE,
+    SD_CONTROL_TYPE_REPAINT,  // inpaint, outpaint
+    SD_CONTROL_TYPE_COUNT,
+};
+
 typedef struct {
     bool enabled;
     enum sd_hires_upscaler_t upscaler;
@@ -444,6 +457,7 @@ typedef struct {
     bool circular_x;
     bool circular_y;
     sd_image_preprocess_params_t image_preprocess;
+    enum sd_control_type_t control_type;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -486,6 +500,15 @@ struct ggml_tensor;
 
 typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);
 typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data);
+typedef struct {
+    const char* stage;
+    const char* detail;
+    const char* unit;
+    int sample_pass;
+    int64_t completed;
+    int64_t total;
+} sd_generation_progress_t;
+typedef void (*sd_generation_progress_cb_t)(const sd_generation_progress_t* progress, void* data);
 typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data);
 
 typedef struct {
@@ -496,6 +519,11 @@ typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void*
 
 SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
 SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
+// Context-local generation events, independent of preview and the legacy progress callback.
+// Strings and the event are valid only during the synchronous callback. Callbacks must not throw.
+// total == 0 means unknown; sample_pass counts from 1 since callback registration.
+// Register/clear only while the context is idle. Passing nullptr clears the callback.
+SD_API void sd_set_generation_progress_callback(sd_ctx_t* ctx, sd_generation_progress_cb_t cb, void* data);
 // In each sampling pass, a positive interval previews every Nth denoiser step, while a
 // negative interval previews only completed logical step -interval. Zero previews the final
 // completed step of the first sampling pass (base-resolution or high-noise).
@@ -531,6 +559,10 @@ SD_API const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode);
 SD_API enum lora_apply_mode_t str_to_lora_apply_mode(const char* str);
 SD_API const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler);
 SD_API enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str);
+SD_API const char* sd_control_type_name(enum sd_control_type_t control_type);
+// Also accepts the names of the kinds a type stands for (hed, lineart, mlsd, ...);
+// returns SD_CONTROL_TYPE_COUNT when the name is unknown.
+SD_API enum sd_control_type_t str_to_sd_control_type(const char* str);
 
 SD_API void sd_cache_params_init(sd_cache_params_t* cache_params);
 SD_API void sd_hires_params_init(sd_hires_params_t* hires_params);
@@ -597,6 +629,42 @@ SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
 // Reads model metadata only; returns 0 if the file is not a recognized RGB ESRGAN model.
 SD_API int get_upscaler_model_scale(const char* model_path);
 
+typedef struct video_upscaler_ctx_t video_upscaler_ctx_t;
+
+typedef struct {
+    const char* diffusion_model_path;
+    const char* vae_path;
+    // Explicit output scale; RelayVSR requires 2 or 4 matching both checkpoints.
+    int upscale_factor;
+    int n_threads;
+    const char* backend;
+    const char* params_backend;
+    // auto/on/off or diffusion/vae assignments; null selects auto.
+    const char* flash_attn_mode;
+    const char* conv_direct_mode;
+} sd_video_upscaler_ctx_params_t;
+
+typedef struct {
+    const sd_image_t* frames;
+    int frame_count;
+    int fps;
+    int64_t seed;
+    // RelayVSR reference window, including both endpoints; defaults to 16.
+    int window_size;
+    // SeedVR2 independent temporal chunks in input frames (1 or 4n+1); defaults to 5.
+    int temporal_chunk_size;
+} sd_video_upscale_params_t;
+
+SD_API void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params);
+SD_API void sd_video_upscale_params_init(sd_video_upscale_params_t* params);
+SD_API video_upscaler_ctx_t* new_video_upscaler_ctx(const sd_video_upscaler_ctx_params_t* params);
+SD_API void free_video_upscaler_ctx(video_upscaler_ctx_t* ctx);
+// Frames are RGB, have equal dimensions, and remain owned by the caller.
+// Free the result with free_sd_video_generation_result().
+SD_API bool upscale_video(video_upscaler_ctx_t* ctx,
+                          const sd_video_upscale_params_t* params,
+                          sd_video_generation_result_t** result_out);
+
 typedef struct adetailer_ctx_t adetailer_ctx_t;
 
 typedef struct {

scripts/ci-cuda-cache.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+mode="${1:-}"
+if [[ "$mode" != restore && "$mode" != save ]]; then
+    echo "Usage: $0 restore|save" >&2
+    exit 1
+fi
+
+cache_dir="${RUNNER_TEMP:?}/cuda-ccache"
+builder="${CUDA_CACHE_BUILDER:?}"
+work_dir="$(mktemp -d "${RUNNER_TEMP}/cuda-cache.XXXXXX")"
+trap 'rm -rf -- "$work_dir"' EXIT
+mkdir -p "$cache_dir"
+
+# Cache mounts are outside image layers; always execute the transfer RUN.
+if [[ "$mode" == restore ]]; then
+    cat > "$work_dir/Dockerfile" <<'DOCKERFILE'
+FROM ubuntu:24.04
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    --mount=type=bind,source=.,target=/cache-input \
+    cp -a /cache-input/. /root/.cache/ccache/
+DOCKERFILE
+    docker buildx build --builder "$builder" --no-cache \
+        --file "$work_dir/Dockerfile" --output type=cacheonly "$cache_dir"
+else
+    cat > "$work_dir/Dockerfile" <<'DOCKERFILE'
+FROM ubuntu:24.04 AS transfer
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    mkdir /cache-output && cp -a /root/.cache/ccache/. /cache-output/
+FROM scratch
+COPY --from=transfer /cache-output/ /
+DOCKERFILE
+    docker buildx build --builder "$builder" --no-cache \
+        --file "$work_dir/Dockerfile" \
+        --output "type=local,dest=$work_dir/export" "$work_dir"
+    # Replace only after export succeeds, including removal of evicted entries.
+    mv "$cache_dir" "$work_dir/previous-cache"
+    mv "$work_dir/export" "$cache_dir"
+fi

scripts/format-code.ps1

@@ -8,8 +8,19 @@ $patterns = @(
     "src/core/*.cpp"
     "src/core/*.h"
     "src/core/*.hpp"
+    "src/core/cpu/*.cpp"
+    "src/core/cpu/*.h"
+    "src/core/metal/*.mm"
+    "src/core/metal/*.metal"
+    "src/core/cuda/*.cu"
+    "src/core/cuda/*.cuh"
+    "src/core/cuda/*.h"
+    "src/core/cuda/sageattention/*.cuh"
     "src/core/probe/*.cpp"
     "src/core/probe/*.h"
+    "src/core/vulkan/*.cpp"
+    "src/core/vulkan/*.h"
+    "src/core/vulkan/shaders/*.comp"
     "src/extensions/*.cpp"
     "src/extensions/*.h"
     "src/extensions/*.hpp"
@@ -23,8 +34,6 @@ $patterns = @(
     "src/model/*/*.hpp"
     "src/tokenizers/*.h"
     "src/tokenizers/*.cpp"
-    "src/tokenizers/vocab/*.h"
-    "src/tokenizers/vocab/*.cpp"
     "src/model_io/*.h"
     "src/model_io/*.cpp"
     "examples/cli/*.cpp"
@@ -46,10 +55,6 @@ try {
         foreach ($file in $files) {
             $relativePath = $file.FullName.Substring($root.Length).TrimStart('\', '/') -replace '\\', '/'
 
-            if ($relativePath -like "vocab*") {
-                continue
-            }
-
             Write-Host "formatting '$relativePath'"
 
             # if ($relativePath -ne "stable-diffusion.h") {

scripts/format-code.sh

@@ -8,16 +8,20 @@ cd "$repo_root"
 for f in src/*.cpp src/*.h src/*.hpp \
          src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \
          src/core/*.cpp src/core/*.h src/core/*.hpp \
+         src/core/cpu/*.cpp src/core/cpu/*.h \
+         src/core/metal/*.mm src/core/metal/*.metal \
+         src/core/cuda/*.cu src/core/cuda/*.cuh src/core/cuda/*.h \
+         src/core/cuda/sageattention/*.cuh \
          src/core/probe/*.cpp src/core/probe/*.h \
+         src/core/vulkan/*.cpp src/core/vulkan/*.h src/core/vulkan/shaders/*.comp \
          src/extensions/*.cpp src/extensions/*.h src/extensions/*.hpp \
          src/pipeline/*.cpp src/pipeline/*.h \
          src/runtime/*.cpp src/runtime/*.h src/runtime/*.hpp \
          src/model/*/*.cpp src/model/*/*.h src/model/*/*.hpp \
-         src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \
+         src/tokenizers/*.h src/tokenizers/*.cpp \
          src/model_io/*.h src/model_io/*.cpp examples/cli/*.cpp examples/cli/*.h examples/server/*.cpp \
          examples/common/*.hpp examples/common/*.h examples/common/*.cpp; do
   [[ -e "$f" ]] || continue
-  [[ "$f" == vocab* ]] && continue
   echo "formatting '$f'"
   # if [ "$f" != "stable-diffusion.h" ]; then
   #   clang-tidy -fix -p build_linux/ "$f"

src/conditioning/conditioner.hpp

@@ -17,6 +17,7 @@
 #include "model/te/llada_image_te.hpp"
 #include "model/te/llm.hpp"
 #include "model/te/ming_image_te.hpp"
+#include "model/te/ph3_image_connector.hpp"
 #include "model/te/t5.hpp"
 #include "model_loader.h"
 #include "tokenizers/sensenova_u1_tokenizer.h"
@@ -1976,6 +1977,7 @@ struct LLMEmbedder : public Conditioner {
     std::shared_ptr<Tokenizer> tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
     std::shared_ptr<T5Runner> byt5;
+    std::shared_ptr<PH3ImageTextConnectorRunner> ph3_connector;
 
     LLMEmbedder(ggml_backend_t backend,
                 const String2TensorStorage& tensor_storage_map      = {},
@@ -2022,12 +2024,18 @@ struct LLMEmbedder : public Conditioner {
         } else if (version == VERSION_MOSS_IMAGE) {
             arch = LLM::LLMArch::QWEN3_5;
         }
-        llm        = std::make_shared<LLM::LLMRunner>(arch,
-                                                      backend,
-                                                      tensor_storage_map,
-                                                      "text_encoders.llm",
-                                                      enable_vision,
-                                                      weight_manager);
+        llm = std::make_shared<LLM::LLMRunner>(arch,
+                                               backend,
+                                               tensor_storage_map,
+                                               "text_encoders.llm",
+                                               enable_vision,
+                                               weight_manager);
+        if (version == VERSION_P_H3_IMAGE) {
+            if (llm->config.num_layers != 36 || llm->config.hidden_size != 2560) {
+                throw std::runtime_error("P-H3-Image requires the Qwen3-VL-4B text encoder");
+            }
+            ph3_connector = std::make_shared<PH3ImageTextConnectorRunner>(backend, tensor_storage_map, weight_manager);
+        }
         int pad_id = 151643;
         if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
             pad_id = 11;
@@ -2063,6 +2071,9 @@ struct LLMEmbedder : public Conditioner {
 
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
+        if (ph3_connector) {
+            ph3_connector->get_param_tensors(tensors);
+        }
         if (byt5) {
             byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
         }
@@ -2074,6 +2085,9 @@ struct LLMEmbedder : public Conditioner {
 
     void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
+        if (ph3_connector) {
+            ph3_connector->set_max_graph_vram_bytes(max_vram_bytes);
+        }
         if (byt5) {
             byt5->set_max_graph_vram_bytes(max_vram_bytes);
         }
@@ -2081,6 +2095,9 @@ struct LLMEmbedder : public Conditioner {
 
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         llm->set_runtime_backends(backends);
+        if (ph3_connector) {
+            ph3_connector->set_runtime_backends(backends);
+        }
         if (byt5) {
             byt5->set_runtime_backends(backends);
         }
@@ -2113,6 +2130,9 @@ struct LLMEmbedder : public Conditioner {
 
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
+        if (ph3_connector) {
+            ph3_connector->set_flash_attention_enabled(enabled);
+        }
         if (byt5) {
             byt5->set_flash_attention_enabled(enabled);
         }
@@ -2120,12 +2140,18 @@ struct LLMEmbedder : public Conditioner {
 
     void set_scale_overrides(float linear_scale, float attn_scale) override {
         llm->set_scale_overrides(linear_scale, attn_scale);
+        if (ph3_connector) {
+            ph3_connector->set_scale_overrides(linear_scale, attn_scale);
+        }
         if (byt5) {
             byt5->set_scale_overrides(linear_scale, attn_scale);
         }
     }
 
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        if (ph3_connector) {
+            ph3_connector->set_weight_adapter(adapter);
+        }
         if (llm) {
             llm->set_weight_adapter(adapter);
         }
@@ -2135,6 +2161,9 @@ struct LLMEmbedder : public Conditioner {
     }
 
     void runner_end() override {
+        if (ph3_connector) {
+            ph3_connector->runner_end();
+        }
         if (llm) {
             llm->runner_end();
         }
@@ -2311,6 +2340,12 @@ struct LLMEmbedder : public Conditioner {
             LOG_ERROR("LLM prompt encoding failed");
             return {};
         }
+        if (ph3_connector) {
+            hidden_states = ph3_connector->compute(n_threads, hidden_states);
+            if (hidden_states.empty()) {
+                return {};
+            }
+        }
         hidden_states = apply_token_weights(std::move(hidden_states), weights);
         GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
 
@@ -2457,8 +2492,9 @@ struct LLMEmbedder : public Conditioner {
 
         if (sd_version_is_minimax_h3(version)) {
             prompt_template_encode_start_idx = 0;
-            out_layers                       = {50};
-            prompt_attn_range                = {0, 0};
+            // The final tap is normalized; LLMRunner exposes it at num_layers + 1.
+            out_layers        = version == VERSION_P_H3_IMAGE ? std::set<int>{13, 19, 25, 31, 37} : std::set<int>{50};
+            prompt_attn_range = {0, 0};
 
             if (llm->enable_vision) {
                 const int patch_size = llm->config.vision.patch_size;
@@ -2564,7 +2600,17 @@ struct LLMEmbedder : public Conditioner {
                     }
                 }
             }
-            prompt += conditioner_params.text;
+            auto text = conditioner_params.text;
+            if (version == VERSION_P_H3_IMAGE && !image_embeds.empty()) {
+                // Preserve explicit picture targets and structured briefs from the upstream upsampler.
+                static const std::regex verbatim(R"(^\s*(Edit <Picture \d+>|subject_definitions:))");
+                if (!std::regex_search(text, verbatim)) {
+                    text = "Edit <Picture 1>. " + text;
+                } else {
+                    text = trim(text);
+                }
+            }
+            prompt += text;
         } else if (sd_version_is_hunyuan_video(version)) {
             prompt_template_encode_start_idx = 98;
             out_layers                       = {26};
@@ -3402,8 +3448,8 @@ struct LTXAVTextProjection : public GGMLBlock {
 
         auto video_projection = std::dynamic_pointer_cast<Linear>(blocks["video_aggregate_embed"]);
         auto audio_projection = std::dynamic_pointer_cast<Linear>(blocks["audio_aggregate_embed"]);
-        auto video_in         = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(4096.f / static_cast<float>(kHiddenSize)));
-        auto audio_in         = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(2048.f / static_cast<float>(kHiddenSize)));
+        auto video_in         = ggml_ext_scale(ctx, x, std::sqrt(4096.f / static_cast<float>(kHiddenSize)));
+        auto audio_in         = ggml_ext_scale(ctx, x, std::sqrt(2048.f / static_cast<float>(kHiddenSize)));
         auto video            = video_projection->forward(ctx, video_in);
         auto audio            = audio_projection->forward(ctx, audio_in);
         return ggml_concat(ctx->ggml_ctx, video, audio, 0);

src/convert.cpp

@@ -463,11 +463,15 @@ static bool load_model_into_memory(ModelLoader& model_loader,
         const ggml_type export_type = get_export_tensor_type(model_loader, storage, output_type, type_rules);
         const bool quant_src        = ggml_is_quantized(storage.type);
         const bool quant_dst        = ggml_is_quantized(export_type);
+        const bool half_src         = storage.type == GGML_TYPE_F16 || storage.type == GGML_TYPE_BF16;
         if (quant_src && quant_dst) {
             load_type = GGML_TYPE_F16;
-        } else if (quant_src || quant_dst) {
-            load_type = quant_src ? export_type : storage.type;
-        } else if (storage.type == GGML_TYPE_F32 || export_type == GGML_TYPE_F32) {
+        } else if (quant_src) {
+            load_type = export_type;
+        } else if (storage.type == GGML_TYPE_F32 || export_type == GGML_TYPE_F32 ||
+                   (half_src && export_type != storage.type)) {
+            // LoRA deltas can be smaller than half an F16/BF16 ulp of the weight. Merge in F32
+            // when the export converts the tensor again, so the merged weight is rounded once.
             load_type = GGML_TYPE_F32;
         }
         ggml_tensor* tensor = ggml_new_tensor(ctx, load_type, storage.n_dims, storage.ne);

src/core/backend_fit.cpp

@@ -42,8 +42,9 @@ namespace sd::backend_fit {
     struct Device {
         std::string name;
         std::string description;
-        int64_t free_bytes   = 0;
-        int64_t budget_bytes = 0;
+        int64_t free_bytes              = 0;
+        int64_t budget_bytes            = 0;
+        enum ggml_backend_dev_type type = GGML_BACKEND_DEVICE_TYPE_GPU;
     };
 
     enum class ParamsLocation {
@@ -133,13 +134,14 @@ namespace sd::backend_fit {
         for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
             ggml_backend_dev_t dev = ggml_backend_dev_get(i);
             const auto type        = ggml_backend_dev_type(dev);
-            if (type != GGML_BACKEND_DEVICE_TYPE_GPU &&
+            if (type != GGML_BACKEND_DEVICE_TYPE_GPU && type != GGML_BACKEND_DEVICE_TYPE_IGPU &&
                 (!include_other_devices || type == GGML_BACKEND_DEVICE_TYPE_CPU)) {
                 continue;
             }
             Device device;
             device.name        = ggml_backend_dev_name(dev);
             device.description = ggml_backend_dev_description(dev);
+            device.type        = type;
             size_t free_bytes = 0, total_bytes = 0;
             ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
             device.free_bytes = (int64_t)free_bytes;
@@ -193,14 +195,23 @@ namespace sd::backend_fit {
     }
 
     static size_t select_main_device(const std::vector<Device>& devices) {
-        size_t main_device = SIZE_MAX;
-        for (size_t di = 0; di < devices.size(); ++di) {
-            if (devices[di].budget_bytes > 0 &&
-                (main_device == SIZE_MAX || devices[di].budget_bytes > devices[main_device].budget_bytes)) {
-                main_device = di;
+        // Shared system memory must not make an integrated GPU outrank a discrete GPU.
+        for (bool discrete : {true, false}) {
+            size_t main_device = SIZE_MAX;
+            for (size_t di = 0; di < devices.size(); ++di) {
+                if ((devices[di].type == GGML_BACKEND_DEVICE_TYPE_GPU) != discrete) {
+                    continue;
+                }
+                if (devices[di].budget_bytes > 0 &&
+                    (main_device == SIZE_MAX || devices[di].budget_bytes > devices[main_device].budget_bytes)) {
+                    main_device = di;
+                }
+            }
+            if (main_device != SIZE_MAX) {
+                return main_device;
             }
         }
-        return main_device;
+        return SIZE_MAX;
     }
 
     static Plan compute_plan(const std::vector<Component>& components,

src/core/compute_workspace.cpp

@@ -1,17 +1,14 @@
 #include "core/compute_workspace.h"
 
 #include <algorithm>
-#include <cstring>
 #include <exception>
 #include <map>
 #include <unordered_map>
 #include <unordered_set>
 
+#include "core/ggml_compat.h"
 #include "core/ggml_extend_backend.h"
-#include "core/ggml_graph_cut.h"
 #include "core/util.h"
-#include "ggml-cpu.h"
-#include "ggml-impl.h"
 
 namespace sd {
     ComputeWorkspace::~ComputeWorkspace() {
@@ -28,18 +25,6 @@ namespace sd {
         }
     }
 
-    bool ComputeWorkspace::needs_scheduler(ggml_cgraph* graph) const {
-        if (!extra_backends_.empty()) {
-            return true;
-        }
-        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
-            if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph, i))) {
-                return true;
-            }
-        }
-        return false;
-    }
-
     ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) {
         std::vector<ggml_backend_t> backends{backend_};
         backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end());
@@ -63,10 +48,6 @@ namespace sd {
                                       graph_size, false, false);
     }
 
-    bool ComputeWorkspace::measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const {
-        return measurement.scheduler == needs_scheduler(graph);
-    }
-
     bool ComputeWorkspace::prepare(const Measurement& measurement) {
         GGML_ASSERT(!active_);
         if (measurement.buffers.empty()) {
@@ -74,7 +55,7 @@ namespace sd {
         }
         const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
                                        [&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); });
-        if (measurement.scheduler != (scheduler_ != nullptr) || grows) {
+        if (grows) {
             release();
         }
         return true;
@@ -88,47 +69,21 @@ namespace sd {
 
     bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
         GGML_ASSERT(!active_);
-        const bool use_scheduler = needs_scheduler(graph);
-        if (use_scheduler) {
-            if (allocator_ != nullptr) {
-                release();
-            }
-            const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
-            if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
-                release();
-                scheduler_          = make_scheduler(capacity);
-                scheduler_capacity_ = capacity;
-            }
-            if (scheduler_ == nullptr) {
-                return false;
-            }
-            ggml_backend_sched_reset(scheduler_);
-            assign_nodes(scheduler_, graph);
-            // Scheduler allocation rewrites sources. Split the execution graph only once.
-            if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
-                release();
-                return false;
-            }
-        } else {
-            if (scheduler_ != nullptr) {
-                release();
-            }
-            if (allocator_ == nullptr) {
-                allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
-            }
-            auto signature = ggml_graph_cut::graph_layout(graph, true);
-            if (signature != reservation_) {
-                if (!ggml_gallocr_reserve(allocator_, graph)) {
-                    release();
-                    return false;
-                }
-                reservation_ = std::move(signature);
-                ++reservations_;
-            }
-            if (!ggml_gallocr_alloc_graph(allocator_, graph)) {
-                release();
-                return false;
-            }
+        const size_t capacity = static_cast<size_t>(ggml_graph_n_nodes(graph) + ggml_graph_n_leafs(graph)) + 8;
+        if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
+            release();
+            scheduler_          = make_scheduler(capacity);
+            scheduler_capacity_ = capacity;
+        }
+        if (scheduler_ == nullptr) {
+            return false;
+        }
+        ggml_backend_sched_reset(scheduler_);
+        assign_nodes(scheduler_, graph);
+        // Scheduler allocation rewrites sources. Split the execution graph only once.
+        if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
+            release();
+            return false;
         }
         active_ = true;
         return true;
@@ -136,12 +91,9 @@ namespace sd {
 
     ComputeWorkspace::Measurement ComputeWorkspace::measure(
         ggml_cgraph* graph,
-        size_t direct_bytes,
         const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
-        const AssignNodes& assign_nodes) {
-        if (!needs_scheduler(graph)) {
-            return {{{backend_, direct_bytes}}, false};
-        }
+        const AssignNodes& assign_nodes,
+        bool cache_measurement) {
         std::vector<const ggml_tensor*> tensors;
         std::unordered_set<const ggml_tensor*> seen;
         auto visit = [&](const ggml_tensor* tensor) {
@@ -149,11 +101,11 @@ namespace sd {
                 tensors.push_back(tensor);
             }
         };
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            visit(graph->nodes[i]);
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            visit(ggml_graph_node(graph, i));
         }
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            visit(graph->leafs[i]);
+        for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) {
+            visit(ggml_graph_leaf(graph, i));
         }
         for (size_t i = 0; i < tensors.size(); ++i) {
             visit(tensors[i]->view_src);
@@ -161,13 +113,58 @@ namespace sd {
                 visit(source);
             }
         }
+        std::vector<uint8_t> key;
+        if (cache_measurement) {
+            key.reserve(tensors.size() * sizeof(ggml_tensor));
+            auto append = [&](const auto& value) {
+                const auto* bytes = reinterpret_cast<const uint8_t*>(&value);
+                key.insert(key.end(), bytes, bytes + sizeof(value));
+            };
+            append(ggml_graph_n_nodes(graph));
+            append(ggml_graph_n_leafs(graph));
+            for (auto tensor : tensors) {
+                append(tensor);
+                append(tensor->type);
+                append(tensor->ne);
+                append(tensor->nb);
+                append(tensor->op);
+                append(tensor->op_params);
+                append(tensor->flags);
+                append(tensor->src);
+                append(tensor->view_src);
+                append(tensor->view_offs);
+                append(tensor->name);
+                const auto external = external_backend(tensor);
+                append(external);
+                if (external == nullptr || tensor->view_src != nullptr) {
+                    append(tensor->buffer);
+                    append(tensor->data);
+                    append(tensor->extra);
+                    if (tensor->buffer != nullptr) {
+                        const auto usage = ggml_backend_buffer_get_usage(tensor->buffer);
+                        const auto type  = ggml_backend_buffer_get_type(tensor->buffer);
+                        append(usage);
+                        append(type);
+                    }
+                } else {
+                    append(tensor->buffer);
+                    append(tensor->data);
+                }
+                const int32_t uses = ggml_graph_use_count(graph, tensor);
+                append(uses);
+            }
+            // Match the full allocation metadata, including topology and external placement.
+            if (key == measurement_key_) {
+                return {measured_buffers_};
+            }
+        }
         const size_t graph_size = tensors.size() + 8;
         auto context            = ggml_init({tensors.size() * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true});
         if (context == nullptr) {
             return {};
         }
         std::unordered_map<const ggml_tensor*, ggml_tensor*> copies;
-        std::map<ggml_backend_t, ggml_backend_buffer_t> external_buffers;
+        std::map<std::pair<ggml_backend_t, const void*>, ggml_backend_buffer_t> external_buffers;
         for (auto tensor : tensors) {
             auto copy      = ggml_dup_tensor(context, tensor);
             *copy          = *tensor;
@@ -182,29 +179,41 @@ namespace sd {
             }
             auto external = external_backend(source);
             if (external != nullptr && source->view_src == nullptr) {
-                auto& buffer = external_buffers[external];
+                // Preserve distinct external allocations so overlap checks do not invent aliases.
+                const void* storage = source->buffer ? static_cast<const void*>(source->buffer) : source;
+                auto& buffer        = external_buffers[{external, storage}];
                 if (buffer == nullptr) {
                     buffer = ggml_backend_alloc_buffer(external, 0);
                     GGML_ASSERT(buffer != nullptr);
                     ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
                 }
                 copy->buffer = buffer;
-                copy->data   = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
-                copy->extra  = nullptr;
+                // Capability checks inspect alignment even before weights are resident.
+                copy->data  = source->data ? source->data : reinterpret_cast<void*>(ggml_backend_get_alignment(external));
+                copy->extra = nullptr;
             }
         }
-        auto copy_graph     = ggml_new_graph_custom(context, graph_size, false);
-        copy_graph->n_nodes = graph->n_nodes;
-        copy_graph->n_leafs = graph->n_leafs;
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            copy_graph->nodes[i] = copies.at(graph->nodes[i]);
-        }
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            copy_graph->leafs[i] = copies.at(graph->leafs[i]);
+        for (const auto& entry : copies) {
+            const auto* source = entry.first;
+            if (source->view_src && source->buffer) {
+                const auto* owner = source->view_src;
+                while (owner->view_src) {
+                    owner = owner->view_src;
+                }
+                if (external_backend(owner) != nullptr) {
+                    entry.second->buffer = copies.at(owner)->buffer;
+                }
+            }
         }
+        auto copy_graph = ggml_new_graph_custom(context, graph_size, false);
+        auto map_tensor = [](const ggml_tensor* tensor, void* user_data) -> ggml_tensor* {
+            const auto& mapped = *static_cast<const std::unordered_map<const ggml_tensor*, ggml_tensor*>*>(user_data);
+            const auto found   = mapped.find(tensor);
+            return found == mapped.end() ? nullptr : found->second;
+        };
+        GGML_ASSERT(ggml_graph_cpy_mapped(graph, copy_graph, map_tensor, &copies));
         Measurement result;
-        result.scheduler = true;
-        auto scheduler   = make_scheduler(graph_size);
+        auto scheduler = make_scheduler(graph_size);
         if (scheduler != nullptr) {
             assign_nodes(scheduler, copy_graph);
             std::vector<size_t> sizes(extra_backends_.size() + 2);
@@ -219,6 +228,10 @@ namespace sd {
             ggml_backend_buffer_free(entry.second);
         }
         ggml_free(context);
+        if (cache_measurement && !result.buffers.empty()) {
+            measurement_key_  = std::move(key);
+            measured_buffers_ = result.buffers;
+        }
         return result;
     }
 
@@ -253,12 +266,11 @@ namespace sd {
         if (active_) {
             return false;
         }
-        ggml_gallocr_free(allocator_);
-        allocator_ = nullptr;
+        measurement_key_.clear();
+        measured_buffers_.clear();
         ggml_backend_sched_free(scheduler_);
         scheduler_          = nullptr;
         scheduler_capacity_ = 0;
-        reservation_.clear();
         return true;
     }
 
@@ -266,6 +278,6 @@ namespace sd {
         if (scheduler_ != nullptr) {
             return ggml_backend_sched_get_buffer_size(scheduler_, backend);
         }
-        return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0;
+        return 0;
     }
 }

src/core/compute_workspace.h

@@ -4,7 +4,6 @@
 #include <functional>
 #include <vector>
 
-#include "ggml-alloc.h"
 #include "ggml-backend.h"
 
 namespace sd {
@@ -17,20 +16,17 @@ namespace sd {
         ggml_backend_t backend_;
         std::vector<ggml_backend_t> extra_backends_;
         ggml_backend_t cpu_backend_     = nullptr;
-        ggml_gallocr_t allocator_       = nullptr;
         ggml_backend_sched_t scheduler_ = nullptr;
         size_t scheduler_capacity_      = 0;
-        std::vector<uint64_t> reservation_;
-        bool active_         = false;
-        size_t reservations_ = 0;
+        bool active_                    = false;
+        std::vector<uint8_t> measurement_key_;
+        std::vector<BackendBufferSize> measured_buffers_;
 
         ggml_backend_sched_t make_scheduler(size_t graph_size);
-        bool needs_scheduler(ggml_cgraph* graph) const;
 
     public:
         struct Measurement {
             std::vector<BackendBufferSize> buffers;
-            bool scheduler = false;
         };
         using AssignNodes = std::function<void(ggml_backend_sched_t, ggml_cgraph*)>;
 
@@ -41,15 +37,14 @@ namespace sd {
         ComputeWorkspace& operator=(const ComputeWorkspace&) = delete;
 
         void set_extra_backends(const std::vector<ggml_backend_t>& backends);
-        bool measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const;
         bool prepare(const Measurement& measurement);
         bool release_excess(const Measurement& measurement);
         bool allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes);
         Measurement measure(
             ggml_cgraph* graph,
-            size_t direct_bytes,
             const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
-            const AssignNodes& assign_nodes);
+            const AssignNodes& assign_nodes,
+            bool cache_measurement = false);
         void synchronize() const;
         bool segment_end() noexcept;
         bool release();
@@ -57,7 +52,6 @@ namespace sd {
         ggml_backend_sched_t scheduler() const { return scheduler_; }
         ggml_backend_t cpu_backend() const { return cpu_backend_; }
         size_t bytes(ggml_backend_t backend) const;
-        size_t reservation_count() const { return reservations_; }
     };
 }
 

src/core/cpu/int8.cpp

@@ -0,0 +1,181 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+
+namespace sd_cpu {
+
+    static void transform(float* values, int group) {
+        const float scale = 1.f / std::sqrt(float(group));
+        for (int i = 0; i < group; ++i)
+            values[i] *= scale;
+        for (int stride = 1; stride < group; stride *= 4)
+            for (int base = 0; base < group; base += 4 * stride)
+                for (int j = 0; j < stride; ++j) {
+                    const int i   = base + j;
+                    const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride];
+                    values[i]              = a + b + c - d;
+                    values[i + stride]     = a + b - c + d;
+                    values[i + 2 * stride] = a - b + c + d;
+                    values[i + 3 * stride] = -a + b + c + d;
+                }
+    }
+    struct I8Execution {
+        SDI8Operation op;
+        int64_t rows;
+        int8_t* quant;
+        float* scales;
+        float* output;
+    };
+    static void quantize(ggml_tensor*, int ith, int nth, void* data) {
+        auto& e           = *static_cast<I8Execution*>(data);
+        const int64_t k   = e.op.input->ne[0];
+        const int group   = e.op.group_size;
+        const auto* input = static_cast<const float*>(e.op.input->data);
+        for (int64_t row = ith; row < e.rows; row += nth) {
+            const float* src = input + row * k;
+            int8_t* dst      = e.quant + row * k;
+            float amax       = 0.f, values[256];
+            auto visit       = [&](auto fn) { if (!group) { for (int64_t i = 0; i < k; ++i) fn(i, src[i]); }
+            else for (int64_t base = 0; base < k; base += group) { std::memcpy(values, src + base, group * sizeof(float)); transform(values, group); for (int j = 0; j < group; ++j) fn(base + j, values[j]); } };
+            visit([&](int64_t, float value) { amax = std::max(amax, std::fabs(value)); });
+            const float scale = amax / 127.f, inverse = scale == 0.f ? 0.f : 1.f / scale;
+            e.scales[row] = scale;
+            if (scale == 0.f)
+                std::memset(dst, 0, k);
+            else
+                visit([&](int64_t i, float value) { dst[i] = static_cast<int8_t>(std::max(-127L, std::min(127L, std::lrintf(value * inverse)))); });
+        }
+    }
+    static void multiply(ggml_tensor*, int ith, int nth, void* data) {
+        auto& e         = *static_cast<I8Execution*>(data);
+        const auto& op  = e.op;
+        const int64_t k = op.weight->ne[0], n = op.weight->ne[1];
+        const auto* w    = static_cast<const int8_t*>(op.weight->data);
+        const auto* ws   = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
+        const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+        for (int64_t index = ith; index < e.rows * n; index += nth) {
+            const int64_t row = index / n, column = index % n;
+            uint32_t sum = 0;
+            for (int64_t i = 0; i < k; ++i)
+                sum += uint32_t(int32_t(w[column * k + i]) * int32_t(e.quant[row * k + i]));
+            int32_t signed_sum;
+            std::memcpy(&signed_sum, &sum, sizeof(sum));
+            float value = float(signed_sum) * e.scales[row];
+            if (ws)
+                value *= ws[column];
+            if (bias)
+                value += bias[column];
+            e.output[index] = value;
+        }
+    }
+    static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
+        if (!src)
+            return false;
+        const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(src->data);
+        return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src);
+    }
+    static size_t quant_workspace(const ggml_tensor* node, const SDI8Operation& op) {
+        return op.input->type == GGML_TYPE_F32 ? (size_t(ggml_nelements(op.input)) + 3) / 4 * 4 + size_t(ggml_nrows(node)) * 4 : 0;
+    }
+    static bool i8_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        SDI8Operation op;
+        return sd_ggml_i8_describe(node, &op);
+    }
+    static bool convrot_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        int group;
+        return sd_ggml_convrot_describe(node, &group);
+    }
+    static ggml_status i8_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        SDI8Operation op;
+        if (!sd_ggml_i8_describe(nodes[0], &op))
+            return GGML_STATUS_FAILED;
+        *bytes = quant_workspace(nodes[0], op) + ggml_nbytes(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static ggml_status convrot_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        *bytes = ggml_nbytes(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static void destroy(void*, void*) {}
+    static ggml_ext_cpu_parallel_fn parallel(ggml_backend_t backend) {
+        return reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)), "ggml_ext_cpu_parallel_v1"));
+    }
+    static ggml_status i8_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+        auto* dst     = nodes[0];
+        I8Execution e = {};
+        if (!sd_ggml_i8_describe(dst, &e.op))
+            return GGML_STATUS_FAILED;
+        e.rows   = ggml_nrows(dst);
+        auto run = parallel(context->backend);
+        if (!run)
+            return GGML_STATUS_FAILED;
+        char* scratch           = static_cast<char*>(ggml_backend_buffer_get_base(context->workspace));
+        const size_t quant_size = quant_workspace(dst, e.op);
+        if (e.op.input->type == GGML_TYPE_F32) {
+            e.quant     = reinterpret_cast<int8_t*>(scratch);
+            e.scales    = reinterpret_cast<float*>(scratch + (size_t(ggml_nelements(e.op.input)) + 3) / 4 * 4);
+            auto status = run(context->backend, dst, quantize, &e);
+            if (status != GGML_STATUS_SUCCESS)
+                return status;
+        } else {
+            e.quant  = static_cast<int8_t*>(e.op.input->data);
+            e.scales = e.op.packed ? reinterpret_cast<float*>(e.quant + e.op.input->ne[0] * ((e.rows + 3) / 4 * 4)) : static_cast<float*>(e.op.input_scale->data);
+        }
+        bool alias = false;
+        for (auto* input : dst->src)
+            alias |= overlaps(dst, input);
+        e.output    = alias ? reinterpret_cast<float*>(scratch + quant_size) : static_cast<float*>(dst->data);
+        auto status = run(context->backend, dst, multiply, &e);
+        if (status == GGML_STATUS_SUCCESS && alias)
+            std::memcpy(dst->data, e.output, ggml_nbytes(dst));
+        return status;
+    }
+    static ggml_status convrot_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+        auto* dst     = nodes[0];
+        I8Execution e = {};
+        if (!sd_ggml_convrot_describe(dst, &e.op.group_size))
+            return GGML_STATUS_FAILED;
+        auto run = parallel(context->backend);
+        if (!run)
+            return GGML_STATUS_FAILED;
+        e.op.input      = dst->src[0];
+        e.rows          = ggml_nrows(e.op.input);
+        bool alias      = overlaps(dst, e.op.input);
+        e.quant         = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(context->workspace) : dst->data);
+        const int64_t k = dst->ne[0], padded = (e.rows + 3) / 4 * 4;
+        e.scales    = reinterpret_cast<float*>(e.quant + k * padded);
+        auto status = run(context->backend, dst, quantize, &e);
+        if (status == GGML_STATUS_SUCCESS) {
+            std::memset(e.quant + k * e.rows, 0, k * (padded - e.rows));
+            if (alias)
+                std::memcpy(dst->data, e.quant, ggml_nbytes(dst));
+        }
+        return status;
+    }
+    void register_i8(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc d = {};
+        d.abi_version          = GGML_EXT_ABI_VERSION;
+        d.priority             = 100;
+        d.destroy              = destroy;
+        d.name                 = "sd.linear_i8.cpu.v1";
+        d.op                   = GGML_OP_CUSTOM;
+        d.custom_op            = sd_ggml_i8_op();
+        d.supports             = i8_supports;
+        d.prepare              = i8_prepare;
+        d.execute              = i8_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+        d.name      = "sd.quantize_i8_convrot.cpu.v1";
+        d.op        = GGML_OP_CUSTOM;
+        d.custom_op = sd_ggml_convrot_op();
+        d.supports  = convrot_supports;
+        d.prepare   = convrot_prepare;
+        d.execute   = convrot_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    }
+
+}
+#endif

src/core/cpu/modulate.cpp

@@ -0,0 +1,82 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <cstring>
+
+namespace sd_cpu {
+    static bool modulate_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        return sd_ggml_modulate_describe(node);
+    }
+
+    static float modulate_value(const ggml_tensor* t, const int64_t* index) {
+        size_t offset = 0;
+        for (int d = 0; d < GGML_MAX_DIMS; ++d) {
+            offset += (index[d] % t->ne[d]) * t->nb[d];
+        }
+        return *reinterpret_cast<const float*>(static_cast<const char*>(t->data) + offset);
+    }
+
+    static void modulate_compute(ggml_tensor* dst, int ith, int nth, void*) {
+        auto* output       = static_cast<float*>(dst->data);
+        const int64_t rows = ggml_nrows(dst);
+        for (int64_t row = ith; row < rows; row += nth) {
+            int64_t index[] = {0, row % dst->ne[1], (row / dst->ne[1]) % dst->ne[2], row / (dst->ne[1] * dst->ne[2])};
+            for (int64_t col = 0; col < dst->ne[0]; ++col) {
+                index[0]      = col;
+                const float x = modulate_value(dst->src[0], index);
+                // Match separate ggml MUL/ADD rounding even with contraction enabled.
+                volatile float product         = x * modulate_value(dst->src[1], index);
+                volatile float scaled          = x + product;
+                output[row * dst->ne[0] + col] = dst->src[2] ? scaled + modulate_value(dst->src[2], index) : scaled;
+            }
+        }
+    }
+
+    static ggml_status modulate_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        *bytes = sd_ggml_modulate_workspace_size(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+
+    static void modulate_destroy(void*, void*) {}
+
+    static ggml_status modulate_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+        auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(context->backend));
+        auto run = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+        GGML_ASSERT(run);
+        auto* dst          = nodes[0];
+        const auto address = reinterpret_cast<uintptr_t>(dst->data);
+        bool alias         = false;
+        for (const auto* input : dst->src) {
+            if (input) {
+                const auto src = reinterpret_cast<uintptr_t>(input->data);
+                alias |= address < src + ggml_nbytes(input) && src < address + ggml_nbytes(dst);
+            }
+        }
+        if (!alias) {
+            return run(context->backend, dst, modulate_compute, nullptr);
+        }
+        ggml_tensor temporary = *dst;
+        temporary.data        = ggml_backend_buffer_get_base(context->workspace);
+        const auto status     = run(context->backend, &temporary, modulate_compute, nullptr);
+        if (status == GGML_STATUS_SUCCESS) {
+            std::memcpy(dst->data, temporary.data, ggml_nbytes(dst));
+        }
+        return status;
+    }
+
+    void register_modulate(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.modulate.cpu.v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_modulate_op();
+        desc.priority             = 100;
+        desc.supports             = modulate_supports;
+        desc.prepare              = modulate_prepare;
+        desc.destroy              = modulate_destroy;
+        desc.execute              = modulate_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+    }
+}
+#endif

src/core/cpu/registration.cpp

@@ -0,0 +1,14 @@
+#include "core/cpu/registration.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+namespace sd_cpu {
+
+    void register_extensions(ggml_backend_dev_t device) {
+        register_i8(device);
+        register_w4(device);
+        register_rope(device);
+        register_modulate(device);
+    }
+
+}
+#endif

src/core/cpu/registration.h

@@ -0,0 +1,16 @@
+#ifndef __SD_CORE_CPU_REGISTRATION_H__
+#define __SD_CORE_CPU_REGISTRATION_H__
+
+#include "ggml-backend.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+namespace sd_cpu {
+    void register_extensions(ggml_backend_dev_t device);
+    void register_rope(ggml_backend_dev_t device);
+    void register_modulate(ggml_backend_dev_t device);
+    void register_i8(ggml_backend_dev_t device);
+    void register_w4(ggml_backend_dev_t device);
+}
+#endif
+
+#endif

src/core/cpu/rope.cpp

@@ -0,0 +1,139 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <cmath>
+#include <cstring>
+
+namespace sd_cpu {
+
+    static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        return sd_ggml_rope_describe(node);
+    }
+
+    struct RopeExecution {
+        const ggml_tensor* rope;
+        const ggml_tensor* norm;
+        const ggml_tensor* weight;
+    };
+
+    static void rope_compute(ggml_tensor* dst, int ith, int nth, void* user_data) {
+        const auto& execution = *static_cast<const RopeExecution*>(user_data);
+        const auto* rope      = execution.rope;
+        const auto* x         = execution.norm ? execution.norm->src[0] : rope->src[0];
+        const auto* pe        = rope->src[1];
+        int32_t interleaved;
+        std::memcpy(&interleaved, ggml_ext_op_attributes(rope), sizeof(interleaved));
+        float eps = 0.0f;
+        if (execution.norm) {
+            std::memcpy(&eps, execution.norm->op_params, sizeof(eps));
+        }
+        const int64_t half_dim = pe->ne[2];
+        const int64_t rows     = x->ne[1] * x->ne[2] * x->ne[3];
+        for (int64_t row = ith; row < rows; row += nth) {
+            const int64_t token = row % x->ne[2];
+            const int64_t head  = (row / x->ne[2]) % x->ne[1];
+            const int64_t batch = row / (x->ne[2] * x->ne[1]);
+            const char* input   = static_cast<const char*>(x->data) + head * x->nb[1] + token * x->nb[2] + batch * x->nb[3];
+            const char* pos     = static_cast<const char*>(pe->data) + token * pe->nb[3];
+            float* output       = static_cast<float*>(dst->data) + row * x->ne[0];
+            float scale         = 1.0f;
+            if (execution.norm) {
+                double sum = 0.0;
+                for (int64_t col = 0; col < x->ne[0]; ++col) {
+                    const float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
+                    sum += static_cast<double>(v * v);
+                }
+                const float mean = static_cast<float>(sum / x->ne[0]);
+                scale            = 1.0f / std::sqrt(mean + eps);
+            }
+            auto value = [&](int64_t col) {
+                float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
+                if (execution.norm) {
+                    const auto* weight = static_cast<const char*>(execution.weight->data);
+                    v                  = (v * scale) * *reinterpret_cast<const float*>(weight + col * execution.weight->nb[0]);
+                }
+                return v;
+            };
+            for (int64_t j = 0; j < half_dim; ++j) {
+                const int64_t i0 = interleaved ? 2 * j : j;
+                const int64_t i1 = interleaved ? 2 * j + 1 : j + half_dim;
+                const float a    = value(i0);
+                const float b    = value(i1);
+                const char* m    = pos + j * pe->nb[2];
+                const float c    = *reinterpret_cast<const float*>(m);
+                const float s    = *reinterpret_cast<const float*>(m + pe->nb[1]);
+                if (interleaved) {
+                    const float neg_s = *reinterpret_cast<const float*>(m + pe->nb[0]);
+                    const float c1    = *reinterpret_cast<const float*>(m + pe->nb[0] + pe->nb[1]);
+                    output[i0]        = a * c + b * neg_s;
+                    output[i1]        = a * s + b * c1;
+                } else {
+                    output[i0] = a * c - b * s;
+                    output[i1] = a * s + b * c;
+                }
+            }
+            for (int64_t col = 2 * half_dim; col < x->ne[0]; ++col) {
+                output[col] = value(col);
+            }
+        }
+    }
+
+    static int rope_fusion_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+        auto* view = const_cast<ggml_cgraph*>(graph);
+        if (first + 2 >= ggml_graph_n_nodes(view)) {
+            return 0;
+        }
+        const auto* norm = ggml_graph_node(view, first);
+        const auto* mul  = ggml_graph_node(view, first + 1);
+        const auto* rope = ggml_graph_node(view, first + 2);
+        if (norm->op != GGML_OP_RMS_NORM || norm->type != GGML_TYPE_F32 || !norm->src[0] ||
+            norm->src[0]->type != GGML_TYPE_F32 || norm->src[0]->nb[0] != sizeof(float) ||
+            mul->op != GGML_OP_MUL || mul->type != GGML_TYPE_F32 || rope->src[0] != mul ||
+            ggml_ext_op_get(rope) != sd_ggml_rope_op() ||
+            !sd_ggml_rope_describe(rope) || !ggml_are_same_shape(norm, mul)) {
+            return 0;
+        }
+        const auto* weight = mul->src[0] == norm ? mul->src[1] : (mul->src[1] == norm ? mul->src[0] : nullptr);
+        if (!weight || weight->type != GGML_TYPE_F32 || weight->ne[0] != norm->ne[0] ||
+            weight->ne[1] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1) {
+            return 0;
+        }
+        return 3;
+    }
+
+    static ggml_status rope_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int count, void*, void*) {
+        auto device = ggml_backend_get_device(context->backend);
+        auto reg    = ggml_backend_dev_backend_reg(device);
+        auto run    = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+        GGML_ASSERT(run);
+        RopeExecution execution = {nodes[count - 1], nullptr, nullptr};
+        if (count == 3) {
+            execution.norm   = nodes[0];
+            execution.weight = nodes[1]->src[0] == nodes[0] ? nodes[1]->src[1] : nodes[1]->src[0];
+        }
+        return run(context->backend, nodes[count - 1], rope_compute, &execution);
+    }
+
+    void register_rope(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.rope_apply.cpu.v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_rope_op();
+        desc.priority             = 100;
+        desc.supports             = rope_supports;
+        desc.execute              = rope_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+
+        desc.name      = "sd.rms_norm_mul_rope.cpu.v1";
+        desc.op        = GGML_OP_RMS_NORM;
+        desc.custom_op = nullptr;
+        desc.priority  = 200;
+        desc.supports  = nullptr;
+        desc.match     = rope_fusion_match;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+    }
+
+}
+#endif

src/core/cpu/w4.cpp

@@ -0,0 +1,95 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <cstring>
+
+namespace sd_cpu {
+    struct W4Execution {
+        SDW4Operation op;
+        float* output;
+        int64_t rows;
+    };
+
+    static void multiply_w4(ggml_tensor*, int ith, int nth, void* data) {
+        auto& e         = *static_cast<W4Execution*>(data);
+        const auto& op  = e.op;
+        const int64_t k = op.input->ne[0], n = op.weight->ne[1];
+        const auto* w       = static_cast<const uint8_t*>(op.weight->data);
+        const auto* x       = static_cast<const int8_t*>(op.input->data);
+        const auto* xs      = reinterpret_cast<const float*>(x + k * ((e.rows + 3) / 4 * 4));
+        const auto* ws      = static_cast<const float*>(op.weight_scale->data);
+        const auto* bias    = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+        const auto* sc      = op.s_channel ? static_cast<const float*>(op.s_channel->data) : nullptr;
+        const auto* sr      = op.s_rel ? static_cast<const uint8_t*>(op.s_rel->data) : nullptr;
+        const auto to_float = ggml_get_type_traits(GGML_TYPE_F8_E4M3)->to_float;
+        for (int64_t index = ith; index < n * e.rows; index += nth) {
+            const int64_t row = index / n, col = index % n;
+            const auto* packed = w + col * (k / 2);
+            float value;
+            if (op.kind == SDW4Kind::W4A4) {
+                uint32_t sum = 0;
+                for (int64_t i = 0; i < k; ++i) {
+                    const int nibble = (packed[i / 2] >> (4 * (i % 2))) & 15;
+                    sum += uint32_t(((nibble ^ 8) - 8) * int32_t(x[row * k + i]));
+                }
+                int32_t signed_sum;
+                std::memcpy(&signed_sum, &sum, sizeof(sum));
+                value = float(signed_sum) * xs[row] * ws[col];
+            } else {
+                float sum = 0.f;
+                for (int64_t g = 0; g < k / 16; ++g) {
+                    float relative;
+                    to_float(sr + col * (k / 16) + g, &relative, 1);
+                    const float scale = sc[col] * relative;
+                    for (int j = 0; j < 16; ++j) {
+                        const int64_t i = g * 16 + j;
+                        const int code  = (packed[i / 2] >> (4 * (i % 2))) & 15;
+                        sum += (ws[code] * scale) * float(x[row * k + i]);
+                    }
+                }
+                value = sum * xs[row];
+            }
+            e.output[index] = value + (bias ? bias[col] : 0.f);
+        }
+    }
+
+    static bool w4_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        SDW4Operation op;
+        return sd_ggml_w4_describe(node, &op);
+    }
+    static ggml_status w4_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        *bytes = ggml_nbytes(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static ggml_status w4_execute(const ggml_ext_execution_context* ctx, ggml_tensor* const* nodes, int, void*, void*) {
+        auto* dst     = nodes[0];
+        W4Execution e = {};
+        if (!sd_ggml_w4_describe(dst, &e.op))
+            return GGML_STATUS_FAILED;
+        auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(ctx->backend));
+        auto run = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+        if (!run)
+            return GGML_STATUS_FAILED;
+        e.rows      = ggml_nrows(dst);
+        e.output    = static_cast<float*>(ggml_backend_buffer_get_base(ctx->workspace));
+        auto status = run(ctx->backend, dst, multiply_w4, &e);
+        if (status == GGML_STATUS_SUCCESS)
+            std::memcpy(dst->data, e.output, ggml_nbytes(dst));
+        return status;
+    }
+    void register_w4(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc d = {};
+        d.abi_version          = GGML_EXT_ABI_VERSION;
+        d.name                 = "sd.linear_w4_convrot.cpu.v1";
+        d.op                   = GGML_OP_CUSTOM;
+        d.custom_op            = sd_ggml_w4_op();
+        d.priority             = 100;
+        d.supports             = w4_supports;
+        d.prepare              = w4_prepare;
+        d.execute              = w4_execute;
+        d.destroy              = [](void*, void*) {};
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    }
+}
+#endif

src/core/cuda/CMakeLists.txt

@@ -0,0 +1,93 @@
+cmake_policy(SET CMP0079 NEW)
+
+function(sd_cuda_sources VENDOR LANGUAGE)
+    set(SOURCES)
+    foreach(OP rope int8 w4 convrot affine modulate bf16 rms_affine registration)
+        set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu")
+        string(TOUPPER "${VENDOR}" UPPER_VENDOR)
+        file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n")
+        set_source_files_properties("${SOURCE}" PROPERTIES LANGUAGE ${LANGUAGE})
+        list(APPEND SOURCES "${SOURCE}")
+    endforeach()
+    set(TARGET "sd-${VENDOR}")
+    add_library(${TARGET} OBJECT ${SOURCES})
+    target_include_directories(${TARGET} PRIVATE "${PROJECT_SOURCE_DIR}/src")
+    target_link_libraries(${TARGET} PRIVATE ggml)
+    set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
+    target_compile_features(${TARGET} PRIVATE cxx_std_17)
+    target_link_libraries(${SD_LIB} PRIVATE ${TARGET})
+    set(SD_CUDA_TARGET ${TARGET} PARENT_SCOPE)
+    set(SD_CUDA_SOURCES ${SOURCES} PARENT_SCOPE)
+endfunction()
+
+if(SD_CUDA)
+    enable_language(CUDA)
+    sd_cuda_sources(cuda CUDA)
+    target_sources(${SD_CUDA_TARGET} PRIVATE sage_attention.cu sol_attention.cu)
+    set_source_files_properties(sage_attention.cu sol_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
+    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math")
+    if(SD_CUDA_CUTLASS_PATH)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_W4_CUTLASS)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${SD_CUDA_CUTLASS_PATH}/include")
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS --expt-relaxed-constexpr)
+        if(MSVC)
+            set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus)
+        endif()
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/bf16_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_BF16_CUTLASS)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/registration_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${PROJECT_SOURCE_DIR}/src/core/cuda")
+    endif()
+    set_target_properties(${SD_CUDA_TARGET} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
+    target_link_libraries(${SD_CUDA_TARGET} PRIVATE CUDA::cudart CUDA::cublas)
+    target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas)
+endif()
+
+if(SD_HIPBLAS)
+    list(APPEND CMAKE_PREFIX_PATH "$ENV{ROCM_PATH}" /opt/rocm /usr)
+    find_package(hip REQUIRED)
+    find_package(hipblas REQUIRED)
+    if(WIN32 OR CMAKE_CXX_COMPILER MATCHES "hipcc")
+        sd_cuda_sources(hip CXX)
+        target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::device)
+    else()
+        sd_cuda_sources(hip HIP)
+        set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON)
+    endif()
+    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math")
+    # W4A8 WMMA uses the wave64 fragment layout on both GFX11 and GFX12.
+    set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_hip.cu" APPEND PROPERTY COMPILE_OPTIONS "-mwavefrontsize64")
+    target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas)
+    target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas)
+endif()
+
+if(SD_MUSA)
+    if(NOT MUSA_PATH)
+        if(DEFINED ENV{MUSA_PATH})
+            set(MUSA_PATH "$ENV{MUSA_PATH}")
+        elseif(EXISTS /opt/musa)
+            set(MUSA_PATH /opt/musa)
+        else()
+            set(MUSA_PATH /usr/local/musa)
+        endif()
+    endif()
+    list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
+    find_package(MUSAToolkit REQUIRED)
+    set(CMAKE_CXX_COMPILER "${MUSA_PATH}/bin/clang++")
+    sd_cuda_sources(musa CXX)
+    set(FLAGS -Od3 -fno-strict-aliasing -ffast-math -fsigned-char -x musa -mtgpu -fmusa-flush-denormals-to-zero)
+    if(NOT DEFINED MUSA_ARCHITECTURES)
+        set(MUSA_ARCHITECTURES "22;31")
+    endif()
+    foreach(ARCH ${MUSA_ARCHITECTURES})
+        list(APPEND FLAGS "--cuda-gpu-arch=mp_${ARCH}")
+    endforeach()
+    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "${FLAGS}")
+    if(GGML_STATIC)
+        target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart_static MUSA::mublas_static)
+        target_link_libraries(${SD_LIB} PRIVATE MUSA::musart_static MUSA::mublas_static)
+    else()
+        target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart MUSA::mublas)
+        target_link_libraries(${SD_LIB} PRIVATE MUSA::musart MUSA::mublas)
+    endif()
+endif()

src/core/cuda/affine.cu

@@ -0,0 +1,183 @@
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include "ggml-ext.h"
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_affine_cuda {
+        static bool disabled() {
+            static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+            return value;
+        }
+        static bool aligned(const ggml_tensor* tensor, size_t alignment) {
+            size_t offset = 0;
+            while (!tensor->data && tensor->view_src) {
+                offset += tensor->view_offs;
+                tensor = tensor->view_src;
+            }
+            return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
+        }
+        static float param(const ggml_tensor* tensor, int index) {
+            float value;
+            std::memcpy(&value, tensor->op_params + index, sizeof(value));
+            return value;
+        }
+        static bool contiguous_f32(const ggml_tensor* tensor) {
+            return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16);
+        }
+        static bool bounded(const ggml_tensor* tensor) {
+            return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024;
+        }
+        static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (disabled() || first + 1 >= ggml_graph_n_nodes(g))
+                return 0;
+            const auto* node = ggml_graph_node(g, first);
+            const auto* next = ggml_graph_node(g, first + 1);
+            if (node->op == GGML_OP_MUL) {
+                if (next->op != GGML_OP_ADD || (next->src[0] != node && next->src[1] != node))
+                    return 0;
+                const auto* x        = node->src[0];
+                const auto* gate     = node->src[1];
+                const auto* residual = next->src[0] == node ? next->src[1] : next->src[0];
+                return residual != node && contiguous_f32(x) && contiguous_f32(gate) && contiguous_f32(residual) && contiguous_f32(next) &&
+                               node->type == GGML_TYPE_F32 && ggml_are_same_shape(node, next) && ggml_are_same_shape(x, next) &&
+                               ggml_are_same_shape(residual, next) && gate->ne[0] == x->ne[0] && ggml_nelements(gate) == x->ne[0] &&
+                               x->ne[0] % 4 == 0 && bounded(next)
+                           ? 2
+                           : 0;
+            }
+            if (node->op == GGML_OP_CONCAT) {
+                const bool permuted = next->op == GGML_OP_PERMUTE;
+                const int count     = permuted ? 3 : 2;
+                if (first + count > ggml_graph_n_nodes(g) || node->type != GGML_TYPE_F32 || node->ne[0] % 4 ||
+                    node->op_params[0] != (permuted ? 2 : 1) || !ggml_is_contiguous(node) ||
+                    !contiguous_f32(node->src[0]) || !contiguous_f32(node->src[1]))
+                    return 0;
+                if (permuted && (next->src[0] != node || next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] ||
+                                 next->ne[2] != node->ne[1] || next->ne[3] != node->ne[3] || next->nb[0] != node->nb[0] ||
+                                 next->nb[1] != node->nb[2] || next->nb[2] != node->nb[1] || next->nb[3] != node->nb[3]))
+                    return 0;
+                const auto* dst = ggml_graph_node(g, first + count - 1);
+                return dst->op == GGML_OP_CPY && dst->src[0] == (permuted ? next : node) && dst->type == GGML_TYPE_F16 &&
+                               ggml_is_contiguous(dst) && ggml_are_same_shape(dst, dst->src[0]) && aligned(dst, 4) && bounded(dst)
+                           ? count
+                           : 0;
+            }
+            if (node->op == GGML_OP_CONT && first + 3 < ggml_graph_n_nodes(g)) {
+                const auto* src     = node->src[0];
+                const auto* permute = ggml_graph_node(g, first + 2);
+                const auto* dst     = ggml_graph_node(g, first + 3);
+                return ggml_are_same_shape(src, node) && src->type == GGML_TYPE_F32 && src->nb[0] == 4 && src->ne[0] % 4 == 0 &&
+                               aligned(src, 16) && (src->nb[1] | src->nb[2] | src->nb[3]) % 16 == 0 &&
+                               next->op == GGML_OP_SCALE && next->src[0] == node && next->type == GGML_TYPE_F32 && param(next, 1) == 0.f &&
+                               permute->op == GGML_OP_PERMUTE && permute->src[0] == next && permute->op_params[0] == 0 &&
+                               permute->op_params[1] == 2 && permute->op_params[2] == 1 && permute->op_params[3] == 3 &&
+                               dst->op == GGML_OP_CPY && dst->src[0] == permute && dst->type == GGML_TYPE_F16 &&
+                               ggml_are_same_shape(dst, permute) && ggml_is_contiguous(dst) && aligned(dst, 4) && bounded(dst)
+                           ? 4
+                           : 0;
+            }
+            return 0;
+        }
+        static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) {
+            const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            if (i >= count4)
+                return;
+            const float4 a = x[i], g = gate[i % n4], r = residual[i];
+            // The graph performs two FP32 operations; contracting them into FMA changes rounding.
+            dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y),
+                                 __fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w));
+        }
+        template <bool Permuted>
+        static __global__ void concat_cast(const float4* prefix, const float4* current, half2* dst, int64_t columns4, int64_t prefix_tokens, int64_t tokens, int64_t heads, int64_t count4) {
+            const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            if (i >= count4)
+                return;
+            const int64_t col = i % columns4, row = i / columns4;
+            const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
+            const bool from_prefix   = token < prefix_tokens;
+            const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens;
+            const int64_t src_token  = from_prefix ? token : token - prefix_tokens;
+            const int64_t src_row    = Permuted ? (batch * src_tokens + src_token) * heads + head
+                                                : (batch * heads + head) * src_tokens + src_token;
+            const float4 value       = (from_prefix ? prefix : current)[src_row * columns4 + col];
+            dst[2 * i]               = __floats2half2_rn(value.x, value.y);
+            dst[2 * i + 1]           = __floats2half2_rn(value.z, value.w);
+        }
+        static __global__ void scale_permute(const float* src, half2* dst, int64_t dim, int64_t tokens, int64_t heads, int64_t count, int64_t nb1, int64_t nb2, int64_t nb3, float scale) {
+            const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4;
+            if (index >= count)
+                return;
+            const int64_t col = index % dim, row = index / dim;
+            const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
+            const float4 value = *reinterpret_cast<const float4*>(src + col + head * nb1 + token * nb2 + batch * nb3);
+            dst[index / 2]     = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale));
+            dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale));
+        }
+        static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+            const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+            return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+        }
+        static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+            *workspace = ggml_nbytes(nodes[count - 1]);
+            return GGML_STATUS_SUCCESS;
+        }
+        static void destroy(void*, void*) {}
+        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context))
+                return GGML_STATUS_FAILED;
+            auto stream           = static_cast<CUDAStream>(context.stream);
+            const auto* node      = nodes[0];
+            const auto* a         = node->src[0];
+            const auto* b         = node->op == GGML_OP_CONT ? nullptr : node->src[1];
+            auto* dst             = nodes[count - 1];
+            const auto* residual  = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr;
+            const bool gate       = residual != nullptr;
+            const bool alias      = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
+                                    (residual && overlaps(dst, residual) && dst->data != residual->data);
+            void* output          = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+            const int64_t count4  = ggml_nelements(dst) / 4;
+            const unsigned blocks = unsigned((count4 + 255) / 256);
+            if (gate) {
+                gated_residual<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
+                                                           static_cast<const float4*>(residual->data), static_cast<float4*>(output), a->ne[0] / 4, count4);
+            } else if (node->op == GGML_OP_CONCAT) {
+                const bool permuted = count == 3;
+                auto kernel         = permuted ? concat_cast<true> : concat_cast<false>;
+                kernel<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
+                                                   static_cast<half2*>(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1],
+                                                   node->ne[permuted ? 1 : 2], count4);
+            } else {
+                scale_permute<<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
+                                                          a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, param(nodes[1], 0));
+            }
+            auto error = cuda_last_error();
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
+    }
+
+    void register_affine(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.priority             = 200;
+        desc.match                = sd_affine_cuda::match;
+        desc.prepare              = sd_affine_cuda::prepare;
+        desc.destroy              = sd_affine_cuda::destroy;
+        desc.execute              = sd_affine_cuda::execute;
+        desc.name                 = "sd.gated_residual." SD_CUDA_NAME ".v1";
+        desc.op                   = GGML_OP_MUL;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name = "sd.concat_cast." SD_CUDA_NAME ".v1";
+        desc.op   = GGML_OP_CONCAT;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name = "sd.cont_scale_permute_cast." SD_CUDA_NAME ".v1";
+        desc.op   = GGML_OP_CONT;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
+}

src/core/cuda/bf16.cu

@@ -0,0 +1,277 @@
+#include "bf16.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hipblas/hipblas.h>
+#elif defined(SD_CUDA_VENDOR_MUSA)
+#include <mublas.h>
+#else
+#include <cublas_v2.h>
+#endif
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include "rms_affine.h"
+#include "runtime.h"
+#ifdef SD_CUDA_BF16_CUTLASS
+#include "bf16_gemm.h"
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_bf16 {
+        static float scale_value(const ggml_tensor* node) {
+            float value;
+            std::memcpy(&value, node->op_params, sizeof(value));
+            return value;
+        }
+
+        static bool matrix(const ggml_tensor* node, ggml_type type) {
+            return node && node->type == type && ggml_is_contiguous(node) &&
+                   node->ne[2] == 1 && node->ne[3] == 1 && node->ne[0] <= INT_MAX && node->ne[1] <= INT_MAX;
+        }
+
+        static bool projection(const ggml_tensor* mat, const ggml_tensor* input) {
+            return mat->op == GGML_OP_MUL_MAT && mat->src[1] == input && mat->op_params[0] == GGML_PREC_DEFAULT &&
+                   matrix(input, GGML_TYPE_F32) && input->ne[1] >= 32 && ggml_nelements(input) <= INT_MAX &&
+                   matrix(mat, GGML_TYPE_F32) && matrix(mat->src[0], GGML_TYPE_BF16) && !mat->view_src;
+        }
+
+        static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+            // An explicit cuBLAS compute override must retain the native conversion path.
+            static const bool disabled = [] {
+                const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
+                return (value && std::atoi(value)) || std::getenv("GGML_CUDA_CUBLAS_COMPUTE_TYPE");
+            }();
+            if (disabled)
+                return 0;
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (first + 1 >= ggml_graph_n_nodes(g))
+                return 0;
+            const auto* glu = ggml_graph_node(g, first);
+            int offset      = 0;
+            if (glu->op == GGML_OP_RMS_NORM) {
+                if (!rms_affine_match(graph, first) || first + 4 >= ggml_graph_n_nodes(g))
+                    return 0;
+                const auto* input = ggml_graph_node(g, first + 3);
+                const auto* gate  = ggml_graph_node(g, first + 4);
+                if (!projection(gate, input))
+                    return 0;
+                if (first + 6 >= ggml_graph_n_nodes(g))
+                    return 5;
+                const auto* up = ggml_graph_node(g, first + 5);
+                glu            = ggml_graph_node(g, first + 6);
+                if (!projection(up, input) || glu->op != GGML_OP_GLU || glu->src[0] != gate || glu->src[1] != up)
+                    return 5;
+                offset = 6;
+            }
+            if (first + offset + 1 >= ggml_graph_n_nodes(g))
+                return 0;
+            if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU ||
+                !matrix(glu, GGML_TYPE_F32) || !matrix(glu->src[0], GGML_TYPE_F32) ||
+                !matrix(glu->src[1], GGML_TYPE_F32) || !ggml_are_same_shape(glu, glu->src[0]) ||
+                !ggml_are_same_shape(glu, glu->src[1]) || glu->ne[1] < 32 || ggml_nelements(glu) > INT_MAX)
+                return 0;
+            int count         = 2;
+            const auto* input = glu;
+            const auto* mat   = ggml_graph_node(g, first + offset + 1);
+            if (mat->op == GGML_OP_SCALE) {
+                float bias;
+                std::memcpy(&bias, mat->op_params + 1, sizeof(bias));
+                if (mat->src[0] != glu || bias != 0.f || first + offset + 2 >= ggml_graph_n_nodes(g))
+                    return 0;
+                input = mat;
+                mat   = ggml_graph_node(g, first + offset + 2);
+                count = 3;
+            }
+            if (!projection(mat, input))
+                return 0;
+            return offset + count;
+        }
+
+        struct Workspace {
+            size_t norm = 0, gate = 0, up = 0, packed = 0, output = 0, size = 0;
+            Workspace(ggml_tensor* const* nodes, int count) {
+                if (nodes[0]->op == GGML_OP_RMS_NORM) {
+                    norm = reserve(ggml_nelements(nodes[3]) * sizeof(CUDABFloat16));
+                    if (count > 5) {
+                        gate   = reserve(ggml_nbytes(nodes[4]));
+                        up     = reserve(ggml_nbytes(nodes[5]));
+                        packed = reserve(ggml_nelements(nodes[6]) * sizeof(CUDABFloat16));
+                    }
+                } else {
+                    packed = reserve(ggml_nelements(nodes[0]) * sizeof(CUDABFloat16));
+                }
+                output = reserve(ggml_nbytes(nodes[count - 1]));
+            }
+            size_t reserve(size_t bytes) {
+                const size_t offset = size;
+                size += GGML_PAD(bytes, 256);
+                return offset;
+            }
+        };
+
+        static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+            *workspace = Workspace(nodes, count).size;
+            return GGML_STATUS_SUCCESS;
+        }
+
+        static void destroy(void*, void*) {}
+
+        static __device__ __forceinline__ float swiglu(float x, float up, float scale) {
+            return __fmul_rn(__fmul_rn(x / (1.f + expf(-x)), up), scale);
+        }
+
+        static __global__ void pack_swiglu(const float* gate, const float* up, CUDABFloat16* dst, int64_t size, float scale) {
+            const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            if (i < size) {
+                dst[i] = __float2bfloat16(swiglu(gate[i], up[i], scale));
+            }
+        }
+
+        static __global__ void pack_swiglu4(const float4* gate, const float4* up, CUDABFloat162* dst, int64_t size, float scale) {
+            const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            if (i < size) {
+                const float4 x = gate[i], y = up[i];
+                dst[2 * i]     = __float22bfloat162_rn(make_float2(swiglu(x.x, y.x, scale), swiglu(x.y, y.y, scale)));
+                dst[2 * i + 1] = __float22bfloat162_rn(make_float2(swiglu(x.z, y.z, scale), swiglu(x.w, y.w, scale)));
+            }
+        }
+
+        static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+            const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+            return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+        }
+
+        static ggml_status gemm(void* handle, const ggml_tensor* w, const void* input, int rows, void* output, CUDAStream stream, bool tiled) {
+#ifdef SD_CUDA_BF16_CUTLASS
+            if (tiled) {
+                const auto result = sd_cuda_bf16_gemm(w->data, input, static_cast<float*>(output), rows, w->ne[1], w->ne[0], stream);
+                if (result == SDBF16GemmResult::Success)
+                    return GGML_STATUS_SUCCESS;
+                if (result == SDBF16GemmResult::Failed)
+                    return GGML_STATUS_FAILED;
+            }
+#else
+            (void)tiled;
+#endif
+            const float alpha = 1.f, beta = 0.f;
+#if defined(SD_CUDA_VENDOR_HIP)
+#if HIP_VERSION >= 60500000 || defined(HIPBLAS_V2)
+            constexpr auto compute_type = HIPBLAS_COMPUTE_32F;
+#else
+            constexpr auto compute_type = HIPBLAS_R_32F;
+#endif
+#if HIP_VERSION >= 70000000 || defined(HIPBLAS_V2)
+            constexpr auto bf16_type = HIP_R_16BF, f32_type = HIP_R_32F;
+#else
+            constexpr auto bf16_type = HIPBLAS_R_16B, f32_type = HIPBLAS_R_32F;
+#endif
+            return hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N,
+                                 int(w->ne[1]), rows, int(w->ne[0]), &alpha,
+                                 w->data, bf16_type, int(w->ne[0]), input, bf16_type, int(w->ne[0]),
+                                 &beta, output, f32_type, int(w->ne[1]), compute_type, HIPBLAS_GEMM_DEFAULT) == HIPBLAS_STATUS_SUCCESS
+                       ? GGML_STATUS_SUCCESS
+                       : GGML_STATUS_FAILED;
+#elif defined(SD_CUDA_VENDOR_MUSA)
+            return mublasGemmEx(static_cast<mublasHandle_t>(handle), MUBLAS_OP_T, MUBLAS_OP_N,
+                                int(w->ne[1]), rows, int(w->ne[0]), &alpha,
+                                w->data, MUSA_R_16BF, int(w->ne[0]), input, MUSA_R_16BF, int(w->ne[0]),
+                                &beta, output, MUSA_R_32F, int(w->ne[1]), MUSA_R_32F, MUBLAS_GEMM_DEFAULT) == MUBLAS_STATUS_SUCCESS
+                       ? GGML_STATUS_SUCCESS
+                       : GGML_STATUS_FAILED;
+#else
+            return cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N, int(w->ne[1]), rows, int(w->ne[0]),
+                                &alpha, w->data, CUDA_R_16BF, int(w->ne[0]), input, CUDA_R_16BF, int(w->ne[0]),
+                                &beta, output, CUDA_R_32F, int(w->ne[1]), CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP) == CUBLAS_STATUS_SUCCESS
+                       ? GGML_STATUS_SUCCESS
+                       : GGML_STATUS_FAILED;
+#endif
+        }
+
+        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void* user_data) {
+            const bool tiled = *static_cast<const bool*>(user_data);
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context))
+                return GGML_STATUS_FAILED;
+            const auto reg    = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+            const auto query  = reinterpret_cast<ggml_ext_gpu_blas_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_blas_v1"));
+            const auto handle = query ? query(execution->backend) : nullptr;
+            if (!handle)
+                return GGML_STATUS_FAILED;
+            const auto stream = static_cast<CUDAStream>(context.stream);
+            const bool norm   = nodes[0]->op == GGML_OP_RMS_NORM;
+            const int offset  = norm ? 6 : 0;
+            const auto* glu   = norm && count == 5 ? nullptr : nodes[offset];
+            auto* dst         = nodes[count - 1];
+            const auto* w     = dst->src[0];
+            auto* base        = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+            const Workspace work(nodes, count);
+            auto* packed     = reinterpret_cast<CUDABFloat16*>(base + work.packed);
+            const bool alias = overlaps(dst, w);
+            void* output     = alias ? base + work.output : dst->data;
+            if (norm) {
+                if (rms_affine_pack(stream, nodes, base + work.norm) != GGML_STATUS_SUCCESS)
+                    return GGML_STATUS_FAILED;
+                if (gemm(handle, nodes[4]->src[0], base + work.norm, int(nodes[3]->ne[1]), count == 5 ? output : base + work.gate, stream, tiled) != GGML_STATUS_SUCCESS)
+                    return GGML_STATUS_FAILED;
+                if (count == 5)
+                    return !alias || cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+                if (gemm(handle, nodes[5]->src[0], base + work.norm, int(nodes[3]->ne[1]), base + work.up, stream, tiled) != GGML_STATUS_SUCCESS)
+                    return GGML_STATUS_FAILED;
+            }
+            const auto size   = ggml_nelements(glu);
+            const void* gate  = norm ? base + work.gate : glu->src[0]->data;
+            const void* up    = norm ? base + work.up : glu->src[1]->data;
+            const float scale = count - offset == 3 ? scale_value(nodes[offset + 1]) : 1.f;
+            if (size % 4 == 0 && (reinterpret_cast<uintptr_t>(gate) | reinterpret_cast<uintptr_t>(up)) % 16 == 0) {
+                pack_swiglu4<<<unsigned((size / 4 + 255) / 256), 256, 0, stream>>>(static_cast<const float4*>(gate),
+                                                                                   static_cast<const float4*>(up),
+                                                                                   reinterpret_cast<CUDABFloat162*>(packed), size / 4, scale);
+            } else {
+                pack_swiglu<<<unsigned((size + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(gate),
+                                                                              static_cast<const float*>(up), packed, size, scale);
+            }
+            if (cuda_last_error() != cuda_success)
+                return GGML_STATUS_FAILED;
+            const auto status = gemm(handle, w, packed, int(glu->ne[1]), output, stream, tiled);
+            if (status != GGML_STATUS_SUCCESS)
+                return GGML_STATUS_FAILED;
+            return !alias || cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
+    }
+
+    void register_bf16(ggml_backend_dev_t device) {
+        const auto info = cuda_device_info(device);
+        if (info.device < 0 || !ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1"))
+            return;
+        bool tiled = false;
+#if defined(SD_CUDA_VENDOR_HIP)
+        hipDeviceProp_t props;
+        if (hipGetDeviceProperties(&props, info.device) != hipSuccess || std::strncmp(props.gcnArchName, "gfx", 3))
+            return;
+        const auto arch = std::strtoul(props.gcnArchName + 3, nullptr, 16);
+        // Match the native backend's BF16 operand policy; older GPUs use FP32 GEMM.
+        if (arch < 0x1100 && !(arch >= 0x908 && arch < 0x1010 && arch != 0x909 && arch != 0x90c))
+            return;
+#elif !defined(SD_CUDA_VENDOR_MUSA)
+        cudaDeviceProp props;
+        if (cudaGetDeviceProperties(&props, info.device) != cuda_success || props.major < 8)
+            return;
+        tiled = props.major >= 8;
+#endif
+        static bool tiled_gemm[]  = {false, true};
+        ggml_ext_kernel_desc desc = {};
+        desc.user_data            = &tiled_gemm[tiled];
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.swiglu_bf16_matmul." SD_CUDA_NAME ".v1";
+        desc.op                   = GGML_OP_GLU;
+        desc.priority             = 200;
+        desc.match                = sd_bf16::match;
+        desc.prepare              = sd_bf16::prepare;
+        desc.destroy              = sd_bf16::destroy;
+        desc.execute              = sd_bf16::execute;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name     = "sd.rms_affine_bf16_matmul." SD_CUDA_NAME ".v1";
+        desc.op       = GGML_OP_RMS_NORM;
+        desc.priority = 400;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+}

src/core/cuda/bf16.h

@@ -0,0 +1,26 @@
+#ifndef __SD_CORE_CUDA_BF16_H__
+#define __SD_CORE_CUDA_BF16_H__
+
+#include "runtime.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hip/hip_bf16.h>
+#elif defined(SD_CUDA_VENDOR_MUSA)
+#include <musa_bf16.h>
+#else
+#include <cuda_bf16.h>
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+#if defined(SD_CUDA_VENDOR_HIP)
+    using CUDABFloat16  = __hip_bfloat16;
+    using CUDABFloat162 = __hip_bfloat162;
+#elif defined(SD_CUDA_VENDOR_MUSA)
+    using CUDABFloat16  = __mt_bfloat16;
+    using CUDABFloat162 = __mt_bfloat162;
+#else
+    using CUDABFloat16  = nv_bfloat16;
+    using CUDABFloat162 = nv_bfloat162;
+#endif
+}
+
+#endif

src/core/cuda/bf16_gemm.cu

@@ -0,0 +1,25 @@
+#include <cutlass/bfloat16.h>
+#include <cutlass/gemm/device/gemm.h>
+#include "bf16_gemm.h"
+
+SDBF16GemmResult sd_cuda_bf16_gemm(const void* weight, const void* input, float* output, int64_t rows, int64_t n, int64_t k, cudaStream_t stream) {
+#ifdef SD_CUDA_FORCE_CUBLAS
+    return SDBF16GemmResult::Unsupported;
+#else
+    if (rows < 512 || rows > 2048 || (n < 4096 && (n < 2048 || k < 2 * n)) || k % 8 || n % 4 ||
+        (reinterpret_cast<uintptr_t>(weight) | reinterpret_cast<uintptr_t>(input) | reinterpret_cast<uintptr_t>(output)) % 16)
+        return SDBF16GemmResult::Unsupported;
+    using Gemm = cutlass::gemm::device::Gemm<
+        cutlass::bfloat16_t, cutlass::layout::RowMajor, cutlass::bfloat16_t, cutlass::layout::ColumnMajor,
+        float, cutlass::layout::RowMajor, float, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80,
+        cutlass::gemm::GemmShape<64, 128, 32>, cutlass::gemm::GemmShape<32, 64, 32>, cutlass::gemm::GemmShape<16, 8, 16>,
+        cutlass::epilogue::thread::LinearCombination<float, 4, float, float>, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 3>;
+    Gemm gemm;
+    const Gemm::Arguments args({int(rows), int(n), int(k)},
+                               {static_cast<const cutlass::bfloat16_t*>(input), k},
+                               {static_cast<const cutlass::bfloat16_t*>(weight), k}, {output, n}, {output, n}, {1.f, 0.f});
+    if (gemm.can_implement(args) != cutlass::Status::kSuccess)
+        return SDBF16GemmResult::Unsupported;
+    return gemm(args, nullptr, stream) == cutlass::Status::kSuccess ? SDBF16GemmResult::Success : SDBF16GemmResult::Failed;
+#endif
+}

src/core/cuda/bf16_gemm.h

@@ -0,0 +1,13 @@
+#ifndef __SD_CORE_CUDA_BF16_GEMM_H__
+#define __SD_CORE_CUDA_BF16_GEMM_H__
+
+#include <cuda_runtime.h>
+#include <cstdint>
+
+enum class SDBF16GemmResult { Unsupported,
+                              Success,
+                              Failed };
+
+SDBF16GemmResult sd_cuda_bf16_gemm(const void* weight, const void* input, float* output, int64_t rows, int64_t n, int64_t k, cudaStream_t stream);
+
+#endif

src/core/cuda/convrot.cu

@@ -0,0 +1,307 @@
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include "convrot_kernel.h"
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_convrot {
+
+        static bool fusion_disabled() {
+            static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+            return disabled;
+        }
+        static float param_f32(const ggml_tensor* tensor, int index) {
+            float value;
+            std::memcpy(&value, tensor->op_params + index, sizeof(value));
+            return value;
+        }
+        static bool f32(const ggml_tensor* tensor) {
+            return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor);
+        }
+        static bool vector(const ggml_tensor* tensor, int64_t k) {
+            return f32(tensor) && tensor->ne[0] == k && ggml_nelements(tensor) == k;
+        }
+        enum Kind { SCALE,
+                    NORM,
+                    MOD,
+                    RMS_MOD,
+                    SWIGLU,
+                    SILU,
+                    SIGMOID };
+        struct Operation {
+            Kind kind;
+            const ggml_tensor* input      = nullptr;
+            const ggml_tensor* up         = nullptr;
+            const ggml_tensor* modulation = nullptr;
+            const ggml_tensor* weight     = nullptr;
+            const ggml_tensor* shift      = nullptr;
+            int64_t k, rows, stride, up_stride, gate_offset = 0, up_offset = 0;
+            float eps = 0.f, scale = 1.f, bias = 0.f;
+        };
+        static bool describe(ggml_tensor* const* nodes, int count, Operation& op) {
+            int group;
+            const auto* first = nodes[0];
+            const auto* quant = nodes[count - 1];
+            if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
+                return false;
+            const auto* value = quant->src[0];
+            op.k              = value->ne[0];
+            op.rows           = ggml_nrows(value);
+            op.input          = first->src[0];
+            if (!f32(op.input) || !f32(first) || op.k > INT_MAX || op.rows > INT_MAX)
+                return false;
+            op.stride    = op.input->ne[0];
+            op.up_stride = op.k;
+            if (count == 2 && first->op == GGML_OP_SCALE && value == first) {
+                op.kind  = SCALE;
+                op.scale = param_f32(first, 0);
+                op.bias  = param_f32(first, 1);
+            } else if (count == 3 && first->op == GGML_OP_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && value == nodes[1]) {
+                op.kind       = NORM;
+                op.modulation = nodes[1]->src[1];
+                op.eps        = param_f32(first, 0);
+                if (op.k < 1024 || !vector(op.modulation, op.k))
+                    return false;
+            } else if (count == 2 && sd_ggml_modulate_describe(first) && value == first) {
+                op.kind       = MOD;
+                op.modulation = first->src[1];
+                op.shift      = first->src[2];
+                if (!vector(op.modulation, op.k) || (op.shift && !vector(op.shift, op.k)))
+                    return false;
+            } else if (count == 4 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first &&
+                       sd_ggml_modulate_describe(nodes[2]) && nodes[2]->src[0] == nodes[1] && value == nodes[2]) {
+                op.kind       = RMS_MOD;
+                op.weight     = nodes[1]->src[1];
+                op.modulation = nodes[2]->src[1];
+                op.shift      = nodes[2]->src[2];
+                op.eps        = param_f32(first, 0);
+                if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || (op.shift && !vector(op.shift, op.k)))
+                    return false;
+            } else if (count == 2 && first->op == GGML_OP_GLU && ggml_get_glu_op(first) == GGML_GLU_OP_SWIGLU && value == first) {
+                op.kind = SWIGLU;
+                op.up   = first->src[1] ? first->src[1] : op.input;
+                if (!f32(op.up))
+                    return false;
+                op.up_stride = op.up->ne[0];
+                if (!first->src[1]) {
+                    if (first->op_params[1])
+                        op.gate_offset = op.k;
+                    else
+                        op.up_offset = op.k;
+                }
+            } else if (count == 3 && first->op == GGML_OP_UNARY && nodes[1]->op == GGML_OP_MUL && value == nodes[1] &&
+                       (nodes[1]->src[0] == first || nodes[1]->src[1] == first)) {
+                auto unary = ggml_get_unary_op(first);
+                if (unary != GGML_UNARY_OP_SILU && unary != GGML_UNARY_OP_SIGMOID)
+                    return false;
+                op.kind = unary == GGML_UNARY_OP_SILU ? SILU : SIGMOID;
+                op.up   = nodes[1]->src[0] == first ? nodes[1]->src[1] : nodes[1]->src[0];
+                if (!f32(op.up) || !ggml_are_same_shape(op.input, op.up))
+                    return false;
+            } else
+                return false;
+            for (int i = 0; i < count - 1; ++i) {
+                if (!f32(nodes[i]) || !ggml_are_same_shape(nodes[i], value))
+                    return false;
+                for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
+                    if (input == nodes[i])
+                        return false;
+            }
+            return true;
+        }
+        static bool supported(const Operation& op, const ggml_ext_gpu_device_info& device) {
+            const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256
+                                                                            : 128;
+            return device.device >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared_memory;
+        }
+        static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+            if (fusion_disabled())
+                return 0;
+            auto* g             = const_cast<ggml_cgraph*>(graph);
+            const int available = ggml_graph_n_nodes(g) - first;
+            const auto* node    = ggml_graph_node(g, first);
+            const int count     = node->op == GGML_OP_RMS_NORM ? 4 : node->op == GGML_OP_NORM || node->op == GGML_OP_UNARY ? 3
+                                                                                                                           : 2;
+            if (available < count)
+                return 0;
+            ggml_tensor* nodes[6];
+            for (int j = 0; j < count; ++j)
+                nodes[j] = ggml_graph_node(g, first + j);
+            Operation op;
+            return describe(nodes, count, op) && supported(op, cuda_device_info(device)) ? count : 0;
+        }
+        static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
+            if (fusion_disabled())
+                return;
+            for (int i = 1; i < ggml_graph_n_nodes(graph); ++i) {
+                auto* quant = ggml_graph_node(graph, i);
+                int group;
+                if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
+                    continue;
+                ggml_tensor* chain[6] = {};
+                int count             = 0;
+                auto* value           = quant->src[0];
+                if (sd_ggml_modulate_describe(value)) {
+                    auto* weight = value->src[0];
+                    if (weight->op == GGML_OP_MUL && weight->src[0]->op == GGML_OP_RMS_NORM) {
+                        chain[0] = weight->src[0];
+                        chain[1] = weight;
+                        chain[2] = value;
+                        chain[3] = quant;
+                        count    = 4;
+                    } else {
+                        chain[0] = value;
+                        chain[1] = quant;
+                        count    = 2;
+                    }
+                } else if (value->op == GGML_OP_MUL) {
+                    for (int src = 0; src < 2; ++src)
+                        if (value->src[src]->op == GGML_OP_UNARY) {
+                            chain[0] = value->src[src];
+                            chain[1] = value;
+                            chain[2] = quant;
+                            count    = 3;
+                            break;
+                        }
+                }
+                Operation op;
+                if (!count || !describe(chain, count, op) || !supported(op, cuda_device_info(device)))
+                    continue;
+                int begin = -1;
+                for (int j = std::max(0, i - 64); j < i; ++j)
+                    if (ggml_graph_node(graph, j) == chain[0]) {
+                        begin = j;
+                        break;
+                    }
+                if (begin < 0)
+                    continue;
+                // Move the owning producer with its in-place consumers so alias checks retain their storage relationship.
+                for (int part = 1; part < count; ++part) {
+                    int next = -1;
+                    for (int j = begin + part; j <= i; ++j)
+                        if (ggml_graph_node(graph, j) == chain[part]) {
+                            next = j;
+                            break;
+                        }
+                    if (next < 0 || (next > begin + part && !ggml_ext_graph_delay(graph, begin, part, next)))
+                        break;
+                    begin = next - part;
+                }
+            }
+        }
+        struct Plan {
+            ggml_ext_gpu_device_info device;
+            Operation op;
+            bool configured = false;
+        };
+        static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+            auto* plan   = new Plan{};
+            *state       = plan;
+            plan->device = cuda_device_info(ggml_backend_get_device(backend));
+            if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device))
+                return GGML_STATUS_FAILED;
+            const auto* dst = nodes[count - 1];
+            *workspace      = !dst->data && !dst->buffer && !dst->view_src ? 0 : ggml_nbytes(dst);
+            return GGML_STATUS_SUCCESS;
+        }
+        static void destroy(void* state, void*) {
+            delete static_cast<Plan*>(state);
+        }
+        static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+            if (!b)
+                return false;
+            const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+            return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+        }
+        template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false, bool Modulate = false>
+        static CUDAError launch(Plan& p, const Operation& op, int8_t* output, CUDAStream stream) {
+            constexpr bool warp = Scaled || RMS || Modulate;
+            constexpr int extra = Normalize ? 384 : RMS ? 256
+                                                        : 128;
+            if (!p.configured) {
+                auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS, Modulate>,
+                                                     int(p.device.shared_memory - extra));
+                if (status != cuda_success)
+                    return status;
+                p.configured = true;
+            }
+            const int threads    = Normalize ? 1024 : RMS                               ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256
+                                                                                                            : op.k <= 8192   ? 512
+                                                                                                                             : 1024)
+                                                  : Swiglu && !Sigmoid && op.k >= 16384 ? 512
+                                                                                        : 256;
+            const int64_t padded = (op.rows + 3) / 4 * 4;
+            auto data            = [](const ggml_tensor* t) { return t ? static_cast<const float*>(t->data) : nullptr; };
+            quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS, Modulate><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
+                data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(output + op.k * padded), op.k, op.rows, op.k,
+                op.up ? data(op.up) + op.up_offset : nullptr, op.stride, op.up_stride, nullptr, nullptr, data(op.modulation), op.eps, op.scale, op.bias, data(op.weight), data(op.shift));
+            return cuda_last_error();
+        }
+        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+            auto& plan     = *static_cast<Plan*>(state);
+            const auto& op = plan.op;
+            auto* dst      = nodes[count - 1];
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context) || context.device != plan.device.device)
+                return GGML_STATUS_FAILED;
+            auto stream = static_cast<CUDAStream>(context.stream);
+            bool alias  = false;
+            for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
+                alias |= overlaps(dst, input);
+            auto* output         = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
+            const int64_t padded = (op.rows + 3) / 4 * 4;
+            if (padded > op.rows && cuda_memset_async(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cuda_success)
+                return GGML_STATUS_FAILED;
+            CUDAError error;
+            switch (op.kind) {
+                case SCALE:
+                    error = launch<false, false, true>(plan, op, output, stream);
+                    break;
+                case NORM:
+                    error = launch<false, true>(plan, op, output, stream);
+                    break;
+                case RMS_MOD:
+                    error = launch<false, false, false, false, true>(plan, op, output, stream);
+                    break;
+                case MOD:
+                    error = launch<false, false, false, false, false, true>(plan, op, output, stream);
+                    break;
+                case SWIGLU:
+                case SILU:
+                    error = launch<true>(plan, op, output, stream);
+                    break;
+                case SIGMOID:
+                    error = launch<true, false, false, true>(plan, op, output, stream);
+                    break;
+                default:
+                    return GGML_STATUS_FAILED;
+            }
+            if (error != cuda_success || (alias && cuda_copy_async(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), stream) != cuda_success))
+                return GGML_STATUS_FAILED;
+            return GGML_STATUS_SUCCESS;
+        }
+    }
+
+    void register_convrot(ggml_backend_dev_t device) {
+        for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY, GGML_OP_CUSTOM}) {
+            ggml_ext_kernel_desc desc = {};
+            desc.abi_version          = GGML_EXT_ABI_VERSION;
+            const std::string name    = std::string("sd.convrot_fusion.") + ggml_op_name(op) + "." SD_CUDA_NAME ".v1";
+            desc.name                 = name.c_str();
+            desc.op                   = op;
+            desc.custom_op            = op == GGML_OP_CUSTOM ? sd_ggml_modulate_op() : nullptr;
+            desc.priority             = 225;
+            desc.optimize             = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr;
+            desc.match                = sd_convrot::match;
+            desc.prepare              = sd_convrot::prepare;
+            desc.execute              = sd_convrot::execute;
+            desc.destroy              = sd_convrot::destroy;
+            GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        }
+    }
+
+}

src/core/cuda/convrot_kernel.h

@@ -0,0 +1,256 @@
+#ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__
+#define __SD_CORE_CUDA_CONVROT_KERNEL_H__
+#include <cstdint>
+#include "runtime.h"
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_convrot {
+        static __device__ __forceinline__ float warp_reduce_sum(float x) {
+#pragma unroll
+            for (int offset = 16; offset > 0; offset >>= 1)
+                x += cuda_shuffle_xor(x, offset, 32);
+            return x;
+        }
+        static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) {
+#pragma unroll
+            for (int offset = 16; offset > 0; offset >>= 1) {
+                x.x += cuda_shuffle_xor(x.x, offset, 32);
+                x.y += cuda_shuffle_xor(x.y, offset, 32);
+            }
+            return x;
+        }
+        template <int Width>
+        static __device__ __forceinline__ float warp_reduce_max(float x) {
+#pragma unroll
+            for (int offset = Width / 2; offset > 0; offset >>= 1)
+                x = fmaxf(x, cuda_shuffle_xor(x, offset, Width));
+            return x;
+        }
+        static __device__ float2 block_sum(float2 x, float2* shared) {
+            x = warp_reduce_sum(x);
+            if (threadIdx.x % 32 == 0)
+                shared[threadIdx.x / 32] = x;
+            __syncthreads();
+            return warp_reduce_sum(shared[threadIdx.x % 32]);
+        }
+        static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+            a = index == 3 ? -a : a;
+            b = index == 2 ? -b : b;
+            c = index == 1 ? -c : c;
+            d = index == 0 ? -d : d;
+            return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+        }
+
+        template <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false, bool modulate = false>
+        __global__ void quantize(
+            const T* src,
+            int8_t* dst,
+            float* scales,
+            int64_t k,
+            int64_t rows,
+            int64_t dst_stride,
+            const float* up,
+            int64_t src_stride,
+            int64_t up_stride,
+            const float* activation_scales = nullptr,
+            const float* weight_scales     = nullptr,
+            const float* modulation        = nullptr,
+            float eps                      = 0.0f,
+            float input_scale              = 1.0f,
+            float input_bias               = 0.0f,
+            const float* rms_weight        = nullptr,
+            const float* rms_shift         = nullptr) {
+            constexpr int group_size = 256;
+
+            const int64_t row = blockIdx.x;
+            const int tid     = threadIdx.x;
+            if (row >= rows) {
+                return;
+            }
+
+            extern __shared__ float values[];
+            __shared__ float maxima[32];
+            const T* src_row = src + row * src_stride;
+            int8_t* dst_row  = dst + row * dst_stride;
+
+            [[maybe_unused]] float rms_scale = 1.0f;
+            if constexpr (rms_modulation) {
+                __shared__ float sums[32];
+                // Preserve the unfused RMSNorm's 1024-thread reduction order with fewer physical warps.
+                for (int virtual_tid = tid; virtual_tid < 1024; virtual_tid += blockDim.x) {
+                    float partial = 0.0f;
+                    for (int64_t i = virtual_tid; i < k; i += 1024) {
+                        const float value = src_row[i];
+                        partial += value * value;
+                    }
+                    partial = warp_reduce_sum(partial);
+                    if (tid % 32 == 0) {
+                        sums[virtual_tid / 32] = partial;
+                    }
+                }
+                __syncthreads();
+                const float sum = warp_reduce_sum(sums[tid % 32]);
+                rms_scale       = rsqrtf(sum / k + eps);
+            }
+
+            if constexpr (swiglu || warp_rotation) {
+                const int lane = tid % 32;
+                for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
+                    float v[8];
+#pragma unroll
+                    for (int j = 0; j < 8; ++j) {
+                        const int64_t i = group * group_size + j * 32 + lane;
+                        if constexpr (!swiglu) {
+                            float value = src_row[i];
+                            if constexpr (rms_modulation) {
+                                value = __fmul_rn(__fmul_rn(value, rms_scale), rms_weight[i]);
+                            }
+                            if constexpr (rms_modulation || modulate) {
+                                value = __fadd_rn(value, __fmul_rn(value, modulation[i]));
+                                if (rms_shift) {
+                                    value = __fadd_rn(value, rms_shift[i]);
+                                }
+                            }
+                            if constexpr (scaled) {
+                                value = __fmaf_rn(value, input_scale, input_bias);
+                            }
+                            v[j] = value * (1.0f / 16.0f);
+                        } else {
+                            float gate_value, up_value;
+                            if constexpr (dequantize) {
+                                gate_value = (float)((const int32_t*)src_row)[i] * activation_scales[row];
+                                up_value   = (float)((const int32_t*)up)[row * up_stride + i] * activation_scales[row];
+                                gate_value *= weight_scales[i];
+                                up_value *= weight_scales[k + i];
+                            } else {
+                                gate_value = src_row[i];
+                                up_value   = up[row * up_stride + i];
+                            }
+                            const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value)));
+                            v[j]                  = (activated * up_value) * (1.0f / 16.0f);
+                        }
+                    }
+                    // One logical warp owns H256, keeping every butterfly in registers.
+#pragma unroll
+                    for (int stride = 1; stride <= 4; stride *= 4) {
+                        const int base = lane & ~(3 * stride);
+#pragma unroll
+                        for (int j = 0; j < 8; ++j) {
+                            const float a = cuda_shuffle(v[j], base, 32);
+                            const float b = cuda_shuffle(v[j], base + stride, 32);
+                            const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
+                            const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
+                            v[j]          = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
+                        }
+                    }
+#pragma unroll
+                    for (int j = 0; j < 8; j += 2) {
+                        const float a = cuda_shuffle(v[j], lane % 16, 32);
+                        const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
+                        const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
+                        const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
+                        v[j]          = convrot_h4_cuda(a, b, c, d, lane / 16);
+                        v[j + 1]      = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
+                    }
+#pragma unroll
+                    for (int j = 0; j < 2; ++j) {
+                        const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
+#pragma unroll
+                        for (int t = 0; t < 4; ++t) {
+                            values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
+                        }
+                    }
+                }
+                __syncthreads();
+            } else {
+                constexpr int group_threads   = group_size / 4;
+                constexpr int groups_per_wave = 1024 / group_threads;
+                if constexpr (normalize) {
+                    __shared__ float2 sums[32];
+                    float2 mean_var = make_float2(0.0f, 0.0f);
+                    for (int64_t i = tid; i < k; i += 1024) {
+                        const float value = src_row[i];
+                        mean_var.x += value;
+                        mean_var.y += value * value;
+                    }
+                    mean_var             = block_sum(mean_var, sums);
+                    const float mean     = mean_var.x / k;
+                    const float variance = mean_var.y / k - mean * mean;
+                    const float inv_std  = rsqrtf(variance + eps);
+                    for (int64_t i = tid; i < k; i += 1024) {
+                        const float normalized = __fmul_rn(__fsub_rn(src_row[i], mean), inv_std);
+                        values[i]              = __fmul_rn(__fmul_rn(normalized, modulation[i]), 1.0f / 16.0f);
+                    }
+                } else {
+                    for (int64_t i = tid; i < k; i += blockDim.x) {
+                        const float value = scaled ? __fmaf_rn(src_row[i], input_scale, input_bias) : src_row[i];
+                        values[i]         = value * (1.0f / 16.0f);
+                    }
+                }
+                __syncthreads();
+
+                const int group_lane = tid % group_threads;
+                const int wave_group = tid / group_threads;
+                const int64_t groups = k / group_size;
+                for (int64_t group_base = 0; group_base < groups; group_base += groups_per_wave) {
+                    const int64_t group = group_base + wave_group;
+#pragma unroll
+                    for (int stride = 1; stride < group_size; stride *= 4) {
+                        if (group < groups) {
+                            float* group_values = values + group * group_size;
+                            const int base      = (group_lane / stride) * 4 * stride + group_lane % stride;
+                            const int i0        = base;
+                            const int i1        = i0 + stride;
+                            const int i2        = i1 + stride;
+                            const int i3        = i2 + stride;
+                            const float a       = group_values[i0];
+                            const float b       = group_values[i1];
+                            const float c       = group_values[i2];
+                            const float d       = group_values[i3];
+                            group_values[i0]    = a + b + c - d;
+                            group_values[i1]    = a + b - c + d;
+                            group_values[i2]    = a - b + c + d;
+                            group_values[i3]    = -a + b + c + d;
+                        }
+                        __syncthreads();
+                    }
+                }
+            }
+
+            float local_max = 0.0f;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                local_max = fmaxf(local_max, fabsf(values[i]));
+            }
+            local_max      = warp_reduce_max<32>(local_max);
+            const int lane = tid & 31;
+            const int warp = tid >> 5;
+            if (lane == 0) {
+                maxima[warp] = local_max;
+            }
+            __syncthreads();
+
+            if (warp == 0) {
+                local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
+                local_max = warp_reduce_max<32>(local_max);
+                if (lane == 0) {
+                    maxima[0] = local_max / 127.0f;
+                    if (scales != nullptr) {
+                        scales[row] = maxima[0];
+                    } else {
+                        *((float*)(dst_row + k)) = maxima[0];
+                    }
+                }
+            }
+            __syncthreads();
+
+            const float row_scale = maxima[0];
+            const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                int value  = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
+                value      = max(-127, min(127, value));
+                dst_row[i] = (int8_t)value;
+            }
+        }
+
+    }
+}
+#endif

src/core/cuda/cute_attention.cu

@@ -0,0 +1,841 @@
+/*
+ * Adapted from FlashAttention v2.8.3:
+ * https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src
+ * Sources: kernel_traits.h, utils.h, softmax.h, flash_fwd_kernel.h.
+ * Copyright (c) 2023-2024, Tri Dao.
+ *
+ * BSD 3-Clause License
+ *
+ * Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * * Redistributions of source code must retain the above copyright notice, this
+ *   list of conditions and the following disclaimer.
+ *
+ * * Redistributions in binary form must reproduce the above copyright notice,
+ *   this list of conditions and the following disclaimer in the documentation
+ *   and/or other materials provided with the distribution.
+ *
+ * * Neither the name of the copyright holder nor the names of its
+ *   contributors may be used to endorse or promote products derived from
+ *   this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cutlass/numeric_conversion.h>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <cute/tensor.hpp>
+#include <mutex>
+#include <type_traits>
+#include <unordered_map>
+#include "core/cuda/i8_projection.h"
+#include "core/cuda/rope.h"
+#include "core/ggml_extensions.h"
+#include "ggml-ext.h"
+
+static __device__ __forceinline__ float sd_warp_max(float v) {
+    for (int mask = 16; mask > 0; mask >>= 1)
+        v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, mask));
+    return v;
+}
+namespace sd_fattn_cute {
+    using namespace cute;
+    template <bool A_in_regs = false, bool B_in_regs = false, typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3, typename Tensor4, typename TiledMma, typename TiledCopyA, typename TiledCopyB, typename ThrCopyA, typename ThrCopyB>
+    __forceinline__ __device__ void gemm(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsA, Tensor4 const& tCsB, TiledMma tiled_mma, TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B, ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) {
+        CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+        CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+        CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
+        Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
+        CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
+        Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
+        CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
+        if (!A_in_regs) {
+            cute::copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
+        }
+        if (!B_in_regs) {
+            cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
+        }
+#pragma unroll
+        for (int i = 0; i < size<2>(tCrA); ++i) {
+            if (i < size<2>(tCrA) - 1) {
+                if (!A_in_regs) {
+                    cute::copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
+                }
+                if (!B_in_regs) {
+                    cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
+                }
+            }
+            cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
+        }
+    }
+
+    template <typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3, typename TiledMma, typename TiledCopy, typename ThrCopy>
+    __forceinline__ __device__ void gemm_rs(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsB, TiledMma tiled_mma, TiledCopy smem_tiled_copy_B, ThrCopy smem_thr_copy_B) {
+        CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+        CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+        CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
+        Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
+        CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
+        cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
+#pragma unroll
+        for (int i = 0; i < size<2>(tCrA); ++i) {
+            if (i < size<2>(tCrA) - 1) {
+                cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
+            }
+            cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
+        }
+    }
+
+    template <typename Layout>
+    __forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
+        static_assert(decltype(size<0>(acc_layout))::value == 4);
+        static_assert(decltype(rank(acc_layout))::value == 3);
+        auto l = logical_divide(acc_layout, Shape<_2>{});
+        return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));
+    };
+
+    template <typename MMA_traits, typename Layout>
+    __forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) {
+        using X = Underscore;
+        static_assert(decltype(size<0>(acc_layout))::value == 4);
+        static_assert(decltype(rank(acc_layout))::value == 3);
+        constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{});
+        static_assert(mma_shape_K == 8 || mma_shape_K == 16);
+        if constexpr (mma_shape_K == 8) {
+            return acc_layout;
+        } else {
+            auto l = logical_divide(acc_layout, Shape<X, X, _2>{});
+            return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
+        }
+    };
+
+    template <typename TensorT>
+    __device__ __forceinline__ auto to_half(const TensorT& src) {
+        constexpr int count = decltype(size(src))::value;
+        auto dst            = make_tensor<cutlass::half_t>(src.layout());
+        cutlass::NumericArrayConverter<cutlass::half_t, float, count> convert;
+        *reinterpret_cast<cutlass::Array<cutlass::half_t, count>*>(dst.data()) =
+            convert(*reinterpret_cast<const cutlass::Array<float, count>*>(src.data()));
+        return dst;
+    }
+
+    template <bool Full, bool Clear, bool FullCols = true, typename Copy, typename Src, typename Dst, typename Id>
+    __device__ __forceinline__ void copy_rows(Copy copier, const Src& src, Dst& dst, const Id& id, int rows, int cols = 0) {
+#pragma unroll
+        for (int m = 0; m < size<1>(src); ++m) {
+#pragma unroll
+            for (int k = 0; k < size<2>(src); ++k) {
+                if ((Full || get<0>(id(0, m, k)) < rows) && (FullCols || get<1>(id(0, m, k)) < cols)) {
+                    cute::copy(copier, src(_, m, k), dst(_, m, k));
+                } else if constexpr (Clear) {
+                    cute::clear(dst(_, m, k));
+                }
+            }
+        }
+    }
+
+    template <int Rows, bool Masked = false>
+    struct Softmax {
+        decltype(make_tensor<float>(Shape<Int<Rows>>{})) maximum, sum;
+        template <bool First, typename Scores, typename Output>
+        __device__ __forceinline__ void update(Scores& acc_s, Output& acc_o) {
+            auto scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
+            auto output = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+#pragma unroll
+            for (int m = 0; m < Rows; ++m) {
+                const float old_max = First ? -INFINITY : maximum(m);
+                float next_max      = First ? scores(m, 0) : fmaxf(old_max, scores(m, 0));
+#pragma unroll
+                for (int n = 1; n < size<1>(scores); ++n)
+                    next_max = fmaxf(next_max, scores(m, n));
+                next_max      = fmaxf(next_max, __shfl_xor_sync(0xffffffff, next_max, 2));
+                next_max      = fmaxf(next_max, __shfl_xor_sync(0xffffffff, next_max, 1));
+                maximum(m)    = next_max;
+                float row_sum = 0.f;
+                if constexpr (!First) {
+                    row_sum = sum(m);
+                    if (old_max != next_max) {
+                        const float factor = exp2f((old_max - next_max) * 1.4426950408889634f);
+                        row_sum *= factor;
+#pragma unroll
+                        for (int n = 0; n < size<1>(output); ++n)
+                            output(m, n) *= factor;
+                    }
+                }
+#pragma unroll
+                for (int n = 0; n < size<1>(scores); ++n) {
+                    const float safe_max = Masked && next_max == -INFINITY ? 0.f : next_max;
+                    const float value    = exp2f(scores(m, n) * 1.4426950408889634f - safe_max * 1.4426950408889634f);
+                    scores(m, n)         = value;
+                    row_sum += value;
+                }
+                // Defer the four-lane sum reduction until the final output normalization.
+                sum(m) = row_sum;
+            }
+        }
+        template <typename Output>
+        __device__ __forceinline__ void normalize(Output& acc_o) {
+            auto output = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+#pragma unroll
+            for (int m = 0; m < Rows; ++m) {
+                float total = sum(m);
+                total += __shfl_xor_sync(0xffffffff, total, 2);
+                total += __shfl_xor_sync(0xffffffff, total, 1);
+                const float inverse = Masked && total == 0.f ? 0.f : 1.f / total;
+#pragma unroll
+                for (int n = 0; n < size<1>(output); ++n)
+                    output(m, n) *= inverse;
+            }
+        }
+    };
+
+    template <int BM, int BN, int D = 128, int Warps = 4>
+    struct Traits {
+        using Element                    = cutlass::half_t;
+        static constexpr int Threads     = Warps * 32;
+        using Mma                        = TiledMMA<MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
+                                                    Layout<Shape<Int<Warps>, _1, _1>>,
+                                                    Tile<Int<16 * Warps>, _16, _16>>;
+        static constexpr int SmemWidth   = D % 64 == 0 ? 64 : D % 32 == 0 ? 32
+                                                                          : 16;
+        static constexpr int SwizzleBits = SmemWidth == 64 ? 3 : SmemWidth == 32 ? 2
+                                                                                 : 1;
+        using SmemAtom                   = decltype(composition(Swizzle<SwizzleBits, 3, 3>{},
+                                                                Layout<Shape<_8, Int<SmemWidth>>, Stride<Int<SmemWidth>, _1>>{}));
+        using QLayout                    = decltype(tile_to_shape(SmemAtom{}, Shape<Int<BM>, Int<D>>{}));
+        using KVLayout                   = decltype(tile_to_shape(SmemAtom{}, Shape<Int<BN>, Int<D>>{}));
+        using VtLayout                   = decltype(composition(KVLayout{}, make_layout(Shape<Int<D>, Int<BN>>{}, GenRowMajor{})));
+        using CopyAtom                   = Copy_Atom<SM75_U32x4_LDSM_N, Element>;
+        using CopyAtomT                  = Copy_Atom<SM75_U16x8_LDSM_T, Element>;
+        using GlobalLayout               = Layout<Shape<Int<Threads / (SmemWidth / 8)>, Int<SmemWidth / 8>>, Stride<Int<SmemWidth / 8>, _1>>;
+        using GlobalCopy                 = decltype(make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, Element>{},
+                                                                    GlobalLayout{},
+                                                                    Layout<Shape<_1, _8>>{}));
+        using OutputCopy                 = decltype(make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
+                                                                    GlobalLayout{},
+                                                                    Layout<Shape<_1, _8>>{}));
+        static constexpr int SharedBytes = (BM + 2 * BN) * D * sizeof(Element);
+    };
+
+    struct Params {
+        const cutlass::half_t *q, *k, *v;
+        cutlass::half_t* o;
+        int nq, nk, heads, kv_heads, dim;
+        float* output_f32  = nullptr;
+        float output_scale = 1.0f, output_bias = 0.0f;
+        bool scaled_output = false;
+    };
+
+    template <int BM, int BN, int D = 128, bool FullDim = true, int Warps = 4>
+    __global__ __launch_bounds__(Warps * 32) void kernel(Params p) {
+        using T       = Traits<BM, BN, D, Warps>;
+        using Element = typename T::Element;
+        extern __shared__ char storage[];
+        const int tid = threadIdx.x, m_block = blockIdx.x, batch = blockIdx.y, head = blockIdx.z;
+        const int64_t q_offset  = (int64_t(batch) * p.heads + head) * p.nq * p.dim;
+        const int64_t kv_offset = (int64_t(batch) * p.kv_heads + head / (p.heads / p.kv_heads)) * p.nk * p.dim;
+        auto mQ                 = make_tensor(make_gmem_ptr(p.q + q_offset), make_shape(p.nq, p.dim), make_stride(p.dim, _1{}));
+        auto mK                 = make_tensor(make_gmem_ptr(p.k + kv_offset), make_shape(p.nk, p.dim), make_stride(p.dim, _1{}));
+        auto mV                 = make_tensor(make_gmem_ptr(p.v + kv_offset), make_shape(p.nk, p.dim), make_stride(p.dim, _1{}));
+        auto gQ                 = local_tile(mQ, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+        auto gK                 = local_tile(mK, Shape<Int<BN>, Int<D>>{}, make_coord(_, 0));
+        auto gV                 = local_tile(mV, Shape<Int<BN>, Int<D>>{}, make_coord(_, 0));
+        auto sQ                 = make_tensor(make_smem_ptr(reinterpret_cast<Element*>(storage)), typename T::QLayout{});
+        auto sK                 = make_tensor(sQ.data() + size(sQ), typename T::KVLayout{});
+        auto sV                 = make_tensor(sK.data() + size(sK), typename T::KVLayout{});
+        auto sVt                = make_tensor(sV.data(), typename T::VtLayout{});
+        auto sVplain            = make_tensor(sV.data().get(), get_nonswizzle_portion(typename T::VtLayout{}));
+        typename T::GlobalCopy global_copy;
+        auto thread_copy = global_copy.get_thread_slice(tid);
+        auto tQgQ        = thread_copy.partition_S(gQ);
+        auto tQsQ        = thread_copy.partition_D(sQ);
+        auto tKgK        = thread_copy.partition_S(gK);
+        auto tKsK        = thread_copy.partition_D(sK);
+        auto tVgV        = thread_copy.partition_S(gV);
+        auto tVsV        = thread_copy.partition_D(sV);
+        auto tQcQ        = thread_copy.partition_S(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
+        auto tKVcKV      = thread_copy.partition_S(make_identity_tensor(Shape<Int<BN>, Int<D>>{}));
+        typename T::Mma mma;
+        auto thread_mma = mma.get_thread_slice(tid);
+        auto rQ         = thread_mma.partition_fragment_A(sQ);
+        auto rK         = thread_mma.partition_fragment_B(sK);
+        auto rV         = thread_mma.partition_fragment_B(sVplain);
+        auto acc_o      = partition_fragment_C(mma, Shape<Int<BM>, Int<D>>{});
+        auto copyQ      = make_tiled_copy_A(typename T::CopyAtom{}, mma);
+        auto copyK      = make_tiled_copy_B(typename T::CopyAtom{}, mma);
+        auto copyV      = make_tiled_copy_B(typename T::CopyAtomT{}, mma);
+        auto threadQ    = copyQ.get_thread_slice(tid);
+        auto threadK    = copyK.get_thread_slice(tid);
+        auto threadV    = copyV.get_thread_slice(tid);
+        auto tSsQ       = threadQ.partition_S(sQ);
+        auto tSsK       = threadK.partition_S(sK);
+        auto tOsV       = threadV.partition_S(sVt);
+        auto score_id   = thread_mma.partition_C(make_identity_tensor(Shape<Int<BM>, Int<BN>>{}));
+        // Process the partial K/V tile first; subsequent tiles need no score masking.
+        int block = (p.nk - 1) / BN;
+        copy_rows<false, true, FullDim>(global_copy, tQgQ, tQsQ, tQcQ, p.nq - m_block * BM, p.dim);
+        auto last_k = tKgK(_, _, _, block);
+        copy_rows<false, true, FullDim>(global_copy, last_k, tKsK, tKVcKV, p.nk - block * BN, p.dim);
+        cute::cp_async_fence();
+        clear(acc_o);
+        Softmax<2 * size<1>(acc_o)> softmax;
+        auto step = [&](auto first_flag, int n) {
+            constexpr bool First = decltype(first_flag)::value;
+            auto acc_s           = partition_fragment_C(mma, Shape<Int<BM>, Int<BN>>{});
+            clear(acc_s);
+            cute::cp_async_wait<0>();
+            __syncthreads();
+            auto v_tile = tVgV(_, _, _, n);
+            copy_rows<!First, First, FullDim>(global_copy, v_tile, tVsV, tKVcKV, p.nk - n * BN, p.dim);
+            cute::cp_async_fence();
+            gemm(acc_s, rQ, rK, tSsQ, tSsK, mma, copyQ, copyK, threadQ, threadK);
+            if constexpr (First) {
+#pragma unroll
+                for (int i = 0; i < size(acc_s); ++i) {
+                    if (get<1>(score_id(i)) >= p.nk - n * BN)
+                        acc_s(i) = -INFINITY;
+                }
+            }
+            cute::cp_async_wait<0>();
+            __syncthreads();
+            // Do not commit an empty async group: the next iteration waits for this K tile.
+            if (n > 0) {
+                auto next_k = tKgK(_, _, _, n - 1);
+                copy_rows<true, false, FullDim>(global_copy, next_k, tKsK, tKVcKV, 0, p.dim);
+                cute::cp_async_fence();
+            }
+            softmax.template update<First>(acc_s, acc_o);
+            auto rP   = to_half(acc_s);
+            auto tOrP = make_tensor(rP.data(), convert_layout_acc_Aregs<typename T::Mma>(rP.layout()));
+            gemm_rs(acc_o, tOrP, rV, tOsV, mma, copyV, threadV);
+        };
+        step(std::true_type{}, block);
+        for (--block; block >= 0; --block)
+            step(std::false_type{}, block);
+        softmax.normalize(acc_o);
+        auto rO      = to_half(acc_o);
+        auto sO      = make_tensor(sQ.data(), typename T::QLayout{});
+        auto copyO   = make_tiled_copy_C(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{}, mma);
+        auto threadO = copyO.get_thread_slice(tid);
+        auto tOrO    = threadO.retile_S(rO);
+        auto tOsO    = threadO.partition_D(sO);
+        cute::copy(copyO, tOrO, tOsO);
+        typename T::OutputCopy output_copy;
+        auto global_threadO = output_copy.get_thread_slice(tid);
+        auto from_sO        = global_threadO.partition_S(sO);
+        auto output_id      = global_threadO.partition_D(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
+        __syncthreads();
+        auto output_regs = make_tensor<Element>(shape(from_sO));
+        cute::copy(output_copy, from_sO, output_regs);
+        const int64_t output_offset = int64_t(batch) * p.nq * p.heads * p.dim + int64_t(head) * p.dim;
+        if (p.output_f32) {
+            auto float_copy = make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, float>{},
+                                              typename T::GlobalLayout{}, Layout<Shape<_1, _8>>{});
+            auto mO         = make_tensor(make_gmem_ptr(p.output_f32 + output_offset), make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
+            auto gO         = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+            auto to_gO      = float_copy.get_thread_slice(tid).partition_D(gO);
+            auto values     = make_tensor<float>(shape(output_regs));
+#pragma unroll
+            for (int i = 0; i < size(values); ++i) {
+                // Preserve the half-rounded attention result before applying the F32 epilogue.
+                const float value = float(output_regs(i));
+                values(i)         = p.scaled_output ? __fmaf_rn(value, p.output_scale, p.output_bias) : value;
+            }
+            copy_rows<false, false, FullDim>(float_copy, values, to_gO, output_id, p.nq - m_block * BM, p.dim);
+        } else {
+            auto mO    = make_tensor(make_gmem_ptr(p.o + output_offset), make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
+            auto gO    = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+            auto to_gO = global_threadO.partition_D(gO);
+            copy_rows<false, false, FullDim>(output_copy, output_regs, to_gO, output_id, p.nq - m_block * BM, p.dim);
+        }
+    }
+}
+
+static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+    a = index == 3 ? -a : a;
+    b = index == 2 ? -b : b;
+    c = index == 1 ? -c : c;
+    d = index == 0 ? -d : d;
+    return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+}
+
+static __global__ void quantize_attention(const half* src, int8_t* dst, float* scales, int64_t k, int64_t rows, float input_scale, float input_bias) {
+    const int64_t row = blockIdx.x;
+    const int tid     = threadIdx.x;
+    if (row >= rows)
+        return;
+    extern __shared__ float values[];
+    __shared__ float maxima[32];
+    const half* src_row = src + row * k;
+    int8_t* dst_row     = dst + row * k;
+    const int lane_id   = tid % 32;
+    {
+        const int lane           = lane_id;
+        constexpr int group_size = 256;
+        for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
+            float v[8];
+#pragma unroll
+            for (int j = 0; j < 8; ++j) {
+                const float value = __fmaf_rn(__half2float(src_row[group * group_size + j * 32 + lane]), input_scale, input_bias);
+                v[j]              = value * (1.0f / 16.0f);
+            }
+            // One logical warp owns H256, keeping every butterfly in registers.
+#pragma unroll
+            for (int stride = 1; stride <= 4; stride *= 4) {
+                const int base = lane & ~(3 * stride);
+#pragma unroll
+                for (int j = 0; j < 8; ++j) {
+                    const float a = __shfl_sync(0xffffffff, v[j], base, 32);
+                    const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
+                    const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
+                    const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32);
+                    v[j]          = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
+                }
+            }
+#pragma unroll
+            for (int j = 0; j < 8; j += 2) {
+                const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32);
+                const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
+                const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
+                const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32);
+                v[j]          = convrot_h4_cuda(a, b, c, d, lane / 16);
+                v[j + 1]      = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
+            }
+#pragma unroll
+            for (int j = 0; j < 2; ++j) {
+                const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
+#pragma unroll
+                for (int t = 0; t < 4; ++t) {
+                    values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
+                }
+            }
+        }
+        __syncthreads();
+    }
+    float local_max = 0.0f;
+    for (int64_t i = tid; i < k; i += blockDim.x) {
+        local_max = fmaxf(local_max, fabsf(values[i]));
+    }
+    local_max      = sd_warp_max(local_max);
+    const int lane = tid & 31;
+    const int warp = tid >> 5;
+    if (lane == 0) {
+        maxima[warp] = local_max;
+    }
+    __syncthreads();
+
+    if (warp == 0) {
+        local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
+        local_max = sd_warp_max(local_max);
+        if (lane == 0) {
+            maxima[0] = local_max / 127.0f;
+            if (scales != nullptr) {
+                scales[row] = maxima[0];
+            } else {
+                *((float*)(dst_row + k)) = maxima[0];
+            }
+        }
+    }
+    __syncthreads();
+
+    const float row_scale = maxima[0];
+    const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+    for (int64_t i = tid; i < k; i += blockDim.x) {
+        int value  = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
+        value      = max(-127, min(127, value));
+        dst_row[i] = (int8_t)value;
+    }
+}
+
+template <typename T>
+static __global__ void copy_to_half(const char* src, half* dst, int64_t count, int64_t dim, int64_t rows, int64_t heads, size_t nb1, size_t nb2, size_t nb3, float scale) {
+    const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (i < count) {
+        const int64_t row = i / dim;
+        const auto* p     = reinterpret_cast<const T*>(src + (row % rows) * nb1 +
+                                                       ((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3);
+        dst[i]            = __float2half_rn(float(p[i % dim]) * scale);
+    }
+}
+
+static __global__ void copy_from_half(const half* src, float* dst, int64_t count, float scale, float bias, bool scaled) {
+    const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (i < count)
+        dst[i] = scaled ? __fmaf_rn(__half2float(src[i]), scale, bias) : __half2float(src[i]);
+}
+
+static float param_f32(const ggml_tensor* node, int index) {
+    float value;
+    std::memcpy(&value, reinterpret_cast<const char*>(node->op_params) + index * sizeof(value), sizeof(value));
+    return value;
+}
+
+static int padded_dim(int64_t dim) {
+    return dim == 80 || dim == 96 || dim == 112 ? int(dim) : int((dim + 63) / 64 * 64);
+}
+
+struct CUDADevice {
+    int ordinal;
+    cudaDeviceProp properties;
+};
+
+static const CUDADevice* cuda_device(ggml_backend_dev_t device) {
+    static std::mutex mutex;
+    static std::unordered_map<ggml_backend_dev_t, CUDADevice> devices;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto found = devices.find(device);
+    if (found != devices.end())
+        return &found->second;
+    auto reg   = ggml_backend_dev_backend_reg(device);
+    auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+    if (!query)
+        return nullptr;
+    CUDADevice info = {};
+    info.ordinal    = query(device);
+    if (info.ordinal < 0 || cudaGetDeviceProperties(&info.properties, info.ordinal) != cudaSuccess)
+        return nullptr;
+    return &devices.emplace(device, info).first->second;
+}
+
+static bool attention_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+    static const bool disabled = std::getenv("GGML_CUDA_DISABLE_CUTE_FATTN") || std::getenv("GGML_CUDA_DISABLE_CUTLASS_FATTN");
+    const auto* info           = cuda_device(device);
+    const auto* q              = node->src[0];
+    const auto* k              = node->src[1];
+    const auto* v              = node->src[2];
+    if (disabled || !info || info->properties.major < 8 || !q || !k || !v ||
+        q->type != GGML_TYPE_F32 || node->type != GGML_TYPE_F32 ||
+        q->ne[0] < 32 || q->ne[0] > 256 || q->ne[0] % 8 ||
+        size_t(192 * padded_dim(q->ne[0]) * sizeof(half)) > info->properties.sharedMemPerBlockOptin ||
+        q->ne[1] < 1 || k->ne[1] < 1 || k->ne[2] < 1 ||
+        node->src[3] || node->src[4] || param_f32(node, 1) != 0.0f || param_f32(node, 2) != 0.0f ||
+        q->ne[2] % k->ne[2] || q->ne[3] != k->ne[3] || !ggml_are_same_shape(k, v) ||
+        q->nb[0] != sizeof(float) || !ggml_is_contiguous(node) ||
+        (k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_F32) ||
+        (v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_F32) ||
+        k->nb[0] != ggml_type_size(k->type) || v->nb[0] != ggml_type_size(v->type))
+        return false;
+    return k->ne[0] == q->ne[0] && q->ne[1] <= INT_MAX / q->ne[0] && k->ne[1] <= INT_MAX / q->ne[0] &&
+           q->ne[2] <= 65535 && q->ne[3] <= 65535 && ggml_nelements(q) <= INT_MAX && ggml_nelements(k) <= INT_MAX;
+}
+
+static bool attention_preferred(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+    const auto* q = node->src[0];
+    const auto* k = node->src[1];
+    if (q->ne[0] == 40 || q->ne[0] == 72)
+        return q->ne[1] >= 512 && k->ne[1] >= 512;
+    if (q->ne[0] == 80 || q->ne[0] == 96 || q->ne[0] == 112)
+        return false;
+    const int minimum = q->ne[0] == 128 ? 2048 : 4096;
+    return q->ne[1] >= minimum && k->ne[1] >= minimum;
+}
+
+static bool attention_selected(ggml_backend_dev_t device, const ggml_tensor* node) {
+    return node->op == GGML_OP_FLASH_ATTN_EXT && attention_supports(device, node, nullptr) &&
+           (attention_preferred(device, node, nullptr) || !ggml_ext_dev_supports_native_op(device, node));
+}
+
+static bool fusion_enabled() {
+    static const bool disabled = [] {
+        const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
+        return value && std::atoi(value) != 0;
+    }();
+    return !disabled;
+}
+
+static const ggml_tensor* graph_node(const ggml_cgraph* graph, int index) {
+    return ggml_graph_node(const_cast<ggml_cgraph*>(graph), index);
+}
+
+static int output_fusion(ggml_backend_dev_t device, const ggml_cgraph* graph, int first) {
+    const auto* attn     = graph_node(graph, first);
+    const auto* previous = attn;
+    const int total      = ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph));
+    bool scaled          = false;
+    for (int count = 1; count < 8 && first + count < total; ++count) {
+        const auto* node = graph_node(graph, first + count);
+        int convrot_group;
+        if (node->src[0] != previous)
+            break;
+        if (node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE) {
+            if (node->view_offs || !ggml_is_contiguous(node) || ggml_nelements(node) != ggml_nelements(attn))
+                break;
+        } else if (node->op == GGML_OP_SCALE && !scaled) {
+            if (!ggml_is_contiguous(node) || !ggml_are_same_shape(previous, node))
+                break;
+            scaled = true;
+        } else if (sd_ggml_convrot_describe(node, &convrot_group)) {
+            if (convrot_group == 256 && previous->ne[0] % 256 == 0 &&
+                size_t(previous->ne[0]) * sizeof(float) + 128 <= cuda_device(device)->properties.sharedMemPerBlockOptin)
+                return count + 1;
+            break;
+        } else
+            break;
+        previous = node;
+    }
+    if (first + 1 < total) {
+        const auto* next = graph_node(graph, first + 1);
+        if (next->op == GGML_OP_SCALE && next->src[0] == attn && ggml_is_contiguous(next) && ggml_are_same_shape(attn, next))
+            return 2;
+    }
+    return 1;
+}
+
+static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+    if (!fusion_enabled())
+        return 0;
+    int offset = 0;
+    if (ggml_ext_op_get(graph_node(graph, first)) == sd_ggml_i8_op()) {
+        if (first + 5 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
+            !sd_cuda_i8_rope_supports(device, graph_node(graph, first), graph_node(graph, first + 1),
+                                      graph_node(graph, first + 2), graph_node(graph, first + 3), graph_node(graph, first + 4)) ||
+            graph_node(graph, first + 5)->src[0] != graph_node(graph, first + 4))
+            return 0;
+        offset = 5;
+    } else if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
+        if (first + 3 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
+            !sd_cuda::rope_fusion_supports(graph_node(graph, first), graph_node(graph, first + 1), graph_node(graph, first + 2)) ||
+            graph_node(graph, first + 3)->src[0] != graph_node(graph, first + 2))
+            return 0;
+        offset = 3;
+    }
+    if (!attention_selected(device, graph_node(graph, first + offset)))
+        return 0;
+    const int count = offset + output_fusion(device, graph, first + offset);
+    return count > 1 ? count : 0;
+}
+
+static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
+    if (!fusion_enabled())
+        return;
+    for (int i = 3; i < ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)); ++i) {
+        const auto* attn = graph_node(graph, i);
+        if (!attention_selected(device, attn))
+            continue;
+        for (int j = i - 3; j >= 0 && j >= i - 32; --j) {
+            if (attn->src[0] == graph_node(graph, j + 2) &&
+                sd_cuda::rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
+                if (j >= 2 && sd_cuda_i8_rope_supports(device, graph_node(graph, j - 2), graph_node(graph, j - 1), graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2)) &&
+                    ggml_ext_graph_delay(graph, j - 2, 5, i))
+                    break;
+                ggml_ext_graph_delay(graph, j, 3, i);
+                break;
+            }
+        }
+    }
+}
+
+struct AttentionPlan {
+    int index;
+    size_t offsets[4];
+    bool direct_output;
+};
+
+static bool attention_direct_input(const ggml_tensor* src) {
+    size_t offset     = 0;
+    const auto* owner = src;
+    while (!owner->data && owner->view_src) {
+        offset += owner->view_offs;
+        owner = owner->view_src;
+    }
+    return src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && (uintptr_t(owner->data) + offset) % 16 == 0;
+}
+
+static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void** state, size_t* workspace_size, void*) {
+    auto* plan  = new AttentionPlan{};
+    *state      = plan;
+    plan->index = ggml_ext_op_get(nodes[0]) == sd_ggml_i8_op() ? 5 : nodes[0]->op == GGML_OP_RMS_NORM ? 3
+                                                                                                      : 0;
+    GGML_ASSERT(plan->index < count);
+    const auto* attn    = nodes[plan->index];
+    const auto* dst     = nodes[count - 1];
+    plan->direct_output = dst->type == GGML_TYPE_F32 && !dst->data && !dst->buffer && !dst->view_src;
+    size_t size         = 0;
+    for (int i = 0; i < 4; ++i) {
+        if ((i > 0 && i < 3 && attention_direct_input(attn->src[i])) || (i == 3 && plan->direct_output)) {
+            plan->offsets[i] = SIZE_MAX;
+            continue;
+        }
+        plan->offsets[i]   = size;
+        const size_t bytes = size_t(ggml_nelements(i < 3 ? attn->src[i] : attn)) * sizeof(half);
+        size += (bytes + 255) / 256 * 256;
+    }
+    *workspace_size = size;
+    return GGML_STATUS_SUCCESS;
+}
+
+template <int D, bool FullDim>
+static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
+    auto kernel                = sd_fattn_cute::kernel<64, 64, D, FullDim>;
+    constexpr int shared_bytes = sd_fattn_cute::Traits<64, 64, D>::SharedBytes;
+    static std::mutex mutex;
+    static std::unordered_map<int, bool> initialized;
+    {
+        std::lock_guard<std::mutex> lock(mutex);
+        if (!initialized[context.device]) {
+            auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes);
+            if (error != cudaSuccess)
+                return error;
+            initialized[context.device] = true;
+        }
+    }
+    kernel<<<dim3((p.nq + 63) / 64, batches, p.heads), 128, shared_bytes, static_cast<cudaStream_t>(context.stream)>>>(p);
+    return cudaGetLastError();
+}
+
+template <int D>
+static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
+    return p.dim == D ? attention_launch<D, true>(context, p, batches) : attention_launch<D, false>(context, p, batches);
+}
+
+static ggml_status attention_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+    const auto& plan = *static_cast<AttentionPlan*>(state);
+    const auto* attn = nodes[plan.index];
+    auto device      = ggml_backend_get_device(execution->backend);
+    auto reg         = ggml_backend_dev_backend_reg(device);
+    auto query       = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+    ggml_ext_cuda_context context;
+    if (!query || !query(execution->backend, &context))
+        return GGML_STATUS_FAILED;
+    auto stream = static_cast<cudaStream_t>(context.stream);
+    auto* base  = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+    half* work[4];
+    for (int i = 0; i < 4; ++i)
+        work[i] = plan.offsets[i] == SIZE_MAX ? nullptr : reinterpret_cast<half*>(base + plan.offsets[i]);
+    const half* inputs[3] = {work[0], work[1], work[2]};
+    for (int i = 0; i < 3; ++i) {
+        const auto* src   = attn->src[i];
+        const float scale = i == 0 ? param_f32(attn, 0) : 1.0f;
+        if (i == 0 && plan.index == 5) {
+            if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess)
+                return GGML_STATUS_FAILED;
+        } else if (i == 0 && plan.index == 3) {
+            if (sd_cuda::rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess)
+                return GGML_STATUS_FAILED;
+        } else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
+            inputs[i] = static_cast<const half*>(src->data);
+        } else {
+            if (!work[i])
+                return GGML_STATUS_FAILED;
+            auto kernel = src->type == GGML_TYPE_F16 ? copy_to_half<half> : copy_to_half<float>;
+            kernel<<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast<const char*>(src->data), work[i],
+                                                                          ggml_nelements(src), src->ne[0], src->ne[1], src->ne[2], src->nb[1], src->nb[2], src->nb[3], scale);
+        }
+        if (cudaGetLastError() != cudaSuccess)
+            return GGML_STATUS_FAILED;
+    }
+    const auto* q = attn->src[0];
+    const auto* k = attn->src[1];
+    auto* dst     = nodes[count - 1];
+    float scale = 1.0f, bias = 0.0f;
+    for (int i = plan.index + 1; i < count; ++i) {
+        if (nodes[i]->op == GGML_OP_SCALE) {
+            scale = param_f32(nodes[i], 0);
+            bias  = param_f32(nodes[i], 1);
+        }
+    }
+    sd_fattn_cute::Params p = {reinterpret_cast<const cutlass::half_t*>(inputs[0]),
+                               reinterpret_cast<const cutlass::half_t*>(inputs[1]), reinterpret_cast<const cutlass::half_t*>(inputs[2]),
+                               reinterpret_cast<cutlass::half_t*>(work[3]), int(q->ne[1]), int(k->ne[1]), int(q->ne[2]), int(k->ne[2]), int(q->ne[0])};
+    if (plan.direct_output) {
+        p.output_f32    = static_cast<float*>(dst->data);
+        p.output_scale  = scale;
+        p.output_bias   = bias;
+        p.scaled_output = dst->op == GGML_OP_SCALE;
+    }
+    cudaError_t error = cudaErrorInvalidValue;
+    switch (padded_dim(p.dim)) {
+        case 64:
+            error = attention_launch<64>(context, p, q->ne[3]);
+            break;
+        case 80:
+            error = attention_launch<80, true>(context, p, q->ne[3]);
+            break;
+        case 96:
+            error = attention_launch<96, true>(context, p, q->ne[3]);
+            break;
+        case 112:
+            error = attention_launch<112, true>(context, p, q->ne[3]);
+            break;
+        case 128:
+            error = attention_launch<128>(context, p, q->ne[3]);
+            break;
+        case 192:
+            error = attention_launch<192>(context, p, q->ne[3]);
+            break;
+        case 256:
+            error = attention_launch<256>(context, p, q->ne[3]);
+            break;
+    }
+    if (error != cudaSuccess)
+        return GGML_STATUS_FAILED;
+    if (plan.direct_output)
+        return GGML_STATUS_SUCCESS;
+    int convrot_group;
+    if (sd_ggml_convrot_describe(dst, &convrot_group)) {
+        const int64_t width = dst->src[0]->ne[0], rows = ggml_nrows(dst->src[0]), padded_rows = (rows + 3) / 4 * 4;
+        auto* packed = static_cast<int8_t*>(dst->data);
+        if (rows != padded_rows && cudaMemsetAsync(packed + width * rows, 0, width * (padded_rows - rows), stream) != cudaSuccess)
+            return GGML_STATUS_FAILED;
+        if (cudaFuncSetAttribute(quantize_attention, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                 int(cuda_device(device)->properties.sharedMemPerBlockOptin - 128)) != cudaSuccess)
+            return GGML_STATUS_FAILED;
+        quantize_attention<<<rows, 256, width * sizeof(float), stream>>>(work[3], packed,
+                                                                         reinterpret_cast<float*>(packed + width * padded_rows), width, rows, scale, bias);
+    } else {
+        copy_from_half<<<(ggml_nelements(attn) + 255) / 256, 256, 0, stream>>>(work[3], static_cast<float*>(dst->data), ggml_nelements(attn), scale, bias, dst->op == GGML_OP_SCALE);
+    }
+    return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+}
+
+static void attention_destroy(void* state, void*) {
+    delete static_cast<AttentionPlan*>(state);
+}
+
+void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
+    if (!cuda_device(device))
+        return;
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.cute_attention.cuda.v1";
+    desc.op                   = GGML_OP_FLASH_ATTN_EXT;
+    desc.priority             = 200;
+    desc.supports             = attention_supports;
+    desc.preferred            = attention_preferred;
+    desc.match                = attention_match;
+    desc.optimize             = attention_optimize;
+    desc.prepare              = attention_prepare;
+    desc.execute              = attention_execute;
+    desc.destroy              = attention_destroy;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    desc.name      = "sd.rms_rope_cute_attention.cuda.v1";
+    desc.op        = GGML_OP_RMS_NORM;
+    desc.priority  = 300;
+    desc.supports  = nullptr;
+    desc.preferred = nullptr;
+    desc.optimize  = nullptr;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    desc.name      = "sd.i8_rope_cute_attention.cuda.v1";
+    desc.op        = GGML_OP_CUSTOM;
+    desc.custom_op = sd_ggml_i8_op();
+    desc.priority  = 400;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}
+
+#include "window_attention.cuh"

src/core/cuda/i8_gated_residual.cu

@@ -0,0 +1,485 @@
+#include <cuda_runtime.h>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
+#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
+#include <unordered_map>
+#include "core/ggml_extensions.h"
+
+namespace sd_i8_gated {
+    struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
+        using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
+        using Base::Base;
+
+        CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
+            const auto shape = tiled_shape();
+            if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
+                // Column bands retain weight tiles in Ada's L2 across successive row tiles.
+                constexpr int band_width = 16;
+                if (index >= shape.m() * shape.n()) {
+                    return {shape.m(), shape.n(), get_batch_idx()};
+                }
+                const int band   = index / (shape.m() * band_width);
+                const int width  = min(band_width, shape.n() - band * band_width);
+                const int within = index - band * shape.m() * band_width;
+                return {within / width, band * band_width + within % width, get_batch_idx()};
+            }
+#endif
+            return Base::get_tile_offset(index);
+        }
+    };
+
+    template <int M, int N, int WM>
+    struct TileConfig {
+        using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+        using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+    };
+
+    enum Tile {
+        TILE_128_256,
+        TILE_128_128,
+        TILE_64_128,
+        TILE_16_128,
+    };
+
+    static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
+        static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
+        return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
+    }
+
+    static Tile select_tile(int64_t rows, int64_t n, int nsm) {
+        if (rows < 32) {
+            return TILE_16_128;
+        }
+        const int64_t columns = (n + 127) / 128;
+        if (((rows + 15) / 16) * columns <= nsm) {
+            return TILE_16_128;
+        }
+        if (((rows + 63) / 64) * columns <= nsm) {
+            return TILE_64_128;
+        }
+        const int64_t wide_tiles   = ((rows + 127) / 128) * ((n + 255) / 256);
+        const int64_t square_tiles = ((rows + 127) / 128) * columns;
+        // A narrower tile costs more per output element, but can waste less of the last wave.
+        const double wide_cost   = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
+        const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
+        return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
+    }
+
+    template <class T>
+    struct MulRound {
+        CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
+            T result;
+            CUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < T::kElements; ++i) {
+                result[i] = __fmul_rn(a[i], b[i]);
+            }
+            return result;
+        }
+    };
+
+    template <class T>
+    struct AddRound {
+        CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
+            T result;
+            CUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < T::kElements; ++i) {
+                result[i] = __fadd_rn(a[i], b[i]);
+            }
+            return result;
+        }
+    };
+
+    template <class Config, bool scaled, bool stream_k>
+    struct KernelTypes {
+        using Tile         = typename Config::Tile;
+        using Warp         = typename Config::Warp;
+        using Swizzle      = std::conditional_t<stream_k, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Map          = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Col          = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Row          = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Mul          = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using ScaleX       = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleW       = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using Residual     = cutlass::epilogue::threadblock::VisitorAuxLoad<Map, float, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using GateMul      = cutlass::epilogue::threadblock::VisitorCompute<MulRound, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using Add          = cutlass::epilogue::threadblock::VisitorCompute<AddRound, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using Scalar       = cutlass::epilogue::threadblock::VisitorScalarBroadcast<float>;
+        using ScaleOut     = cutlass::epilogue::threadblock::Sm80EVT<GateMul, ScaleW, Scalar>;
+        using Scaled       = std::conditional_t<scaled, ScaleOut, ScaleW>;
+        using Gated        = cutlass::epilogue::threadblock::Sm80EVT<GateMul, Scaled, Row>;
+        using WithResidual = cutlass::epilogue::threadblock::Sm80EVT<Add, Gated, Residual>;
+        using Output       = WithResidual;
+        using Store        = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using Epilogue     = cutlass::epilogue::threadblock::Sm80EVT<Store, Output>;
+        using Kernel       = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+            int8_t,
+            cutlass::layout::RowMajor,
+            cutlass::ComplexTransform::kNone,
+            16,
+            int8_t,
+            cutlass::layout::ColumnMajor,
+            cutlass::ComplexTransform::kNone,
+            16,
+            float,
+            cutlass::layout::RowMajor,
+            4,
+            int,
+            float,
+            cutlass::arch::OpClassTensorOp,
+            cutlass::arch::Sm80,
+            Tile,
+            Warp,
+            cutlass::gemm::GemmShape<16, 8, 32>,
+            Epilogue,
+            Swizzle,
+            4,
+            cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
+        using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+    };
+
+    struct DeviceInfo {
+        int ordinal        = -1;
+        int major          = 0;
+        int sms            = 0;
+        int threads_per_sm = 0;
+        size_t shared      = 0;
+    };
+    static DeviceInfo device_info(ggml_backend_dev_t device) {
+        static std::mutex mutex;
+        static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
+        std::lock_guard<std::mutex> lock(mutex);
+        auto found = cache.find(device);
+        if (found != cache.end())
+            return found->second;
+        DeviceInfo info;
+        auto reg             = ggml_backend_dev_backend_reg(device);
+        auto query           = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+        int ordinal          = query ? query(device) : -1;
+        cudaDeviceProp props = {};
+        if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
+            info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+        cache.emplace(device, info);
+        return info;
+    }
+    static float param_f32(const ggml_tensor* tensor, int index) {
+        float value;
+        std::memcpy(&value, tensor->op_params + index, sizeof(value));
+        return value;
+    }
+    static bool aligned(const ggml_tensor* tensor) {
+        if (!tensor)
+            return false;
+        size_t offset = 0;
+        while (tensor->view_src) {
+            offset += tensor->view_offs;
+            tensor = tensor->view_src;
+        }
+        return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % 16 == 0;
+    }
+    struct Operation {
+        SDI8Operation mat;
+        const ggml_tensor* gate;
+        const ggml_tensor* residual;
+        const ggml_tensor* prefix;
+        float scale;
+        int64_t rows, n, k;
+    };
+    static Operation describe(ggml_tensor* const* nodes, int count) {
+        Operation op = {};
+        sd_ggml_i8_describe(nodes[0], &op.mat);
+        const bool concat = nodes[count - 1]->op == GGML_OP_CONCAT;
+        const auto* add   = nodes[count - (concat ? 2 : 1)];
+        const auto* mul   = nodes[count - (concat ? 3 : 2)];
+        op.gate           = mul->src[1];
+        op.residual       = add->src[0] == mul ? add->src[1] : add->src[0];
+        op.prefix         = concat ? nodes[count - 1]->src[0] : nullptr;
+        op.scale          = nodes[1]->op == GGML_OP_SCALE ? param_f32(nodes[1], 0) : 1.f;
+        op.rows           = ggml_nrows(nodes[0]);
+        op.n              = nodes[0]->ne[0];
+        op.k              = op.mat.weight->ne[0];
+        return op;
+    }
+    template <class Types>
+    static typename Types::Gemm::Arguments arguments(const Operation& op, float* output, bool metadata) {
+        const auto* x                            = metadata ? nullptr : static_cast<const int8_t*>(op.mat.input->data);
+        const auto* w                            = metadata ? nullptr : static_cast<const int8_t*>(op.mat.weight->data);
+        const auto* xs                           = metadata ? nullptr : op.mat.packed ? reinterpret_cast<const float*>(x + op.k * ((op.rows + 3) / 4 * 4))
+                                                                                      : static_cast<const float*>(op.mat.input_scale->data);
+        const auto* ws                           = metadata ? nullptr : static_cast<const float*>(op.mat.weight_scale->data);
+        const auto* gate                         = metadata ? nullptr : static_cast<const float*>(op.gate->data);
+        auto* residual                           = metadata ? nullptr : static_cast<float*>(op.residual->data);
+        typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
+        typename Types::Scaled::Arguments scaled;
+        if constexpr (std::is_same<typename Types::Scaled, typename Types::ScaleW>::value)
+            scaled = scales;
+        else
+            scaled = {scales, {{op.scale}}, {}};
+        typename Types::Epilogue::Arguments epilogue = {{{scaled, {gate, 0.f, {}}, {}}, {residual, 0.f, {op.n, cute::_1{}, op.rows * op.n}}, {}}, {output, {op.n, cute::_1{}, op.rows * op.n}}};
+        return {cutlass::gemm::GemmUniversalMode::kGemm, {int(op.rows), int(op.n), int(op.k)}, 1, epilogue, x, w, nullptr, nullptr, op.rows * op.k, op.n * op.k, 0, 0, op.k, op.k, 0, 0};
+    }
+    struct Plan {
+        DeviceInfo device;
+        size_t kernel_workspace                                                     = 0;
+        size_t output_offset                                                        = 0;
+        int occupancy                                                               = 0;
+        bool output_scratch                                                         = true;
+        ggml_status (*launch)(Plan&, const Operation&, float*, void*, cudaStream_t) = nullptr;
+    };
+    template <class Types>
+    static ggml_status launch(Plan& plan, const Operation& op, float* output, void* scratch, cudaStream_t stream) {
+        using Kernel = typename Types::Kernel;
+        auto args    = arguments<Types>(op, output, false);
+        if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+            return GGML_STATUS_FAILED;
+        if (plan.occupancy == 0) {
+            if (sizeof(typename Kernel::SharedStorage) >= (48 << 10) &&
+                cudaFuncSetAttribute(cutlass::Kernel2<Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess)
+                return GGML_STATUS_FAILED;
+            int occupancy = 0;
+            if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&occupancy, cutlass::Kernel2<Kernel>,
+                                                                       Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), cudaOccupancyDisableCachingOverride) != cudaSuccess ||
+                occupancy <= 0)
+                return GGML_STATUS_FAILED;
+            plan.occupancy = occupancy;
+        }
+        typename Kernel::Params params(args, plan.device.sms, plan.occupancy);
+        if (params.get_workspace_size() > plan.kernel_workspace || params.init_workspace(scratch, stream) != cutlass::Status::kSuccess)
+            return GGML_STATUS_FAILED;
+        cutlass::arch::synclog_setup();
+        cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
+        return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+    template <class Config, bool Scaled, bool StreamK>
+    static bool configure(const Operation& op, Plan& plan) {
+        using Types  = KernelTypes<Config, Scaled, StreamK>;
+        using Kernel = typename Types::Kernel;
+        if (sizeof(typename Kernel::SharedStorage) > plan.device.shared)
+            return false;
+        auto args = arguments<Types>(op, nullptr, true);
+        if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+            return false;
+        size_t workspace = 0;
+        // Bound every possible occupancy without setting a CUDA device or initializing a kernel during graph measurement.
+        for (int occupancy = 1; occupancy <= plan.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
+            typename Kernel::Params params(args, plan.device.sms, occupancy);
+            const auto grid = params.get_grid_dims();
+            if (grid.y > 65535 || grid.z > 65535)
+                return false;
+            workspace = std::max(workspace, params.get_workspace_size());
+        }
+        plan.kernel_workspace = workspace;
+        plan.output_offset    = (workspace + 255) / 256 * 256;
+        plan.launch           = launch<Types>;
+        return true;
+    }
+    template <class Config>
+    static bool configure(const Operation& op, Plan& plan) {
+        if (op.scale != 1.f)
+            return op.rows >= 8192 ? configure<Config, true, true>(op, plan) : configure<Config, true, false>(op, plan);
+        return op.rows >= 8192 ? configure<Config, false, true>(op, plan) : configure<Config, false, false>(op, plan);
+    }
+    static bool configure(const Operation& op, Plan& plan) {
+        auto try_tile = [&](Tile tile) {
+            switch (tile) {
+                case TILE_128_256:
+                    return configure<TileConfig<128, 256, 64>>(op, plan);
+                case TILE_128_128:
+                    return configure<TileConfig<128, 128, 64>>(op, plan);
+                case TILE_64_128:
+                    return configure<TileConfig<64, 128, 32>>(op, plan);
+                case TILE_16_128:
+                    return configure<TileConfig<16, 128, 16>>(op, plan);
+            }
+            return false;
+        };
+        const auto tile = select_tile(op.rows, op.n, plan.device.sms);
+        if (try_tile(tile))
+            return true;
+        for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
+            if (fallback != tile && try_tile(fallback))
+                return true;
+        return false;
+    }
+    static bool fusion_disabled() {
+        static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+        return disabled;
+    }
+    static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
+#ifndef SD_CUDA_FORCE_CUBLAS
+        if (fusion_disabled())
+            return;
+        const int total = ggml_graph_n_nodes(graph);
+        for (int i = 2; i + 1 < total; ++i) {
+            const auto* mul = ggml_graph_node(graph, i);
+            const auto* add = ggml_graph_node(graph, i + 1);
+            if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
+                continue;
+            const auto* producer = mul->src[0];
+            const bool scaled    = producer->op == GGML_OP_SCALE;
+            const auto* mat      = scaled ? producer->src[0] : producer;
+            SDI8Operation op;
+            if (!sd_ggml_i8_describe(mat, &op) || op.input->type != GGML_TYPE_I8 || !op.weight_scale || op.bias)
+                continue;
+            int scale_index = -1;
+            for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
+                if (scaled && ggml_graph_node(graph, j) == producer)
+                    scale_index = j;
+                if (ggml_graph_node(graph, j) != mat)
+                    continue;
+                if (scaled) {
+                    if (scale_index < 0 || (j + 1 < scale_index && !ggml_ext_graph_delay(graph, j, 1, scale_index)))
+                        break;
+                    j = scale_index - 1;
+                }
+                ggml_ext_graph_delay(graph, j, scaled ? 2 : 1, i);
+                break;
+            }
+        }
+#endif
+    }
+    static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+#ifdef SD_CUDA_FORCE_CUBLAS
+        return 0;
+#else
+        const int available = ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) - first;
+        if (fusion_disabled() || available < 3)
+            return 0;
+        ggml_tensor* nodes[5] = {};
+        for (int i = 0; i < std::min(available, 5); ++i)
+            nodes[i] = ggml_graph_node(const_cast<ggml_cgraph*>(graph), first + i);
+        SDI8Operation mat;
+        if (!sd_ggml_i8_describe(nodes[0], &mat) || mat.input->type != GGML_TYPE_I8 || !mat.weight_scale || mat.bias)
+            return 0;
+        const int64_t rows = ggml_nrows(nodes[0]), n = nodes[0]->ne[0], k = mat.weight->ne[0];
+        if (rows < (small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n > INT_MAX || k > INT_MAX || n % 128 || k % 64 ||
+            ((rows + 15) / 16) > INT_MAX / ((n + 127) / 128))
+            return 0;
+        for (const auto* t : {mat.weight, mat.input, mat.weight_scale, mat.packed ? mat.weight_scale : mat.input_scale})
+            if (!aligned(t))
+                return 0;
+        const bool scaled = nodes[1]->op == GGML_OP_SCALE;
+        int count         = scaled ? 4 : 3;
+        if (available < count || (scaled && (nodes[1]->src[0] != nodes[0] || param_f32(nodes[1], 1) != 0.f)))
+            return 0;
+        const auto* mul = nodes[count - 2];
+        const auto* add = nodes[count - 1];
+        if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || mul->src[0] != nodes[count - 3] ||
+            (add->src[0] != mul && add->src[1] != mul))
+            return 0;
+        const auto* gate     = mul->src[1];
+        const auto* residual = add->src[0] == mul ? add->src[1] : add->src[0];
+        if (!gate || !residual || gate->ne[0] != n || ggml_nelements(gate) != n || !aligned(gate) || !aligned(residual))
+            return 0;
+        for (const auto* t : {gate, residual, mul, add})
+            if (t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || (t != gate && !ggml_are_same_shape(nodes[0], t)))
+                return 0;
+        if (available > count) {
+            const auto* concat = nodes[count];
+            if (concat->op == GGML_OP_CONCAT && concat->src[1] == add && concat->op_params[0] == 1 &&
+                concat->type == GGML_TYPE_F32 && concat->ne[2] == 1 && concat->ne[3] == 1 && ggml_is_contiguous(concat) &&
+                concat->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(concat->src[0]))
+                ++count;
+        }
+        const auto* prefix = nodes[count - 1]->op == GGML_OP_CONCAT ? nodes[count - 1]->src[0] : nullptr;
+        for (int i = 0; i < count; ++i)
+            if (nodes[i] == gate || nodes[i] == residual || nodes[i] == prefix)
+                return 0;
+        Plan plan;
+        plan.device = device_info(device);
+        if (plan.device.major < 8 || plan.device.sms <= 0)
+            return 0;
+        return configure(describe(nodes, count), plan) ? count : 0;
+#endif
+    }
+    static uintptr_t address(const ggml_tensor* tensor) {
+        size_t offset = 0;
+        while (tensor && !tensor->data && tensor->view_src) {
+            offset += tensor->view_offs;
+            tensor = tensor->view_src;
+        }
+        return tensor && tensor->data ? reinterpret_cast<uintptr_t>(tensor->data) + offset : 0;
+    }
+    static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+        if (!b)
+            return false;
+        const auto x = address(a), y = address(b);
+        if (!x || !y)
+            return false;
+        return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+    }
+    static bool needs_scratch(const ggml_tensor* dst, const Operation& op) {
+        if (address(dst) % 16 != 0)
+            return true;
+        for (const auto* input : {op.mat.weight, op.mat.input, op.mat.input_scale, op.mat.weight_scale, op.gate, op.residual, op.prefix}) {
+            if (input == op.residual && !op.prefix && address(input) == address(dst))
+                continue;
+            if (overlaps(dst, input))
+                return true;
+        }
+        return false;
+    }
+    static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+        auto* plan    = new Plan;
+        *state        = plan;
+        plan->device  = device_info(ggml_backend_get_device(backend));
+        const auto op = describe(nodes, count);
+        if (!configure(op, *plan))
+            return GGML_STATUS_FAILED;
+        // Fusion dependencies keep external inputs live until this output is allocated.
+        const auto* dst         = nodes[count - 1];
+        const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src;
+        plan->output_scratch    = !fresh_output && !(dst->data && !needs_scratch(dst, op));
+        *workspace              = plan->output_offset + (plan->output_scratch ? ggml_nbytes(dst) : 0);
+        return GGML_STATUS_SUCCESS;
+    }
+    static void destroy(void* state, void*) {
+        delete static_cast<Plan*>(state);
+    }
+    static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+        auto& plan    = *static_cast<Plan*>(state);
+        const auto op = describe(nodes, count);
+        auto* dst     = nodes[count - 1];
+        auto reg      = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+        auto query    = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+        ggml_ext_cuda_context context;
+        if (!query || !query(execution->backend, &context) || context.device != plan.device.ordinal)
+            return GGML_STATUS_FAILED;
+        auto stream      = static_cast<cudaStream_t>(context.stream);
+        auto* scratch    = execution->workspace ? static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace)) : nullptr;
+        const bool alias = needs_scratch(dst, op);
+        if (alias && !plan.output_scratch)
+            return GGML_STATUS_FAILED;
+        auto* output              = alias ? scratch + plan.output_offset : static_cast<char*>(dst->data);
+        const size_t prefix_bytes = op.prefix ? ggml_nbytes(op.prefix) : 0;
+        auto status               = plan.launch(plan, op, reinterpret_cast<float*>(output + prefix_bytes), scratch, stream);
+        if (status != GGML_STATUS_SUCCESS)
+            return status;
+        if (prefix_bytes && cudaMemcpyAsync(output, op.prefix->data, prefix_bytes, cudaMemcpyDeviceToDevice, stream) != cudaSuccess)
+            return GGML_STATUS_FAILED;
+        if (alias && cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream) != cudaSuccess)
+            return GGML_STATUS_FAILED;
+        return GGML_STATUS_SUCCESS;
+    }
+}
+
+void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device) {
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.i8_gated_residual.cuda.v1";
+    desc.op                   = GGML_OP_CUSTOM;
+    desc.custom_op            = sd_ggml_i8_op();
+    desc.priority             = 325;
+    desc.optimize             = sd_i8_gated::optimize;
+    desc.match                = sd_i8_gated::match;
+    desc.prepare              = sd_i8_gated::prepare;
+    desc.execute              = sd_i8_gated::execute;
+    desc.destroy              = sd_i8_gated::destroy;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}

src/core/cuda/i8_gemm.cu

@@ -0,0 +1,405 @@
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
+#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
+#include <unordered_map>
+#include "i8_gemm.h"
+
+namespace sd_i8_cutlass {
+    struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
+        using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
+        using Base::Base;
+
+        CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
+            const auto shape = tiled_shape();
+            if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
+                // Column bands retain weight tiles in Ada's L2 across successive row tiles.
+                constexpr int band_width = 16;
+                if (index >= shape.m() * shape.n()) {
+                    return {shape.m(), shape.n(), get_batch_idx()};
+                }
+                const int band   = index / (shape.m() * band_width);
+                const int width  = min(band_width, shape.n() - band * band_width);
+                const int within = index - band * shape.m() * band_width;
+                return {within / width, band * band_width + within % width, get_batch_idx()};
+            }
+#endif
+            return Base::get_tile_offset(index);
+        }
+    };
+
+    template <int M, int N, int WM>
+    struct TileConfig {
+        using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+        using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+    };
+
+    enum Tile {
+        TILE_128_256,
+        TILE_128_128,
+        TILE_64_128,
+        TILE_16_128,
+    };
+
+    static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
+        static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
+        return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
+    }
+
+    static Tile select_tile(int64_t rows, int64_t n, int nsm) {
+        if (rows < 32) {
+            return TILE_16_128;
+        }
+        const int64_t columns = (n + 127) / 128;
+        if (((rows + 15) / 16) * columns <= nsm) {
+            return TILE_16_128;
+        }
+        if (((rows + 63) / 64) * columns <= nsm) {
+            return TILE_64_128;
+        }
+        const int64_t wide_tiles   = ((rows + 127) / 128) * ((n + 255) / 256);
+        const int64_t square_tiles = ((rows + 127) / 128) * columns;
+        // A narrower tile costs more per output element, but can waste less of the last wave.
+        const double wide_cost   = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
+        const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
+        return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
+    }
+
+    struct DeviceInfo {
+        int ordinal        = -1;
+        int major          = 0;
+        int sms            = 0;
+        int threads_per_sm = 0;
+        size_t shared      = 0;
+    };
+    static DeviceInfo device_info(ggml_backend_dev_t device) {
+        static std::mutex mutex;
+        static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
+        std::lock_guard<std::mutex> lock(mutex);
+        auto found = cache.find(device);
+        if (found != cache.end())
+            return found->second;
+        DeviceInfo info;
+        auto reg             = ggml_backend_dev_backend_reg(device);
+        auto query           = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+        int ordinal          = query ? query(device) : -1;
+        cudaDeviceProp props = {};
+        if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
+            info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+        cache.emplace(device, info);
+        return info;
+    }
+    // The iterator pairs gate/up columns while reading the original checkpoint layout.
+    struct WeightLayout : cutlass::layout::PermuteBase {
+        int stride, half_n;
+        CUTLASS_HOST_DEVICE WeightLayout(cutlass::PitchLinearCoord extent, int64_t ld)
+            : stride(static_cast<int>(ld)), half_n(extent.strided() / 2) {}
+        CUTLASS_HOST_DEVICE WeightLayout(cutlass::MatrixCoord extent, int64_t ld)
+            : stride(static_cast<int>(ld)), half_n(extent.column() / 2) {}
+        CUTLASS_HOST_DEVICE int64_t operator()(cutlass::PitchLinearCoord c) const { return c.contiguous() + int64_t(c.strided() / 2 + (c.strided() % 2) * half_n) * stride; }
+        CUTLASS_HOST_DEVICE int64_t operator()(cutlass::MatrixCoord c) const { return c.row() + int64_t(c.column() / 2 + (c.column() % 2) * half_n) * stride; }
+    };
+
+    template <class Base>
+    struct SplitWeightIterator : Base {
+        using Layout      = typename Base::Layout;
+        using Pointer     = typename Base::Pointer;
+        using TensorCoord = typename Base::TensorCoord;
+        using AccessType  = typename Base::AccessType;
+        struct Params : Base::Params {
+            const int8_t* up = nullptr;
+            int64_t k        = 0;
+            Params()         = default;
+            CUTLASS_HOST_DEVICE Params(const Layout& layout)
+                : Base::Params(layout), k(layout.stride(0)) {}
+        };
+        uintptr_t gate, up;
+        size_t gate_bytes;
+        CUTLASS_DEVICE SplitWeightIterator(const Params& params, Pointer pointer, TensorCoord extent, int thread, TensorCoord offset, const int* indices = nullptr)
+            : Base(params, pointer, extent, thread, offset, indices),
+              gate(reinterpret_cast<uintptr_t>(pointer)),
+              up(reinterpret_cast<uintptr_t>(params.up)),
+              gate_bytes(size_t(params.k) * (extent.column() / 2)) {}
+
+        CUTLASS_DEVICE AccessType* get() const {
+            auto* address = Base::get();
+            if (!address)
+                return nullptr;
+            // Keep the packed SwiGLU iterator's column pairing, but redirect its
+            // virtual up half to a separate allocation. Params.k stays full-sized
+            // even when Stream-K restricts this iterator's K extent.
+            const uintptr_t offset = reinterpret_cast<uintptr_t>(address) - gate;
+            return offset < gate_bytes ? address : reinterpret_cast<AccessType*>(up + offset - gate_bytes);
+        }
+    };
+
+    template <class Map>
+    struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
+        using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using Base::Base;
+        template <class CB>
+        struct Callbacks : CB {
+            CUTLASS_DEVICE Callbacks(CB const& base)
+                : CB(base) {}
+            CUTLASS_DEVICE void end_step(int step) {
+                auto src    = cute::filter(this->tC_rAux);
+                auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step));
+                CUTLASS_PRAGMA_UNROLL
+                for (int i = 0; i < cute::size(src); ++i) {
+                    auto coord = coords(i);
+                    if (cute::elem_less(coord, this->problem_shape)) {
+                        // Each vector contains two gate/up pairs; only their SwiGLU results are stored.
+                        const float4 a = *reinterpret_cast<const float4*>(&src(i));
+                        float2 value   = make_float2(__fmul_rn(a.x / (1.f + expf(-a.x)), a.y), __fmul_rn(a.z / (1.f + expf(-a.z)), a.w));
+                        int64_t n = cute::get<1>(this->problem_shape), m = cute::get<0>(this->problem_shape);
+                        int64_t offset                                                 = cute::get<0>(coord) * (n / 2) + cute::get<1>(coord) / 2 + cute::get<2>(coord) * m * (n / 2);
+                        *reinterpret_cast<float2*>(this->params_ptr->ptr_aux + offset) = value;
+                    }
+                }
+            }
+        };
+        template <class Shape>
+        CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) {
+            auto cb = Base::get_callbacks(tile, tid, shape);
+            return Callbacks<decltype(cb)>(cb);
+        }
+    };
+
+    static __global__ void pair_scales(const float* src, const float* up, float* dst, int n) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i < n)
+            dst[i] = up ? (i % 2 ? up[i / 2] : src[i / 2]) : src[i / 2 + (i % 2) * (n / 2)];
+    }
+
+    template <class Config, bool StreamK, bool PackedW4 = false>
+    struct DenseTypes {
+        static constexpr bool kSplit = false;
+        using Weight                 = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
+        using Tile                   = typename Config::Tile;
+        using Warp                   = typename Config::Warp;
+        using Swizzle                = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Map                    = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Col                    = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Row                    = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Mul                    = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using ScaleX                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleW                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using Store                  = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using Epilogue               = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+        using Kernel                 = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+            int8_t,
+            cutlass::layout::RowMajor,
+            cutlass::ComplexTransform::kNone,
+            16,
+            Weight,
+            cutlass::layout::ColumnMajor,
+            cutlass::ComplexTransform::kNone,
+            PackedW4 ? 32 : 16,
+            float,
+            cutlass::layout::RowMajor,
+            4,
+            int,
+            float,
+            cutlass::arch::OpClassTensorOp,
+            cutlass::arch::Sm80,
+            Tile,
+            Warp,
+            cutlass::gemm::GemmShape<16, 8, 32>,
+            Epilogue,
+            Swizzle,
+            4,
+            std::conditional_t<PackedW4, cutlass::arch::OpMultiplyAddMixedInputUpcast, cutlass::arch::OpMultiplyAddSaturate>>::GemmKernel;
+        using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+    };
+    template <class Config, bool StreamK, bool Split = false>
+    struct SwigluTypes {
+        static constexpr bool kSplit = Split;
+        using Weight                 = int8_t;
+        using Tile                   = typename Config::Tile;
+        using Warp                   = typename Config::Warp;
+        using Swizzle                = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Map                    = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Col                    = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Row                    = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Mul                    = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using ScaleX                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleW                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using Store                  = SwigluStore<Map>;
+        using Epilogue               = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+        using GemmBase               = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
+        using VisitorEpilogue        = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
+        using BaseMma                = typename GemmBase::Mma;
+        using SplitMma               = cutlass::gemm::threadblock::MmaMultistage<
+            typename BaseMma::Shape,
+            typename BaseMma::IteratorA,
+            typename BaseMma::SmemIteratorA,
+            BaseMma::kCacheOpA,
+            SplitWeightIterator<typename BaseMma::IteratorB>,
+            typename BaseMma::SmemIteratorB,
+            BaseMma::kCacheOpB,
+            typename BaseMma::ElementC,
+            typename BaseMma::LayoutC,
+            typename BaseMma::Policy,
+            BaseMma::Detail::kStages>;
+        using Mma    = std::conditional_t<Split, SplitMma, BaseMma>;
+        using Kernel = std::conditional_t<StreamK,
+                                          cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<Mma, VisitorEpilogue, Swizzle>,
+                                          cutlass::gemm::kernel::GemmWithEpilogueVisitor<Mma, VisitorEpilogue, Swizzle>>;
+        using Gemm   = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+    };
+    struct Plan {
+        DeviceInfo device;
+        int64_t rows, n, k;
+        bool swiglu;
+        bool packed_w4;
+        bool split;
+        size_t kernel_bytes = 0, scratch_offset = 0;
+        int occupancy                                                                                                                                    = 0;
+        ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*) = nullptr;
+    };
+    template <class Types>
+    static typename Types::Gemm::Arguments arguments(const Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out) {
+        typename Types::ScaleW::Arguments scales     = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
+        typename Types::Epilogue::Arguments epilogue = {scales, {out, {p.n, cute::_1{}, p.rows * p.n}}};
+        return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, reinterpret_cast<const typename Types::Weight*>(w), nullptr, nullptr, p.rows * p.k, p.n * p.k, 0, 0, p.k, p.k, 0, 0};
+    }
+    template <class Types>
+    static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const int8_t* up, const float* up_scales) {
+        using Kernel = typename Types::Kernel;
+        if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
+             reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out) | reinterpret_cast<uintptr_t>(up) | reinterpret_cast<uintptr_t>(up_scales)) %
+            16)
+            return GGML_STATUS_FAILED;
+        auto* paired = static_cast<float*>(scratch);
+        auto args    = arguments<Types>(p, x, w, xs, p.swiglu ? paired : ws, out);
+        if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+            return GGML_STATUS_FAILED;
+        if (!p.occupancy) {
+            if (sizeof(typename Kernel::SharedStorage) >= (48 << 10) &&
+                cudaFuncSetAttribute(cutlass::Kernel2<Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess)
+                return GGML_STATUS_FAILED;
+            if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&p.occupancy, cutlass::Kernel2<Kernel>, Kernel::kThreadCount,
+                                                                       sizeof(typename Kernel::SharedStorage), cudaOccupancyDisableCachingOverride) != cudaSuccess ||
+                p.occupancy <= 0)
+                return GGML_STATUS_FAILED;
+        }
+        typename Kernel::Params params(args, p.device.sms, p.occupancy);
+        if constexpr (Types::kSplit) {
+            if (!up || !up_scales)
+                return GGML_STATUS_FAILED;
+            params.params_B.up = up;
+        }
+        if (params.get_workspace_size() > p.kernel_bytes)
+            return GGML_STATUS_FAILED;
+        if (p.swiglu) {
+            pair_scales<<<unsigned((p.n + 255) / 256), 256, 0, stream>>>(ws, up_scales, paired, int(p.n));
+            if (cudaGetLastError() != cudaSuccess)
+                return GGML_STATUS_FAILED;
+        }
+        if (params.init_workspace(scratch ? static_cast<char*>(scratch) + p.scratch_offset : nullptr, stream) != cutlass::Status::kSuccess)
+            return GGML_STATUS_FAILED;
+        cutlass::arch::synclog_setup();
+        cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
+        return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+    template <class Types>
+    static bool configure_kernel(Plan& p) {
+        using Kernel = typename Types::Kernel;
+        if (sizeof(typename Kernel::SharedStorage) > p.device.shared)
+            return false;
+        auto args = arguments<Types>(p, nullptr, nullptr, nullptr, nullptr, nullptr);
+        if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+            return false;
+        size_t bytes = 0;
+        for (int occupancy = 1; occupancy <= p.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
+            typename Kernel::Params params(args, p.device.sms, occupancy);
+            const auto grid = params.get_grid_dims();
+            if (grid.y > 65535 || grid.z > 65535)
+                return false;
+            bytes = std::max(bytes, params.get_workspace_size());
+        }
+        p.kernel_bytes   = bytes;
+        p.scratch_offset = p.swiglu ? (size_t(p.n) * sizeof(float) + 255) / 256 * 256 : 0;
+        p.launch         = launch<Types>;
+        return true;
+    }
+    template <class Config>
+    static bool configure_tile(Plan& p) {
+        if (p.packed_w4)
+            return configure_kernel<DenseTypes<Config, false, true>>(p);
+        if (p.split)
+            return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true, true>>(p) : configure_kernel<SwigluTypes<Config, false, true>>(p);
+        if (p.swiglu)
+            return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true>>(p) : configure_kernel<SwigluTypes<Config, false>>(p);
+        return p.rows >= 8192 ? configure_kernel<DenseTypes<Config, true>>(p) : configure_kernel<DenseTypes<Config, false>>(p);
+    }
+    static bool configure(Plan& p) {
+        auto try_tile = [&](Tile tile) {
+            switch (tile) {
+                case TILE_128_256:
+                    return configure_tile<TileConfig<128, 256, 64>>(p);
+                case TILE_128_128:
+                    return configure_tile<TileConfig<128, 128, 64>>(p);
+                case TILE_64_128:
+                    return configure_tile<TileConfig<64, 128, 32>>(p);
+                case TILE_16_128:
+                    return configure_tile<TileConfig<16, 128, 16>>(p);
+            }
+            return false;
+        };
+        const auto tile = select_tile(p.rows, p.n, p.device.sms);
+        if (try_tile(tile))
+            return true;
+        for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
+            if (fallback != tile && try_tile(fallback))
+                return true;
+        return false;
+    }
+}
+struct SDI8CudaGemmPlan : sd_i8_cutlass::Plan {};
+static SDI8CudaGemmPlan* prepare_gemm(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4, bool split) {
+#ifdef SD_CUDA_FORCE_CUBLAS
+    return nullptr;
+#else
+    const auto info = sd_i8_cutlass::device_info(device);
+    if ((packed_w4 && (swiglu || k > INT_MAX / 1024)) || info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
+        rows > INT_MAX || n < 1 || n > INT_MAX || k < 1 || k > INT_MAX || k % 64 || n % (swiglu ? 256 : 128))
+        return nullptr;
+    auto* p      = new SDI8CudaGemmPlan;
+    p->device    = info;
+    p->rows      = rows;
+    p->n         = n;
+    p->k         = k;
+    p->swiglu    = swiglu;
+    p->packed_w4 = packed_w4;
+    p->split     = split;
+    if (!sd_i8_cutlass::configure(*p)) {
+        delete p;
+        return nullptr;
+    }
+    *workspace = p->scratch_offset + p->kernel_bytes;
+    return p;
+#endif
+}
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4) {
+    return prepare_gemm(device, rows, n, k, swiglu, workspace, packed_w4, false);
+}
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace) {
+    if (n > INT_MAX / 2 || k > INT_MAX / 16384)
+        return nullptr;
+    return prepare_gemm(device, rows, 2 * n, k, true, workspace, false, true);
+}
+ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
+    return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr);
+}
+ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream) {
+    return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales);
+}
+void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan) {
+    delete plan;
+}

src/core/cuda/i8_gemm.h

@@ -0,0 +1,12 @@
+#ifndef __SD_CORE_CUDA_I8_GEMM_H__
+#define __SD_CORE_CUDA_I8_GEMM_H__
+#include <cuda_runtime.h>
+#include "ggml-ext.h"
+struct SDI8CudaGemmPlan;
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4 = false);
+ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream);
+// n is the width of each branch; gate/up weights and scales remain separate.
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace);
+ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream);
+void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan);
+#endif

src/core/cuda/i8_projection.cu

@@ -0,0 +1,394 @@
+#include "core/cuda/i8_projection.h"
+#include "core/cuda/rope.h"
+#include "core/ggml_extensions.h"
+
+#include <cuda_fp16.h>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
+#include <type_traits>
+#include <unordered_map>
+
+static __device__ __forceinline__ float warp_sum8(float value) {
+#pragma unroll
+    for (int offset = 4; offset > 0; offset >>= 1)
+        value = __fadd_rn(value, __shfl_xor_sync(0xffffffff, value, offset, 8));
+    return value;
+}
+
+template <class Map, class Out, bool rms = true>
+struct I8RopeStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
+    using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+    struct Arguments : Base::Arguments {
+        const float* weight;
+        const float* pe;
+        float eps;
+        float scale;
+        int64_t tokens;
+        half* half_output;
+        int64_t prefix_tokens;
+    };
+    using Params        = Arguments;
+    using SharedStorage = typename Base::SharedStorage;
+    const Params* extra;
+    template <class Shape>
+    static constexpr Params to_underlying_arguments(Shape const&, Arguments const& args, void*) { return args; }
+    CUTLASS_HOST_DEVICE I8RopeStore(Params const& p, SharedStorage const& s)
+        : Base(p, s), extra(&p) {}
+    template <class CB>
+    struct Callbacks : CB {
+        const Params* extra;
+        CUTLASS_DEVICE Callbacks(CB const& base, const Params* p)
+            : CB(base), extra(p) {}
+        CUTLASS_DEVICE float sum32(float4 v) {
+            float a = warp_sum8(__fmul_rn(v.x, v.x));
+            float b = warp_sum8(__fmul_rn(v.y, v.y));
+            float c = warp_sum8(__fmul_rn(v.z, v.z));
+            float d = warp_sum8(__fmul_rn(v.w, v.w));
+            return __fadd_rn(__fadd_rn(a, c), __fadd_rn(b, d));
+        }
+        CUTLASS_DEVICE void end_step(int step) {
+            static_assert(Map::Detail::kAccessWidth == 16);
+            auto src    = cute::filter(this->tC_rAux);
+            auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step));
+            static_assert(cute::size(src) == 4);
+            CUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < cute::size(src); i += 2) {
+                float4 va = *reinterpret_cast<const float4*>(&src(i)), vb = *reinterpret_cast<const float4*>(&src(i + 1));
+                float norm = 1.0f;
+                if constexpr (rms) {
+                    // Preserve the standalone RMSNorm reduction tree across four 32-element groups.
+                    float sum = __fadd_rn(sum32(va), sum32(vb));
+                    sum       = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, 8));
+                    norm      = rsqrtf(sum / 128.f + extra->eps);
+                }
+                CUTLASS_PRAGMA_UNROLL
+                for (int j = 0; j < 2; ++j) {
+                    auto coord = coords(i + j);
+                    int row = cute::get<0>(coord), col = cute::get<1>(coord);
+                    int heads = cute::get<1>(this->problem_shape) / 128;
+                    if (cute::elem_less(coord, this->problem_shape)) {
+                        float4 value  = j ? vb : va;
+                        int local     = col % 128;
+                        int64_t token = row % extra->tokens, batch = row / extra->tokens, head = col / 128;
+                        float4 r = value;
+                        if constexpr (rms) {
+                            float4 weight = ((const float4*)extra->weight)[local / 4];
+                            value.x       = __fmul_rn(__fmul_rn(norm, value.x), weight.x);
+                            value.y       = __fmul_rn(__fmul_rn(norm, value.y), weight.y);
+                            value.z       = __fmul_rn(__fmul_rn(norm, value.z), weight.z);
+                            value.w       = __fmul_rn(__fmul_rn(norm, value.w), weight.w);
+                            float4 p0 = ((const float4*)extra->pe)[token * 64 + local / 2], p1 = ((const float4*)extra->pe)[token * 64 + local / 2 + 1];
+                            r.x = __fadd_rn(__fmul_rn(value.x, p0.x), __fmul_rn(value.y, p0.y));
+                            r.y = __fadd_rn(__fmul_rn(value.x, p0.z), __fmul_rn(value.y, p0.w));
+                            r.z = __fadd_rn(__fmul_rn(value.z, p1.x), __fmul_rn(value.w, p1.y));
+                            r.w = __fadd_rn(__fmul_rn(value.z, p1.z), __fmul_rn(value.w, p1.w));
+                        }
+                        int64_t offset = ((batch * heads + head) * (extra->tokens + extra->prefix_tokens) + token + extra->prefix_tokens) * 128 + local;
+                        if constexpr (std::is_same<Out, half>::value) {
+                            half2 a    = __floats2half2_rn(__fmul_rn(r.x, extra->scale), __fmul_rn(r.y, extra->scale));
+                            half2 b    = __floats2half2_rn(__fmul_rn(r.z, extra->scale), __fmul_rn(r.w, extra->scale));
+                            half2* dst = (half2*)(extra->half_output + offset);
+                            dst[0]     = a;
+                            dst[1]     = b;
+                        } else {
+                            *((float4*)(this->params_ptr->ptr_aux + offset)) = r;
+                        }
+                    }
+                }
+            }
+        }
+    };
+    template <class Shape>
+    CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) {
+        auto cb = Base::get_callbacks(tile, tid, shape);
+        return Callbacks<decltype(cb)>(cb, extra);
+    }
+};
+
+template <class Out, bool RMS = true>
+struct ProjectionKernel {
+    using Tile    = cutlass::gemm::GemmShape<128, 256, 64>;
+    using Warp    = cutlass::gemm::GemmShape<64, 64, 64>;
+    using Swizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
+    using Map     = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+    using Col     = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+    using Row     = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+    using Mul     = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+    using ScaleX  = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+    using ScaleW  = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+    using Store   = I8RopeStore<Map, Out, RMS>;
+    using EVT     = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+    using Kernel  = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+        int8_t,
+        cutlass::layout::RowMajor,
+        cutlass::ComplexTransform::kNone,
+        16,
+        int8_t,
+        cutlass::layout::ColumnMajor,
+        cutlass::ComplexTransform::kNone,
+        16,
+        float,
+        cutlass::layout::RowMajor,
+        4,
+        int,
+        float,
+        cutlass::arch::OpClassTensorOp,
+        cutlass::arch::Sm80,
+        Tile,
+        Warp,
+        cutlass::gemm::GemmShape<16, 8, 32>,
+        EVT,
+        Swizzle,
+        4,
+        cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
+    using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+};
+
+static float param_f32(const ggml_tensor* tensor, int index) {
+    float result;
+    std::memcpy(&result, tensor->op_params + index, sizeof(result));
+    return result;
+}
+
+static bool device_supports(ggml_backend_dev_t device) {
+#ifdef SD_CUDA_FORCE_CUBLAS
+    (void)device;
+    return false;
+#else
+    static std::mutex mutex;
+    static std::unordered_map<ggml_backend_dev_t, bool> cache;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto found = cache.find(device);
+    if (found != cache.end())
+        return found->second;
+    auto reg             = ggml_backend_dev_backend_reg(device);
+    auto query           = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+    int ordinal          = query ? query(device) : -1;
+    cudaDeviceProp props = {};
+    size_t shared        = std::max({sizeof(typename ProjectionKernel<float>::Kernel::SharedStorage),
+                                     sizeof(typename ProjectionKernel<half>::Kernel::SharedStorage), sizeof(typename ProjectionKernel<half, false>::Kernel::SharedStorage)});
+    bool supported       = ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess && props.major >= 8 && props.sharedMemPerBlockOptin >= shared;
+    cache.emplace(device, supported);
+    return supported;
+#endif
+}
+
+static bool aligned(const ggml_tensor* tensor) {
+    size_t offset = 0;
+    while (tensor->view_src) {
+        offset += tensor->view_offs;
+        tensor = tensor->view_src;
+    }
+    return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % 16 == 0;
+}
+
+static bool projection_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view) {
+    SDI8Operation op;
+    if (!sd_ggml_i8_describe(mat, &op) || !op.packed || !op.weight_scale || op.bias)
+        return false;
+    const auto* w  = op.weight;
+    const auto* x  = op.input;
+    const auto* ws = op.weight_scale;
+    if (mat->type != GGML_TYPE_F32 ||
+        w->type != GGML_TYPE_I8 || x->type != GGML_TYPE_I8 || ws->type != GGML_TYPE_F32 ||
+        w->ne[2] != 1 || w->ne[3] != 1 || mat->ne[0] != w->ne[1] || ggml_nelements(ws) != mat->ne[0] ||
+        !ggml_is_contiguous(mat) || !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || !ggml_is_contiguous(ws) ||
+        !aligned(w) || !aligned(x) || !aligned(ws) || !aligned(mat) ||
+        view->op != GGML_OP_RESHAPE || view->src[0] != mat || view->view_offs != 0 || !ggml_is_contiguous(view) ||
+        view->ne[0] != 128 || view->ne[0] * view->ne[1] != mat->ne[0] || ggml_nelements(view) != ggml_nelements(mat))
+        return false;
+    const int64_t rows = ggml_nrows(mat), n = mat->ne[0], k = w->ne[0];
+    if (rows < 4096 || rows > INT_MAX || n < 4096 || n > INT_MAX || k < 4096 || k > INT_MAX || n % 256 || k % 64 || n / 256 > 65535 ||
+        x->ne[0] != k || x->ne[1] != (rows + 3) / 4 * 4 + (rows * 4 + k - 1) / k || x->ne[2] != 1 || x->ne[3] != 1)
+        return false;
+    return device_supports(device);
+}
+
+bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
+    if (!projection_supports(device, mat, view) || norm->src[0] != view ||
+        !sd_cuda::rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
+        return false;
+    int32_t interleaved;
+    std::memcpy(&interleaved, ggml_ext_op_attributes(rope), sizeof(interleaved));
+    return interleaved != 0;
+}
+
+template <class Out, bool RMS = true>
+static cudaError_t projection_launch(cudaStream_t stream, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope, void* output, float scale, int64_t prefix_tokens) {
+    using P        = ProjectionKernel<Out, RMS>;
+    using Gemm     = typename P::Gemm;
+    const int rows = int(ggml_nrows(mat)), n = int(mat->ne[0]), k = int(mat->src[0]->ne[0]);
+    const auto* x  = static_cast<const int8_t*>(mat->src[1]->data);
+    const auto* w  = static_cast<const int8_t*>(mat->src[0]->data);
+    const auto* xs = reinterpret_cast<const float*>(x + int64_t(k) * ((int64_t(rows) + 3) / 4 * 4));
+    const auto* ws = static_cast<const float*>(mat->src[2]->data);
+    const auto* nw = RMS ? static_cast<const float*>(mul->src[1]->data) : nullptr;
+    const auto* pe = RMS ? static_cast<const float*>(rope->src[1]->data) : nullptr;
+    // The visitor forms FP32 addresses even when its only writes target the FP16 output.
+    float* address_space = std::is_same<Out, half>::value ? static_cast<float*>(mat->data) : static_cast<float*>(output);
+    if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
+         reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(nw) | reinterpret_cast<uintptr_t>(pe) |
+         reinterpret_cast<uintptr_t>(address_space)) %
+            16 ||
+        reinterpret_cast<uintptr_t>(output) % 8)
+        return cudaErrorInvalidValue;
+    typename P::EVT::Arguments evt = {{{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}},
+                                      {{address_space, {n, cute::_1{}, int64_t(rows) * n}}, nw, pe, RMS ? param_f32(norm, 0) : 0.f, scale, view->ne[2], static_cast<half*>(output), prefix_tokens}};
+    typename Gemm::Arguments args(cutlass::gemm::GemmUniversalMode::kGemm, {rows, n, k}, 1, evt, x, w, nullptr, nullptr,
+                                  int64_t(rows) * k, int64_t(n) * k, 0, 0, k, k, 0, 0);
+    Gemm gemm;
+    if (gemm.can_implement(args) != cutlass::Status::kSuccess)
+        return cudaErrorInvalidValue;
+    if (gemm.initialize(args, nullptr, stream) != cutlass::Status::kSuccess || gemm(stream) != cutlass::Status::kSuccess)
+        return cudaErrorUnknown;
+    return cudaGetLastError();
+}
+
+cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale) {
+    return projection_launch<half>(stream, nodes[0], nodes[1], nodes[2], nodes[3], nodes[4], output, scale, 0);
+}
+
+static const ggml_tensor* node_at(const ggml_cgraph* graph, int index) {
+    return ggml_graph_node(const_cast<ggml_cgraph*>(graph), index);
+}
+
+static int concat_match(const ggml_cgraph* graph, int first, bool key) {
+    const int offset = key ? 5 : 2, count = key ? 7 : 5;
+    if (first + count > ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)))
+        return 0;
+    const auto* concat = node_at(graph, first + offset);
+    const auto* copy   = node_at(graph, first + count - 1);
+    if (concat->op != GGML_OP_CONCAT || concat->op_params[0] != (key ? 1 : 2) || concat->src[1] != node_at(graph, first + offset - 1) ||
+        copy->op != GGML_OP_CPY || copy->src[0] != node_at(graph, first + count - 2) || copy->type != GGML_TYPE_F16 ||
+        !ggml_is_contiguous(copy) || !ggml_are_same_shape(copy, copy->src[0]))
+        return 0;
+    for (int i = 0; i < 2; ++i)
+        if (!concat->src[i] || concat->src[i]->type != GGML_TYPE_F32 || !ggml_is_contiguous(concat->src[i]))
+            return 0;
+    const auto* view   = node_at(graph, first + 1);
+    const auto* prefix = concat->src[0];
+    if (prefix->ne[key ? 1 : 2] > INT_MAX - view->ne[2] || (ggml_nelements(prefix) + 255) / 256 > INT_MAX)
+        return 0;
+    if (!key) {
+        const auto* permute = node_at(graph, first + 3);
+        if (permute->op != GGML_OP_PERMUTE || permute->src[0] != concat || permute->op_params[0] != 0 || permute->op_params[1] != 2 ||
+            permute->op_params[2] != 1 || permute->op_params[3] != 3)
+            return 0;
+    }
+    return count;
+}
+
+static int projection_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+    static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION");return value&&std::atoi(value)!=0; }();
+    if (disabled || first + 4 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)))
+        return 0;
+    const auto* mat  = node_at(graph, first);
+    const auto* view = node_at(graph, first + 1);
+    const auto* norm = node_at(graph, first + 2);
+    const bool key   = norm->op == GGML_OP_RMS_NORM;
+    if (key) {
+        if (!sd_cuda_i8_rope_supports(device, mat, view, norm, node_at(graph, first + 3), node_at(graph, first + 4)))
+            return 0;
+        if (int count = concat_match(graph, first, true))
+            return count;
+        if (first + 6 < ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph))) {
+            const auto* rope  = node_at(graph, first + 4);
+            const auto* scale = node_at(graph, first + 5);
+            const auto* copy  = node_at(graph, first + 6);
+            if (scale->op == GGML_OP_SCALE && scale->src[0] == rope && param_f32(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale &&
+                copy->type == GGML_TYPE_F16 && ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+                return 7;
+        }
+        return 5;
+    }
+    return projection_supports(device, mat, view) ? concat_match(graph, first, false) : 0;
+}
+
+template <bool Aligned>
+static __global__ void prefix_copy(const float* input, half2* output, int64_t count4, int64_t prefix, int64_t tokens, int64_t heads, bool permuted) {
+    const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (i >= count4)
+        return;
+    const int64_t col = i % 32, row = i / 32;
+    const int64_t token = row % prefix, head = (row / prefix) % heads, batch = row / (prefix * heads);
+    const int64_t source_row = permuted ? (batch * prefix + token) * heads + head : (batch * heads + head) * prefix + token;
+    const float* source      = input + source_row * 128 + col * 4;
+    float4 value;
+    if constexpr (Aligned)
+        value = *reinterpret_cast<const float4*>(source);
+    else
+        value = make_float4(source[0], source[1], source[2], source[3]);
+    const int64_t target   = ((batch * heads + head) * (prefix + tokens) + token) * 32 + col;
+    output[2 * target]     = __floats2half2_rn(value.x, value.y);
+    output[2 * target + 1] = __floats2half2_rn(value.z, value.w);
+}
+
+static ggml_status projection_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+    *workspace = ggml_nbytes(nodes[count - 1]);
+    return GGML_STATUS_SUCCESS;
+}
+static void projection_destroy(void*, void*) {}
+
+static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+    const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+    return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+}
+
+static ggml_status projection_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+    const auto* mat    = nodes[0];
+    const auto* view   = nodes[1];
+    const bool key     = nodes[2]->op == GGML_OP_RMS_NORM;
+    const auto* concat = nodes[key ? count == 7 ? 5 : 4 : 2];
+    const bool joined  = concat->op == GGML_OP_CONCAT;
+    auto* dst          = nodes[count - 1];
+    auto reg           = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+    auto query         = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+    ggml_ext_cuda_context context;
+    if (!query || !query(execution->backend, &context))
+        return GGML_STATUS_FAILED;
+    auto stream = static_cast<cudaStream_t>(context.stream);
+    bool alias  = reinterpret_cast<uintptr_t>(dst->data) % (dst->type == GGML_TYPE_F32 ? 16 : 8) != 0;
+    for (const auto* input : {mat->src[0], mat->src[1], mat->src[2], key ? nodes[3]->src[1] : nullptr, key ? nodes[4]->src[1] : nullptr, joined ? concat->src[0] : nullptr})
+        if (input)
+            alias |= overlaps(dst, input);
+    void* output   = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+    int64_t prefix = joined ? concat->src[0]->ne[key ? 1 : 2] : 0;
+    float scale    = key && count == 7 && !joined ? param_f32(nodes[5], 0) : 1.f;
+    cudaError_t error;
+    if (dst->type == GGML_TYPE_F32)
+        error = projection_launch<float>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, 0);
+    else if (key)
+        error = projection_launch<half>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, prefix);
+    else
+        error = projection_launch<half, false>(stream, mat, view, nullptr, nullptr, nullptr, output, scale, prefix);
+    if (error == cudaSuccess && joined && prefix > 0) {
+        const auto* input    = concat->src[0];
+        const int64_t count4 = ggml_nelements(input) / 4;
+        auto kernel          = reinterpret_cast<uintptr_t>(input->data) % 16 == 0 ? prefix_copy<true> : prefix_copy<false>;
+        kernel<<<static_cast<unsigned int>((count4 + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(input->data), static_cast<half2*>(output), count4, prefix, view->ne[2], view->ne[1], !key);
+        error = cudaGetLastError();
+    }
+    if (error == cudaSuccess && alias)
+        error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
+    return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+}
+
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device) {
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.i8_projection.cuda.v1";
+    desc.op                   = GGML_OP_CUSTOM;
+    desc.custom_op            = sd_ggml_i8_op();
+    desc.priority             = 350;
+    desc.match                = projection_match;
+    desc.prepare              = projection_prepare;
+    desc.execute              = projection_execute;
+    desc.destroy              = projection_destroy;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}

src/core/cuda/i8_projection.h

@@ -0,0 +1,11 @@
+#ifndef __SD_CORE_CUDA_I8_PROJECTION_H__
+#define __SD_CORE_CUDA_I8_PROJECTION_H__
+
+#include <cuda_runtime.h>
+#include "ggml-ext.h"
+
+bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
+cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale);
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
+
+#endif

src/core/cuda/int8.cu

@@ -0,0 +1,928 @@
+#include <cstdlib>
+#include "convrot_kernel.h"
+#include "core/ggml_extensions.h"
+#ifdef SD_CUDA_I8_CUTLASS
+#include "core/cuda/i8_gemm.h"
+#else
+struct SDI8CudaGemmPlan;
+#endif
+
+#include "runtime.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hipblas/hipblas.h>
+#elif !defined(SD_CUDA_VENDOR_MUSA)
+#include <cublas_v2.h>
+#endif
+#include <algorithm>
+#include <climits>
+#include <mutex>
+#include <unordered_map>
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_i8_cuda {
+        template <int Width = 32>
+        static __device__ __forceinline__ float warp_reduce_max(float value) {
+#pragma unroll
+            for (int offset = Width / 2; offset > 0; offset >>= 1)
+                value = fmaxf(value, cuda_shuffle_xor(value, offset, Width));
+            return value;
+        }
+        static __global__ void quantize_rowwise_i8_cuda(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows) {
+            const int64_t row = blockIdx.x;
+            const int tid     = threadIdx.x;
+            if (row >= rows) {
+                return;
+            }
+
+            __shared__ float maxima[256];
+            const float* src_row = src + row * k;
+            int8_t* dst_row      = dst + row * k;
+            float local_max      = 0.0f;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                local_max = fmaxf(local_max, fabsf(src_row[i]));
+            }
+            maxima[tid] = local_max;
+            __syncthreads();
+
+            for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+                if (tid < stride) {
+                    maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
+                }
+                __syncthreads();
+            }
+
+            const float scale = maxima[0] / 127.0f;
+            if (tid == 0) {
+                scales[row] = scale;
+            }
+            if (scale == 0.0f) {
+                for (int64_t i = tid; i < k; i += blockDim.x) {
+                    dst_row[i] = 0;
+                }
+                return;
+            }
+
+            const float inv_scale = 1.0f / scale;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                int value  = __float2int_rn(src_row[i] * inv_scale);
+                value      = max(-127, min(127, value));
+                dst_row[i] = (int8_t)value;
+            }
+        }
+        static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+            a = index == 3 ? -a : a;
+            b = index == 2 ? -b : b;
+            c = index == 1 ? -c : c;
+            d = index == 0 ? -d : d;
+            return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+        }
+        static __global__ void quantize_convrot(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows, int64_t dst_stride) {
+            constexpr int group_size = 256;
+            const int64_t row        = blockIdx.x;
+            const int tid            = threadIdx.x;
+            if (row >= rows)
+                return;
+            extern __shared__ float values[];
+            __shared__ float maxima[32];
+            const float* src_row = src + row * k;
+            int8_t* dst_row      = dst + row * dst_stride;
+            {
+                const int lane = tid % 32;
+                for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
+                    float v[8];
+#pragma unroll
+                    for (int j = 0; j < 8; ++j) {
+                        const int64_t i   = group * group_size + j * 32 + lane;
+                        const float value = src_row[i];
+                        v[j]              = value * (1.0f / 16.0f);
+                    }
+                    // One logical warp owns H256, keeping every butterfly in registers.
+#pragma unroll
+                    for (int stride = 1; stride <= 4; stride *= 4) {
+                        const int base = lane & ~(3 * stride);
+#pragma unroll
+                        for (int j = 0; j < 8; ++j) {
+                            const float a = cuda_shuffle(v[j], base, 32);
+                            const float b = cuda_shuffle(v[j], base + stride, 32);
+                            const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
+                            const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
+                            v[j]          = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
+                        }
+                    }
+#pragma unroll
+                    for (int j = 0; j < 8; j += 2) {
+                        const float a = cuda_shuffle(v[j], lane % 16, 32);
+                        const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
+                        const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
+                        const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
+                        v[j]          = convrot_h4_cuda(a, b, c, d, lane / 16);
+                        v[j + 1]      = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
+                    }
+#pragma unroll
+                    for (int j = 0; j < 2; ++j) {
+                        const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
+#pragma unroll
+                        for (int t = 0; t < 4; ++t) {
+                            values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
+                        }
+                    }
+                }
+                __syncthreads();
+            }
+            float local_max = 0.0f;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                local_max = fmaxf(local_max, fabsf(values[i]));
+            }
+            local_max      = warp_reduce_max<32>(local_max);
+            const int lane = tid & 31;
+            const int warp = tid >> 5;
+            if (lane == 0) {
+                maxima[warp] = local_max;
+            }
+            __syncthreads();
+
+            if (warp == 0) {
+                local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
+                local_max = warp_reduce_max<32>(local_max);
+                if (lane == 0) {
+                    maxima[0] = local_max / 127.0f;
+                    if (scales != nullptr) {
+                        scales[row] = maxima[0];
+                    } else {
+                        *((float*)(dst_row + k)) = maxima[0];
+                    }
+                }
+            }
+            __syncthreads();
+
+            const float row_scale = maxima[0];
+            const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                int value  = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
+                value      = max(-127, min(127, value));
+                dst_row[i] = (int8_t)value;
+            }
+        }
+
+        static __global__ void convrot_group_amax_i8_cuda(
+            const float* src,
+            float* partial_max,
+            int64_t k,
+            int64_t groups,
+            int64_t total_groups) {
+            const int64_t index = blockIdx.x;
+            const int tid       = threadIdx.x;
+            if (index >= total_groups) {
+                return;
+            }
+
+            __shared__ float values[256];
+            __shared__ float maxima[2];
+            const int64_t row   = index / groups;
+            const int64_t group = index - row * groups;
+            const float* input  = src + row * k + group * 256;
+            values[tid]         = input[tid] * (1.0f / 16.0f);
+            values[tid + 64]    = input[tid + 64] * (1.0f / 16.0f);
+            values[tid + 128]   = input[tid + 128] * (1.0f / 16.0f);
+            values[tid + 192]   = input[tid + 192] * (1.0f / 16.0f);
+            __syncthreads();
+
+#pragma unroll
+            for (int stride = 1; stride < 256; stride *= 4) {
+                const int base = (tid / stride) * 4 * stride + tid % stride;
+                const int i0   = base;
+                const int i1   = i0 + stride;
+                const int i2   = i1 + stride;
+                const int i3   = i2 + stride;
+                const float a  = values[i0];
+                const float b  = values[i1];
+                const float c  = values[i2];
+                const float d  = values[i3];
+                values[i0]     = a + b + c - d;
+                values[i1]     = a + b - c + d;
+                values[i2]     = a - b + c + d;
+                values[i3]     = -a + b + c + d;
+                __syncthreads();
+            }
+
+            float value = fmaxf(fabsf(values[tid]), fabsf(values[tid + 64]));
+            value       = fmaxf(value, fabsf(values[tid + 128]));
+            value       = fmaxf(value, fabsf(values[tid + 192]));
+            value       = warp_reduce_max<32>(value);
+            if ((tid & 31) == 0) {
+                maxima[tid >> 5] = value;
+            }
+            __syncthreads();
+            if (tid == 0) {
+                partial_max[index] = fmaxf(maxima[0], maxima[1]);
+            }
+        }
+
+        static __global__ void reduce_convrot_row_amax_i8_cuda(
+            const float* partial_max,
+            float* scales,
+            int64_t groups,
+            int64_t rows) {
+            const int64_t row = blockIdx.x;
+            const int tid     = threadIdx.x;
+            if (row >= rows) {
+                return;
+            }
+
+            __shared__ float maxima[256];
+            float value = 0.0f;
+            for (int64_t group = tid; group < groups; group += blockDim.x) {
+                value = fmaxf(value, partial_max[row * groups + group]);
+            }
+            maxima[tid] = value;
+            __syncthreads();
+            for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+                if (tid < stride) {
+                    maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
+                }
+                __syncthreads();
+            }
+            if (tid == 0) {
+                scales[row] = maxima[0] / 127.0f;
+            }
+        }
+
+        static __global__ void convrot_group_quantize_i8_cuda(
+            const float* src,
+            int8_t* dst,
+            const float* scales,
+            int64_t k,
+            int64_t groups,
+            int64_t total_groups,
+            int64_t dst_stride) {
+            const int64_t index = blockIdx.x;
+            const int tid       = threadIdx.x;
+            if (index >= total_groups) {
+                return;
+            }
+
+            __shared__ float values[256];
+            const int64_t row   = index / groups;
+            const int64_t group = index - row * groups;
+            const float* input  = src + row * k + group * 256;
+            int8_t* output      = dst + row * dst_stride + group * 256;
+            values[tid]         = input[tid] * (1.0f / 16.0f);
+            values[tid + 64]    = input[tid + 64] * (1.0f / 16.0f);
+            values[tid + 128]   = input[tid + 128] * (1.0f / 16.0f);
+            values[tid + 192]   = input[tid + 192] * (1.0f / 16.0f);
+            __syncthreads();
+
+#pragma unroll
+            for (int stride = 1; stride < 256; stride *= 4) {
+                const int base = (tid / stride) * 4 * stride + tid % stride;
+                const int i0   = base;
+                const int i1   = i0 + stride;
+                const int i2   = i1 + stride;
+                const int i3   = i2 + stride;
+                const float a  = values[i0];
+                const float b  = values[i1];
+                const float c  = values[i2];
+                const float d  = values[i3];
+                values[i0]     = a + b + c - d;
+                values[i1]     = a + b - c + d;
+                values[i2]     = a - b + c + d;
+                values[i3]     = -a + b + c + d;
+                __syncthreads();
+            }
+
+            const float row_scale = scales[row];
+            const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+            const int indices[4]  = {tid, tid + 64, tid + 128, tid + 192};
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                int value          = row_scale == 0.0f ? 0 : __float2int_rn(values[indices[i]] * inv_scale);
+                value              = max(-127, min(127, value));
+                output[indices[i]] = (int8_t)value;
+            }
+            if (group == 0 && tid == 0 && dst_stride != k) {
+                *((float*)(dst + row * dst_stride + k)) = row_scale;
+            }
+        }
+
+        static __global__ void dequantize_i32_rows_cuda(
+            const int32_t* src,
+            const float* scales,
+            const float* weight_scales,
+            const float* bias,
+            float* dst,
+            int64_t n,
+            int64_t rows) {
+            const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+            const int64_t row    = blockIdx.y + int64_t(blockIdx.z) * 65535;
+            if (output < n && row < rows) {
+                const int64_t index          = row * n + output;
+                const float activation_scale = scales[row];
+                float value                  = (float)src[index] * activation_scale;
+                if (weight_scales != nullptr) {
+                    value *= weight_scales[output];
+                }
+                if (bias != nullptr) {
+                    value += bias[output];
+                }
+                dst[index] = value;
+            }
+        }
+
+        static __global__ void dequantize_i32_rows_vec4_cuda(
+            const int4* src,
+            const float* scales,
+            const float4* weight_scales,
+            const float4* bias,
+            float4* dst,
+            int64_t n4,
+            int64_t rows) {
+            const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+            const int64_t row    = blockIdx.y + int64_t(blockIdx.z) * 65535;
+            if (output >= n4 || row >= rows) {
+                return;
+            }
+            const int64_t index = row * n4 + output;
+            const int4 v        = src[index];
+            const float scale   = scales[row];
+            float4 value        = make_float4(v.x * scale, v.y * scale, v.z * scale, v.w * scale);
+            if (weight_scales != nullptr) {
+                const float4 w = weight_scales[output];
+                value.x *= w.x;
+                value.y *= w.y;
+                value.z *= w.z;
+                value.w *= w.w;
+            }
+            if (bias != nullptr) {
+                const float4 b = bias[output];
+                value.x += b.x;
+                value.y += b.y;
+                value.z += b.z;
+                value.w += b.w;
+            }
+            dst[index] = value;
+        }
+
+        static bool aligned(const ggml_tensor* tensor, size_t alignment = 4) {
+            if (!tensor)
+                return true;
+            size_t offset = 0;
+            while (tensor->view_src) {
+                offset += tensor->view_offs;
+                tensor = tensor->view_src;
+            }
+            return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
+        }
+        static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+            SDI8Operation op;
+            if (!sd_ggml_i8_describe(node, &op) || !aligned(node))
+                return false;
+            const int64_t k = op.weight->ne[0], n = op.weight->ne[1], rows = ggml_nrows(node);
+            if (k > INT_MAX || n > INT_MAX || rows > INT_MAX - 3 || k % 4 || n % 4 || (op.input->type == GGML_TYPE_F32 && op.group_size && (op.group_size != 256 || rows > INT_MAX / (k / 256))))
+                return false;
+            for (auto* input : node->src)
+                if (!aligned(input))
+                    return false;
+            const auto info = cuda_device_info(device);
+            return info.int8_blas &&
+                   ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1");
+        }
+        template <bool Quantize>
+        static __global__ void convrot_h64(const float* src, int8_t* dst, float* partial, const float* scales, int64_t k) {
+            const int64_t group = blockIdx.x, groups = k / 64, row = group / groups;
+            const int tid = threadIdx.x;
+            __shared__ float values[64];
+            __shared__ float maxima[16];
+#pragma unroll
+            for (int j = 0; j < 4; ++j)
+                values[tid + j * 16] = src[group * 64 + tid + j * 16] * (1.f / 8.f);
+            __syncthreads();
+#pragma unroll
+            for (int stride = 1; stride < 64; stride *= 4) {
+                const int i   = (tid / stride) * 4 * stride + tid % stride;
+                const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride];
+                values[i]              = convrot_h4_cuda(a, b, c, d, 0);
+                values[i + stride]     = convrot_h4_cuda(a, b, c, d, 1);
+                values[i + 2 * stride] = convrot_h4_cuda(a, b, c, d, 2);
+                values[i + 3 * stride] = convrot_h4_cuda(a, b, c, d, 3);
+                __syncthreads();
+            }
+            if constexpr (Quantize) {
+                const float inverse = scales[row] == 0.f ? 0.f : 1.f / scales[row];
+#pragma unroll
+                for (int j = 0; j < 4; ++j) {
+                    const int value                = __float2int_rn(values[tid + j * 16] * inverse);
+                    dst[group * 64 + tid + j * 16] = int8_t(max(-127, min(127, value)));
+                }
+            } else {
+                float amax = 0.f;
+#pragma unroll
+                for (int j = 0; j < 4; ++j)
+                    amax = fmaxf(amax, fabsf(values[tid + j * 16]));
+                maxima[tid] = amax;
+                __syncthreads();
+                for (int stride = 8; stride; stride >>= 1) {
+                    if (tid < stride)
+                        maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
+                    __syncthreads();
+                }
+                if (!tid)
+                    partial[group] = maxima[0];
+            }
+        }
+        static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+            int group;
+            if (!sd_ggml_convrot_describe(node, &group) || (group != 1 && group != 64 && group != 256) || !aligned(node) || !aligned(node->src[0]))
+                return false;
+            const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]);
+            return k <= INT_MAX && rows <= INT_MAX && (group == 1 || rows <= INT_MAX / (k / group)) && cuda_device_info(device).device >= 0;
+        }
+        enum Kind { MAT,
+                    QUANTIZE,
+                    SWIGLU,
+                    SWIGLU_QUANTIZE,
+                    CONCAT };
+        static bool fusion_disabled() {
+            static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+            return disabled;
+        }
+#ifdef SD_CUDA_I8_CUTLASS
+        static bool cutlass_compatible(const SDI8Operation& op, const ggml_tensor* dst) {
+            return op.weight_scale && !op.bias && op.weight->ne[0] <= INT_MAX / 16384 &&
+                   aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
+                   (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16)));
+        }
+#endif
+        static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (fusion_disabled() || first + 1 >= ggml_graph_n_nodes(g))
+                return 0;
+            const auto* mat = ggml_graph_node(g, first);
+            SDI8Operation op;
+            if (!supports(device, mat, nullptr) || !sd_ggml_i8_describe(mat, &op) || !op.weight_scale || op.bias)
+                return 0;
+            const auto* next = ggml_graph_node(g, first + 1);
+            if (next->op == GGML_OP_CONCAT && next->src[1] == mat && next->src[0] != mat && next->op_params[0] == 1 &&
+                next->type == GGML_TYPE_F32 && mat->ne[2] == 1 && mat->ne[3] == 1 &&
+                next->ne[2] == 1 && next->ne[3] == 1 && ggml_is_contiguous(next) && ggml_is_contiguous(next->src[0]))
+                return 2;
+            if (next->op != GGML_OP_GLU || ggml_get_glu_op(next) != GGML_GLU_OP_SWIGLU || next->op_params[1] != 0 ||
+                next->src[0] != mat || next->src[1] || next->type != GGML_TYPE_F32 || !ggml_is_contiguous(next) ||
+                ggml_nelements(next) > int64_t(INT_MAX) * 256)
+                return 0;
+            if (first + 2 < ggml_graph_n_nodes(g)) {
+                const auto* quant = ggml_graph_node(g, first + 2);
+                int group;
+                if (quant->src[0] == next && sd_ggml_convrot_describe(quant, &group) && group == 256 && convrot_supports(device, quant, nullptr) &&
+                    size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory) {
+#ifdef SD_CUDA_I8_CUTLASS
+                    if (!next->data && !next->buffer && !next->view_src && cutlass_compatible(op, next)) {
+                        size_t workspace          = 0;
+                        auto* gemm                = sd_cuda_i8_gemm_prepare(device, ggml_nrows(mat), mat->ne[0], op.weight->ne[0], true, &workspace);
+                        const bool fused_epilogue = gemm != nullptr;
+                        sd_cuda_i8_gemm_free(gemm);
+                        // ConvRot is a separate launch either way. Keep its input in the
+                        // tensor arena so other operations can reuse this large allocation.
+                        if (fused_epilogue)
+                            return 2;
+                    }
+#endif
+                    return 3;
+                }
+            }
+            return 2;
+        }
+        struct Plan {
+            ggml_ext_gpu_device_info device;
+            Kind kind;
+            int group;
+            int64_t k, n, rows, padded;
+            size_t quant, scales, accum, partial, output, gemm_offset;
+            SDI8CudaGemmPlan* gemm = nullptr;
+            bool fast_swiglu       = false;
+        };
+        static size_t reserve(size_t& total, size_t bytes) {
+            const size_t offset = total;
+            total += (bytes + 255) / 256 * 256;
+            return offset;
+        }
+        static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+            const auto* dst = nodes[count - 1];
+            SDI8Operation op;
+            const bool mat = sd_ggml_i8_describe(nodes[0], &op);
+            auto* p        = new Plan{};
+            *state         = p;
+            auto device    = ggml_backend_get_device(backend);
+            p->device      = cuda_device_info(device);
+            p->kind        = !mat ? QUANTIZE : count == 3              ? SWIGLU_QUANTIZE
+                                           : count == 1                ? MAT
+                                           : dst->op == GGML_OP_CONCAT ? CONCAT
+                                                                       : SWIGLU;
+            p->k           = mat ? op.weight->ne[0] : dst->ne[0];
+            p->n           = mat ? nodes[0]->ne[0] : 0;
+            p->rows        = mat ? ggml_nrows(nodes[0]) : ggml_nrows(dst->src[0]);
+            p->padded      = (p->rows + 3) / 4 * 4;
+            p->group       = mat ? op.group_size : 0;
+            if (!mat)
+                sd_ggml_convrot_describe(dst, &p->group);
+            size_t total = 0;
+            if (mat) {
+                p->quant  = reserve(total, op.input->type == GGML_TYPE_F32 || (!op.packed && p->rows != p->padded) ? size_t(p->k) * p->padded : 0);
+                p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0);
+#ifdef SD_CUDA_I8_CUTLASS
+                size_t gemm_bytes = 0;
+                // Retain cuBLAS when I32 accumulation can overflow.
+                if (cutlass_compatible(op, dst) &&
+                    (p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) {
+                    if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) {
+                        p->gemm        = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes);
+                        p->fast_swiglu = p->gemm != nullptr;
+                    }
+                    if (!p->gemm)
+                        p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, false, &gemm_bytes);
+                }
+                p->gemm_offset = reserve(total, gemm_bytes);
+#endif
+                const bool intermediate = !p->gemm || p->kind == SWIGLU_QUANTIZE || (p->kind == SWIGLU && !p->fast_swiglu);
+                p->accum                = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0);
+            }
+            const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory;
+            p->partial                = reserve(total, p->group == 64 ? size_t(p->rows) * (p->k / 64) * sizeof(float) : large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float)
+                                                                                                                                       : 0);
+            // Custom/concat outputs and packed GLU cannot reuse their input storage.
+            const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src &&
+                                      (dst->op == GGML_OP_CUSTOM || dst->op == GGML_OP_CONCAT ||
+                                       (dst->op == GGML_OP_GLU && !dst->src[1]));
+            p->output               = fresh_output || (p->gemm && p->kind == SWIGLU_QUANTIZE) ? SIZE_MAX : reserve(total, ggml_nbytes(dst));
+            *workspace              = total;
+            return GGML_STATUS_SUCCESS;
+        }
+        static void destroy(void* state, void*) {
+            auto* p = static_cast<Plan*>(state);
+#ifdef SD_CUDA_I8_CUTLASS
+            if (p)
+                sd_cuda_i8_gemm_free(p->gemm);
+#endif
+            delete p;
+        }
+        static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+            if (!b)
+                return false;
+            const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+            return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+        }
+#ifdef SD_CUDA_I8_CUTLASS
+        static bool split_describe(ggml_backend_dev_t device, ggml_tensor* const* nodes, SDI8Operation& gate, SDI8Operation& up) {
+            const auto* glu = nodes[2];
+            if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || glu->op_params[1] != 0 ||
+                !glu->src[0] || !glu->src[1] || glu->src[0] == glu->src[1] || glu->view_src ||
+                glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu) || !aligned(glu, 16) ||
+                !((nodes[0] == glu->src[0] && nodes[1] == glu->src[1]) || (nodes[1] == glu->src[0] && nodes[0] == glu->src[1])))
+                return false;
+            for (const auto* mat : {glu->src[0], glu->src[1]})
+                if (mat->view_src || !supports(device, mat, nullptr) || !ggml_are_same_shape(mat, glu))
+                    return false;
+            if (!sd_ggml_i8_describe(glu->src[0], &gate) || !sd_ggml_i8_describe(glu->src[1], &up) ||
+                gate.input != up.input || gate.input_scale != up.input_scale || gate.packed != up.packed ||
+                gate.input->type != GGML_TYPE_I8 || gate.bias || up.bias || !gate.weight_scale || !up.weight_scale ||
+                !ggml_are_same_shape(gate.weight, up.weight))
+                return false;
+            for (const auto* t : {gate.input, gate.input_scale, gate.weight, up.weight, gate.weight_scale, up.weight_scale})
+                if (!aligned(t, 16))
+                    return false;
+            return true;
+        }
+        static int split_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (fusion_disabled() || first + 2 >= ggml_graph_n_nodes(g))
+                return 0;
+            ggml_tensor* nodes[] = {ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)};
+            SDI8Operation gate, up;
+            if (!split_describe(device, nodes, gate, up))
+                return 0;
+            size_t workspace     = 0;
+            auto* gemm           = sd_cuda_i8_gemm_prepare_split(device, ggml_nrows(nodes[2]), nodes[2]->ne[0], gate.weight->ne[0], &workspace);
+            const bool supported = gemm != nullptr;
+            sd_cuda_i8_gemm_free(gemm);
+            return supported ? 3 : 0;
+        }
+        struct SplitPlan {
+            SDI8CudaGemmPlan* gemm = nullptr;
+            size_t output          = SIZE_MAX;
+            int device;
+        };
+        static ggml_status split_prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+            auto* p     = new SplitPlan;
+            *state      = p;
+            auto device = ggml_backend_get_device(backend);
+            p->device   = cuda_device_info(device).device;
+            SDI8Operation gate, up;
+            if (count != 3 || !split_describe(device, nodes, gate, up))
+                return GGML_STATUS_FAILED;
+            auto* dst    = nodes[2];
+            size_t bytes = 0;
+            p->gemm      = sd_cuda_i8_gemm_prepare_split(device, ggml_nrows(dst), dst->ne[0], gate.weight->ne[0], &bytes);
+            if (!p->gemm)
+                return GGML_STATUS_FAILED;
+            size_t total = (bytes + 255) / 256 * 256;
+            if (dst->data || dst->buffer || dst->view_src)
+                p->output = reserve(total, ggml_nbytes(dst));
+            *workspace = total;
+            return GGML_STATUS_SUCCESS;
+        }
+        static void split_destroy(void* state, void*) {
+            auto* p = static_cast<SplitPlan*>(state);
+            if (p)
+                sd_cuda_i8_gemm_free(p->gemm);
+            delete p;
+        }
+        static ggml_status split_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void* state, void*) {
+            auto& p   = *static_cast<SplitPlan*>(state);
+            auto* dst = nodes[2];
+            SDI8Operation gate, up;
+            if (!sd_ggml_i8_describe(dst->src[0], &gate) || !sd_ggml_i8_describe(dst->src[1], &up))
+                return GGML_STATUS_FAILED;
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context) || context.device != p.device)
+                return GGML_STATUS_FAILED;
+            auto stream   = static_cast<CUDAStream>(context.stream);
+            auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+            bool alias    = false;
+            for (const auto* t : {gate.input, gate.input_scale, gate.weight, up.weight, gate.weight_scale, up.weight_scale})
+                alias |= overlaps(dst, t);
+            if (alias && p.output == SIZE_MAX)
+                return GGML_STATUS_FAILED;
+            auto* output         = alias ? reinterpret_cast<float*>(scratch + p.output) : static_cast<float*>(dst->data);
+            const auto* input    = static_cast<const int8_t*>(gate.input->data);
+            const int64_t padded = (ggml_nrows(dst) + 3) / 4 * 4;
+            const auto* scales   = gate.packed ? reinterpret_cast<const float*>(input + gate.weight->ne[0] * padded) : static_cast<const float*>(gate.input_scale->data);
+            auto status          = sd_cuda_i8_gemm_execute_split(p.gemm, input, static_cast<const int8_t*>(gate.weight->data), static_cast<const int8_t*>(up.weight->data),
+                                                                 scales, static_cast<const float*>(gate.weight_scale->data), static_cast<const float*>(up.weight_scale->data),
+                                                                 output, scratch, stream);
+            if (status != GGML_STATUS_SUCCESS || (alias && cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) != cuda_success))
+                return GGML_STATUS_FAILED;
+            return GGML_STATUS_SUCCESS;
+        }
+#endif
+        static CUDAError convrot_launch(const Plan& p, CUDAStream stream, const float* src, int8_t* dst, float* scales, float* partial) {
+            if (p.padded > p.rows) {
+                auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                if (error != cuda_success)
+                    return error;
+            }
+            const size_t shared = size_t(p.k) * sizeof(float);
+            if (p.group == 1) {
+                quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(src, dst, scales, p.k, p.rows);
+                return cuda_last_error();
+            }
+            if (p.group == 64) {
+                const auto groups = p.k / 64;
+                convrot_h64<false><<<unsigned(p.rows * groups), 16, 0, stream>>>(src, dst, partial, scales, p.k);
+                reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
+                convrot_h64<true><<<unsigned(p.rows * groups), 16, 0, stream>>>(src, dst, partial, scales, p.k);
+                return cuda_last_error();
+            }
+            if (shared + 128 <= p.device.shared_memory) {
+                static std::mutex mutex;
+                static std::unordered_map<int, bool> configured;
+                {
+                    std::lock_guard<std::mutex> lock(mutex);
+                    if (!configured[p.device.device]) {
+                        auto error = cuda_set_shared_memory(quantize_convrot, int(p.device.shared_memory - 128));
+                        if (error != cuda_success)
+                            return error;
+                        configured[p.device.device] = true;
+                    }
+                }
+                int threads = p.k <= 1024 ? 128 : p.k <= 8192 ? 256
+                                                              : 1024;
+                quantize_convrot<<<unsigned(p.rows), threads, shared, stream>>>(src, dst, scales, p.k, p.rows, p.k);
+            } else {
+                const int64_t groups = p.k / 256, total_groups = p.rows * groups;
+                convrot_group_amax_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, partial, p.k, groups, total_groups);
+                reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
+                convrot_group_quantize_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, dst, scales, p.k, groups, total_groups, p.k);
+            }
+            return cuda_last_error();
+        }
+        template <bool Dequantize>
+        static __global__ void swiglu(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows) {
+            const int64_t i    = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            const int64_t half = n / 2;
+            if (i >= rows * half)
+                return;
+            const int64_t row = i / half, col = i % half;
+            float gate, up;
+            if constexpr (Dequantize) {
+                gate = float(reinterpret_cast<const int32_t*>(src)[row * n + col]) * xs[row];
+                up   = float(reinterpret_cast<const int32_t*>(src)[row * n + half + col]) * xs[row];
+                gate *= ws[col];
+                up *= ws[half + col];
+            } else {
+                gate = src[row * n + col];
+                up   = src[row * n + half + col];
+            }
+            dst[i] = __fmul_rn(gate / (1.f + expf(-gate)), up);
+        }
+        template <bool Swiglu, bool Dequantize, bool Rotated>
+        static CUDAError quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, CUDAStream stream) {
+            const int64_t hidden = p.n / 2;
+            auto kernel          = sd_convrot::quantize<Swiglu, Dequantize, false, false, Rotated>;
+            static std::mutex mutex;
+            static std::unordered_map<int, bool> initialized;
+            {
+                std::lock_guard<std::mutex> lock(mutex);
+                if (!initialized[p.device.device]) {
+                    auto error = cuda_set_shared_memory(kernel, int(p.device.shared_memory - 128));
+                    if (error != cuda_success)
+                        return error;
+                    initialized[p.device.device] = true;
+                }
+            }
+            if (p.padded > p.rows) {
+                auto error = cuda_memset_async(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
+                if (error != cuda_success)
+                    return error;
+            }
+            const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256
+                                                                                                           : 1024)
+                                                                  : 256;
+            kernel<<<unsigned(p.rows), threads, size_t(hidden) * sizeof(float), stream>>>(src, output,
+                                                                                          reinterpret_cast<float*>(output + hidden * p.padded), hidden, p.rows, hidden,
+                                                                                          Swiglu ? src + hidden : nullptr, Swiglu ? p.n : hidden, Swiglu ? p.n : hidden, xs, ws, nullptr, 0.f, 1.f, 0.f, nullptr, nullptr);
+            return cuda_last_error();
+        }
+        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+            const auto& p = *static_cast<const Plan*>(state);
+            auto reg      = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context) || context.device != p.device.device)
+                return GGML_STATUS_FAILED;
+            auto stream     = static_cast<CUDAStream>(context.stream);
+            auto* scratch   = execution->workspace ? static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace)) : nullptr;
+            auto scratch_at = [scratch](size_t offset) { return scratch ? scratch + offset : nullptr; };
+            auto* dst       = nodes[count - 1];
+            SDI8Operation op;
+            const bool mat = sd_ggml_i8_describe(nodes[0], &op);
+            bool alias     = false;
+            for (int i = 0; i < count; ++i)
+                for (auto* input : nodes[i]->src) {
+                    bool internal = false;
+                    for (int j = 0; j < count; ++j)
+                        internal |= input == nodes[j];
+                    if (!internal)
+                        alias |= overlaps(dst, input);
+                }
+            // CUTLASS quantization epilogues read only the workspace intermediate after GEMM completes.
+            if (p.gemm && p.kind == SWIGLU_QUANTIZE)
+                alias = false;
+            if (alias && p.output == SIZE_MAX)
+                return GGML_STATUS_FAILED;
+            void* output    = alias ? scratch_at(p.output) : dst->data;
+            CUDAError error = cuda_success;
+            if (!mat) {
+                auto* quant  = static_cast<int8_t*>(output);
+                auto* scales = reinterpret_cast<float*>(quant + p.k * p.padded);
+                error        = convrot_launch(p, stream, static_cast<const float*>(dst->src[0]->data), quant, scales, reinterpret_cast<float*>(scratch_at(p.partial)));
+            } else {
+                auto* quant = reinterpret_cast<int8_t*>(scratch_at(p.quant));
+                const float* scales;
+                if (op.input->type == GGML_TYPE_F32) {
+                    auto* temporary_scales = reinterpret_cast<float*>(scratch_at(p.scales));
+                    scales                 = temporary_scales;
+                    if (op.group_size)
+                        error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch_at(p.partial)));
+                    else {
+                        if (p.padded > p.rows)
+                            error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                        if (error == cuda_success)
+                            quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(static_cast<const float*>(op.input->data), quant, temporary_scales, p.k, p.rows);
+                    }
+                } else {
+                    auto* data = static_cast<int8_t*>(op.input->data);
+                    scales     = op.packed ? reinterpret_cast<const float*>(data + p.k * p.padded) : static_cast<const float*>(op.input_scale->data);
+                    if (op.packed || p.rows == p.padded)
+                        quant = data;
+                    else {
+                        error = cuda_copy_async(quant, data, size_t(p.k) * p.rows, stream);
+                        if (error == cuda_success)
+                            error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                    }
+                }
+                if (error != cuda_success || cuda_last_error() != cuda_success)
+                    return GGML_STATUS_FAILED;
+                const auto* ws          = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
+                const auto* bias        = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+                const size_t prefix     = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0;
+                auto* gemm_output       = reinterpret_cast<float*>(static_cast<char*>(output) + prefix);
+                auto* accum             = reinterpret_cast<int32_t*>(scratch_at(p.accum));
+                const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu);
+#ifdef SD_CUDA_I8_CUTLASS
+                if (p.gemm) {
+                    auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast<const int8_t*>(op.weight->data), scales, ws,
+                                                          intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch_at(p.gemm_offset), stream);
+                    if (status != GGML_STATUS_SUCCESS)
+                        return status;
+                } else
+#endif
+                {
+                    auto blas_query = reinterpret_cast<ggml_ext_gpu_blas_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_blas_v1"));
+                    auto handle     = blas_query ? blas_query(execution->backend) : nullptr;
+                    if (!handle)
+                        return GGML_STATUS_FAILED;
+#if defined(SD_CUDA_VENDOR_MUSA)
+                    return GGML_STATUS_FAILED;
+#else
+                    const int32_t alpha = 1, beta = 0;
+#if defined(SD_CUDA_VENDOR_HIP)
+#if HIP_VERSION >= 60500000
+                    constexpr auto compute_type = HIPBLAS_COMPUTE_32I;
+#else
+                    constexpr auto compute_type = HIPBLAS_R_32I;
+#endif
+                    const auto status = hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N,
+                                                      int(p.n), int(p.padded), int(p.k), &alpha,
+                                                      op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k),
+                                                      &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT);
+                    if (status != HIPBLAS_STATUS_SUCCESS)
+#else
+                    const auto status = cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N,
+                                                     int(p.n), int(p.padded), int(p.k), &alpha,
+                                                     op.weight->data, CUDA_R_8I, int(p.k), quant, CUDA_R_8I, int(p.k),
+                                                     &beta, accum, CUDA_R_32I, int(p.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+                    if (status != CUBLAS_STATUS_SUCCESS)
+#endif
+                        return GGML_STATUS_FAILED;
+#endif
+                }
+                const auto* values = reinterpret_cast<const float*>(accum);
+                if (p.kind == SWIGLU_QUANTIZE) {
+                    if (p.fast_swiglu)
+                        error = quantize_epilogue<false, false, true>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
+                    else if (p.gemm)
+                        error = quantize_epilogue<true, false, false>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
+                    else
+                        error = quantize_epilogue<true, true, false>(p, values, scales, ws, static_cast<int8_t*>(output), stream);
+                } else if (p.kind == SWIGLU && !p.fast_swiglu) {
+                    const unsigned blocks = unsigned((p.rows * (p.n / 2) + 255) / 256);
+                    if (p.gemm)
+                        swiglu<false><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
+                    else
+                        swiglu<true><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
+                    error = cuda_last_error();
+                } else if (!p.gemm) {
+                    const unsigned gy = unsigned(std::min<int64_t>(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535);
+                    if ((reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(bias) | reinterpret_cast<uintptr_t>(gemm_output)) % 16 == 0) {
+                        dequantize_i32_rows_vec4_cuda<<<dim3(unsigned((p.n + 1023) / 1024), gy, gz), 256, 0, stream>>>(
+                            reinterpret_cast<const int4*>(accum), scales, reinterpret_cast<const float4*>(ws), reinterpret_cast<const float4*>(bias), reinterpret_cast<float4*>(gemm_output), p.n / 4, p.rows);
+                    } else {
+                        dequantize_i32_rows_cuda<<<dim3(unsigned((p.n + 255) / 256), gy, gz), 256, 0, stream>>>(accum, scales, ws, bias, gemm_output, p.n, p.rows);
+                    }
+                    error = cuda_last_error();
+                }
+                if (error == cuda_success && prefix)
+                    error = cuda_copy_async(output, dst->src[0]->data, prefix, stream);
+            }
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
+
+    }
+
+    void register_i8(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.linear_i8." SD_CUDA_NAME ".v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_i8_op();
+        desc.priority             = 100;
+        desc.supports             = sd_i8_cuda::supports;
+        desc.prepare              = sd_i8_cuda::prepare;
+        desc.execute              = sd_i8_cuda::execute;
+        desc.destroy              = sd_i8_cuda::destroy;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name     = "sd.linear_i8.epilogue." SD_CUDA_NAME ".v1";
+        desc.priority = 240;
+        desc.supports = nullptr;
+        desc.match    = sd_i8_cuda::match;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.priority  = 100;
+        desc.match     = nullptr;
+        desc.name      = "sd.quantize_i8_convrot." SD_CUDA_NAME ".v1";
+        desc.custom_op = sd_ggml_convrot_op();
+        desc.supports  = sd_i8_cuda::convrot_supports;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+#ifdef SD_CUDA_I8_CUTLASS
+        desc.name      = "sd.linear_i8.swiglu_split.cuda.v1";
+        desc.custom_op = sd_ggml_i8_op();
+        desc.priority  = 260;
+        desc.supports  = nullptr;
+        desc.match     = sd_i8_cuda::split_match;
+        desc.prepare   = sd_i8_cuda::split_prepare;
+        desc.execute   = sd_i8_cuda::split_execute;
+        desc.destroy   = sd_i8_cuda::split_destroy;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+#endif
+    }
+
+}

src/core/cuda/modulate.cu

@@ -0,0 +1,95 @@
+#include <algorithm>
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_modulate {
+        struct Parameters {
+            int64_t ne[4], scale_ne[4], shift_ne[4];
+            size_t x_stride[4], scale_stride[4], shift_stride[4];
+            int64_t count;
+        };
+
+        static __global__ void modulate(const float* x, const float* scale, const float* shift, float* dst, Parameters p) {
+            for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; i < p.count; i += int64_t(blockDim.x) * gridDim.x) {
+                int64_t index = i;
+                size_t xi = 0, si = 0, hi = 0;
+#pragma unroll
+                for (int d = 0; d < 4; ++d) {
+                    const int64_t coord = index % p.ne[d];
+                    index /= p.ne[d];
+                    xi += coord * p.x_stride[d];
+                    si += (coord % p.scale_ne[d]) * p.scale_stride[d];
+                    if (shift) {
+                        hi += (coord % p.shift_ne[d]) * p.shift_stride[d];
+                    }
+                }
+                float value = __fadd_rn(x[xi], __fmul_rn(x[xi], scale[si]));
+                dst[i]      = shift ? __fadd_rn(value, shift[hi]) : value;
+            }
+        }
+
+        static bool supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+            return sd_ggml_modulate_describe(node);
+        }
+
+        static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+            *bytes = sd_ggml_modulate_workspace_size(nodes[0]);
+            return GGML_STATUS_SUCCESS;
+        }
+
+        static void destroy(void*, void*) {}
+
+        static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void*, void*) {
+            ggml_ext_gpu_context context;
+            if (!cuda_context(e->backend, &context)) {
+                return GGML_STATUS_FAILED;
+            }
+            auto* dst          = nodes[0];
+            const auto* x      = dst->src[0];
+            const auto* scale  = dst->src[1];
+            const auto* shift  = dst->src[2];
+            Parameters p       = {};
+            p.count            = ggml_nelements(dst);
+            bool alias         = false;
+            const auto address = reinterpret_cast<uintptr_t>(dst->data);
+            for (const auto* input : {x, scale, shift}) {
+                if (input) {
+                    const auto src = reinterpret_cast<uintptr_t>(input->data);
+                    alias |= address < src + ggml_nbytes(input) && src < address + ggml_nbytes(dst);
+                }
+            }
+            for (int d = 0; d < 4; ++d) {
+                p.ne[d]           = x->ne[d];
+                p.scale_ne[d]     = scale->ne[d];
+                p.shift_ne[d]     = shift ? shift->ne[d] : 1;
+                p.x_stride[d]     = x->nb[d] / sizeof(float);
+                p.scale_stride[d] = scale->nb[d] / sizeof(float);
+                p.shift_stride[d] = shift ? shift->nb[d] / sizeof(float) : 0;
+            }
+            auto stream           = static_cast<CUDAStream>(context.stream);
+            auto* output          = static_cast<float*>(alias ? ggml_backend_buffer_get_base(e->workspace) : dst->data);
+            const unsigned blocks = unsigned(std::min<int64_t>((p.count + 255) / 256, 65535));
+            modulate<<<blocks, 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(scale->data),
+                                                 shift ? static_cast<const float*>(shift->data) : nullptr, output, p);
+            if (cuda_last_error() != cuda_success || (alias && cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) != cuda_success)) {
+                return GGML_STATUS_FAILED;
+            }
+            return GGML_STATUS_SUCCESS;
+        }
+    }
+
+    void register_modulate(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.modulate." SD_CUDA_NAME ".v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_modulate_op();
+        desc.priority             = 100;
+        desc.supports             = sd_modulate::supports;
+        desc.prepare              = sd_modulate::prepare;
+        desc.destroy              = sd_modulate::destroy;
+        desc.execute              = sd_modulate::execute;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+}

src/core/cuda/registration.cu

@@ -0,0 +1,46 @@
+#include "runtime.h"
+
+#ifdef SD_CUDA_CUTE_ATTENTION
+void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
+void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
+void sd_ggml_register_cuda_window_attention(ggml_backend_dev_t device);
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+    void register_rope(ggml_backend_dev_t device);
+    void register_i8(ggml_backend_dev_t device);
+    void register_w4(ggml_backend_dev_t device);
+    void register_convrot(ggml_backend_dev_t device);
+    void register_affine(ggml_backend_dev_t device);
+    void register_modulate(ggml_backend_dev_t device);
+    void register_bf16(ggml_backend_dev_t device);
+    void register_rms_affine(ggml_backend_dev_t device);
+#if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA)
+    void register_sage_attention(ggml_backend_dev_t device);
+    void register_sol_attention(ggml_backend_dev_t device);
+#endif
+
+    void register_extensions(ggml_backend_dev_t device) {
+        if (cuda_device_info(device).device < 0)
+            return;
+        register_affine(device);
+        register_modulate(device);
+        register_i8(device);
+        register_w4(device);
+        register_convrot(device);
+        register_rope(device);
+        register_bf16(device);
+        register_rms_affine(device);
+#if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA)
+        register_sage_attention(device);
+        register_sol_attention(device);
+#endif
+#ifdef SD_CUDA_CUTE_ATTENTION
+        sd_ggml_register_cuda_i8_projection(device);
+        sd_ggml_register_cuda_i8_gated_residual(device);
+        sd_ggml_register_cuda_attention(device);
+        sd_ggml_register_cuda_window_attention(device);
+#endif
+    }
+}

src/core/cuda/registration.h

@@ -0,0 +1,13 @@
+#ifndef __SD_CORE_CUDA_REGISTRATION_H__
+#define __SD_CORE_CUDA_REGISTRATION_H__
+#include "ggml-backend.h"
+namespace sd_cuda {
+    void register_extensions(ggml_backend_dev_t device);
+}
+namespace sd_hip {
+    void register_extensions(ggml_backend_dev_t device);
+}
+namespace sd_musa {
+    void register_extensions(ggml_backend_dev_t device);
+}
+#endif

src/core/cuda/rms_affine.cu

@@ -0,0 +1,286 @@
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include "bf16.h"
+#include "core/ggml_extensions.h"
+#include "rms_affine.h"
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_rms_affine {
+        static bool disabled() {
+            static const bool value = [] {
+                const char* fusion = std::getenv("GGML_CUDA_DISABLE_FUSION");
+                return fusion && std::atoi(fusion);
+            }();
+            return value;
+        }
+
+        static bool f32(const ggml_tensor* t) {
+            return t && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t);
+        }
+
+        static int modulate_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (disabled() || first + 2 >= ggml_graph_n_nodes(g))
+                return 0;
+            const auto* norm     = ggml_graph_node(g, first);
+            const auto* affine   = ggml_graph_node(g, first + 1);
+            const auto* modulate = ggml_graph_node(g, first + 2);
+            if (norm->op != GGML_OP_RMS_NORM || affine->op != GGML_OP_MUL || affine->src[0] != norm ||
+                !sd_ggml_modulate_describe(modulate) || modulate->src[0] != affine)
+                return 0;
+            const auto* x = norm->src[0];
+            if (!f32(x) || !f32(norm) || !f32(affine) || x->ne[0] < 32 || x->ne[0] > 8192 ||
+                ggml_nrows(x) > INT_MAX || !ggml_are_same_shape(x, norm) || !ggml_are_same_shape(x, affine))
+                return 0;
+            for (const auto* parameter : {affine->src[1], modulate->src[1], modulate->src[2]}) {
+                if (parameter && (!f32(parameter) || parameter == norm || parameter == affine ||
+                                  parameter->ne[0] != x->ne[0] || ggml_nelements(parameter) != x->ne[0]))
+                    return 0;
+            }
+            return affine->src[1] ? 3 : 0;
+        }
+
+        static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (disabled() || first + 3 >= ggml_graph_n_nodes(g))
+                return 0;
+            const auto* norm   = ggml_graph_node(g, first);
+            const auto* affine = ggml_graph_node(g, first + 1);
+            const auto* mul    = ggml_graph_node(g, first + 2);
+            const auto* add    = ggml_graph_node(g, first + 3);
+            if (norm->op != GGML_OP_RMS_NORM || affine->op != GGML_OP_MUL || affine->src[0] != norm ||
+                mul->op != GGML_OP_MUL || mul->src[0] != affine || add->op != GGML_OP_ADD ||
+                (add->src[0] != mul && add->src[1] != mul))
+                return 0;
+            const auto* x        = norm->src[0];
+            const auto* residual = add->src[0] == mul ? add->src[1] : add->src[0];
+            return f32(x) && f32(affine->src[1]) && f32(mul->src[1]) && f32(residual) && f32(add) &&
+                           x->ne[0] >= 32 && x->ne[0] <= 8192 && ggml_nrows(x) <= INT_MAX &&
+                           ggml_nelements(affine->src[1]) == x->ne[0] && ggml_nelements(mul->src[1]) == x->ne[0] &&
+                           affine->src[1]->ne[0] == x->ne[0] && mul->src[1]->ne[0] == x->ne[0] &&
+                           residual != mul && ggml_are_same_shape(residual, x) && ggml_are_same_shape(add, x)
+                       ? 4
+                       : 0;
+        }
+
+        static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
+            if (disabled())
+                return;
+            for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
+                const auto* mul = ggml_graph_node(graph, i);
+                if (!sd_ggml_modulate_describe(mul)) {
+                    if (i + 1 >= ggml_graph_n_nodes(graph))
+                        continue;
+                    const auto* add = ggml_graph_node(graph, i + 1);
+                    if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
+                        continue;
+                }
+                const auto* affine = mul->src[0];
+                if (!affine || affine->op != GGML_OP_MUL || !affine->src[0] || affine->src[0]->op != GGML_OP_RMS_NORM)
+                    continue;
+                for (int j = i - 2; j >= 0 && j >= i - 32; --j) {
+                    if (ggml_graph_node(graph, j) == affine->src[0] && ggml_graph_node(graph, j + 1) == affine) {
+                        if (j + 2 < i)
+                            ggml_ext_graph_delay(graph, j, 2, i);
+                        break;
+                    }
+                }
+            }
+        }
+
+        static __device__ __forceinline__ float warp_sum(float value) {
+            for (int offset = 16; offset > 0; offset >>= 1)
+                value = __fadd_rn(value, cuda_shuffle_xor(value, offset, 32));
+            return value;
+        }
+
+        enum class Kind { Residual,
+                          SelfResidual,
+                          Modulate };
+
+        struct Operation {
+            Kind kind = Kind::Residual;
+            const ggml_tensor* input;
+            const ggml_tensor* weight;
+            const ggml_tensor* gate;
+            const ggml_tensor* residual = nullptr;
+            const ggml_tensor* shift    = nullptr;
+            ggml_tensor* output;
+            float eps;
+        };
+
+        static Operation describe(ggml_tensor* const* nodes, int count) {
+            Operation op;
+            op.input  = nodes[0]->src[0];
+            op.weight = nodes[1]->src[1];
+            op.gate   = nodes[2]->src[1];
+            op.output = nodes[count - 1];
+            std::memcpy(&op.eps, nodes[0]->op_params, sizeof(op.eps));
+            if (count == 3) {
+                op.kind  = Kind::Modulate;
+                op.shift = op.output->src[2];
+            } else {
+                op.residual = op.output->src[0] == nodes[2] ? op.output->src[1] : op.output->src[0];
+                if (op.residual == nodes[1]) {
+                    op.kind     = Kind::SelfResidual;
+                    op.residual = nullptr;
+                }
+            }
+            return op;
+        }
+
+        template <int Block, int Groups, Kind Op, typename T>
+        static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, const float* shift, T* dst, int width, float eps) {
+            const int col0     = threadIdx.x;
+            const int64_t base = int64_t(blockIdx.x) * width;
+            float values[Groups][8];
+            __shared__ float sums[32];
+            const int lane = col0 % 32;
+#pragma unroll
+            for (int group = 0; group < Groups; ++group) {
+                float sum = 0.f;
+                for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) {
+                    values[group][j] = x[base + col];
+                    sum              = __fmaf_rn(values[group][j], values[group][j], sum);
+                }
+                sum = warp_sum(sum);
+                if (lane == 0)
+                    sums[(col0 + group * Block) / 32] = sum;
+            }
+            // Groups preserve the native reduction tree when using fewer physical threads.
+            __syncthreads();
+            const float sum   = warp_sum(lane < Block * Groups / 32 ? sums[lane] : 0.f);
+            const float scale = rsqrtf(sum / width + eps);
+#pragma unroll
+            for (int group = 0; group < Groups; ++group) {
+                for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) {
+                    float value = __fmul_rn(__fmul_rn(scale, values[group][j]), weight[col]);
+                    if constexpr (Op == Kind::Modulate) {
+                        value = __fadd_rn(value, __fmul_rn(value, gate[col]));
+                        if (shift)
+                            value = __fadd_rn(value, shift[col]);
+                    } else {
+                        const float r = Op == Kind::SelfResidual ? value : residual[base + col];
+                        value         = __fadd_rn(__fmul_rn(value, gate[col]), r);
+                    }
+                    dst[base + col] = T(value);
+                }
+            }
+        }
+
+        static uintptr_t address(const ggml_tensor* tensor) {
+            size_t offset = 0;
+            while (tensor && !tensor->data && tensor->view_src) {
+                offset += tensor->view_offs;
+                tensor = tensor->view_src;
+            }
+            return tensor && tensor->data ? reinterpret_cast<uintptr_t>(tensor->data) + offset : 0;
+        }
+        static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+            if (!b)
+                return false;
+            const auto x = address(a), y = address(b);
+            if (!x || !y)
+                return false;
+            return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+        }
+        static bool needs_scratch(const Operation& op) {
+            const auto* dst = op.output;
+            return (overlaps(dst, op.input) && address(dst) != address(op.input)) ||
+                   (overlaps(dst, op.residual) && address(dst) != address(op.residual)) ||
+                   overlaps(dst, op.weight) || overlaps(dst, op.gate) || overlaps(dst, op.shift);
+        }
+
+        static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+            const auto op    = describe(nodes, count);
+            const auto* dst  = op.output;
+            const bool fresh = !dst->data && !dst->buffer && !dst->view_src;
+            // Fresh outputs may reuse the residual elementwise; external inputs stay live.
+            *workspace = fresh || (dst->data && !needs_scratch(op)) ? 0 : ggml_nbytes(dst);
+            return GGML_STATUS_SUCCESS;
+        }
+        static void destroy(void*, void*) {}
+
+        template <typename T, Kind Op>
+        static CUDAError launch(CUDAStream stream, const Operation& op, T* output) {
+            auto compute = kernel<256, 1, Op, T>;
+            int threads  = 256;
+            if (op.input->ne[0] >= 1024) {
+                if constexpr (Op == Kind::Modulate) {
+                    compute = kernel<512, 2, Op, T>;
+                    threads = 512;
+                } else {
+                    compute = kernel<1024, 1, Op, T>;
+                    threads = 1024;
+                }
+            }
+            compute<<<unsigned(ggml_nrows(op.input)), threads, 0, stream>>>(
+                static_cast<const float*>(op.input->data), static_cast<const float*>(op.weight->data),
+                static_cast<const float*>(op.gate->data), op.residual ? static_cast<const float*>(op.residual->data) : nullptr,
+                op.shift ? static_cast<const float*>(op.shift->data) : nullptr, output, int(op.input->ne[0]), op.eps);
+            return cuda_last_error();
+        }
+
+        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context))
+                return GGML_STATUS_FAILED;
+            const auto op    = describe(nodes, count);
+            auto* dst        = op.output;
+            const bool alias = needs_scratch(op);
+            if (alias && !execution->workspace)
+                return GGML_STATUS_FAILED;
+            auto* output      = static_cast<float*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
+            const auto stream = static_cast<CUDAStream>(context.stream);
+            CUDAError error   = cuda_invalid_value;
+            switch (op.kind) {
+                case Kind::Residual:
+                    error = launch<float, Kind::Residual>(stream, op, output);
+                    break;
+                case Kind::SelfResidual:
+                    error = launch<float, Kind::SelfResidual>(stream, op, output);
+                    break;
+                case Kind::Modulate:
+                    error = launch<float, Kind::Modulate>(stream, op, output);
+                    break;
+            }
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
+    }
+
+    int rms_affine_match(const ggml_cgraph* graph, int first) {
+        return sd_rms_affine::match(nullptr, graph, first, nullptr);
+    }
+
+    ggml_status rms_affine_pack(CUDAStream stream, ggml_tensor* const* nodes, void* output) {
+        const auto op    = sd_rms_affine::describe(nodes, 4);
+        auto* dst        = static_cast<CUDABFloat16*>(output);
+        const auto error = op.kind == sd_rms_affine::Kind::SelfResidual
+                               ? sd_rms_affine::launch<CUDABFloat16, sd_rms_affine::Kind::SelfResidual>(stream, op, dst)
+                               : sd_rms_affine::launch<CUDABFloat16, sd_rms_affine::Kind::Residual>(stream, op, dst);
+        return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+
+    void register_rms_affine(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.rms_affine_residual." SD_CUDA_NAME ".v1";
+        desc.op                   = GGML_OP_RMS_NORM;
+        desc.priority             = 300;
+        desc.match                = sd_rms_affine::match;
+        desc.optimize             = sd_rms_affine::optimize;
+        desc.prepare              = sd_rms_affine::prepare;
+        desc.destroy              = sd_rms_affine::destroy;
+        desc.execute              = sd_rms_affine::execute;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name = "sd.rms_affine_modulate." SD_CUDA_NAME ".v1";
+        // Prefer the existing fusion that also produces ConvRot input when available.
+        desc.priority = 200;
+        desc.match    = sd_rms_affine::modulate_match;
+        desc.optimize = nullptr;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+}

src/core/cuda/rms_affine.h

@@ -0,0 +1,11 @@
+#ifndef __SD_CORE_CUDA_RMS_AFFINE_H__
+#define __SD_CORE_CUDA_RMS_AFFINE_H__
+
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    int rms_affine_match(const ggml_cgraph* graph, int first);
+    ggml_status rms_affine_pack(CUDAStream stream, ggml_tensor* const* nodes, void* output);
+}
+
+#endif

src/core/cuda/rope.cu

@@ -0,0 +1,347 @@
+#include "core/cuda/rope.h"
+#include "core/ggml_extensions.h"
+
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <type_traits>
+
+namespace SD_CUDA_NAMESPACE {
+    static __device__ __forceinline__ float sd_warp_sum(float v) {
+        for (int mask = 16; mask > 0; mask >>= 1)
+            v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
+        return v;
+    }
+    template <int D, bool interleaved, typename T>
+    static __global__ void rms_norm_mul_rope_apply_fused(
+        const float* x,
+        const float* weight,
+        const float* pe,
+        T* dst,
+        int64_t heads,
+        int64_t tokens,
+        int64_t rows,
+        int64_t nb1,
+        int64_t nb2,
+        int64_t nb3,
+        int rot_half,
+        float eps,
+        float output_scale) {
+        const int lane    = threadIdx.x;
+        const int64_t row = int64_t(blockIdx.x) * blockDim.y + threadIdx.y;
+        if (row >= rows) {
+            return;
+        }
+        const int64_t head  = row % heads;
+        const int64_t token = (row / heads) % tokens;
+        const int64_t batch = row / (heads * tokens);
+        const float* x_row  = x + head * nb1 + token * nb2 + batch * nb3;
+        float v[D / 32], sums[8] = {};
+#pragma unroll
+        for (int j = 0; j < D / 32; ++j) {
+            v[j]    = x_row[j * 32 + lane];
+            sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j]));
+        }
+        // Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps.
+        const float sum = __fadd_rn(
+            __fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])),
+            __fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7])));
+        const float scale = rsqrtf(sum / float(D) + eps);
+        T* out            = dst + ((batch * heads + head) * tokens + token) * D;
+#pragma unroll
+        for (int j = 0; j < D / 32; ++j) {
+            const int col     = j * 32 + lane;
+            const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
+            float rotated     = value;
+            if constexpr (interleaved) {
+                const float other = cuda_shuffle_xor(value, 1, 32);
+                if (col < 2 * rot_half) {
+                    const float* m      = pe + 4 * (token * rot_half + col / 2);
+                    const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
+                    rotated             = lane % 2 == 0
+                                              ? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y))
+                                              : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w));
+                }
+            } else {
+                if (col < 2 * rot_half) {
+                    const int partner   = col < rot_half ? col + rot_half : col - rot_half;
+                    const float other   = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
+                    const float* m      = pe + 4 * (token * rot_half + col % rot_half);
+                    const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
+                    rotated             = col < rot_half
+                                              ? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z))
+                                              : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x));
+                }
+            }
+            if constexpr (std::is_same<T, half>::value) {
+                out[col] = __float2half_rn(__fmul_rn(rotated, output_scale));
+            } else {
+                out[col] = rotated;
+            }
+        }
+    }
+
+    struct RopeApplyStrides {
+        int64_t x[4];
+        int64_t pe[4];
+    };
+
+    template <bool interleaved>
+    static __global__ void rope_apply_f32(
+        const float* x,
+        const float* pe,
+        float* dst,
+        int64_t half_dim,
+        int64_t rot_half,
+        int64_t heads,
+        int64_t tokens,
+        int64_t pairs,
+        RopeApplyStrides strides) {
+        const int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+        if (index >= pairs) {
+            return;
+        }
+        const int64_t j     = index % half_dim;
+        const int64_t row   = index / half_dim;
+        const int64_t token = row % tokens;
+        const int64_t head  = (row / tokens) % heads;
+        const int64_t batch = row / (tokens * heads);
+        const int64_t i0    = interleaved || j >= rot_half ? 2 * j : j;
+        const int64_t i1    = interleaved || j >= rot_half ? 2 * j + 1 : j + rot_half;
+        const float* x_row  = x + head * strides.x[1] + token * strides.x[2] + batch * strides.x[3];
+        const float a       = x_row[i0 * strides.x[0]];
+        const float b       = x_row[i1 * strides.x[0]];
+        float* out          = dst + row * half_dim * 2;
+        if (j >= rot_half) {
+            out[i0] = a;
+            out[i1] = b;
+            return;
+        }
+        const float* matrix = pe + j * strides.pe[2] + token * strides.pe[3];
+        const float c       = matrix[0];
+        const float s       = matrix[strides.pe[1]];
+        // Match the separate mul/add rounding of the unfused graph.
+        if (interleaved) {
+            out[i0] = __fadd_rn(__fmul_rn(a, c), __fmul_rn(b, matrix[strides.pe[0]]));
+            out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, matrix[strides.pe[0] + strides.pe[1]]));
+        } else {
+            out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s));
+            out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c));
+        }
+    }
+
+    static bool rope_interleaved(const ggml_tensor* rope) {
+        int32_t value;
+        std::memcpy(&value, ggml_ext_op_attributes(rope), sizeof(value));
+        return value != 0;
+    }
+
+    static float param_f32(const ggml_tensor* node, int index) {
+        float value;
+        std::memcpy(&value, node->op_params + index, sizeof(value));
+        return value;
+    }
+
+    static bool rope_is_node(const ggml_tensor* node) {
+        return node && ggml_ext_op_get(node) == sd_ggml_rope_op();
+    }
+
+    static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
+        const auto* x  = rope->src[0];
+        const auto* pe = rope->src[1];
+        if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 ||
+            x->ne[0] % 2 || pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 ||
+            pe->ne[3] != x->ne[2] || !ggml_is_contiguous(rope) || rope->ne[0] != x->ne[0] ||
+            rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 ||
+            (ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) {
+            return false;
+        }
+        for (int i = 0; i < 4; ++i) {
+            if (x->nb[i] % sizeof(float) || pe->nb[i] % sizeof(float))
+                return false;
+        }
+        return true;
+    }
+
+    bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
+        if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !rope_is_node(rope) ||
+            mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr))
+            return false;
+        const auto* x      = norm->src[0];
+        const auto* weight = mul->src[1];
+        return x && weight && weight != norm && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
+               x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
+               ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
+               weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
+               x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 &&
+               x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) &&
+               (ggml_nrows(x) + 3) / 4 <= INT_MAX;
+    }
+
+    template <typename T, bool Interleaved>
+    static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
+        const auto* x      = nodes[0]->src[0];
+        const auto* pe     = nodes[2]->src[1];
+        const int64_t rows = ggml_nrows(x);
+#define SD_ROPE(D)                                                                                                               \
+    case D:                                                                                                                      \
+        rms_norm_mul_rope_apply_fused<D, Interleaved, T><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \
+            static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data),                               \
+            static_cast<const float*>(pe->data), output, x->ne[1], x->ne[2], rows,                                               \
+            x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale);                            \
+        break
+        switch (x->ne[0]) {
+            SD_ROPE(32);
+            SD_ROPE(64);
+            SD_ROPE(96);
+            SD_ROPE(128);
+            SD_ROPE(160);
+            SD_ROPE(192);
+            SD_ROPE(224);
+            SD_ROPE(256);
+            default:
+                return cuda_invalid_value;
+        }
+#undef SD_ROPE
+        return cuda_last_error();
+    }
+
+    CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
+        if (fp16) {
+            return rope_interleaved(nodes[2]) ? rms_rope_launch<half, true>(stream, nodes, static_cast<half*>(output), scale)
+                                              : rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
+        }
+        return rope_interleaved(nodes[2]) ? rms_rope_launch<float, true>(stream, nodes, static_cast<float*>(output), scale)
+                                          : rms_rope_launch<float, false>(stream, nodes, static_cast<float*>(output), scale);
+    }
+
+    static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+        static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }();
+        if (disabled)
+            return 0;
+        auto* g = const_cast<ggml_cgraph*>(graph);
+        if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports(
+                                                      ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
+            return 0;
+        if (first + 4 < ggml_graph_n_nodes(g)) {
+            const auto* rope  = ggml_graph_node(g, first + 2);
+            const auto* scale = ggml_graph_node(g, first + 3);
+            const auto* copy  = ggml_graph_node(g, first + 4);
+            if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
+                scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
+                ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+                return 5;
+        }
+        if (first + 3 < ggml_graph_n_nodes(g)) {
+            const auto* rope = ggml_graph_node(g, first + 2);
+            const auto* copy = ggml_graph_node(g, first + 3);
+            if (copy->op == GGML_OP_CPY && copy->src[0] == rope && copy->type == GGML_TYPE_F16 &&
+                ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+                return 4;
+        }
+        return 3;
+    }
+
+    static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
+        static const bool disabled = [] {
+            const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
+            return value && std::atoi(value) != 0;
+        }();
+        if (disabled)
+            return;
+        const int total = ggml_graph_n_nodes(graph);
+        for (int i = 2; i + 1 < total; ++i) {
+            const auto* mul  = ggml_graph_node(graph, i);
+            const auto* rope = ggml_graph_node(graph, i + 1);
+            if (mul->op != GGML_OP_MUL || !rope_is_node(rope) || rope->src[0] != mul ||
+                !mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM)
+                continue;
+            const auto* norm = mul->src[0];
+            for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
+                if (ggml_graph_node(graph, j) != norm)
+                    continue;
+                if (j + 1 < i) {
+                    if (j >= 2) {
+                        const auto* mat  = ggml_graph_node(graph, j - 2);
+                        const auto* view = ggml_graph_node(graph, j - 1);
+                        if (ggml_ext_op_get(mat) == sd_ggml_i8_op() &&
+                            view->op == GGML_OP_RESHAPE && view->src[0] == mat && norm->src[0] == view &&
+                            ggml_ext_graph_delay(graph, j - 2, 3, i))
+                            break;
+                    }
+                    ggml_ext_graph_delay(graph, j, 1, i);
+                }
+                break;
+            }
+        }
+    }
+
+    static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+        const auto* dst = nodes[count - 1];
+        // Owning casts are fresh outputs; copies into existing storage retain view_src.
+        const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src;
+        *workspace              = fresh_output ? 0 : ggml_nbytes(dst);
+        return GGML_STATUS_SUCCESS;
+    }
+
+    static void rope_destroy(void*, void*) {}
+
+    static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+        const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+        return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+    }
+
+    static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+        ggml_ext_gpu_context context;
+        if (!cuda_context(execution->backend, &context))
+            return GGML_STATUS_FAILED;
+        auto stream      = static_cast<CUDAStream>(context.stream);
+        const auto* rope = nodes[count == 1 ? 0 : 2];
+        const auto* x    = count == 1 ? rope->src[0] : nodes[0]->src[0];
+        auto* dst        = nodes[count - 1];
+        const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
+        void* output     = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+        CUDAError error;
+        if (count > 1) {
+            error = rms_rope_launch(stream, nodes, output, count > 3, count == 5 ? param_f32(nodes[3], 0) : 1.f);
+        } else {
+            RopeApplyStrides strides;
+            const auto* pe = rope->src[1];
+            for (int i = 0; i < 4; ++i) {
+                strides.x[i]  = x->nb[i] / 4;
+                strides.pe[i] = pe->nb[i] / 4;
+            }
+            const int64_t pairs = ggml_nelements(x) / 2;
+            auto kernel         = rope_interleaved(rope) ? rope_apply_f32<true> : rope_apply_f32<false>;
+            kernel<<<static_cast<unsigned int>((pairs + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(pe->data),
+                                                                                       static_cast<float*>(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides);
+            error = cuda_last_error();
+        }
+        if (error == cuda_success && alias)
+            error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+        return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+
+    void register_rope(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.rope_apply." SD_CUDA_NAME ".v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_rope_op();
+        desc.priority             = 100;
+        desc.supports             = rope_supports;
+        desc.prepare              = rope_prepare;
+        desc.destroy              = rope_destroy;
+        desc.execute              = rope_execute;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name      = "sd.rms_norm_mul_rope." SD_CUDA_NAME ".v1";
+        desc.op        = GGML_OP_RMS_NORM;
+        desc.custom_op = nullptr;
+        desc.priority  = 250;
+        desc.supports  = nullptr;
+        desc.match     = rope_match;
+        desc.optimize  = rope_optimize;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
+}

src/core/cuda/rope.h

@@ -0,0 +1,11 @@
+#ifndef __SD_CORE_CUDA_ROPE_H__
+#define __SD_CORE_CUDA_ROPE_H__
+
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
+    CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
+
+}
+#endif

src/core/cuda/runtime.h

@@ -0,0 +1,111 @@
+#ifndef __SD_CORE_CUDA_RUNTIME_H__
+#define __SD_CORE_CUDA_RUNTIME_H__
+
+#include "ggml-ext.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#define SD_CUDA_NAMESPACE sd_hip
+#define SD_CUDA_NAME "hip"
+#elif defined(SD_CUDA_VENDOR_MUSA)
+#include <musa_fp16.h>
+#include <musa_runtime.h>
+#define SD_CUDA_NAMESPACE sd_musa
+#define SD_CUDA_NAME "musa"
+#else
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#define SD_CUDA_NAMESPACE sd_cuda
+#define SD_CUDA_NAME "cuda"
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+#if defined(SD_CUDA_VENDOR_HIP)
+    using CUDAError                          = hipError_t;
+    using CUDAStream                         = hipStream_t;
+    static constexpr auto cuda_success       = hipSuccess;
+    static constexpr auto cuda_invalid_value = hipErrorInvalidValue;
+    static constexpr auto cuda_kind          = GGML_EXT_GPU_HIP;
+    static CUDAError cuda_last_error() {
+        return hipGetLastError();
+    }
+    static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+        return hipMemsetAsync(dst, value, bytes, stream);
+    }
+    static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+        return hipMemcpyAsync(dst, src, bytes, hipMemcpyDeviceToDevice, stream);
+    }
+    template <typename Kernel>
+    static CUDAError cuda_set_shared_memory(Kernel, int) {
+        // HIP uses the device's shared-memory limit without CUDA's opt-in attribute.
+        return cuda_success;
+    }
+#elif defined(SD_CUDA_VENDOR_MUSA)
+    using CUDAError                          = musaError_t;
+    using CUDAStream                         = musaStream_t;
+    static constexpr auto cuda_success       = musaSuccess;
+    static constexpr auto cuda_invalid_value = musaErrorInvalidValue;
+    static constexpr auto cuda_kind          = GGML_EXT_GPU_MUSA;
+    static CUDAError cuda_last_error() {
+        return musaGetLastError();
+    }
+    static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+        return musaMemsetAsync(dst, value, bytes, stream);
+    }
+    static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+        return musaMemcpyAsync(dst, src, bytes, musaMemcpyDeviceToDevice, stream);
+    }
+    template <typename Kernel>
+    static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
+        return musaFuncSetAttribute(kernel, musaFuncAttributeMaxDynamicSharedMemorySize, bytes);
+    }
+#else
+    using CUDAError                          = cudaError_t;
+    using CUDAStream                         = cudaStream_t;
+    static constexpr auto cuda_success       = cudaSuccess;
+    static constexpr auto cuda_invalid_value = cudaErrorInvalidValue;
+    static constexpr auto cuda_kind          = GGML_EXT_GPU_CUDA;
+    static CUDAError cuda_last_error() {
+        return cudaGetLastError();
+    }
+    static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+        return cudaMemsetAsync(dst, value, bytes, stream);
+    }
+    static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+        return cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, stream);
+    }
+    template <typename Kernel>
+    static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
+        return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
+    }
+#endif
+    template <typename T>
+    static __device__ __forceinline__ T cuda_shuffle(T value, int lane, int width = 32) {
+#if defined(SD_CUDA_VENDOR_HIP)
+        return __shfl(value, lane, width);
+#else
+        return __shfl_sync(0xffffffff, value, lane, width);
+#endif
+    }
+    template <typename T>
+    static __device__ __forceinline__ T cuda_shuffle_xor(T value, int lane, int width = 32) {
+#if defined(SD_CUDA_VENDOR_HIP)
+        return __shfl_xor(value, lane, width);
+#else
+        return __shfl_xor_sync(0xffffffff, value, lane, width);
+#endif
+    }
+    static bool cuda_context(ggml_backend_t backend, ggml_ext_gpu_context* context) {
+        auto reg   = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        auto query = reinterpret_cast<ggml_ext_gpu_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1"));
+        return query && query(backend, context) && context->kind == cuda_kind;
+    }
+    static ggml_ext_gpu_device_info cuda_device_info(ggml_backend_dev_t device) {
+        ggml_ext_gpu_device_info info = {};
+        auto query                    = reinterpret_cast<ggml_ext_gpu_device_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_device_v1"));
+        if (!query || !query(device, &info) || info.kind != cuda_kind)
+            return {cuda_kind, -1, 0, 0, false};
+        return info;
+    }
+}
+#endif

src/core/cuda/sage_attention.cu

@@ -0,0 +1,387 @@
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+#if __CUDACC_VER_MAJOR__ >= 12
+
+#include <algorithm>
+#include <climits>
+#include <cstring>
+#include <map>
+#include "sageattention/attn.cuh"
+#include "sageattention/attn_fp8.cuh"
+
+static constexpr int sage_cuda_version = __CUDACC_VER_MAJOR__ * 1000 + __CUDACC_VER_MINOR__ * 10;
+
+namespace sd_cuda {
+
+    struct SageDevice {
+        int cc;
+        int compiled_cc;
+    };
+
+    static ggml_sage_attn_mode sage_mode(const SageDevice& device, ggml_sage_attn_mode requested) {
+        if (requested != GGML_SAGE_ATTN_AUTO)
+            return requested;
+        if (sage_cuda_version >= 12040 && device.cc >= 89 && device.compiled_cc >= 89) {
+            if (sage_cuda_version >= 12080 && device.cc != 90)
+                return GGML_SAGE_ATTN_2_PLUS_PLUS;
+            return GGML_SAGE_ATTN_2;
+        }
+        return GGML_SAGE_ATTN_AUTO;
+    }
+
+    static __device__ __forceinline__ float warp_reduce_max(float x) {
+#pragma unroll
+        for (int offset = 16; offset > 0; offset >>= 1)
+            x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset));
+        return x;
+    }
+
+    static __global__ void sage_output(const half* input, float* output, int count) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i < count)
+            output[i] = __half2float(input[i]);
+    }
+
+    struct SageWorkspace {
+        size_t q;
+        size_t k;
+        size_t q_scale;
+        size_t k_scale;
+        size_t k_mean;
+        size_t v;
+        size_t v_scale;
+        size_t output;
+        size_t size;
+
+        explicit SageWorkspace(const ggml_tensor* dst) {
+            const ggml_tensor* qt = dst->src[0];
+            const ggml_tensor* kt = dst->src[1];
+            size                  = 0;
+            q                     = reserve(ggml_nelements(qt));
+            k                     = reserve(ggml_nelements(kt));
+            q_scale               = reserve(GGML_PAD(qt->ne[1], 128) / 32 * 8 * qt->ne[2] * qt->ne[3] * sizeof(float));
+            k_scale               = reserve(GGML_PAD(kt->ne[1], 64) / 64 * 4 * kt->ne[2] * kt->ne[3] * sizeof(float));
+            k_mean                = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float));
+            v                     = reserve(kt->ne[0] * GGML_PAD(kt->ne[1], 64) * kt->ne[2] * kt->ne[3]);
+            v_scale               = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float));
+            output                = reserve(ggml_nelements(dst) * sizeof(half));
+        }
+
+        size_t reserve(size_t bytes) {
+            const size_t offset = size;
+            size += GGML_PAD(bytes, 256);
+            return offset;
+        }
+    };
+
+    template <int D>
+    static __global__ void sage_key_mean(const float* k, float* mean, int length) {
+        const int channel = blockIdx.x * 32 + threadIdx.x % 32;
+        const int row     = threadIdx.x / 32;
+        const size_t base = size_t(blockIdx.y) * length * D;
+        float sum         = 0.0f;
+        for (int token = row; token < length; token += 8) {
+            sum += k[base + size_t(token) * D + channel];
+        }
+        __shared__ float sums[256];
+        sums[threadIdx.x] = sum;
+        __syncthreads();
+        if (row == 0) {
+            for (int i = 1; i < 8; ++i) {
+                sum += sums[threadIdx.x + i * 32];
+            }
+            mean[size_t(blockIdx.y) * D + channel] = sum / length;
+        }
+    }
+
+    template <int D, int TOKENS, bool SMOOTH>
+    static __global__ void sage_quantize(const float* input, const float* mean, int8_t* output, float* scales, int length) {
+        constexpr int VALUES = D * TOKENS / 256;
+        const int head       = blockIdx.y;
+        const int first      = blockIdx.x * TOKENS * D;
+        const size_t base    = size_t(head) * length * D;
+        float values[VALUES];
+        float amax = 1e-7f;
+#pragma unroll
+        for (int i = 0; i < VALUES; ++i) {
+            const int index = first + threadIdx.x + i * 256;
+            float value     = 0.0f;
+            if (index < length * D) {
+                value = input[base + index];
+                if constexpr (SMOOTH) {
+                    value -= mean[size_t(head) * D + index % D];
+                }
+            }
+            values[i] = value;
+            amax      = fmaxf(amax, fabsf(value));
+        }
+        amax = warp_reduce_max(amax);
+        __shared__ float maxima[8];
+        if (threadIdx.x % 32 == 0) {
+            maxima[threadIdx.x / 32] = amax;
+        }
+        __syncthreads();
+        if (threadIdx.x < 32) {
+            amax = threadIdx.x < 8 ? maxima[threadIdx.x] : 1e-7f;
+            amax = warp_reduce_max(amax);
+            if (threadIdx.x == 0) {
+                maxima[0]                                     = amax;
+                scales[size_t(head) * gridDim.x + blockIdx.x] = amax / 127.0f;
+            }
+        }
+        __syncthreads();
+        const float inverse_scale = 127.0f / maxima[0];
+#pragma unroll
+        for (int i = 0; i < VALUES; ++i) {
+            const int index = first + threadIdx.x + i * 256;
+            if (index < length * D) {
+                output[base + index] = int8_t(__float2int_rn(values[i] * inverse_scale));
+            }
+        }
+    }
+
+    template <int D, bool KEY>
+    static __global__ void sage_quantize_per_thread(const float* input, const float* mean, int8_t* output, float* scales, int length) {
+        constexpr int GROUPS = KEY ? 4 : 8;
+        constexpr int TOKENS = KEY ? 64 : 32;
+        constexpr int VALUES = D * TOKENS / GROUPS / 32;
+        const int group      = threadIdx.x / 32;
+        const int lane       = threadIdx.x % 32;
+        const size_t base    = size_t(blockIdx.y) * length * D;
+        float values[VALUES];
+        float amax = 0.0f;
+#pragma unroll
+        for (int i = 0; i < VALUES; ++i) {
+            const int index = lane + i * 32;
+            const int row   = index / D;
+            const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group);
+            float value     = token < length ? input[base + size_t(token) * D + index % D] : 0.0f;
+            if constexpr (KEY) {
+                if (token < length) {
+                    value -= mean[size_t(blockIdx.y) * D + index % D];
+                }
+            }
+            values[i] = value;
+            amax      = fmaxf(amax, fabsf(value));
+        }
+        const float scale = warp_reduce_max(amax) / 127.0f + 1e-7f;
+        if (lane == 0) {
+            scales[(size_t(blockIdx.y) * gridDim.x + blockIdx.x) * GROUPS + group] = scale;
+        }
+#pragma unroll
+        for (int i = 0; i < VALUES; ++i) {
+            const int index = lane + i * 32;
+            const int row   = index / D;
+            const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group);
+            if (token < length) {
+                const float value                            = values[i] / scale;
+                output[base + size_t(token) * D + index % D] = int8_t(__float2int_rz(value + copysignf(0.5f, value)));
+            }
+        }
+    }
+
+    template <int D>
+    static __global__ void sage_value_scale(const half* v, float* scales, int length, float scale_max) {
+        const int channel = blockIdx.x * 32 + threadIdx.x % 32;
+        const int row     = threadIdx.x / 32;
+        const size_t base = size_t(blockIdx.y) * length * D;
+        float amax        = 1e-7f;
+        for (int token = row; token < length; token += 8) {
+            amax = fmaxf(amax, fabsf(__half2float(v[base + size_t(token) * D + channel])));
+        }
+        __shared__ float maxima[256];
+        maxima[threadIdx.x] = amax;
+        __syncthreads();
+        if (row == 0) {
+            for (int i = 1; i < 8; ++i) {
+                amax = fmaxf(amax, maxima[threadIdx.x + i * 32]);
+            }
+            scales[size_t(blockIdx.y) * D + channel] = amax / scale_max;
+        }
+    }
+
+    template <int D>
+    static __global__ void sage_quantize_value(const half* v, const float* scales, int8_t* output, int length, int padded) {
+        __shared__ uint8_t tile[32][D + 4];
+        const int head    = blockIdx.y;
+        const int first   = blockIdx.x * 32;
+        const size_t base = size_t(head) * length * D;
+        for (int i = threadIdx.x; i < 32 * D; i += 256) {
+            const int token   = first + i / D;
+            const int channel = i % D;
+            const float value = token < length ? __half2float(v[base + size_t(token) * D + channel]) / scales[size_t(head) * D + channel] : 0.0f;
+            // FP8 MMA consumes pairs of tokens in a permuted, transposed V tile.
+            const int row           = i / D;
+            const int permuted      = row / 16 * 16 + row % 16 / 8 * 2 + row % 8 / 2 * 4 + row % 2;
+            tile[permuted][channel] = __nv_cvt_float_to_fp8(value, __NV_SATFINITE, __NV_E4M3);
+        }
+        __syncthreads();
+        for (int i = threadIdx.x; i < 32 * D; i += 256) {
+            output[(size_t(head) * D + i / 32) * padded + first + i % 32] = tile[i % 32][i / 32];
+        }
+    }
+
+    template <int D, bool PLUS_PLUS>
+    static void sage_attn_fp8(cudaStream_t stream, ggml_tensor* dst, char* data, const SageWorkspace& work, float scale) {
+        const ggml_tensor* q = dst->src[0];
+        const ggml_tensor* k = dst->src[1];
+        const ggml_tensor* v = dst->src[2];
+        const int lq = q->ne[1], lk = k->ne[1], padded = GGML_PAD(lk, 64);
+        const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3];
+        auto* qi  = reinterpret_cast<int8_t*>(data + work.q);
+        auto* ki  = reinterpret_cast<int8_t*>(data + work.k);
+        auto* vi  = reinterpret_cast<int8_t*>(data + work.v);
+        auto* qs  = reinterpret_cast<float*>(data + work.q_scale);
+        auto* ks  = reinterpret_cast<float*>(data + work.k_scale);
+        auto* vs  = reinterpret_cast<float*>(data + work.v_scale);
+        auto* km  = reinterpret_cast<float*>(data + work.k_mean);
+        auto* out = reinterpret_cast<half*>(data + work.output);
+        sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
+        sage_quantize_per_thread<D, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
+            static_cast<const float*>(q->data), nullptr, qi, qs, lq);
+        sage_quantize_per_thread<D, true><<<dim3(padded / 64, hk * batch), 128, 0, stream>>>(
+            static_cast<const float*>(k->data), km, ki, ks, lk);
+        // The smaller V range keeps the FP16 instruction accumulator finite in SageAttention2++.
+        sage_value_scale<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(v->data), vs, lk, PLUS_PLUS ? 2.25f : 448.0f);
+        sage_quantize_value<D><<<dim3(padded / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(v->data), vs, vi, lk, padded);
+        auto kernel = sageattention::qk_int_sv_f8_attn_kernel<128, 64, 32, 64, D,
+                                                              sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerThread,
+                                                              sageattention::QuantGranularity::kPerThread, float, true, half,
+                                                              sageattention::ComputeUnit::kCudaCore, sageattention::MaskMode::kNone, false, true, false, PLUS_PLUS>;
+        kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), 256 * D, stream>>>(
+            qi, ki, vi, out, nullptr, qs, ks, vs, nullptr, lq, lk, hq / hk,
+            D * lq * hq, D, D * lq, D * lk * hk, D, D * lk,
+            D * padded * hk, D * padded, padded, D * lq * hq, D * hq, D, scale);
+        sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast<float*>(dst->data), ggml_nelements(dst));
+    }
+
+    template <int D>
+    static void sage_attn(cudaStream_t stream, ggml_tensor* dst, char* data, ggml_sage_attn_mode mode, float scale) {
+        const ggml_tensor* q = dst->src[0];
+        const ggml_tensor* k = dst->src[1];
+        const ggml_tensor* v = dst->src[2];
+        const int lq = q->ne[1], lk = k->ne[1];
+        const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3];
+        const SageWorkspace work(dst);
+        if (mode == GGML_SAGE_ATTN_2_PLUS_PLUS) {
+            sage_attn_fp8<D, true>(stream, dst, data, work, scale);
+            return;
+        }
+        if (mode == GGML_SAGE_ATTN_2) {
+            sage_attn_fp8<D, false>(stream, dst, data, work, scale);
+            return;
+        }
+        auto* qi  = reinterpret_cast<int8_t*>(data + work.q);
+        auto* ki  = reinterpret_cast<int8_t*>(data + work.k);
+        auto* qs  = reinterpret_cast<float*>(data + work.q_scale);
+        auto* ks  = reinterpret_cast<float*>(data + work.k_scale);
+        auto* km  = reinterpret_cast<float*>(data + work.k_mean);
+        auto* out = reinterpret_cast<half*>(data + work.output);
+
+        sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
+        sage_quantize<D, 32, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
+            static_cast<const float*>(q->data), nullptr, qi, qs, lq);
+        sage_quantize<D, 64, true><<<dim3(GGML_PAD(lk, 64) / 64, hk * batch), 256, 0, stream>>>(
+            static_cast<const float*>(k->data), km, ki, ks, lk);
+
+        constexpr size_t shared_bytes = (128 + 64 + 64 * 2) * D;
+        auto kernel                   = sageattention::qk_int_sv_f16_attn_kernel<128, 64, 32, 64, D,
+                                                                                 sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp,
+                                                                                 sageattention::QuantGranularity::kPerWarp, float, false, half,
+                                                                                 sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>;
+        kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), shared_bytes, stream>>>(
+            qi, ki, static_cast<half*>(v->data), out, nullptr, qs, ks, nullptr,
+            lq, lk, hq / hk,
+            D * lq * hq, D, D * lq,
+            D * lk * hk, D, D * lk,
+            D * lk * hk, D, D * lk,
+            D * lq * hq, D * hq, D,
+            scale);
+        sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast<float*>(dst->data), ggml_nelements(dst));
+    }
+
+    static bool sage_supports(ggml_backend_dev_t, const ggml_tensor* op, void* user_data) {
+        const auto& device = *static_cast<const SageDevice*>(user_data);
+        SDSageAttentionParams params;
+        if (!sd_ggml_sage_attn_describe(op, &params) || device.cc < 80 || device.compiled_cc < 80)
+            return false;
+        const auto mode = sage_mode(device, params.mode);
+        if (mode != GGML_SAGE_ATTN_AUTO && (sage_cuda_version < 12040 || device.cc < 89 || device.compiled_cc < 89))
+            return false;
+        if (mode == GGML_SAGE_ATTN_2_PLUS_PLUS && sage_cuda_version < 12080)
+            return false;
+        const auto* q = op->src[0];
+        const auto* k = op->src[1];
+        const auto* v = op->src[2];
+        if (q->ne[0] != 64 && q->ne[0] != 128)
+            return false;
+        for (const auto* t : {q, k, v, op}) {
+            if (ggml_nelements(t) > INT_MAX || t->ne[2] * t->ne[3] > 65535 ||
+                (reinterpret_cast<uintptr_t>(t->data) & 15))
+                return false;
+        }
+        return k->ne[0] * GGML_PAD(k->ne[1], 64) * k->ne[2] * k->ne[3] <= INT_MAX;
+    }
+
+    static ggml_status sage_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
+        *workspace = SageWorkspace(nodes[0]).size;
+        return GGML_STATUS_SUCCESS;
+    }
+
+    static void sage_destroy(void*, void*) {}
+
+    static ggml_status sage_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void* user_data) {
+        ggml_ext_gpu_context context;
+        if (!cuda_context(execution->backend, &context))
+            return GGML_STATUS_FAILED;
+        SDSageAttentionParams params;
+        if (!sd_ggml_sage_attn_describe(nodes[0], &params))
+            return GGML_STATUS_FAILED;
+        const auto mode = sage_mode(*static_cast<const SageDevice*>(user_data), params.mode);
+        auto* data      = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+        auto stream     = static_cast<cudaStream_t>(context.stream);
+        if (nodes[0]->src[0]->ne[0] == 64)
+            sage_attn<64>(stream, nodes[0], data, mode, params.scale);
+        else
+            sage_attn<128>(stream, nodes[0], data, mode, params.scale);
+        return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+
+    void register_sage_attention(ggml_backend_dev_t device) {
+        const auto info = cuda_device_info(device);
+        if (info.device < 0)
+            return;
+        int previous;
+        cudaDeviceProp properties;
+        cudaFuncAttributes attributes;
+        if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || cudaGetDevice(&previous) != cudaSuccess ||
+            cudaSetDevice(info.device) != cudaSuccess)
+            return;
+        const auto error = cudaFuncGetAttributes(&attributes, sage_output);
+        if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess)
+            return;
+        // PTX compiled for an older target retains that target's instruction guards after JIT.
+        static std::map<int, SageDevice> devices;
+        auto& data                = devices[info.device];
+        data                      = {properties.major * 10 + properties.minor, std::min(attributes.binaryVersion, attributes.ptxVersion)};
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.sage_attn.cuda.v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_sage_attn_op();
+        desc.priority             = 100;
+        desc.supports             = sage_supports;
+        desc.prepare              = sage_prepare;
+        desc.destroy              = sage_destroy;
+        desc.execute              = sage_execute;
+        desc.user_data            = &data;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
+}
+#else
+namespace sd_cuda {
+    void register_sage_attention(ggml_backend_dev_t) {}
+}
+#endif

src/core/cuda/sageattention/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2024 Jintao Zhang, Haofeng Huang
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

src/core/cuda/sageattention/README.md

@@ -0,0 +1,28 @@
+# SageAttention CUDA device code
+
+Adapted from [thu-ml/SageAttention](https://github.com/thu-ml/SageAttention),
+revision `d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5`, under Apache-2.0 (see LICENSE).
+
+`attn.cuh` contains the device kernel from
+`csrc/qattn/qk_int_sv_f16_cuda_sm80.cu`; `attn_fp8.cuh` adapts
+`csrc/qattn/qk_int_sv_f8_cuda_sm89.cuh`. The remaining headers come from `csrc/`
+and `csrc/qattn/attn_utils.cuh`. PyTorch includes and host wrappers were removed;
+include paths and guards were adjusted, device helpers were placed in the
+`sageattention` namespace, and `div_ceil` was renamed to avoid macro collisions.
+Unused `<cuda/pipeline>` includes were removed so multi-architecture builds can
+also compile targets below SM70; the backend rejects SageAttention below SM80.
+Predicated loads now zero unused FP16 V lanes: masked probabilities cannot
+suppress NaNs left in shared memory by previous kernels.
+
+The sd.cpp adapter in `../sage_attention.cu` supplies FP32 K smoothing and native CUDA
+per-thread INT8 quantization, matching the token grouping in upstream
+`sageattention/triton/quant_per_thread.py`. V is quantized per channel and
+transposed, padded with zeros, and permuted according to
+`csrc/fused/fused.cu`. Its FP8 range is 448 for SageAttention2 and 2.25 for
+SageAttention2++. Zero channels use a positive scale to avoid division by zero.
+SageAttention2 uses an FP32 instruction buffer; 2++ uses an FP16 instruction
+buffer added to FP32 accumulators. The SM80 compatibility path retains per-warp
+INT8 Q/K and FP16 V with FP32 accumulation.
+
+The adapter requests workspace and the current stream through public ggml extension APIs; it has no Python, PyTorch,
+or Triton dependency. The public output is converted back to FP32.

src/core/cuda/sageattention/attn.cuh

@@ -0,0 +1,592 @@
+/*
+ * Copyright (c) 2024 by SageAttention team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH
+#define SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH
+
+#include "attn_utils.cuh"
+#include "math.cuh"
+
+namespace sageattention {
+
+#define PACK_SIZE_QK 16  // as if it is int8
+#define PACK_SIZE_V 8    // fp16
+#define PACK_SIZE_O 8    // fp16
+
+// treat as if int8 tensor core
+#define MMA_QK_M 16
+#define MMA_QK_N 16
+#define MMA_QK_K 32
+
+// fp16 tensor core
+#define MMA_SV_M 16
+#define MMA_SV_N 16
+#define MMA_SV_K 16
+
+    template <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_mean = false>
+    __global__ void qk_int_sv_f16_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, half* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, DTypeOut* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
+        // Older CUDA headers omit BF16 intrinsics in pre-Ampere device passes.
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
+        // compile time check
+        static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4");
+        static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread");
+        static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread");
+        static_assert(std::is_same<DTypeSVAccum, float>::value || !use_inst_buffer, "use_inst_buffer only supports DTypeSVAccum as float");
+        static_assert(std::is_same<DTypeSVAccum, float>::value || std::is_same<DTypeSVAccum, half>::value, "DTypeSVAccum must be float or half");
+        static_assert(std::is_same<DTypeOut, half>::value || std::is_same<DTypeOut, nv_bfloat16>::value, "DTypeOut must be half or nv_bfloat16");
+        static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64");
+        static_assert(!fuse_v_mean || std::is_same<DTypeSVAccum, half>::value, "fuse_v_mean only supports half");
+        static_assert(CTA_Q / CTA_K <= 2);  // for efficient causal implementation
+
+        using DTypeOut2 = typename std::conditional<std::is_same<DTypeOut, half>::value, half2, nv_bfloat162>::type;
+
+        constexpr uint32_t num_warps_q        = CTA_Q / WARP_Q;
+        constexpr uint32_t num_warps_k        = CTA_K / WARP_K;
+        constexpr uint32_t num_warps          = num_warps_q * num_warps_k;
+        constexpr uint32_t num_tiles_q        = WARP_Q / MMA_QK_M;
+        constexpr uint32_t num_tiles_k        = WARP_K / MMA_QK_N;
+        constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K);
+        constexpr uint32_t num_tiles_v        = head_dim / MMA_SV_N;
+
+        constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2);
+        constexpr uint32_t O_SMEM_STRIDE  = head_dim;
+        constexpr uint32_t V_SMEM_STRIDE  = head_dim;
+
+        extern __shared__ int8_t smem[];
+
+        const uint32_t lane_id = get_lane_id();
+        const uint32_t warp_id = get_warp_id();
+
+        // maximize L2 hit rate
+        const uint32_t batch_id     = blockIdx.z;
+        const uint32_t bx           = blockIdx.x;
+        const uint32_t num_qo_heads = gridDim.y;
+        const uint32_t head_id      = blockIdx.y;
+
+        // transfer to base 2 instead of base e with better numerical efficiency
+        sm_scale *= math::log2e;
+
+        // RS holds the fragment of S
+        int32_t RS[num_tiles_q][num_tiles_k][8];
+        DTypeSVAccum RO[num_tiles_q][num_tiles_v][8];
+        float m[num_tiles_q][2];  // max
+        float d[num_tiles_q][2];  // denominator
+
+        uint32_t q_scale_idx, k_scale_idx;
+
+        if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
+            const uint32_t num_block_q = gridDim.x;
+            q_scale_idx                = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx;
+        } else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) {
+            const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
+            q_scale_idx                     = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>();
+        } else if constexpr (Q_GRAN == QuantGranularity::kPerThread) {
+            const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
+            q_scale_idx                     = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4;
+        }
+
+        if constexpr (K_GRAN == QuantGranularity::kPerBlock) {
+            const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K);
+            k_scale_idx                = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k;
+        } else if constexpr (K_GRAN == QuantGranularity::kPerWarp) {
+            const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
+            k_scale_idx                     = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>();
+        } else if constexpr (K_GRAN == QuantGranularity::kPerThread) {
+            const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
+            k_scale_idx                     = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4;
+        }
+
+        constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K)
+                                                                                                                                         : (CTA_K / WARP_K) * 4;
+
+        // initialize o, m, d
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 8; k++) {
+                        RO[fq][fv][k] = 0.0f;
+                    }
+                } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 4; k++) {
+                        ((int32_t*)RO[fq][fv])[k] = 0;
+                    }
+                }
+            }
+        }
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t k = 0; k < 2; k++) {
+                m[fq][k] = -5000000.0f;
+                d[fq][k] = 1.0f;
+            }
+        }
+
+        constexpr uint32_t K_smem_idx_offset = CTA_Q;
+        constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K;
+
+        constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B
+                                                                                                                    : SwizzleMode::k128B;
+        smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
+        smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
+        constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
+        smem_t<swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V> smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE);
+        constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
+        smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem);
+
+        constexpr uint32_t global_to_shared_line_lanes_QK          = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
+                                                                                                                         : 8;
+        constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
+                                                                                                                          : 4;
+        constexpr uint32_t global_to_shared_line_lanes_V           = (V_SMEM_STRIDE == 32) ? 4 : 8;
+        constexpr uint32_t global_to_shared_copy_lines_per_warp_V  = (V_SMEM_STRIDE == 32) ? 8 : 4;
+        constexpr uint32_t global_to_shared_line_lanes_O           = (O_SMEM_STRIDE == 32) ? 4 : 8;
+        constexpr uint32_t global_to_shared_copy_lines_per_warp_O  = (O_SMEM_STRIDE == 32) ? 8 : 4;
+
+        constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK);
+        constexpr uint32_t Q_smem_iters_col  = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK);
+        constexpr uint32_t K_smem_iters_col  = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK);
+        constexpr uint32_t V_smem_iters_row  = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V);
+        constexpr uint32_t V_smem_iters_col  = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_V);
+        constexpr uint32_t O_smem_iters_row  = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O);
+        constexpr uint32_t O_smem_iters_col  = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O);
+
+        int8_t* Q_lane_base_ptr     = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
+        int8_t* K_lane_base_ptr     = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
+        half* V_lane_base_ptr       = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V) * stride_seq_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
+        uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
+        uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
+        uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V);
+
+        uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id % 16, lane_id / 16);
+        uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2);
+        uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 16, lane_id / 16);
+
+        // for causal masking
+        uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+        uint32_t K_idx_lane_base = get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + 2 * (lane_id % 4);
+
+        // for loading
+        uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
+        uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
+        uint32_t V_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V;
+
+        const uint32_t num_iterations = SAGE_DIV_CEIL(
+            mask_mode == MaskMode::kCausal
+                ? min(kv_len, (bx + 1) * CTA_Q)
+                : kv_len,
+            CTA_K);
+
+        // load Q with predicate
+        load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, Q_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_Q>(
+            &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len);
+        cp_async::commit_group();
+        cp_async::wait_group<0>();
+        __syncthreads();
+
+        // for num_tiles_qk_inner = 1, we load all Qs in register
+        uint32_t RQ[num_tiles_q][4];
+        if constexpr (num_tiles_qk_inner == 1) {
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]);
+                Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma);
+            }
+        }
+
+        // load K with predicate
+        load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+            &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
+        cp_async::commit_group();
+
+        float q_scale = Q_scale[q_scale_idx];
+
+        float original_sm_scale = sm_scale;
+        float dequant_scale     = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset];
+
+        sm_scale = original_sm_scale * dequant_scale;
+
+        // load V with predicate
+        load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+            &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len);
+        cp_async::commit_group();
+
+        K_load_idx_lane_base += CTA_K;
+        V_load_idx_lane_base += CTA_K;
+
+#pragma unroll
+        for (uint32_t iter = 1; iter < num_iterations - 1; iter++) {
+            // ensure K is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            // compute QK^T
+            if constexpr (num_tiles_qk_inner == 1) {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_K, RS, RQ, K_smem_offset_mma);
+            } else {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+            }
+
+            float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+                for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 8; k++) {
+                        RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]);
+                    }
+                }
+            }
+
+            // do not apply causal mask and out of bound mask for these iterations
+            K_idx_lane_base += CTA_K;
+
+            if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, sm_scale);
+            } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, sm_scale);
+            }
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+                accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+            }
+
+            uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
+            RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+                accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
+            }
+
+            __syncthreads();
+
+            // load K
+            load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+                &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K);
+            cp_async::commit_group();
+
+            dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset];
+            sm_scale      = original_sm_scale * dequant_scale;
+
+            // ensure V is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            if constexpr (!use_inst_buffer) {
+                compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+                    smem_V, RS_f16, RO, d, V_smem_offset_mma);
+            } else {
+                compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+                    smem_V, RS_f16, RO, d, V_smem_offset_mma);
+            }
+
+            __syncthreads();
+            // load V
+            load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+                &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V);
+            cp_async::commit_group();
+            K_load_idx_lane_base += CTA_K;
+            V_load_idx_lane_base += CTA_K;
+        }
+
+        // second last iter, apply causal mask
+        if (num_iterations > 1) {
+            // ensure K is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            // compute QK^T
+            if constexpr (num_tiles_qk_inner == 1) {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_K, RS, RQ, K_smem_offset_mma);
+            } else {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+            }
+
+            float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+                for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 8; k++) {
+                        RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
+                    }
+                }
+            }
+
+            if constexpr (mask_mode == MaskMode::kCausal) {
+                apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+            }
+            // apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+            K_idx_lane_base += CTA_K;
+
+            if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
+            } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
+            }
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+                accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+            }
+
+            uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
+            RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+                accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
+            }
+
+            __syncthreads();
+
+            // load K with predicate
+            load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+                &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
+            cp_async::commit_group();
+
+            dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset];
+            sm_scale      = original_sm_scale * dequant_scale;
+
+            // ensure V is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            if constexpr (!use_inst_buffer) {
+                compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+                    smem_V, RS_f16, RO, d, V_smem_offset_mma);
+            } else {
+                compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+                    smem_V, RS_f16, RO, d, V_smem_offset_mma);
+            }
+
+            __syncthreads();
+            // load V with predicate
+            load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+                &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len);
+            cp_async::commit_group();
+            K_load_idx_lane_base += CTA_K;
+            V_load_idx_lane_base += CTA_K;
+        }
+
+        // last iter, apply causal mask and out of bound mask
+        {
+            // ensure K is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            // compute QK^T
+            if constexpr (num_tiles_qk_inner == 1) {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_K, RS, RQ, K_smem_offset_mma);
+            } else {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+            }
+
+            float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+                for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 8; k++) {
+                        RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
+                    }
+                }
+            }
+
+            if constexpr (mask_mode == MaskMode::kCausal) {
+                apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+            }
+            // check out of bound in the last iter
+            apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+            K_idx_lane_base += CTA_K;
+
+            if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
+            } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
+            }
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+                accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+            }
+
+            uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
+            RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+                accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
+            }
+
+            // ensure V is ready
+            cp_async::wait_group<0>();
+            __syncthreads();
+
+            if constexpr (!use_inst_buffer) {
+                compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+                    smem_V, RS_f16, RO, d, V_smem_offset_mma);
+            } else {
+                compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+                    smem_V, RS_f16, RO, d, V_smem_offset_mma);
+            }
+
+            __syncthreads();
+        }
+
+        // TODO: thread block sync mdo state for num_warps_k > 0
+
+        normalize_d<num_tiles_q, num_tiles_v, DenominatorAccumUnit>(RO, m, d);
+
+        // save the result
+        // if (get_warp_idx_k<num_warps_q, num_warps_k>() == 0)
+        // {
+
+        // convert half to bfloat16
+        if constexpr (std::is_same<DTypeSVAccum, half>::value && std::is_same<DTypeOut, nv_bfloat16>::value) {
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+                for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                    ((nv_bfloat162*)RO[fq][fv])[0] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[0]));
+                    ((nv_bfloat162*)RO[fq][fv])[1] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[1]));
+                    ((nv_bfloat162*)RO[fq][fv])[2] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[2]));
+                    ((nv_bfloat162*)RO[fq][fv])[3] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[3]));
+                }
+            }
+        }
+
+        // add v_mean
+        if constexpr (fuse_v_mean) {
+            DTypeOut2 v_mean[2];
+            DTypeOut* V_mean_lane_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + lane_id % 4 * 2;
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                v_mean[0] = *((DTypeOut2*)(V_mean_lane_ptr + fv * 16));
+                v_mean[1] = *((DTypeOut2*)(V_mean_lane_ptr + 8 + fv * 16));
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    ((DTypeOut2*)RO[fq][fv])[0] = __hadd2(((DTypeOut2*)RO[fq][fv])[0], v_mean[0]);
+                    ((DTypeOut2*)RO[fq][fv])[1] = __hadd2(((DTypeOut2*)RO[fq][fv])[1], v_mean[0]);
+                    ((DTypeOut2*)RO[fq][fv])[2] = __hadd2(((DTypeOut2*)RO[fq][fv])[2], v_mean[1]);
+                    ((DTypeOut2*)RO[fq][fv])[3] = __hadd2(((DTypeOut2*)RO[fq][fv])[3], v_mean[1]);
+                }
+            }
+        }
+
+        // save the result to shared memory
+        uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O));
+
+                if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                    // convert RO to half
+                    uint32_t RO_f16[4];
+#pragma unroll
+                    for (uint32_t k = 0; k < 4; k++) {
+                        if constexpr (std::is_same<DTypeOut, half>::value) {
+                            ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
+                        } else if constexpr (std::is_same<DTypeOut, nv_bfloat16>::value) {
+                            ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]);
+                        }
+                    }
+
+                    ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4]                                     = RO_f16[0];
+                    ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1];
+
+                    // ! permuted, make sure you know what you are doing
+                    ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4]                                     = RO_f16[2];
+                    ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3];
+                } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                    ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4]                                     = ((uint32_t*)RO[fq][fv])[0];
+                    ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[1];
+
+                    // ! permuted, make sure you know what you are doing
+                    ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4]                                     = ((uint32_t*)RO[fq][fv])[2];
+                    ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[3];
+                }
+            }
+        }
+
+        // ! do we need to sync here?
+        __syncwarp();
+
+        // shared memory to global memory
+        DTypeOut* O_lane_ptr          = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
+        uint32_t offset_O             = smem_O.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O);
+        uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O;
+
+#pragma unroll
+        for (uint32_t i = 0; i < O_smem_iters_col; i++) {
+#pragma unroll
+            for (uint32_t j = 0; j < O_smem_iters_row; j++) {
+                if (O_load_idx_lane_base < qo_len) {
+                    smem_O.store_128b(offset_O, O_lane_ptr);
+                }
+                O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O);
+                offset_O = smem_O.advance_offset_by_column<global_to_shared_line_lanes_O>(offset_O);
+            }
+
+            offset_O = smem_O.advance_offset_by_row<global_to_shared_copy_lines_per_warp_O>(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O));
+            O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O));
+            O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O;
+        }
+
+        if constexpr (return_lse) {
+            uint32_t lse_idx    = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>();
+            float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx;
+            uint32_t fq         = (lane_id % 4) / 2;
+            uint32_t k          = (lane_id % 4) % 2;
+
+            if (lse_idx < qo_len && (lane_id % 4) < 2 * num_tiles_q) {
+                lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]);
+            }
+        }
+
+        // }
+#else
+        RUNTIME_ASSERT("SageAttention requires SM80 or newer");
+#endif
+    }
+
+#undef PACK_SIZE_QK
+#undef PACK_SIZE_V
+#undef PACK_SIZE_O
+#undef MMA_QK_M
+#undef MMA_QK_N
+#undef MMA_QK_K
+#undef MMA_SV_M
+#undef MMA_SV_N
+#undef MMA_SV_K
+
+}  // namespace sageattention
+
+#endif

src/core/cuda/sageattention/attn_fp8.cuh

@@ -0,0 +1,622 @@
+/*
+ * Copyright (c) 2024 by SageAttention team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH
+#define SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH
+
+#include "attn_utils.cuh"
+#include "math.cuh"
+
+namespace sageattention {
+
+#define PACK_SIZE_QK 16  // as if it is int8
+#define PACK_SIZE_V 16   // fp8
+#define PACK_SIZE_O 8    // fp16
+
+// treat as if int8 tensor core
+#define MMA_QK_M 16
+#define MMA_QK_N 16
+#define MMA_QK_K 32
+
+// fp8 tensor core
+#define MMA_SV_M 16
+#define MMA_SV_N 16
+#define MMA_SV_K 32
+
+    template <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale = false, bool fuse_v_mean = false, bool use_pv_fp16_accu = false>
+    __global__ void qk_int_sv_f8_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, int8_t* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, float* __restrict__ V_scale, float* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
+        // Keep unsupported device passes from parsing architecture-specific intrinsics.
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 890
+        // compile time check
+        static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4");
+        static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread");
+        static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread");
+        static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64");
+        static_assert(std::is_same<DTypeSVAccum, float>::value, "DTypeSVAccum must be float, half is WIP");
+        static_assert(std::is_same<DTypeOut, half>::value || std::is_same<DTypeOut, nv_bfloat16>::value, "DTypeOut must be half or nv_bfloat16");
+        static_assert(CTA_K % 64 == 0);
+        static_assert(CTA_Q / CTA_K <= 2);  // for efficient causal implementation
+
+        constexpr uint32_t num_warps_q        = CTA_Q / WARP_Q;
+        constexpr uint32_t num_warps_k        = CTA_K / WARP_K;
+        constexpr uint32_t num_warps          = num_warps_q * num_warps_k;
+        constexpr uint32_t num_tiles_q        = WARP_Q / MMA_QK_M;
+        constexpr uint32_t num_tiles_k        = WARP_K / MMA_QK_N;
+        constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K);
+        constexpr uint32_t num_tiles_v        = head_dim / MMA_SV_N;
+
+        constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2);
+        constexpr uint32_t O_SMEM_STRIDE  = head_dim;
+        //                       for fp16: head_dim
+        constexpr uint32_t V_SMEM_STRIDE = CTA_K;
+
+        extern __shared__ int8_t smem[];
+
+        const uint32_t lane_id = get_lane_id();
+        const uint32_t warp_id = get_warp_id();
+
+        // maximize L2 hit rate
+        const uint32_t batch_id     = blockIdx.z;
+        const uint32_t bx           = blockIdx.x;
+        const uint32_t num_qo_heads = gridDim.y;
+        const uint32_t head_id      = blockIdx.y;
+
+        // transfer to base 2 instead of base e with better numerical efficiency
+        sm_scale *= math::log2e;
+
+        // RS holds the fragment of S
+        int32_t RS[num_tiles_q][num_tiles_k][8];
+        DTypeSVAccum RO[num_tiles_q][num_tiles_v][8];
+        float m[num_tiles_q][2];  // max
+        float d[num_tiles_q][2];  // denominator
+
+        uint32_t q_scale_idx, k_scale_idx;
+
+        if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
+            const uint32_t num_block_q = gridDim.x;
+            q_scale_idx                = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx;
+        } else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) {
+            const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
+            q_scale_idx                     = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>();
+        } else if constexpr (Q_GRAN == QuantGranularity::kPerThread) {
+            const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
+            q_scale_idx                     = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4;
+        }
+
+        if constexpr (K_GRAN == QuantGranularity::kPerBlock) {
+            const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K);
+            k_scale_idx                = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k;
+        } else if constexpr (K_GRAN == QuantGranularity::kPerWarp) {
+            const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
+            k_scale_idx                     = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>();
+        } else if constexpr (K_GRAN == QuantGranularity::kPerThread) {
+            const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
+            k_scale_idx                     = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4;
+        }
+
+        constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K)
+                                                                                                                                         : (CTA_K / WARP_K) * 4;
+
+        // initialize o, m, d
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 8; k++) {
+                        RO[fq][fv][k] = 0.0f;
+                    }
+                } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 4; k++) {
+                        ((int32_t*)RO[fq][fv])[k] = 0;
+                    }
+                }
+            }
+        }
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t k = 0; k < 2; k++) {
+                m[fq][k] = -5000000.0f;
+                d[fq][k] = 1.0f;
+            }
+        }
+
+        constexpr uint32_t K_smem_idx_offset = CTA_Q;
+        constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K;
+
+        constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B
+                                                                                                                    : SwizzleMode::k128B;
+        smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
+        smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
+        //                                             for fp16: 32
+        constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B;
+        smem_t<swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V> smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE);
+        constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
+        smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem);
+
+        constexpr uint32_t global_to_shared_line_lanes_QK          = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
+                                                                                                                         : 8;
+        constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
+                                                                                                                          : 4;
+        //                                                         for fp16: 32
+        constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 64) ? 4 : 8;
+        //                                                                  for fp16: 32
+        constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 64) ? 8 : 4;
+        constexpr uint32_t global_to_shared_line_lanes_O          = (O_SMEM_STRIDE == 32) ? 4 : 8;
+        constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4;
+
+        constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK);
+        constexpr uint32_t Q_smem_iters_col  = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK);
+        constexpr uint32_t K_smem_iters_col  = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK);
+        constexpr uint32_t V_smem_iters_row  = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V);
+        //                          for fp16: CTA_K
+        constexpr uint32_t V_smem_iters_col = head_dim / (num_warps * global_to_shared_copy_lines_per_warp_V);
+        constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O);
+        constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O);
+
+        int8_t* Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
+        int8_t* K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
+        //                                                                for fp16: CTA_K / num_warps * warp_id * stride_seq_v + lane_id / global_to_shared_line_lanes_V * stride_seq_v
+        int8_t* V_lane_base_ptr     = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + head_dim / num_warps * warp_id * stride_d_v + lane_id / global_to_shared_line_lanes_V * stride_d_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
+        uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
+        uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
+        uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V);
+
+        uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id % 16, lane_id / 16);
+        uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2);
+        // for fp 16:
+        // uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 16, lane_id / 16);
+        uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(lane_id % 8 + (lane_id / 16) * 8, get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K / PACK_SIZE_V + (lane_id / 8) % 2);
+
+        // for causal masking
+        uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+        uint32_t K_idx_lane_base = get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + 2 * (lane_id % 4);
+
+        // for loading
+        uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
+        uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
+
+        const uint32_t num_iterations = SAGE_DIV_CEIL(
+            mask_mode == MaskMode::kCausal
+                ? min(kv_len, (bx + 1) * CTA_Q)
+                : kv_len,
+            CTA_K);
+
+        // load Q with predicate
+        load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, Q_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_Q>(
+            &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len);
+        cp_async::commit_group();
+        cp_async::wait_group<0>();
+        __syncthreads();
+
+        // for num_tiles_qk_inner = 1, we load all Qs in register
+        uint32_t RQ[num_tiles_q][4];
+        if constexpr (num_tiles_qk_inner == 1) {
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]);
+                Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma);
+            }
+        }
+
+        // load K with predicate
+        load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+            &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
+        cp_async::commit_group();
+
+        float q_scale = Q_scale[q_scale_idx];
+
+        float original_sm_scale = sm_scale;
+        float dequant_scale     = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset];
+
+        sm_scale = original_sm_scale * dequant_scale;
+
+        // load V
+        // ! we assume that V is padded. If not, there might be illegal memory access or nan issue.
+        // for fp16:
+        // load_global_to_share                stride_seq_v
+        load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+            &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V);
+        cp_async::commit_group();
+
+        K_load_idx_lane_base += CTA_K;
+
+#pragma unroll
+        for (uint32_t iter = 1; iter < num_iterations - 1; iter++) {
+            // ensure K is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            // compute QK^T
+            if constexpr (num_tiles_qk_inner == 1) {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_K, RS, RQ, K_smem_offset_mma);
+            } else {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+            }
+            float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+                for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 8; k++) {
+                        RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]);
+                    }
+                }
+            }
+
+            K_idx_lane_base += CTA_K;
+
+            if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, sm_scale);
+            } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, sm_scale);
+            }
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+                accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+            }
+
+            uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
+            RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+                accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+            }
+
+            __syncthreads();
+
+            // load K without predicate
+            load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+                &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K);
+            cp_async::commit_group();
+
+            dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset];
+            sm_scale      = original_sm_scale * dequant_scale;
+
+            // ensure V is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            // for fp16:
+            // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+            //   smem_V, RS_f16, RO, d, V_smem_offset_mma);
+            if constexpr (!use_inst_buffer) {
+                compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+                    smem_V, RS_f8, RO, d);
+            } else {
+                if constexpr (!use_pv_fp16_accu) {
+                    compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+                        smem_V, RS_f8, RO, d);
+                } else {
+                    compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+                        smem_V, RS_f8, RO, d);
+                }
+            }
+            __syncthreads();
+            // load V
+            // for fp16:
+            // load_global_to_share                stride_seq_v
+            load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+                &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V);
+            cp_async::commit_group();
+
+            K_load_idx_lane_base += CTA_K;
+        }
+
+        // second last iter, apply causal mask
+        if (num_iterations > 1) {
+            // ensure K is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            // compute QK^T
+            if constexpr (num_tiles_qk_inner == 1) {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_K, RS, RQ, K_smem_offset_mma);
+            } else {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+            }
+
+            float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+                for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 8; k++) {
+                        RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
+                    }
+                }
+            }
+
+            if constexpr (mask_mode == MaskMode::kCausal) {
+                apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+            }
+            // apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+            K_idx_lane_base += CTA_K;
+
+            if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
+            } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
+            }
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+                accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+            }
+
+            uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
+            RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+                accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+            }
+
+            __syncthreads();
+
+            // load K with predicate
+            load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+                &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
+            cp_async::commit_group();
+
+            dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset];
+            sm_scale      = original_sm_scale * dequant_scale;
+
+            // ensure V is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            // for fp16:
+            // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+            //   smem_V, RS_f16, RO, d, V_smem_offset_mma);
+            if constexpr (!use_inst_buffer) {
+                compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+                    smem_V, RS_f8, RO, d);
+            } else {
+                if constexpr (!use_pv_fp16_accu) {
+                    compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+                        smem_V, RS_f8, RO, d);
+                } else {
+                    compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+                        smem_V, RS_f8, RO, d);
+                }
+            }
+
+            __syncthreads();
+            // load V
+            // for fp16:
+            // load_global_to_share                stride_seq_v
+            load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+                &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V);
+            cp_async::commit_group();
+            K_load_idx_lane_base += CTA_K;
+        }
+
+        // last iter, apply causal mask and out of bound mask
+        {
+            // ensure K is ready
+            cp_async::wait_group<1>();
+            __syncthreads();
+
+            // compute QK^T
+            if constexpr (num_tiles_qk_inner == 1) {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_K, RS, RQ, K_smem_offset_mma);
+            } else {
+                compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+                    smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+            }
+
+            float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+                for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+                    for (uint32_t k = 0; k < 8; k++) {
+                        RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
+                    }
+                }
+            }
+
+            if constexpr (mask_mode == MaskMode::kCausal) {
+                apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+            }
+            apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+            K_idx_lane_base += CTA_K;
+
+            if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
+            } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
+            }
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+                accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+            }
+
+            uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
+            RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+
+            if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+                accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+            }
+
+            // ensure V is ready
+            cp_async::wait_group<0>();
+            __syncthreads();
+
+            // for fp16:
+            // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+            //   smem_V, RS_f16, RO, d, V_smem_offset_mma);
+            if constexpr (!use_inst_buffer) {
+                compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+                    smem_V, RS_f8, RO, d);
+            } else {
+                if constexpr (!use_pv_fp16_accu) {
+                    compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+                        smem_V, RS_f8, RO, d);
+                } else {
+                    compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+                        smem_V, RS_f8, RO, d);
+                }
+            }
+
+            __syncthreads();
+        }
+
+        // TODO: thread block sync mdo state for num_warps_k > 0. Then only one thread block needs to do the final saving.
+
+        normalize_d<num_tiles_q, num_tiles_v, ComputeUnit::kCudaCore>(RO, m, d);
+
+        // ! here we just implement the case for fp32 acumulation
+        if constexpr (fuse_v_scale) {
+            float v_scale[4];
+            float* V_scale_base_ptr = V_scale + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2;
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                ((float2*)v_scale)[0] = *((float2*)(V_scale_base_ptr + fv * 16));
+                ((float2*)v_scale)[1] = *((float2*)(V_scale_base_ptr + fv * 16 + 8));
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    RO[fq][fv][0] *= v_scale[0];
+                    RO[fq][fv][1] *= v_scale[1];
+                    RO[fq][fv][2] *= v_scale[0];
+                    RO[fq][fv][3] *= v_scale[1];
+                    RO[fq][fv][4] *= v_scale[2];
+                    RO[fq][fv][5] *= v_scale[3];
+                    RO[fq][fv][6] *= v_scale[2];
+                    RO[fq][fv][7] *= v_scale[3];
+                }
+            }
+        }
+
+        if constexpr (fuse_v_mean) {
+            float v_mean[4];
+            float* V_mean_base_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2;
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                ((float2*)v_mean)[0] = *((float2*)(V_mean_base_ptr + fv * 16));
+                ((float2*)v_mean)[1] = *((float2*)(V_mean_base_ptr + fv * 16 + 8));
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    RO[fq][fv][0] += v_mean[0];
+                    RO[fq][fv][1] += v_mean[1];
+                    RO[fq][fv][2] += v_mean[0];
+                    RO[fq][fv][3] += v_mean[1];
+                    RO[fq][fv][4] += v_mean[2];
+                    RO[fq][fv][5] += v_mean[3];
+                    RO[fq][fv][6] += v_mean[2];
+                    RO[fq][fv][7] += v_mean[3];
+                }
+            }
+        }
+
+        // save the result to shared memory
+        uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O));
+
+                if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                    // convert RO to half
+                    uint32_t RO_f16[4];
+#pragma unroll
+                    for (uint32_t k = 0; k < 4; k++) {
+                        if constexpr (std::is_same<DTypeOut, half>::value) {
+                            ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
+                        } else {
+                            ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]);
+                        }
+                    }
+
+                    ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4]                                     = RO_f16[0];
+                    ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1];
+
+                    offset_O                                                                               = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O) + 1);
+                    ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4]                                     = RO_f16[2];
+                    ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3];
+                } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                    // TODO: not implement
+                }
+            }
+        }
+
+        // ! do we need to sync here?
+        __syncwarp();
+
+        // shared memory to global memory
+        DTypeOut* O_lane_ptr          = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
+        uint32_t offset_O             = smem_O.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O);
+        uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O;
+
+#pragma unroll
+        for (uint32_t i = 0; i < O_smem_iters_col; i++) {
+#pragma unroll
+            for (uint32_t j = 0; j < O_smem_iters_row; j++) {
+                if (O_load_idx_lane_base < qo_len) {
+                    smem_O.store_128b(offset_O, O_lane_ptr);
+                }
+                O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O);
+                offset_O = smem_O.advance_offset_by_column<global_to_shared_line_lanes_O>(offset_O);
+            }
+
+            offset_O = smem_O.advance_offset_by_row<global_to_shared_copy_lines_per_warp_O>(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O));
+            O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O));
+            O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O;
+        }
+
+        if constexpr (return_lse) {
+            // ! this only works for num_tiles_q = 2
+            uint32_t lse_idx    = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>();
+            float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx;
+            uint32_t fq         = (lane_id % 4) / 2;
+            uint32_t k          = (lane_id % 4) % 2;
+
+            if (lse_idx < qo_len) {
+                lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]);
+            }
+        }
+#else
+        RUNTIME_ASSERT("FP8 SageAttention requires SM89 or newer");
+#endif
+    }
+
+#undef PACK_SIZE_QK
+#undef PACK_SIZE_V
+#undef PACK_SIZE_O
+#undef MMA_QK_M
+#undef MMA_QK_N
+#undef MMA_QK_K
+#undef MMA_SV_M
+#undef MMA_SV_N
+#undef MMA_SV_K
+
+}  // namespace sageattention
+
+#endif

src/core/cuda/sageattention/attn_utils.cuh

@@ -0,0 +1,814 @@
+/*
+ * Copyright (c) 2024 by SageAttention team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_UTILS_CUH
+#define SD_CORE_CUDA_SAGEATTENTION_ATTN_UTILS_CUH
+#include <cuda_fp16.h>
+#include <cuda_pipeline_primitives.h>
+
+#include "cp_async.cuh"
+#include "math.cuh"
+#include "mma.cuh"
+#include "numeric_conversion.cuh"
+#include "permuted_smem.cuh"
+
+namespace sageattention {
+
+#define WARP_SIZE 32
+
+#define S_FP8_OFFSET 8.807f
+#define S_FP8_OFFSET_EXP 6680.8477f
+#define S_FP8_OFFSET_EXP_INV 0.0022326917f
+
+#define SAGE_DIV_CEIL(M, N) (((M) + (N) - 1) / (N))
+
+    enum class MaskMode {
+        kNone   = 0,
+        kCausal = 1,
+    };
+
+    enum class DataType {
+        kHalf,
+        kInt8,
+        kInt4,
+        kE4M3,
+        kE5M2,
+    };
+
+    enum class QuantGranularity {
+        kPerTensor = 0,
+        kPerBlock  = 1,
+        kPerWarp   = 2,
+        kPerThread = 3,
+    };
+
+    enum class ComputeUnit {
+        kTensorCore,
+        kCudaCore,
+    };
+
+    __device__ __forceinline__ uint32_t get_warp_id() {
+        return threadIdx.y;
+    }
+
+    __device__ __forceinline__ uint32_t get_lane_id() {
+        return threadIdx.x;
+    }
+
+    template <uint32_t num_warps_q, uint32_t num_warps_k>
+    __device__ __forceinline__ uint32_t get_warp_idx_q() {
+        return get_warp_id() / num_warps_k;
+    }
+
+    template <uint32_t num_warps_q, uint32_t num_warps_k>
+    __device__ __forceinline__ uint32_t get_warp_idx_k() {
+        return get_warp_id() % num_warps_k;
+    }
+
+    template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
+    __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem) {
+        static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
+        static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value);
+
+        constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16;
+
+#pragma unroll
+        for (uint32_t i = 0; i < smem_iters_col; i++) {
+#pragma unroll
+            for (uint32_t j = 0; j < smem_iters_row; j++) {
+                smem.load_128b_async(smem_offset, *lane_ptr);
+                *lane_ptr += (global_to_shared_line_lanes * pack_size);
+                smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
+            }
+
+            smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(smem_offset - (smem_iters_row * global_to_shared_line_lanes));
+            *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size));
+        }
+        smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride);
+        *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
+    }
+
+    // with predicate
+    template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
+    __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem, uint32_t base_idx, uint32_t max_len) {
+        static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
+        static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value);
+
+        constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16;
+
+#pragma unroll
+        for (uint32_t i = 0; i < smem_iters_col; i++) {
+#pragma unroll
+            for (uint32_t j = 0; j < smem_iters_row; j++) {
+                // Masked probabilities still propagate NaNs from stale V lanes through 0 * NaN.
+                smem.load_128b_async<std::is_same<T, half>::value ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill>(smem_offset, *lane_ptr, base_idx < max_len);
+                *lane_ptr += (global_to_shared_line_lanes * pack_size);
+                smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
+            }
+
+            smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(smem_offset - (smem_iters_row * global_to_shared_line_lanes));
+            *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size));
+            base_idx += global_to_shared_copy_lines_per_warp_per_iter;
+        }
+        smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride);
+        *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
+    }
+
+    template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA>
+    __device__ __forceinline__ void load_fp8_V_global_to_share(int8_t** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem) {
+        static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
+        constexpr uint32_t pack_size_fp8 = 16;
+
+#pragma unroll
+        for (uint32_t i = 0; i < smem_iters_col; i++) {
+#pragma unroll
+            for (uint32_t j = 0; j < smem_iters_row; j++) {
+                smem.load_128b_async(smem_offset, *lane_ptr);
+                *lane_ptr += (global_to_shared_line_lanes * pack_size_fp8);
+                smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
+            }
+
+            smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(smem_offset - (smem_iters_row * global_to_shared_line_lanes));
+            *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size_fp8));
+        }
+        smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride);
+        // for QK: *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
+        *lane_ptr += CTA;  // ! prevent underflow
+        *lane_ptr -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
+    }
+
+    template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
+    __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride>& smem_Q, const smem_t<swizzle_mode, stride>& smem_K, int32_t RS[][num_tiles_k][8], uint32_t& offset_Q, uint32_t& offset_K) {
+        static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4);
+
+        uint32_t RQ[num_tiles_q][4];
+        uint32_t RK[4];
+
+        // the first iteration, mma mode is kInit
+#pragma unroll
+        for (uint32_t iter = 0; iter < 1; iter++) {
+            // load RQ
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]);
+                offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q);
+            }
+            // ! using permutation invariance
+            offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter);
+
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+                // load RK
+                smem_K.ldmatrix_m8n8x4(offset_K, RK);
+                offset_K = smem_K.advance_offset_by_row<16>(offset_K);
+
+                // mma
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    if constexpr (DTypeQK == DataType::kInt8) {
+                        mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+                    } else if constexpr (DTypeQK == DataType::kInt4) {
+                        mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+                    }
+                }
+            }
+            offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter);
+        }
+
+        // following iteration, mma mode is kInplace
+#pragma unroll
+        for (uint32_t iter = 1; iter < num_tiles_qk_inner; iter++) {
+            // load RQ
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]);
+                offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q);
+            }
+            offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter);
+
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+                // load RK
+                smem_K.ldmatrix_m8n8x4(offset_K, RK);
+                offset_K = smem_K.advance_offset_by_row<16>(offset_K);
+
+                // mma
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    if constexpr (DTypeQK == DataType::kInt8) {
+                        mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK);
+                    } else if constexpr (DTypeQK == DataType::kInt4) {
+                        mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK);
+                    }
+                }
+            }
+            offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter);
+        }
+
+        offset_Q -= (2 * num_tiles_qk_inner);
+        offset_K -= (2 * num_tiles_qk_inner);
+    }
+
+    // for case when num_tiles_qk_inner = 1
+    template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
+    __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride>& smem_K, int32_t RS[][num_tiles_k][8], uint32_t RQ[][4], uint32_t offset_K) {
+        static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4);
+        static_assert(num_tiles_qk_inner == 1);
+
+        uint32_t RK[4];
+
+        // mma mode is kInit
+#pragma unroll
+        for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+            // load RK
+            smem_K.ldmatrix_m8n8x4(offset_K, RK);
+            offset_K = smem_K.advance_offset_by_row<16>(offset_K);
+
+            // mma
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                if constexpr (DTypeQK == DataType::kInt8) {
+                    mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+                } else if constexpr (DTypeQK == DataType::kInt4) {
+                    mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+                }
+            }
+        }
+    }
+
+    template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
+    __device__ __forceinline__ void apply_causal_mask(const uint32_t& Q_idx_lane_base, const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8]) {
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+                for (uint32_t k = 0; k < 8; k++) {
+                    const uint32_t q_idx       = Q_idx_lane_base + fq * 16 + 8 * ((k % 4) / 2);
+                    const uint32_t kv_idx      = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2;
+                    const bool out_of_boundary = (kv_idx > q_idx);
+
+                    if constexpr (std::is_same<DTypeQKAccum, float>::value) {
+                        RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
+                    } else if constexpr (std::is_same<DTypeQKAccum, half>::value) {
+                        RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
+                    }
+                }
+            }
+        }
+    }
+
+    template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
+    __device__ __forceinline__ void apply_out_of_bound_mask(const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8], const uint32_t& kv_len) {
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+                for (uint32_t k = 0; k < 8; k++) {
+                    const uint32_t kv_idx      = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2;
+                    const bool out_of_boundary = (kv_idx >= kv_len);
+
+                    if constexpr (std::is_same<DTypeQKAccum, float>::value) {
+                        RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
+                    } else if constexpr (std::is_same<DTypeQKAccum, half>::value) {
+                        RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
+                    }
+                }
+            }
+        }
+    }
+
+    // for DTypeQKAccum float
+    template <uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, bool use_half_o_scale, bool exp_offset, bool fuse_scale = false, typename DTypeSVAccum>
+    __device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAccum RO[][num_tiles_v][8], float m[][2], float d[][2], const float& sm_scale) {
+        static_assert(std::is_same<DTypeSVAccum, half>::value || (!use_half_o_scale));
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t k = 0; k < 2; k++) {
+                // assign the smallest value possible
+                float m_prev = m[fq][k];
+                float m_temp = -5000000.0f;
+#pragma unroll
+                for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+                    float m_local = max(max(RS[fq][fk][k * 2 + 0], RS[fq][fk][k * 2 + 1]),
+                                        max(RS[fq][fk][k * 2 + 4], RS[fq][fk][k * 2 + 5]));
+                    m_temp        = max(m_temp, m_local);
+                }
+
+                if constexpr (!fuse_scale) {
+                    if constexpr (exp_offset) {
+                        m_temp = fmaf(m_temp, sm_scale, -S_FP8_OFFSET);
+                    } else {
+                        m_temp *= sm_scale;
+                    }
+                } else if constexpr (exp_offset) {
+                    m_temp += (-S_FP8_OFFSET);
+                }
+
+                // exchange element with the 4 threads in the row
+                m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x1));  // 0 exchange with 1, 2 exchange with 3
+                m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x2));  // 0 exchange with 2, 1 exchange with 3
+
+                m[fq][k] = max(m[fq][k], m_temp);
+
+                float o_scale = math::ptx_exp2(m_prev - m[fq][k]);
+
+                // update denominator
+                d[fq][k] *= o_scale;
+
+                half2 o_scale2;
+                if constexpr (use_half_o_scale) {
+                    o_scale2 = __floats2half2_rn(o_scale, o_scale);
+                }
+
+                // update RO
+#pragma unroll
+                for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                    if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                        RO[fq][fv][k * 2 + 0] *= o_scale;
+                        RO[fq][fv][k * 2 + 1] *= o_scale;
+                        RO[fq][fv][k * 2 + 4] *= o_scale;
+                        RO[fq][fv][k * 2 + 5] *= o_scale;
+                    } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                        if constexpr (use_half_o_scale) {
+                            ((half2*)RO[fq][fv])[k]     = __hmul2(((half2*)RO[fq][fv])[k], o_scale2);
+                            ((half2*)RO[fq][fv])[k + 2] = __hmul2(((half2*)RO[fq][fv])[k + 2], o_scale2);
+                        } else {
+                            RO[fq][fv][k * 2 + 0] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 0]) * o_scale);
+                            RO[fq][fv][k * 2 + 1] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 1]) * o_scale);
+                            RO[fq][fv][k * 2 + 4] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 4]) * o_scale);
+                            RO[fq][fv][k * 2 + 5] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 5]) * o_scale);
+                        }
+                    }
+                }
+
+                // raise RS to exponent
+                float negative_m = -m[fq][k];
+#pragma unroll
+                for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+                    if constexpr (fuse_scale) {
+                        RS[fq][fk][k * 2 + 0] = math::ptx_exp2(RS[fq][fk][k * 2 + 0] + negative_m);
+                        RS[fq][fk][k * 2 + 1] = math::ptx_exp2(RS[fq][fk][k * 2 + 1] + negative_m);
+                        RS[fq][fk][k * 2 + 4] = math::ptx_exp2(RS[fq][fk][k * 2 + 4] + negative_m);
+                        RS[fq][fk][k * 2 + 5] = math::ptx_exp2(RS[fq][fk][k * 2 + 5] + negative_m);
+                    } else {
+                        RS[fq][fk][k * 2 + 0] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 0], sm_scale, negative_m));
+                        RS[fq][fk][k * 2 + 1] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 1], sm_scale, negative_m));
+                        RS[fq][fk][k * 2 + 4] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 4], sm_scale, negative_m));
+                        RS[fq][fk][k * 2 + 5] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 5], sm_scale, negative_m));
+                    }
+                }
+            }
+        }
+    }
+
+    template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename T>
+    __device__ __forceinline__ void RS_32_to_16(T RS[][num_tiles_k][8], uint32_t RS_16[][num_tiles_k][4]) {
+        static_assert(sizeof(T) == 4);
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+                ((half2*)RS_16[fq][fk])[0] = __float22half2_rn(((float2*)RS[fq][fk])[0]);
+                ((half2*)RS_16[fq][fk])[1] = __float22half2_rn(((float2*)RS[fq][fk])[1]);
+                ((half2*)RS_16[fq][fk])[2] = __float22half2_rn(((float2*)RS[fq][fk])[2]);
+                ((half2*)RS_16[fq][fk])[3] = __float22half2_rn(((float2*)RS[fq][fk])[3]);
+            }
+        }
+    }
+
+    template <uint32_t num_tiles_q, uint32_t num_tiles_k>
+    __device__ __forceinline__ void RS_32_to_8(float RS[][num_tiles_k][8], uint32_t RS_8[][num_tiles_k / 2][4]) {
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+                floatx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 4);
+                floatx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 2, RS[fq][fk * 2 + 0] + 6);
+                floatx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 4);
+                floatx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 2, RS[fq][fk * 2 + 1] + 6);
+            }
+        }
+    }
+
+    template <uint32_t num_tiles_q, uint32_t num_tiles_k>
+    __device__ __forceinline__ void RS_16_to_8(uint32_t RS[][num_tiles_k][4], uint32_t RS_8[][num_tiles_k / 2][4]) {
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+                halfx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2);
+                halfx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3);
+                halfx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2);
+                halfx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3);
+            }
+        }
+    }
+
+    template <uint32_t num_tiles_q, uint32_t num_tiles_k>
+    __device__ __forceinline__ void RS_8_to_16(uint32_t RS_8[][num_tiles_k / 2][4], uint32_t RS[][num_tiles_k][4]) {
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+                e4m3x4_to_halfx4(RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2, RS_8[fq][fk]);
+                e4m3x4_to_halfx4(RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3, RS_8[fq][fk] + 1);
+                e4m3x4_to_halfx4(RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2, RS_8[fq][fk] + 2);
+                e4m3x4_to_halfx4(RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3, RS_8[fq][fk] + 3);
+            }
+        }
+    }
+
+    template <uint32_t num_tiles_q, uint32_t num_tiles_k, ComputeUnit compute_unit = ComputeUnit::kTensorCore, typename T>
+    __device__ __forceinline__ void accumulate_d(T RS[][num_tiles_k][(compute_unit == ComputeUnit::kTensorCore) ? 4 : 8], float d[][2]) {
+        // for compute unit cuda core, RS is float
+        // for compute unit tensor core, RS is packed half
+        static_assert((std::is_same<T, float>::value && compute_unit == ComputeUnit::kCudaCore) ||
+                      (std::is_same<T, uint32_t>::value && compute_unit == ComputeUnit::kTensorCore));
+
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+                if constexpr (compute_unit == ComputeUnit::kTensorCore) {
+                    // full accumulate with tensor core
+                    mma::rowsum_f16f16f32(d[fq], (uint32_t*)(RS[fq][fk]));
+                } else if constexpr (compute_unit == ComputeUnit::kCudaCore) {
+                    // partial accumulate with cuda core
+                    d[fq][0] += RS[fq][fk][0] + RS[fq][fk][1] + RS[fq][fk][4] + RS[fq][fk][5];
+                    d[fq][1] += RS[fq][fk][2] + RS[fq][fk][3] + RS[fq][fk][6] + RS[fq][fk][7];
+                }
+            }
+        }
+    }
+
+    template <uint32_t num_tiles_q, uint32_t num_tiles_k>
+    __device__ __forceinline__ void accumulate_d_f8(uint32_t RS[][num_tiles_k / 2][4], float d[][2]) {
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+                mma::rowsum_f8f8f32(d[fq], RS[fq][fk]);
+            }
+        }
+    }
+
+    template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+    __device__ __forceinline__ void compute_fp16_sv(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f16[][num_tiles_k][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
+        uint32_t smem_V_row_base = get_warp_idx_k<num_warps_q, num_warps_k>() * (num_tiles_k * 16) + get_lane_id() % 16;
+        uint32_t smem_V_col_base = get_lane_id() / 16;
+#pragma unroll
+        for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                // load RV
+                uint32_t RV[4];
+                uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fk * 16, smem_V_col_base + fv * 2);
+                smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                        mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV);
+                    } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                        mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV);
+                    }
+                }
+            }
+        }
+    }
+
+    template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width = 4, typename T, typename DTypeSVAccum>
+    __device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t<swizzle_mode, stride>& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) {
+        static_assert(sizeof(T) == 4);
+
+        // ! be sure you know what you are doing
+#pragma unroll
+        for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                // load RV
+                uint32_t RV[4];
+                smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                        mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
+                    } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                        mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
+                    }
+                }
+
+                offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv);
+            }
+            offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v));
+        }
+
+        // make offset_V their original value
+        offset_V -= (16 * num_tiles_k * stride);
+    }
+
+    template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width = 4, typename T, typename DTypeSVAccum>
+    __device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<swizzle_mode, stride>& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) {
+        static_assert(sizeof(T) == 4);
+        static_assert(std::is_same<DTypeSVAccum, float>::value);
+
+        uint32_t RO_inst_buf[num_tiles_q][num_tiles_v][4];
+
+        // ! be sure you know what you are doing
+#pragma unroll
+        for (uint32_t fk = 0; fk < 1; fk++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                // load RV
+                uint32_t RV[4];
+                smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    {
+                        mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInit>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
+                    }
+                }
+
+                offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv);
+            }
+            offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v));
+        }
+
+#pragma unroll
+        for (uint32_t fk = 1; fk < num_tiles_k; fk++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                // load RV
+                uint32_t RV[4];
+                smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    {
+                        mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInplaceUpdate>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
+                    }
+                }
+
+                offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv);
+            }
+            offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v));
+        }
+
+        // accumulate into RO
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                RO[fq][fv][0] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].x);
+                RO[fq][fv][1] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].y);
+                RO[fq][fv][2] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].x);
+                RO[fq][fv][3] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].y);
+                RO[fq][fv][4] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].x);
+                RO[fq][fv][5] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].y);
+                RO[fq][fv][6] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].x);
+                RO[fq][fv][7] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].y);
+            }
+        }
+
+        // make offset_V their original value
+        offset_V -= (16 * num_tiles_k * stride);
+    }
+
+    template <uint32_t num_tiles_q, uint32_t num_tiles_v,
+              ComputeUnit compute_unit = ComputeUnit::kTensorCore,  // compute unit for accumulate_d
+              typename DTypeQKAccum,
+              typename DTypeSVAccum>
+    __device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], DTypeQKAccum m[][2], float d[][2]) {
+        if constexpr (compute_unit == ComputeUnit::kCudaCore) {
+            // accumulate_d performs partial accumulation with cuda core
+            // aggregate d
+#pragma unroll
+            for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+                for (uint32_t k = 0; k < 2; k++) {
+                    d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x1);  // sum 0 and 1, 2 and 3
+                    d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x2);  // sum 0 and 2, 1 and 3
+                }
+            }
+        }
+
+        // divide O by d
+        float d_rcp[num_tiles_q][2];
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t k = 0; k < 2; k++) {
+                // TODO: check m to prevent nan
+                d_rcp[fq][k] = math::ptx_rcp(d[fq][k]);
+            }
+        }
+
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+#pragma unroll
+                for (uint32_t k = 0; k < 8; k++) {
+                    if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                        RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2];
+                    } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                        RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]);
+                    }
+                }
+            }
+        }
+    }
+
+    template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+    __device__ __forceinline__ void compute_fp8_sv(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
+        uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
+        // uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
+        uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
+#pragma unroll
+        for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+            uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                // load RV
+                uint32_t RV[4];
+                // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+                smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                        mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV);
+                    } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                        // ! Not Implemented
+                    }
+                }
+                offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+            }
+        }
+    }
+
+    template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+    __device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
+        uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
+        // uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
+        uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
+
+        float RO_inst_buf[num_tiles_q][num_tiles_v][8];
+
+#pragma unroll
+        for (uint32_t fk = 0; fk < 1; fk++) {
+            uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                // load RV
+                uint32_t RV[4];
+                // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+                smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                        mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+                    } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                        // ! Not Implemented
+                    }
+                }
+                offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+            }
+        }
+
+#pragma unroll
+        for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) {
+            uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                // load RV
+                uint32_t RV[4];
+                // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+                smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                        mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+                    } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                        // ! Not Implemented
+                    }
+                }
+                offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+            }
+        }
+
+#pragma unroll
+        for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                RO[fq][fv][0] += RO_inst_buf[fq][fv][0];
+                RO[fq][fv][1] += RO_inst_buf[fq][fv][1];
+                RO[fq][fv][2] += RO_inst_buf[fq][fv][2];
+                RO[fq][fv][3] += RO_inst_buf[fq][fv][3];
+                RO[fq][fv][4] += RO_inst_buf[fq][fv][4];
+                RO[fq][fv][5] += RO_inst_buf[fq][fv][5];
+                RO[fq][fv][6] += RO_inst_buf[fq][fv][6];
+                RO[fq][fv][7] += RO_inst_buf[fq][fv][7];
+            }
+        }
+    }
+
+    template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+    __device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
+        uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
+        // uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
+        uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
+
+        uint32_t RO_int32[num_tiles_q][num_tiles_v][4];
+
+#pragma unroll
+        for (uint32_t fk = 0; fk < 1; fk++) {
+            uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                // load RV
+                uint32_t RV[4];
+                // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+                smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                        // mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+                        mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInit>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
+                    } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                        // ! Not Implemented
+                    }
+                }
+                offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+            }
+        }
+
+#pragma unroll
+        for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) {
+            uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+                // load RV
+                uint32_t RV[4];
+                // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+                smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+                for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+                    if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+                        // mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+                        mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInplaceUpdate>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
+                    } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+                        // ! Not Implemented
+                    }
+                }
+                offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+            }
+        }
+        float RO_tmp_float[2];
+#pragma unroll
+        for (int i = 0; i < num_tiles_q; i++) {
+#pragma unroll
+            for (int j = 0; j < num_tiles_v; j++) {
+#pragma unroll
+                for (int k = 0; k < 4; k++) {
+                    unpack_half2_from_uint32_to_float(RO_tmp_float, RO_int32[i][j][k]);
+                    RO[i][j][k * 2 + 0] += RO_tmp_float[0];
+                    RO[i][j][k * 2 + 1] += RO_tmp_float[1];
+                }
+            }
+        }
+
+        // #pragma unroll
+        //   for (uint32_t fq = 0; fq < num_tiles_q; fq++)
+        //   {
+        // #pragma unroll
+        //     for (uint32_t fv = 0; fv < num_tiles_v; fv++)
+        //     {
+        //       RO[fq][fv][0] += RO_inst_buf[fq][fv][0];
+        //       RO[fq][fv][1] += RO_inst_buf[fq][fv][1];
+        //       RO[fq][fv][2] += RO_inst_buf[fq][fv][2];
+        //       RO[fq][fv][3] += RO_inst_buf[fq][fv][3];
+        //       RO[fq][fv][4] += RO_inst_buf[fq][fv][4];
+        //       RO[fq][fv][5] += RO_inst_buf[fq][fv][5];
+        //       RO[fq][fv][6] += RO_inst_buf[fq][fv][6];
+        //       RO[fq][fv][7] += RO_inst_buf[fq][fv][7];
+        //     }
+        //   }
+    }
+
+}  // namespace sageattention
+
+#endif

src/core/cuda/sageattention/cp_async.cuh

@@ -0,0 +1,148 @@
+/*
+ * Copyright (c) 2024 by SageAttention team.
+ *
+ * This file is based on code from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/cp_async.cuh
+ * Copyright (c) 2023 by FlashInfer team.
+ * Small modifications made by SageAttention team, 2024 (e.g., renamed namespace).
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SD_CORE_CUDA_SAGEATTENTION_CP_ASYNC_CUH
+#define SD_CORE_CUDA_SAGEATTENTION_CP_ASYNC_CUH
+#include <cuda_runtime.h>
+
+namespace sageattention {
+
+    namespace cp_async {
+
+        enum class SharedMemFillMode {
+            kFillZero,  // Fill zero to shared memory when predicate is false
+            kNoFill     // Do not fill zero to shared memory when predicate is false
+        };
+
+        enum class PrefetchMode {
+            kNoPrefetch,  // Do not fetch additional data from global memory to L2
+            kPrefetch     // Fetch additional data from global memory to L2
+        };
+
+#if (__CUDACC_VER_MAJOR__ >= 11)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
+#define CP_ASYNC_ENABLED
+#endif
+#endif
+
+        /*!
+         * \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted
+         *   cp.async instructions to a group
+         */
+        __device__ __forceinline__ void commit_group() {
+#ifdef CP_ASYNC_ENABLED
+            asm volatile("cp.async.commit_group;\n" ::);
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of PTX cp.async.wait_group instruction
+         * \tparam n Wait till most recent n groups are committed
+         */
+        template <size_t n>
+        __device__ __forceinline__ void wait_group() {
+#ifdef CP_ASYNC_ENABLED
+            asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
+         *   global memory to shared memory
+         * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
+         * \tparam T Data type
+         * \param smem_ptr Pointer to shared memory
+         * \param gmem_ptr Pointer to global memory
+         */
+        template <PrefetchMode prefetch_mode, typename T>
+        __device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
+#ifdef CP_ASYNC_ENABLED
+            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+            if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+                asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+                             "l"(gmem_ptr), "n"(16), "r"(16));
+            } else {
+                asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+                             "l"(gmem_ptr), "n"(16), "r"(16));
+            }
+#else
+            *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
+         *   global memory to shared memory with predicate.
+         * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
+         * \tparam fill_mode Whether to fill zero to shared memory when predicate is false
+         * \tparam T Data type
+         * \param smem_ptr Pointer to shared memory
+         * \param gmem_ptr Pointer to global memory
+         * \param predicate Predicate value
+         * \note fill zero is slower than not fill zero
+         */
+        template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
+        __device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
+#ifdef CP_ASYNC_ENABLED
+            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+            if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
+                int src_in_bytes = predicate ? 16 : 0;
+                if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+                    asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+                                 "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
+                } else {
+                    asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+                                 "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
+                }
+            } else {
+                if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+                    asm volatile(
+                        "{\n"
+                        " .reg .pred p;\n"
+                        " setp.ne.b32 p, %0, 0;\n"
+                        " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n"
+                        "}\n" ::"r"((int)predicate),
+                        "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
+                } else {
+                    asm volatile(
+                        "{\n"
+                        " .reg .pred p;\n"
+                        " setp.ne.b32 p, %0, 0;\n"
+                        " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
+                        "}\n" ::"r"((int)predicate),
+                        "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
+                }
+            }
+#else
+            if (predicate) {
+                *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
+            } else {
+                if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
+                    *((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0);
+                }
+            }
+#endif
+        }
+
+    }  // namespace cp_async
+
+}  // namespace sageattention
+
+#endif

src/core/cuda/sageattention/math.cuh

@@ -0,0 +1,166 @@
+/*
+ * Copyright (c) 2024 by SageAttention team.
+ *
+ * This file is based on code from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/math.cuh
+ * Copyright (c) 2023 by FlashInfer team.
+ * Small modifications made by SageAttention team, 2024 (e.g., renamed namespace).
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SD_CORE_CUDA_SAGEATTENTION_MATH_CUH
+#define SD_CORE_CUDA_SAGEATTENTION_MATH_CUH
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+namespace sageattention {
+#ifndef USHORT_TYPE
+#define USHORT_TYPE
+    typedef unsigned short ushort;
+#endif
+
+    namespace math {
+
+        // log2(e)
+        constexpr float log2e      = 1.44269504088896340736f;
+        constexpr float log2e_recp = 1.0f / log2e;
+
+        __forceinline__ __device__ half2 uint32_as_half2(uint32_t x) {
+            return *(half2*)&x;
+        }
+
+        __forceinline__ __device__ uint32_t half2_as_uint32(half2 x) {
+            return *(uint32_t*)&x;
+        }
+
+        /*!
+         * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x
+         * \param x input
+         */
+        __forceinline__ __device__ float ptx_exp2(float x) {
+            float y;
+            asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
+
+        /*!
+         * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x)
+         * \param x input
+         */
+        __forceinline__ __device__ float ptx_log2(float x) {
+            float y;
+            asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
+
+        /*!
+         * \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x
+         * \param x input
+         */
+        __forceinline__ __device__ half2 ptx_exp2(half2 x) {
+            uint32_t y_u32;
+            uint32_t x_u32 = half2_as_uint32(x);
+            asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
+            return uint32_as_half2(y_u32);
+        }
+
+        /*!
+         * \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x
+         * \param x input
+         */
+        __forceinline__ __device__ half ptx_exp2(half x) {
+            ushort y_u16;
+            asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
+            return __ushort_as_half(y_u16);
+        }
+
+        /*!
+         * \brief Wrapper of PTX rcp.approx instruction, which computes 1/x
+         * \param x input
+         */
+        __forceinline__ __device__ float ptx_rcp(float x) {
+            float y;
+            asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
+
+        /*!
+         * \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle
+         *   between threads in a warp.
+         * \param x The value in the source lane
+         * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta]
+         */
+        __forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) {
+            float y;
+            asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;"
+                         : "=f"(y)
+                         : "f"(x), "r"(lane_mask));
+            return y;
+        }
+
+        /*!
+         * \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly
+         *   shuffle between threads in a warp.
+         * \param x The value in the source lane
+         * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask]
+         */
+        __forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) {
+            return __shfl_xor_sync(0xffffffff, x, lane_mask);
+        }
+
+        /*!
+         * \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x)
+         * \param x input
+         */
+        __forceinline__ __device__ float rsqrt(float x) {
+            float y;
+            asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
+
+        /*!
+         * \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x)
+         * \param x input
+         */
+        __forceinline__ __device__ float tanh(float x) {
+            float y;
+            asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x));
+            return y;
+        }
+
+        /*!
+         * \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x)
+         * \param x input
+         */
+        __forceinline__ __device__ half2 tanh(half2 x) {
+            uint32_t y_u32;
+            uint32_t x_u32 = half2_as_uint32(x);
+            asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
+            return uint32_as_half2(y_u32);
+        }
+
+        /*!
+         * \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x)
+         * \param x input
+         */
+        __forceinline__ __device__ half tanh(half x) {
+            ushort y_u16;
+            asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
+            return __ushort_as_half(y_u16);
+        }
+
+    }  // namespace math
+
+}  // namespace sageattention
+
+#endif

src/core/cuda/sageattention/mma.cuh

@@ -0,0 +1,684 @@
+/*
+ * Adapted from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/mma.cuh
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Modifications copyright (c) 2024 by SageAttention team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SD_CORE_CUDA_SAGEATTENTION_MMA_CUH
+#define SD_CORE_CUDA_SAGEATTENTION_MMA_CUH
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cassert>
+#include <type_traits>
+
+namespace sageattention {
+
+    namespace mma {
+
+#if (__CUDACC_VER_MAJOR__ >= 11)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
+#define MMA_F16F16F32_M16N8K16_ENABLED
+#define MMA_F16F16F16_M16N8K16_ENABLED
+#define MMA_S8S8S32_M16N8K32_ENABLED
+#define MMA_S4S4S32_M16N8K64_ENABLED
+#endif
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 750))
+#define MMA_F16F16F32_M16N8K8_ENABLED
+#define MMA_F16F16F16_M16N8K8_ENABLED
+#define LDMATRIX_M8N8X2_ENABLED
+#define LDMATRIX_M8N8X4_ENABLED
+#endif
+#endif
+
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
+#define MMA_F8F8F32_M16N8K16_ENABLED
+#endif
+#endif
+
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
+#define MMA_F8F8F16_M16N8K16_ENABLED
+#endif
+#endif
+
+#if defined(__CUDA_ARCH__)
+#define RUNTIME_ASSERT(x) __brkpt()
+#else
+
+#define RUNTIME_ASSERT(x) assert(0 && x)
+#endif
+
+        enum class MMAMode {
+            kInit          = 0U,
+            kInplaceUpdate = 1U,
+        };
+
+        /*!
+         * \brief Wrapper of PTX ldmatrix m8n8.x2 instruction, loads data from shared memory
+         *   to fragment
+         * \tparam T data type of the fragment
+         * \param R pointer to the fragment
+         * \param smem_ptr pointer to the shared memory
+         */
+        template <typename T>
+        __device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) {
+#ifdef LDMATRIX_M8N8X2_ENABLED
+            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+            asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
+                         : "=r"(R[0]), "=r"(R[1])
+                         : "r"(smem_int_ptr));
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory
+         *   to fragment
+         * \tparam T data type of the fragment
+         * \param R pointer to the fragment
+         * \param smem_ptr pointer to the shared memory
+         */
+        template <typename T>
+        __device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) {
+#ifdef LDMATRIX_M8N8X4_ENABLED
+            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+            asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+                         : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
+                         : "r"(smem_int_ptr));
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from
+         *   shared memory to fragment and transposes the fragment
+         * \tparam T data type of the fragment
+         * \param R pointer to the fragment
+         * \param smem_ptr pointer to the shared memory
+         */
+        template <typename T>
+        __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) {
+#ifdef LDMATRIX_M8N8X4_ENABLED
+            uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+            asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+                         : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
+                         : "r"(smem_int_ptr));
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
+         *   multiplication, accumulated in f32.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F16F16F32_M16N8K16_ENABLED
+            // ! only support half dtype now
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
+                      "f"(C[2]), "f"(C[3]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
+                      "f"(0.f), "f"(0.f));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
+         *   multiplication, accumulated in f32.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F16F16F32_M16N8K16_ENABLED
+            // ! only support half dtype now
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
+                      "f"(C[2]), "f"(C[3]));
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
+                      "f"(C[6]), "f"(C[7]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
+                      "f"(0.f), "f"(0.f));
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f),
+                      "f"(0.f), "f"(0.f));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
+         *   multiplication, accumulated in f16.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F16F16F16_M16N8K16_ENABLED
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C[0]), "=r"(C[1])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C[0]), "=r"(C[1])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
+         *   multiplication, accumulated in f16.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F16F16F16_M16N8K16_ENABLED
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C[0]), "=r"(C[1])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]));
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C[0]), "=r"(C[1])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
+                asm volatile(
+                    "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int8 matrix
+         *   multiplication, accumulated in int32.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_S8S8S32_M16N8K32_ENABLED
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
+                      "r"(C[2]), "r"(C[3]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
+                      "r"(0), "r"(0));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n16k32 instruction for row major and column major int8 matrix
+         *   multiplication, accumulated in int32.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_S8S8S32_M16N8K32_ENABLED
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
+                      "r"(C[2]), "r"(C[3]));
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
+                      "r"(C[6]), "r"(C[7]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
+                      "r"(0), "r"(0));
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0),
+                      "r"(0), "r"(0));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int4 matrix
+         *   multiplication, accumulated in int32.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_S4S4S32_M16N8K64_ENABLED
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
+                      "r"(C[2]), "r"(C[3]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
+                      "r"(0), "r"(0));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n16k64 instruction for row major and column major int4 matrix
+         *   multiplication, accumulated in int32.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_S4S4S32_M16N8K64_ENABLED
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
+                      "r"(C[2]), "r"(C[3]));
+                asm volatile(
+                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
+                      "r"(C[6]), "r"(C[7]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
+                      "r"(0), "r"(0));
+                asm volatile(
+                    "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0),
+                      "r"(0), "r"(0));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n8k32 instruction for row major and column major fp8 e4m3 matrix
+         *   multiplication, accumulated in fp32.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F8F8F32_M16N8K16_ENABLED
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
+                      "f"(C[2]), "f"(C[3]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
+                      "f"(0.f), "f"(0.f));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
+         *   multiplication, accumulated in fp16.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, uint32_t* B) {
+            // uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C);
+#ifdef MMA_F8F8F16_M16N8K16_ENABLED
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C_uint32[0]), "=r"(C_uint32[1])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C_uint32[0]), "r"(C_uint32[1]));
+
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C_uint32[2]), "=r"(C_uint32[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C_uint32[2]), "r"(C_uint32[3]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C_uint32[0]), "=r"(C_uint32[1])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
+
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
+                    "{%0,  %1},"
+                    "{%2,  %3,  %4,  %5},"
+                    "{%6,  %7},"
+                    "{%8,  %9};\n"
+                    : "=r"(C_uint32[2]), "=r"(C_uint32[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
+         *   multiplication, accumulated in fp32.
+         * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+         * \param C pointer to the accumulator
+         * \param A pointer to the fragment of matrix A
+         * \param B pointer to the fragment of matrix B
+         */
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+        __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F8F8F32_M16N8K16_ENABLED
+            if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
+                      "f"(C[2]), "f"(C[3]));
+
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
+                      "f"(C[6]), "f"(C[7]));
+            } else if constexpr (mma_mode == MMAMode::kInit) {
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
+                      "f"(0.f), "f"(0.f));
+
+                asm volatile(
+                    "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+                    "{%0,  %1,  %2,  %3},"
+                    "{%4,  %5,  %6,  %7},"
+                    "{%8,  %9},"
+                    "{%10, %11, %12, %13};\n"
+                    : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+                    : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f),
+                      "f"(0.f), "f"(0.f));
+            }
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Use mma instructions to compute rowsum.
+         */
+        __device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) {
+#ifdef MMA_F16F16F32_M16N8K16_ENABLED
+            asm volatile(
+                "{\n"
+                "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+                "{%0,  _,  %1,  _},"
+                "{%2,  %3,  %4,  %5},"
+                "{%6,  %7},"
+                "{%8,  0.,  %9,  0.};\n"
+                "}\n"
+                : "=f"(d[0]), "=f"(d[1])
+                : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(1006648320),  // 1006648320 packs two 1.0f in half precision
+                  "r"(1006648320), "f"(d[0]), "f"(d[1]));
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+        /*!
+         * \brief Use mma instructions to compute rowsum.
+         */
+        __device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) {
+#ifdef MMA_F8F8F32_M16N8K16_ENABLED
+            asm volatile(
+                "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+                "{%0,  _,  %1,  _},"
+                "{%2,  %3,  %4,  %5},"
+                "{%6,  %7},"
+                "{%8,  0.,  %9,  0.};\n"
+                : "=f"(d[0]), "=f"(d[1])
+                : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(943208504), "r"(943208504),  // 943208504 packs four 1.0f in e4m3
+                  "f"(d[0]), "f"(d[1]));
+#else
+            RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+        }
+
+    }  // namespace mma
+
+}  // namespace sageattention
+
+#endif

src/core/cuda/sageattention/numeric_conversion.cuh

@@ -0,0 +1,149 @@
+/*
+ * Copyright (c) 2024 by SageAttention team.
+ *
+ * Inspired by CUTLASS, https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/numeric_conversion.h
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SD_CORE_CUDA_SAGEATTENTION_NUMERIC_CONVERSION_CUH
+#define SD_CORE_CUDA_SAGEATTENTION_NUMERIC_CONVERSION_CUH
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cassert>
+
+namespace sageattention {
+
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
+#define FP8_CAST_ENABLED
+#endif
+#endif
+
+#if defined(__CUDA_ARCH__)
+#define RUNTIME_ASSERT(x) __brkpt()
+#else
+
+#define RUNTIME_ASSERT(x) assert(0 && x)
+#endif
+
+    __device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) {
+        uint16_t h0 = source & 0xFFFF;
+        uint16_t h1 = (source >> 16) & 0xFFFF;
+        asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0));
+        asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1));
+    }
+
+    __device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t* dest, float* source0, float* source1) {
+#ifdef FP8_CAST_ENABLED
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo;\n"
+            ".reg .b16 hi;\n"
+            "cvt.rn.satfinite.e4m3x2.f32   lo, %2, %1;\n"
+            "cvt.rn.satfinite.e4m3x2.f32   hi, %4, %3;\n"
+            "mov.b32 %0, {lo, hi};\n"
+            "}"
+            : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1]));
+#else
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+    }
+
+    __device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t* dest, float* source0, float* source1) {
+#ifdef FP8_CAST_ENABLED
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo;\n"
+            ".reg .b16 hi;\n"
+            "cvt.rn.satfinite.e5m2x2.f32   lo, %2, %1;\n"
+            "cvt.rn.satfinite.e5m2x2.f32   hi, %4, %3;\n"
+            "mov.b32 %0, {lo, hi};\n"
+            "}"
+            : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1]));
+#else
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+    }
+
+    __device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
+#ifdef FP8_CAST_ENABLED
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo;\n"
+            ".reg .b16 hi;\n"
+            "cvt.rn.satfinite.e4m3x2.f16x2   lo, %1;\n"
+            "cvt.rn.satfinite.e4m3x2.f16x2   hi, %2;\n"
+            "mov.b32 %0, {lo, hi};\n"
+            "}"
+            : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
+#else
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+    }
+
+    __device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
+#ifdef FP8_CAST_ENABLED
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo;\n"
+            ".reg .b16 hi;\n"
+            "cvt.rn.satfinite.e5m2x2.f16x2   lo, %1;\n"
+            "cvt.rn.satfinite.e5m2x2.f16x2   hi, %2;\n"
+            "mov.b32 %0, {lo, hi};\n"
+            "}"
+            : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
+#else
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+    }
+
+    __device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
+#ifdef FP8_CAST_ENABLED
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo, hi;\n"
+            "mov.b32 {lo, hi}, %2;\n"
+            "cvt.rn.f16x2.e4m3x2 %0, lo;\n"
+            "cvt.rn.f16x2.e4m3x2 %1, hi;\n"
+            "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
+#else
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+    }
+
+    __device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
+#ifdef FP8_CAST_ENABLED
+        asm volatile(
+            "{\n"
+            ".reg .b16 lo, hi;\n"
+            "mov.b32 {lo, hi}, %2;\n"
+            "cvt.rn.f16x2.e5m2x2 %0, lo;\n"
+            "cvt.rn.f16x2.e5m2x2 %1, hi;\n"
+            "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
+#else
+        RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+    }
+
+    __device__ __forceinline__ int8_t float_to_int8_rn(float x) {
+        uint32_t dst;
+        asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
+        return reinterpret_cast<const int8_t&>(dst);
+    }
+
+}  // namespace sageattention
+
+#endif

src/core/cuda/sageattention/permuted_smem.cuh

@@ -0,0 +1,201 @@
+/*
+ * Adapted from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/permuted_smem.cuh
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Modifications copyright (c) 2024 by SageAttention team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SD_CORE_CUDA_SAGEATTENTION_PERMUTED_SMEM_CUH
+#define SD_CORE_CUDA_SAGEATTENTION_PERMUTED_SMEM_CUH
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include "cp_async.cuh"
+#include "mma.cuh"
+
+namespace sageattention {
+
+    enum class SwizzleMode {
+        k32B,   // for k32B mode, a line of shared memory must have 32B (16 half value)
+        k64B,   // for k64B mode, a line of shared memory must have 64B (32 half value)
+        k128B,  // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B.
+    };
+
+    // Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth
+    using b128_t = uint4;
+
+    /*!
+     * \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure
+     * that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through
+     * permutation in the granularity of 128bit.
+     *
+     * This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless,
+     * meaning it does not maintain any information about the current pointer position. The offset returnd by
+     * the struct can be used to access the shared memory through the provided interface.
+     *
+     * The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j).
+     * We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp.
+     * Under this assumption, we do not permute for k32B mode.
+     */
+    template <SwizzleMode swizzle_mode, uint32_t stride>
+    struct smem_t {
+        // The base pointer.
+        b128_t* base;
+        // How many b128_t value a row contains
+        // uint32_t stride;
+
+        __device__ __forceinline__ smem_t() : base(nullptr) {}
+        template <typename T>
+        __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode");
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode");
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode");
+            } else {
+                static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode");
+            }
+        }
+
+        /*!
+         * \brief Set the base pointer.
+         */
+        template <typename T>
+        __device__ __forceinline__ void set_base(T* new_base) {
+            base = (b128_t*)new_base;
+        }
+
+        /*!
+         * \brief Compute the element offset given coordinates in a permuted shared memory.
+         * \param i The row index.
+         * \param j The column index.
+         */
+        static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t& i, const uint32_t& j) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                return i * stride + (j ^ (i % 8));
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                return i * stride + (j ^ ((i / 2) % 4));
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                return i * stride + j;
+            }
+        }
+
+        /*!
+         * \tparam step_size The step size to advance the offset in the permuted shared memory.
+         * \param offset The current offset.
+         */
+        template <uint32_t step_size>
+        static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                static_assert(step_size % 8 == 0,
+                              "Unsupported step size");
+                return offset + step_size;
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                static_assert(step_size == 4, "Unsupported step size");
+                return offset + step_size;
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                static_assert(step_size == 2, "Unsupported step size");
+                return offset + step_size;
+            }
+        }
+
+        // ! use with care
+        template <uint32_t step_size>
+        static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset, const uint32_t& step_idx) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0,
+                              "Unsupported step size");
+                if constexpr (step_size == 2) {
+                    return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8;
+                } else if constexpr (step_size == 4) {
+                    return (offset ^ 0x4) + (step_idx % 2 == 1) * 8;
+                } else {
+                    // step_size % 8 == 0
+                    return offset + step_size;
+                }
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                static_assert(step_size == 2 || step_size == 4, "Unsupported step size");
+                if constexpr (step_size == 2) {
+                    return (offset ^ 0x2) + (step_idx % 2 == 1) * 4;
+                } else {
+                    return offset + step_size;
+                }
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                return offset + step_size;
+            }
+        }
+
+        template <uint32_t step_size>
+        static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t& offset) {
+            if constexpr (swizzle_mode == SwizzleMode::k128B) {
+                static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
+                if constexpr (step_size == 4) {
+                    return (offset ^ 0x4) + step_size * stride;
+                } else {
+                    // step_size % 8 == 0
+                    return offset + step_size * stride;
+                }
+            } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+                static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
+                if constexpr (step_size == 4) {
+                    return (offset ^ 0x2) + step_size * stride;
+                } else {
+                    // step_size % 8 == 0
+                    return offset + step_size * stride;
+                }
+            } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+                return offset + step_size * stride;
+            }
+        }
+
+        __device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t& offset, uint32_t* R) const {
+            b128_t* smem_ptr = base + offset;
+            mma::ldmatrix_m8n8x2(R, smem_ptr);
+        }
+
+        __device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t& offset, uint32_t* R) const {
+            b128_t* smem_ptr = base + offset;
+            mma::ldmatrix_m8n8x4(R, smem_ptr);
+        }
+
+        __device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t& offset, uint32_t* R) const {
+            b128_t* smem_ptr = base + offset;
+            mma::ldmatrix_m8n8x4_trans(R, smem_ptr);
+        }
+
+        template <cp_async::SharedMemFillMode fill_mode, typename T>
+        __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr, bool predicate) const {
+            b128_t* smem_ptr = base + offset;
+            cp_async::pred_load_128b<cp_async::PrefetchMode::kPrefetch, fill_mode>(
+                smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate);
+        }
+
+        template <typename T>
+        __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr) const {
+            b128_t* smem_ptr = base + offset;
+            cp_async::load_128b<cp_async::PrefetchMode::kPrefetch>(smem_ptr, reinterpret_cast<const b128_t*>(gptr));
+        }
+
+        template <typename T>
+        __device__ __forceinline__ void store_128b(const uint32_t& offset, T* gptr) const {
+            *reinterpret_cast<b128_t*>(gptr) = *(base + offset);
+        }
+    };
+
+}  // namespace sageattention
+
+#endif

src/core/cuda/sol_attention.cu

@@ -0,0 +1,409 @@
+// Sol-Attn diagonal-threshold algorithm:
+// https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+#if __CUDACC_VER_MAJOR__ >= 12
+
+#include <cuda_bf16.h>
+#include <algorithm>
+#include <climits>
+
+namespace sd_cuda {
+
+    static __device__ __forceinline__ float warp_reduce_sum(float x) {
+#pragma unroll
+        for (int offset = 16; offset > 0; offset >>= 1)
+            x += __shfl_xor_sync(0xffffffff, x, offset);
+        return x;
+    }
+
+    static constexpr int SOL_DIM   = 128;
+    static constexpr int SOL_BLOCK = 64;
+
+    struct SolWorkspace {
+        size_t q, k, v, kc, vc, qm, stats, threshold, size;
+        int tokens, blocks;
+
+        explicit SolWorkspace(const ggml_tensor* dst) {
+            const auto* input  = dst->src[0];
+            const size_t heads = input->ne[2] * input->ne[3];
+            tokens             = GGML_PAD(input->ne[1], SOL_BLOCK);
+            blocks             = GGML_PAD(tokens / SOL_BLOCK, SOL_BLOCK);
+            size               = 0;
+            q                  = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+            k                  = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+            v                  = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+            kc                 = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
+            vc                 = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
+            qm                 = reserve(heads * (tokens / SOL_BLOCK) * SOL_DIM * sizeof(float));
+            stats              = reserve(heads * 2 * SOL_DIM * sizeof(float));
+            threshold          = reserve(heads * (tokens / SOL_BLOCK) * sizeof(float));
+        }
+
+        size_t reserve(size_t bytes) {
+            const size_t offset = size;
+            size += GGML_PAD(bytes, 256);
+            return offset;
+        }
+    };
+
+    static __global__ void sol_pack(const float* q, const float* k, const float* v, nv_bfloat16* qb, nv_bfloat16* kb, nv_bfloat16* vb, int length, int padded) {
+        const int index = blockIdx.x * blockDim.x + threadIdx.x;
+        if (index >= padded * SOL_DIM) {
+            return;
+        }
+        const int token   = index / SOL_DIM;
+        const int channel = index % SOL_DIM;
+        const size_t src  = size_t(blockIdx.y) * length * SOL_DIM + index;
+        const size_t dst  = size_t(blockIdx.y) * padded * SOL_DIM + index;
+        qb[dst]           = __float2bfloat16(token < length ? q[src] : 0.f);
+        kb[dst]           = __float2bfloat16(token < length ? k[src] : 0.f);
+        vb[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token] =
+            __float2bfloat16(token < length ? v[src] : 0.f);
+    }
+
+    static __global__ void sol_summaries(const nv_bfloat16* q, const nv_bfloat16* k, const nv_bfloat16* v, nv_bfloat16* kc, nv_bfloat16* vc, float* qm, int length, int padded, int padded_blocks) {
+        const int channel = threadIdx.x;
+        const int block   = blockIdx.x;
+        const int count   = min(SOL_BLOCK, max(0, length - block * SOL_BLOCK));
+        float sum_q = 0.f, sum_k = 0.f, sum_v = 0.f;
+        for (int i = 0; i < count; ++i) {
+            const int token    = block * SOL_BLOCK + i;
+            const size_t index = (size_t(blockIdx.y) * padded + token) * SOL_DIM + channel;
+            sum_q += __bfloat162float(q[index]);
+            sum_k += __bfloat162float(k[index]);
+            sum_v += __bfloat162float(v[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token]);
+        }
+        kc[(size_t(blockIdx.y) * padded_blocks + block) * SOL_DIM + channel] =
+            __float2bfloat16(count ? sum_k / count : 0.f);
+        vc[(size_t(blockIdx.y) * SOL_DIM + channel) * padded_blocks + block] = __float2bfloat16(sum_v);
+        if (count) {
+            qm[(size_t(blockIdx.y) * (padded / SOL_BLOCK) + block) * SOL_DIM + channel] = sum_q / count;
+        }
+    }
+
+    static __global__ void sol_statistics(const nv_bfloat16* kc, float* stats, int blocks, int padded_blocks) {
+        const int channel = threadIdx.x;
+        float sum = 0.f, sum_sq = 0.f;
+        for (int i = 0; i < blocks; ++i) {
+            const float value = __bfloat162float(kc[(size_t(blockIdx.x) * padded_blocks + i) * SOL_DIM + channel]);
+            sum += value;
+            sum_sq += value * value;
+        }
+        const float mean                                            = sum / blocks;
+        stats[size_t(blockIdx.x) * 2 * SOL_DIM + channel]           = mean;
+        stats[size_t(blockIdx.x) * 2 * SOL_DIM + SOL_DIM + channel] = fmaxf(sum_sq / blocks - mean * mean, 0.f);
+    }
+
+    static __global__ void sol_thresholds(const float* qm, const float* stats, float* threshold, int blocks, float scale, float tau) {
+        const int channel = threadIdx.x;
+        const float q     = qm[(size_t(blockIdx.y) * blocks + blockIdx.x) * SOL_DIM + channel];
+        float mean        = q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + channel];
+        float var         = q * q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + SOL_DIM + channel];
+        mean              = warp_reduce_sum(mean);
+        var               = warp_reduce_sum(var);
+        __shared__ float means[4], vars[4];
+        if (channel % 32 == 0) {
+            means[channel / 32] = mean;
+            vars[channel / 32]  = var;
+        }
+        __syncthreads();
+        if (channel == 0) {
+            mean = means[0] + means[1] + means[2] + means[3];
+            var  = vars[0] + vars[1] + vars[2] + vars[3];
+            threshold[size_t(blockIdx.y) * blocks + blockIdx.x] =
+                mean * scale + tau * sqrtf(fmaxf(var * scale * scale, 0.f) + 1.e-6f);
+        }
+    }
+
+    struct SolQTile {
+        uint32_t x[4];
+    };
+    struct SolKTile {
+        uint32_t x[2];
+    };
+    struct SolTile {
+        float x[4] = {};
+        static __device__ __forceinline__ int get_i(int i) {
+            return (i / 2) * 8 + threadIdx.x / 4;
+        }
+        static __device__ __forceinline__ int get_j(int i) {
+            return (threadIdx.x % 4) * 2 + i % 2;
+        }
+    };
+
+    static __device__ __forceinline__ void sol_load(SolQTile& tile, const nv_bfloat162* input, int stride) {
+        const auto* data = reinterpret_cast<const uint32_t*>(input);
+#pragma unroll
+        for (int i = 0; i < 4; ++i)
+            tile.x[i] = data[((i % 2) * 8 + threadIdx.x / 4) * stride + (i / 2) * 4 + threadIdx.x % 4];
+    }
+    static __device__ __forceinline__ void sol_load(SolKTile& tile, const nv_bfloat162* input, int stride) {
+        const auto* data = reinterpret_cast<const uint32_t*>(input);
+#pragma unroll
+        for (int i = 0; i < 2; ++i)
+            tile.x[i] = data[(threadIdx.x / 4) * stride + i * 4 + threadIdx.x % 4];
+    }
+    static __device__ __forceinline__ void sol_mma(SolTile& output, const SolQTile& q, const SolKTile& k) {
+#if __CUDA_ARCH__ >= 800
+        auto* dst = reinterpret_cast<uint32_t*>(output.x);
+        asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+            : "+r"(dst[0]), "+r"(dst[1]), "+r"(dst[2]), "+r"(dst[3])
+            : "r"(q.x[0]), "r"(q.x[1]), "r"(q.x[2]), "r"(q.x[3]), "r"(k.x[0]), "r"(k.x[1]));
+#else
+        asm("trap;");
+#endif
+    }
+
+    static __device__ __forceinline__ void sol_scores(const SolQTile (&q)[8], const nv_bfloat16* k, SolTile (&scores)[8], float scale) {
+#pragma unroll
+        for (int n = 0; n < 8; ++n) {
+            scores[n] = SolTile{};
+#pragma unroll
+            for (int d = 0; d < 8; ++d) {
+                SolKTile kt;
+                sol_load(kt, reinterpret_cast<const nv_bfloat162*>(k + n * 8 * SOL_DIM + d * 16), SOL_DIM / 2);
+                sol_mma(scores[n], q[d], kt);
+            }
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                scores[n].x[i] *= scale;
+            }
+        }
+    }
+
+    template <bool PROXY>
+    static __device__ __forceinline__ void sol_update(SolTile (&scores)[8], SolTile (&output)[16], float (&row_max)[2], float (&row_sum)[2], nv_bfloat16* probabilities, const nv_bfloat16* values, int stride, int start, int length, const int* exact) {
+        float maximum[2] = {-INFINITY, -INFINITY};
+#pragma unroll
+        for (int n = 0; n < 8; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const int column = n * 8 + SolTile::get_j(i);
+                const bool valid = PROXY ? ((start + column) * SOL_BLOCK < length && !exact[column])
+                                         : (start + column < length);
+                scores[n].x[i]   = valid ? scores[n].x[i] : -INFINITY;
+                maximum[i / 2]   = fmaxf(maximum[i / 2], scores[n].x[i]);
+            }
+        }
+        float alpha[2], sums[2] = {0.f, 0.f};
+#pragma unroll
+        for (int row = 0; row < 2; ++row) {
+            maximum[row]     = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 1));
+            maximum[row]     = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 2));
+            const float next = fmaxf(row_max[row], maximum[row]);
+            alpha[row]       = row_max[row] == next ? 1.f : exp2f(row_max[row] - next);
+            row_max[row]     = next;
+        }
+#pragma unroll
+        for (int n = 0; n < 8; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const int column = n * 8 + SolTile::get_j(i);
+                const float p    = scores[n].x[i] == -INFINITY ? 0.f : exp2f(scores[n].x[i] - row_max[i / 2]);
+                const int count  = PROXY ? min(SOL_BLOCK, max(0, length - (start + column) * SOL_BLOCK)) : 1;
+                sums[i / 2] += p * count;
+                probabilities[SolTile::get_i(i) * SOL_BLOCK + column] = __float2bfloat16(p);
+            }
+        }
+#pragma unroll
+        for (int row = 0; row < 2; ++row) {
+            sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 1);
+            sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 2);
+            row_sum[row] = row_sum[row] * alpha[row] + sums[row];
+        }
+        __syncwarp();
+#pragma unroll
+        for (int n = 0; n < 16; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                output[n].x[i] *= alpha[i / 2];
+            }
+#pragma unroll
+            for (int d = 0; d < 4; ++d) {
+                SolQTile pt;
+                SolKTile vt;
+                sol_load(pt, reinterpret_cast<const nv_bfloat162*>(probabilities + d * 16), SOL_BLOCK / 2);
+                sol_load(vt, reinterpret_cast<const nv_bfloat162*>(values + n * 8 * stride + d * 16), stride / 2);
+                sol_mma(output[n], pt, vt);
+            }
+        }
+        __syncwarp();
+    }
+
+    static __global__ __launch_bounds__(128) void sol_forward(
+        const nv_bfloat16* q,
+        const nv_bfloat16* k,
+        const nv_bfloat16* v,
+        const nv_bfloat16* kc,
+        const nv_bfloat16* vc,
+        const float* thresholds,
+        float* dst,
+        int length,
+        int padded,
+        int padded_blocks,
+        int heads,
+        float scale) {
+#if __CUDA_ARCH__ >= 800
+        const int warp    = threadIdx.y;
+        const int lane    = threadIdx.x;
+        const int tid     = warp * 32 + lane;
+        const int blocks  = padded / SOL_BLOCK;
+        const int q_start = blockIdx.x * SOL_BLOCK;
+        const int q_len   = min(SOL_BLOCK, length - q_start);
+        const size_t base = size_t(blockIdx.y) * padded * SOL_DIM;
+        SolQTile qt[8];
+#pragma unroll
+        for (int d = 0; d < 8; ++d) {
+            sol_load(qt[d], reinterpret_cast<const nv_bfloat162*>(q + base + (q_start + warp * 16) * SOL_DIM + d * 16), SOL_DIM / 2);
+        }
+        SolTile output[16];
+        float row_max[2] = {-INFINITY, -INFINITY};
+        float row_sum[2] = {0.f, 0.f};
+        __shared__ nv_bfloat16 probabilities[SOL_BLOCK * SOL_BLOCK];
+        __shared__ float routes[4][SOL_BLOCK];
+        __shared__ int exact[SOL_BLOCK];
+        const float threshold = thresholds[size_t(blockIdx.y) * blocks + blockIdx.x];
+        for (int group = 0; group < blocks; group += SOL_BLOCK) {
+            SolTile scores[8];
+            sol_scores(qt, kc + (size_t(blockIdx.y) * padded_blocks + group) * SOL_DIM, scores, scale);
+#pragma unroll
+            for (int n = 0; n < 8; ++n) {
+#pragma unroll
+                for (int column = 0; column < 2; ++column) {
+                    float sum = scores[n].x[column] + scores[n].x[column + 2];
+                    sum += __shfl_xor_sync(0xffffffff, sum, 4);
+                    sum += __shfl_xor_sync(0xffffffff, sum, 8);
+                    sum += __shfl_xor_sync(0xffffffff, sum, 16);
+                    if (lane < 4) {
+                        routes[warp][n * 8 + 2 * lane + column] = sum;
+                    }
+                }
+            }
+            __syncthreads();
+            if (tid < SOL_BLOCK) {
+                const int block   = group + tid;
+                const float score = (routes[0][tid] + routes[1][tid] + routes[2][tid] + routes[3][tid]) / q_len;
+                exact[tid]        = block < blocks && (score > threshold || abs(int(blockIdx.x) - block) <= 1);
+            }
+            __syncthreads();
+            sol_update<true>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
+                             vc + size_t(blockIdx.y) * padded_blocks * SOL_DIM + group,
+                             padded_blocks, group, length, exact);
+            for (int block = 0; block < SOL_BLOCK && group + block < blocks; ++block) {
+                if (!exact[block]) {
+                    continue;
+                }
+                const int start = (group + block) * SOL_BLOCK;
+                sol_scores(qt, k + base + start * SOL_DIM, scores, scale);
+                sol_update<false>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
+                                  v + base + start, padded, start, length, nullptr);
+            }
+            __syncthreads();
+        }
+        const int head  = blockIdx.y % heads;
+        const int batch = blockIdx.y / heads;
+#pragma unroll
+        for (int n = 0; n < 16; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const int token = q_start + warp * 16 + SolTile::get_i(i);
+                if (token < length) {
+                    const int channel = n * 8 + SolTile::get_j(i);
+                    dst[((size_t(batch) * length + token) * heads + head) * SOL_DIM + channel] =
+                        __bfloat162float(__float2bfloat16(output[n].x[i] / row_sum[i / 2]));
+                }
+            }
+        }
+#else
+        asm("trap;");
+#endif
+    }
+
+    static bool sol_supports(ggml_backend_dev_t, const ggml_tensor* op, void*) {
+        SDSolAttentionParams params;
+        if (!sd_ggml_sol_attn_describe(op, &params))
+            return false;
+        const auto* q = op->src[0];
+        if (q->ne[2] * q->ne[3] > 65535 || GGML_PAD(q->ne[1], SOL_BLOCK) > INT_MAX / SOL_DIM)
+            return false;
+        for (const auto* tensor : {op->src[0], op->src[1], op->src[2]})
+            if (reinterpret_cast<uintptr_t>(tensor->data) & 3)
+                return false;
+        return !(reinterpret_cast<uintptr_t>(op->data) & 3);
+    }
+
+    static ggml_status sol_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
+        *workspace = SolWorkspace(nodes[0]).size;
+        return GGML_STATUS_SUCCESS;
+    }
+    static void sol_destroy(void*, void*) {}
+
+    static ggml_status sol_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
+        ggml_ext_gpu_context context;
+        if (!cuda_context(execution->backend, &context))
+            return GGML_STATUS_FAILED;
+        auto* dst = nodes[0];
+        SDSolAttentionParams params;
+        if (!sd_ggml_sol_attn_describe(dst, &params))
+            return GGML_STATUS_FAILED;
+        const auto* q = dst->src[0];
+        const SolWorkspace work(dst);
+        char* data       = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+        auto* qb         = reinterpret_cast<nv_bfloat16*>(data + work.q);
+        auto* kb         = reinterpret_cast<nv_bfloat16*>(data + work.k);
+        auto* vb         = reinterpret_cast<nv_bfloat16*>(data + work.v);
+        auto* kc         = reinterpret_cast<nv_bfloat16*>(data + work.kc);
+        auto* vc         = reinterpret_cast<nv_bfloat16*>(data + work.vc);
+        auto* qm         = reinterpret_cast<float*>(data + work.qm);
+        auto* stats      = reinterpret_cast<float*>(data + work.stats);
+        auto* threshold  = reinterpret_cast<float*>(data + work.threshold);
+        const int length = q->ne[1], heads = q->ne[2], batch_heads = heads * q->ne[3];
+        const int blocks  = work.tokens / SOL_BLOCK;
+        const float scale = params.scale * 1.4426950408889634f;
+        const auto stream = static_cast<cudaStream_t>(context.stream);
+        sol_pack<<<dim3((work.tokens * SOL_DIM + 255) / 256, batch_heads), 256, 0, stream>>>(
+            static_cast<const float*>(q->data), static_cast<const float*>(dst->src[1]->data), static_cast<const float*>(dst->src[2]->data),
+            qb, kb, vb, length, work.tokens);
+        sol_summaries<<<dim3(work.blocks, batch_heads), SOL_DIM, 0, stream>>>(qb, kb, vb, kc, vc, qm, length, work.tokens, work.blocks);
+        sol_statistics<<<batch_heads, SOL_DIM, 0, stream>>>(kc, stats, blocks, work.blocks);
+        sol_thresholds<<<dim3(blocks, batch_heads), SOL_DIM, 0, stream>>>(qm, stats, threshold, blocks, scale, params.tau);
+        sol_forward<<<dim3(blocks, batch_heads), dim3(32, 4), 0, stream>>>(qb, kb, vb, kc, vc, threshold,
+                                                                           static_cast<float*>(dst->data), length, work.tokens, work.blocks, heads, scale);
+        return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+
+    void register_sol_attention(ggml_backend_dev_t device) {
+        const auto info = cuda_device_info(device);
+        if (info.device < 0)
+            return;
+        int previous;
+        cudaDeviceProp properties;
+        cudaFuncAttributes attributes;
+        if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || properties.major < 8 ||
+            cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info.device) != cudaSuccess)
+            return;
+        const auto error = cudaFuncGetAttributes(&attributes, sol_forward);
+        if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess ||
+            std::min(attributes.binaryVersion, attributes.ptxVersion) < 80)
+            return;
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.sol_attn.cuda.v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_sol_attn_op();
+        desc.priority             = 100;
+        desc.supports             = sol_supports;
+        desc.prepare              = sol_prepare;
+        desc.execute              = sol_execute;
+        desc.destroy              = sol_destroy;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
+}
+#else
+namespace sd_cuda {
+    void register_sol_attention(ggml_backend_dev_t) {}
+}
+#endif

src/core/cuda/w4.cu

@@ -0,0 +1,939 @@
+#include <climits>
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+#ifdef SD_CUDA_W4_CUTLASS
+#include <cstdlib>
+#include <mutex>
+#include <unordered_map>
+#include "i8_gemm.h"
+#include "w4_mma.h"
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_w4_cuda {
+        static __device__ __forceinline__ int shuffle_sum(int value, int lane, int width) {
+#if defined(SD_CUDA_VENDOR_HIP)
+            // ROCm LLVM 24 can crash lowering __shfl's wave_shuffle on RDNA2.
+            // These 8/16-lane groups stay within one ds_bpermute half-wave.
+            const int index = (lane & (width - 1)) + (__lane_id() & ~(width - 1));
+            return __builtin_amdgcn_ds_bpermute(index * 4, value);
+#else
+            return cuda_shuffle(value, lane, width);
+#endif
+        }
+
+        // Signed byte dots preserve the INT4 grid after the XOR-8 bias correction.
+        static __device__ __forceinline__ int dot4(int a, int b, int c) {
+#if defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
+            return __builtin_amdgcn_sudot4(true, a, true, b, c, false);
+#elif defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx906__) || defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__))
+            return __builtin_amdgcn_sdot4(a, b, c, false);
+#elif defined(SD_CUDA_VENDOR_MUSA) || (!defined(SD_CUDA_VENDOR_HIP) && __CUDA_ARCH__ >= 610)
+            return __dp4a(a, b, c);
+#else
+            const auto* av = reinterpret_cast<const signed char*>(&a);
+            const auto* bv = reinterpret_cast<const signed char*>(&b);
+            for (int i = 0; i < 4; ++i)
+                c += int(av[i]) * int(bv[i]);
+            return c;
+#endif
+        }
+        static __device__ __forceinline__ float fp8_e4m3(uint8_t v) {
+            const int exponent = (v >> 3) & 15, mantissa = v & 7;
+            float result = exponent == 15 && mantissa == 7 ? __int_as_float(0x7fc00000) : exponent == 0 ? ldexpf(float(mantissa), -9)
+                                                                                                        : ldexpf(1.f + float(mantissa) / 8.f, exponent - 7);
+            return v & 128 ? -result : result;
+        }
+        // A CTA shares decoded weights across 32 activation rows. Each double-buffered
+        // stage covers 64 K values; W4A8 restores the codebook scale per 16-value group.
+
+        static __device__ __forceinline__ uint32_t perm_bytes(uint32_t lo, uint32_t hi, uint32_t sel) {
+#if defined(SD_CUDA_VENDOR_HIP)
+            return __builtin_amdgcn_perm(hi, lo, sel);
+#elif !defined(SD_CUDA_VENDOR_MUSA)
+            return __byte_perm(lo, hi, sel);
+#else
+            uint32_t r = 0;
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const uint32_t b    = (sel >> (4 * i)) & 0xF;
+                const uint32_t byte = (b < 4 ? lo : hi) >> (8 * (b & 3)) & 0xFF;
+                r |= byte << (8 * i);
+            }
+            return r;
+#endif
+        }
+
+#if defined(SD_CUDA_VENDOR_HIP)
+#define PERM_SEL_EVEN 0x06040200u
+#define PERM_SEL_ODD 0x07050301u
+#else
+#define PERM_SEL_EVEN 0x6420u
+#define PERM_SEL_ODD 0x7531u
+#endif
+
+        static __device__ __forceinline__ void interleave_act_pair(
+            const int lo,
+            const int hi,
+            int* __restrict__ dst) {
+            dst[0] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_EVEN);
+            dst[1] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_ODD);
+        }
+
+#define MUL_MAT_W4_BLOCK 16
+#define MUL_MAT_W4_ROWS (2 * MUL_MAT_W4_BLOCK)
+#define MUL_MAT_W4_GPS 4
+#define MUL_MAT_W4_SLOTS (4 * MUL_MAT_W4_GPS)
+#define MUL_MAT_W4_OUT (MUL_MAT_W4_BLOCK * 4)
+
+        static __global__ void __launch_bounds__(MUL_MAT_W4_BLOCK* MUL_MAT_W4_BLOCK) mul_mat_w4a4_convrot_dp4a(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            int* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t packed_k) {
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4_OUT;
+            const uint32_t r0     = blockIdx.y * MUL_MAT_W4_ROWS + threadIdx.y;
+            const uint32_t r1     = r0 + MUL_MAT_W4_BLOCK;
+            const uint32_t lx     = threadIdx.x;
+            const uint32_t ly     = threadIdx.y;
+            const uint32_t idx    = ly * MUL_MAT_W4_BLOCK + lx;
+            const uint32_t dc     = idx & (MUL_MAT_W4_OUT - 1);
+            const uint32_t dw     = idx >> 6;
+
+            __shared__ int w_dec[2][MUL_MAT_W4_SLOTS * MUL_MAT_W4_OUT];
+            __shared__ int x_tile[2][MUL_MAT_W4_ROWS * MUL_MAT_W4_SLOTS];
+
+            const uint32_t w_words = packed_k / 4;
+            const uint32_t x_words = (packed_k * 2) / 4;
+            const uint32_t groups  = packed_k / 8;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            const uint32_t fl = lx & 7;
+            const uint32_t fr = (lx >> 3) ? r1 : r0;
+            const uint32_t tr = (lx >> 3) * MUL_MAT_W4_BLOCK + ly;
+            int my_sum        = 0;
+
+            {
+                const uint32_t ocol = o_base + dc;
+                int2 wv             = make_int2(0, 0);
+                if (ocol < n && dw * 2 < w_words) {
+                    wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + dw];
+                }
+                w_dec[0][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808;
+                w_dec[0][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+                w_dec[0][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808;
+                w_dec[0][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+
+                const uint32_t gi = fl >> 1;
+                const uint32_t xb = gi * 4 + (fl & 1) * 2;
+                const int a       = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0;
+                const int b       = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0;
+                interleave_act_pair(a, b, &x_tile[0][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]);
+                my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0);
+            }
+            __syncthreads();
+
+            int acc0a = 0, acc0b = 0;
+            int acc1a = 0, acc1b = 0;
+            int acc2a = 0, acc2b = 0;
+            int acc3a = 0, acc3b = 0;
+            int acc4a = 0, acc4b = 0;
+            int acc5a = 0, acc5b = 0;
+            int acc6a = 0, acc6b = 0;
+            int acc7a = 0, acc7b = 0;
+
+            for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4_GPS) {
+                const uint32_t buf = (gs >> 2) & 1;
+
+                if (gs + MUL_MAT_W4_GPS < groups) {
+                    const uint32_t nbuf = buf ^ 1;
+                    const uint32_t ngs  = gs + MUL_MAT_W4_GPS;
+                    {
+                        const uint32_t ocol = o_base + dc;
+                        const uint32_t widx = ngs * 2 + dw * 2;
+                        int2 wv             = make_int2(0, 0);
+                        if (ocol < n && widx < w_words) {
+                            wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + ngs + dw];
+                        }
+                        w_dec[nbuf][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808;
+                        w_dec[nbuf][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+                        w_dec[nbuf][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808;
+                        w_dec[nbuf][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+                        const uint32_t gi                                         = ngs + (fl >> 1);
+                        const uint32_t xb                                         = gi * 4 + (fl & 1) * 2;
+                        const int a                                               = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0;
+                        const int b                                               = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0;
+                        interleave_act_pair(a, b, &x_tile[nbuf][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]);
+                        my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0);
+                    }
+                }
+#pragma unroll
+                for (uint32_t gl = 0; gl < MUL_MAT_W4_GPS; ++gl) {
+                    const int u0[4] = {
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 0],
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 1],
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 2],
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 3],
+                    };
+                    const int u1[4] = {
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 0],
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 1],
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 2],
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 3],
+                    };
+#pragma unroll
+                    for (uint32_t jp = 0; jp < 2; ++jp) {
+                        const int ue0     = jp == 0 ? u0[0] : u0[2];
+                        const int uo0     = jp == 0 ? u0[1] : u0[3];
+                        const int ue1     = jp == 0 ? u1[0] : u1[2];
+                        const int uo1     = jp == 0 ? u1[1] : u1[3];
+                        const uint32_t we = (gl * 4 + jp * 2) * MUL_MAT_W4_OUT;
+                        const uint32_t wo = we + MUL_MAT_W4_OUT;
+
+                        acc0a = dot4(w_dec[buf][we + lx], ue0, acc0a);
+                        acc1a = dot4(w_dec[buf][we + 16 + lx], ue0, acc1a);
+                        acc2a = dot4(w_dec[buf][we + 32 + lx], ue0, acc2a);
+                        acc3a = dot4(w_dec[buf][we + 48 + lx], ue0, acc3a);
+                        acc0b = dot4(w_dec[buf][wo + lx], uo0, acc0b);
+                        acc1b = dot4(w_dec[buf][wo + 16 + lx], uo0, acc1b);
+                        acc2b = dot4(w_dec[buf][wo + 32 + lx], uo0, acc2b);
+                        acc3b = dot4(w_dec[buf][wo + 48 + lx], uo0, acc3b);
+                        acc4a = dot4(w_dec[buf][we + lx], ue1, acc4a);
+                        acc5a = dot4(w_dec[buf][we + 16 + lx], ue1, acc5a);
+                        acc6a = dot4(w_dec[buf][we + 32 + lx], ue1, acc6a);
+                        acc7a = dot4(w_dec[buf][we + 48 + lx], ue1, acc7a);
+                        acc4b = dot4(w_dec[buf][wo + lx], uo1, acc4b);
+                        acc5b = dot4(w_dec[buf][wo + 16 + lx], uo1, acc5b);
+                        acc6b = dot4(w_dec[buf][wo + 32 + lx], uo1, acc6b);
+                        acc7b = dot4(w_dec[buf][wo + 48 + lx], uo1, acc7b);
+                    }
+                }
+                __syncthreads();
+            }
+
+            int v = my_sum;
+            v += shuffle_sum(v, (int(threadIdx.x) % 8 + 4) % 8, 8);
+            v += shuffle_sum(v, (int(threadIdx.x) % 8 + 2) % 8, 8);
+            v += shuffle_sum(v, (int(threadIdx.x) % 8 + 1) % 8, 8);
+            const int sum0 = shuffle_sum(v, 0, 16);
+            const int sum1 = shuffle_sum(v, 8, 16);
+
+            const int accv[2][4] = {
+                {acc0a + acc0b, acc1a + acc1b, acc2a + acc2b, acc3a + acc3b},
+                {acc4a + acc4b, acc5a + acc5b, acc6a + acc6b, acc7a + acc7b},
+            };
+            const int sumv[2] = {sum0, sum1};
+#pragma unroll
+            for (uint32_t t = 0; t < 2; ++t) {
+                const uint32_t rt = t == 0 ? r0 : r1;
+                if (rt < rows) {
+#pragma unroll
+                    for (uint32_t j = 0; j < 4; ++j) {
+                        const uint32_t oj = o_base + j * 16 + lx;
+                        if (oj < n) {
+                            y[(size_t)rt * n + oj] = accv[t][j] - 8 * sumv[t];
+                        }
+                    }
+                }
+            }
+        }
+
+#define MUL_MAT_W4A8_BLOCK 16
+#define MUL_MAT_W4A8_ROWS (2 * MUL_MAT_W4A8_BLOCK)
+#define MUL_MAT_W4A8_GPS 4
+#define MUL_MAT_W4A8_SLOTS (4 * MUL_MAT_W4A8_GPS)
+#define MUL_MAT_W4A8_OUT (MUL_MAT_W4A8_BLOCK * 4)
+
+#if defined(SD_CUDA_VENDOR_HIP) && defined(__HIP_DEVICE_COMPILE__)
+// __has_builtin checks the active device target's features, not just whether
+// the compiler recognizes the name. Prefer the non-replicated fragment layout.
+#if __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12)
+#define SD_W4_WMMA_GFX12
+#elif __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64)
+#define SD_W4_WMMA_GFX11
+#endif
+#endif
+
+// Keep each intrinsic paired with its fragment layout, including on newer GPUs.
+#if defined(SD_W4_WMMA_GFX11)
+        typedef int v4si32_w4a8 __attribute__((ext_vector_type(4)));
+
+        static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            const float* __restrict__ codebook,
+            const float* __restrict__ s_channel,
+            const uint8_t* __restrict__ s_rel,
+            const float* __restrict__ act_scales,
+            const float* __restrict__ bias,
+            float* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t k,
+            const uint32_t r_base) {
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
+            const uint32_t idx    = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x;
+            const uint32_t wv     = idx >> 6;
+            const uint32_t lane   = idx & 63;
+            const uint32_t j      = lane & 15;
+            const uint32_t q      = lane >> 4;
+            const uint32_t ocol   = o_base + wv * 16 + j;
+            const uint32_t ra     = r_base + j;
+
+            __shared__ __align__(4) int8_t cb8[16];
+            __shared__ float s_cb;
+            __shared__ int lut[256];
+            __shared__ float s_lut[256];
+
+            const uint32_t w_words = k / 8;
+            const uint32_t x_words = k / 4;
+            const uint32_t groups  = k / 16;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            if (idx == 0) {
+                float amax = 0.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    amax = fmaxf(amax, fabsf(codebook[i]));
+                }
+                s_cb = amax / 127.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb);
+                    cb8[i]      = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8);
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    s_lut[b] = fp8_e4m3(uint8_t(b));
+                }
+            }
+            __syncthreads();
+
+            const bool o_valid = ocol < n;
+            const float s_chan = o_valid ? s_channel[ocol] : 0.0f;
+            const float scb    = s_cb;
+
+            float acc[4]           = {0.0f, 0.0f, 0.0f, 0.0f};
+            const v4si32_w4a8 zero = {};
+
+            int w0c = 0, w1c = 0, xc0 = 0, xc1 = 0, xc2 = 0, xc3 = 0, sc = 0;
+            int w0n = 0, w1n = 0, xn0 = 0, xn1 = 0, xn2 = 0, xn3 = 0, sn = 0;
+            {
+                const int* xr = xi + (size_t)ra * x_words;
+                xc0           = ra < rows ? xr[0] : 0;
+                xc1           = ra < rows ? xr[1] : 0;
+                xc2           = ra < rows ? xr[2] : 0;
+                xc3           = ra < rows ? xr[3] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words;
+                    w0c           = wp[0];
+                    w1c           = wp[1];
+                    sc            = s_rel[(size_t)ocol * groups];
+                }
+            }
+            if (groups > 1) {
+                const int* xr = xi + (size_t)ra * x_words + 4;
+                xn0           = ra < rows ? xr[0] : 0;
+                xn1           = ra < rows ? xr[1] : 0;
+                xn2           = ra < rows ? xr[2] : 0;
+                xn3           = ra < rows ? xr[3] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words + 2;
+                    w0n           = wp[0];
+                    w1n           = wp[1];
+                    sn            = s_rel[(size_t)ocol * groups + 1];
+                }
+            }
+
+            for (uint32_t g = 0; g < groups; ++g) {
+                v4si32_w4a8 bfrag = {};
+                float s           = 0.0f;
+                if (o_valid) {
+                    bfrag = {
+                        (lut[w0c & 0xFF] & 0xFFFF) | ((lut[(w0c >> 8) & 0xFF] & 0xFFFF) << 16),
+                        (lut[(w0c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w0c >> 24) & 0xFF] & 0xFFFF) << 16),
+                        (lut[w1c & 0xFF] & 0xFFFF) | ((lut[(w1c >> 8) & 0xFF] & 0xFFFF) << 16),
+                        (lut[(w1c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w1c >> 24) & 0xFF] & 0xFFFF) << 16),
+                    };
+
+                    s = scb * (s_chan * s_lut[sc & 0xFF]);
+                }
+                const v4si32_w4a8 afrag = {xc0, xc1, xc2, xc3};
+                const v4si32_w4a8 dfrag =
+                    __builtin_amdgcn_wmma_i32_16x16x16_iu8_w64(true, afrag, true, bfrag, zero, false);
+
+                w0c = w0n;
+                w1c = w1n;
+                sc  = sn;
+                xc0 = xn0;
+                xc1 = xn1;
+                xc2 = xn2;
+                xc3 = xn3;
+                if (g + 2 < groups) {
+                    const int* xr = xi + (size_t)ra * x_words + (g + 2) * 4;
+                    xn0           = ra < rows ? xr[0] : 0;
+                    xn1           = ra < rows ? xr[1] : 0;
+                    xn2           = ra < rows ? xr[2] : 0;
+                    xn3           = ra < rows ? xr[3] : 0;
+                    if (o_valid) {
+                        const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2;
+                        w0n           = wp[0];
+                        w1n           = wp[1];
+                        sn            = s_rel[(size_t)ocol * groups + g + 2];
+                    }
+                }
+
+                acc[0] = fmaf((float)dfrag.x, s, acc[0]);
+                acc[1] = fmaf((float)dfrag.y, s, acc[1]);
+                acc[2] = fmaf((float)dfrag.z, s, acc[2]);
+                acc[3] = fmaf((float)dfrag.w, s, acc[3]);
+            }
+
+            if (o_valid) {
+#pragma unroll
+                for (uint32_t d = 0; d < 4; ++d) {
+                    const uint32_t r = r_base + 4 * d + q;
+                    if (r < rows) {
+                        float value = acc[d] * act_scales[r];
+                        if (bias != nullptr) {
+                            value += bias[ocol];
+                        }
+                        y[(size_t)r * n + ocol] = value;
+                    }
+                }
+            }
+        }
+#endif
+
+#if defined(SD_W4_WMMA_GFX12)
+        typedef int v4si32_w4a8_g12 __attribute__((ext_vector_type(4)));
+
+        static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma_gfx12(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            const float* __restrict__ codebook,
+            const float* __restrict__ s_channel,
+            const uint8_t* __restrict__ s_rel,
+            const float* __restrict__ act_scales,
+            const float* __restrict__ bias,
+            float* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t k,
+            const uint32_t r_base) {
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
+            const uint32_t idx    = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x;
+            const uint32_t wv     = idx >> 6;
+            const uint32_t lane   = idx & 63;
+            const uint32_t j      = lane & 15;
+            const uint32_t h      = lane >> 4;
+            const uint32_t kh     = ((h & 1) << 1) | (h >> 1);
+            const uint32_t ocol   = o_base + wv * 16 + j;
+            const uint32_t ra     = r_base + j;
+
+            __shared__ __align__(4) int8_t cb8[16];
+            __shared__ float s_cb;
+            __shared__ int lut[256];
+            __shared__ float s_lut[256];
+
+            const uint32_t w_words = k / 8;
+            const uint32_t x_words = k / 4;
+            const uint32_t groups  = k / 16;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            if (idx == 0) {
+                float amax = 0.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    amax = fmaxf(amax, fabsf(codebook[i]));
+                }
+                s_cb = amax / 127.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb);
+                    cb8[i]      = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8);
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    s_lut[b] = fp8_e4m3(uint8_t(b));
+                }
+            }
+            __syncthreads();
+
+            const bool o_valid = ocol < n;
+            const float s_chan = o_valid ? s_channel[ocol] : 0.0f;
+            const float scb    = s_cb;
+
+            float acc[4]               = {0.0f, 0.0f, 0.0f, 0.0f};
+            const v4si32_w4a8_g12 zero = {};
+
+            int w0c = 0, w1c = 0, xc = 0, sc = 0;
+            int w0n = 0, w1n = 0, xn = 0, sn = 0;
+            {
+                const int* xr = xi + (size_t)ra * x_words;
+                xc            = ra < rows ? xr[kh] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words;
+                    w0c           = wp[0];
+                    w1c           = wp[1];
+                    sc            = s_rel[(size_t)ocol * groups];
+                }
+            }
+            if (groups > 1) {
+                const int* xr = xi + (size_t)ra * x_words;
+                xn            = ra < rows ? xr[4 + kh] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words + 2;
+                    w0n           = wp[0];
+                    w1n           = wp[1];
+                    sn            = s_rel[(size_t)ocol * groups + 1];
+                }
+            }
+
+            for (uint32_t g = 0; g < groups; ++g) {
+                int bfrag = 0;
+                float s   = 0.0f;
+                if (o_valid) {
+                    const int bwd = (kh & 2) ? w1c : w0c;
+                    const int bsh = (kh & 1) * 16;
+                    bfrag         = (lut[(bwd >> bsh) & 0xFF] & 0xFFFF) |
+                                    ((lut[(bwd >> (bsh + 8)) & 0xFF] & 0xFFFF) << 16);
+
+                    s = scb * (s_chan * s_lut[sc & 0xFF]);
+                }
+                const v4si32_w4a8_g12 dfrag =
+                    __builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12(true, xc, true, bfrag, zero, false);
+
+                w0c = w0n;
+                w1c = w1n;
+                sc  = sn;
+                xc  = xn;
+                if (g + 2 < groups) {
+                    const int* xr = xi + (size_t)ra * x_words;
+                    xn            = ra < rows ? xr[(g + 2) * 4 + kh] : 0;
+                    if (o_valid) {
+                        const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2;
+                        w0n           = wp[0];
+                        w1n           = wp[1];
+                        sn            = s_rel[(size_t)ocol * groups + g + 2];
+                    }
+                }
+
+                acc[0] = fmaf((float)dfrag.x, s, acc[0]);
+                acc[1] = fmaf((float)dfrag.y, s, acc[1]);
+                acc[2] = fmaf((float)dfrag.z, s, acc[2]);
+                acc[3] = fmaf((float)dfrag.w, s, acc[3]);
+            }
+
+            if (o_valid) {
+#pragma unroll
+                for (uint32_t d = 0; d < 4; ++d) {
+                    const uint32_t r = r_base + 4 * kh + d;
+                    if (r < rows) {
+                        float value = acc[d] * act_scales[r];
+                        if (bias != nullptr) {
+                            value += bias[ocol];
+                        }
+                        y[(size_t)r * n + ocol] = value;
+                    }
+                }
+            }
+        }
+#endif
+
+        static __global__ void __launch_bounds__(MUL_MAT_W4A8_BLOCK* MUL_MAT_W4A8_BLOCK) mul_mat_w4a8_convrot(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            const float* __restrict__ codebook,
+            const float* __restrict__ s_channel,
+            const uint8_t* __restrict__ s_rel,
+            const float* __restrict__ act_scales,
+            const float* __restrict__ bias,
+            float* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t k) {
+#if defined(SD_W4_WMMA_GFX11) || defined(SD_W4_WMMA_GFX12)
+            // Retain the 32-row launch grid shared with DP4A. Each WMMA wave
+            // computes a 16-row tile; the barrier protects the reused LUTs.
+            for (uint32_t tile = 0; tile < 2; ++tile) {
+#if defined(SD_W4_WMMA_GFX12)
+                mul_mat_w4a8_convrot_wmma_gfx12(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k,
+                                                blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK);
+#else
+                mul_mat_w4a8_convrot_wmma(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k,
+                                          blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK);
+#endif
+                __syncthreads();
+            }
+#else
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
+            const uint32_t r0     = blockIdx.y * MUL_MAT_W4A8_ROWS + threadIdx.y;
+            const uint32_t r1     = r0 + MUL_MAT_W4A8_BLOCK;
+            const uint32_t lx     = threadIdx.x;
+            const uint32_t ly     = threadIdx.y;
+            const uint32_t idx    = ly * MUL_MAT_W4A8_BLOCK + lx;
+            const uint32_t dc     = idx & (MUL_MAT_W4A8_OUT - 1);
+            const uint32_t dw     = idx >> 6;
+
+            __shared__ int w_dec[2][MUL_MAT_W4A8_SLOTS * MUL_MAT_W4A8_OUT];
+            __shared__ int x_tile[2][MUL_MAT_W4A8_ROWS * MUL_MAT_W4A8_SLOTS];
+            __shared__ float s_dec[2][MUL_MAT_W4A8_GPS * MUL_MAT_W4A8_OUT];
+            __shared__ __align__(4) int8_t cb8[16];
+            __shared__ float s_cb;
+            __shared__ int lut[256];
+
+            const uint32_t w_words = k / 8;
+            const uint32_t x_words = k / 4;
+            const uint32_t groups  = k / 16;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            if (ly == 0 && lx == 0) {
+                float amax = 0.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    amax = fmaxf(amax, fabsf(codebook[i]));
+                }
+
+                s_cb = amax / 127.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int v = s_cb == 0.f ? 0 : __float2int_rn(codebook[i] / s_cb);
+                    cb8[i]      = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
+                }
+            }
+            __syncthreads();
+
+            lut[idx] = (cb8[idx & 0xF] & 0xFF) | ((cb8[idx >> 4] & 0xFF) << 8);
+            __syncthreads();
+
+#pragma unroll
+            for (uint32_t j = 0; j < 2; ++j) {
+                const uint32_t pw                              = dw * 2 + j;
+                const uint32_t ocol                            = o_base + dc;
+                const int wv                                   = (ocol < n && pw < w_words) ? wi[(size_t)ocol * w_words + pw] : 0;
+                w_dec[0][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16);
+                w_dec[0][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16);
+            }
+            {
+                const uint32_t ocol = o_base + dc;
+                float s             = 0.0f;
+                if (ocol < n && dw < groups) {
+                    s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + dw]);
+                }
+                s_dec[0][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s;
+            }
+            if (lx < 16) {
+                const uint32_t aidx = lx;
+                x_tile[0][ly * MUL_MAT_W4A8_SLOTS + lx] =
+                    (r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0;
+                x_tile[0][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] =
+                    (r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0;
+            }
+            __syncthreads();
+
+            float acc0 = 0.0f;
+            float acc1 = 0.0f;
+            float acc2 = 0.0f;
+            float acc3 = 0.0f;
+            float acc4 = 0.0f;
+            float acc5 = 0.0f;
+            float acc6 = 0.0f;
+            float acc7 = 0.0f;
+
+            for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4A8_GPS) {
+                const uint32_t buf = (gs >> 2) & 1;
+
+                if (gs + MUL_MAT_W4A8_GPS < groups) {
+                    const uint32_t nbuf = buf ^ 1;
+                    const uint32_t ngs  = gs + MUL_MAT_W4A8_GPS;
+#pragma unroll
+                    for (uint32_t j = 0; j < 2; ++j) {
+                        const uint32_t pw                                 = dw * 2 + j;
+                        const uint32_t ocol                               = o_base + dc;
+                        const uint32_t widx                               = ngs * 2 + pw;
+                        const int wv                                      = (ocol < n && widx < w_words) ? wi[(size_t)ocol * w_words + widx] : 0;
+                        w_dec[nbuf][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16);
+                        w_dec[nbuf][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16);
+                    }
+                    {
+                        const uint32_t ocol = o_base + dc;
+                        float s             = 0.0f;
+                        if (ocol < n && ngs + dw < groups) {
+                            s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + ngs + dw]);
+                        }
+                        s_dec[nbuf][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s;
+                    }
+                    if (lx < 16) {
+                        const uint32_t aidx = ngs * 4 + lx;
+                        x_tile[nbuf][ly * MUL_MAT_W4A8_SLOTS + lx] =
+                            (r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0;
+                        x_tile[nbuf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] =
+                            (r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0;
+                    }
+                }
+
+#pragma unroll
+                for (uint32_t gl = 0; gl < MUL_MAT_W4A8_GPS; ++gl) {
+                    const int u0[4] = {
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 0],
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 1],
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 2],
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 3],
+                    };
+                    const int u1[4] = {
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 0],
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 1],
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 2],
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 3],
+                    };
+                    int accg0 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u0[0], 0);
+                    accg0     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u0[1], accg0);
+                    accg0     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u0[2], accg0);
+                    accg0     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u0[3], accg0);
+                    int accg1 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u0[0], 0);
+                    accg1     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u0[1], accg1);
+                    accg1     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u0[2], accg1);
+                    accg1     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u0[3], accg1);
+                    int accg2 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u0[0], 0);
+                    accg2     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u0[1], accg2);
+                    accg2     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u0[2], accg2);
+                    accg2     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u0[3], accg2);
+                    int accg3 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u0[0], 0);
+                    accg3     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u0[1], accg3);
+                    accg3     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u0[2], accg3);
+                    accg3     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u0[3], accg3);
+                    int accg4 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u1[0], 0);
+                    accg4     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u1[1], accg4);
+                    accg4     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u1[2], accg4);
+                    accg4     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u1[3], accg4);
+                    int accg5 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u1[0], 0);
+                    accg5     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u1[1], accg5);
+                    accg5     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u1[2], accg5);
+                    accg5     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u1[3], accg5);
+                    int accg6 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u1[0], 0);
+                    accg6     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u1[1], accg6);
+                    accg6     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u1[2], accg6);
+                    accg6     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u1[3], accg6);
+                    int accg7 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u1[0], 0);
+                    accg7     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u1[1], accg7);
+                    accg7     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u1[2], accg7);
+                    accg7     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u1[3], accg7);
+                    acc0      = fmaf((float)accg0, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc0);
+                    acc1      = fmaf((float)accg1, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc1);
+                    acc2      = fmaf((float)accg2, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc2);
+                    acc3      = fmaf((float)accg3, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc3);
+                    acc4      = fmaf((float)accg4, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc4);
+                    acc5      = fmaf((float)accg5, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc5);
+                    acc6      = fmaf((float)accg6, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc6);
+                    acc7      = fmaf((float)accg7, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc7);
+                }
+                __syncthreads();
+            }
+
+            const float accv[8] = {acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7};
+#pragma unroll
+            for (uint32_t t = 0; t < 2; ++t) {
+                const uint32_t rt = t == 0 ? r0 : r1;
+                if (rt < rows) {
+#pragma unroll
+                    for (uint32_t j = 0; j < 4; ++j) {
+                        const uint32_t oj = o_base + j * 16 + lx;
+                        if (oj < n) {
+                            float value = accv[t * 4 + j] * act_scales[rt];
+                            if (bias != nullptr) {
+                                value += bias[oj];
+                            }
+                            y[(size_t)rt * n + oj] = value;
+                        }
+                    }
+                }
+            }
+#endif
+        }
+
+        static __global__ void dequantize(const int32_t* accum, const float* xs, const float* ws, const float* bias, float* dst, int64_t n, int64_t total) {
+            const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            if (i < total) {
+                float value = float(accum[i]) * xs[i / n] * ws[i % n];
+                dst[i]      = value + (bias ? bias[i % n] : 0.f);
+            }
+        }
+        static bool aligned(const ggml_tensor* t, size_t alignment) {
+            if (!t)
+                return true;
+            size_t offset = 0;
+            while (!t->data && t->view_src) {
+                offset += t->view_offs;
+                t = t->view_src;
+            }
+            return (reinterpret_cast<uintptr_t>(t->data) + offset) % alignment == 0;
+        }
+        static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+            SDW4Operation op;
+            if (!sd_ggml_w4_describe(node, &op) || cuda_device_info(device).device < 0)
+                return false;
+            const int64_t k = op.input->ne[0], n = node->ne[0], rows = ggml_nrows(node);
+            if (k % 64 || k > INT_MAX || n > INT_MAX - 63 || rows > 65535 * 32 ||
+                ggml_nelements(node) > int64_t(INT_MAX) * 256 || !aligned(op.weight, 8) || !aligned(node, 4))
+                return false;
+            for (auto* t : node->src)
+                if (!aligned(t, 4))
+                    return false;
+            return true;
+        }
+        struct Plan {
+            bool mma           = false;
+            size_t gemm_offset = 0;
+#ifdef SD_CUDA_W4_CUTLASS
+            SDI8CudaGemmPlan* gemm = nullptr;
+            ~Plan() { sd_cuda_i8_gemm_free(gemm); }
+#endif
+        };
+        static bool use_mma(ggml_backend_t backend, const SDW4Operation& op) {
+#ifdef SD_CUDA_W4_CUTLASS
+            if (!aligned(op.input, 16) || !aligned(op.weight, 16))
+                return false;
+            static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_W4_FORCE_DP4A"); return v && std::atoi(v); }();
+            if (disabled)
+                return false;
+            static std::mutex mutex;
+            static std::unordered_map<int, bool> devices;
+            const auto device = cuda_device_info(ggml_backend_get_device(backend)).device;
+            std::lock_guard<std::mutex> lock(mutex);
+            auto found = devices.find(device);
+            if (found != devices.end())
+                return found->second;
+            cudaDeviceProp prop  = {};
+            const bool available = cudaGetDeviceProperties(&prop, device) == cudaSuccess && prop.major >= 8;
+            devices.emplace(device, available);
+            return available;
+#else
+            return false;
+#endif
+        }
+        static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int, void** state, size_t* bytes, void*) {
+            SDW4Operation op;
+            if (!sd_ggml_w4_describe(nodes[0], &op))
+                return GGML_STATUS_FAILED;
+            const auto* dst = nodes[0];
+            bool may_alias  = dst->data || dst->view_src;
+            if (dst->data && dst->buffer && !dst->view_src) {
+                may_alias = false;
+                for (const auto* input : dst->src) {
+                    if (!input)
+                        continue;
+                    size_t offset     = 0;
+                    const size_t size = ggml_nbytes(input);
+                    while (!input->data && input->view_src) {
+                        offset += input->view_offs;
+                        input = input->view_src;
+                    }
+                    if (!input->data)
+                        continue;
+                    const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(input->data) + offset;
+                    may_alias |= !input->buffer || (input->buffer == dst->buffer && a < b + size && b < a + ggml_nbytes(dst));
+                }
+            }
+            auto* plan        = new Plan;
+            plan->mma         = use_mma(backend, op);
+            *state            = plan;
+            size_t gemm_bytes = 0;
+#ifdef SD_CUDA_W4_CUTLASS
+            if (plan->mma && op.kind == SDW4Kind::W4A4) {
+                if (!op.bias && aligned(dst, 16) && aligned(op.weight_scale, 16))
+                    plan->gemm = sd_cuda_i8_gemm_prepare(ggml_backend_get_device(backend), ggml_nrows(dst), dst->ne[0], op.input->ne[0], false, &gemm_bytes, true);
+                plan->mma = plan->gemm != nullptr;
+                if (plan->gemm && gemm_bytes == 0)
+                    gemm_bytes = 256;
+            }
+#endif
+            const size_t tensor_bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 && !plan->mma ? 1 : 0) + (may_alias ? 1 : 0));
+            plan->gemm_offset         = (tensor_bytes + 255) / 256 * 256;
+            *bytes                    = plan->gemm_offset + gemm_bytes;
+            return GGML_STATUS_SUCCESS;
+        }
+        static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
+            if (!src)
+                return false;
+            const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(src->data);
+            return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src);
+        }
+        static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void* state, void*) {
+            auto* dst = nodes[0];
+            SDW4Operation op;
+            ggml_ext_gpu_context context;
+            if (!sd_ggml_w4_describe(dst, &op) || !cuda_context(e->backend, &context))
+                return GGML_STATUS_FAILED;
+            const auto stream = static_cast<CUDAStream>(context.stream);
+            const int64_t n = dst->ne[0], rows = ggml_nrows(dst), k = op.input->ne[0];
+            auto* scratch = e->workspace ? static_cast<float*>(ggml_backend_buffer_get_base(e->workspace)) : nullptr;
+            bool alias    = false;
+            for (auto* input : dst->src)
+                alias |= overlaps(dst, input);
+            auto* output     = alias ? scratch : static_cast<float*>(dst->data);
+            const auto* w    = static_cast<const char*>(op.weight->data);
+            const auto* x    = static_cast<const char*>(op.input->data);
+            const auto* xs   = reinterpret_cast<const float*>(x + k * ((rows + 3) / 4 * 4));
+            const auto* ws   = static_cast<const float*>(op.weight_scale->data);
+            const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+            const dim3 grid(unsigned((n + 63) / 64), unsigned((rows + 31) / 32)), block(16, 16);
+#ifdef SD_CUDA_W4_CUTLASS
+            auto* plan = static_cast<Plan*>(state);
+            if (plan->mma) {
+                if (plan->gemm) {
+                    auto* workspace = scratch ? reinterpret_cast<char*>(scratch) + plan->gemm_offset : nullptr;
+                    if (sd_cuda_i8_gemm_execute(plan->gemm, reinterpret_cast<const int8_t*>(x), reinterpret_cast<const int8_t*>(w), xs, ws, output, workspace, stream) != GGML_STATUS_SUCCESS)
+                        return GGML_STATUS_FAILED;
+                } else {
+                    const dim3 tiles(unsigned((n + 63) / 64), unsigned((rows + 127) / 128));
+                    sd_w4_mma::w4a8<128, 64, 128><<<tiles, 256, 0, stream>>>(reinterpret_cast<const uint8_t*>(w), reinterpret_cast<const int8_t*>(x), xs, ws,
+                                                                             static_cast<const float*>(op.s_channel->data), static_cast<const uint8_t*>(op.s_rel->data), bias, output, int(rows), int(n), int(k));
+                }
+            } else
+#endif
+                if (op.kind == SDW4Kind::W4A4) {
+                auto* accum = reinterpret_cast<int32_t*>(scratch + (alias ? ggml_nelements(dst) : 0));
+                mul_mat_w4a4_convrot_dp4a<<<grid, block, 0, stream>>>(w, x, accum, uint32_t(n), uint32_t(rows), uint32_t(k / 2));
+                if (cuda_last_error() != cuda_success)
+                    return GGML_STATUS_FAILED;
+                dequantize<<<unsigned((n * rows + 255) / 256), 256, 0, stream>>>(accum, xs, ws, bias, output, n, n * rows);
+            } else {
+                mul_mat_w4a8_convrot<<<grid, block, 0, stream>>>(w, x, ws,
+                                                                 static_cast<const float*>(op.s_channel->data), static_cast<const uint8_t*>(op.s_rel->data),
+                                                                 xs, bias, output, uint32_t(n), uint32_t(rows), uint32_t(k));
+            }
+            auto error = cuda_last_error();
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
+    }
+    void register_w4(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc d = {};
+        d.abi_version          = GGML_EXT_ABI_VERSION;
+        d.name                 = "sd.linear_w4_convrot." SD_CUDA_NAME ".v1";
+        d.op                   = GGML_OP_CUSTOM;
+        d.custom_op            = sd_ggml_w4_op();
+        d.priority             = 100;
+        d.supports             = sd_w4_cuda::supports;
+        d.prepare              = sd_w4_cuda::prepare;
+        d.execute              = sd_w4_cuda::execute;
+        d.destroy              = [](void* state, void*) { delete static_cast<sd_w4_cuda::Plan*>(state); };
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &d, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+}

src/core/cuda/w4_mma.h

@@ -0,0 +1,138 @@
+#ifndef __SD_CORE_CUDA_W4_MMA_H__
+#define __SD_CORE_CUDA_W4_MMA_H__
+
+#include <cuda_runtime.h>
+#include <cutlass/arch/mma_sm80.h>
+#include <cstdint>
+
+namespace sd_w4_mma {
+    static __device__ __forceinline__ float fp8_e4m3(uint8_t v) {
+        const int exponent = (v >> 3) & 15, mantissa = v & 7;
+        const float result = exponent == 15 && mantissa == 7 ? __int_as_float(0x7fc00000) : exponent == 0 ? ldexpf(float(mantissa), -9)
+                                                                                                          : ldexpf(1.f + float(mantissa) / 8.f, exponent - 7);
+        return v & 128 ? -result : result;
+    }
+
+    template <int M, int N, int K>
+    static __global__ void __launch_bounds__(256) w4a8(
+        const uint8_t* __restrict__ w,
+        const int8_t* __restrict__ x,
+        const float* __restrict__ xs,
+        const float* __restrict__ codebook,
+        const float* __restrict__ channel,
+        const uint8_t* __restrict__ relative,
+        const float* __restrict__ bias,
+        float* __restrict__ dst,
+        int rows,
+        int n,
+        int k) {
+        const int row_base = blockIdx.y * M, col_base = blockIdx.x * N;
+        const int tid = threadIdx.x, lane = tid % 32, warp = tid / 32;
+        __shared__ int cb[16];
+        __shared__ float cbscale;
+        if (tid == 0) {
+            float amax = 0.f;
+#pragma unroll
+            for (int i = 0; i < 16; ++i)
+                amax = fmaxf(amax, fabsf(codebook[i]));
+            cbscale = amax / 127.f;
+#pragma unroll
+            for (int i = 0; i < 16; ++i) {
+                const int value = cbscale == 0.f ? 0 : __float2int_rn(codebook[i] / cbscale);
+                cb[i]           = max(-127, min(127, value));
+            }
+        }
+        __syncthreads();
+#if __CUDA_ARCH__ >= 800
+        constexpr int stride = K / 4 + 4;
+        __shared__ __align__(16) uint32_t a[M][stride], b[N][stride];
+        __shared__ uint32_t lut[256];
+        __shared__ float scales[K / 16][N];
+        lut[tid] = (cb[tid & 15] & 255) | ((cb[tid >> 4] & 255) << 8);
+        __syncthreads();
+        constexpr int warp_cols = N / 32;
+        constexpr int row_tiles = M / (16 * (8 / warp_cols));
+        const int wr = (warp / warp_cols) * 16 * row_tiles, wc = (warp % warp_cols) * 32;
+        float acc[row_tiles][4][4] = {};
+        for (int base = 0; base < k; base += K) {
+            for (int i = tid; i < M * K / 16; i += 256) {
+                const int r = i / (K / 16), t = (i % (K / 16)) * 16;
+                const int4 value                       = row_base + r < rows && base + t < k ? *reinterpret_cast<const int4*>(x + int64_t(row_base + r) * k + base + t) : make_int4(0, 0, 0, 0);
+                *reinterpret_cast<int4*>(&a[r][t / 4]) = value;
+            }
+            for (int i = tid; i < N * K / 8; i += 256) {
+                const int c = i / (K / 8), t = i % (K / 8);
+                const uint32_t value = col_base + c < n && base + t * 8 < k ? *reinterpret_cast<const uint32_t*>(w + int64_t(col_base + c) * (k / 2) + base / 2 + t * 4) : 0;
+                b[c][2 * t]          = lut[value & 255] | (lut[(value >> 8) & 255] << 16);
+                b[c][2 * t + 1]      = lut[(value >> 16) & 255] | (lut[value >> 24] << 16);
+            }
+            for (int i = tid; i < N * K / 16; i += 256) {
+                const int c = i % N, g = i / N;
+                scales[g][c] = col_base + c < n && base + g * 16 < k ? cbscale * (channel[col_base + c] * fp8_e4m3(relative[int64_t(col_base + c) * (k / 16) + base / 16 + g])) : 0.f;
+            }
+            __syncthreads();
+#pragma unroll
+            for (int g = 0; g < K / 16; ++g) {
+#pragma unroll
+                for (int r = 0; r < row_tiles; ++r) {
+                    const int ar = wr + r * 16 + lane / 4, ac = g * 4 + lane % 4;
+                    const uint32_t a0 = a[ar][ac], a1 = a[ar + 8][ac];
+#pragma unroll
+                    for (int j = 0; j < 4; ++j) {
+                        const uint32_t b0 = b[wc + j * 8 + lane / 4][ac];
+                        using Mma         = cutlass::arch::Mma<cutlass::gemm::GemmShape<16, 8, 16>, 32,
+                                                               int8_t, cutlass::layout::RowMajor, int8_t, cutlass::layout::ColumnMajor,
+                                                               int, cutlass::layout::RowMajor, cutlass::arch::OpMultiplyAddSaturate>;
+                        Mma::FragmentA af;
+                        Mma::FragmentB bf;
+                        Mma::FragmentC d;
+                        reinterpret_cast<uint32_t*>(&af)[0] = a0;
+                        reinterpret_cast<uint32_t*>(&af)[1] = a1;
+                        reinterpret_cast<uint32_t*>(&bf)[0] = b0;
+                        d.clear();
+                        // Each K=16 MMA has its own scale; combining groups before
+                        // dequantization would change the W4A8 representation.
+                        Mma{}(d, af, bf, d);
+                        const float s0 = scales[g][wc + j * 8 + (lane % 4) * 2], s1 = scales[g][wc + j * 8 + (lane % 4) * 2 + 1];
+#pragma unroll
+                        for (int v = 0; v < 4; ++v)
+                            acc[r][j][v] = fmaf(float(d[v]), v % 2 ? s1 : s0, acc[r][j][v]);
+                    }
+                }
+            }
+            __syncthreads();
+        }
+#pragma unroll
+        for (int r = 0; r < row_tiles; ++r) {
+#pragma unroll
+            for (int j = 0; j < 4; ++j) {
+#pragma unroll
+                for (int v = 0; v < 4; ++v) {
+                    const int row = row_base + wr + r * 16 + lane / 4 + (v / 2) * 8;
+                    const int col = col_base + wc + j * 8 + (lane % 4) * 2 + v % 2;
+                    if (row < rows && col < n)
+                        dst[int64_t(row) * n + col] = __fmul_rn(acc[r][j][v], xs[row]) + (bias ? bias[col] : 0.f);
+                }
+            }
+        }
+#else
+        for (int i = tid; i < M * N; i += 256) {
+            const int row = row_base + i / N, col = col_base + i % N;
+            if (row < rows && col < n) {
+                float sum = 0.f;
+                for (int g = 0; g < k / 16; ++g) {
+                    int d = 0;
+                    for (int j = 0; j < 16; ++j) {
+                        const int t = g * 16 + j;
+                        d += cb[(w[int64_t(col) * (k / 2) + t / 2] >> (4 * (t % 2))) & 15] * int(x[int64_t(row) * k + t]);
+                    }
+                    sum = fmaf(float(d), cbscale * (channel[col] * fp8_e4m3(relative[int64_t(col) * (k / 16) + g])), sum);
+                }
+                dst[int64_t(row) * n + col] = __fmul_rn(sum, xs[row]) + (bias ? bias[col] : 0.f);
+            }
+        }
+#endif
+    }
+
+}
+#endif